NVIDIA NeMo AutoModel을 통한 Transformer 미세 조정(Fine-Tuning) 가속화

NVIDIA NeMo AutoModel은 대규모 맞춤형 생성형 AI 모델을 구축하기 위한 NVIDIA NeMo 프레임워크의 일부인 오픈 라이브러리입니다. NeMo AutoModel은 v5를 기반으로 깔끔하게 구축되었으며, 전문가 병렬성 (Expert Parallelism), DeepEP fused all-to-all dispatch, 그리고 TransformerEngine 커널을 추가합니다. 또한 v5의 동적 가중치 로딩 (dynamic weight loading) 기능을 활용하여 이러한 최적화 기술을 광범위하고 성장하는 다양한 모델 제품군에 적용합니다. 그 결과, 동일한 from_pretrained() API를 사용하면서도(코드 변경 없이 단 한 줄의 import만으로), 네이티브 Transformers v5 대비 MoE 모델 미세 조정 시 **3.4-3.7배 더 높은 학습 처리량 (training throughput)**과 29-32% 적은 GPU 메모리 사용량을 달기할 수 있습니다.

이 블로그에서는 이러한 조합이 어떻게 작동하는지, 그리고 사용자가 API를 변경하지 않고 어떻게 MoE 모델을 더 빠르게 미세 조정할 수 있는지에 대해 자세히 설명합니다.

MoE 모델의 부상은 효율적인 학습에 있어 새로운 과제들을 던져주었습니다. 수백 명의 전문가 (experts) 사이에서 토큰을 라우팅하고, 전문가 행렬 곱셈 (expert matmuls)을 단일 커널로 융합하며, GPU 간에 가중치를 샤딩 (sharding)하고, 통신과 연산을 중첩 (overlapping)시키는 작업은 모두 범용 라이브러리가 기본적으로 제공하는 수준을 넘어서는 인프라를 요구합니다.

Transformers v5 (“v5”)는 전문가 백엔드 (expert backends), 동적 가중치 로딩 (dynamic weight loading), 분산 실행을 위한 텐서 병렬 계획 (tensor parallel plans)과 같은 일급 MoE 지원을 도입했습니다. 또한, v5는 PyTorch의 DeviceMesh를 from_pretrained()에 직접 통합함으로써 분산 학습을 일급 기능으로 만들었습니다.

NeMo AutoModel은 AutoModelForCausalLM을 서브클래싱 (subclassing)하고 전문가 병렬성 (Expert Parallelism, EP), DeepEP fused all-to-all dispatch, 그리고 TransformerEngine 커널을 추가함으로써 v5 위에 구축됩니다. DeepEP는 v5에는 아직 없는 요소로, 통신과 전문가 연산을 중첩시킵니다. 또한 NeMo AutoModel은 각 모델을 로드할 때 v5의 가역적 가중치 변환 (reversible weight conversion) 기능을 활용하기 때문에, 모델별 체크포인트 배관 작업 대신 이러한 재사용 가능한 핵심 연산에 엔지니어링 역량을 집중할 수 있습니다. 동시에 save_pretrained()는 여전히 vLLM 및 SGLang과 같은 도구들이 로드할 수 있는 표준 HF 체크포인트를 생성합니다.

다음 섹션에서는 두 방식이 어떻게 함께 작동하는지, 그리고 16개 노드에 걸친 NVIDIA Nemotron 3 Ultra 550B A55B의 전체 미세 조정 (Full Fine-Tuning)부터 Qwen3-30B-A3B 및 Nemotron 3 Nano 30B A3B와 같은 단일 노드 모델에 이르기까지 우리가 측정한 성능 이득에 대해 살펴봅니다.

NeMo AutoModel의 목표 중 하나는 오픈 소스 커뮤니티를 지원하기 위해 HuggingFace Transformers와의 API 호환성을 확보하는 것입니다. NeMoAutoModelForCausalLM은 AutoModelForCausalLM을 상속받으므로, HF 모델과 작동하는 모든 코드는 AutoModel에서도 동일하게 작동합니다.

