Investigación de ingeniería: llevar el transformer de Laya al ANE

Entorno de investigación: Apple M3 Max (40 núcleos de GPU, 128 GiB de memoria unificada), macOS 27.2, Core ML Tools 9.0, PyTorch 2.7.0 y NumPy 2.1.3. Es un prototipo independiente bajo experiments/ane_engineering/; el runtime publicado no cambia.

Resultado actual

Un prototipo multilingüe de B=1, L=96 fijo asigna con éxito el codificador completo, la cabeza de decisión y el scorer al Neural Engine en el plan de cómputo previsto de Core ML: 6,390 operaciones no constantes prefieren ANE, con pesos de coste estimados que suman aproximadamente 1. Las 3,809 entradas restantes son constantes. La exportación original de formas enumeradas/SDPA prefería la CPU para las 1,318 operaciones asignadas bajo CPU_AND_NE, pese a que 988 operaciones individuales listan ANE como dispositivo admitido.

La ruta de predicción completa del prototipo midió 5.167 ms p50 en 50 llamadas de cribado, incluidas tokenización, búsqueda de embeddings, máscaras de atención, inferencia en ANE, cómputo de la cabeza de acción en CPU, calibración y formateo. El cuerpo aislado de Core ML midió 4.403 ms p50 en 30 llamadas con embeddings sintéticos. Este último es una medición de componente y no es una afirmación de velocidad de extremo a extremo. La carga y compilación del modelo quedan excluidas de ambas mediciones en caliente.

En el subconjunto de longitud fija de la referencia golden FP32 original, 59/59 comparaciones de respuesta coinciden, incluidas ocho lenguajes y preguntas choice/score/noul. El mayor cambio de probabilidad calibrada es 0.002925, y 100 llamadas públicas repetidas son finitas y devuelven resultados redondeados idénticos. El fixture original de 63 preguntas contiene tres entradas largas de 1,024 tokens y una entrada de 147 tokens con 20 opciones; estas cuatro evaluaciones quedan explícitamente omitidas por la exportación L96. En este fixture hay rúbricas repetidas. Son comparaciones de regresión, no 59 ejemplos etiquetados independientes ni prueba de una precisión general en tareas sin cambios.

Exportaciones aparte L192 y L1024 también colocan las 6,390 operaciones asignadas del cuerpo en ANE. L192 supera 60/60 comparaciones; L1024 supera el fixture golden completo de 63/63. Ambas superan 100 llamadas públicas repetidas y el error máximo de probabilidad calibrada sigue siendo 0.002925. El subconjunto de entradas largas tiene por sí mismo un error máximo de 0.001128. Todos los recuentos de uso de tokens evaluados coinciden con la referencia.

Son ejecuciones de cribado en serie, no comparaciones emparejadas entre backends. Los tiempos del cuerpo usan 5 llamadas de calentamiento y 30 medidas; los tiempos de pregunta corta completa usan 10 de calentamiento y 50 medidas. La ruta completa incluye trabajo del host y comprobaciones de entrada. El cribado L96 es anterior a las comprobaciones adicionales finales de validación de entrada; la comparación controlada final usa el adaptador actual y registra su huella del código fuente. Los tiempos de compilación/carga se miden en cada proceso después de la conversión, no son una promesa sobre el primer arranque del sistema con las cachés de framework vacías. Los grafos fijos grandes realizan trabajo rellenado (padding) incluso para solicitudes cortas. Un adaptador práctico seleccionaría buckets de longitud separados; enrutar cada solicitud por L1024 descartaría la ventaja de entradas cortas.

Una ejecución aparte del fixture real workload(1, long=True) confirma una solicitud de 1,024 tokens, en lugar de una solicitud corta rellenada a ese tamaño. Mide 91.703 ms p50 / 94.776 ms p95 en 50 predicciones completas tras diez llamadas de calentamiento; todas las salidas redondeadas se mantienen estables. El benchmark histórico de entradas largas de MLX es de 51.98 ms p50. No son mediciones emparejadas de la misma ronda, pero este cribado no aporta evidencia de que el grafo ANE actual acelere entradas largas. El hash de entrada, la longitud de token real, las huellas actuales del experimento y los tiempos sin procesar se conservan en long1024-performance.json.

Evidencia sin procesar:

• Plan de cómputo del cuerpo completo y tiempos de componentes
• Validación con referencia FP32 de entradas reales y tiempos de predicción completa
• Validación L192
• Validación del fixture completo L1024
• Plan de una sola capa y comprobaciones numéricas
• Plan de un solo MLP y comprobaciones numéricas

