Абляция, которая сравнивает две смеси данных, может на самом деле измерять количество видеокарт. Звучит как ошибка стажёра, но именно так работает большинство загрузчиков данных: при смене числа GPU меняется порядок семплов, а вместе с ним и кривая лосса. Команда PyTorch выпустила Zephon, загрузчик для TorchTitan с эластичным детерминизмом, который убирает эту переменную из эксперимента.

Что такое Zephon

Zephon это загрузчик данных для обучения больших моделей, который готовит данные прямо внутри тренировочного задания и гарантирует одинаковые глобальные батчи при любом числе GPU. Токенизация, упаковка, шаффл и смешивание датасетов идут на свободных CPU-ядрах параллельно с обучением, поэтому отдельный офлайн-проход по корпусу становится не нужен.

Скрытая переменная в каждом эксперименте с данными

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

Дальше начинается тонкий момент. Сделать абляцию мало, надо ещё понять, что именно ты измерил. Обычный загрузчик раздаёт семплы по рангам, и при смене числа GPU меняется порядок, в котором данные попадают в модель. Значит, разница в лоссе между двумя прогонами может быть разницей не между смесями, а между железом, которое случайно досталось под задачу.

С приходом агентов проблема обостряется. Агент хорошо запускает десятки экспериментов параллельно и плохо замечает, что результат зависит от чего-то помимо изменённой переменной, например от количества GPU. Загрузчик, который меняет порядок данных при смене железа, подсовывает такому агенту ложный сигнал, а дальше решение принимается по шуму.

Чекпоинт на 16 GPU, продолжение на 8, 32 и 64

Авторы показывают масштаб проблемы на простом сценарии. Берут модель на 1B параметров в TorchTitan, обучают её на 16 GPU, снимают чекпоинт на шаге 150 и продолжают с этой точки на 8, 32 и 64 картах, рядом с исходным прогоном, отдельно со стандартным загрузчиком и отдельно с Zephon. Сценарий соответствует реальной практике: прогон словил отказ железа и поднялся на тех картах, которые оказались свободны.

Со стандартным загрузчиком каждая ветка после резюма получает свою последовательность семплов и свою кривую лосса. С Zephon кривые ложатся друг на друга, а SHA-256 хеш каждого глобального батча совпадает с исходным прогоном. Совпадает не «примерно», а побайтово.

Сколько стоил разброс: 0,82 против 0,05

Отдельный эксперимент показывает цену нестабильности. Ту же модель на 1B обучают с нуля на 20B токенов при 8, 16, 32 и 64 GPU со стандартным загрузчиком. Оценка DCLM Core v1 разъезжается до 0,82 балла в зависимости от числа карт, оценка по набору FineWeb до 0,64.

Чтобы понять масштаб, авторы сравнивают это с порогами, которые использовали сами создатели DCLM и FineWeb, когда оценивали методы отбора данных: разброс примерно вчетверо больше. То есть абляция двух смесей на двух разных аллокациях запросто меряет число GPU. Zephon на том же сценарии даёт разброс 0,05 и 0,011 балла соответственно, то есть практически укладывается в шум оценки.

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

Как выглядит пайплайн

Zephon просит описать три вещи: датасеты, WorkSource, который решает, как эти датасеты смешиваются и упорядочиваются в учебный план, и Pipeline из операторов, превращающих семплы в тензоры для тренировочного цикла.

Минимальный пример смешивает два текстовых датасета прямо из Hugging Face, а потом токенизирует, упаковывает, шаффлит и собирает в батчи.

from zephon import Pipeline
from zephon.io import Dataset
from zephon.work import StaticMixtureWorkSource, TokenEstimation

source = StaticMixtureWorkSource(
    datasets=[
        Dataset.from_path("wiki", "hf://Salesforce/wikitext/wikitext-2-raw-v1/train"),
        Dataset.from_path("news", "hf://fancyzhx/ag_news/default/train"),
    ],
    mixture={"wiki": 0.5, "news": 0.5},  # доля токенов, а не семплов
    token_estimation=TokenEstimation(),
    seed=42,
)
pipeline = (
    Pipeline(source)
    .options(io_options={"cache": {"enabled": True}})
    .tokenize(tokenizer_id="gpt2", field="text", parallelism=4)
    .pack_flat(max_length=2048, algorithm="wrap")
    .shuffle(buffer_size=1024, seed=42)
    .batch(microbatch_size=8)
)
for batch in pipeline:
    train_step(batch.to_training())

