Обучение на последовательностях длиннее 1 млн токенов

Агент способен работать над задачей в течение нескольких часов: читать файлы, выполнять команды и сохранять контекст каждого хода сессии. Накопленная сессия достигает сотен тысяч токенов, именно поэтому передовые модели теперь заявляют о контекстных окнах в миллион токенов.

Чтобы модель хорошо справлялась с таким объёмом контекста, её нужно обучать на последовательностях такой длины, и именно здесь начинаются трудности: последовательность из миллиона токенов не помещается на одном GPU, и даже на восьми — без дополнительной помощи.

Пример в этом руководстве обучается ровно на одной такой последовательности за шаг на одном узле с 8 GPU.

Запустите

Вам понадобится один узел с 8 H100 (или лучше) и transformers>=5.17.0: в примере используется offload для gradient checkpointing, добавленный в этом релизе.

torchrun --nproc_per_node 8 examples/sft_qwen3_8b_1m_context/sft_qwen3_8b_1m_context.py

Скрипт выполняет дообучение Qwen3-8B на книгах из PG-19, которые склеиваются последовательно, пока каждая не достигнет примерно миллиона токенов. После примерно десяти минут загрузки и токенизации выполняется первый шаг:

{'loss': 4.311, 'grad_norm': 29.25, 'num_tokens': 1049000.0, 'epoch': 0.25}
  8%|▊ | 1/12 [06:20<1:09:41, 380.10s/it]

Стоит обратить внимание на три числа:

• num_tokens равно 1.049e6. Это одна последовательность, а не батч из коротких. Каждый токен обращается к каждому предыдущему токену на протяжении всего миллиона.
• Значение функции потерь начинается с 4,3, что является нормальным начальным значением. Запуск, в котором незначительно неправильно настроена работа с длинным контекстом, начинается примерно со значения 10, и мы вернёмся к объяснению причин.
• 380 с на шаг. Один шаг — это одна последовательность, то есть чуть более шести минут на последовательность.

Наивный подсчёт показывает, что для этой последовательности требуется 288 ГБ на каждый GPU. Остальная часть этого руководства посвящена тому, как эта цифра снижается до 56.

Что только что произошло

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

Потеря — это первое, что ломается

Возьмём Qwen3-4B на одной карте с 80 ГБ памяти и увеличим длину последовательности. Первым, что исчерпает память, будет не модель, а loss. Проверьте профиль использования памяти.

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

Последний слой декодера выводит скрытое состояние формы$L \times H$, где$L$ — длина последовательности, а$H$ — размер скрытого слоя. Затем это скрытое состояние умножается на веса головы языковой модели, чтобы получить матрицу логитов формы$L \times V$, где$V$ — размер словаря. Эта матрица огромна, так как последовательность длинная, а словарь большой (более ста тысяч элементов). Затем на основе этой матрицы и целевых меток вычисляется функция потерь с использованием перекрёстной энтропии, и именно её материализация является пиком.

Можно ли вычислить потерю, не материализуя всю матрицу логитов? Да, можно обрабатывать данные по частям.

Массивная матрица логитов: разбейте её на части

Вместо того чтобы вычислять логиты для всей последовательности сразу, мы берем по несколько сотен строк за раз, вычисляем потерю для них и суммируем полученные значения. Матрица логитов целиком никогда не удерживается в памяти. Строки — это измерение, которое разделяется без проблем: потеря представляет собой сумму по строкам, а softmax внутри нее выполняется вдоль словаря одной строки, поэтому блок строк содержит все, что нужно для вычисления собственной потери.

Одно большое умножение превращается в несколько меньших, а экономия памяти настолько значительна, что потеря производительности больше не становится пиковой. TRL использует 256 строк на чанк, что позволяет каждому умножению оставаться достаточно большим для эффективного использования GPU.

Новый профиль выглядит так.

Пик исчез.

Как включить это в TRL? Ничего делать не нужно: chunked loss используется по умолчанию. Если вы хотите использовать обычный loss, установите loss_type="nll" в SFTConfig.

Насколько это позволяет нам увеличить длину последовательности? С точки зрения памяти, только это одно изменение позволяет нам перейти от 32k токенов к 160k.

Модель никогда не видела позицию 100 000: перемасштабируйте позиции с помощью YaRN

Теперь в памяти помещается 160k токенов. Но может ли модель действительно их прочитать?

Qwen3-4B обучался на последовательностях из 40 960 токенов. Об этом сказано в конфигурации:

>>> from transformers import AutoConfig
>>> AutoConfig.from_pretrained("Qwen/Qwen3-4B").max_position_embeddings
40960

