Как превратить Gemma 4 31B Dense в MoE: дообучение роутера и экспертов на Hugging Face

Вы когда-нибудь смотрели на плотную модель и думали: «Эх, вот бы ей добавить экспертов и роутер, чтобы она думала как Mixture of Experts, но не переплачивать за обучение с нуля»? Я — да. И, кажется, Google (а точнее сообщество) придумало нечто гениальное: additive MoE. Берём готовую плотную Gemma 4 31B, включаем флаг enable_moe_block — и вуаля, модель превращается в MoE. Осталось только доучить роутер и экспертов на паре тысяч шагов. Звучит как магия? Давайте разбираться на реальном примере.

Зачем вообще мутить плотную модель в MoE?

MoE (Mixture of Experts) позволяет при том же количестве параметров использовать только часть нейронов на каждом токене. Это даёт прирост скорости инференса без потери качества. Но обучение MoE с нуля — это ад. Нужно много данных, много карт, терпение. А если у вас уже есть плотная Gemma 4 31B, вы можете просто «нарастить» на неё дополнительные экспертные слои, оставив старые веса нетронутыми. Аддитивный подход: вы вставляете между существующими FFN новые эксперты и обучаете только их и роутер. Веса оригинальной модели заморожены. Это радикально дешевле.

Кстати, похожий трюк недавно показывали в статье про склеивание Gemma и DeepSeek — там без дообучения, но тоже через перекомпоновку слоёв. Наш метод сложнее, но зато даёт настоящий MoE.

Что такое enable_moe_block и как он работает?

В конфигурации Gemma 4 есть скрытый параметр enable_moe_block. По умолчанию он выключен. Если его включить, модель перестаёт использовать стандартные FFN-слои и начинает работать как MoE: каждая позиция теперь имеет роутер, который направляет токен к одному из нескольких экспертов. Сами эксперты при этом поначалу пусты (или инициализируются копией оригинального FFN). Самое смешное: физически количество параметров не меняется, пока вы не добавите новые эксперты. Именно для этого и нужен additive approach — вы добавляете параллельные expert layers к существующему FFN.

Знакомство с AIOne-Agent-52B

Команда JDONE-Research уже выложила на Hugging Face модель JDONE-Research/AIOne-Agent-52B. Это Gemma 4 31B Dense, мутированная в MoE с 4 экспертами на слой. Они дообучали только роутер и экспертов на датасете агентных диалогов. Результат: модель с 52B параметров (31B старых + 21B новых экспертов) активных только 18B на токен. Скорость инференса — около 120 TPS на двух RTX 3090, как писали в этом обзоре. Мы же пойдём по их стопам, но соберём свою версию.

Пошаговая инструкция: как превратить Gemma 4 31B в MoE

1 Подготовка окружения

Устанавливаем зависимости. Всё стандартно: transformers, peft, accelerate, bitsandbytes, datasets.

pip install transformers peft accelerate bitsandbytes datasets

2 Загружаем плотную модель и включаем MoE режим

Загружаем Gemma 4 31B через AutoModelForCausalLM, но перед этим подменим конфиг. Добавляем параметр enable_moe_block=True и указываем num_experts=4. Также нужно задать moe_type="additive" — это скажет модели, что эксперты будут добавлены к существующим FFN, а не заменят их.

from transformers import AutoModelForCausalLM, AutoConfig
import torch

model_name = "google/gemma-4-31b"
config = AutoConfig.from_pretrained(model_name, trust_remote_code=True)
config.enable_moe_block = True
config.num_experts = 4
config.moe_type = "additive"  # ключевой флаг
config.moe_topk = 2  # каждый токен активирует 2 эксперта

model = AutoModelForCausalLM.from_pretrained(
    model_name,
    config=config,
    torch_dtype=torch.bfloat16,
    device_map="auto",
    trust_remote_code=True
)

3 Добавляем новые эксперты и роутер

В аддитивном режиме модель создаст пустые эксперты. Но нам нужно их проинициализировать. Лучший способ — скопировать веса оригинального FFN в каждый экспертный слой, а затем немного расшатать (добавить шум). Иначе обучение будет слишком долгим.

import copy

