NVIDIA NeMo AutoModel 가 Mixture-of-Experts 모델의 파인튜닝을 가속화합니다

NVIDIA NeMo AutoModel은 Mixture-of-Experts 모델의 파인튜닝을 가속화하여 Hugging Face Transformers v5 대비 3.4‑3.7배 높은 학습 처리량과 29‑32% 낮은 GPU 메모리 사용량을 제공하며, 단 한 줄의 import 변경만으로 사용할 수 있습니다.

배경

Mixture-of-Experts 모델의 급증은 학습에 어려움을 초래했으며, Transformers v5는 전문가 백엔드, 동적 가중치 로딩, 분산 실행을 통해 이를 해결했지만 여전히 통합된 통신 커널이 부족했습니다.

NeMo AutoModel: 동일한 API, 더 높은 성능

NeMo AutoModel은 AutoModelForCausalLM을 상속하므로 사용자는 import 라인 하나만 바꾸어 Hugging Face Transformers에서 NeMo AutoModel로 전환하고, Expert Parallelism, DeepEP 융합 all‑to‑all 디스패치, TransformerEngine 커널을 얻을 수 있으며 학습 코드를 수정할 필요가 없습니다.

NeMo AutoModel로 모델을 로드하는 방식은 import만 제외하고 Hugging Face 버전과 동일합니다:

from nemo_automodel import NeMoAutoModelForCausalLM
model = NeMoAutoModelForCausalLM.from_pretrained(
    "nvidia/NVIDIA-Nemotron-3-Nano-30B-A3B-BF16",
    dtype=torch.bfloat16,
)

Qwen3, NVIDIA Nemotron, GPT‑OSS, DeepSeek V3 등 인기 MoE 아키텍처에 대해 NeMo AutoModel은 TransformerEngine attention, 융합 선형 레이어, 맞춤형 전문가 커널을 갖춘 손‑튜닝 구현을 제공합니다. 다른 모델에 대해서는 기본 Hugging Face 구현을 사용하면서도 Liger 커널 패치와 같은 최적화를 적용합니다. 결과 모델은 device_mesh를 전달하여 다중 GPU 학습으로 확장됩니다.

8 GPU에 걸쳐 Expert Parallelism을 사용해 Nemotron 3 Nano 30B A3B를 학습하려면, 사용자는 분산 mesh 구성을 추가합니다:

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(0)
torch.cuda.set_device(int(os.environ.get("LOCAL_RANK", 0)))

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()

이렇게 하면 from_pretrained() 호출 하나만으로 FSDP2, Expert Parallelism, TransformerEngine 커널, DeepEP 디스패치에서 얻는 속도, 확장성, 메모리 최적화를 모두 활용할 수 있습니다.

성능 비교

벤치마크 결과, NeMo AutoModel은 550B 모델을 16노드에 걸쳐 실행한 경우와 단일 노드에서 두 개의 30B MoE 모델을 실행한 경우 모두 GPU당 토큰‑처리량이 3.4‑3.7배 높고 피크 메모리가 29‑32% 낮습니다.

Nemotron 3 Ultra 550B A55B (전체 파인튜닝, 다중 노드)

Expert Parallelism을 통해 16개의 H100 노드에서 550B 파라미터를 가진 Nemotron 3 Ultra 모델을 전체 파인튜닝할 수 있으며, 이 규모에서는 Transformers v5가 메모리 부족으로 실행되지 못합니다.

방법론:

하드웨어: 16x H100 80GB (128 GPU)
Expert Parallelism: EP=64
로컬 배치 크기: 2
시퀀스 길이: 4,096
특징: MTP, 활성화 체크포인팅, fused linear cross-entropy
커널: DeepEP dispatch + torch_mm experts + TransformerEngine

결과:

지표	NeMo AutoModel (EP=64)
TPS/GPU (평균)	815
TFLOP/s/GPU	~293
최대 메모리	58.2 GiB

Transformers v5는 이 규모에서 메모리 부족으로 실행되지 않으므로 v5 수치는 제공되지 않습니다.

단일 노드 30B MoE 벤치마크

단일 8‑GPU H100 노드에서 NeMo AutoModel은 Transformers v5보다 3.36‑3.69배 높은 처리량을 보이며 Qwen3‑30B‑A3B와 Nemotron 3 Nano 30B‑A3B의 피크 메모리를 29‑32% 줄입니다.

방법론:

