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 — це власні ШІ-прискорювачі 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. У посібнику Diffusers для TPU наведено повний рецепт.

Ось кілька прикладів, запущених на Cloud TPU v6e-4:

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

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

_{«Пухнастий плюшевий талісман Hugging Face із зображення — кругле жовте усміхнене обличчя з рожевими щоками та двома маленькими розкритими руками — стоїть на носі давньогрецької триреми […] Звук розбиваючих хвиль, скрип деревини, вітер у такелажі, далекі крики чайок і тихе епічне наростання оркестру». · 1344×768, 5,2 с, 49 кроків · 188 с у скомпільованому режимі, 262 с у eager-режимі}	_{«Капібара в капелюсі чарівника читає книгу при свічках, олійний живопис» · 1024×1024, 40 кроків · 3,5 с у скомпільованому режимі, 8,0 с у eager-режимі}	_{«Неонова вивіска магазину з написом “QWEN IMAGE 2.1”, дощова ніч, відображення на мокрому тротуарі» · 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 декодування на Qwen3-30B-A3B відбувається у 21,5 раза швидше, а на Gemma 4 26B-A4B — у 12,8 раза.

Пропускна здатність декодування на 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. Ми виправили їх у головній гілці.

Якщо щось ламається на TPU, ми хочемо дізнатися про це раніше за вас.