OLMo — моделька от AI2, разработанная учёными для учёных. Если вы занимаетесь экспериментами над нейронками, это идеальный вариант: весь код и данные, необходимые для тренировки, открыто лежат на GitHub. Более того, выложены даже промежуточные чекпоинты, с очень высокой гранулярностью. Это отличает ее от якобы "открытых" нейронок, которые обычно приходят к тебе монолитным финальным бинарём.
Эта статья — короткий гайд, по большей части взятый с GitHub создателей нейросети. Гайд проверен в боевых условиях и действительно работает.
Установка
Сначала установливаем Git, Python и PyTorch, следуя инструкциям для вашей операционной системы. Инструкции для PyTorch есть здесь.
Для обучения и дообучения рекомендуется установка из исходного кода:
git clone https://github.com/allenai/OLMo.git
cd OLMo
pip install -e .[all]Также можно установить из PyPI:
pip install ai2-olmoShameless plug. Хотите свежего мяса и пейперов про нейронки? Телега @1red2black.
Предварительное обучение
Предварительное обучение OLMo проходит в два этапа. На первом этапе мы обучаем модель на больших объёмах данных, преимущественно из интернета: OLMo-mix-1124 На втором этапе мы обучаем на меньшем объёме высококачественных целевых данных: Dolmino-mix-1124
Вы можете найти все контрольные точки моделей, как минимум каждые 1000 шагов обучения, в форматах OLMo и Hugging Face:
Вариант | Формат OLMo (Этап 1) | Формат OLMo (Этап 2) | Формат Hugging Face |
|---|---|---|---|
OLMo-2 1B | |||
OLMo-2 7B | |||
OLMo-2 13B | |||
OLMo-2 32B |
Примечание: для варианта 32B использовались новые скрипты для тренинга, которые тоже открыто опубликованы. Если хотите пройти весь этот путь, для обучения или дообучения OLMo-2 32B посетите OLMo-core.
Шаги для воспроизведения
Чтобы воспроизвести любой из описанных ниже процессов обучения, выполните:
torchrun --nproc_per_node=8 scripts/train.py {путь_к_файлу_конфигурации}Параметр --nproc_per_node , как и следует из названия — количество одновременных процессов PyTorch. Если у вас несколько GPU, то стоит установить значение, равное количеству GPU. При тренинге целиком на CPU, там должно быть разумное количество ядер процессора (не стоит на 16-ядерном процессоре занимать все 16 ядер).
Для настройки обучения, используйте любую из конфигураций, перечисленных ниже.
Если хотите переопределить какие-то настройки в конфигурации, без необходимости создавать новую конфигурацию каждый раз, можно сделать следующее:
torchrun --nproc_per_node=8 scripts/train.py {путь_к_файлу_конфигурации} \
--настройка=значение \
--настройка=значение \
--настройка.поднастройка=значениеВ конфигурации ниже, данные передаются в реальном времени по HTTP. Если нужно провести более полномасштабный эксперимент, рекомендуется загрузить файлы локально и изменить пути, чтобы они указывали на вашу локальную файловую систему.
Для запуска на устройствах Mac с процессорами Apple:
python scripts/train.py {путь_к_файлу_конфигурации}Пример:
python scripts/train.py configs/tiny/OLMo-20M.yaml --save_overwriteПримечание: Для запуска необходимо обновить PyTorch до версии 2.5.x.
Этап 1
Этап 1 — самый масштабный и дорогой этап, где нейронка обучается на 4-5 триллионах токенов. В таблице ниже приведены ссылки на файлы, которые нужно передавать для консольной команды pytorch:
OLMo2 1B | OLMo2 7B | OLMo2 13B | |
|---|---|---|---|
Number of tokens | 4 Trillion | 4 Trillion | 5 Trillion |
Checkpoint | |||
Training config | |||
WandB |
Данные, в основном, загружаются из веба. Посмотрите, как выглядит раздел data.paths, и не забудьте пересохранить его в виде локальных файлов. В качестве лайфхака, можно локально развернуть веб-сервер и подменить домен olmo-data.org на 127.0.0.1. В общем, развлекайтесь, оптимизация — дело добровольное.
data:
pad_direction: right
# generate_doc_lengths: true
num_workers: 32
drop_last: true
pin_memory: true
prefetch_factor: 8
persistent_workers: true
memmap_dtype: uint32
timeout: 0
instance_filter:
repetition_max_period: 13
repetition_min_period: 1
repetition_max_count: 32
paths:
# ProofPile 2: Algebraic Stack Data
- http://olmo-data.org/preprocessed/proof-pile-2/v0_decontaminated/algebraic-stack/train/allenai/dolma2-tokenizer/part-00-00000.npy
- http://olmo-data.org/preprocessed/proof-pile-2/v0_decontaminated/algebraic-stack/train/allenai/dolma2-tokenizer/part-01-00000.npy
- http://olmo-data.org/preprocessed/proof-pile-2/v0_decontaminated/algebraic-stack/train/allenai/dolma2-tokenizer/part-02-00000.npyЭтап 2 для 1B
Для модели 1B можно сделать три прохода обучения, с разным порядком данных, на 50 миллиардах высококачественных токенов, а затем усреднили и сварили модельный суп.
"суп"/"soup" — удобный термин, появившийся где-то в 2022-23 годах. Это метафора: когда варишь суп, нужно в кастрюлю сложить много ингридиентов. В данном случае, это модели, которые усредняются и склеиваются для достижения общего результата.
Контрольная точка | Конфигурация обучения | WandB | |
|---|---|---|---|
случайное зерно 42049 | |||
случайное зерно 666 | |||
случайное зерно 42 | |||
итоговая смешанная модель | нет конфигурации, просто усреднили веса в Python |
Этап 2 для 7B
Для модели 7B проводилось три этапа обучения, с разным порядком данных, на 50 миллиардах высококачественных токенов, а затем всё усреднено в модельном супе.
Контрольная точка | Конфигурация обучения | WandB | |
|---|---|---|---|
случайное зерно 42 | |||
случайное зерно 42069 | |||
случайное зерно 666 | |||
итоговая смешанная модель | нет конфигурации, просто усреднили веса в Python |
Указанные здесь конфигурации настроены для загрузки последней контрольной точки после этапа 1, чтобы обучение началось прямо с этого момента. Как всегда, можно подшаманить и что-нибудь оптимизировать.
Этап 2 для 13B
Для модели 13B проводилось три этапа обучения с разным порядком данных на 100 миллиардах высококачественных токенов и одно дополнительное обучение на 300 миллиардах высококачественных токенов. Как всегда, результаты сварили в супе.
Контрольная точка | Конфигурация обучения | WandB | |
|---|---|---|---|
случайное зерно 1110, 100B | |||
случайное зерно 2662, 100B | |||
случайное зерно 6209, 100B | |||
случайное зерно 2662, 300B | |||
итоговая смешанная модель | нет конфигурации, мы просто усреднили веса в Python |
Так же как и на прошлом этапе, конфигурации настроены для загрузки последнего чекпоинта после этапа 1, чтобы обучение началось прямо с этого момента. Как всегда, можно подшаманить и что-нибудь оптимизировать.
Примечание: Всю информацию о модели 32B можно найти в репозитории OLMo-core.
Варианты с инструкциями
По ссылкам доступен полный набор instruct моделей:
Примечание: эти модели специально тренируются, чтобвы выполнять инструкции. Если базовые модели натренированы только на предсказание следующего токена и отлично подходит как база для чего-то большего, то OLMo2 Instruct лучше подходит для ответов на вопросы, суммаризации текстов и так далее. Во флагманских моделях компаний типа Anthropic подразумевается тренировка на соответствие человеческим ценностям, но конкретно в OLMo2 Instruct этим особо не заморачивались. Поэтому, если вы собираетесь использовать модельку в продакшен среде, подумайте что будете делать, когда она ответит пользователю какую-нибудь невероятную дичь.
Инференс
Вы можете использовать интеграцию с Hugging Face для запуска инференса на контрольных точках OLMo Transformers:
from transformers import AutoModelForCausalLM, AutoTokenizer
olmo = AutoModelForCausalLM.from_pretrained("allenai/OLMo-2-0425-1B")
tokenizer = AutoTokenizer.from_pretrained("allenai/OLMo-2-0425-1B")
message = ["Сшит колпак не по колпаковски. Что с колпаком станет, коли колпак переколпаковать?"]
inputs = tokenizer(message, return_tensors='pt', return_token_type_ids=False)
# необязательная проверка cuda
# inputs = {k: v.to('cuda') for k,v in inputs.items()}
# olmo = olmo.to('cuda')
response = olmo.generate(**inputs, max_new_tokens=100, do_sample=True, top_k=50, top_p=0.95)
print(tokenizer.batch_decode(response, skip_special_tokens=True)[0])
Альтернатива — абстракция pipeline от Hugging Face:
from transformers import pipeline
olmo_pipe = pipeline("text-generation", model="allenai/OLMo-2-0425-1B")
print(olmo_pipe("Ехал Грека через реку, видит Грека — в реке что?"))
Квантизация
olmo = AutoModelForCausalLM.from_pretrained("allenai/OLMo-2-0425-1B", torch_dtype=torch.float16, load_in_8bit=True) # требуется bitsandbytes
Квантизированная модель чувствительна к типам входных данных и обработке CUDA. Чтобы избежать потенциальных проблем, рекомендуется явно преобразовывать входные идентификаторы в CUDA, используя: inputs.input_ids.to('cuda')
Оценка
Дополнительные инструменты для оценки моделей OLMo доступны в репозиториях OLMo Eval и olmes.
Хостинг на Modal.com
Modal - довольно популярный зарубежный хостинг для моделей, который используют разработчики OLMo. Купить его карточкой МИР, конечно, нельзя — кроме хостинга типа Cloud.ru у нас не так много опций. Но вот турецкой карточкой — уже можно попробовать.
Пример скрипта для размещения модели OLMo 2 на Modal.com с использованием OpenAI API предоставлен в ./scripts/olmo2_modal_openai.py. Для его запуска:
Следуйте инструкциям в разделе "Getting Started" в руководстве Modal.com, чтобы установить библиотеку Modal и необходимые консольные утилиты.
Следуйте инструкциям в разделе Secrets руководства Modal.com, чтобы создать секрет Modal с именем "example-secret-token", который определяет значение переменной MODAL_TOKEN для вашего сервера.
Затем выполните:
modal deploy ./scripts/olmo2_modal_openai.pyЭндпоинт можно проверить с помощью curl:
curl -X POST \
-H "Authorization: Bearer [секретный токен из предыдущего шага]" \
-H "Content-Type: application/json" \
-d @body.json \
https://[домен эндпоинта из предыдущей команды]/v1/chat/completions
где body.json имеет формат:
{
"model": "OLMo-2-1124-13B-Instruct",
"messages": [
{
"role": "user",
"content": "Что общего у письменного стола и Дональда Трампа?"
}
],
"max_tokens": 100,
"temperature": 0.9,
"stream": true
}
Выводы
Миниатюрная (доступная не только арабским шейхам), качественная, полностью открытая кодовая база, бесплатно. Идеальная штука, чтобы ставить эксперименты и тренироваться в алгоритмах. Лицензия Apache 2.0. Идеалити.
Разработчики не собираются останавливаться на достигнутом, следите за обновлениями их странички на Hugging Face.
