Hugging Face ❤️ TorchTPU: TPU-та, нативно

Назад към статиите

Hugging Face ❤️ TorchTPU: TPU, нативно

Публикувано 20 октомври 2026 г.

Актуализиране в GitHub

Глас „за“

Ако някога сте опитвали да стартирате PyTorch на TPU, вероятно сте пренасяли кода си към torch_xla или в крайна сметка ви се е налагало да учите JAX. С TorchTPU, нативен бекенд за PyTorch, TPU вече е просто още едно устройство. Това отваря възможност за нативна поддръжка на TPU в екосистемата на Hugging Face. Transformers, Diffusers и Kernels работят нативно на TPU. Изберете модел от Hub, изпратете го към "tpu" и сте готови.

Чакайте, какво е TPU? 🤔

TPU означава Tensor Processing Unit (тензорен процесор). TPU са собствени AI ускорители на Google. Това са чипове, създадени за една задача: матричните изчисления, с които невронните мрежи прекарват по-голямата част от времето си. Google ги разработва от около десет години и ги използва за обучение и обслужване на собствените си модели, включително Gemini.

TPU се предлагат в Google Cloud — от един чип до модули с хиляди чипове. Те са отлични, когато ви е необходима голяма изчислителна мощ за обучение или обслужване в мащаб.

Уловката винаги е била софтуерът. Стекът за TPU се разви около TensorFlow и JAX, а потребителите на PyTorch трябваше да положат повече усилия, за да се включат.

Старият начин

Досега стартирането на PyTorch върху TPU означаваше използване на PyTorch/XLA (torch_xla). Работеше, но имаше своите особености. Тензорите бяха лениви (тензорите, не вие), така че кодът ви не винаги се изпълняваше, когато очаквахте. Малки промени можеха да задействат повторно компилиране на графа, а отстраняването на грешки беше трудно за разбиране.

За да улесним живота, създадохме Optimum TPU. Той предостави на моделите на Hugging Face поддръжка за инференс и обучение върху TPU. Сега, когато TorchTPU предлага нативна поддръжка, Optimum TPU вече не е необходим за стартиране на модели на Hugging Face с Transformers и Diffusers.

На сцената излиза TorchTPU 🚀

TorchTPU добавя TPU като нов бекенд към PyTorch. Пишете обикновен PyTorch код. Той включва Pallas ядра, оптимизирани за TPU, поддръжка за разпределено обучение и интеграция с torch.compile, когато искате по-висока скорост. Инсталирането му на TPU VM изисква една команда:

pip install torch-tpu --extra-index-url https://download.pytorch.org/whl/cpu

Допълнителният индекс ви предоставя CPU версията на torch. Без него pip изтегля много по-голямата CUDA версия, от която нямате нужда при TPU.

Любимите ви библиотеки — вече и на TPU

Transformers

Transformers v5 е създаден около една идея: да бъде източникът на истина за дефинициите на моделите, така че те да могат да работят навсякъде. „Навсякъде“ вече включва и TPU. За да заредите модел на TPU, подайте device_map="tpu". Това е всичко.

from transformers import AutoModelForCausalLM, AutoTokenizer

model_id = "Qwen/Qwen3.5-9B"
model = AutoModelForCausalLM.from_pretrained(
    model_id,
    device_map="tpu", 
)
tokenizer = AutoTokenizer.from_pretrained(model_id)
inputs = tokenizer("The future of AI is", return_tensors="pt").to("tpu")
out = model.generate(**inputs, max_new_tokens=64)
print(tokenizer.decode(out[0], skip_special_tokens=True))

По подразбиране моделът работи в eager режим. Можете също да използвате torch.compile за по-добра производителност или да го стартирате в разпределена конфигурация.

Изображения? Видеоклипове? Diffusers ги рисува върху TPU 🎨

TPU вече е пълноправно устройство в Diffusers. За прототипиране или отстраняване на грешки в eager режим е необходим само pipe.to("tpu"). За внедряване компилираният режим добавя едно извикване на compile. Конвейерите и кодът ви остават същите и няма TPU разклонение, което да поддържате.

import torch
from diffusers import StableDiffusion3Pipeline

