TorchTitan: один стек вместо двух для претрейна и agentic RL

TorchTitan: один стек вместо двух для претрейна и agentic RL

В обычном RL-пайплайне живут две копии одной модели: одна учится, вторая генерирует роллауты. Пока всё идёт гладко, расхождение между ними никто не замечает. Когда награда проседает, начинается самая неприятная часть отладки: непонятно, виноват алгоритм, функция награды или две копии, которые тихо разошлись в численных деталях. Обновление TorchTitan этого года отвечает на это радикально: копия должна быть одна.

TorchTitan это PyTorch-native платформа для экспериментов и крупномасштабного пре- и пост-обучения. Она задумана как минимальная clean-room реализация техник PyTorch с точками расширения, без форков фреймворка и без вендорских патчей. За год команда добавила поддержку свежих открытых архитектур, RL-контур TitanRL, библиотеку шардинга FlexShard, графовый режим обучения и новую схему распределённых чекпоинтов.

TitanRL: одна модель от претрейна до vLLM

Обычная схема RL-стека выглядит так: рядом с тренером стоит отдельный инференс-движок, и модель в нём приходится держать синхронной с обучаемой. Любое архитектурное изменение нужно протаскивать в оба места, а при падении награды вы тратите дни на выяснение того, какая из копий отклонилась. TitanRL убирает вторую копию: та же модель, что училась на претрейне, проходит пост-обучение и попадает внутрь vLLM без пересборки отдельной реализации.

Роллауты при этом асинхронные. Обучение и генерация масштабируются независимо, каждое своими формами параллелизма, поэтому узкое место на одной стороне не останавливает другую. Команда проверила контур на плотных моделях и на MoE, на сотнях шагов, на задачах с длинным контекстом, исполнением кода, компакцией контекста, собственным harness и self-play.

Два опубликованных прогона показывают, о каких масштабах речь. Первый: Qwen3.5-35B-A3B base, контекст 64K, 40 ускорителей GB300, задача SWE-Rebench easy, то есть правка реального кода в репозитории. Второй: обучение на Tmax 15k, оценка на terminal bench 2.1, тот же базовый чекпоинт, контекст 128K и 20 GB300. Это уже не игрушечный RL на математических примерах, а сценарии, где модель часами работает с инструментами.

Бит-в-бит логиты и конец спора о причине

Главная инженерная деталь TitanRL выглядит скромно, но меняет правила отладки. В поддерживаемых on-policy конфигурациях тренер и генератор выдают бит-в-бит идентичные логарифмы вероятностей, то есть расхождение между ними равно нулю на каждом шаге. Если вы включили новый алгоритм или ядро и награда поехала, причина точно в вашем изменении, а не в численном дрейфе между двумя движками.

Это не абстрактная гарантия. Совместное исследование с командой TorchTitan использовало такой режим, чтобы измерить численные расхождения в RL на математике и на терминальных агентах. Практический вывод простой: часть различий, которые раньше списывали на «шум RL», имеет вполне конкретный источник в виде разной арифметики трейнера и генератора.

Семь архитектур и рецепты вместо форков

Список поддержанных моделей вырос заметно: Muse Spark Open, Muse Glimmer, Qwen3-VL, Qwen3.5, Kimi K2.7, Kimi K3 и DeepSeek-V4. Каждая интеграция сделана по всему стеку 5D-параллелизма, а не только на уровне определения модели, поэтому новую архитектуру можно запускать сразу с тензорным, контекстным, пайплайновым и экспертным параллелизмом.

Вместо отдельных форков под каждую модель в TorchTitan появилась композируемая система конфигурации: типизированные Python-рецепты, трансформации конфига модели и точечные переопределения компонентов. Менять стратегию параллелизма, отдельный блок модели или аппаратно-специфичное ядро можно, не поддерживая собственную ветку фреймворка. Через те же интерфейсы подключаются ядра Attention Gym, включая Gated DeltaNet, Kimi Delta Attention и разреженное внимание, а также опциональные ядра Helion вроде fused RoPE, при совместимости с torch.compile.

Побочный эффект такого подхода полезнее самих моделей. Каждая новая архитектура пополняет общий набор переиспользуемых компонентов: пайплайны мультимодальных данных, упаковку сэмплов, общие Vision-блоки, совместимость с чекпоинтами Hugging Face, декларативный шардинг и численную валидацию. Следующая модель приезжает в уже собранный конвейер.

DeepSeek V3 671B на MLPerf: цифры

