Обучение speculative decoding в промышленном масштабе на PyTorch
За последний год большие языковые модели стремительно выросли в масштабах и возможностях. Флагманские модели вроде Kimi K2.5, GLM 5 и Qwen 3.5 насчитывают сотни миллиардов параметров и контекстные окна в миллионы токенов. Это позволяет работать с длинным контекстом, агентными workflow и сложным использованием инструментов. По мере роста возможностей моделей эффективный инференс стал одной из ключевых системных проблем при развёртывании LLM.
Speculative decoding — один из самых эффективных способов ускорения генерации. Лёгкая draft-модель предлагает несколько токенов вперёд, а более крупная target-модель проверяет их за один проход данных через модель (forward pass). Когда предсказания принимаются, несколько токенов генерируются одновременно. Это улучшает throughput и снижает latency. Недавние подходы вроде MTP (Multi Token Prediction) и EAGLE-3 показывают, что хорошо обученные draft-модели дают стабильное ускорение.
Важный аспект обучения draft-модели — передача информации от target-модели через промежуточные скрытые состояния (внутренние числовые данные в слоях нейросети). По мере роста флагманских моделей возникает новый системный бутылочный горлышко: эффективная передача огромного объёма скрытых состояний от target-модели к draft-модели. Например, EAGLE-3 опирается на 3 слоя скрытых состояний target-модели. При обучении EAGLE-3 draft-модели для Kimi K2.5 один обучающий пример на 128K токенов требует ~7 ГБ скрытых состояний от target-модели. На масштабах датасета это становится неподъёмным.
Существующие пайплайны обычно идут одним из двух путей. Первый — предварительно вычислить скрытые состояния и сохранить на диск. Это ведёт к огромным требованиям к хранилищу и серьёзной I/O-нагрузке. Второй — совместить инференс и обучение, генерируя скрытые состояния на лету. Это избавляет от записи на диск, но требует размещения target-модели на тех же GPU, что и обучающий воркер. Возникает огромное давление на видеопамять.
Чтобы решить эти проблемы, мы представляем TorchSpec — нативный для PyTorch фреймворк для disaggregated (разнесённого) speculative decoding training. TorchSpec разделяет инференс-систему, генерирующую скрытые состояния, и обучающую систему, которая их потребляет. Вместо записи на диск скрытые состояния стримятся напрямую от группы инференс-движков к группе обучающих воркеров через центральное хранилище Mooncake по RDMA (Remote Direct Memory Access) или TCP. Это исключает дисковое хранилище и позволяет независимо масштабировать ресурсы инференса и обучения.
С помощью TorchSpec мы успешно обучили EAGLE-3 draft-модель для Kimi K2.5, затратив 1500 GPU-часов на H200 и обработав 600K обучающих примеров — 6 миллиардов токенов. Draft-модель показывает сильные результаты на различных бенчмарках:
*draft-модель обучена с lookahead=4 (модель заглядывает на 4 токена вперёд)
С обученной draft-моделью throughput вывода увеличивается более чем на +60% при batch size 1, +30% при batch size 8 и +26% при batch size 16 при lookahead из 3 токенов.
Предпосылки
Сегодня распространены два подхода к обучению speculative decoding:
- Обучение совместно с инференсом
- Офлайн-подготовка скрытых состояний
Оба работают на умеренных масштабах, но сталкиваются с проблемами при росте размера draft-модели и длины контекста.
Обучение совместно с инференсом
При совместном размещении target-модель и draft-модель делят одни и те же GPU. Target-модель делает forward pass для генерации скрытых состояний и логитов (сырых выходных оценок модели), которые сразу потребляются draft-моделью для обучения. Из-за тесной связанности возникают ограничения:
- Жёсткое шардирование (разбиение модели на части): стратегия параллелизма draft-модели привязана к target-модели. Если target использует TP=4 (Tensor Parallelism — модель разбита на 4 GPU), draft тоже должен использовать ровно 4 ранга. Даже если другая конфигурация была бы эффективнее для его меньшей архитектуры.
- Невозможность независимого масштабирования: текущие фреймворки совместного размещения обычно не поддерживают кросс-нодовое шардирование, ограничивая обучение GPU внутри одной ноды. Инференс и обучение привязаны к одному и тому же объёму ресурсов.
- Давление на память: target-модель занимает значительную часть GPU-памяти, оставляя draft-модели мало места для обучения.
Анализ памяти при совместном обучении с Kimi K2.5 (1T параметров, MoE — Mixture-of-Experts, 384 эксперта, ~575 ГБ весов модели):
| GPU | Общая память (8 GPU) | Веса модели | Шард на GPU | Остаток на GPU |
|---|---|---|---|---|
| 8×H200 | 1 128 ГБ | ~575 ГБ | ~72 ГБ | ~69 ГБ |
| 8×H100 | 640 ГБ | ~575 ГБ | ~72 ГБ | ~8 ГБ |
Хотя draft-модель обычно невелика, современные методы вроде Training-Time Testing (TTT) требуют много памяти. Они сохраняют промежуточные активации для нескольких speculative-шагов. Накопление активаций увеличивает общий расход памяти. С 8 ГБ можно обучаться только при длине контекста 4096.
Офлайн-подготовка скрытых состояний
Офлайн-подход предварительно вычисляет скрытые состояния target-модели, сериализует их на диск и загружает позже для обучения draft-модели. Это разделяет инференс и обучение, но создаёт серьёзную проблему хранения — особенно для больших моделей с длинным контекстом.
Анализ хранения для Kimi K2.5 (hidden_size=7168, vocab_size=163 840):
На один пример при длине контекста = 131 072 токенов:
| Тензор | Форма | Dtype | Размер |
|---|---|---|---|
| Скрытые состояния (3 aux-слоя) | (131072, 21504) | bf16 | 5.25 ГБ |
| Последние скрытые состояния | (131072, 7168) | bf16 | 1.75 ГБ |
| Input IDs | (131072,) | int64 | 1 МБ |
| Итого на пример | ~7.0 ГБ |
Примечание: target-логиты можно пересчитать из последних скрытых состояний через lm_head (выходной слой, преобразующий скрытые состояния в вероятности токенов), поэтому их не нужно хранить. Но даже без них требования к хранилищу быстро растут:
| Размер датасета | Требуется хранения |
|---|---|
| 10 000 примеров | 70 ТБ |
| 30 000 примеров | 210 ТБ |
| 100 000 примеров | 700 ТБ |
На таких масштабах распределённые файловые системы испытывают серьёзную нагрузку. Особенно когда несколько обучающих запусков работают параллельно и конкурируют за I/O-пропускную способность. Сериализация и десериализация также существенно замедляют обучение.
TorchSpec: disaggregated обучение draft-модели
TorchSpec идёт другим путём: полное разделение инференса и обучения. Target-модель работает на выделенных GPU для инференса. Draft-модель обучается на отдельных GPU для обучения. Тензорные данные стримятся между ними через высокоскоростной протокол RDMA или TCP через хранилище Mooncake.
Эта архитектура решает ключевые проблемы:
- Гибкое, независимое масштабирование. Количество GPU для инференса и обучения полностью независимо. Можно добавить инференс-движки для большей пропускной способности генерации скрытых состояний. Либо увеличить количество обучающих GPU для большего FSDP-шардирования (разбиения данных и весов по GPU через Fully Sharded Data Parallel) и более крупных global batch.
- Вся память — для обучения. Обучающие GPU полностью выделены под draft-модель, максимизируя доступную память для длинных последовательностей и больших батчей.
- Нулевые затраты на хранение. Скрытые состояния стримятся напрямую от инференса к обучению через RDMA/TCP. Данные не выгружаются на диск, что устраняет нагрузку на файловую систему и издержки сериализации.
Почему Mooncake?
Mooncake изначально разработан Moonshot AI и Университетом Цинхуа как движок передачи данных для управления KV cache (кэшем ключей и значений — промежуточными данными, ускоряющими генерацию) в production-сервисе LLM. С тех пор проект вырос в активное сообщество в экосистеме PyTorch. Mooncake обрабатывает высокопроизводительную кросс-нодовую передачу данных через различные сетевые протоколы и управляет жизненным циклом памяти. Именно эти возможности нужны TorchSpec для надёжной и эффективной стриминговой передачи скрытых состояний.
Ключевые свойства Mooncake:
- RDMA + TCP с единым API. Трансферы на скорости, близкой к линейной, на InfiniBand/RoCE-кластерах. Работает из коробки по TCP, когда RDMA недоступен — без изменений в коде.
- GPU Direct RDMA. Передаёт данные напрямую в GPU-память, минуя CPU-стейджинг (промежуточное копирование через оперативную память процессора). Это критично, когда каждый обучающий пример содержит гигабайты скрытых состояний.
- Zero-copy трансферы. Тензоры упаковываются в зарегистрированные pinned-memory буферы (блоки оперативной памяти, закреплённые от перемещения ОС) и передаются напрямую — без сериализации и промежуточных копирований.
- Production-надёжность. Проверен в крупномасштабных production-развертываниях, что даёт TorchSpec стабильную основу для длительного мульти-нодового обучения.
Поддержка длинного контекста
Когда вся память выделена под draft-модель, TorchSpec поддерживает длины последовательностей, недостижимые при совместном размещении в обучении EAGLE-3. Например, Kimi K2.5 потребляет 72 ГБ памяти при совместном обучении. С lookahead=4 и disaggregated-обучением один H100 может обучаться на последовательностях до 44K токенов, а один B200 — масштабироваться до 200K токенов.
Помимо разделения ресурсов, TorchSpec использует реализацию, нативную для инференс-движков: скрытые состояния генерируются непосредственно production-инференс-движками. Это даёт два ключевых преимущества:
- Инференс-обучение согласованы. Форматирование шаблонов, токенизация и ядра полностью совпадают. Нет разрыва между обучающей средой и средой развёртывания.
- Нативная поддержка моделей через движок. Для поддержки новой архитектуры target-модели минимальны изменения на стороне обучения. Сейчас TorchSpec поддерживает vLLM и SGLang, поддержка TensorRT LLM скоро появится. Если инференс-движок поддерживает модель, TorchSpec может обучить draft-модель из коробки. Это включает:
- Новые архитектуры (MoE, мультимодальные и т.д.)
- Квантованные модели (FP8, INT4 и т.д.)
- Sparse attention, варианты RoPE (Rotary Positional Embeddings) и другие специфичные фичи моделей
Обучение на декоде.
Draft-модели обычно работают лучше всего, когда обучены на токенном распределении target-модели. Распространённый подход — оставить оригинальные промпты датасета и перегенерировать ответы через target-модель как подготовительный шаг. Но этот двухэтапный процесс неудобен для исследователей и инженеров. Благодаря нативной интеграции с движком мы можем генерировать выходы авторегрессионно (пошагово, токен за токеном) из промптов прямо во время обучения.
Кейс: обучение EAGLE-3 модели для Kimi K2.5
Kimi K2.5 представляет сложный сценарий обучения, хорошо иллюстрирующий ценность disaggregated-подхода.
Проблема:
- Масштаб модели: Kimi K2.5 требует минимум 8×H200 или 16×H100 GPU для инференса target-модели. При совместном размещении остаётся очень мало памяти для обучения draft-модели.
- Длинный контекст: Kimi K2.5 нацелена на агентные и reasoning-задачи с длинным контекстом, требуя обучения на последовательностях до 200 000 токенов.
- Большой словарь: 163 840 токенов при hidden dimension 7 168.
Решение TorchSpec
С TorchSpec мы рекомендуем развернуть Kimi K2.5 на 8×H200 GPU как выделенный инференс-движок. EAGLE-3 draft-модель при этом обучается на отдельном кластере из 8×H200 GPU. Инференс-кластер располагает полной памятью для обслуживания и генерации скрытых состояний. Обучающий кластер имеет полную GPU-память для draft-модели, что позволяет обучаться на контексте в 100 000 токенов по 600K примерам.
Скрипты: предоставляем два готовых скрипта для обучения draft-модели Kimi K2.5:
- 3 ноды по 8×H100 с TP=16 для инференса и TP=8 для обучения: kimi-k25-3node-h100
- 2 ноды по 8×H200 с TP=8 для инференса и TP=8 для обучения: kimi-k25-2node-h200
Обучающий датасет: мы открыли наш миксированный датасет: kimi-600k-training-dataset.
Draft-модель: мы открыли нашу draft-модель: kimi-k2.5-eagle3.
Roadmap
TorchSpec активно развивается. Ключевые направления:
- Расширение покрытия моделей: поддержка популярных моделей вроде Minimax M2.5, Qwen 3.5 и дообучение MTP-слоя от GLM 5.
- Packed sequence training: упаковка нескольких коротких последовательностей в один обучающий пример для максимизации утилизации GPU и снижения потерь от padding. Особенно полезно для датасетов с переменной длиной.
- Дополнительные алгоритмы обучения: выход за рамки EAGLE-3 в сторону других подходов вроде DFlash и MTP, расширяя спектр архитектур draft-моделей.
- Интеграции с движками: подключение других популярных инференс-движков (например, TensorRT LLM), чтобы пользователи могли выбрать любой подходящий движок.
Благодарности
Мы благодарим следующие команды и соавторы:
Команда TorchSpec и сообщество: *Yubo Wang, *Yinghui Liu, Shirley Wu, Junxiong Wang, Qingyang Wu, Bobbie Bie, Fan Yin, Chao Wang, Weicong Wu, Jue Wang
Команда Mooncake: *Jiaqi Liao, Mingxing Zhang