OLMoE: Открытые языковые модели на основе Mixture-of-Experts.OLMoE (Open Mixtur…

🌟 OLMoE: Открытые языковые модели на основе Mixture-of-Experts.

OLMoE (Open Mixture-of-Experts Language Models) – это метод создания языковых моделей, использующий разреженно активированные смеси экспертов (Mixture-of-Experts, MoE). Следуя этому методу была создана и опубликована базовая модель и 2 ее вариации:

🟢OLMoE-1B-7B;

🟠OLMoE-1B-7B-SFT – промежуточная версия базовой модели, прошедшая этап SFT-посттрейна, но без Instruct-дообучения;

🟢OLMoE-1B-7B-Instruct

Модель предварительно обучалась на 5,1 трлн. токенов с 7 млрд. общих параметров, из которых только 1,3 млрд. активируются для каждого входного токена.
Такая конфигурация дает аналогичную стоимость инференса, как при использовании моделей с 1 млрд. параметров (напр. TinyLlama 1B), но требует больше VRAM для хранения 7 млрд общих параметров.

В создании OLMoE-1B-7B была использована архитектура decoder-only трансформера, в которой feedforward network была заменена на Mixture-of-Experts модуль, состоящий из набора меньших feedforward network, собранных в массивы k-экспертов. Ими управляет маршрутизатор (его роль выполняет обученный линейный слой), определяющий, какие из k-экспертов будут активированы для каждого входящего токена. Инференсы k-экспертов умножаются на вероятности маршрутизатора и потом суммируются, образуя общий инференс MoE-модуля.

Претрейн модели выполнялся на очищенных от дубликатов датасетах DCLM и Dolma 1.7 (Github, Wikipedia, наборы научных статей). На следующем этапе инструктивного дообучения были добавлены данные по программированию и математике, чтобы сбалансировать эти области знаний для повышения точности MoE.

Проведенные после обучения эксперименты показывают, что OLMoE-модели обучаются примерно в 2 раза быстрее, чем обычные LLM с эквивалентными активными параметрами.

▶️Пример инференса базовой модели на Transformers:

# NOTE! Install the `transformers` & `torch` libraries first
from transformers import OlmoeForCausalLM, AutoTokenizer
import torch

DEVICE = "cuda" if torch.cuda.is_available() else "cpu"

# Load different ckpts via passing e.g. `revision=step10000-tokens41B`
# also check allenai/OLMoE-1B-7B-0924-SFT & allenai/OLMoE-1B-7B-0924-Instruct
model = OlmoeForCausalLM.from_pretrained("allenai/OLMoE-1B-7B-0924").to(DEVICE)
tokenizer = AutoTokenizer.from_pretrained("allenai/OLMoE-1B-7B-0924")
inputs = tokenizer("Bitcoin is", return_tensors="pt")
inputs = {k: v.to(DEVICE) for k, v in inputs.items()}
out = model.generate(**inputs, max_length=64)
print(tokenizer.decode(out[0]))
# > # Bitcoin is a digital currency that is created and held electronically. 
#No one controls it. Bitcoins aren’t printed, like dollars or euros – they’re produced by people and businesses running computers all around the world, using software that solves mathematical

📌Лицензирование : Apache 2.0 License.

🟡Arxiv
🟡Dataset
🟡Набор моделей
🖥Github

@ai_machinelearning_big_data

#AI #ML #LLM #MoE