OLMo: ускорение науки о языковых моделях

OLMo - первая модель в семействе olmo от Allen Institute for AI. Отчёт, написанный Дирком Грёневельдом (Dirk Groeneveld), Изом Белтаджи (Iz Beltagy), Питом Уолшем (Pete Walsh) и ещё примерно сорока соавторами, включая Нейтана Ламберта (Nathan Lambert), вышел в феврале 2024 года (arXiv 2402.00838). Основной тезис краток: по мере того как языковые модели приобретали коммерческую ценность, сильнейшие из них закрывались, а необходимые учёным детали (обучающие данные, архитектурные решения, ход разработки) перестали раскрывать. OLMo - это попытка AI2 создать модель, конкурентоспособную на фоне тогдашних open-weight моделей на 7B параметров, выложив при этом в открытый доступ абсолютно всё.

Степени открытости

Авторы статьи позиционируют свою работу на шкале существующих релизов. Mixtral 8x7B вышел только с весами и коротким отчётом. LLaMA сопровождалась подробными инструкциями по адаптации. MPT от MosaicML описывала распределение данных в датасете, но сами данные не публиковала. Falcon раскрыла часть данных предобучения. Наиболее открытыми инициативами были Pythia и BLOOM, выложившие код, чекпоинты и данные. В качестве современника с аналогичными целями упоминается LLM360. Заявка OLMo заключается в том, чтобы сократить разрыв между этим полностью открытым эшелоном и возможностями моделей уровня Llama 2.

Опубликованные AI2 материалы охватывают весь пайплайн. Для предобучения: код обучения и моделирования, веса OLMo-7B, 7B-"близнеца", обученного на другом железе, и OLMo-1B - каждый с более чем 500 промежуточными чекпоинтами с интервалом в 1000 шагов, опубликованными как ревизии на HuggingFace, плюс полные метрики обучения из Weights & Biases. Для данных: полный корпус Dolma, код для его сборки, инструмент анализа датасетов WIMBD, а также утилиты, восстанавливающие порядок обучения и показывающие, какие именно данные модель видела на каждом шаге. Для адаптации: код и данные для instruction- и preference-обучения, а также веса моделей SFT и SFT+DPO. Для оценки: фреймворки Catwalk и Paloma и набор тестов Tülu. Код и веса распространяются под лицензией Apache 2.0. Именно это в статье практически понимается под "по-настоящему открытым" (truly open) подходом, и именно этого не дают релизы с одними лишь весами вроде Mixtral или Llama: сторонний исследователь может воспроизвести прогон, изучить, что именно видела модель, или переобучить собственный вариант.

Модели и обучение

Релиз включает четыре варианта 7B, различающихся архитектурой, оптимизатором и аппаратной платформой обучения, а также одну модель на 1B параметров. Все они обучались как минимум на 2T токенов (одна эпоха по выборке из Dolma объёмом 2T токенов), а некоторые пошли на вторую эпоху с другим перемешиванием (shuffle). Чекпоинт 7B, использованный для итоговой оценки, обучался до 2.46T токенов, после чего дообучался ещё 1000 шагов с линейным спадом learning rate до нуля, что дало заметный скачок на большинстве задач в финальной точке кривой обучения. Таблица с количеством слоёв, скрытой размерностью и learning rate для каждого размера (Таблица 1) не сохранилась при извлечении текста в локальной копии источника.

Архитектура представляет собой декодерный трансформер со стандартными для LLaMA, PaLM и Falcon модификациями: отсутствие bias-параметров (ради стабильности), активации SwiGLU со скрытым размером 11 008 у модели 7B и поворотные позиционные эмбеддинги (RoPE). Одно решение было нетипичным: непараметрическая нормализация слоёв (layer norm) без обучаемых параметров gain и bias, выбранная вместо параметрического layer norm и RMSNorm как "наиболее безопасный вариант" и самый быстрый. Токенизатор - модифицированный BPE от GPT-NeoX-20B с дополнительными токенами для маскирования персональных данных; словарь на 50 280 токенов размещён в матрице эмбеддингов на 50 304 строки, округлённой вверх до кратного 128 ради пропускной способности.

При обучении использовались PyTorch FSDP с шардингом ZeRO, смешанная точность bfloat16 с сохранением softmax в полной точности (FP32), глобальный размер батча около 4M токенов (2048 последовательностей длиной 2048), AdamW с разогревом (warmup) на 5000 шагов (около 21B токенов), линейным спадом до одной десятой от пикового значения и отсечением градиентов (gradient clipping) на уровне 1.0. Документы конкатенируются через токен EOS и нарезаются на фрагменты по 2048 токенов, перемешиваясь одинаково для каждого запуска, благодаря чему точный состав каждого батча можно восстановить по выложенным артефактам.

