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

Ambiente de pesquisa: 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 lançado está inalterado.

Resultado atual

Um protótipo multilíngue fixo B=1, L=96 atribui com sucesso o encoder completo, a decision head e o scorer ao Neural Engine no compute plan previsto do Core ML: 6,390 operações não constantes preferem o ANE, com pesos de custo estimados somando aproximadamente 1. As 3,809 entradas restantes são constantes. A exportação original de 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 listarem o ANE como dispositivo suportado.

O caminho de predição completo do protótipo mediu 5.167 ms p50 em 50 chamadas de triagem, incluindo tokenização, lookup de embedding, máscaras de atenção, inferência no ANE, computação da action-head na CPU, calibração e formatação. O corpo Core ML isolado mediu 4.403 ms p50 em 30 chamadas com embeddings sintéticos. Este último é uma medição de componente e não é uma afirmação de velocidade ponta a ponta. O carregamento e a compilação do modelo são excluídos das duas medições a quente.

No subconjunto de comprimento fixo da golden reference FP32 original, 59/59 comparações de resposta concordam, incluindo oito idiomas e perguntas choice/score/noul. A maior variação de probabilidade calibrada é 0.002925, e 100 chamadas públicas repetidas são finitas e retornam resultados arredondados idênticos. A 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; essas quatro avaliações são explicitamente puladas pela exportação L96. Rubricas repetidas ocorrem nesta fixture. Essas são comparações de regressão, não 59 exemplos rotulados independentes nem prova de acurácia geral de tarefa inalterada.

Exportações separadas L192 e L1024 também colocam todas as 6,390 operações de corpo atribuídas no ANE. L192 passa em 60/60 comparações; L1024 passa na golden fixture completa de 63/63. Ambas passam em 100 chamadas públicas repetidas e o erro máximo de probabilidade calibrada permanece 0.002925. O próprio subconjunto de entrada longa tem erro máximo 0.001128. Todas as contagens de uso de tokens avaliadas correspondem à referência.

Essas são execuções seriais de triagem, não comparações pareadas entre backends. As medições de corpo usam 5 chamadas de aquecimento e 30 medidas; as medições de pergunta curta completa usam 10 de aquecimento e 50 medidas. O caminho completo inclui trabalho no host e verificações de entrada. A triagem L96 é anterior às últimas verificações extras de validação de entrada; a comparação controlada final usa o adapter atual e registra sua fingerprint de origem. Os tempos de compilação/carregamento são medidos em cada processo após a conversão, não uma promessa sobre a primeira inicialização do sistema com caches de framework vazios. Grafos fixos grandes realizam trabalho com preenchimento mesmo para requisições curtas. Um adapter prático selecionaria buckets de comprimento separados; rotear toda requisição por L1024 descartaria a vantagem de entrada curta.

Uma execução separada da fixture real workload(1, long=True) confirma uma requisição de 1,024 tokens, em vez de uma requisição curta preenchida até esse tamanho. Ela mede 91.703 ms p50 / 94.776 ms p95 em 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. Essas não são medições pareadas na mesma rodada, mas esta triagem não fornece evidência de que o grafo ANE atual acelere entradas longas. O hash de entrada, o comprimento real de tokens, as fingerprints do experimento atual e as medições brutas são mantidos em long1024-performance.json.

Evidência bruta:

• Compute plan do corpo completo e tempos dos componentes
• Validação contra a referência FP32 com entradas reais e tempos de predição completa
• Validação L192
• Validação da fixture completa L1024
• Plano e verificações numéricas de uma única camada
• Plano e verificações numéricas de um único MLP