MLComputePlan describe una colocación prevista, no una traza de ejecución de hardware. CPU_AND_NE permite CPU y ANE; no es un interruptor solo para ANE. En este experimento todas las operaciones pesadas asignadas del cuerpo prefieren ANE, pero la telemetría de hardware en runtime se evalúa por separado. El diagnóstico posterior de Instruments registró actividad de hardware del Neural Engine; su tabla es global y no puede atribuir todos los eventos a este modelo. La comparación emparejada con MLX, los resultados de integración de potencia y las limitaciones de la traza se reportan en ANE_BENCHMARKS.md. Un contador de ANE con valor cero de una herramienta de monitorización no puede establecer la ausencia de actividad de ANE sin validar ese contador en este sistema operativo/dispositivo.

Por qué el grafo original era un mal objetivo para ANE

La línea base conserva una disposición transformer B×L×C convencional, operaciones de forma dinámica, atención por lotes de cabezas completas y el operador SDPA de Core ML. Bajo selección CPU/ANE, su plan contiene 24 operaciones SDPA sin una asignación de dispositivo reportada, además de muchos casts, slices, consultas de forma, gathers y transposiciones. Algunas operaciones individuales admiten ANE, pero el grafo en su conjunto no se particiona sobre él. Por tanto, el soporte de dispositivo para operadores individuales es evidencia insuficiente de una ruta de ejecución de ANE útil.

El prototipo exitoso cambia varias cosas a la vez. Es evidencia de que la combinación habilita la colocación en ANE, no una ablación completa que identifique un único operador culpable. Los controles con SDPA de disposición original de forma fija y con atención explícita son los siguientes experimentos discriminantes útiles.

La guía publicada de Apple sobre Transformer recomienda tensores 4D channel-first, convoluciones 1×1 para las proyecciones, atención por cabeza y menos copias de disposición. Estos principios motivaron la implementación; las aceleraciones históricas de DistilBERT de Apple no demuestran una mejora de 10× frente a la línea base MLX FP16 ya rápida de este proyecto. Artículo de Apple sobre Transformer en ANE, implementación de referencia de Apple.

Arquitectura del prototipo y contrato numérico

model.py es un modelo de exportación aparte construido a partir de los parámetros del checkpoint original:

• Las activaciones ocultas usan B,C,1,L. Cada peso denso W[out,in] se convierte en un kernel de convolución 1×1 K[out,in,0,0] sin reentrenamiento ni aproximación de pesos.
• La atención se divide en cabezas individuales de 64 canales. El tensor key se transpone una vez, y dos einsums explícitos calculan QK y AV manteniendo la disposición 4D. El softmax se ejecuta sobre el eje key, la dimensión 1.
• RoPE divide cada cabeza en sus dos mitades de 32 canales. Sus cosenos, senos y bases provienen del modelo original, incluida la theta local multilingüe 160000.
• La normalización por canales preserva el orden original normalized * weight + bias y el epsilon. No copia la expresión afín con orden distinto ni el clipping opcional de la LayerNorm de referencia de Apple.
• La primera norma de atención del codificador sigue siendo la identidad; el codificador usa erf GELU exacto, mientras que las dos FFN de la cabeza de decisión conservan ReLU.
• La atención completa y las máscaras de ventana deslizante preservan el enmascarado de keys válidas y la regla de consultas rellenadas. El radio local se lee de local_attention // 2.
• El gather final de marcadores se expresa con un selector one-hot preparado externamente y un einsum 4D. El grafo mantiene 32 ranuras de marcador, incluidas las inactivas; el host reemplaza los logits inactivos con el valor original -1e4.

Una capa de codificador FP32 original completa y su contraparte BC1S difirieron como máximo en 2.29e-5 en la comprobación de disposición de PyTorch. El error de salida de capa de Core ML FP16 fue mayor, como es de esperar en este cambio de precisión/backend. La sonda del cuerpo contenía originalmente una autocomparación; esa evidencia inválida se eliminó y su comprobación de disposición del cuerpo completo en PyTorch es explícitamente not_measured. La validación real del modelo completo es, en cambio, contra los logits y las decisiones FP32 originales almacenados.

Las regresiones independientes de CPU en test_ane_layout.py comparan además un ConvBody diminuto completo con el DecisionModel original, usando oráculos de atención explícita y SDPA, tres tipos de pregunta, valores de padding cambiados, sesgos de norma distintos de cero, múltiples bases de RoPE y un epsilon de norma no predeterminado. Estos tests cubren la semántica de disposición y enmascarado sin confundirla con la precisión de hardware FP16 del checkpoint completo. Los cinco tests de disposición pasaron en local.

