Долгие годы экосистема PyTorch и специализированные ускорители Google (TPU) существовали в параллельных вселенных. Разработчикам приходилось использовать PyTorch/XLA — сложный слой абстракции, который транслировал вычислительные графы в формат, понятный TPU. Это часто приводило к головной боли: отладка была затруднена из-за «ленивой» (lazy) природы выполнения, а поддержка распределенных вычислений ограничивалась строгими правилами SPMD (Single Program, Multiple Data). Однако ситуация кардинально меняется. На конференции PyTorch Conference 2026, состоявшейся 20 октября 2026 года, Google и Meta анонсировали TorchTPU — полностью открытый, нативный бэкенд, который интегрирует TPU в PyTorch так же естественно, как GPU или CPU.
TorchTPU — это не просто обновление старого драйвера. Это фундаментальный сдвиг в архитектуре взаимодействия фреймворка и железа. Google официально заявляет, что TorchTPU заменяет PyTorch/XLA и рекомендует использовать его для всех новых рабочих нагрузок. Теперь тензоры TPU — это обычные тензоры PyTorch, без обертки или наследования специальных классов. Вы просто пишете torch.device("tpu"), и все работает. Это открывает путь к более простому, быстрому и предсказуемому машинному обучению в облаке Google.
01Что такое TorchTPU и почему это важно?
В основе TorchTPU лежит интерфейс PrivateUse1 в PyTorch. Это технический механизм, позволяющий регистрировать новые типы устройств без изменения ядра фреймворка. Однако Google пошла дальше простой регистрации. Они создали бэкенд, который понимает семантику PyTorch «из коробки».
Главное отличие от предшественника, PyTorch/XLA, заключается в модели выполнения. PyTorch/XLA записывал ленивые графы тензоров. Это означало, что код не выполнялся сразу, а накапливался в граф, который затем компилировался. Хотя это давало оптимизации, оно делало отладку кошмаром: вы не могли просто поставить print() или проверить значение переменной в середине обучения, не сломав граф. Кроме того, PyTorch/XLA поддерживал только чистый код SPMD, где все процессы выполняли одну и ту же логику. TorchTPU же поддерживает дивергентное MPMD-исполнение (Multiple Program, Multiple Data), что критически важно для современных сложных архитектур, где разные части модели или разные ранги могут выполнять разный код.
02Как работает выполнение кода: три режима и компиляция
Одной из самых сильных сторон TorchTPU является гибкость режимов выполнения. Google реализовала три режима eager mode, которые позволяют разработчикам балансировать между скоростью отладки и производительностью:
- Strict (Строгий) — режим по умолчанию. В этом режиме каждая операция диспетчеризируется асинхронно, но с сохранением стандартной семантики PyTorch. Это обеспечивает предсказуемое поведение и хорошую производительность без необходимости глубокой оптимизации.
- Debug (Отладочный). Этот режим запускает операции синхронно с проверкой границ тензоров (tensor bounds checking). Он идеален для поиска ошибок, так как позволяет мгновенно видеть результаты вычислений и отлавливать несоответствия размеров тензоров.
- DeferAndFuse (Отложить и объединить). Это режим оптимизации, который объединяет последовательности операций eager mode, чтобы сократить трафик между памятью HBM (High Bandwidth Memory) и вычислительными ядрами. Согласно техническому отчету Google, этот режим может быть заметно быстрее режима Strict, так как минимизирует накладные расходы на передачу данных.
Для достижения пиковой производительности в продакшене используется torch.compile(backend="tpu"). Этот процесс захватывает вычислительный граф с помощью TorchDynamo, затем понижает его через стандарт StableHLO в OpenXLA, который генерирует машинный код для TPU. Скомпилированные исполняемые файлы кэшируются между процессами, хостами и даже перезапусками заданий, что значительно ускоряет повторные запуски.
Кроме того, TorchTPU поддерживает пользовательские ядра Pallas, которые можно вызывать как стандартные функции PyTorch. Также работают предварительно собранные ядра Tokamax и Helion, что дает разработчикам доступ к низкоуровневой оптимизации, не отказываясь от высокого уровня абстракции Python.
03Производительность: цифры и масштабирование
Теория — это хорошо, но в индустрии AI важны цифры. Google и их партнеры предоставили впечатляющие данные о производительности TorchTPU, особенно в контексте обучения больших языковых моделей (LLM).
В тестах с использованием фреймворка TorchTitan, модель Qwen 3 8B достигла показателя MFU (Model FLOPs Utilization) на уровне 52% на чипах Ironwood (TPU7x). Это очень высокий показатель эффективности использования вычислительных мощностей. Более того, Google сообщает о почти линейном масштабировании до 8192 чипов Ironwood. Это означает, что если вы удвоите количество TPU, время обучения удвоится пропорционально, что критически важно для экономии времени при обучении гигантских моделей.
Интересный кейс оптимизации: в сценарии тонкой настройки (fine-tuning) с помощью TorchTitan время одного шага обучения сократилось с 11.49 секунд до 2.78 секунд. Само по себе использование режима DeferAndFuse дало ускорение в 1.15 раза, а полная компиляция через torch.compile дала ускорение в 3.80 раза. Эти цифры демонстрируют, что даже «из коробки» TorchTPU работает неплохо, но с правильной настройкой он может превзойти ожида
Источник: MarkTechPost ↗