Investigação de engenharia: mover o transformer do Laya para o ANE

Ambiente de investigação: Apple M3 Max (40 núcleos de GPU, 128 GiB de memória unificada), macOS 27.2, Core ML Tools 9.0, PyTorch 2.7.0 e NumPy 2.1.3. Este é um protótipo independente em experiments/ane_engineering/; o runtime distribuído está inalterado.

Resultado atual

Um protótipo multilingue fixo B=1, L=96 atribui com êxito o codificador completo, a cabeça de decisão e o scorer ao Neural Engine no plano de cálculo previsto do Core ML: 6,390 operações não constantes preferem o ANE, com pesos de custo estimados a somar aproximadamente 1. As restantes 3,809 entradas são constantes. A exportação original com formas enumeradas/SDPA preferia a CPU para todas as 1,318 operações atribuídas sob CPU_AND_NE, apesar de 988 operações individuais indicarem o ANE como dispositivo suportado.

O caminho de predição completo do protótipo mediu 5.167 ms p50 ao longo de 50 chamadas de triagem, incluindo tokenização, consulta de embedding, máscaras de atenção, inferência no ANE, cálculo da cabeça de ação na CPU, calibração e formatação. O corpo Core ML isolado mediu 4.403 ms p50 ao longo de 30 chamadas com embeddings sintéticos. O último é uma medição de componente e não é uma afirmação de velocidade ponta a ponta. O carregamento do modelo e a compilação estão excluídos de ambos os tempos a quente.

No subconjunto de comprimento fixo da referência golden FP32 original, 59/59 comparações de resposta coincidem, incluindo oito idiomas e perguntas choice/score/noul. A maior alteração de probabilidade calibrada é 0.002925, e 100 chamadas públicas repetidas são finitas e devolvem resultados arredondados idênticos. O fixture original de 63 perguntas contém três entradas longas de 1,024 tokens e uma entrada de 147 tokens com 20 opções; estas quatro avaliações são explicitamente ignoradas pela exportação L96. Ocorrem rubricas repetidas neste fixture. Estas são comparações de regressão, não 59 exemplos etiquetados independentes nem prova de exatidão geral na tarefa inalterada.

Exportações separadas L192 e L1024 também colocam todas as 6,390 operações atribuídas do corpo no ANE. O L192 passa 60/60 comparações; o L1024 passa o fixture golden completo de 63/63. Ambos passam 100 chamadas públicas repetidas e o erro máximo de probabilidade calibrada mantém-se 0.002925. O próprio subconjunto de entradas longas tem um erro máximo de 0.001128. Todas as contagens de utilização de tokens avaliadas coincidem com a referência.

Estas são execuções de triagem em série, não comparações pareadas entre backends. Os tempos do corpo usam 5 chamadas de aquecimento e 30 medidas; os tempos de pergunta curta completa usam 10 de aquecimento e 50 medidas. O caminho completo inclui trabalho no anfitrião e verificações de entrada. A triagem L96 é anterior às últimas verificações extra de validação de entrada; a comparação controlada final usa o adaptador atual e regista a sua impressão digital de código-fonte. Os tempos de compilação/carregamento são medidos em cada processo após a conversão, não são uma promessa sobre o primeiro arranque do sistema com caches de framework vazias. Os grafos fixos grandes realizam trabalho preenchido mesmo para pedidos curtos. Um adaptador prático selecionaria buckets de comprimento separados; encaminhar todos os pedidos pelo L1024 descartaria a vantagem das entradas curtas.

Uma execução separada do fixture real workload(1, long=True) confirma um pedido de 1,024 tokens, e não um pedido curto preenchido até esse tamanho. Mede 91.703 ms p50 / 94.776 ms p95 ao longo de 50 predições completas após dez chamadas de aquecimento; todas as saídas arredondadas permanecem estáveis. O benchmark histórico de entrada longa do MLX é 51.98 ms p50. Estas não são medições pareadas na mesma ronda, mas esta triagem não fornece provas de que o grafo ANE atual acelere entradas longas. O hash da entrada, o comprimento real em tokens, as impressões digitais atuais da experiência e os tempos em bruto são mantidos em long1024-performance.json.

