Lila Sciences и Anyscale обучили GLM-5.3-Flash, открытую модель, на 1 миллионе токенов на последовательность, используя 64 графических процессора B200. Это в 64 раза больше, чем стек обучения мог выдержать в самом начале, и каждое изменение доступно с открытым исходным кодом в SkyRL. Наша команда SkyRL в Anyscale тесно сотрудничала с исследователями Lila Science над их инфраструктурой на базе открытого стека, чтобы изучить исследовательский путь по увеличению длины контекста для этой конкретной модели.
Научные данные по своей природе протяженны. Геном, семейство родственных белков или последовательностей РНК, либо запись эксперимента, проводимого в течение нескольких дней, теряют важную структуру, если их разбить на несколько тысяч токенов. ИИ, управляющий научными исследованиями в замкнутом цикле с автономной лабораторией, имеет ту же потребность: на протяжении циклов гипотез, протоколов и показаний приборов его контекст вырастает до сотен тысяч токенов.
Многие открытые модели теперь поддерживают контекст в 1 млн токенов, но они могут не обладать производительностью из коробки для целевой задачи при больших значениях контекста. Устранение этого разрыва означает обучение на собственных данных задачи во всей их полноте, без сокращения или сильного сжатия контекста. Эта работа была направлена на достижение этой цели с помощью открытой модели и открытого стека обучения.
Результат вкратце:
- 1 млн токенов на последовательность помещается в 144,5 ГиБ на GPU, что ниже лимита B200 в 178 ГиБ.
1 млн токенов на последовательность помещается в 144,5 ГиБ на GPU, что ниже лимита B200 в 178 ГиБ.
- 13,7 тыс. токенов/с при 1М и пиковое значение 24,9 тыс. токенов/с при 256k на 64 GPU B200.
13,7 тыс. токенов/с при 1М и пиковое значение 24,9 тыс. токенов/с при 256k на 64 GPU B200.
- Модель использует контекст: после дообучения на геноме человека ее функция потерь улучшается сильнее всего на токенах, перед которыми находятся сотни тысяч токенов контекста.
Модель использует контекст: после дообучения на геноме человека ее функция потерь улучшается сильнее всего на токенах, перед которыми находятся сотни тысяч токенов контекста.
- Открытый исходный код: рецепт находится в SkyRL, и несколько исправлений переносятся в основную ветку Megatron-LM.
Открытый исходный код: рецепт находится в , и несколько исправлений переносятся в основную ветку Megatron-LM.
Ссылка: Как мы достигли контекста в 1M токенов для GLM 5.3 Flash?
GLM-5.3-Flash — это гибридная модель: 34 из ее 45 слоев используют линейное внимание (Kimi Delta Attention), а 11 — разреженное внимание DeepSeek (DSA). В ней также используется остаточный поток, который в 4 раза шире обычного, и слои смеси экспертов (MoE). Каждое из этих решений потребовало отдельной работы. Объем памяти растет с длиной последовательности в каждом слое, а во внимании он растет пропорционально квадрату длины, поэтому каждая оптимизация лишь отодвигает предел до тех пор, пока не столкнется со следующим. Мы выполнили несколько оптимизаций, каждая из которых подробно описана с собственными экспериментами и PR в оригинальной .
Большинство столкнувшихся с нами проблем не специфичны для GLM и могут быть применены к моделям с аналогичными типами компонентов:
- Убедитесь, что каждый слой выполняет именно то ядро, которое вы предполагаете: слои разреженного внимания GLM молча переключались на плотное внимание, поскольку слитные ядра принимали только точные формы тензоров DeepSeek. Ничто не вызывало ошибки; запуск просто потреблял 114 ГиБ на GPU при 16k и исчерпывал память при 32k. Дополнение тензоров до соответствия разреженному ядру уменьшило потребление памяти вдвое.
Убедитесь, что каждый слой выполняет именно то ядро, которое вы предполагаете: слои разреженного внимания GLM молча переключались на плотное внимание, поскольку слитные ядра принимали только точные формы тензоров DeepSeek. Ничто не вызывало ошибки; запуск просто потреблял 114 ГиБ на GPU при 16k и исчерпывал память при 32k. Дополнение тензоров до соответствия разреженному ядру уменьшило потребление памяти вдвое.
- Разреженное внимание переносит квадратичные затраты, но не устраняет их. DSA учитывает только 2048 токенов на запрос, но выбор этих токенов означает оценку каждого запроса по всей последовательности. Вычисление этой матрицы оценок по несколько сотен запросов за раз позволяет удерживать объем памяти ниже 2 ГиБ, а разделение оценки между GPU вместо ее повторения на каждом из них сделало полный шаг обучения в 2,1 раза быстрее при 256k.
Разреженное внимание переносит квадратичные затраты, но не устраняет их. DSA учитывает только 2048 токенов на запрос, но выбор этих токенов означает оценку каждого запроса по всей последовательности. Вычисление этой матрицы оценок по несколько сотен запросов за раз позволяет удерживать объем памяти ниже 2 ГиБ, а разделение оценки между GPU вместо ее повторения на каждом из них сделало полный шаг обучения в 2,1 раза быстрее при 256k.
- Каждое исправление памяти снижает пропускную способность, поэтому используйте только то, что необходимо для текущей длины. Полное перевычисление, выгрузка на CPU и контекстный параллелизм дают больше контекста за счет некоторого снижения скорости. Выгрузка, например, снижает пиковое потребление памяти при 256k со 142 до 98 ГиБ, но увеличивает время шага на 21%. Мы включаем каждое из них только на той длине, при которой обучение иначе бы не поместилось.
Каждое исправление памяти снижает пропускную способность, поэтому используйте только то, что необходимо для текущей длины. Полное перевычисление, выгрузка на CPU и контекстный параллелизм дают больше контекста за счет некоторого снижения скорости. Выгрузка, например, снижает пиковое потребление памяти при 256k со 142 до 98 ГиБ, но увеличивает время шага на 21%. Мы включаем каждое из них только на той длине, при которой обучение иначе бы не поместилось.
- Бенчмарки на случайных токенах занижают требования к памяти для MoE-моделей. Реальное распределение текста направляет потоки неравномерно по экспертам: на реальном датасете с длинным контекстом самый загруженный GPU получал примерно в 4 раза больше среднего числа токенов, и запуск при 512k исчерпывал память там, где запуски со случайными токенами помещались легко. Обработка слоя MoE по последовательности частями решила эту проблему. Тестируйте память на реальных данных перед определением размера кластера.
Бенчмарки на случайных токенах занижают требования к памяти для MoE-моделей. Реальное распределение текста направляет потоки неравномерно по экспертам: на реальном датасете с длинным контекстом самый загруженный GPU получал примерно в 4 раза больше среднего числа токенов, и запуск при 512k исчерпывал память там, где запуски со случайными токенами помещались легко. Обработка слоя MoE по последовательности частями решила эту проблему. Тестируйте память на реальных данных перед определением размера кластера.
- Планируйте память хоста, а не только память GPU. При 1M токенов входные данные слоев, выгружаемые на CPU, составляют около 180 ГБ на GPU до контекстного параллелизма (более 1 ТБ на узел с 8 GPU), и они разделяют оперативную память хоста с состоянием оптимизатора. Используйте закрепленную память для копий (они работали на 22% быстрее страничной памяти) и освобождайте ее на каждом шаге.
Планируйте память хоста, а не только память GPU. При 1M токенов входные данные слоев, выгружаемые на CPU, составляют около 180 ГБ на GPU до контекстного параллелизма (более 1 ТБ на узел с 8 GPU), и они разделяют оперативную память хоста с состоянием оптимизатора. Используйте закрепленную память для копий (они работали на 22% быстрее страничной памяти) и освобождайте ее на каждом шаге.
- Проверяйте правильность отдельно от скорости. Показатели скорости и памяти могут быть получены на случайных весах, но правильность — нет. В запусках с подкрепленным обучением (RL) с реальными весами логарифмические вероятности тренажера совпадали с вероятностями vLLM с отклонением в пределах 0,04 — 0,07 (средняя абсолютная разница) от 128k до 1M токенов, а тест на переобучение снизил функцию потерь с 1,65 до 0,17 за 20 шагов.
Проверяйте правильность отдельно от скорости. Показатели скорости и памяти можно получить на случайных весах, а правильность — нет. В запусках RL с реальными весами лог-вероятности тренера совпадали с таковыми у vLLM с отклонением в пределах 0,04 — 0,07 (средняя абсолютная разница) в диапазоне от 128 тыс. до 1 млн токенов, а тест на переобучение снизил потери с 1,65 до 0,17 за 20 шагов.
Благодаря сочетанию этих идей нам удалось поддержать хорошую пропускную способность при различных длинах контекста
LinkTraining Results
Мы провели дообучение на геноме человека с длиной последовательности 512 тыс. токенов, потери модели на отложенной выборке снизились на 4,8% за 80 шагов, причем наибольший прирост наблюдался на токенах с наибольшим контекстом перед ними. Последнее служит доказательством того, что модель научилась использовать длинный контекст, а не только локальные паттерны.
Мы провели обучение на хромосомах 1 и 2 эталонного генома человека (hg38) и выполнили оценку на хромосоме 22, которую модель никогда не видела. ДНК — отличный тест, поскольку она обладает структурой на любом масштабе: от коротких мотивов до паттернов, охватывающих очень длинные участки.
Чтобы отделить обучение на больших расстояниях от локального обучения, мы разделили потери на отложенной выборке по позициям токенов (с шагом 32 тыс. токенов) и сравнили их до и после тонкой настройки:
- У ранних токенов перед ними мало контекста, поэтому любое улучшение там обусловлено локальными паттернами, такими как короткие мотивы и нуклеотидный состав.
У ранних токенов перед ними мало контекста, поэтому любое улучшение там обусловлено локальными паттернами, такими как короткие мотивы и нуклеотидный состав.
- У поздних токенов перед ними находится до 500 тыс. токенов контекста.
У поздних токенов перед ними находится до 500 тыс. токенов контекста.
- Локальное обучение помогает каждой позиции в равной степени, поэтому более значительное улучшение на поздних позициях должно быть обусловлено именно длинным контекстом, доступным только им.
Локальное обучение помогает каждой позиции в равной степени, поэтому более значительное улучшение на поздних позициях должно быть обусловлено именно длинным контекстом, доступным только им.
LinkRunning on Anyscale
Каждый эксперимент в этой статье выполнялся на Ray под управлением платформы Anyscale в кластере из 64 графических процессоров B200 в облаке Lila. SkyRL использует Ray для размещения и координации своих рабочих процессов обучения и инференса, поэтому одна и та же конфигурация выполняла тесты памяти на случайных токенах, тонкую настройку на реальных данных и запуски RL для проверки тренера относительно vLLM.
Более подробно о конкретных методах, используемых для поддержки длинного контекста, см. в этом подробном блог-посте.