O MLComputePlan descreve o posicionamento previsto, não um trace de execução de hardware. CPU_AND_NE permite CPU e ANE; não é um interruptor somente ANE. Neste experimento todas as operações atribuídas ao corpo pesado preferem o ANE, mas a telemetria de hardware em tempo de execução é avaliada separadamente. O diagnóstico posterior do Instruments registrou atividade de hardware do Neural Engine; 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 trace são relatados em ANE_BENCHMARKS.md. Um contador ANE com valor zero de uma ferramenta de monitoramento não pode estabelecer a ausência de atividade do ANE sem validar esse contador neste sistema operacional/dispositivo.

Por que o grafo original era um alvo ruim para o ANE

A linha de base preserva um layout de transformer B×L×C convencional, operações de forma dinâmica, atenção em lote por cabeça inteira e o operador SDPA do Core ML. Sob a seleção CPU/ANE, seu plano contém 24 operações SDPA sem atribuição de dispositivo relatada, 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 é particionado nele. 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 muda várias coisas em conjunto. É evidência de que a combinação viabiliza o posicionamento no ANE, não uma ablação concluída que identifica um único operador problemático. SDPA de forma fixa no layout original e controles de atenção explícita são os próximos experimentos discriminantes úteis.

A orientação publicada da Apple para Transformer recomenda tensores 4D channel-first, convoluções 1×1 para projeções, atenção por cabeça e menos cópias de layout. Esses princípios motivaram a implementação; os ganhos históricos de velocidade do DistilBERT da Apple não estabelecem uma melhoria de 10× contra a linha de base MLX FP16 já rápida deste projeto. Artigo da Apple sobre Transformer 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] se torna um kernel de convolução 1×1 K[out,in,0,0] sem retreinamento nem aproximação de pesos.
• A atenção é dividida em cabeças individuais de 64 canais. O tensor key é transposto uma vez, e dois einsums explícitos calculam QK e AV mantendo o layout 4D. O softmax roda sobre o eixo key, dimensão 1.
• O RoPE divide cada cabeça em suas duas metades de 32 canais. Seus cossenos, senos e bases vêm do modelo original, incluindo o theta local multilíngue 160000.
• A normalização de canal preserva a ordenação original normalized * weight + bias e o epsilon. Ela não copia a expressão afim com ordenação diferente nem o clipping opcional do LayerNorm de referência da Apple.
• A primeira normalização de atenção do encoder permanece a identidade; o encoder usa GELU erf exato, enquanto os dois FFNs da decision head mantêm o ReLU.
• As máscaras de atenção completa e de janela deslizante preservam o mascaramento de chaves válidas e a regra de query com preenchimento. O raio local é lido de local_attention // 2.
• O gather final de marcadores é expresso usando um seletor one-hot preparado externamente e um einsum 4D. O grafo mantém 32 slots de marcador, incluindo slots inativos; o host substitui logits inativos pelo valor original -1e4.

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

As regressões independentes de CPU em test_ane_layout.py também comparam um ConvBody pequeno 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 diferentes de zero, múltiplas bases de RoPE e um epsilon de norma não padrão. Esses testes cobrem a semântica de layout e mascaramento sem confundi-las 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. Isso tem o comportamento de máscara pretendido nas ativações validadas finitas, mas não é uma identidade bit a bit com substituir pontuações mascaradas por -1e4 ou -infinito para entradas extremas arbitrárias. Da mesma forma, não se afirma que a execução FP16 do Core ML 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, reúne apenas as linhas de embedding solicitadas e constrói máscaras aditivas de forma fixa, vetores de tipo e seletores de marcador. Ela não reutiliza hidden states contextuais nem K/V entre perguntas.
2. Uma chamada ao Core ML realiza a normalização de embedding, todas as 22 camadas do encoder, as duas camadas da decision head e as convoluções de pontuação.
3. A CPU deriva os recursos de ação do softmax bruto de logits não calibrado e roda a pequena action head original em FP32 com GELU erf. A calibração pública e a formatação da saída então usam a implementação existente.

