Анонс поддержки Ray для PyTorch на Google TPU

Ранее в этом году на конференции Ray Summit 2026 мы рассказывали о том, как стек PARK (PyTorch, модели ИИ, Ray и Kubernetes) становится отраслевым стандартом для промышленного ИИ. Ray уже давно является универсальной распределенной средой выполнения для рабочих нагрузок PyTorch на графических процессорах (GPU).

Сегодня Anyscale и Google Cloud анонсируют нативную поддержку Ray для PyTorch на Google TPU — базовой суперкомпьютерной инфраструктуре, на которой работают Gemini и крупномасштабные рабочие нагрузки ИИ в Google Cloud.

Опираясь на проект Google TorchTPU и KubeRay в Google Kubernetes Engine (GKE), эта интеграция обеспечивает нативную работу с PyTorch из расчета «одно устройство на воркер» для Cloud TPU. Команды платформ и специалисты по машинному обучению, работающие на GKE, теперь могут планировать, масштабировать и обслуживать распределенное обучение PyTorch, инференс vLLM и рабочие нагрузки с подкрепленным обучением (RL) как на GPU, так и на TPU в рамках единой плоскости управления Kubernetes.

Нативный PyTorch на Google TPU

Исторически сложилось так, что запуск рабочих нагрузок TPU на Ray требовал адаптации к модели выполнения Single-Program, Multiple Data (SPMD), разработанной в первую очередь для JAX, где один рабочий процесс Ray на хост-ВМ управлял всеми чипами TPU, подключенными к машине. , представленный в прошлом году в качестве нативного бэкенда PyTorch от Google для тензорных процессоров, обеспечивает нативную интеграцию с TPU с поддержкой нетерпеливого (eager) выполнения и расширение torch.distributed для запуска распределенных рабочих нагрузок на TPU.

С появлением TorchTPU мы усовершенствовали Ray так, чтобы он запускал по одному воркеру на каждый чип TPU, что соответствует модели «одно устройство на воркер», к которой разработчики PyTorch привыкли на GPU.

При планировании задач или акторов с помощью Ray и TorchTPU каждый воркер запрашивает одно логическое устройство TPU (resources={"TPU": 1}), точно так же, как воркер GPU запрашивает num_gpus=1:

• Изоляция устройств: точно так же, как Ray автоматически устанавливает переменную CUDA_VISIBLE_DEVICES для воркеров GPU, Ray изолирует и назначает отдельные устройства TPU воркерам с помощью TPU_VISIBLE_CHIPS.

Изоляция устройств: точно так же, как Ray автоматически устанавливает переменную CUDA_VISIBLE_DEVICES для воркеров GPU, Ray изолирует и назначает отдельные устройства TPU воркерам с помощью TPU_VISIBLE_CHIPS.

• Нативные устройства PyTorch и коллективные примитивы: внутри задачи или актора Ray перенос существующего кода PyTorch с GPU на TPU сводится к простой замене "cuda" на "tpu" и переключению бэкенда распределенной связи с "nccl" на "tpu_dist".

Нативные устройства PyTorch и коллективные примитивы: внутри задачи или актора Ray перенос существующего кода PyTorch с GPU на TPU сводится к простой замене "cuda" на "tpu" и переключению бэкенда распределенной связи с "nccl" на "tpu_dist".

• Групповое планирование с учетом топологии: GKE подготавливает многохостовые слайсы TPU в виде пулов узлов, соединенных через высокоскоростные межчиповые соединения (ICI). Чтобы использовать эту физическую топологию, API Ray предоставляет реализацию нативного планирования, которая гарантирует, что акторы планируются групповым образом на одном и том же слайсе TPU.

Групповое планирование с учетом топологии: GKE подготавливает многохостовые слайсы TPU в виде пулов узлов, соединенных через высокоскоростные межчиповые соединения (ICI). Чтобы использовать эту физическую топологию, API Ray предоставляет реализацию нативного планирования, которая гарантирует, что акторы планируются групповым образом на одном и том же слайсе TPU.

Ниже приведен пример распределенной операции all_reduce в PyTorch с использованием акторов Ray на TPU. За исключением запроса ресурсов и строк устройства/бэкенда, код идентичен актору Ray, запущенному на GPU:

TorchTrainer на TPU

Для распределенного обучения и тонкой настройки многие команды полагаются на TorchTrainer из Ray Train, чтобы организовывать рабочие нагрузки Distributed Data Parallel (DDP), Fully Sharded Data Parallel (FSDPv2) и DTensor без необходимости вручную писать шаблонный код для групп распределенных процессов.

Благодаря нативной поддержке TorchTPU в Ray Train вы теперь можете запускать существующие скрипты TorchTrainer на однохостовых и многохостовых слайсах Cloud TPU без каких-либо изменений в цикле обучения PyTorch.

