Ускорение тонкой настройки MoE-моделей с NVIDIA NeMo AutoModel

Экосистема открытых искусственных интеллектуальных моделей переживает фундаментальный сдвиг. Долгое время стандартом де-факто для работы с трансформерами оставалась библиотека Hugging Face Transformers, которая стала фундаментом для тысяч проектов. Однако с выходом версии 5 (Transformers v5) ландшафт изменился: теперь эта библиотека предлагает первоклассную поддержку архитектур Mixture-of-Experts (MoE), которые стремительно становятся доминирующим стандартом для моделей передового уровня (frontier models). Версия 5 привнесла необходимые основы: бэкенды для экспертов, динамическую загрузку весов и распределенное выполнение, сделав MoE расширяемыми и удобными для построения.

Но поддержка архитектуры — это лишь половина дела. Настоящая боль разработчиков и исследователей заключается в эффективности обучения. Модели MoE требуют сложной маршрутизации токенов, слияния матричных умножений и перекрытия коммуникаций с вычислениями. Здесь на сцену выходит NVIDIA NeMo AutoModel — открытая библиотека, часть фреймворка NVIDIA NeMo, предназначенная для создания пользовательских генеративных ИИ-моделей в промышленных масштабах. NeMo AutoModel не просто дополняет Hugging Face, он надстраивается поверх него, добавляя Expert Parallelism (EP), слиянный механизм dispatch DeepEP и ядра TransformerEngine.

Результат этого симбиоза впечатляет: увеличение пропускной способности обучения в 3.4–3.7 раза и снижение потребления памяти GPU на 29–32% при тонкой настройке (fine-tuning) MoE-моделей по сравнению с нативным Transformers v5. И самое главное — это достигается без изменения вашего кода. Вы используете тот же самый API from_pretrained(), меняя лишь одну строку импорта. В этой статье мы подробно разберем, как работает эта комбинация, какие технические инновации стоят за такими цифрами и как вы можете применить их для ускорения своих проектов, даже работая из России с локальными ресурсами.

01Почему MoE-модели требуют новой инфраструктуры

Рост популярности архитектур Mixture-of-Experts (MoE) принес с собой новые вызовы для эффективного обучения. В отличие от плотных (dense) моделей, где каждый токен проходит через все слои, в MoE-моделях токены маршрутизируются к различным «экспертам» (подмоделям). При наличии сотен экспертов это создает огромную нагрузку на инфраструктуру:

• Маршрутизация токенов: Необходимо эффективно распределять токены между сотнями экспертов, избегая узких мест в коммуникации.
• Слияние вычислений: Матричные умножения (matmuls) для каждого эксперта должны быть слиты в единое ядро для максимальной эффективности GPU.
• Шардирование весов: Веса экспертов должны быть распределены между множеством GPU, что требует сложной координации.
• Перекрытие коммуникаций: Обмен данными между GPU должен происходить параллельно с вычислениями, чтобы процессоры не простаивали.

Обычные библиотеки общего назначения, такие как стандартный PyTorch или даже базовая версия Hugging Face Transformers, не предоставляют готовых инструментов для решения этих задач «из коробки». Именно здесь вступает в игру Hugging Face Transformers v5, которая ввела первоклассную поддержку MoE через expert backends, dynamic weight loading и планы тензорного параллелизма. Кроме того, v5 интегрировала PyTorch's DeviceMesh прямо в метод from_pretrained(), сделав распределенное обучение более доступным.

Однако NVIDIA NeMo AutoModel идет дальше. Он наследует класс AutoModelForCausalLM от Hugging Face, добавляя слои Expert Parallelism, слияный механизм dispatch DeepEP и оптимизированные ядра TransformerEngine. Ключевым отличием является использование DeepEP — компонента, которого еще нет в нативном Transformers v5. DeepEP позволяет перекрыть коммуникацию (обмен данными между GPU) с вычислениями экспертов, что критически важно для скорости. При этом NeMo AutoModel использует динамическую загрузку весов v5, что позволяет инженерам сосредоточиться на оптимизации ядра, а не на написании специфичного кода для каждой модели. Итоговые чекпоинты сохраняются в стандартном формате Hugging Face, что гарантирует совместимость с инструментами инференса, такими как vLLM и SGLang.