O corpo exportado tem computação FP16, enquanto a pequena action head no host é FP32. O lookup de embedding usa os pesos originais armazenados. O arquivo safetensors de origem contém 169 tensores FP16 e um tensor FP32: “referência FP32” descreve a execução original do PyTorch, não uma afirmação de que o checkpoint original é armazenado inteiramente em FP32. Essa fronteira de precisão mista faz parte do contrato numérico do protótipo. Diferenças de probabilidade de ação iguais a zero em fixtures saturadas não provam identidade de action-logits.

O adapter 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 excedem sua capacidade fixa. Seu limite de entrada padrão é de 96 tokens e seu tamanho de batch é um; várias perguntas executam sequencialmente. Ele nunca trunca uma requisição silenciosamente para caber na exportação mais curta.

Reprodução, proveniência e portões de qualidade

Execute a partir da raiz do repositório no .venv fixado. Os pacotes gerados são ignorados pelo Git; nenhum NPZ de embedding duplicado nem artefato grande de pesos é necessário. O probe.py se recusa a usar um diretório de saída existente. Novas exportações registram os valores SHA256 originais de peso/config, hashes de conteúdo do pacote, forma, versões de ferramentas e fingerprints da origem do experimento em um manifesto.

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

Instale as dependências de conversão, desenvolvimento e pesquisa de compressão com pip install -e '.[convert,dev,research]' (ou os extras correspondentes do uv sync). O extra de pesquisa fixa kmeans1d==0.4.0; essa dependência opcional é usada para os experimentos de K-means agrupado em FP16 e registrada em novos manifestos.

Por padrão, a conversão resolve laya-multilingual para o checkpoint fixado do Hugging Face do repositório e o baixa quando necessário. Passe --source /path/to/checkpoint para usar arquivos locais existentes. A validação usa como padrão o diretório de origem registrado no manifesto do pacote. O próprio runtime ANEAgent só aceita arquivos locais; ele verifica os hashes originais de peso/config, a forma fixa e o hash de conteúdo do pacote antes de carregar. A validação também rejeita uma golden reference com um hash de pesos de origem diferente. O SHA256 medido dos pesos de origem multilíngues é 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 correspondam, erros de probabilidade calibrada e de ação <=0.02, saídas finitas, uso de tokens inalterado e saídas públicas repetidas idênticas. Um candidato com falha grava passed: false e sai sem sucesso. Os casos pulados permanecem não validados mesmo que o subconjunto de forma fixa passe. O relatório L96 inicial teve esses portões aplicados explicitamente após a medição; ele é marcado de acordo, sem alterar seus tempos registrados.

Exportações fixas mais longas e seus comandos completos de validação contra a golden reference:

.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

Esses comandos reproduzem as exportações mais longas verificadas independentemente listadas acima. Seu posicionamento e fidelidade numérica foram verificados separadamente do grafo L96; preencher requisições curtas 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 independentes de paleta somente de pesos, com tabelas de lookup uniformes de 8 bits como o candidato de triagem econômico inicial. Apenas pesos de convolução maiores que 2048 elementos são selecionados; constantes de RoPE, normalização, matemática de ativação e pesos de ação no host permanecem inalterados. Canais de saída agrupados usam tabelas de lookup separadas. O K-means é um modo mais caro. Ele usa a implementação kmeans1d instalada para esses grupos FP16. O Core ML Tools 9.0 paraleliza grupos independentes via um starmap de process-pool quando num_kmeans_workers > 1; os experimentos usam oito workers para K-means offline e uma thread de biblioteca matemática por worker. A contagem de workers muda o throughput da exportação, não o objetivo de codebook pretendido. Os candidatos de 6/4 bits são artefatos aproximados separados, não implementações exatas.

Os tamanhos de compressão se referem ao pacote corpo do transformer exportado. A tabela de embedding original de 196.608 milhões de entradas permanece no host, com apenas as linhas solicitadas consultadas para cada requisição, e os pequenos pesos de ação no host permanecem inalterados. Um pacote de corpo encolhendo cerca de duas vezes não é uma redução de duas vezes no checkpoint inteiro, na memória de runtime ou na energia por requisição.