Каждая позиция за этим пределом — это позиция, с которой модель никогда не сталкивалась. Ничего не падает, но запуск просто обучается на токенах, которые модель не может «разместить». Вот потери по каждому токену последовательности из 160k токенов:

Модель не «падает» на отметке 40 960. Она продолжает работать ещё какое-то время и начинает терять нить повествования лишь после примерно 70k. После этого значение стабильно растёт, в итоге превышая энтропию униграмм для того же текста (6.45): это хуже, чем у модели, которая полностью игнорирует контекст и предсказывает на основе только частот токенов.

Каждый токен сопровождается своей позицией, и модель выучила, что означают позиции в диапазоне от 0 до 40 960. Если попросить её о позиции 100 000, у неё не будет никакой опоры.

Выход состоит в том, чтобы никогда не передавать модели позицию, которой она не знает. Мы перемасштаблируем: позиция 100 000 становится позицией 25 000, и каждая позиция в последовательности возвращается в диапазон, известный модели. Токены не меняются, меняется только позиция, присваиваемая каждому из них. Именно это делает YaRN, а коэффициент — это просто то, насколько дальше мы хотим уйти: обучено на 40 960, требуется 163 840, значит, 4.

Как включить это в TRL? Передайте конфигурацию RoPE при загрузке модели:

training_args = SFTConfig(
    ...,
    model_init_kwargs={
        "rope_parameters": {
            "rope_type": "yarn",
            "rope_theta": 1_000_000,  # has to match what the model ships with
            "factor": 4.0,
            "original_max_position_embeddings": 40960,
        },
    },
)

Та же самая величина, с перемасштабированным запуском сверху:

После масштабирования кривая обучения остаётся плоской вплоть до 160k. На последних 20k токенов средняя величина loss составляет 2.8, тогда как без масштабирования она равна 7.3.

Остались активации: выгрузите их

После разбиения потерь на чанки и перемасштабирования позиций одна GPU позволяет обрабатывать 160k токенов. Затем память снова исчерпывается. Что теперь заполняет видеокарту?

Модель и её оптимизатор занимают фиксированные 22 ГБ независимо от длины последовательности. Сверх этого каждый красный сегмент представляет собой сохранённые активации одного слоя, по 0,47 ГБ каждый. Они накапливаются по мере записи во время прямого прохода, остаются в памяти и удаляются по одному по мере их потребления во время обратного прохода. Серая область сверху — это всё, что пересчитывается и отбрасывается в течение миллисекунд.

TRL уже уменьшает их за вас, и эта функция включена по умолчанию. Gradient checkpointing сохраняет только то, что поступает на вход каждого слоя, и отбрасывает всё, что слой вычисляет внутри себя: attention scores, промежуточные значения широкого MLP, нормы. Когда обратный проход нуждается в них, слой выполняется повторно, чтобы получить их обратно. Сохраняется один тензор на слой, каждый размером sequence x hidden, и при длинных последовательностях это по-прежнему самый большой объект на карте.

Полезное наблюдение состоит в том, что эти тензоры вообще не нужны на этапе прямого прохода. Они записываются, остаются нетронутыми и читаются значительно позже. Поэтому им не обязательно находиться на GPU в промежутке: их можно переместить в память CPU и вернуть обратно, когда обратный проход дойдёт до них.

training_args = SFTConfig(..., gradient_checkpointing_kwargs={"offload": True})

Тот же шаг, та же длина: вместо 38 накопленных полос одновременно резидентны только 4, а пиковое потребление снижается с 59,9 ГБ до 48,8 ГБ. Это позволяет увеличить длину последовательности:

Ниже 48k оба варианта совпадают, так как на такой длине практически нечего перемещать. Выше этого порога прогон с offload остаётся ниже, и там, где последовательность раньше останавливалась на 160k, теперь она достигает 256k на той же карте.

Этот вариант не является бесплатным. Каждое сохраняемое активационное значение передаётся на хост и обратно, и этот трафик должен укладываться в промежутки между вычислениями, которые он сопровождает.

Одного GPU недостаточно: распределите последовательность по нескольким

Потери по чанкам, выгруженные активации: теперь одна карта обучается на 256k токенов. Как достичь миллиона? Чтобы понять, что мешает, посмотрите внутрь одного слоя во время обратного прохода (серые пики на профиле выше).

Gradient checkpointing означает, что прямой проход сохраняет почти ничего: каждый слой хранит только свой вход и отбрасывает остальное. Слой пересчитывается позже, когда обратный проход доходит до него, и именно в этот момент всё существует одновременно. В MLP, одна строка модели,

