Anunciando PyTorch/XLA 2.4: Um... Nota

Anunciando PyTorch/XLA 2.4: Um Pallas melhorado e experiência do desenvolvedor, além do 'modo ansioso'

O PyTorch/XLA 2.4 oferece melhorias significativas para o aprendizado profundo em TPU e GPUs. Esta versão apresenta melhorias no Pallas, uma linguagem de kernel personalizada, que agora suporta tanto TPU quanto GPUs, melhorando o desempenho com código Python. Novas chamadas de API, como torch_xla.sync(), simplificam a integração em workflows PyTorch existentes. Um modo eager experimental permite a execução imediata de operações no hardware-alvo, embora os TPU exijam uma chamada "mark_step" para emulação. As melhorias no Pallas incluem suporte a Flash Attention e Paged Attention, e kernels de matriz esparsa em blocos integrados para multiplicação de matriz em grupo. Além disso, uma nova interface de linha de comando TPU, tpu-info, facilita o depuração exibindo informações de utilização e dispositivo, semelhante à ferramenta nvidia-smi da Nvidia. Apesar dessas mudanças, o código existente permanece compatível, tornando a atualização indolor para os desenvolvedores. Essas atualizações visam coletivamente melhorar a usabilidade e o desempenho no PyTorch/XLA.