.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 mantêm todas as 59 decisões argmax da fixture e passam em 100 chamadas repetidas, mas falham no portão de probabilidade: tamanho de grupo 32 alcança erro 0.023612 e tamanho de grupo 4 alcança 0.033858. A triagem K-means/group32 W8 passa com erro máximo 0.014393 e 59/59 decisões. Probabilidades de ação saturadas escondem diferenças de action-logits de até 16.24 para esse candidato K-means; passar nessa pequena fixture de regressão não estabelece calibração geral ou acurácia de tarefa preservadas. Os candidatos comprimidos continuam sendo modelos aproximados identificados separadamente.

O candidato fixo W6 K-means/group32 também mantém 59/59 decisões argmax e saídas repetidas estáveis, mas falha com erro máximo de probabilidade 0.052243. Seu erro máximo de action-logits é 76.62. Reduzir os pesos para seis bits portanto não satisfaz o portão de aceitação inalterado, mesmo que sua latência de triagem de requisição curta permaneça próxima do FP16.

O W4 K-means/group32 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 action-logits para 605.30. Ele falha no mesmo portão. A ausência de mudanças de argmax nas cinco triagens comprimidas mostra por que as decisões saturadas desta fixture sozinhas são um teste de aceitação inadequado.

Todas as variantes comprimidas mantêm 6,390 operações com dispositivo atribuído preferidas por NE, 59/59 concordâncias de argmax da fixture e 100 chamadas repetidas estáveis. As entradas restantes do plano incluem constantes e expressões de reconstrução de LUT de pesos; apenas os metadados de posicionamento não provam quanto dado comprimido trafega da DRAM durante uma requisição. As três exportações K-means levam 186.50, 56.13 e 23.90 segundos, respectivamente, com oito workers offline. O setup e os processos de worker terminam antes de cada medição de inferência.

Essas triagens seriais de 50 chamadas não estabelecem razões de aceleração. A triagem FP16 é anterior às últimas verificações de validação de entrada, e as triagens não sã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 implantações recomendadas. A compressão só foi validada para este subconjunto multilíngue 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 ponto flutuante a partir de tabelas de lookup indexadas; tensores armazenados menores não estabelecem por si sós inferência mais rápida ou menor energia. Toda variante precisa dos mesmos portões de precisão, um compute plan novo e uma comparação pareada de velocidade/energia ponta a ponta. Documentação de paletização do Core ML.

O posicionamento inicial bem-sucedido no ANE estabelece um caminho de otimização crível, não um resultado de 10×. A comparação deve usar MLX FP16, incluindo sua variante compilada onde ela for mais rápida, e relatar o mesmo escopo de tarefa. A potência deve ser integrada ao longo de requisições completas. Tanto a energia bruta do sistema quanto qualquer estimativa com idle subtraído devem ser relatadas com 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 computação para um layout que o Core ML consegue mapear para o ANE. Isso muda o posicionamento de execução mantendo os parâmetros treinados. É substancialmente mais eficaz do que mudar compute_units no grafo original. Não remove os 24 blocos sequenciais de atenção/MLP nem seu trabalho de projeção densa.

O próximo candidato exato para requisições longas é uma atenção que na verdade visita apenas janelas locais, implementada usando tiles fixos de query/key e as regras originais de preenchimento e RoPE. O grafo atual ainda calcula uma matriz de pontuação 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; seu posicionamento, precisão e benefício ponta a ponta permanecem não medidos. Mais compressão de pesos exige calibração ou recuperação de qualidade após as falhas acima. Destilação ou menos camadas introduziria um novo modelo e exigiria uma avaliação de qualidade de tarefa mais ampla. Nenhuma dessas direções não implementadas fornece evidência de um ganho de 10× hoje.