rwkv-quant: una sola matriz en bf16 elimina el 39% de las pérdidas al cuantizar RWKV-7 2.9B

La cuantización de los modelos de la familia RWKV, con su estructura recurrente lineal, está menos estudiada que la de los transformers. El autor de la herramienta rwkv-quant, que comprime RWKV-7 con dos presets (aproximadamente 2.1x y 3.1x), detectó una anomalía. El preset agresivo COMPRESSION da +4.04% de perplejidad en 2.9B, pero solo +2.85% en 1.5B. El análisis de las causas mostró que buena parte de las pérdidas se concentra en una única matriz de 2560×2560. Guardarla en bf16 cuesta el 0.5% del archivo y elimina el 39% del KL.

Dos presets y seis tamaños

El estudio abarca la serie RWKV-7 en seis tamaños. Tres pares permiten separar el efecto del ancho y el de la profundidad: 0.4B/1.5B (misma profundidad, el doble de ancho), 2.9B/7.2B (profundidad 32, ancho 2560 frente a 4096) y 7.2B/13.3B (ancho 4096, profundidad 32 frente a 61).

• 0.1B: 12 capas, ancho 768
• 0.4B: 24 capas, ancho 1024
• 1.5B: 24 capas, ancho 2048
• 2.9B: 32 capas, ancho 2560
• 7.2B: 32 capas, ancho 4096
• 13.3B: 61 capas, ancho 4096

El preset REDUCTION apenas introduce pérdidas: proyecciones en 8 bits con disposición simétrica y grupo de 16 pesos, cmix (FFN) en 6 bits, embedding y cabeza en 8 bits. COMPRESSION es más agresivo: proyecciones en 4 bits con disposición asimétrica por superbloques y grupo de 32, cmix en 4 bits (ffn.key en 5), cabeza en 5 bits y embedding en 6 bits. Las ramas LoRA van en 6 bits en ambos presets (g en 8). En los dos casos, la escala y el desplazamiento de los grupos se eligen mediante una búsqueda ponderada por el cuadrado medio de la activación de entrada del canal, E[x²].

Dónde se rompe el 2.9B

La descomposición del KL por grupos de matrices mediante leave-one-out mostró que, en el 2.9B, o_proj aporta el 46% del KL total y el embedding el 17%. En el 1.5B esas proporciones son del 16% y el 3%, y la mayor contribución corresponde a cmix.value (37%). Dentro de o_proj, entre el 85% y el 86% del daño recae en una sola matriz: la proyección de salida del bloque cero.

Un experimento directo en el 2.9B (ruta real, MLX) lo confirmó:

• preset base: KL 0.052048, +4.04% de ppl;
• o_proj de la capa 0 en bf16: KL 0.031966, +2.43% de ppl, archivo 9.4 MB mayor;
• o_proj de la capa 31 en bf16: sin efecto (+4.07%);
• todo o_proj en 5 bits: +3.54%, +26.2 MB;
• todo o_proj en 6 bits: +2.89%, +52.4 MB.

Una sola matriz en bf16 funciona mejor que pasar todo o_proj a 6 bits y cuesta más de cinco veces menos. En otros tamaños el efecto es más débil: el KL baja un 7.6% en 0.1B, un 1.5% en 0.4B, un 0.4% en 1.5B y un 3.9% en 7.2B. Solo son significativos los de 0.1B y 2.9B.

La entrada de esta matriz lleva un «canal masivo». En el 2.9B, ocho canales concentran el 98% de la energía, y el número de canal 517 coincide en dos series de checkpoints; en el 1.5B es el canal 1202. Es decir, lo determinan el ancho y la arquitectura, no un entrenamiento concreto. Pero la fuerza del canal no predice el daño: el 7.2B tiene la mayor concentración (99%) y, sin embargo, el KL base más bajo (0.023 frente a 0.078 del 2.9B). Sin ponderación AW, 4 bits en o_proj destruyen el 2.9B: KL 0.340, +37.2% de ppl.

El autor muestra que tres métricas locales de error naturales no predicen el daño. Lo que sí funcionó fue el error en el flujo residual multiplicado por la ganancia medida.

Fuga de calibración y mediciones fiables

El primer conjunto de evaluación (38 ventanas de 512 tokens, en/ru/sr) estaba íntegramente dentro del corpus de calibración: el 100% de los 16-gramas eran comunes. La magnitud de la fuga se midió con GPTQ. En el 1.5B, la mejora de KL fue de −61.8% en el conjunto antiguo y de −27.1% en el limpio, es decir, la fuga inflaba el resultado 2.3 veces. La selección automática de bits no ganó en él: −16.1% frente a −17.6%.

Por eso, las secciones 7–10 del trabajo usan solo conjuntos limpios: Wikipedia en/ru/sr (36 ventanas), un conjunto cerrado de código Python/Swift (24 ventanas) y uno abierto (48 ventanas). Los intervalos se calculan con bootstrap pareado por ventanas, 20 000 repeticiones; la semiamplitud del intervalo del 95% para el KL es del 7–12% del valor.

Además, el trabajo describe la selección automática de bits con la regla «daño por byte», que funciona en ambos sentidos, y GPTQ sobre la cuadrícula del preset sin cambiar el formato, los kernels ni el tamaño del archivo. Estos métodos se combinan de forma casi multiplicativa. Aparte se incluye una comparación con GGUF y MLX a igual tamaño de archivo.

La conclusión práctica es sencilla. Si cuantiza RWKV-7 para ejecutarlo en local, sobre todo en Mac con MLX, mantenga o_proj de la primera capa en bf16: son 6–10 MB en modelos de 1.5B a 2.9B. El embedding conviene dejarlo como mínimo en 6 bits. Antes de fiarse de las cifras de un método con calibración, compruebe que la intersección de 16-gramas entre el texto de calibración y el de evaluación es nula. Es mejor medir la contribución de cada grupo de matrices con leave-one-out que suponer que, en un modelo grande, las pérdidas se reparten de manera uniforme.