Evidência em bruto:

• Plano de cálculo do corpo completo e tempos de componente
• Validação com entradas reais face à referência FP32 e tempos de predição completa
• Validação L192
• Validação do fixture completo L1024
• Plano de camada única e verificações numéricas
• Plano de MLP única e verificações numéricas

MLComputePlan descreve a colocação prevista, não um traço de execução de hardware. CPU_AND_NE permite CPU e ANE; não é um interruptor exclusivo do ANE. Nesta experiência, todas as operações atribuídas do corpo pesado preferem o ANE, mas a telemetria de hardware em runtime é avaliada em separado. O diagnóstico Instruments subsequente registou atividade de hardware do Neural Engine; a sua tabela é global e não consegue atribuir cada evento a este modelo. A comparação pareada com o MLX, os resultados de integração de potência e as limitações do traço são reportados em ANE_BENCHMARKS.md. Um contador ANE a zero de uma ferramenta de monitorização não pode estabelecer a ausência de atividade do ANE sem validar esse contador neste SO/dispositivo.

Porque o grafo original era um mau alvo para o ANE

A linha de base preserva um layout transformer convencional B×L×C, operações de forma dinâmica, atenção agrupada por cabeça inteira e o operador SDPA do Core ML. Sob a seleção CPU/ANE, o seu plano contém 24 operações SDPA sem atribuição de dispositivo reportada, além de muitos casts, slices, consultas de forma, gathers e transposes. Algumas operações individuais suportam o ANE, mas o grafo como um todo não está particionado para ele. O suporte de dispositivo para operadores individuais é, portanto, evidência insuficiente de um caminho de execução útil no ANE.

O protótipo bem-sucedido altera várias coisas em conjunto. É evidência de que a combinação possibilita a colocação no ANE, não uma ablação concluída que identifique um único operador culpado. Controlos com SDPA de layout original de forma fixa e atenção explícita são as próximas experiências discriminantes úteis.

As orientações publicadas da Apple sobre Transformers recomendam tensores 4D com o canal primeiro, convoluções 1×1 para projeções, atenção por cabeça e menos cópias de layout. Estes princípios motivaram a implementação; as acelerações históricas do DistilBERT da Apple não estabelecem uma melhoria de 10× face à linha de base MLX FP16 já rápida deste projeto. Artigo da Apple sobre Transformers no ANE, implementação de referência da Apple.

Arquitetura do protótipo e contrato numérico

model.py é um modelo de exportação separado, construído a partir dos parâmetros do checkpoint original:

• As ativações ocultas usam B,C,1,L. Cada peso denso W[out,in] torna-se um kernel de convolução 1×1 K[out,in,0,0] sem re-treino nem aproximação de pesos.
• A atenção é dividida em cabeças individuais de 64 canais. O tensor chave é transposto uma vez, e dois einsums explícitos calculam QK e AV mantendo o layout 4D. O softmax corre sobre o eixo chave, a dimensão 1.
• O RoPE divide cada cabeça nas suas duas metades de 32 canais. Os seus cossenos, senos e bases vêm do modelo original, incluindo o theta local multilingue 160000.
• A normalização por canal preserva a ordem original normalized * weight + bias e o epsilon. Não copia a expressão afim com ordem diferente nem o clipping opcional da LayerNorm de referência da Apple.
• A primeira norma de atenção do codificador permanece a identidade; o codificador usa o GELU erf exato, enquanto as duas FFN da cabeça de decisão mantêm o ReLU.
• A atenção completa e as máscaras de janela deslizante preservam a máscara de chaves válidas e a regra de consulta preenchida. O raio local é lido de local_attention // 2.
• O gather final de marcadores é expresso com um seletor one-hot preparado externamente e um einsum 4D. O grafo mantém 32 slots de marcador, incluindo slots inativos; o anfitrião substitui os logits inativos pelo valor original -1e4.

