Как собрать свою языковую модель из кирпичиков OLMo-core

В мире больших языковых моделей (LLM) сейчас происходит странная вещь. С одной стороны, мы можем скачать готовые веса популярных моделей в пару кликов. С другой — если нужно не просто «поиграться» с промптами, а глубоко кастомизировать процесс обучения или архитектуру, мы упираемся в закрытый код или перегруженные фреймворки, где шаг влево, шаг вправо карается бесконечным дебагом.

Команда из Allen Institute for AI (AI2) выпустила проект OLMo-core. По сути, это набор чертежей и проверенных деталей, из которых они сами собирают свои флагманские открытые модели серии OLMo. Если вы когда-нибудь задумывались, что находится под капотом у профессиональных моделей и как организовать обучение на кластере H100 так, чтобы ничего не развалилось, этот проект для вас.

Что внутри коробки

OLMo-core — это не просто библиотека, а полноценный фреймворк для претрейнинга. Главная ценность здесь не в самой архитектуре трансформера (хотя и она тут вылизана), а в инструментах для стабильного обучения.

Проект закрывает сразу несколько «болей» ML-инженера:

• Готовые конфиги для обучения моделей разных размеров: от 7B до 32B параметров.
• Интеграция с современными оптимизаторами памяти вроде Liger-Kernel и torchao для работы с float8.
• Поддержка сложных механизмов внимания (Flash Attention, Ring Flash Attention).
• Инфраструктурная обвязка для запуска на распределенных кластерах через torchrun.

Я часто вижу, как исследователи пытаются адаптировать тренировочные скрипты из репозиториев пятилетней давности. В OLMo-core код свежий и, что важно, он прошел проверку в «боевых» условиях при создании последних моделей OLMo-2 и OLMo-3.

Как запустить инференс на скорую руку

Самый быстрый способ потрогать результат работы этих «кирпичиков» — использовать Hugging Face Transformers. Команда AI2 позаботилась о полной совместимости.

from transformers import AutoModelForCausalLM, AutoTokenizer

# Загружаем одну из последних моделей
model_name = "allenai/Olmo-3-1125-32B"
olmo = AutoModelForCausalLM.from_pretrained(model_name)
tokenizer = AutoTokenizer.from_pretrained(model_name)

inputs = tokenizer(["Language modeling is "], return_tensors='pt')
response = olmo.generate(**inputs, max_new_tokens=50)
print(tokenizer.batch_decode(response, skip_special_tokens=True)[0])

Если же вам нужна серьезная пропускная способность, библиотека дружит с vLLM. Это позволяет развернуть инференс для пакетной обработки текстов буквально за несколько минут.

Глубокое погружение в обучение

Самое интересное в репозитории скрыто в директории src/scripts/official/. Там лежат оригинальные скрипты обучения. Это не абстрактные примеры из учебника, а те самые файлы, которые запускали авторы на своих кластерах.

Меня зацепило то, как просто реализовано переопределение конфигов. Можно запустить обучение огромной модели в 32 миллиарда параметров и прямо из командной строки подкрутить скорость обучения или путь к чекпоинтам:

torchrun --nproc-per-node=8 src/scripts/official/OLMo2/OLMo-2-0325-32B-train.py \
  --save-folder=/my/checkpoints \
  --train_module.optim.lr=6e-3

Кстати, для тех, кто хочет сэкономить видеопамять, проект поддерживает Liger-Kernel. Это специальные CUDA-ядра, которые объединяют линейный слой и функцию потерь (fused cross-entropy), что сильно снижает аппетиты модели при обучении.

Кому это пригодится

OLMo-core вряд ли подойдет для тех, кто только начинает учить Python. Но если вы работаете над своим проектом в области NLP или пытаетесь дообучить (fine-tune) модель под узкие задачи, здесь можно подсмотреть много полезных решений.

1. Исследователям, которым важна воспроизводимость. Весь цикл обучения прозрачен.
2. ML-инженерам, которые ищут стабильный пайплайн для работы на нескольких GPU.
3. Разработчикам инструментов инференса. Бета-чат в самом OLMo-core (olmo_core.generate.chat) показывает, как работать с генерацией напрямую.

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

Вердикт

Проект выглядит как отличный фундамент для тех, кто перерос стандартные туториалы и хочет строить LLM-решения «по-взрослому». Это не магия, скрытая за API, а честный инженерный инструмент. Даже если вы не планируете обучать модель с нуля, стоит заглянуть в scripts, чтобы увидеть, как организована работа с памятью и чекпоинтами в топовых AI-лабораториях.

Попробовать стоит хотя бы ради того, чтобы увидеть, насколько далеко ушли технологии обучения со времен первого GPT. Только не забудьте про GPU — без мощного железа здесь делать особо нечего.