하드웨어: 8x H100 80GB (단일 노드)
시퀀스 길이: 4,096
로컬 배치 크기: 1

Qwen3-30B-A3B

지표	v4	v5 (FA2 + grouped_mm)	NeMo AutoModel (EP=8)	v5 → NeMo AutoModel
TPS/GPU (평균)	deadlock	3,075	11,340	3.69x
최대 메모리	—	68.2 GiB	48.1 GiB	-29%
평균 Forward+Loss	—	582 ms	194 ms	3.00x
평균 Backward	—	758 ms	178 ms	4.26x

Transformers v4는 Qwen3 MoE 전문가를 128개의 개별 MLP 모듈을 담은 ModuleList로 저장하고 각각을 별도로 FSDP‑wrap 하여 컬렉티브 불일치가 발생해 deadlock이 발생합니다. Transformers v5는 전문가를 융합된 3D 파라미터 텐서로 저장함으로써 이를 해결했습니다.

Nemotron 3 Nano 30B A3B

지표	v4 (hub code)	v5 (FA2 + grouped_mm + Mamba CUDA)	NeMo AutoModel (EP=8)	v5 → NeMo AutoModel
TPS/GPU (평균)	1,807	4,583	15,421	3.36x
최대 메모리	61.9 GiB	62.1 GiB	42.5 GiB	-32%
평균 Forward+Loss	1,024 ms	283 ms	109 ms	2.60x
평균 Backward	1,246 ms	611 ms	157 ms	3.89x

v4는 trust_remote_code=True(NVIDIA hub 모델링 코드)를 사용하며, 토큰 할당과 관계없이 모든 전문가를 순회하는 루프를 수행해 Qwen3 v4에서 발생한 deadlock을 회피합니다.

속도 향상의 원인

3.4‑3.7배의 속도 향상은 Expert Parallelism이 전문가 가중치를 분할하고, DeepEP가 토큰 라우팅을 연산과 융합하며, TransformerEngine 커널이 attention과 선형 레이어를 가속화한 결과입니다.

Expert Parallelism은 메모리 압력을 감소시킵니다: EP=8은 전문가 가중치를 GPU에 고르게 분산시켜 GPU당 MoE footprint를 8배 축소합니다. Qwen3에서는 피크 메모리가 68.2 GiB에서 48.1 GiB(‑29%)로 감소하고, Nemotron Nano에서는 62.1 GiB에서 42.5 GiB(‑32%)로 감소해 더 큰 배치나 긴 시퀀스를 위한 여유 공간을 제공합니다.
DeepEP는 통신과 연산을 융합합니다: 전문가 라우팅을 위한 별도 AllGather/ReduceScatter 컬렉티브 대신, DeepEP는 토큰 디스패치를 최적화된 GPU 커널로 융합해 통신과 전문가 연산을 겹치게 합니다.
TransformerEngine 커널은 핵심 연산을 가속합니다: TE의 융합 attention, 선형 레이어, RMSNorm 구현은 PyTorch/Flash Attention 대비 일관된 속도 향상을 제공하며, 이는 MoE 레이어뿐 아니라 모든 레이어에 적용됩니다.

Transformers v5 기능을 활용한 HuggingFace AutoModel

NeMo AutoModel은 Transformers v5의 전문가 백엔드(특히 grouped_mm), 동적 가중치 로딩, 텐서‑패럴렐 Expert Parallelism을 기반으로 하며, 여기에 DeepEP와 TransformerEngine을 추가합니다.

Expert 백엔드

Transformers v5에서 가장 영향력 있는 기능 중 하나는 experts_implementation 파라미터이며, 세 가지 전문가 백엔드를 포함합니다:

백엔드	설명	최적 용도
eager	선택된 전문가에 대한 for‑loop	디버깅, 호환성, 정확성. v4에서도 사용 가능.
batched_mm	전문가 파라미터 복제, `torch.bmm`을 통한 단일 배치 GEMM	작은 입력, `torch.compile`과 함께 빠름. v5에서 추가됨
grouped_mm	토큰을 전문가별로 정렬하고 `torch.nn.functional.grouped_mm`을 통한 단일 그룹 GEMM	학습(메모리 효율, 파라미터 복제 없음). v5에서 추가됨

grouped_mm 백엔드는 핵심 학습 최적화입니다. 전문가를 하나씩 순회하는 대신, 할당된 전문가별로 토큰을 정렬하고 하나의 융합된 그룹 매트릭스 곱셈을 실행합니다.

