Google та Meta представляють TorchTPU, нативний бекенд PyTorch для Google TPU

TorchTPU - це відкрита співпраця між Google та Meta, що представляє нативний бекенд PyTorch для Google TPU. Вона дозволяє запускати моделі «з коробки» та переносити їх на TPU, орієнтуючись на конкурентоспроможну продуктивність інференсу й навчання завдяки оптимізаціям eager execution, інтеграції torch.compile та нативному прискоренню обладнання.

Google TPU використовували для тренування та розгортання великих моделей ML у величезному масштабі. Хоча TPU історично були прив'язані до TensorFlow та JAX, ця ініціатива розпочалася в Google у середині 2025 року за участі команди PyTorch у Meta. Попередня інтеграція PyTorch/XLA, представлена у 2020 році й заснована на Lazy Tensors, вимагала глибоких втручань у вихідний код, страждала від значної затримки компіляції, не мала підтримки eager execution і відставала у розподілених примітивах.

Принципи проектування та зручність використання

TorchTPU був створений як принциповий TPU-бекенд для PyTorch, а не як адаптація JAX. Його принципи включають громадянство PyTorch (відсутність залежностей від JAX у runtime, надійний eager-режим, torch.compile на базі Torch Dynamo та підтримку DDP, FSDP і DTensor), безшовну інтеграцію з екосистемою та залучення спільноти. Фокус на користувачах - це простота портування: зміна model.to("cuda") на model.to("tpu") повинна дозволити моделі працювати «з коробки», хоча для пікової продуктивності може знадобитися додаткове налаштування. Проект також прагне до обмеженої динаміки, щоб уникнути перекомпіляції XLA при зміні форм, і досліджує динамічні ядра TPU, які можна попередньо скомпілювати. Імпорт PyTorch на машині з TPU із встановленим пакетом wheel TorchTPU вимагає лише незначних змін через механізм автозавантаження PyTorch.

Eager execution та компіляція

TorchTPU розширює диспетчер PyTorch бекендом C++, зареєстрованим як PrivateUse1, і визначає зберігання тензорів TPU через стандартний механізм at::Tensor. Він пропонує три eager-режими: Strict (за замовчуванням) відправляє одну Aten-операцію за раз асинхронно; Debug (TORCH_TPU_LAUNCH_BLOCKING=1) є синхронним з додатковими перевірками меж; та DeferAndFuse (TORCH_TPU_DEFER_AND_FUSE=1) відкладає та об'єднує кілька операцій; за повідомленнями, він на 50% або більше швидший за Strict eager і водночас зменшує пікове споживання пам'яті HBM. У точці матеріалізації, наприклад при виклику .to("cpu"), TorchTPU застосовує аналіз життєвості, будує ключ кешу для графа операцій і відправляє закешований виконуваний файл або компілює асинхронно. Евристика повторюваних операцій (Repeated Ops Heuristic) виявляє повторювані послідовності, такі як стеки блоків трансформера, і запускає матеріалізацію завчасно. Трирівневий кеш компіляції охоплює локальну пам'ять процесу, спільну локальну пам'ять машини та постійний віддалений кеш файлової системи для повторного використання між хостами та перезапусками.

Підтримувані робочі навантаження та точність

TorchTPU підтримує близько 450 Aten-операцій, включаючи прямий та зворотний прохід для основних ATen-операцій, необхідних для головних архітектур, а також власні ядра через JAX/Pallas або Helion DSL. Він був перевірений «з коробки» на 5 000 найбільш завантажених моделях Hugging Face, що охоплюють LLM, аудіо та мовлення, комп'ютерний зір і текстові енкодери. Також перевірено шардинг на базі DTensor, багатохостову інференцію з vLLM та SGLang, тонке налаштування та тренування з TorchTitan, а також інтеграцію з Ray та TorchRec. Матричні помножувачі TPU використовують вхідні дані bfloat16 з накопиченням FP32, тоді як операції без згортання виконуються у FP32 на векторному блоці. TorchTPU відображає API точності PyTorch на це обладнання: medium (за замовчуванням) використовує множення BF16 з накопиченням FP32; high застосовує 3-прохідний алгоритм за приблизно 3-кратними витратами обчислень; highest використовує 6 проходів за приблизно 6-кратними накладними витратами. Вибір high або highest викликає попередження про продуктивність.

Джерела: PyTorch · Google Cloud Blog · Hugging Face Blog