Uma camada de codificador FP32 original completa e a sua contraparte BC1S diferiram no máximo 2.29e-5 na verificação de layout do PyTorch. O erro de saída das camadas FP16 do Core ML foi maior, como esperado para esta mudança de precisão/backend. A sonda do corpo continha originalmente uma auto-comparação; essa evidência inválida foi removida e a sua verificação de layout PyTorch do corpo completo é explicitamente not_measured. A validação real do modelo completo é feita, em vez disso, contra os logits e decisões FP32 originais guardados.

As regressões independentes de CPU em test_ane_layout.py comparam ainda um ConvBody minúsculo completo com o DecisionModel original, usando oráculos de atenção explícita e SDPA, três tipos de pergunta, valores de preenchimento alterados, vieses de norma não nulos, múltiplas bases RoPE e um epsilon de norma não predefinido. Estes testes cobrem a semântica de layout e de mascaramento sem a confundir com a precisão de hardware FP16 do checkpoint completo. Todos os cinco testes de layout passaram localmente.

O protótipo de atenção adiciona um viés de máscara finito de -1e4. Isto tem o comportamento de máscara pretendido nas ativações finitas validadas, mas não é uma identidade bit a bit com a substituição de pontuações mascaradas por -1e4 ou -infinito para entradas extremas arbitrárias. Do mesmo modo, não se afirma que a execução Core ML FP16 seja bit a bit idêntica ao modelo FP32 original.

runtime.py fornece uma única fronteira CPU→ANE→CPU para todo o transformer, em vez de uma transição de dispositivo por camada:

1. A CPU tokeniza cada pergunta, recolhe apenas as linhas de embedding pedidas e constrói máscaras aditivas de forma fixa, vetores de tipo e seletores de marcador. Não reutiliza estados ocultos contextuais nem K/V entre perguntas.
2. Uma chamada Core ML realiza a normalização de embedding, todas as 22 camadas do codificador, as duas camadas da cabeça de decisão e as convoluções de pontuação.
3. A CPU deriva as características de ação a partir do softmax de logits em bruto não calibrado e corre a pequena cabeça de ação original em FP32 com GELU erf. A calibração pública e a formatação da saída usam depois a implementação existente.

O corpo exportado tem cálculo em FP16, enquanto a pequena cabeça de ação no anfitrião é FP32. A consulta de embedding usa os pesos originais guardados. O ficheiro safetensors de origem contém 169 tensores FP16 e um tensor FP32: «referência FP32» descreve a execução PyTorch original, não uma afirmação de que o checkpoint original esteja guardado inteiramente em FP32. Esta fronteira de precisão mista faz parte do contrato numérico do protótipo. Diferenças nulas de probabilidade de ação em fixtures saturados não provam identidade dos logits de ação.

O adaptador rejeita dimensões de pacote incompatíveis, IDs/marcadores fora do intervalo, valores de máscara inválidos, linhas de chave de atenção vazias e comprimentos de entrada que excedam a sua capacidade fixa. O seu limite de entrada predefinido é 96 tokens e o seu tamanho de lote é um; várias perguntas executam sequencialmente. Nunca trunca silenciosamente um pedido para caber na exportação mais curta.

Reprodução, proveniência e gates de qualidade

Corre a partir da raiz do repositório no .venv fixado. Os pacotes gerados são ignorados pelo Git; não é preciso um NPZ de embedding duplicado nem um artefacto de pesos grande. O probe.py recusa um diretório de saída já existente. As novas exportações registam num manifesto os valores SHA256 dos pesos/configuração originais, os hashes do conteúdo do pacote, a forma, as versões das ferramentas e as impressões digitais do código-fonte da experiência.

