OLMo: (Миниатюрная) Открытая Языковая Модель

OLMo — моделька от AI2, разработанная учёными для учёных. Если вы занимаетесь экспериментами над нейронками, это идеальный вариант: весь код и данные, необходимые для тренировки, открыто лежат на GitHub ^[1]. Более того, выложены даже промежуточные чекпоинты, с очень высокой гранулярностью. Это отличает ее от якобы “открытых” нейронок, которые обычно приходят к тебе монолитным финальным бинарём.

Эта статья — короткий гайд, по большей части взятый с GitHub создателей нейросети. Гайд проверен в боевых условиях и действительно работает.

Установка

Сначала установливаем Git, Python и PyTorch ^[2], следуя инструкциям для вашей операционной системы. Инструкции для PyTorch есть здесь ^[3].

Для обучения ^[4] и дообучения рекомендуется установка из исходного кода:

git clone https://github.com/allenai/OLMo.git
cd OLMo
pip install -e .[all]

Также можно установить из PyPI:

pip install ai2-olmo

Shameless plug. Хотите свежего мяса и пейперов про нейронки? Телега @1red2black ^[5].

Предварительное обучение

Предварительное обучение OLMo проходит в два этапа. На первом этапе мы обучаем модель на больших объёмах данных, преимущественно из интернета: OLMo-mix-1124 ^[6] На втором этапе мы обучаем на меньшем объёме высококачественных целевых данных: Dolmino-mix-1124 ^[7]

Вы можете найти все контрольные точки моделей, как минимум каждые 1000 шагов обучения, в форматах OLMo и Hugging Face:

Примечание: для варианта 32B использовались новые скрипты для тренинга, которые тоже открыто опубликованы. Если хотите пройти весь этот путь, для обучения или дообучения OLMo-2 32B посетите OLMo-core ^[20].

Шаги для воспроизведения

Чтобы воспроизвести любой из описанных ниже процессов обучения, выполните:

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:

Данные, в основном, загружаются из веба. Посмотрите, как выглядит раздел 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 годах. Это метафора: когда варишь суп, нужно в кастрюлю сложить много ингридиентов. В данном случае, это модели, которые усредняются и склеиваются для достижения общего результата.

Этап 2 для 7B

Для модели 7B проводилось три этапа обучения, с разным порядком данных, на 50 миллиардах высококачественных токенов, а затем всё усреднено в модельном супе.

Указанные здесь конфигурации настроены для загрузки последней контрольной точки после этапа 1, чтобы обучение началось прямо с этого момента. Как всегда, можно подшаманить и что-нибудь оптимизировать.

Этап 2 для 13B

Для модели 13B проводилось три этапа обучения с разным порядком данных на 100 миллиардах высококачественных токенов и одно дополнительное обучение на 300 миллиардах высококачественных токенов. Как всегда, результаты сварили в супе.

Так же как и на прошлом этапе, конфигурации настроены для загрузки последнего чекпоинта после этапа 1, чтобы обучение началось прямо с этого момента. Как всегда, можно подшаманить и что-нибудь оптимизировать.

Примечание: Всю информацию о модели 32B можно найти в репозитории OLMo-core ^[20].

Варианты с инструкциями

По ссылкам доступен полный набор instruct моделей:

• OLMo2 1B Instruct ^[54]

• OLMo2 7B Instruct ^[55]

• OLMo2 13B Instruct ^[56]

• OLMo2 32B Instruct ^[57]

Примечание: эти модели специально тренируются, чтобвы выполнять инструкции. Если базовые модели натренированы только на предсказание следующего токена и отлично подходит как база для чего-то большего, то 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 ^[58]('cuda')

Оценка

Дополнительные инструменты для оценки моделей OLMo доступны в репозиториях OLMo Eval ^[59] и olmes ^[60].

Хостинг на Modal.com

Modal – довольно популярный зарубежный хостинг для моделей, который используют разработчики OLMo. Купить его карточкой МИР, конечно, нельзя — кроме хостинга типа Cloud.ru ^[61] у нас не так много опций. Но вот турецкой карточкой — уже можно попробовать.

Пример скрипта для размещения модели OLMo 2 на Modal.com ^[62] с использованием OpenAI API предоставлен в ./scripts/olmo2_modal_openai.py ^[63]. Для его запуска:

1. Следуйте инструкциям в разделе “Getting Started” в руководстве ^[64]Modal.com ^[62], чтобы установить библиотеку Modal и необходимые консольные утилиты.

2. Следуйте инструкциям в разделе Secrets ^[65] руководства Modal.com ^[62], чтобы создать секрет Modal с именем “example-secret-token”, который определяет значение переменной MODAL_TOKEN для вашего сервера.

3. Затем выполните:

modal deploy ./scripts/olmo2_modal_openai.py

Эндпоинт можно проверить с помощью curl:

curl -X POST 
  -H "Authorization: Bearer [секретный токен из предыдущего шага]" 
  -H "Content-Type: application/json" 
  -d @body.json 
  https://[веб-точка, созданная modal выше]/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 ^[66].