다음은 두 방식에서 모델을 로드하는 모습입니다. 오직 import 부분만 달라집니다:

그 단 하나의 import가 많은 작업을 수행합니다. Qwen3, NVIDIA Nemotron, GPT-OSS, DeepSeek V3와 같은 인기 있는 MoE (Mixture-of-Experts) 아키텍처의 경우, NeMo AutoModel은 TransformerEngine 어텐션 (Attention), 융합된 선형 레이어 (Fused Linear Layers), 그리고 커스텀 전문가 커널 (Custom Expert Kernels)이 포함된 수동 튜닝된 구현체를 제공합니다. 그 외의 모든 경우에는 Liger 커널 패칭 (Liger kernel patching)과 같은 최적화를 적용하면서 기본 HF 방식으로 전환됩니다. 어떤 경로를 택하든 결과 모델은 확장할 준비가 되어 있습니다. device_mesh를 전달하기만 하면 추가적인 코드 수정 없이 멀티 GPU 학습이 가능합니다.

NeMo AutoModel이 진정으로 빛을 발하는 지점은 MoE 모델을 멀티 GPU 학습으로 확장하는 것입니다. 8개의 GPU에서 전문가 병렬성 (Expert Parallelism)을 사용하여 Nemotron 3 Nano 30B A3B를 학습시키려면, 분산 메시 (Distributed Mesh) 설정을 추가하면 됩니다:

import os
import torch
import torch.distributed as dist
...

이를 통해 from_pretrained() 호출 한 번만으로 FSDP2, 전문가 병렬성 (Expert Parallelism), TransformerEngine 커널, 그리고 DeepEP 디스패치 (Dispatch)를 통한 속도, 확장성 및 메모리 최적화를 얻을 수 있습니다.

우리는 두 가지 환경에서 NeMo AutoModel을 평가했습니다: 16개 노드에 걸친 프런티어 규모의 550B 모델 전체 미세 조정 (Full Fine-Tuning), 그리고 단일 노드에서의 두 가지 30B MoE 모델 학습입니다. 550B 결과는 대규모 환경에서 왜 전문가 병렬성 (Expert Parallelism)이 필수적인지를 보여주며, 30B 결과는 Transformers v5 대비 GPU당 속도 향상을 수치로 보여줍니다.

Nemotron 3 Ultra 550B A55B는 Mamba2, LatentMoE, 그리고 Multi-Token Prediction (MTP) 기술이 적용된 550B 파라미터 하이브리드 모델입니다. 우리는 모든 파라미터가 업데이트되고 Adam 옵티마이저 (optimizer) 상태가 실체화되는 **전체 미세 조정 (full fine-tune)**을 벤치마크하였으며, 이 규모에서는 **16개의 H100 노드 (128개의 GPU)**가 사용됩니다.

방법론 (Methodology):

파라미터	값
하드웨어	16x H100 80GB (128 GPUs)
...
지표	NeMo AutoModel (EP=64)
---	---
TPS/GPU (평균)	815
...

Transformers v5 컬럼이 없는 이유. Transformers v5는 이 규모에서 메모리 부족 (out of memory) 현상이 발생하므로, 여기에 보고할 v5 수치가 없습니다. AutoModel의 전문가 병렬성 (Expert Parallelism)은 전문가들을 여러 GPU에 분산(shard)하여 메모리 사용량을 예산 범위 내로 유지하며, 이것이 전체 미세 조정 (full fine-tune)을 실행할 수 있게 해주는 핵심입니다. 아래의 30B 비교 결과는 v5가 수용 가능한 범위 내에서 동일한 이점을 보여줍니다.

우리는 8x H100 80GB GPU가 장착된 단일 노드에서 세 가지 접근 방식을 벤치마크했습니다: HF Transformers v4 (허브 코드), HF Transformers v5 (사용 가능한 최적의 최적화 적용), 그리고 NeMo AutoModel (EP=8 + 커스텀 커널).