NeMo AutoModel은 이를 한 단계 더 발전시킵니다. 맞춤 구현이 있는 모델에 대해 DeepEP 융합 all‑to‑all 디스패치와 그룹 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
)

Expert Parallelism과 DeepEP

Transformers v5는 전문가 가중치를 GPU에 분할하는 Expert Parallelism 경로를 제공하지만, NeMo AutoModel은 다중 GPU MoE 학습에 맞게 보완된 접근 방식을 취합니다. EP를 자체 병렬 차원으로 만들고, 데이터‑패럴렐 메시와는 별개의 moe_mesh를 사용해 PyTorch DTensor의 Shard(0)으로 구현합니다. 8 GPU에서 EP=8과 DP=8을 동시에 실행해 각 GPU는 자체 데이터 샤드를 학습하면서 전문가 파라미터의 1/8만 보유합니다.

# From nemo_automodel/components/moe/parallelizer.py
from torch.distributed.tensor import Shard, distribute_tensor

# 각 GPU는 전문가 가중치의 1/ep_size만 보유
 distribute_tensor(param, device_mesh, [Shard(0)])

ep_size=8인 8 GPU 환경에서, 예를 들어 Nemotron‑3‑Nano‑30B‑A3B는 약 55 GiB의 전문가 가중치를 가지고 있는데, EP를 적용하면 GPU당 전문가 footprint가 약 55 GiB에서 6.8 GiB로 감소해 FSDP‑only 방식이 메모리 부족으로 실패하던 상황에서도 학습이 가능해집니다.

EP 위에 DeepEP를 통합해 토큰 라우팅을 최적화된 GPU 커널로 융합하고, 그룹 GEMM과 결합해 큰 속도 향상을 달성합니다.

동적 가중치 로딩

Transformers v5는 WeightConverter와 WeightRenaming을 통한 동적 가중치 로딩 시스템을 도입했습니다. 이를 통해 MoE 체크포인트를 융합된 3D 텐서 형태로 저장해 실행 효율을 높일 수 있습니다. WeightConverter는 from_pretrained() 중에 체크포인트 텐서를 변환하는 조합 가능한 연산을 적용합니다.

NeMo AutoModel은 이 v5 API를 직접 활용합니다. MODELS_REQUIRING_TENSOR_MERGING을 통해 20개 이상의 모델(예: Mixtral, Qwen2 MoE, Qwen3 MoE, DeepSeek V2/V3, OLMoE 등)이 이 메커니즘을 사용합니다. 변환은 완전히 역전 가능하며, save_pretrained()는 표준 HF‑format 체크포인트를 생성해 downstream 도구가 그대로 로드할 수 있게 합니다.

시작하기

사용자는 NVIDIA 문서에서 NeMo AutoModel을 설치하고, import 한 줄만 바꾸어 동일한 from_pretrained() API를 사용할 수 있습니다.

자세한 내용은 다음을 참고하십시오:

NeMo AutoModel HuggingFace API 호환성 가이드
NeMo AutoModel 모델 커버리지
NeMo AutoModel 성능 요약
HuggingFace에서의 NeMo AutoModel

결론

NeMo AutoModel은 Hugging Face 사용자가 MoE 학습을 확장할 때 마찰 없는 업그레이드 경로를 제공하며, 다중 배수의 속도 향상과 메모리 절감을 제공하고 vLLM 및 SGLang과 호환되는 표준 Hugging Face 체크포인트를 생성합니다.

코드, 설정 파일, 벤치마크 스크립트는 모두 NeMo AutoModel 저장소에서 확인할 수 있습니다.

NVIDIA NeMo AutoModel 가 Mixture-of-Experts 모델의 파인튜닝을 가속화합니다

NVIDIA NeMo AutoModel 가 Mixture-of-Experts 모델의 파인튜닝을 가속화합니다

배경

NeMo AutoModel: 동일한 API, 더 높은 성능

성능 비교

Nemotron 3 Ultra 550B A55B (전체 파인튜닝, 다중 노드)

단일 노드 30B MoE 벤치마크

Qwen3-30B-A3B

Nemotron 3 Nano 30B A3B

속도 향상의 원인

Transformers v5 기능을 활용한 HuggingFace AutoModel

Expert 백엔드

Expert Parallelism과 DeepEP

동적 가중치 로딩

시작하기

결론

Sources