Annonce de PyTorch/XLA 2.4 : Une Pallas améliorée et une expérience de développeur, plus le mode « eager »
PyTorch/XLA 2.4 propose des améliorations significatives pour l'apprentissage automatique en profondeur sur les TPU et les GPU. Cette version introduit des améliorations à Pallas, un langage de noyau personnalisé, qui prend désormais en charge à la fois les TPU et les GPU, améliorant les performances avec le code Python. De nouveaux appels d'API, comme torch_xla.sync(), simplifient l'intégration dans les workflows PyTorch existants. Un mode d'exécution immédiat expérimental permet l'exécution immédiate des opérations sur le matériel cible, bien que les TPU nécessitent un appel "mark_step" pour l'émulation. Les améliorations apportées à Pallas incluent le support de l'attention Flash et de l'attention paginée, ainsi que des noyaux de multiplication matricielle groupée denses intégrés pour les méga-blocs. De plus, une nouvelle interface de ligne de commande TPU, tpu-info, facilite le débogage en affichant l'utilisation et les informations du périphérique, similaires à l'outil nvidia-smi de Nvidia. Malgré ces changements, le code existant reste compatible, ce qui rend la mise à jour transparente pour les développeurs. Ces mises à jour visent collectivement à améliorer l'utilisabilité et les performances dans PyTorch/XLA.