🌟 NVIDIA Jet-Nemotron: гибридная архитектура, которая быстрее SOTA-моделей в 53 раза.

Jet-Nemotron - новая архитектура языковых моделей, которая, по заявлениям NVIDIA, превосходит по эффективности топовые опенсорс-модели.

На H100 обещают ускорение пропускной способности при генерации до 53.6 раз, если работать с контекстом в 256 тыс. токенов и максимальным размером батча. Такой скачок производительности стал возможен благодаря двум ключевым инновациям: пайплайну PostNAS и новому блоку линейного внимания JetBlock.

🟡PostNAS (Post Neural Architecture Search).

Суть PostNAS состоит в отказе от дорогостоящего обучения новых архитектур с нуля. Вместо этого берут уже предварительно обученную модель и запускают процесс постобработки. Пайплайн сначала анализирует модель и определяет, какие слои внимания вносят наибольший вклад в её работу, а какие - не так уж и важны. Дальше он ищет оптимальное расположение для слоёв полного внимания и подбирает улучшенный дизайн для остальных блоков.

🟡JetBlock - модуль линейного внимания.

Его фишка - динамические сверточные ядра, генерируемые на лету в зависимости от входных данных и применяемые к value-токенам.

Прямое сравнение с Mamba2 Block, проведенное на идентичных данных и с одинаковыми параметрами обучения, показало существенный прирост в точности при сохранении той же пропускной способности во время обучения и инференса.

🟡Третий элемент успеха - аппаратно-ориентированный поиск архитектуры.

Вместо того чтобы использовать количество параметров в качестве прокси-метрики для эффективности, авторы напрямую оптимизируют архитектуру под целевое железо (H100), используя в качестве цели именно пропускную способность генерации.

Ключевое открытие тут в том, что размер KV-кэша, а не количество параметров, является критическим фактором, ограничивающим скорость генерации на длинных контекстах, поскольку декодирование упирается в пропускную способность памяти.

Фиксируя размер кэша, они провели поиск по размерности ключей/значений и числу голов внимания, обнаружив конфигурации, которые при том же объеме кэша и схожей пропускной способности используют больше параметров для достижения более высокой точности.

Итоговый дизайн Jet-Nemotron, построенный на базе Qwen 2.5, включает всего 2 full-attention слоя (для retrieval) и 2 слоя со скользящим вниманием (SWA, для MMLU), остальные — JetBlock.

Что касается конкретных моделей, то уже есть Jet-Nemotron-2B и Jet-Nemotron-4B. По результатам тестов, они как минимум не уступают по точности ведущим эффективным моделям, например, Qwen3, на целом ряде бенчмарков. При этом младшая модель Jet-Nemotron-2B работает в 21 раз быстрее, чем Qwen3-1.7B-Base, а старшая, Jet-Nemotron-4B, обгоняет её уже в 47 раз.

▶️ Код и веса моделей обещают опубликовать сразу после завершения юридической проверки.


🟡Страница проекта
🟡Arxiv
🖥GitHub (Coming Soon)


@ai_machinelearning_big_data

#AI #ML #LLM #NVIDIA #JetNemotron