2 OLMo 2 Furious

OLMo 2 - второе поколение моделей olmo от Allen Institute for AI. Отчёт, подготовленный командой OLMo под руководством Пита Уолша (Pete Walsh), Луки Солдаини (Luca Soldaini), Дирка Грёневельда (Dirk Groeneveld) и Кайла Ло (Kyle Lo) при участии Нейтана Ламберта (Nathan Lambert) в числе основных авторов, впервые был опубликован в декабре 2024 года (arXiv 2501.00656); версия в vault также охватывает модель на 32B, выпущенную в марте 2025 года. Семейство состоит из плотных (dense) decoder-only моделей размером 7B, 13B и 32B, плюс вариант на 1B, использовавшийся для отработки рецепта обучения и позже также выложенный в открытый доступ. Если первая OLMo была главным образом доказательством того, что по-настоящему открытая модель в принципе возможна, то главный тезис OLMo 2 - это производительность на единицу вычислительных затрат: базовые модели находятся на границе Парето по соотношению среднего балла в бенчмарках и FLOPs на обучение, не уступая или превосходя Llama 3.1, Qwen 2.5 и Gemma 2 в своих весовых категориях при меньших затратах.

Что здесь понимается под "полной открытостью"

Во введении авторы проводят чёткую границу между своим подходом и релизами открытых весов (open-weight) 2024 года (Llama 3, Qwen 2, Gemma, Mistral, Phi и др.): такие веса - "лишь финальные артефакты" пайплайна, и их недостаточно для изучения того, как именно формируется поведение моделей. В лагерь полностью открытых проектов авторы включают Pythia, первую OLMo, Amber, DCLM, MAP Neo и SmolLM. Для OLMo 2 опубликованы веса всех размеров и стадий (base, SFT, DPO и финальный instruct), все данные pretraining и mid-training, код обучения как в исходном репозитории OLMo, так и в переписанном OLMo-core, код пост-тренинга open-instruct, набор для оценки OLMES и инструменты фильтрации Dolma, логи обучения Weights & Biases или Comet для 7B, 13B и 32B, а также тысячи промежуточных чекпоинтов. Всё распространяется под лицензией Apache 2.0, где это возможно, либо под максимально разрешительными из доступных лицензий.

Модели и объёмы токенов

Все модели используют контекст 4096 токенов. 32B - единственная модель с grouped-query attention; авторы отмечают, что это решение было вдохновлено Qwen 3. Общее число токенов для 1B выведено из 4T на этапе pretraining плюс один anneal на 50B; единой цифры в отчёте не приводится.

Список исправлений для стабильности

Самый большой технический раздел посвящён тому, почему OLMo-0424, модель апреля 2024 года, не удавалось масштабировать. В процессе её обучения регулярно случались скачки loss (loss spikes), которым часто предшествовали всплески нормы градиента, а также наблюдался медленный рост нормы градиента на протяжении всего обучения; на более крупных моделях всплесков было больше, что неизбежно вело к расходимости. В OLMo 2 изменили восемь составляющих, оценивая большинство правок через "spike score": процент значений, отклоняющихся более чем на семь стандартных отклонений от скользящего среднего по 1 000 шагов.

Первым делом занялись данными. В батчах, вызывавших сбои, часто встречались длинные повторяющиеся n-граммы (мусор в base64, цепочки 255, 255, 255). Связь не строго детерминирована: одна и та же последовательность могла уронить одну модель и не затронуть другую, либо вызвать скачок при одном порядке данных и пройти гладко после перемешивания. Удаление документов с 32 и более повторами любых n-грамм длиной от 1 до 13 токенов, а также маскирование таких последовательностей в функции потерь трейнера устранили множество скачков, но никак не повлияли на плавный рост нормы градиента.

Инициализация стала главным точечным исправлением. В OLMo-0424 веса каждого слоя уменьшались по мере глубины; в OLMo 2 каждый параметр инициализируется из нормального распределения со средним 0 и стандартным отклонением 0.02. В заведомо нестабильном тестовом прогоне с коротким warmup показатель всплесков нормы градиента упал с 0.40 до 0.03. Отчёт подкрепляет это расчётом "growth exponent", измеряющего рост норм активаций и градиентов по слоям при инициализации (новая схема ближе к нулю на любой ширине), а также тем, что нормы градиентов масштабируются пропорционально квадратному корню из ширины - свойство, важное для переноса гиперпараметров между размерами моделей.