Os comandos de reprodução escrevem sob artifacts/ane-repro/, porque os diretórios de experiência já consolidados contêm relatórios e manifestos. Escolhe outro diretório novo ao repetir uma exportação; nenhum dos exportadores reutiliza silenciosamente um pacote existente.

Instala as dependências de conversão, desenvolvimento e investigação de compressão com pip install -e '.[convert,dev,research]' (ou os extras correspondentes do uv sync). O extra de investigação fixa kmeans1d==0.4.0; esta dependência opcional é usada nas experiências de K-means agrupado FP16 e registada nos novos manifestos.

Por predefinição, a conversão resolve laya-multilingual para o checkpoint do Hugging Face fixado do repositório e descarrega-o quando necessário. Passa --source /path/to/checkpoint para usar ficheiros locais existentes. A validação usa por predefinição o diretório de origem registado no manifesto do pacote. O próprio runtime ANEAgent só aceita ficheiros locais; verifica os hashes originais de pesos/configuração, a forma fixa e o hash do conteúdo do pacote antes de carregar. A validação também rejeita uma referência golden com um hash de pesos de origem diferente. O SHA256 medido dos pesos de origem multilingues é 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

O validador exige que todas as decisões argmax avaliadas coincidam, erros de probabilidade de calibração e de ação <=0.02, saídas finitas, utilização de tokens inalterada e saídas públicas repetidas idênticas. Um candidato falhado escreve passed: false e termina sem sucesso. Os casos ignorados permanecem não validados mesmo que o subconjunto de forma fixa passe. O relatório L96 inicial aplicou estes gates explicitamente após a medição; está marcado em conformidade, sem alterar os seus tempos registados.

Exportações fixas mais longas e os seus comandos de validação da referência 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

Estes comandos reproduzem as exportações mais longas verificadas de forma independente, listadas acima. A sua colocação e fidelidade numérica foram verificadas em separado do grafo L96; preencher pedidos curtos até 1024 tokens não é a política de produção proposta.

Triagem de compressão e o objetivo de 10×

palettize.py prepara variantes de paleta independentes apenas de pesos, com tabelas de consulta uniformes de 8 bits como o primeiro candidato de triagem barato. Apenas são selecionados pesos de convolução maiores do que 2048 elementos; as constantes RoPE, a normalização, a matemática das ativações e os pesos de ação no anfitrião permanecem inalterados. Os canais de saída agrupados usam tabelas de consulta separadas. O K-means é um modo mais caro. Usa a implementação kmeans1d instalada para estes grupos FP16. O Core ML Tools 9.0 paraleliza grupos independentes através de um starmap de pool de processos quando num_kmeans_workers > 1; as experiências usam oito workers para K-means offline e um thread de biblioteca matemática por worker. O número de workers altera o débito da exportação, não o objetivo do livro de códigos pretendido. Os candidatos de 6/4 bits são artefactos aproximados separados, não implementações exatas.

Os tamanhos de compressão referem-se ao pacote do corpo do transformer exportado. A tabela de embedding original de 196.608 milhões de entradas permanece no anfitrião, com apenas as linhas pedidas consultadas em cada pedido, e os pequenos pesos de ação no anfitrião permanecem inalterados. Um pacote do corpo a encolher cerca de duas vezes não é uma redução para metade de todo o checkpoint, da memória do runtime ou da energia por pedido.

.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

As duas triagens uniformes W8 retêm todas as 59 decisões argmax do fixture e passam 100 chamadas repetidas, mas falham o gate de probabilidade: o tamanho de grupo 32 atinge um erro de 0.023612 e o tamanho de grupo 4 atinge 0.033858. A triagem W8 K-means/grupo32 passa com um erro máximo de 0.014393 e 59/59 decisões. As probabilidades de ação saturadas escondem diferenças de logits de ação até 16.24 para esse candidato K-means; passar este pequeno fixture de regressão não estabelece calibração geral nem exatidão na tarefa preservadas. Os candidatos comprimidos continuam a ser modelos aproximados identificados em separado.