Смесь задаётся в долях токенов, а не семплов, и это не косметика. Датасеты хранятся как сырые строки, и соотношение токенов на байт у них разное. Утилита TokenEstimation замеряет это соотношение на небольшой выборке из каждого датасета и корректирует план, чтобы пропорции сошлись именно по токенам.

WorkSource отделён от операторов обработки, поэтому смесь можно менять, не трогая пайплайн, и наоборот. У каждого оператора своя степень параллелизма: в примере токенизация получает четыре воркера, потому что это самый дорогой шаг. Именно распараллеливание препроцессинга позволяет держать GPU загруженными. На полном vision-language пайплайне, где нужно декодировать картинки, токенизировать, упаковывать и собирать батчи, Zephon не отстаёт от Google Grain, самого быстрого загрузчика из протестированных, и при этом даёт более сильные гарантии детерминизма.

Три решения, на которых всё держится

Когда загрузчик токенизирует документы, упаковывает их в последовательности фиксированной длины и смешивает несколько датасетов на лету, порядок данных больше нельзя просто посмотреть по индексу на диске. Авторы называют такие пайплайны stateful n:m. Порядок зависит от состояния пайплайна в момент, когда семпл понадобился, и в большинстве загрузчиков это состояние нельзя перенести на другое число GPU, не изменив выход.

Существующие решения закрывают задачу частично. Mosaic Streaming и Google Grain держат глобальные батчи одинаковыми при разном числе GPU, но только когда порядок семплов можно вычислить заранее по индексу, то есть упаковка на лету отпадает. Megatron-Energon воспроизводит батчи на другом числе GPU, но только пока суммарное число воркеров загрузчика по всем GPU остаётся прежним.

Zephon снимает эти ограничения тремя решениями, которые надстраиваются друг над другом.

Первое: аккумуляторы принимают решения, воркеры выполняют работу. Каждый оператор разделён на две части. Однопоточный аккумулятор видит все записи по порядку и принимает все решения, зависящие от состояния, например какие последовательности упаковать вместе. Функция без состояния делает дорогую работу на любом числе параллельных воркеров. Воркеры могут закончить в произвольном порядке, поэтому Zephon нумерует каждую группу записей перед выдачей и восстанавливает порядок результатов. Из этого следует важное свойство: степень параллелизма становится чистым параметром производительности. Удвоили число воркеров токенизации, выход не изменился.

Второе: полосы вместо рангов. Данные делятся не между доступными GPU напрямую, а на фиксированное число логических потоков, которые называются lane. У каждой полосы своя упорядоченная последовательность семплов и своя копия состояния операторов, включая частично заполненные корзины упаковки. Полоса выдаёт одни и те же записи независимо от того, какой GPU ей достался. Когда число data-parallel рангов меняется, Zephon просто перераспределяет полосы между рангами, и глобальные батчи выходят теми же. Аккумуляторы делают каждую полосу детерминированной, полосы делают детерминированным весь загрузчик. Плата за это одна: число полос надо выбрать заранее.

Третье: сброс и повтор вместо сохранения состояния. Чекпоинт Zephon не содержит содержимого корзин упаковки или буферов шаффла, в нём записано только то, какие куски входа обрабатывались в момент снятия. Чтобы список оставался коротким, Zephon периодически прогоняет через пайплайн специальный sentinel, который велит каждому оператору с состоянием выдать всё накопленное и начать с чистого листа. При восстановлении Zephon проигрывает вход с последнего сброса, точно восстанавливая состояние операторов, и выбрасывает всё, что тренировочный цикл уже потребил. Чекпоинты получаются размером в несколько мегабайт независимо от того, как далеко ушёл прогон.

Подключение к TorchTitan