Архитектурные изменения вводились комплексно. Беспараметрический layer norm, бывший отличительной чертой OLMo 1, заменили обратно на RMSNorm: абляции не выявили разницы, а баги в сторонних библиотеках, из-за которых его когда-то выбрали, к этому моменту исправили. Нормализацию перенесли со входов подслоёв внимания и MLP на их выходы, так что каждый блок вычисляет h = x + RMSNorm(Attention(x)), а запросы (queries) и ключи (keys) получили собственный RMSNorm (QK-norm) перед блоком внимания. По отдельности эти меры не помогают; вместе они снизили показатель всплесков градиента с 0.108 до 0.069. Регуляризация z-loss сдерживает рост выходных логитов. Параметр RoPE theta подняли с 10 000 до 500 000, как в Llama 3. В оптимизаторе AdamW параметр epsilon снизили с 1e-5 до стандартного для PyTorch 1e-8: это позволило делать первые шаги обновления крупнее, благодаря чему норма градиента быстрее стабилизируется и остаётся ниже. Weight decay отключили для эмбеддингов, поскольку их затухание уменьшало норму самих эмбеддингов и раздувало градиенты ранних слоёв (spike score составил 0.16 с затуханием против 0.092 без него).

Реализация fused z-loss из Flash Attention и стандартная реализация в PyTorch совпадают на прямом проходе (forward pass), но расходятся на обратном (backward pass), вероятно, из-за погрешностей точности. На значениях loss и бенчмарках это никак не отразилось, однако AI2 отменила ветку обучения, перешедшую на другую реализацию, и перезапустила прогон с точки расхождения.

Mid-training и Dolmino

Базовое обучение делится на два этапа. Pretraining (90-95% FLOPs) идёт на смеси OLMo 2 Mix 1124: 3.90T токенов, из которых более 95% - веб-тексты из DCLM-Baseline (3.71T), а также 83B токенов кода StarCoder, отфильтрованного по репозиториям минимум с 2 звёздами, 58.6B токенов научных статей peS2o, 20.8B из arXiv, OpenWebMath и Algebraic Stack примерно по 12B каждый, и 3.7B из Википедии и Wikibooks. Тот же микс используется в OLMoE. Косинусный график learning rate рассчитан на 5T токенов и усекается: прогон 7B останавливается на 4T, а 13B доходит до 5T.

Mid-training (5-10% FLOPs) линейно снижает learning rate до нуля на наборе Dolmino Mix 1124 - меньшей по объёму смеси, созданной для накачки STEM-знаний и подтягивания математики. Высококачественная часть содержит 832.6B токенов: веб-документы DCLM из верхних 7% по классификатору FastText с оценкой не менее 2 по классификатору FineWeb-Edu (752B), деконтаминированные инструкции FLAN, peS2o, Википедия и Q&A со Stack Exchange, отфильтрованные по голосам. Математическая часть занимает 10.7B токенов и в основном синтетическая: TinyGSM-MIND (6.5B токенов задач типа GSM8K с кодовыми решениями, переписанными в диалоги на естественном языке с помощью Qwen2.5-7B-Instruct), синтетические учебники в стиле MathCoder2, 230M токенов персонализированных задач и решений от GPT-4o, а также 28M токенов целевой арифметики и GSM8K с подменёнными числами. Выборки объёмом 50B, 100B и 300B токенов формируются так, чтобы фильтрованный веб составлял около половины; небольшие источники повторяются до 4 раз в крупных выборках.

Эффект этого этапа - главная цифра отчёта. Средний балл 7B по набору бенчмарков вырос на 10.6 пункта (с 53.0 до 62.9), превзойдя результат, который 13B показывала после одного лишь pretraining, а GSM8K подскочил с 24.1 до 67.5. Модель 13B прибавила 10.3 пункта и выросла с 37.3 до 75.1 на GSM8K. Относительный прирост уменьшается с размером: 37% на 1B, 12.3% на 32B.

Три вывода из процесса создания этого этапа заслуживают внимания. Во-первых, перебор learning rate показал, что более высокие пиковые значения лучше выглядят на ранних этапах, но затем уступают; после отжига (annealing) прогоны от 3e-4 до 12e-4 приходят практически к одинаковому loss и баллам в бенчмарках: высокий LR делает отжиг ровно настолько эффективнее, насколько ухудшал pretraining. Авторы отмечают, что это противоречит расхожему мнению о критической важности площади под кривой обучения, и предупреждают, что точка пересечения наступила далеко за отметкой 200B токенов, поэтому короткий перебор привёл бы к неверному выбору. Во-вторых, "микроотжиги" (microanneals) позволяют недорого протестировать источник данных, проводя отжиг на смеси 50/50 из этого источника и веб-текста: 19 таких тестов потребовали 130B токенов - меньше, чем три отжига по 50B для финальной 7B. Они показали, что 10% математики дают почти весь прирост GSM8K, характерный для доли в 35%, что двукратный повтор дефицитных математических данных помогает, а TinyGSM с решениями в виде кода портил GSM8K, тогда как те же задачи, переписанные прозой, поднимали его с 28.5 до 65.5. В-третьих, усреднение чекпоинтов ("souping") от прогонов отжига с разным порядком данных показало результаты на уровне или выше лучшего одиночного прогона на всех шести протестированных смесях. Финальная 7B усредняет три отжига по 50B; модели 13B и 32B усредняют три отжига по 100B и один на 300B.

