◧ Base de Conocimiento ANE
EN ES

Capítulos · 07

Compresión de modelos para la ANE: palettization, recetas mixed-bit, outlier decomposition

Parte de la Base de conocimiento sobre la ANE. Fuente: argmaxtools.compress (references/argmaxtools-0.1.23-pypi-snapshot/argmaxtools/compress/), tal como la usa el pipeline --generate-quantized-variants de whisperkittools. Complementa el caso de estudio de WhisperKit (doc 06).

1. Por qué la compresión es una optimización de latencia en la ANE#

Principio 4 de Apple (doc 02): la inferencia de transformers con lotes pequeños en la ANE está limitada por el ancho de banda — la latencia está dominada por la carga de los pesos, no por la aritmética. Reducir el tamaño de los pesos, por tanto, reduce la latencia, no solo el tamaño de la app. La decodificación con KV cache (un token por llamada) es el caso extremo. Por eso la suite de pruebas de argmaxtools afirma que un modelo palettized a 1 bit debe ejecutarse a ≥ 0.95× la velocidad de FP16 (TEST_COMPRESSION_MIN_SPEEDUP) — la compresión debería ser gratuita o mejor en tiempo de ejecución.

2. Palettization (la técnica nativa de Core ML)#

Palettization = agrupar mediante k-means cada tensor de pesos en una tabla de búsqueda (LUT) de 2^nbits entradas y almacenar índices por peso. Anchos de bits admitidos: {1, 2, 3, 4, 6, 8} (SUPPORTED_NBITS); las activaciones permanecen en FP16 — solo se reduce el almacenamiento de los pesos, la descompresión la gestiona el hardware/runtime. Se aplica post-entrenamiento sobre el modelo convertido:

config = ct.optimize.coreml.OptimizationConfig(op_name_configs={
    op_name: ct.optimize.coreml.OpPalettizerConfig(
        mode="kmeans", nbits=nbits,
        granularity="per_tensor",          # or "per_grouped_channel" + group_size
    ) ...
})
mlmodel = ct.optimize.coreml.palettize_weights(mlmodel, config=config)

Restricciones de SO (afirmadas en el código): la palettization a 3 bits y la granularidad per_grouped_channel (tamaños de grupo 4–256) requieren iOS 18 / macOS 15; la CLI de whisperkittools expone --palettization-group-size {4,16,32,64,128,256}.

No todo se palletiza (_get_compressible_modules, _find_nbits_in_recipe):

  • tensores con numel < 1e5 (insignificantes),
  • tensores con sparsity > 0.8 (mejor atendidos por pruning),
  • tensores no FP16,
  • opcionalmente, las top-K capas más sensibles (DONT_PALETTIZE_TOP_K = 3) permanecen en FP16.

3. Recetas mixed-bit: perfilado de sensibilidad por capa#

La compresión uniforme a pocos bits degrada la calidad de forma desigual — algunas capas toleran 2 bits, otras se rompen por debajo de 8. La clase Palettizer (compress/palettize.py) automatiza la búsqueda de una asignación de bits por capa ("receta"):

La maquinaria#

  1. Fake palettization (_fake_palettize): ejecuta k-means exactamente como lo haría coremltools, pero escribe los valores decuantizados de vuelta en el tensor FP16 de torch — simulando la compresión mientras permanece en PyTorch para una evaluación rápida. Los resultados se cachean en disco por (layer, nbits) para reinicios en caliente.
  2. Métrica de divergencia: una función abstracta divergence_fn(reference, proxy) compara las salidas end-to-end del modelo contra la referencia sin comprimir sobre un lote de prueba reservado (TEST_BATCH_SIZE = 32). Las subclases específicas de cada modelo definen qué significa "salida" (por ejemplo, los logits del decodificador).
  3. Respuesta por capa (profile_per_layer_response): palletiza una capa a la vez en cada nbits; registra la divergencia. Este es el mapa de sensibilidad.
  4. Comprobación de sanidad (_sanity_check_per_layer_results): la divergencia debe disminuir con más bits; si > 10 % de los resultados de (capa, par de bits) están invertidos, el lote de prueba es demasiado pequeño — aborta.
  5. Generación de recetas (profile_mixed_bit_response): barre umbrales de divergencia con np.geomspace; para cada umbral asigna a cada capa el menor nbits cuya divergencia por capa esté por debajo de él. Cada receta se indexa por su precisión media de bits (por ejemplo "3.7"), y luego se evalúa end-to-end.
  6. Respuesta acumulada (profile_cumulative_response): palletiza las capas acumulativamente en orden ascendente de divergencia para visualizar la frontera tamaño-vs-calidad (plot() produce curvas de respuesta).