02API-совместимость: одна строка кода

Одной из главных целей NeMo AutoModel является обеспечение бесшовной интеграции с экосистемой Hugging Face. Библиотека NeMoAutoModelForCausalLM является подклассом AutoModelForCausalLM. Это означает, что логика загрузки модели остается идентичной, меняется только источник импорта.

Рассмотрим, как выглядит загрузка модели в обоих случаях. В стандартном Hugging Face вы пишете:

from transformers import AutoModelForCausalLM

model = AutoModelForCausalLM.from_pretrained("model_name")

В NeMo AutoModel вы просто меняете импорт:

from nemo_automodel import NeMoAutoModelForCausalLM

model = NeMoAutoModelForCausalLM.from_pretrained("model_name")

За этой простой строкой скрывается огромная работа. Для популярных архитектур MoE, таких как Qwen3, NVIDIA Nemotron, GPT-OSS и DeepSeek V3, NeMo AutoModel поставляется с ручно настроенными реализациями. Эти реализации включают внимание TransformerEngine, слитые линейные слои и пользовательские ядра для экспертов. Для остальных моделей библиотека использует стандартный путь Hugging Face, но все равно применяет оптимизации, такие как патчинг Liger kernel. Независимо от пути, полученная модель готова к масштабированию: достаточно передать device_mesh, и вы получаете много-GPU обучение без дополнительных переписываний кода.

03Масштабирование с Expert Parallelism и DeepEP

Где NeMo AutoModel действительно сияет, так это при масштабировании MoE-моделей на несколько GPU. Давайте разберем пример обучения модели Nemotron 3 Nano 30B A3B с использованием Expert Parallelism (EP) на 8 GPU. Для этого необходимо настроить распределенную среду:

import os
import torch
import torch.distributed as dist
from nemo_automodel import NeMoAutoModelForCausalLM
from nemo_automodel.recipes._dist_utils import create_distributed_setup_from_config

dist.init_process_group(backend="nccl")
torch.manual_seed(42)
dist_setup = create_distributed_setup_from_config(
    {
        "strategy": "fsdp2",
        "ep_size": 8,
    },
)
model = NeMoAutoModelForCausalLM.from_pretrained(
    "nvidia/NVIDIA-Nemotron-3-Nano-30B-A3B-BF16",
    dtype=torch.bfloat16,
    distributed_setup=dist_setup,
)
dist.destroy_process_group()

Этот код активирует FSDP2, Expert Parallelism, ядра TransformerEngine и dispatch DeepEP. Но как именно это работает технически?

Expert Parallelism (EP)

Традиционные подходы, такие как Data Parallelism (DP) или даже FSDP (Fully Sharded Data Parallel), могут сталкиваться с нехваткой памяти при работе с огромными моделями MoE. Expert Parallelism решает эту проблему, разделяя веса экспертов между GPU. В NeMo AutoModel EP реализован как отдельное измерение параллелизма, ортогональное к Data Parallelism. Это означает, что на 8 GPU вы можете одновременно использовать ep=8 и dp=8. Каждый GPU обучается на своей части данных, но хранит только 1/8 весов экспертов.

В коде это выглядит так:

from torch.distributed.tensor import Shard, distribute_tensor

# Каждый GPU хранит только 1/ep_size весов экспертов
distribute_tensor(param, device_mesh, [Shard(0)])

Для модели Nemotron-3-Nano-30B-A3B, имеющей около 55 ГБ весов экспертов, EP снижает нагрузку на каждый GPU с ~55 ГБ до ~6.8 ГБ. Это делает обучение возможным в тех случаях, когда подходы, основанные только на FSDP, вылетают с ошибкой нехватки памяти (OOM).