Авторы аккуратно подходят к оценке. Они разделяют тестовые бенчмарки и отложенную выборку (AGIEval, GSM8K, MMLU-Pro, TriviaQA), никогда не использовавшуюся для принятия решений, и признают, что GSM8K был отложен лишь частично: 200 из 1 319 его задач использовались для калибровки математической смеси, и только остальные 1 119 идут в зачёт. Авторы призывают других разработчиков открыто указывать, за какими задачами они следили в процессе.

Post-training

Для OLMo 2-Instruct применяется рецепт Tülu 3 от AI2 в три этапа с двумя отличиями: данные предпочтений генерируются только моделями с разрешительными лицензиями, а этап RL расширен до нескольких раундов. SFT использует порядка 939K промптов для 7B и 13B и 866K для 1B и 32B; во втором миксе отфильтрованы синтетические ответы с упоминанием knowledge cutoff (из-за них модель начинала галлюцинировать датами отсечки знаний и начинать фразы с "As an AI language model"), а в математике оставлены только ответы, победившие по большинству голосов из 5 сэмплов. Удаление мультиязычных данных из смеси Tülu 3 стоило около полубалла, поэтому их оставили. Для OLMo 2 потребовались заметно более высокие learning rate на этапе SFT, чем для Llama 3.1 при том же рецепте.

DPO выполняется on-policy: примерно 367K-378K промптов, ответы сэмплируются из чекпоинтов OLMo SFT и пула из 20 других моделей, оцениваются GPT-4o по полезности, правдивости, честности и следованию инструкциям, после чего бинаризуются. Финальный этап - обучение с подкреплением на основе проверяемых наград (RLVR), где политика вознаграждается, только если математический ответ или форматирование прошли строгую проверку. Версии 7B и 13B обучаются через PPO, где value function инициализируется из reward-модели; для 13B потребовалось два дополнительных раунда (сначала на GSM8K, затем только на MATH) после того, как первый проход просел по математике. Варианты 1B и 32B используют GRPO, которому reward-модель не требуется. Первый instruct-релиз пришлось переобучать, когда в AI2 обнаружили, что instruct-токенизатор потерял логику претокенизации базовой модели; те версии теперь помечены как "preview".

По результатам OLMo 2 13B Instruct (средний балл 63.5) обходит Llama 3.1 8B Instruct и Tülu 3 8B и вплотную приближается к Qwen 2.5 14B Instruct (65.3). OLMo 2 32B Instruct набирает в среднем 68.8, находясь на одном уровне с Qwen 2.5 72B и опережая GPT-3.5 Turbo (60.5) и GPT-4o mini (65.7) на том же наборе тестов, хотя и уступает Llama 3.3 70B (72.7). Код не входил в число целевых навыков и исключён из оценки instruct-моделей.

Инфраструктура и стоимость

В отчёте подчёркивается, что детали инфраструктуры - неотъемлемая часть воспроизводимости, которую обычно опускают. Обучение велось на двух кластерах AI2: Jupiter в Остине (1 024 ускорителя H100 в 128 узлах с 8x400 Gbps InfiniBand на узел и PUE 1.2) и Augusta (кластер Google Cloud на 160 узлов в Айове; точные характеристики GPU в выгрузку не попали). Планировщик Beaker переносил задачи между ними изменением одной строки, запускал проверку работоспособности через перемножение тензоров перед каждым запуском и автоматически изолировал сбойные узлы. Четыре практики в PyTorch описаны достаточно подробно для прямого повторения: использование torch.compile, отлов синхронизаций между хостом и устройством (с помощью torch.cuda.set_sync_debug_mode("warn")), сбор метрик и сохранение чекпоинтов в отдельном потоке на бэкенде GLOO (поскольку NCCL не потокобезопасен), а также отключение сборщика мусора Python, чтобы все ранги запускали сборку на одном и том же шаге. Pretraining для 7B и 13B потребовал около 391 МВт·ч и привёл к выбросам порядка 154 т CO2-экв.; кроме того, впервые в отчётах оценён расход воды - около 1.1 миллиона литров. Обучение одной лишь 7B заняло 131 МВт·ч против 1 022 МВт·ч, заявленных для Llama 3.1 8B.

Слабые места, отмеченные в отчёте

Модель 1B плохо масштабировалась по токенам. Она остаётся конкурентоспособной на фоне SmolLM2, но отстаёт от младших базовых версий Gemma 2 и Qwen 2.5, а после одного лишь pretraining не могла превзойти случайное угадывание на MMLU или ARC-Challenge. Авторы предполагают, что ниже определённого размера pretraining требует данных, заточенных под конкретные задачи, либо дистилляции из более крупной модели (именно так создавались малые модели Gemma 2), при этом отмечая, что стандартный instruct-рецепт всё же выводит 1B на конкурентный уровень с Qwen 2.5 1.5B и Gemma 3 1B. Сетка перебора learning rate была недостаточно широкой, чтобы нащупать границы плато. olmo-3-technical-report продолжает это развитие с увеличенным контекстом, reasoning-моделями и пайплайном данных, перестроенным с учётом лицензионных ограничений, описанных в consent-in-crisis-ai-data-commons.