방법론 (Methodology):

파라미터	값
하드웨어	8x H100 80GB (단일 노드)
...

라우팅 게이트 (routing gate)에 관한 참고 사항. 아래의 NeMo AutoModel 수치는 토큰이 전문가들에게 균등하게 분산되도록 강제하는 균형 잡힌 라우팅 게이트 (balanced routing gate)를 사용합니다. 이는 MoE가 학습을 지향하는 이상적인 동작 지점을 모사합니다. 잘 학습된 모델의 부하 분산 손실 (load-balancing loss)은 전문가 활용도를 거의 균등하게 유도하므로, 균형 잡힌 라우팅은 실제 워크로드가 수렴하는 정상 상태 (steady-state)를 반영합니다 (또한 무작위 더미 토큰이 전문가 병렬성에 주입할 수 있는 지연 노드 (straggler) 노이즈를 제거합니다). v4/v5는 동일한 더미 토큰에 대해 자체적인 라우터를 실행합니다. 따라서 균형 잡힌 게이트는 목표로 하는 MoE 동작 지점에서의 NeMo AutoModel을 측정하며, v4/v5 컬럼은 해당 모델들의 기본 동작 (out-of-the-box behavior)을 반영합니다.

지표 (Metric)	v4	v5 (FA2 + grouped_mm)	NeMo AutoModel (EP=8)	v5 → NeMo AutoModel
TPS/GPU (평균)	데드락 (deadlock)	3,075	11,340	3.69x
...
v4가 데드락에 걸리는 이유: Transformers v4는 Qwen3 MoE 전문가(experts)를 128개의 개별 MLP 모듈로 구성된 ModuleList로 저장하며, 각 모듈은 별도로 FSDP로 래핑(wrapped)됩니다. 순전파(forward pass) 과정에서 토큰을 받은 전문가들만 반복하는 데이터 의존적 루프(data-dependent loop)를 사용합니다. 랭크(rank)마다 데이터가 다르기 때문에, 서로 다른 랭크가 서로 다른 전문가를 건너뛰게 되며, 이로 인해 FSDP AllGather/ReduceScatter 컬렉티브(collectives) 간의 불일치가 발생하여 무한 대기(indefinite hang) 상태에 빠집니다. Transformers v5는 전문가를 융합된 3D 파라미터 텐서(fused 3D parameter tensors)로 저장함으로써 이 문제를 해결합니다 (개별 전문가 모듈 및 개별 전문가 FSDP 컬렉티브를 사용하지 않음).

지표 (Metric)	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
...
v4 설정: trust_remote_code=True (NVIDIA의 hub 모델링 코드). hub 코드의 전문가 루프는 FSDP 안전(FSDP-safe)하므로 (토큰 할당과 관계없이 모든 전문가를 반복함), Qwen3 v4처럼 데드락에 걸리지 않습니다.

NeMo AutoModel이 Transformers v5 대비 3.4~3.7배의 속도 향상을 보이는 이유는 세 가지 요소에서 기인합니다:

전문가 병렬성 (Expert Parallelism)이 메모리 압박을 줄입니다. EP=8은 전문가 가중치를 여러 GPU에 분산하여 GPU당 MoE 점유 공간을 8배 줄입니다. Qwen3의 경우, 이는 피크 메모리(peak memory)를 68.2 GiB에서 48.1 GiB로 감소시킵니다 (-29%). Nemotron Nano의 경우, 62.1 GiB에서 42.5 GiB로 감소시켜 (-32%), 더 큰 배치 크기(batch sizes)나 더 긴 시퀀스(sequences)를 위한 여유 공간을 확보합니다.

DeepEP는 통신과 연산을 융합합니다. 전문가 라우팅(expert routing)을 위해 별도의 AllGather/ReduceScatter 컬렉티브를 사용하는 대신, DeepEP는 토큰 디스패치(token dispatch)를 융합하고 이를 최적화된 GPU 커널(kernels)로 결합하여 통신을 전문가 연산과 중첩(overlapping)시킵니다.

