PyTorch 2.15 acelera IA no Mac com FP8 e Metal até 2x mais rápido

Tecnologia

Nova versão do PyTorch traz melhorias significativas para desempenho de IA em Macs com Apple Silicon.

Leia também: Cloudflare lança Clairic, modelo de extração focado em dados

O PyTorch 2.15 chega nesta semana com uma guinada importante para quem roda inteligência artificial em Macs com Apple Silicon: o backend MPS ganha suporte a tensores FP8, novos kernels Metal para operações complexas e aceleração direta na decodificação de modelos de linguagem, aproximando o Mac de um papel de protagonista em treinamentos e inferência.

As melhorias foram conduzidas pela equipe PyTorch da Hugging Face, em parceria com mantenedores do framework e colaboradores da comunidade, e valem para máquinas com macOS 15 ou superior e Python a partir da versão 3.11.

FP8 reduz uso de memória pela metade no Mac

A principal novidade é a chegada dos tensores torch.float8_e4m3fn ao MPS, o backend que usa a GPU dos chips Apple Silicon. Cada valor FP8 ocupa 1 byte, o que significa metade do espaço de formatos como FP16 ou BF16 para armazenar pesos e ativações.

“FP8 gives you another way to save memory on your Mac.”, resume a equipe PyTorch da Hugging Face. Na prática, essa compressão abre espaço para carregar modelos maiores ou aumentar tamanho de contexto e batch sem recorrer a placas externas ou a servidores na nuvem.

Os tensores FP8 podem ser combinados com F.scaled_mm, que aplica estratégias de escala por tensor ou por linha para manter a precisão numérica sob controle, um ponto sensível quando se reduz tanto a quantidade de bits por valor.

Decodificação de LLMs e MoE fica mais rápida

Outra frente de ganho direto é a geração de texto em modelos de linguagem. A fase de decodificação, em que cada novo token atravessa o modelo como um único vetor, recebeu caminhos mais eficientes tanto para camadas lineares densas quanto para arquiteturas Mixture-of-Experts (MoE).

“PyTorch 2.15 speeds up both kinds of layers that dominate this step: dense linear layers and mixture-of-experts layers.”, destaca a equipe. A mudança-chave foi ligar o kernel GEMV já existente no MPS para torch.mm diretamente à operação F.linear, com fusão de bias quando possível.

Em testes com 96 formatos reais de camadas BF16, isso resultou em um ganho geométrico médio de 1,23× na etapa de decodificação, com vitórias maiores em projeções pequenas — como os gates de modelos MoE, que antes desperdiçavam largura de banda de memória.

Para MoE, o PyTorch 2.15 também introduz uma implementação nativa de grouped matmul no MPS, capaz de tratar na própria GPU os grupos desiguais de tokens enviados a cada expert. O que antes exigia copiar offsets para a CPU e iterar grupo a grupo agora é feito por kernels Metal dedicados, com caminhos específicos para grupos pequenos, cenário típico na geração token a token.

O reflexo aparece diretamente na vazão de geração: “The decoding kernel shows up directly in token throughput. On an M3 Ultra with torch.compile, MoE models generate about 1.3× more tokens per second.”, afirma a equipe da Hugging Face.

Metal assume operações complexas e reduz latência

A versão 2.15 aprofunda a estratégia de substituir o MPSGraph por kernels Metal sob medida em partes críticas do cálculo numérico e do processamento multimídia.

No campo da álgebra linear com números complexos, operações como fatoração LU, resolução LU e resolução triangular em complex64 passam a rodar diretamente no MPS, liberando também determinantes complexos e gradientes associados. Pequenas matrizes FP32 até 4×4, comuns em geometria de câmera e transformações de imagem, agora são invertidas em um único kernel Metal que combina fatoração e solução.

Problemas maiores ganham QR com atualizações em blocos e versões otimizadas de triangular solve que convertem boa parte do trabalho em multiplicações de matrizes eficientes. Cargas de trabalho específicas, como Gated DeltaNet, recebem caminhos rápidos dedicados para lotes de sistemas triangulares FP32 nos tamanhos 16×16, 32×32 e 64×64 com múltiplos vetores do lado direito.

