Google открыла TorchTPU — код PyTorch запускается на чипах TPU напрямую

Google открыла исходный код TorchTPU — родной бэкенд PyTorch для собственных чипов TPU. До сих пор для обучения на TPU приходилось переходить на JAX или идти через прослойку PyTorch/XLA, которая переводит вычисления в чужой формат и теряет часть скорости.

TorchTPU позволяет запускать привычный код PyTorch на TPU напрямую — так же, как он запускается на видеокартах Nvidia. Для команд это снимает главный барьер переноса: не нужно переписывать обучение под другой фреймворк ради доступа к чужому железу.

В тот же день Amazon представила TorchNeuron — такой же мост между PyTorch и своими чипами Trainium. Обе компании метят в одно место: привязку обучения к железу и инструментам Nvidia.