TransformerEngine 커널이 핵심 연산을 가속화합니다. TE의 융합된 어텐션(fused attention), 선형 레이어(linear layers), 그리고 RMSNorm 구현은 MoE 레이어뿐만 아니라 모든 레이어 유형에 대해 PyTorch/Flash Attention 대응물보다 일관된 속도 향상을 제공합니다.

Transformers v5에서 가장 영향력 있는 기능 중 하나는 세 가지 전문가 백엔드(expert backends)를 포함하는 experts_implementation 파라미터입니다.

백엔드 (Backend)	설명 (Description)	최적 용도 (Best for)
eager	선택된 전문가들에 대한 For-loop 수행	디버깅(Debugging), 호환성 및 정확성. v4에서도 사용 가능.
...

grouped_mm 백엔드는 핵심적인 학습 최적화 요소입니다. 전문가들을 하나씩 루프를 돌며 처리하는 대신, 할당된 전문가별로 토큰을 정렬한 뒤 단일 융합 그룹 행렬 곱셈(fused grouped matrix multiplication)을 실행합니다.

NeMo AutoModel은 이를 한 단계 더 발전시켰습니다. 커스텀 구현(custom implementations)을 사용하는 모델의 경우, DeepEP 융합 all-to-all dispatch를 그룹화된 GEMM 커널(grouped GEMM kernels) 및 TransformerEngine 선형 레이어(linear layers)와 결합하여 사용합니다. 발전 과정은 다음과 같습니다:

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
...

Transformers v5는 전문가 병렬화(Expert Parallelism) 경로도 함께 제공합니다. 이는 전문가 가중치(expert weights)를 여러 GPU에 샤딩(sharding)합니다. GroupedGemmParallel 스타일은 각 장치의 로컬 전문가만 로드하며, RouterParallel은 토큰을 라우팅하고 all_reduce를 통해 결과를 결합합니다. 이는 v5의 기존 텐서 병렬화(tensor-parallel) 메커니즘 위에 깔끔하게 구축되어 있습니다. 이를 활성화하면 모델의 tp_plan이 전문가 계획(expert plan)을 반환하므로, 전문가 병렬화가 데이터 병렬화(data parallelism)와 장치 예산을 공유하게 됩니다 (ep × dp = world_size). 여기에서 진행한 단일 노드 30B 벤치마크의 경우, 일반적인 데이터 병렬 v5 (dp=8, ep=1)가 가장 빠른 v5 구성임을 확인하였으며, 이에 따라 해당 v5 설정을 보고합니다.

NeMo AutoModel은 멀티 GPU MoE (Mixture-of-Experts) 학습에 최적화된 상호 보완적인 접근 방식을 취합니다. 이 방식은 PyTorch의 DTensor와 Shard(0)를 사용하여, 데이터 병렬 (data-parallel) 메쉬에서 분할하는 대신 데이터 병렬 메쉬와 나란히 존재하는 전용 moe_mesh로서 EP (Expert Parallelism)를 독자적인 병렬 차원으로 만듭니다. 전문가 메쉬 (expert mesh)가 데이터 병렬화와 직교하기 때문에, 두 방식은 동일한 장치에서 결합됩니다. 8개의 GPU에서 NeMo AutoModel은 ep=8과 dp=8을 함께 실행하므로, 각 GPU는 전문가의 1/8만 보유하면서 각자의 데이터 샤드 (data shard)로 학습을 진행합니다. 전문가 가중치 (expert weights)는 전문가 차원을 따라 GPU 전체에 물리적으로 샤딩 (sharding)됩니다.

# From nemo_automodel/components/moe/parallelizer.py
from torch.distributed.tensor import Shard, distribute_tensor
# 각 GPU는 전문가 가중치의 1/ep_size만 보유합니다
...

