🔬🧪 Метод и Эксперименты

Семейство Nemotron 3 содержит в себе 3 модели - Nano, Super, Ultra. Из них на текущий момент выложены только веса Nano. В Nano 3B активных параметров из 30B, размеры остальных моделей неизвестны (известно только, что они больше).

Из архитектурных особенностей следует отметить следующее:

📌 В модели для смешивания контекста преобладают Mamba-2 слои, линейные по времени и со скрытым состоянием фиксированного размера. Чтобы не терять в выразительности, в небольшой части блоков оставляют GQA attention с 2 kv-головами.
📌 В Super / Ultra применяют LatentMoE - скрытое состояние проецируется в пространство более низкой размерности перед подачей в экспертов, при этом пропорционально увеличивают количество экспертов (как общее количество, так и активных на токен). Показывают, что на бенчах это дает прирост.
Super / Ultra обучаются на Multi Token Prediction, что позволяет за раз предсказывать не один, а несколько токен на инференсе, и, как утверждается, еще и улучшает качество.

Для повышения эффективности обучения и инференса учат модели в NVFP4 формате (см. блогпост). Ради большей стабильности и качества немногочисленные слои внимания держат в bf16 и выходные проекции в Mamba-2 слоях в MXFP8.

Кроме того, не квантизуют последние 15% слоев. Также используют Адамаровы вращения и стохастическое округление на обратном проходе.

Модели поддерживают контекст до 1М токенов. Утверждается, что слои Mamba-2 неявно задают позиционную информацию и для слоев внимания не требуют позиционные эмбеддинги. RoPE ограничивает обобщаемость модели на контексты длиннее тех, на которых обучалась модель, потому от них решили отказаться. Обучали на последовательностях до 512к токенов, но неплохо обобщаются на 1M, судя по замерам на RULER.

На RL стадии подают задачи не поодиночке, а все разом. Это помогает избежать reward hacking и дает более-менее равномерный рост качества с течением обучения. Для RL используется модифицированный GRPO.

Из приятного, код обучения и данные выложены в публичный доступ.

💡 Выводы

Выглядит как серьезная инженерная работа с использованием SOTA известных рабочих архитектурных и квантизационных приемов. Вполне вероятно, MoE-шные гибриды, обученные в низкой точности станут стандартом в ближайшее время.