pipe = StableDiffusion3Pipeline.from_pretrained(
    "stabilityai/stable-diffusion-3.5-medium", dtype=torch.bfloat16
)
pipe.to("tpu")  

pipe.transformer.compile(backend="tpu", fullgraph=True, dynamic=False)

image = pipe("a corgi astronaut planting a flag on a pizza moon").images[0]

Първото извикване компилира, а всяко следващо със същия размер на изображението използва повторно графа.

Модели, твърде големи за един чип? Подайте TensorParallelConfig към from_pretrained и трансформърът ще бъде разделен между чиповете при зареждането, така че целият модел никога не трябва да се побира на едно място. Много популярни конвейери го поддържат готово, а при тези, които не го правят, е достатъчно да добавите tp_plan. Ръководството за TPU на Diffusers съдържа пълната рецепта.

Ето няколко примера, стартирани на Cloud TPU v6e-4:

_{TP=4 (всеки модел е разделен между четири чипа) · bf16 · batch 1}

MiniMax-H3 (видео + аудио)	Qwen-Image-2.1

_{"The fluffy plush Hugging Face mascot from the picture, a round yellow smiling face with rosy cheeks and two little open hands, stands at the prow of an ancient Greek trireme […] Sound of crashing waves, creaking timber, wind in the rigging, distant seagulls, and a low epic orchestral swell." · 1344×768, 5.2 s, 49 стъпки · 188 s компилирано, 262 s eager}	_{"A capybara wearing a wizard hat, reading a book by candlelight, oil painting" · 1024×1024, 40 стъпки · 3.5 s компилирано, 8.0 s eager}	_{"A neon shop sign that reads 'QWEN IMAGE 2.1', rainy night, reflections on wet pavement" · 2048×2048, 40 стъпки}

Kernels

Kernels Hub предлага готови за употреба TPU ядра. Заредете едно чрез библиотеката Kernels и ускорете модела си.

TPU ядрата са написани на Pallas, която ви дава ниско ниво на контрол върху хардуера. Можете също да използвате ядра от библиотеката Tokamax. Ето ядро ragged-dot от Tokamax, включено в модел Mixture-of-Experts:

from kernels import get_kernel
from transformers import AutoModelForCausalLM
from transformers.integrations.moe import ALL_EXPERTS_FUNCTIONS

repo_id = "tengomucho/ragged-dot-tpu"
kernel = get_kernel(repo_id, version=1, trust_remote_code=[repo_id])
ALL_EXPERTS_FUNCTIONS.register("tokamax_ragged_dot", kernel.experts_forward)
model = AutoModelForCausalLM.from_pretrained(
    "Qwen/Qwen3-30B-A3B-Instruct-2507",
    experts_implementation="tokamax_ragged_dot",
)

Колкото по-голям е пакетът, толкова по-голяма е печалбата. При размер на пакета 64 декодирането е 21,5 пъти по-бързо при Qwen3-30B-A3B и 12,8 пъти по-бързо при Gemma 4 26B-A4B.

Пропускателна способност при декодиране на TPU v6e-8, ядро от Hub спрямо реализацията по подразбиране.

Тестваме го, за да не се налага на вас 🧪

Непрекъснато подобряваме поддръжката на TPU в transformers, като стартираме тестовия пакет на библиотеката. Работим по добавянето на нощна CI система на TPU и публикуването на резултатите в (таблото за управление на transformers)[https://huggingface.co/spaces/transformers-community/transformers-ci-dashboard].

CI таблото за управление на Transformers с резултатите за TPU редом до тези на други ускорители. Работата продължава!

Все още не всичко е зелено и това ни устройва. От тестовете, които се изпълняват на TPU, 96,5% преминават успешно. Всеки открит проблем се превръща в малък скрипт за възпроизвеждане и issue в хранилището на TorchTPU. Досега сме докладвали 34 проблема, а 19 вече са отстранени. Стартирането на TPU откри и няколко грешки в самите Transformers, например в StaticCache и плановете за tensor-parallel на Qwen3. Отстранихме ги upstream.

Ако нещо се повреди на TPU, искаме да разберем за това преди вас.