8개의 GPU에서 ep_size=8을 사용할 경우, 각 GPU는 전문가 파라미터의 1/8만 보유하게 됩니다. 약 55 GiB의 전문가 가중치를 가진 Nemotron-3-Nano-30B-A3B와 같은 모델의 경우, EP는 GPU당 전문가 점유 공간 (expert footprint)을 약 55 GiB에서 약 6.8 GiB로 줄여주어, FSDP (Fully Sharded Data Parallel) 전용 방식이 메모리 부족(OOM)으로 실행 불가능한 환경에서도 학습을 가능하게 합니다.

EP 외에도 NeMo AutoModel은 토큰 라우팅 (token routing)을 최적화된 GPU 커널로 융합(fuse)하는 DeepEP를 통합하였으며, 그룹화된 전문가 연산을 위한 grouped GEMM과 결합될 때 상당한 속도 향상을 제공합니다. 당사의 대규모 MoE 벤치마크 결과에 따르면, 전체 DeepSeek V3 671B 모델에서 DeepEP + grouped GEMM 조합은 all-gather + looped expert 베이스라인과 비교하여 반복당 비용 (cost per iteration)을 47% 절감했습니다.

Transformers v5는 또한 WeightConverter 및 WeightRenaming을 통해 동적 가중치 로딩 시스템을 도입했습니다. 이를 통해 MoE 체크포인트를 융합된 3D 텐서 (fused 3D tensors)로 저장하여 더욱 효율적인 실행이 가능해집니다. WeightConverter는 from_pretrained() 호출 중에 체크포인트 텐서를 변환하기 위해 결합 가능한 연산 (composable operations)을 적용합니다.

NeMo AutoModel은 이 v5 API를 직접적으로 사용하는 소비자입니다. Mixtral, Qwen2 MoE, Qwen3 MoE, DeepSeek V2/V3, OLMoE 등을 포함하여 20개 이상의 모델 유형이 MODELS_REQUIRING_TENSOR_MERGING를 통해 이 메커니즘을 사용합니다. 이러한 변환은 완전히 가역적입니다. 즉, save_pretrained()를 통해 생성된 표준 HF(HuggingFace) 형식의 체크포인트는 어떤 다운스트림(downstream) 도구에서도 로드할 수 있습니다.

NeMo AutoModel을 사용해 보려면, 공식 문서 페이지를 방문하여 시작해 주세요.

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

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

NVIDIA NeMo AutoModel은 모델 학습을 확장하려는 HuggingFace 사용자들에게 자연스러운 다음 단계입니다. Transformers v5를 기반으로 직접 구축됨으로써, AutoModel은 마찰 없는(zero-friction) 업그레이드 경로를 제공합니다. 즉, 임포트(import) 한 줄만 변경하면 3배 이상 빠른 모델 인스턴스를 얻을 수 있습니다.

Qwen3-30B-A3B 및 Nemotron 3 Nano 30B-A3B에서, 이는 최상의 Transformers v5 구성과 비교했을 때 GPU 메모리를 29-32% 적게 사용하면서도 3.4-3.7배 더 높은 학습 처리량(throughput)을 제공합니다. 또한 진정한 전문가 병렬화(Expert Parallelism)가 GPU 전반에 걸쳐 전문가(experts)를 샤딩(shard)하기 때문에, 동일한 경로를 통해 16개 노드에 걸쳐 Nemotron 3 Ultra와 같은 550B 모델의 전체 미세 조정(full fine-tuning)까지 확장할 수 있습니다. 이 단계는 모델을 메모리에 맞추기 위해 전문가 병렬화가 필수적인 영역입니다. NeMo AutoModel 체크포인트는 표준 HF 형식의 safetensors이므로, vLLM 및 SGLang과 같은 추론 프레임워크에서 배포할 수 있습니다.

코드, 설정(configs), 벤치마크 스크립트는 모두 NeMo AutoModel 리포지토리에서 확인할 수 있습니다.