Как быстро адаптировать Nemotron без обучения с нуля

Itinai.com it company office background blured chaos 50 v f97f418d fd83 4456 b07e 2de7f17e20f9 1

TL;DR

TwoTower разделяет диффузию на замороженную AR‑башню контекста и обучаемую башню денойзера.
Сохраняет 98,7 % качества автокодировочного базиса при 2,42 × ускорении генерации (γ=0.8, S=16, 2×H100).
Денойзер обучён на ≈2,1 трлн токенов; основной бэкбо́н предварительно обучен на 25 трлн токенов.
Один чекпоинт поддерживает режимы диффузии, «mock‑AR» и классического AR.

Nemotron‑Labs‑TwoTower

TwoTower — блочно‑авторегрессивная диффузионная модель, построенная на Nemotron‑3‑Nano‑30B‑A3B (открытый гибридный бэкбо́н). В каждый из двух «башен» входит 52 слоя: 23 Mamba‑2, 6 self‑attention и 23 Mixture‑of‑Experts. Общее количество параметров ≈ 60 млрд, активных ≈ 3 млрд на токен в каждой башне. MoE использует 128 экспертов, из них активируются 6 + 2 общих.

Обучена только денойзер‑башня; контекст‑башня остаётся замороженной и сохраняет способности автокодировочного предобучения.

Как работает TwoTower

Контекст‑башня обрабатывает подсказку и уже «коммитнутые» токены, формируя KV‑кеши и состояния Mamba‑2. Денойзер‑башня получает «шумные» блоки, использует двунаправленное внимание внутри блока и кросс‑attention к соответствующим слоям контекст‑башни. Это даёт многомасштабный доступ к репрезентациям бэкбо́на.

Каждый блок начинается с S [MASK] токенов, которые денойзер уточняет за T шагов, после чего часть токенов, удовлетворяющих порогу уверенности γ, «коммитится». Затем контекст‑башня обновляет свой кэш, и процесс переходит к следующему блоку. Благодаря «коммиту» сразу нескольких токенов, скорость превышает традиционный AR‑декодер, где каждый токен генерируется отдельно.

Бенчмарки

Оценка в BF16 на 2 × H100, γ=0.8, S=16. Таблица сравнивает базовый AR‑вариант с TwoTower‑диффузией.

MMLU (5‑shot): 78.56 → 78.24
ARC‑Challenge (25‑shot): 91.72 → 92.66
HumanEval (0‑shot): 79.27 → 75.58
Общее качество: 98.7 % от AR‑базы
Пропускная способность: 2.42 × ускорения

Общая информированность практически не меняется, заметно снижается только в коде и математике.

Три режима генерации

generate_mask_diffusion() — полный TwoTower‑режим (2 GPU, ~59 GB на GPU в BF16).
generate_mock_ar() — имитация AR, один токен за шаг.
generate_ar() — классический автокодировочный вывод, один GPU.

Пример запуска:

import torch
from transformers import AutoTokenizer, AutoModelForCausalLM

model_name = "nvidia/Nemotron-Labs-TwoTower-30B-A3B-Base-BF16"
tokenizer = AutoTokenizer.from_pretrained(model_name)
model = AutoModelForCausalLM.from_pretrained(
    model_name, torch_dtype=torch.bfloat16, trust_remote_code=True,
)
model.place_towers_on_devices("cuda:0", "cuda:1")
model.eval()

prompt = "France is a country "
inputs = tokenizer(prompt, return_tensors="pt").to("cuda:0")

outputs = model.generate_mask_diffusion(
    inputs["input_ids"], max_new_tokens=128,
    block_size=16, steps_per_block=16, mask_token_id=3,
    temperature=0.1, confidence_threshold=0.8,
    eos_token_id=tokenizer.eos_token_id,
)
print(tokenizer.decode(outputs[0][inputs["input_ids"].shape[1]:], skip_special_tokens=True))

Где применить

Быстрая пакетная генерация: уменьшение качества на 1.3 % в пользу 2.42‑кратного ускорения идеальны для синтетических датасетов.
Тонкая регулировка качества‑скорости: изменение γ позволяет находить компромисс под конкретные SLA.
Единый чекпоинт для разных задач: одна модель сразу поддерживает диффузию, спекулятивный декодинг и традиционный AR‑скори́нг.

Сильные и слабые стороны

Плюсы

Открытый вес под NVIDIA Nemotron Open Model License — можно использовать в коммерции.
98.7 % качества при 2.42 × ускорении (γ=0.8).
Одни чекпоинты покрывают три режима вывода.
Денойзер обучён лишь на части предобучения (≈2.1 трлн токенов).
Память кэша масштабируется как в AR‑базе.

Минусы