DeepEP: Слияние коммуникации и вычислений

Проблема маршрутизации токенов в MoE-моделях заключается в необходимости собирать токены, отправленные к разным экспертам, и затем объединять результаты обратно. В стандартных реализациях это требует отдельных операций AllGather и ReduceScatter, которые создают задержки. DeepEP решает эту проблему, сливая диспетчеризацию токенов и объединение результатов в оптимизированные ядра GPU. Это позволяет перекрыть коммуникацию с вычислениями экспертов, что дает значительный прирост скорости.

04Сравнение производительности: цифры говорят сами за себя

Чтобы оценить реальную пользу NeMo AutoModel, NVIDIA провела бенчмарки в двух сценариях: полная тонкая настройка (full fine-tuning) модели масштаба 550B параметров на 16 узлах и обучение двух 30B MoE-моделей на одном узле с 8 GPU H100.

Сценарий 1: Nemotron 3 Ultra 550B A55B (Мульти-нод)

Эта гибридная модель с 550 миллиардами параметров использует Mamba2, LatentMoE и Multi-Token Prediction. Бенчмарк проводился на 16 узлах H100 (128 GPU) с Expert Parallelism EP=64.

Ключевой момент: Transformers v5 не смог запустить этот бенчмарк из-за нехватки памяти. Модель просто не поместилась в доступные ресурсы. NeMo AutoModel, благодаря Expert Parallelism, успешно выполнил полную тонкую настройку. Результаты:

• TPS/GPU (токенов в секунду на GPU): 815
• TFLOP/s/GPU: ~293
• Пиковая память: 58.2 GiB

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

Сценарий 2: Однонодовые бенчмарки 30B MoE

Здесь мы сравниваем три подхода на одном узле с 8x H100 80GB: HF Transformers v4, HF Transformers v5 (с лучшими оптимизациями) и NeMo AutoModel (EP=8 + кастомные ядра).

Qwen3-30B-A3B

Интересный факт: Transformers v4 здесь вообще зависает (deadlock). Это происходит потому, что v4 хранит экспертов как список отдельных модулей, и при FSDP-параллелизме разные ранги пропускают разных экспертов, что приводит к рассинхронизации коллективных операций. Transformers v5 исправляет это, храня экспертов как слитые тензоры. Но NeMo AutoModel бьет их обоих:

Примечание: В исходном тексте для Qwen3 указано, что пиковая память NeMo составляет 42.5 GiB, что на 29% меньше, чем у v4 (68.2 GiB). По сравнению с v5 (48.1 GiB) экономия составляет около 11.6%. Однако в заголовке статьи и выводе упоминается экономия 29-32% по сравнению с v5 для Nemotron Nano. Для Qwen3 ускорение по TPS составляет 3.69x.

Nemotron 3 Nano 30B A3B

Для этой модели результаты еще более впечатляющие:

Здесь NeMo AutoModel показывает снижение пиковой памяти на 32% по сравнению с Transformers v5 и ускорение пропускной способности в 3.36 раза. Это достигается за счет трех факторов:

1. Expert Parallelism: Снижает давление на память, распределяя веса экспертов. Для Nemotron Nano это снижает объем с 62.1 ГБ до 42.5 ГБ.
2. DeepEP: Сливает коммуникацию и вычисления, устраняя задержки от отдельных операций AllGather/ReduceScatter.
3. TransformerEngine kernels: Оптимизированные ядра для внимания, линейных слоев и RMSNorm обеспечивают стабильное ускорение по сравнению с аналогами PyTorch/Flash Attention.

05Как Transformers v5 и NeMo AutoModel используют Expert Backends

Transformers v5 ввел параметр experts_implementation, который предлагает три бэкенда для работы с экспертами:

