Интеграция TorchTPU в экосистему Hugging Face: нативная поддержка ускорения

Запуск нейросетей PyTorch на аппаратных ускорителях Google долгое время требовали специализированных инструментов вроде torch_xla или фреймворка JAX. Появление проекта TorchTPU меняет привычный подход, добавляя нативный бэкенд TPU непосредственно в PyTorch без необходимости применять дополнительные прослойки.

Теперь популярные библиотеки экосистемы Hugging Face способны работать с тензорными процессорами напрямую. Достаточно выбрать нужную модель в репозитории и передать её на устройство через стандартные команды конфигурации, существенно упростив развертывание и масштабирование ИИ-решений.

Что такое TPU и почему важна программная часть

Тензорные процессоры Tensor Processing Unit представляют собой специализированные чипы разработки Google, ориентированные на выполнение матричных вычислений. Нейронные сети тратят на подобные операции львиную долю времени, поэтому аппаратные ускорители эффективно справляются с обучением и обслуживанием моделей на масштабных серверах Google Cloud.

Однако главным камнем преткновения исторически являлось программное обеспечение. Архитектура TPU изначально развивалась вокруг TensorFlow и JAX, из-за чего пользователям PyTorch приходилось прикладывать гораздо больше усилий для интеграции в эту среду.

Прежний подход к работе

До недавнего времени запуск кода PyTorch на чипах TPU реализовывался через прослойку PyTorch/XLA. Данный метод работал, но обладал своими специфическими особенностями. В частности, тензоры обрабатывались по принципу ленивых вычислений, из-за чего выполнение кода не всегда происходило в ожидаемый момент.

Небольшие изменения в структуре могли провоцировать повторную компиляцию графов, что заметно усложняло отладку. Для решения этой проблемы ранее создавался инструмент Optimum TPU, обеспечивавший поддержку инференса и обучения моделей Hugging Face. Теперь же, с появлением нативной поддержки в TorchTPU, необходимость в Optimum TPU отпала.

Появление TorchTPU

Новый бэкенд интегрирует TPU напрямую в структуру PyTorch, позволяя разработчикам использовать привычный синтаксис. Решение поставляется со специализированными ядрами Pallas, оптимизированными под конкретное железо, поддерживает распределенное обучение и включает интеграцию с инструментом torch.compile для повышения производительности.

Инсталляция пакета на виртуальной машине TPU выполняется с помощью одной команды через менеджер pip с указанием специального индекса для загрузки CPU-сборки PyTorch, что исключает скачивание тяжелых библиотек CUDA, не востребованных на таких чипах.

Популярные библиотеки экосистемы на TPU: Transformers

Пятое поколение библиотеки Transformers проектировалось как универсальный источник архитектур моделей, способный функционировать на любом оборудовании. Теперь в этот список вошли и тензорные процессоры. Для загрузки модели на TPU достаточно прописать параметр конфигурации с указанием устройства, после чего инициализация происходит автоматически.

По умолчанию код выполняется в режиме eager mode. При необходимости можно задействовать компиляцию для ускорения работы или развернуть систему в распределенной конфигурации.

Генерация изображений и видео в Diffusers

Фреймворк Diffusers также получил статус платформы с первоклассной поддержкой TPU. В ходе отладки или прототипирования достаточно перенаправить пайплайн на устройство одной командой. Для этапов развертывания применяется дополнительный вызов компиляции графа.

Повторные запуски с аналогичным разрешением изображений используют скомпилированные графы повторно. Если модель превышает объем памяти одного чипа, механизм TensorParallelConfig распределяет слои между несколькими процессорами в процессе загрузки.

Оптимизированные вычислительные ядра

Специальный хаб ядер предлагает готовые компоненты для ускорения расчетов. Низкоуровневое управление аппаратной частью обеспечивается за счет ядер, написанных на языке Pallas, а также библиотек вроде Tokamax. Использование подобных ядер в составе моделей Mixture-of-Experts позволяет кратно увеличить скорость декодирования при работе с крупными батчами.

Тестирование и контроль качества

Разработчики непрерывно совершенствуют поддержку TPU в библиотеке transformers путем прогона автоматизированных тестов. Команда внедряет ночные сборки CI-систем для отслеживания результатов наряду с другими ускорителями вычислений.

Несмотря на то, что часть проверок пока завершается с ошибками, каждая проблема фиксируется в виде отдельного репозиторного отчета и оперативно устраняется. Подобный подход позволяет находить и исправлять внутренние баги в самом коде Transformers на ранних этапах.

Начало работы

Аппаратные ускорители TPU перестали быть второстепенным направлением для пользователей PyTorch. Установка пакета TorchTPU открывает прямой доступ к моделям из открытых репозиториев и их запуску на специализированном оборудовании.

Сообщество продолжает развивать экосистему и призывает разработчиков делиться возникающими трудностями для их скорейшего решения в рамках открытых репозиториев.