Aplicar una receta al modelo de Core ML#

Las recetas se calculan sobre módulos de torch pero deben aplicarse a operaciones de MIL cuyos nombres difieren. El puente es un hash de contenido: 4 elementos fijos de cada tensor de pesos FP16 se empaquetan en bits en una clave float64 (get_tensor_hash), lo que permite a apply_recipe_coreml emparejar cada peso de Core ML de vuelta con su módulo de torch — robusto ante la alteración de nombres, con detección explícita de colisiones de hash.

Escalera de validación (de CoreMLPalettizerTestsMixin)#

Cada receta pasa tres comprobaciones de PSNR (umbral 35 dB) antes de guardar un asset:

  1. torch fake-palettized vs Core ML fake-palettized (la conversión es fiel),
  2. Core ML fake-palettized vs Core ML real-palettized (la LUT real ≈ simulación),
  3. torch fake-palettized vs Core ML real-palettized (la más estricta, end-to-end).

Lo que se publica#

Para cada versión de Whisper, solo se publican las variantes de receta más pequeña y más grande, y las carpetas se nombran por el tamaño total del artefacto en MB (openai_whisper-tiny_216MB), no por bits — el tamaño es lo que les importa a los usuarios (generate_model.py::rearrange_quantized_variants). Los artefactos completos de perfilado van a hf.co/argmaxinc/compression_artifacts.

4. Sparse outlier decomposition#

La palettization a pocos bits sufre cuando un tensor de pesos tiene unos pocos valores extremos: los outliers estiran la paleta de k-means y desperdician entradas de la LUT. compress/sparse_outlier.py implementa la clásica descomposición inlier + outlier (cf. ideas al estilo de LLM.int8()), adaptada a la ANE:

outlier_inds = (w - w.mean()).abs() > w.std() * 3        # OUTLIER_NUM_STD = 3
w_inlier  = w with outliers zeroed     → palettized (dense, low-bit)
w_outlier = w with inliers zeroed      → kept FP16, stored SPARSE
  • DecomposedModule reemplaza la capa original con dos capas paralelas cuyas salidas se suman (inlier_module(x) + outlier_module(x)) — descomposición a nivel de grafo, sin kernels personalizados.
  • La rama de outliers se comprime en Core ML con prune_weights + OpThresholdPrunerConfig(threshold=1e-6) — dado que ya está enmascarada a cero, aplicar el umbral la convierte en una representación sparse sin pérdidas.
  • El sobrecoste estimado se registra como bits_overhead = 1 + outlier_fraction × 16 bits extra por parámetro — con outliers a 3σ (~0.3 % de los pesos) eso es ≈ 1.05 bits, un seguro barato para la calidad de la paleta.
  • Se activa mediante SPARSE_OUTLIER_DECOMPOSITION / el flag de CLI --outlier-decomp; el tronco convolucional de AudioEncoder gestiona los pesos descompuestos explícitamente (audio_encoder.py::pre_transformer_proj).

5. Guía de decisión de compresión (objetivos ANE)#

Situación Técnica
El modelo cabe, la latencia está limitada por el ancho de banda Palettization uniforme (4–6 bits suele ser seguro)
La calidad cae de forma desigual a pocos bits Receta mixed-bit a partir del perfilado de divergencia por capa
Unas pocas capas dominan el error Mantener las top-K capas sensibles en FP16 (DONT_PALETTIZE_TOP_K)
Distribuciones de pesos de cola pesada Sparse outlier decomposition + palletizar el inlier
Tensores muy sparse (>80 % ceros) Podar, no palletizar
iOS 18+/macOS 15+ disponible Paletas de 3 bits, granularidad per_grouped_channel para LUTs más finas
Verificar cualquiera de lo anterior Escalera de PSNR ≥ 35 dB + prueba de velocidad ≥ 0.95× (doc 06 §8)

Documentos upstream relacionados: guía de palettization de coremltools (citada en la fuente).

Generado desde el markdown de la base de conocimiento — cada afirmación traza a una fuente citada.