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

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

OLMoE-1B-7B; (https://huggingface.co/allenai/OLMoE-1B-7B-0924)

OLMoE-1B-7B-SFT (https://huggingface.co/allenai/OLMoE-1B-7B-0924-SFT) - промежуточная версия базовой модели, прошедшая этап SFT-посттрейна, но без Instruct-дообучения;

OLMoE-1B-7B-Instruct (https://huggingface.co/allenai/OLMoE-1B-7B-0924-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 (https://arxiv.org/pdf/2409.02060)

Dataset (https://huggingface.co/datasets/allenai/OLMoE-mix-0924)

Набор моделей (https://huggingface.co/collections/allenai/olmoe-66cf678c047657a30c8cd3da)

Github (https://github.com/allenai/OLMoE)

Источник: github.com

Комментарии: