Google пусна TorchTPU с отворен код, предоставяйки на разработчиците на PyTorch нативен начин за изпълнение на модели върху тензорните процесори (TPU) на компанията. Google също така обяви поддръжка за TPU в машините за обслужване vLLM и SGLang, рамката за обучение TorchTitan, Ray и библиотеките на Hugging Face.
Екипите, които пишат моделите си на PyTorch и по подразбиране ги изпълняват на графични процесори (GPU), имат най-голяма полза. Google твърди, че тези работни натоварвания вече се прехвърлят към TPU – техните персонализирани чипове за изкуствен интелект – с минимални промени в кода. Ако това твърдение се потвърди на практика, екип, който преценява разходите за GPU, може да тества хардуера на Google, без предварително да пренаписва своя софтуер.
Какво променят разработчиците
TorchTPU следва стандартния работен процес на PyTorch. Разработчиците дебъгват в режим eager, където всяка операция се изпълнява веднага щом бъде извикана. След това превключват към torch.compile, което превръща модела в оптимизиран граф за продукционна среда. Примитивите за разпределено обучение са вградени.
Hugging Face сведе превключването до един аргумент. В Transformers v5 подаването на device_map=„tpu“ изпраща модела към TPU, а Google изпълнява тестовия пакет на Transformers върху TPU всяка вечер. Библиотеката Diffusers за генеративни модели също работи нативно с тензорен паралелизъм между множество чипове.
Meta заявява, че вече изпълнява работни натоварвания за големи езикови модели (LLM), реклами и препоръки върху TPU чрез TorchTPU в ключови производствени сценарии. „Не можем да бъдем ограничавани от една-единствена хардуерна архитектура“, заяви Яна ван Грьонен, директор по PyTorch инженеринг в Meta.
Данните на Google
Обучението с TorchTitan върху TorchTPU постигна до 52% оползотворяване на моделните FLOPs при Qwen 3 8B с почти линейно мащабиране върху до 8 192 чипа Ironwood TPU. Почти линейното мащабиране означава, че пропускателната способност е нараствала почти пропорционално на броя на чиповете. Оползотворяването на FLOPs на модела измерва каква част от теоретичния математически капацитет на чипа отива за полезна работа по обучение, така че 52% означава около половината.
Отделен резултат идва от llm-d – рутер с отворен код, който разпределя заявките за инференция в клъстери. Google пусна llm-d 1.0 съвместно с Red Hat и IBM. Zenteiq.ai съобщава, че llm-d в Google Kubernetes Engine е увеличил дневния добив от 1,47 милиарда на 5,97 милиарда токена на всеки 1 000 чипа TPU v6e, което представлява приблизително четирикратен ръст. Zenteiq отдава този напредък на маршрутизиране, което изпраща всяка заявка към реплика, която вече съхранява нейния контекст.
Все още в предварителен преглед
Два компонента остават във фаза на предварителен преглед (preview). Единият е бекендът за TPU, който Google предостави на VeRL – широко използвана рамка за обучение с подсилване (reinforcement learning). Другият е пакетът от агенти за TPU, който се интегрира в инструментите за програмиране с изкуствен интелект, за да автоматизира първоначалното внедряване, диагностиката и настройката на инференцията за модели на PyTorch. Google също така внедри LiteRT (бивш TensorFlow Lite) в екосистемата на PyTorch за изпълнение на модели на телефони и други устройства.
Кодът на TorchTPU е наличен безплатно в GitHub.