Интеграция заменяет стандартный загрузчик TorchTitan на Pipeline, собранный из рецепта данных: датасеты, веса смеси, настройки упаковки, шаффла и параллелизма. Всё остальное в прогоне остаётся прежним. Состояние Zephon попадает в чекпоинт TorchTitan рядом с моделью и покрывает все полосы, поэтому прогон можно продолжить на другом числе GPU из того же чекпоинта.

При запуске интеграция проверяет, что полосы делятся нацело между data-parallel репликами и что батч каждого шага оптимизатора делится нацело между полосами. Ошибка в конфигурации падает сразу, а не через несколько часов, когда выяснится, что данные молча поменялись. В репозитории есть CPU-демо, которое снимает чекпоинт потока данных, продолжает его на другом числе реплик и сверяет каждый глобальный шаг с непрерывным прогоном, плюс GPU smoke-тест с настоящей, хоть и крошечной, моделью.

Что учесть перед стартом

Число полос стоит выбирать как общее кратное тех data-parallel размеров, на которых вы собираетесь работать. Если в планах 8, 16 и 32 GPU, то 32 полосы покрывают все три варианта. Увлекаться не стоит: несколько тысяч полос при максимуме в 16 карт добавят накладных расходов на восстановление из чекпоинта, потому что лишние полосы на GPU замедляют резюм.

При смене числа рангов рецепт данных должен оставаться тем же: датасеты, токенизатор, seed, смесь, длина контекста и размер глобального батча. Менять можно только параметры производительности вроде параллелизма.

И главное ограничение: Zephon гарантирует данные, а не арифметику. Глобальные батчи будут идентичными при любом числе GPU, а лоссы очень близкими, но не обязательно совпадут бит в бит, поскольку порядок редукций с плавающей точкой в тренировочном цикле зависит от того, сколько карт делят работу.

Куда это движется

Одно направление авторы выделяют особо. Zephon позволяет выбрать runner, то есть способ распараллеливания операторов. Один из реализованных вариантов использует потоки. На Python 3.14t он забирает vision-language семплы примерно в 3,7 раза быстрее, чем процессный runner на стандартном Python. Авторы считают, что за этим будущее загрузки данных.

Второе направление, динамические учебные планы. Поскольку план отделён от пайплайна обработки, WorkSource может менять веса смеси по ходу прогона, не трогая ничего ниже по течению. Сейчас команда работает над конфигурируемым компонентом, который делает это, сохраняя описанные гарантии детерминизма.

Часто задаваемые вопросы

Zephon заменяет DataLoader в PyTorch?

Нет. Это отдельный загрузчик для тренировочных прогонов, интегрированный с TorchTitan. Для локальных экспериментов стандартный DataLoader остаётся на месте. Zephon нацелен на большие прогоны, где порядок данных важен, а конфигурация железа меняется.

Почему нельзя просто зафиксировать seed?

Seed фиксирует генерацию случайной последовательности, но при упаковке на лету и разном числе воркеров состояние операторов различается: частично заполненные корзины, буферы шаффла, границы упаковки. Одинаковый seed не даёт одинаковых глобальных батчей при смене числа воркеров или рангов.

Насколько большие чекпоинты у Zephon?

В экспериментах речь о нескольких мегабайтах. Чекпоинт хранит не содержимое буферов, а только те куски входа, которые обрабатывались в момент снятия. Периодический сброс через sentinel не даёт этому списку разрастаться, поэтому размер почти не зависит от длины прогона.

Итог

Zephon превращает порядок данных из скрытой переменной в управляемую константу. Аккумуляторы отделяют решения от работы, полосы отделяют данные от рангов, сброс с повтором заменяет сохранение состояния. Результат измеримый: разброс DCLM Core v1 на разных аллокациях падает с 0,82 до 0,05 балла.

Если вы запускаете абляции смесей и сравниваете прогоны на разных конфигурациях GPU, проверьте, не меряете ли вы железо вместо данных. В репозитории есть CPU-демо, которое сверяет шаги после резюма без единой видеокарты: его стоит прогнать до того, как тратить GPU-часы на очередную «улучшенную» смесь.