for layer in model.model.layers:
    if hasattr(layer, "moe_ffn"):
        original_ffn = layer.mlp  # оригинальный FFN
        for expert_idx in range(config.num_experts):
            expert = layer.moe_ffn.experts[expert_idx]
            # копируем веса из оригинального FFN
            expert.gate_proj.weight.data = original_ffn.gate_proj.weight.data.clone()
            expert.up_proj.weight.data = original_ffn.up_proj.weight.data.clone()
            expert.down_proj.weight.data = original_ffn.down_proj.weight.data.clone()
            # добавляем небольшой шум
            expert.gate_proj.weight.data += torch.randn_like(expert.gate_proj.weight.data) * 0.01

Роутер (gate) обычно инициализируется случайно. Но можно сделать хитрее: инициализировать его так, чтобы первые токены равномерно распределялись по экспертам. В Transformers это делается автоматически при config.moe_type="additive", но бывают баги. Лучше перепроверить.

4 Замораживаем всё, кроме роутера и новых экспертов

Это ключевой момент. Мы не хотим трогать оригинальные веса. Замораживаем все параметры, затем размораживаем только те, что относятся к экспертам и роутеру. Самый простой способ — пройтись по именам модулей.

for name, param in model.named_parameters():
    param.requires_grad = False
    if "moe_ffn" in name or "moe_gate" in name:
        param.requires_grad = True

5 Дообучаем на маленьком датасете

Берём любой датасет для instruction tuning (например, databricks/databricks-dolly-15k). Обучаем 500-1000 шагов с маленьким learning rate (1e-4). Используем QLoRA для экономии памяти, хотя можно и полный fine-tune.

from transformers import Trainer, TrainingArguments
from datasets import load_dataset

dataset = load_dataset("databricks/databricks-dolly-15k", split="train")
training_args = TrainingArguments(
    output_dir="./gemma4-additive-moe",
    per_device_train_batch_size=1,
    gradient_accumulation_steps=4,
    num_train_epochs=1,
    learning_rate=1e-4,
    bf16=True,
    logging_steps=10,
    save_steps=200,
    max_steps=500,
    report_to="none"
)
trainer = Trainer(
    model=model,
    args=training_args,
    train_dataset=dataset,
    tokenizer=tokenizer
)
trainer.train()

6 Сохраняем и публикуем на Hugging Face

После обучения сохраняем модель через model.save_pretrained(). Важно сохранить и конфиг с enable_moe_block=True. При загрузке на инференс нужно передать тот же config. Публикуем на HF, и готово — у вас персональный MoE.

Сравнение с альтернативами

Давайте честно: почему не взять готовый Mixtral 8x7B или DeepSeek-MoE? Потому что они не имеют агентной специализации. А наш метод позволяет взять именно ту dense модель, которая уже хорошо решает вашу задачу (например, Gemma 4 отлично справляется с инструментами, как показано в FunctionGemma 270M), и добавить ей «экспертность» без потери накопленного знания. Плюс это в разы дешевле полного обучения MoE. Единственный минус — вы ограничены архитектурой исходной модели, и если она не поддерживает additive MoE, фокус не пройдёт. К счастью, Gemma 4 — поддерживает.

Есть ещё гибридный подход из статьи Cactus Hybrid Router — там вообще роутер поверх нескольких готовых моделей. Но это другая история.

Типичные грабли (и как их избежать)

Во время экспериментов я наткнулся на те же ошибки, что описаны в статье Три неудачных попытки дообучения Gemma 4. Вот главные:

• Забыли заморозить оригинальные веса — градиенты затирают исходные знания.
• Не добавили load balancing loss — роутер забивает на новых экспертов.
• Слишком большой learning rate — новые эксперты переобучаются и убивают выход.
• Инициализация экспертов нулями или случайно — обучение идёт в 10 раз дольше. Лучше копировать оригинальный FFN с шумом.

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

Тем, кто хочет получить MoE-модель для агентов, инструментов или специализированных диалогов, но не имеет ресурсов на обучение с нуля. Исследователям, которые изучают поведение экспертов. И просто фанатам Gemma 4, которых уже прорвало от количества её модификаций — взгляните хотя бы на статью 20 финтюнов Gemma 3 от DavidAU, там народ творит безумные вещи.

И последнее: не пытайтесь добавить сразу 64 эксперта на слой. Начните с 2-4, проверьте качество. А то получите, как в шутке про «Тёмную цепочку мыслей» (статья) — модель думает, но ответы — каша. MoE требует баланса.

Подписаться на канал