Самый тяжёлый опубликованный прогон это DeepSeek V3 671B с модулем MTP1 на бенчмарке MLPerf training, на внутреннем кластере GB300, со схемой квантования MXFP8. Масштабировались только оси PP, DP и EP.

На 256 GPU конфигурация PP2 с восемью виртуальными стадиями, DP 128 и EP 64 дала 5 324 токена в секунду на ускоритель при глобальных 1,36 млн токенов в секунду и 1 553,89 TFLOPs на GPU. На 512 GPU та же схема держит 5 262 токена в секунду на устройство и выходит на 2,69 млн токенов в секунду суммарно. Рост пайплайна до PP4 и PP8 стоит примерно 10% производительности на ускоритель: 4 780 токенов в секунду на GPU в конфигурации PP8 на 512 устройствах.

Отдельно проверяли численность. Кривая лосса на MLPerf совпадает с Megatron, и это важнее любого прироста: обвязка из десятка оптимизаций не должна менять сходимость, иначе все замеры throughput бессмысленны.

Muse Spark Open: 192K контекста и цена балансировки

Для релиза MuseSpark 1.2 OSS показали эффективность обучения на длинном контексте через масштабирование TP, EP и CP. Здесь интересна не столько пропускная способность, сколько разница между выученным роутингом и принудительно сбалансированным.

На 256 GPU и контексте 96K выученный роутинг дал 1 500 токенов в секунду на ускоритель при MFU 25,84%. Принудительная балансировка на том же железе выжала 1 839 токенов в секунду и MFU 31,68%. На 512 GPU с контекстом 192K и восемью стадиями контекстного параллелизма картина та же: 1 335 против 1 588 токенов в секунду на устройство, MFU 31,27% против 37,17%.

Разница почти в шесть процентных пунктов MFU это цена свободы роутинга. Для исследователя она окупается качеством, для продакшн-кластера превращается в счёт за электричество, и выбор здесь честнее делать с цифрами перед глазами, а не по интуиции.

FlexShard: Muon, блочная квантизация и минус 79% времени оптимизатора

FSDP2 шардит состояние модели построчно или постолбцово. Этого хватает большинству обучений, но не хватает распределённому оптимизатору Muon и блочной квантизации параметров: им нужны схемы раскладки, которых в FSDP2 просто нет. FlexShard это новая библиотека data parallelism с bucket-centric API, где каждый бакет параметров получает собственный Placement policy, а обмен между рангами перекрывается вычислениями.

Помимо знакомых раскладок FSDP1 и FSDP2, FlexShard предлагает Owned, когда один ранг владеет параметром целиком, и BlockShard с двухмерным выравниванием по плиткам. Для Muon это означает, что параметры шардятся через Owned и BlockShard, и перед обновлением оптимизатора больше не нужен All Gather. На DeepSeek V3 и 256 H100 такое изменение сокращает время работы оптимизатора до 79% по сравнению с чистой реализацией на FSDP.

Вторая история ещё интереснее по эффекту. При блочно-квантованных All Gather параметры собираются уже в FP8, и объём коммуникации падает примерно вдвое. Проблема FSDP2 в том, что границы шардинга не совпадают с границами блоков квантизации: приходилось либо собирать параметры в BF16 и терять экономию, либо делать FP8-сборку с дополнительными коллективами для масштабирующих множителей. BlockShard выравнивает раскладку по плиткам, то есть по определению совпадает со схемой квантизации. На Llama 3 405B и 256 H100 это даёт до 20% сокращения времени шага обучения.

GraphTrainer и два режима с одинаковой арифметикой

GraphTrainer это графовый режим обучения в TorchTitan, построенный вокруг явного графа тренировки. Он сосуществует с обычным eager-путём, и оба режима дают бит-в-бит совпадающие результаты, что позволяет переключаться без перепроверки экспериментов.

Зачем граф нужен: кросс-модульные слияния, отладка на уровне графа и поддержка неоднородного железа плохо выражаются в eager-исполнении. Для команды, которая гоняет десятки конфигураций параллелизма, возможность поднять граф и увидеть реальную структуру вычислений экономит время сильнее, чем ещё один процент throughput.

Связь, симметричная память и чекпоинты

TorchTitan заодно служит полигоном для коммуникационного стека PyTorch, потому что обучение на сотнях узлов нагружает инициализацию, коллективы, обработку отказов и регистрацию памяти одновременно. В PyTorch 2.14 NCCL-реализацию из torchcomms перенесли в c10d как новый опциональный бэкенд, который регистрируется под стандартным именем nccl и работает как drop-in замена. Он добавляет односторонние RMA-окна, переконфигурацию process group на месте для отказоустойчивости, жадную инициализацию коммуникатора в 2.15 и прямой доступ к операциям shrink, grow, suspend и resume.