El prototipo de atención añade un sesgo de máscara finito de -1e4. Esto tiene el comportamiento de máscara previsto sobre las activaciones finitas validadas, pero no es una identidad bit a bit con reemplazar los scores enmascarados por -1e4 o -infinity para entradas extremas arbitrarias. Del mismo modo, no se afirma que la ejecución FP16 de Core ML sea bit a bit idéntica al modelo FP32 original.

runtime.py proporciona un único límite CPU→ANE→CPU para todo el transformer, en lugar de una transición de dispositivo por capa:

1. La CPU tokeniza cada pregunta, reúne solo las filas de embedding solicitadas y construye máscaras aditivas de forma fija, vectores de tipo y selectores de marcador. No reutiliza estados ocultos contextuales ni K/V entre preguntas.
2. Una llamada a Core ML realiza la normalización de embeddings, las 22 capas del codificador, ambas capas de la cabeza de decisión y las convoluciones de scoring.
3. La CPU deriva las características de acción de un softmax de logits sin calibrar y ejecuta la pequeña cabeza de acción original en FP32 con erf GELU. La calibración pública y el formateo de salida usan luego la implementación existente.

El cuerpo exportado tiene cómputo FP16, mientras que la pequeña cabeza de acción del host es FP32. La búsqueda de embeddings usa los pesos originales almacenados. El archivo safetensors de origen contiene 169 tensores FP16 y un tensor FP32: «FP32 reference» describe la ejecución original de PyTorch, no una afirmación de que el checkpoint original esté almacenado por completo en FP32. Este límite de precisión mixta forma parte del contrato numérico del prototipo. Diferencias de probabilidad de acción de cero en fixtures saturados no prueban la identidad de los logits de acción.

El adaptador rechaza dimensiones de paquete incompatibles, IDs/marcadores fuera de rango, valores de máscara inválidos, filas de atención-key vacías y longitudes de entrada que superan su capacidad fija. Su límite de entrada por defecto es de 96 tokens y su tamaño de lote es uno; varias preguntas se ejecutan de forma secuencial. Nunca trunca una solicitud en silencio para ajustarla a la exportación más corta.

Reproducción, procedencia y puertas de calidad

Ejecuta desde la raíz del repositorio en el .venv fijado. Los paquetes generados los ignora Git; no se requiere ningún NPZ de embeddings duplicado ni artefacto de pesos grande. probe.py se niega a usar un directorio de salida existente. Las exportaciones nuevas registran en un manifiesto los valores SHA256 del peso/config originales, los hashes de contenido del paquete, la forma, las versiones de las herramientas y las huellas del código fuente del experimento.

Los comandos de reproducción escriben bajo artifacts/ane-repro/, porque los directorios de experimento comprometidos ya contienen informes y manifiestos. Elige otro directorio nuevo al repetir una exportación; ninguno de los exportadores reutiliza en silencio un paquete existente.

Instala las dependencias de conversión, desarrollo e investigación de compresión con pip install -e '.[convert,dev,research]' (o los extras correspondientes de uv sync). El extra de investigación fija kmeans1d==0.4.0; esta dependencia opcional se usa para los experimentos de K-means agrupado en FP16 y se registra en los manifiestos nuevos.

Por defecto, la conversión resuelve laya-multilingual al checkpoint fijado de Hugging Face del repositorio y lo descarga cuando es necesario. Pasa --source /path/to/checkpoint para usar archivos locales existentes. La validación usa por defecto el directorio de origen registrado en el manifiesto del paquete. El propio runtime ANEAgent solo acepta archivos locales; verifica los hashes del peso/config originales, la forma fija y el hash de contenido del paquete antes de cargar. La validación también rechaza una referencia golden con un hash de peso de origen distinto. El SHA256 medido del peso de origen multilingüe es 9d628fd971b700382ac6f65920a86f149777b2e748e0c955fb3b19695aa8f204.

# Small placement probes, then the complete model.
.venv/bin/python -m experiments.ane_engineering.probe \
  --kind mlp --length 96 --output artifacts/ane-repro/mlp96
.venv/bin/python -m experiments.ane_engineering.probe \
  --kind layer --length 96 --output artifacts/ane-repro/layer96
.venv/bin/python -m experiments.ane_engineering.probe \
  --kind body --length 96 --output artifacts/ane-repro/body96
