AWS и Google выпустили бэкенды PyTorch для Trainium и TPU

AWS представила TorchNeuron для чипов Trainium, а Google и Meta — TorchTPU для TPU. Оба проекта открыты и подключаются к PyTorch через один механизм регистрации PrivateUse1. Google проверила 4 917 моделей Hugging Face: 98% запустились без сбоев, 97% из проверенных дали те же числа, что на GPU.

• TorchTPU поддерживает около 450 операций PyTorch
• Опция Skinny Collectives стоит 30–40% скорости end to end
• Тест на 64 чипах: 1 828 против 1 937 токенов/с у JAX/MaxText
• TorchTPU заменил прежний путь через JAX на нативный бэкенд