При настройке TorchTrainer для TPU через ScalingConfig Ray Train автоматически управляет всей базовой оркестрацией TPU:

• Атомарное резервирование слайса TPU: указав use_tpu=True вместе с целевой топологией (например, "4x4") и типом ускорителя (например, "TPU-V6E" или "TPU-V7X"), TorchTrainer автоматически подготавливает SlicePlacementGroup под капотом. Это гарантирует, что все рабочие процессы обучения планируются групповым образом на смежных физических слайсах TPU через высокоскоростные ICI, предотвращая фрагментацию ресурсов.

Атомарное резервирование слайса TPU: указав use_tpu=True вместе с целевой топологией (например, "4x4") и типом ускорителя (например, "TPU-V6E" или "TPU-V7X"), TorchTrainer автоматически подготавливает SlicePlacementGroup под капотом. Это гарантирует, что все рабочие процессы обучения планируются групповым образом на смежных физических слайсах TPU через высокоскоростные ICI, предотвращая фрагментацию ресурсов.

• Автоматическая инициализация tpu_dist: Ray Train запускает по одному воркеру на каждое устройство TPU (resources_per_worker={"TPU": 1}), автоматически настраивает переменные среды распределенного окружения PyTorch (RANK, WORLD_SIZE, LOCAL_RANK, MASTER_ADDR, MASTER_PORT) наряду с адресами построителя слайсов TorchTPU и топологией сетки, а также инициализирует группу процессов tpu_dist.

Автоматическая инициализация tpu_dist: Ray Train запускает по одному воркеру на каждое устройство TPU (resources_per_worker={"TPU": 1}), автоматически настраивает переменные среды распределенного окружения PyTorch (RANK, WORLD_SIZE, LOCAL_RANK, MASTER_ADDR, MASTER_PORT) наряду с адресами построителя слайсов TorchTPU и топологией сетки, а также инициализирует группу процессов tpu_dist.

• Встроенные утилиты Ray Train: стандартные вспомогательные функции Ray Train, такие как ray.train.torch.get_device() и ray.train.torch.prepare_model(), автоматически определяют среду выполнения TorchTPU, перемещают вашу модель на назначенное устройство TPU и оборачивают ее в класс DistributedDataParallel от PyTorch.

Встроенные утилиты Ray Train: стандартные вспомогательные функции Ray Train, такие как ray.train.torch.get_device() и ray.train.torch.prepare_model(), автоматически определяют среду выполнения TorchTPU, перемещают вашу модель на назначенное устройство TPU и оборачивают ее в класс DistributedDataParallel от PyTorch.

Вот полный пример масштабирования цикла обучения PyTorch на 16-чиповом многохостовом слайсе Trillium (TPU-V6E, топология 4x4) с использованием TorchTrainer:

Подробное пошаговое руководство для Google Kubernetes Engine (GKE) см. в разделе Обучение моделей на многохостовых TPU с помощью Ray и TorchTPU.

vLLM-TorchTPU с Ray Serve

Для онлайнового инференса

регистрирует TorchTPU в качестве плагина платформы vLLM, чтобы модели PyTorch компилировались и выполнялись на Cloud TPU без переписывания определений моделей на JAX. В сочетании с

вы можете развертывать совместимые с OpenAI конечные точки vLLM на однохостовых и многохостовых слайсах TPU, используя тот же декларативный API LLMConfig, который применяется на GPU.

Когда accelerator_type и accelerator_config задают слайс TPU, Ray Serve резервирует SlicePlacementGroup с учетом топологии для каждой реплики и запускает воркеры движка vLLM на всем слайсе:

Подробное пошаговое руководство для Google Kubernetes Engine (GKE) см. в разделе Обслуживание LLM на многохостовых TPU с помощью Ray и TorchTPU.

Расширение стека PARK на TPU

По мере того как нагрузки ИИ эволюционируют от автономных задач обучения до составных систем, объединяющих обучение, высокопроизводительный вывод и обучение с подкреплением, организации стандартизируют стек с открытым исходным кодом PARK (PyTorch, модели ИИ, Ray и Kubernetes). Интеграция TorchTPU в Ray и GKE воплощает это видение в Google Cloud: разработчики PyTorch сохраняют уже знакомую им модель нетерпеливого выполнения (eager execution) и API Ray, а платформенные команды получают топологически-освещенное автомасштабирование, отказоустойчивость и управление несколькими ускорителями в GKE.

Чтобы начать работу уже сегодня, ознакомьтесь с пошаговыми руководствами по GKE для развертывания LLM или обучения моделей с использованием Ray и TorchTPU, и следите за новостями, поскольку мы расширяем эту основу для крупномасштабных фреймворков пост-обучения на основе RL и сегментов Ironwood (TPU v7x) нового поколения в GKE.