.venv/bin/python -m experiments.ane_engineering.validate \
  --package artifacts/ane-repro/body96/model.mlpackage \
  --length 96 --repeats 100 \
  --output artifacts/ane-repro/validation96.json

El validador exige que todas las decisiones argmax evaluadas coincidan, errores de probabilidad calibrada y de acción <=0.02, salidas finitas, uso de tokens sin cambios y salidas públicas repetidas idénticas. Un candidato fallido escribe passed: false y sale con error. Los casos omitidos siguen sin validar aunque el subconjunto de forma fija pase. El informe inicial L96 tenía estas puertas aplicadas explícitamente después de la medición; se marca en consecuencia, sin cambiar sus tiempos registrados.

Exportaciones fijas más largas y sus comandos de validación con la referencia golden completa:

.venv/bin/python -m experiments.ane_engineering.probe \
  --kind body --length 192 --output artifacts/ane-repro/body192
.venv/bin/python -m experiments.ane_engineering.validate \
  --package artifacts/ane-repro/body192/model.mlpackage --length 192 \
  --output artifacts/ane-repro/validation192.json
.venv/bin/python -m experiments.ane_engineering.probe \
  --kind body --length 1024 --output artifacts/ane-repro/body1024
.venv/bin/python -m experiments.ane_engineering.validate \
  --package artifacts/ane-repro/body1024/model.mlpackage --length 1024 \
  --output artifacts/ane-repro/validation1024.json
.venv/bin/python -m experiments.ane_engineering.benchmark \
  --package artifacts/ane-repro/body1024/model.mlpackage --length 1024 --long \
  --output artifacts/ane-repro/long1024-performance.json

Estos comandos reproducen las exportaciones más largas comprobadas de forma independiente que se listan arriba. Su colocación y fidelidad numérica se comprobaron por separado del grafo L96; rellenar solicitudes cortas a 1024 tokens no es la política de producción propuesta.

Cribado de compresión y el objetivo de 10×

palettize.py prepara variantes de paleta independientes solo de pesos, con tablas de búsqueda uniformes de 8 bits como el candidato de cribado inicial barato. Solo se seleccionan los pesos de convolución mayores de 2048 elementos; las constantes de RoPE, la normalización, la aritmética de activaciones y los pesos de acción del host no cambian. Los canales de salida agrupados usan tablas de búsqueda separadas. K-means es un modo más caro. Usa la implementación instalada de kmeans1d para estos grupos FP16. Core ML Tools 9.0 paraleliza grupos independientes mediante un starmap de process pool cuando num_kmeans_workers > 1; los experimentos usan ocho workers para el K-means offline y un hilo de biblioteca matemática por worker. El número de workers cambia el rendimiento de la exportación, no el objetivo previsto del codebook. Los candidatos de menos bits, 6/4 bits, son artefactos aproximados aparte, no implementaciones exactas.

Los tamaños de compresión se refieren al paquete del cuerpo transformer exportado. La tabla de embeddings original de 196.608 millones de entradas permanece en el host, y solo se buscan las filas solicitadas en cada petición, mientras que los pequeños pesos de acción del host no cambian. Un paquete del cuerpo que se reduce aproximadamente a la mitad no es una reducción a la mitad de todo el checkpoint, de la memoria del runtime ni de la energía por solicitud.

.venv/bin/python -m experiments.ane_engineering.palettize \
  --package artifacts/ane-repro/body96/model.mlpackage \
  --bits 8 --mode uniform --group-size 32 \
  --output artifacts/ane-repro/body96-w8
.venv/bin/python -m experiments.ane_engineering.validate \
  --package artifacts/ane-repro/body96-w8/model.mlpackage --length 96 \
  --output artifacts/ane-repro/validation96-w8.json

# Independent K-means candidates; inspect each validation exit status.
for bits in 8 6 4; do
  VECLIB_MAXIMUM_THREADS=1 OPENBLAS_NUM_THREADS=1 OMP_NUM_THREADS=1 \
    .venv/bin/python -m experiments.ane_engineering.palettize \
    --package artifacts/ane-repro/body96/model.mlpackage \
    --bits "$bits" --mode kmeans --group-size 32 --workers 8 \
    --output "artifacts/ane-repro/body96-w${bits}km"
  .venv/bin/python -m experiments.ane_engineering.validate \
    --package "artifacts/ane-repro/body96-w${bits}km/model.mlpackage" --length 96 \
    --output "artifacts/ane-repro/validation96-w${bits}km.json"
done