A decomposição SVD também foi retrabalhada para lidar melhor com extremos numéricos, de valores singulares minúsculos a magnitudes muito grandes, além de garantir completude de base. Isso torna mais confiáveis operações de pseudoinversa e mínimos quadrados em cenários de pesquisa e produção.

Softmax até 8× mais rápido acelera amostragem de tokens

Na borda da geração de texto, onde cada microssegundo conta, o PyTorch 2.15 traz ganhos expressivos em softmax e log_softmax, operações usadas na amostragem de tokens.

“softmax and log_softmax now run on dedicated Metal kernels, and large vocabularies benefit the most: an FP32 softmax over a [1, 128256] row of logits drops from 57 µs to 7 µs, about 8× faster than the MPSGraph version.”, detalha a equipe. Esse salto reduz gargalos na etapa final da inferência, algo especialmente relevante em aplicações com grandes vocabulários, como chatbots multilíngues e sistemas de recomendação sofisticados.

Outras operações custosas em áudio, vídeo e treinamento também migram para caminhos otimizados via Metal, com impacto direto em decodificadores de vídeo, modelos de áudio e loops de treino que antes sofriam com latência e uso excessivo de memória no driver.

Mais operações no MPS e controle fino de streams

Além dos destaques de desempenho, o PyTorch 2.15 amplia a lista de operações compatíveis com MPS: funções como torch.frexp, torch.logspace, torch.take, Tensor.put_, torch.heaviside, torch.special.laguerre_polynomial_l e comparações com inteiros sem sinal passam a rodar na GPU, reduzindo quedas inesperadas para a CPU.

No mecanismo de atenção, a implementação matemática de scaled dot-product attention no MPS ganha correção para lidar corretamente com broadcasting em lotes e cabeças, eliminando inconsistências de resultados em alguns formatos de entrada.

Desenvolvedores também ganham mais controle sobre paralelismo com torch.mps.Stream e o context manager torch.mps.stream, que permitem escolher a stream onde cada trabalho roda e registrar o uso de tensores em streams alternativas por meio de Tensor.record_stream. Isso ajuda a sobrepor operações e a extrair melhor uso da GPU em pipelines complexos.

Impacto para desenvolvedores e limite de compatibilidade

Para a comunidade que hoje usa transformers, diffusers, ComfyUI e outras ferramentas em Macs com Apple Silicon, o efeito prático é a possibilidade de treinar e servir modelos mais pesados, com menor consumo de memória e maior throughput, sem abandonar o ecossistema macOS.

Pesquisadores, startups de IA e empresas que adotam o Mac para prototipagem tendem a ser os maiores beneficiados, já que o ambiente historicamente ficava atrás de GPUs dedicadas em termos de suporte e otimização. As melhorias no MPS e o reforço em torch.compile reforçam o Mac como opção viável até para cargas de trabalho mais intensivas.

Há, porém, uma barreira clara de entrada: para tirar proveito das novidades, é obrigatório ter um Mac com Apple Silicon rodando macOS 15 ou mais recente e Python 3.11 ou superior. Máquinas mais antigas ou sistemas que não puderem ser atualizados ficarão de fora da nova leva de otimizações, o que pode pressionar equipes a planejar migrações de ambiente.

Mensagens de diagnóstico mais detalhadas ajudam a identificar quando o MPS não está disponível ou quando a versão do sistema não atende aos requisitos, reduzindo o tempo gasto em tentativa e erro por parte dos usuários.

Próximos passos já miram nova leva de otimizações

Segundo a equipe da Hugging Face, o trabalho em torno do MPS segue em ritmo acelerado para as próximas versões do PyTorch. A meta é continuar removendo dependências do MPSGraph e migrando cada vez mais operações para kernels Metal especializados, incluindo camadas de normalização de batch, funções de perda e estatísticas de variância, além de estratégias mais agressivas de envio de comandos à GPU para reduzir o tempo ocioso entre tokens.

Com o PyTorch 2.15, a Apple ganha um argumento mais forte para posicionar Macs com Apple Silicon como plataforma de primeira classe para IA. A adoção em massa ainda depende de testes em produção, correção de eventuais bugs e mais educação técnica sobre como explorar FP8 e as novas APIs de streaming, mas os primeiros relatos da comunidade apontam ganhos reais de performance e estabilidade.

Fonte: matéria original