down_proj(act_fn(gate_proj(x)) * up_proj(x))

создаются четыре тензора на полной промежуточной ширине, и стандартная реализация удерживает все четыре до тех пор, пока обратное распространение не пройдет их. Для миллиона токенов на модели, где MLP расширяет каждый токен с 2560 до 9728:

gate_proj(x)                             19 GB = 1,048,576 x 9728 x 2 bytes
act_fn(gate_proj(x))                     19 GB
up_proj(x)                               19 GB
act_fn(gate_proj(x)) * up_proj(x)        19 GB
                                  total  76 GB

76 ГБ, внутри одного слоя, на карте с 80 ГБ, которая уже отдала 30 ГБ модели и её оптимизатору. На этом этапе ответом является не более умное хранение последовательности. Ответом является дать каждой GPU меньше её.

Каждый GPU получает свой срез последовательности. При наличии четырёх GPU каждый хранит 262 144 токена вместо 1 048 576. Четыре тензора уменьшаются соответственно: 19 ГБ вместо 76, что вместе с 30 ГБ модели помещается на карте.

Очевидное возражение касается внимания. Каждый токен должен обращать внимание на каждый предыдущий токен, но ни один GPU не хранит все предыдущие токены. Поэтому GPU обмениваются недостающими данными, и на выходе получается ровно то, что вычислил бы один огромный GPU.

Существует два способа выполнить этот обмен: параллелизм контекста (CP) и последовательное параллелизм Ulysses (SP).

В FSDP2 используйте CP. Каждый GPU хранит собственный срез, а остальные срезы поочерёдно передаются на него, благодаря чему каждый токен в итоге видит все предшествующие токены.

from accelerate import ParallelismConfig

training_args = SFTConfig(..., fsdp=True, parallelism_config=ParallelismConfig(cp_size=4))

В DeepSpeed используйте SP. Вместо перемещения фрагментов он перемешивает батч непосредственно перед attention, так что каждый GPU хранит все токены, но только четверть attention-голов, а затем возвращает исходное распределение.

training_args = SFTConfig(
    ..., deepspeed="examples/deepspeed_configs/zero3.json", parallelism_config=ParallelismConfig(sp_size=4, sp_backend="deepspeed")
)

Эти два параметра сегодня не пересекаются: cp_size требует FSDP2, а sp_size работает только в DeepSpeed. Вы выбираете бэкенд, и метод определяется автоматически. Остальная часть этого раздела посвящена пути FSDP2, который используется в примере в начале этого руководства: его GPU не являются независимыми рабочими процессами, обрабатывающими разные батчи, а представляют собой одну группу, разделяющую единую последовательность.

После включения меняются две вещи:

1. Последовательности должны быть дополнены до кратного cp_size * 2, поэтому pad_to_multiple_of=8 (в SFTConfig) для четырёх GPU.
2. Требование causal-SDPA исключает использование packing, который опирается на блочно-диагональную маску, чтобы документы не читали друг друга, и TRL вызывает ошибку, если вы запросите оба режима. Это также исключает модели, слои которых используют скользящее окно или чанкованное внимание, которые accelerate отклоняет: OpenAI GPT-OSS, Gemma 3 и 4, Qwen3.5 и более поздние версии. Qwen3 и Qwen3-MoE используют полное внимание на всех слоях, именно поэтому они и представлены здесь.

Передача срезов обходится дешевле, чем можно было бы ожидать. При 131k токенах на одном узле один шаг занимает 34,6 с на двух GPU, 17,8 с на четырёх и 9,5 с на восьми. Каждое удвоение группы почти вдвое сокращает время шага, а переход с двух GPU на восемь восстанавливает 3,7x из возможных 4x.

И это последний из четырёх рычагов:

Одна карта останавливается чуть дальше 256k. Четыре из них берут весь миллион!, укладываясь в 63,6 ГБ с запасом. Последовательность, с которой начиналось это руководство, помещается.

Установите PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True. При миллионе токенов описанный выше запуск требует 63,6 ГБ на карте с 80 ГБ памяти и всё равно завершается с ошибкой без этой настройки: запрашиваемые тензоры большие и непрерывные, а аллокатор не может найти для них место среди уже занятых блоков.

Дополнительные материалы

• Контекстная параллельность, руководство Accelerate по cp_size и сетке устройств, которую оно строит.
• Ulysses и ring attention, о том, чем отличаются два типа обмена и какова стоимость коммуникации для каждого.
• YaRN, метод масштабирования позиций, описанный в разделе RoPE.