AI2 проводил обучение на двух кластерах, чтобы показать работоспособность кода на железе обоих вендоров: суперкомпьютере LUMI (до 256 узлов по четыре AMD MI250X в каждом) и MosaicML от Databricks (27 узлов по восемь 40GB NVIDIA A100 в каждом). К отметке в 2T токенов оба кластера дали практически идентичные результаты на валидации. В приложении затраты энергии на предобучение 7B-моделей оцениваются в 239 МВт-ч, а выбросы - в 69.78 т в CO2-эквиваленте, при этом гидроэлектроэнергия LUMI считается безуглеродной согласно официальным данным LUMI (при углеродоёмкости 0.024 кг/кВт-ч для ГЭС прогон на LUMI добавил бы 3.54 т CO2-экв.); авторы называют эти цифры нижней границей, поскольку отладка, подбор гиперпараметров и простои здесь не учитываются.

Dolma

Данные предобучения в 2024 году оставались самой закрытой частью индустрии, и в статье Dolma рассматривается как самостоятельный полноценный релиз. Это мультидоменный корпус объёмом в триллионы токенов, собранный из общедоступных источников, типичных для предобучения LLM. Пайплайн включает языковую фильтрацию, фильтрацию по качеству, контентную фильтрацию, дедупликацию, смешивание источников и токенизацию, причём документы из каждого источника хранятся раздельно как на этапе подготовки, так и в финальном релизе. Подробности вынесены в отдельный отчёт по Dolma; таблица состава данных в текущей статье также пострадала при извлечении текста. Набор инструментов для курации данных открыт. Последующий кризис лицензирования веб-данных для подобных задач описан в consent-in-crisis-ai-data-commons.

Оценка

Оценка проводится в двух режимах. Внутри цикла обучения (in-loop) набор прикладных бенчмарков запускается каждые 1000 шагов (около 4B токенов) и направляет решения по архитектуре, инициализации, оптимизаторам, графикам learning rate и смеси данных. В офлайн-режиме фреймворк Catwalk оценивает финальные чекпоинты на восьми zero-shot задачах (ARC easy и challenge, BoolQ, OpenBookQA, SciQ, HellaSwag, PIQA, WinoGrande) с нормализацией правдоподобия, выбранной индивидуально под каждую задачу. В сравнении с LLaMA-7B, Llama-2-7B, MPT-7B, Pythia-6.9B, Falcon-7B и RPJ-INCITE-7B модель OLMo-7B в совокупности выглядит "конкурентоспособной". Показатели по отдельным задачам приведены в таблице, которая не сохранилась при извлечении.

Оценка перплексии использует Paloma - 585 текстовых доменов из 18 источников с метрикой в битах на байт, что позволяет корректно сравнивать модели с разными словарями. OLMo-7B называют крупнейшей моделью с явной деконтаминацией против Paloma: любой обучающий документ, содержавший абзацы из тестовых данных Paloma, удалялся. Качество предсказания точно отражает распределение обучающей выборки. OLMo, где 88.8% данных приходится на Common Crawl, обходит все остальные модели на C4, но уступает на WikiText-103 и наборах M2D2 (академические тексты и Википедия). MPT-7B в целом показывает более быстрый прогресс, возможно, потому что 27% её данных получены не из Common Crawl (против 11.2% у OLMo). Авторы делают из этого общий вывод: качественный курируемый текст вроде Википедии и arXiv встречается гораздо реже сырого веб-текста, поэтому удерживать эффективность выборки (sample efficiency) на нём по мере роста корпусов будет всё труднее. На 100 языках программирования из Dolma OLMo далеко впереди остальных, что в статье объясняется отчасти идентичной постобработкой кода, отчасти контаминацией, которую Paloma не может вычистить из исходников.

Адаптация

Чтобы показать применимость OLMo в качестве базовой модели, AI2 применил рецепт Tülu 2: instruction finetuning на смеси синтетических (дистиллированных) и написанных людьми данных (3 эпохи с learning rate 2e-6), а затем Direct Preference Optimization на наборе предпочтений UltraFeedback (3 эпохи, 5e-7, beta 0.1). Дообучение заметно подняло MMLU ("by a wide margin") и улучшило результаты на ToxiGen и TruthfulQA, сильнее всего после DPO. Модель OLMo+SFT+DPO превосходит большинство других 7B-чат-моделей, но уступает Tülu 2 на базе Llama 2. Авторы называют две причины: о Llama 2 сообщалось, что она контаминирована тестовыми данными MMLU, а сама смесь Tülu изначально подбиралась под Llama. В разделе ограничений добавляется, что в смеси используются данные, дистиллированные из других моделей, и от этой зависимости AI2 планирует уходить.

Что признаёт отчёт

В статье перечислено и то, что сделано не было. Данные - только на английском языке. Из-за ограничений на объём статьи не вошли логи разошедшихся или неудачных прогонов. Большинство прикладных задач мало похожи на реальные сценарии использования чат-ботов, а сами оценки зашумлены. В этическом заявлении аргумент в пользу открытости приведён прямо: публикация позволяет другим исследователям развивать модель, а не переобучать её с нуля ценой экологических затрат; кроме того, сопоставимые модели под разрешительными лицензиями уже существовали, поэтому более строгая лицензия не устранила бы риски для отрасли. В заключении отмечается, что последующее обновление OLMo 1.7 с улучшенными данными и обучением подняло результат на MMLU на 24 процентных пункта до 52% - именно этот разрыв затем взялась окончательно устранить olmo-2-furious.