Desenvolvedor realiza engenharia reversa da API privada do Apple Neural Engine, implementando pela primeira vez o treino de redes neurais no ANE

Segundo as notícias da BlockBeats, a 3 de março, o programador Manjeet Singh (GitHub: maderix) colaborou com Claude Opus para fazer engenharia reversa das APIs privadas não divulgadas da Apple, para implementar treino de rede neural contendo retropropagação no Apple Neural Engine (ANE) do chip M4 pela primeira vez. O ANE é o acelerador de design de inferência da Apple, e a capacidade oficial de formação nunca esteve aberta, e os programadores só podem chamar indiretamente a sua função de inferência através do framework CoreML.

O projeto contornou o CoreML e mapeou diretamente para toda a stack de software a partir de mais de 40 classes privadas como _ANEClient e _ANECompiler para o kernel IOKit, e descobriu a interface _ANEInMemoryModelDescriptor que compila o modelo diretamente na memória – a chave para alcançar o treino, pois cada atualização de peso precisa de ser recompilada. Atualmente, uma única camada de transformador (dim=768, seq=512) é treinada, com 9,3ms por passo em M4, utilização de 11,2% de ANE (1,78 TFLOPS, pico teórico de 15,8 TFLOPS), gradientes de entrada de propagação direta e retroalimentada são calculados em ANE, e gradientes de peso e otimizadores Adam são feitos na CPU.

O projeto também concluiu que as primitivas computacionais centrais do ANE são convoluções em vez de multiplicação de matrizes, e que expressar multiplicação de matrizes com convoluções 1x1 pode alcançar cerca de 3 vezes o aumento de rendimento, e contornar chamadas diretas CoreML tem um ganho adicional de 2 a 4 vezes, e a promoção oficial da Apple do “38 TOPS” é enganadora. O projeto ainda está numa fase inicial: apenas treino em camada única é utilizado, dados sintéticos, existem cerca de 119 fugas de recursos compiladas que precisam de ser reiniciadas para evitar o processo, e o treino multicamada e o suporte a dados reais ainda estão em desenvolvimento. O projeto é open source sob o protocolo MIT e recebeu cerca de 2800 estrelas em 5 dias após o lançamento.