O candidato fixo W6 K-means/grupo32 também mantém 59/59 decisões argmax e saídas repetidas estáveis, mas falha com um erro máximo de probabilidade de 0.052243. O seu erro máximo de logits de ação é 76.62. Reduzir os pesos a seis bits não satisfaz, portanto, o gate de aceitação inalterado, embora a sua latência de triagem de pedidos curtos permaneça próxima do FP16.

O W4 K-means/grupo32 também mantém 59/59 decisões, mas o erro máximo de probabilidade sobe para 0.200221 e o erro máximo de logits de ação para 605.30. Falha o mesmo gate. A ausência de alterações de argmax nas cinco triagens comprimidas mostra porque as decisões saturadas deste fixture, por si sós, são um teste de aceitação inadequado.

Todas as variantes comprimidas retêm 6,390 operações atribuídas a dispositivo preferidas pelo NE, 59/59 concordâncias de argmax no fixture e 100 chamadas repetidas estáveis. As restantes entradas do plano incluem constantes e expressões de reconstrução de LUT de pesos; os metadados de colocação, por si sós, não provam quantos dados comprimidos viajam da DRAM durante um pedido. As três exportações K-means demoram 186.50, 56.13 e 23.90 segundos respetivamente com oito workers offline. A configuração e os processos worker terminam antes de cada medição de inferência.

Estas triagens em série de 50 chamadas não estabelecem rácios de aceleração. A triagem FP16 é anterior às últimas verificações de validação de entrada, e as triagens não estão intercaladas. O W8 K-means é o único finalista comprimido para a comparação mais forte na mesma sessão em ANE_BENCHMARKS.md. As variantes rejeitadas são mantidas como evidência da fronteira de precisão, não como implementações recomendadas. A compressão só foi validada para este subconjunto multilingue L96; o resultado L1024 de 63 perguntas acima diz respeito à exportação FP16 separada.

A compressão de paleta do Core ML reconstrói pesos de vírgula flutuante a partir de tabelas de consulta indexadas; tensores armazenados mais pequenos não estabelecem por si sós uma inferência mais rápida ou menor energia. Cada variante precisa dos mesmos gates de precisão, de um plano de cálculo novo e de uma comparação pareada de velocidade/energia ponta a ponta. Documentação de paletização do Core ML.

A colocação inicial bem-sucedida no ANE estabelece um caminho de otimização credível, não um resultado de 10×. A comparação tem de usar o MLX FP16, incluindo a sua variante compilada quando for mais rápida, e reportar um âmbito de tarefa igual. A potência tem de ser integrada ao longo de pedidos completos. Tanto a energia total do sistema como qualquer estimativa subtraindo o inativo devem ser reportadas com as suas limitações de medição. A revisão matemática independente é ANE_MATH.md.

O que a implementação manual estabelece

O trabalho manual útil aqui é uma reescrita completa e independentemente validada do grafo de cálculo para um layout que o Core ML consegue mapear para o ANE. Isto altera a colocação da execução mantendo os parâmetros treinados. É substancialmente mais eficaz do que alterar compute_units no grafo original. Não remove os 24 blocos sequenciais de atenção/MLP nem o seu trabalho de projeção densa.

O próximo candidato exato para pedidos longos é uma atenção que visite realmente apenas janelas locais, implementada com tiles fixos de consulta/chave e as regras originais de preenchimento e de RoPE. O grafo atual ainda calcula uma matriz de pontuações densa e aplica uma máscara local. Um grafo em tiles poderia reduzir esse trabalho, mas mais slices, fronteiras e pequenas contrações podem prejudicar o escalonamento do ANE; a sua colocação, precisão e benefício ponta a ponta permanecem por medir. Mais compressão de pesos exige calibração ou recuperação de qualidade após as falhas acima. A destilação ou menos camadas introduziriam um novo modelo e exigiriam uma avaliação mais ampla da qualidade na tarefa. Nenhuma dessas direções não implementadas fornece hoje evidência de um ganho de 10×.