Los dos cribados W8 uniformes conservan las 59 decisiones argmax del fixture y superan 100 llamadas repetidas, pero fallan la puerta de probabilidad: el grupo de tamaño 32 alcanza un error de 0.023612 y el grupo de tamaño 4 alcanza 0.033858. El cribado W8 K-means/group32 pasa con un error máximo de 0.014393 y 59/59 decisiones. Las probabilidades de acción saturadas ocultan diferencias de logits de acción de hasta 16.24 en ese candidato K-means; pasar este pequeño fixture de regresión no establece una calibración general ni una precisión de tareas preservadas. Los candidatos comprimidos siguen siendo modelos aproximados identificados por separado.

El candidato fijo W6 K-means/group32 también mantiene 59/59 decisiones argmax y salidas repetidas estables, pero falla con un error máximo de probabilidad de 0.052243. Su error máximo de logits de acción es 76.62. Por tanto, reducir los pesos a seis bits no satisface la puerta de aceptación sin cambios, aunque su latencia de cribado con solicitudes cortas sigue siendo cercana a FP16.

W4 K-means/group32 también mantiene 59/59 decisiones, pero el error máximo de probabilidad sube a 0.200221 y el error máximo de logits de acción a 605.30. Falla la misma puerta. La ausencia de cambios de argmax en los cinco cribados comprimidos muestra por qué las decisiones saturadas de este fixture por sí solas son una prueba de aceptación inadecuada.

Todas las variantes comprimidas conservan 6,390 operaciones asignadas a dispositivo preferidas por NE, 59/59 acuerdos de argmax del fixture y 100 llamadas repetidas estables. Las entradas restantes del plan incluyen constantes y expresiones de reconstrucción de LUT de pesos; los metadatos de colocación por sí solos no prueban cuántos datos comprimidos viajan desde la DRAM durante una solicitud. Las tres exportaciones K-means tardan 186.50, 56.13 y 23.90 segundos respectivamente con ocho workers offline. La configuración y los procesos worker terminan antes de cada medición de inferencia.

Estos cribados en serie de 50 llamadas no establecen relaciones de aceleración. El cribado FP16 es anterior a las últimas comprobaciones de validación de entrada, y los cribados no están intercalados. W8 K-means es el único finalista comprimido para la comparación más fuerte de la misma sesión en ANE_BENCHMARKS.md. Las variantes rechazadas se conservan como evidencia del límite de precisión, no como despliegues recomendados. La compresión solo se ha validado para este subconjunto multilingüe L96; el resultado L1024 de 63 preguntas de arriba se refiere a la exportación FP16 aparte.

La compresión de paleta de Core ML reconstruye pesos en coma flotante a partir de tablas de búsqueda indexadas; tensores almacenados más pequeños no establecen por sí solos una inferencia más rápida ni una energía menor. Todas las variantes necesitan las mismas puertas de precisión, un plan de cómputo nuevo y una comparación emparejada de velocidad/energía de extremo a extremo. Documentación de paletización de Core ML.

La colocación inicial exitosa en ANE establece una ruta de optimización creíble, no un resultado de 10×. La comparación debe usar MLX FP16, incluida su variante compilada cuando sea más rápida, y reportar un alcance de tareas igual. La potencia debe integrarse sobre solicitudes completas. Tanto la energía bruta del sistema como cualquier estimación con reposo restado deben reportarse con sus limitaciones de medición. La revisión matemática independiente es ANE_MATH.md.

Qué establece la implementación manual

El trabajo manual útil aquí es una reescritura completa y validada de forma independiente del grafo de cómputo a una disposición que Core ML puede mapear a ANE. Esto cambia la colocación de la ejecución y conserva los parámetros entrenados. Es sustancialmente más eficaz que cambiar compute_units en el grafo original. No elimina los 24 bloques secuenciales de atención/MLP ni su trabajo de proyección densa.

El siguiente candidato exacto para solicitudes largas es una atención que visite realmente solo ventanas locales, implementada con tiles fijos de query/key y las reglas originales de padding y RoPE. El grafo actual sigue calculando una matriz de scores densa y aplicando una máscara local. Un grafo por tiles podría reducir ese trabajo, pero más slices, límites y contracciones pequeñas pueden socavar la planificación de ANE; su colocación, precisión y beneficio de extremo a extremo siguen sin medirse. Una mayor compresión de pesos requiere calibración o recuperación de calidad tras los fallos anteriores. La destilación o menos capas introducirían un modelo nuevo y requerirían una evaluación más amplia de la calidad de las tareas. Ninguna de esas direcciones no implementadas aporta hoy evidencia de una ganancia de 10×.