• eager: Цикл for-loop по выбранным экспертам. Используется для отладки и совместимости. Доступен также в v4.
• batched_mm: Дублирует параметры экспертов и выполняет одно большое матричное умножение (GEMM) через torch.bmm. Хорошо работает с малыми входами и torch.compile.
• grouped_mm: Сортирует токены по экспертам и выполняет одно слитое групповое матричное умножение через torch.nn.functional.grouped_mm. Это ключевая оптимизация для обучения, так как она эффективна по памяти и не требует дублирования параметров.

NeMo AutoModel берет оптимизацию grouped_mm на новый уровень. Для моделей с кастомными реализациями он использует DeepEP (слиянный all-to-all dispatch) в сочетании с групповым GEMM и линейными слоями TransformerEngine. Прогрессия выглядит так:

v4 (eager for-loop) → v5 (grouped_mm) → NeMo AutoModel (DeepEP + GMM + TE)

В NeMo AutoModel бэкенд настраивается через BackendConfig:

from nemo_automodel.components.models.common.utils import BackendConfig

backend = BackendConfig(
    attn="te",          # TransformerEngine attention
    linear="te",        # TransformerEngine linear layers
    experts="torch_mm", # Grouped expert matmul
    dispatcher="deepep" # DeepEP fused all-to-all
)

06Динамическая загрузка весов (Dynamic Weight Loading)

Transformers v5 также представила систему динамической загрузки весов через WeightConverter и WeightRenaming. Это позволяет хранить чекпоинты MoE в виде слитых 3D-тензоров для более эффективного выполнения. WeightConverter применяет композиционные операции для преобразования тензоров чекпоинта «на лету» во время вызова from_pretrained().

NeMo AutoModel является прямым потребителем этого API v5. Более 20 типов моделей используют этот механизм через MODELS_REQUIRING_TENSOR_MERGING, включая Mixtral, Qwen2 MoE, Qwen3 MoE, DeepSeek V2/V3, OLMoE и другие. Важно, что эти преобразования полностью обратимы: метод save_pretrained() производит стандартные чекпоинты в формате Hugging Face (safetensors), которые могут быть загружены любыми downstream-инструментами.

07Что это значит на практике

Для исследователей и инженеров, работающих с большими языковыми моделями, внедрение NVIDIA NeMo AutoModel означает несколько ключевых преимуществ:

1. Экономия времени и ресурсов: Ускорение обучения в 3.4–3.7 раза позволяет проводить больше экспериментов за то же время. Снижение потребления памяти на 29–32% позволяет использовать меньшие GPU-конфигурации или увеличивать размер батча, что часто критично для качества обучения.
2. Масштабируемость: Возможность полной тонкой настройки моделей размером 550B параметров, которые ранее были недоступны для такого типа обучения из-за ограничений памяти. Expert Parallelism делает это возможным, разделяя нагрузку между узлами.
3. Простота интеграции: Благодаря полной совместимости API с Hugging Face Transformers, порог входа минимален. Не нужно переписывать пайплайны обучения, достаточно изменить импорт и настроить distributed_setup. Это снижает риск ошибок и упрощает поддержку кода.
4. Доступность из РФ: Поскольку NeMo AutoModel является открытой библиотекой (open-source), доступ к ней не ограничен санкциями. Вы можете скачать код с GitHub, установить зависимости через pip и запустить обучение на локальных GPU-кластерах. Важно лишь убедиться, что у вас есть необходимые версии CUDA и PyTorch, совместимые с вашими GPU (NVIDIA H100, A100, A6000 и др.).

В заключение, NVIDIA NeMo AutoModel — это естественный следующий шаг для пользователей Hugging Face, которые хотят масштабировать обучение моделей. Он предоставляет путь без трения: измените одну строку импорта и получите модель, которая работает более чем в три раза быстрее. С учетом растущей популярности архитектур MoE, такие инструменты, как NeMo AutoModel, становятся необходимыми для эффективной работы с современными ИИ-моделями.

Код, конфигурации и скрипты для бенчмарков доступны в репозитории NeMo AutoModel. Рекомендуется начать с изучения документации и примеров, чтобы оценить потенциал этой технологии для ваших конкретных задач.

Источник: оригинал ↗