MoE-слои в TorchTitan теперь умеют работать на Dist-MoE, новой распределённой библиотеке mixture-of-experts с готовыми рецептами DeepSeek V3 для Blackwell. Токены ходят между экспертами по симметричной памяти, и эта передача сливается с вычислениями экспертов. FSDP тоже может брать буферы для коммуникации из пула симметричной памяти, а NCCL в таком случае включает свои ядерные реализации, включая copy-engine all-gather при активной политике zero-CTA.

Чекпоинты изменились сильнее всего, если смотреть на цифры. Новая библиотека torch_checkpointing убирает синхронизацию при каждом сохранении из DCPv1: один коллектив перед первым сохранением для обмена метаданными, дальше коллективов нет вообще. На 512 H100 с DeepSeek V3 671B и интервалом сохранения 50 шагов блокировка тренера падает с 841,2 мс до 34,1 мс в среднем и с 941,46 мс до 36,6 мс по P99. Запись чекпоинта ускоряется с 384,4 секунды до 13,1, загрузка без перешардирования с 194,7 секунды до 79.

Причина не в микрооптимизациях, а в архитектуре. Классическая схема использует коллективы GLOO для синхронизации, и на масштабе больше 16 тысяч рангов это плохо масштабируется по числу узлов. Лёгкий барьер на каждое сохранение решает проблему, попутно позволяя писать чекпоинт целиком одним torch.save и подключать собственный storage-бэкенд.

Кто уже использует

Экосистема вокруг TorchTitan выросла за пределы одного репозитория. veRL взял его как тренировочный бэкенд, Miles от RadixArk добавил поддержку, Meshy использует как тренер по умолчанию, а тренировочный движок Archon в Areal во многом повторяет его архитектуру. Для фреймворка, который начинался как эталонная реализация техник PyTorch, это неплохой показатель: чужая инфраструктура доверяет ему обучение, а не только чтение кода.

Есть и обратная сторона. Стек быстро обрастает подсистемами: FlexShard, Dist-MoE, GraphTrainer, torch_checkpointing, TitanRL, Helion-ядра. Для команды, у которой 8 GPU, большая часть этого не нужна, а часть возможностей раскрывается только на сотнях ускорителей. Минимализм TorchTitan теперь относителен: он минимален по сравнению с промышленными комбайнами, но уже не помещается в голову за один вечер.

FAQ

TorchTitan заменяет Megatron-LM?

Нет, он решает другую задачу. Megatron-LM это зрелый промышленный стек с собственными CUDA-ядрами и своей кривой обучения, TorchTitan же остаётся чистым PyTorch без форков и повторяет его поведение. На MLPerf кривая лосса совпадает с Megatron, так что выбор идёт не по качеству, а по тому, готовы ли вы жить внутри PyTorch-экосистемы.

Зачем TitanRL, если уже есть veRL и Miles?

Они не конкуренты. veRL и Miles используют TorchTitan как тренировочный бэкенд, а TitanRL добавляет то, чего у них внутри не было: одну модель от претрейна до vLLM и бит-в-бит совпадение логитов между тренером и генератором. Если вы уже собрали RL на veRL, миграция не обязательна, но отладка станет проще.

Можно ли запустить TorchTitan на одной GPU?

Да, репозиторий рассчитан на старт с малого и постепенное масштабирование, рецепты лежат в открытом доступе. Другое дело, что FlexShard, распределённый Muon и Dist-MoE раскрываются на сотнях ускорителей, а на одной карте вы получите обычное обучение без преимуществ, ради которых эти подсистемы писались.

Итог

Главное в обновлении TorchTitan не очередная поддержанная модель и не рекорд на MLPerf, а исчезновение второй копии модели в RL. Когда тренер и генератор это один и тот же код с идентичной арифметикой, отладка перестаёт быть гаданием, а эксперименты с наградой и алгоритмом получают нормальную точку отсчёта. Всё остальное, включая минус 79% времени оптимизатора и чекпоинт за 13 секунд, делает эту работу дешевле.

Репозиторий открыт: github.com/pytorch/torchtitan. Если у вас есть кластер под RL, самый быстрый способ оценить обновление это взять готовый рецепт для DeepSeek V3 или Qwen3.5, прогнать его на своём железе и сравнить время шага с текущим стеком.

← Все записи