NVIDIA NeMo AutoModel Enterprise로 Transformers MoE 파인튜닝 가속화

이 요약은 AI가 원문을 분석해 생성했습니다. 정확한 내용은 원문 기준으로 확인하세요.

TL;DR

이 글은 HuggingFace Transformers v5의 MoE 기반 인프라 위에 NVIDIA NeMo AutoModel을 얹어 MoE 모델의 파인튜닝 처리량을 크게 끌어올리고 GPU 메모리 사용을 줄이는 방법과 성능 결과를 정량적으로 제시한다. NeMo AutoModel은 AutoModelForCausalLM을 서브클래싱해 from_pretrained() API 호환성을 유지하면서 Expert Parallelism, DeepEP fused all-to-all dispatch, TransformerEngine 커널 등 네이티브 수준의 최적화를 적용해 사용자 코드 변경을 최소화한 채 성능을 향상시킨다.

벤치마크에서 NeMo AutoModel은 single-node 8x H100 환경의 대표적 30B MoE 모델들에서 평균 TPS/GPU를 3.4~3.7배 높였고 피크 메모리는 29~32% 줄였으며, v5가 메모리 초과로 실행하지 못하는 550B 전수 파인튜닝을 Expert Parallelism으로 16노드(128 GPU) 규모에서 가능하게 했다. 성능 향상은 전문가 파라미터의 샤딩으로 메모리를 낮추는 EP, 라우팅과 통신을 연산과 겹치게 처리하는 DeepEP, 그리고 TransformerEngine의 fused 연산 커널이 결합되어 얻어진 결과이다.

이 접근법은 from_pretrained() 한 줄만 바꾸는 '제로 프릭션' 마이그레이션 경로를 제공하므로 기존 Transformers 기반 코드베이스에 최소한의 변화만으로 대규모 MoE 학습을 확장할 수 있다. 다만 최적의 이득을 보려면 H100급 GPU와 분산 설정, Balanced routing 등 MoE의 작동점을 고려한 환경 구성이 필요하며 모델별 사용자 정의 커널 적용 여부에 따라 성능 차이가 발생할 수 있다.

섹션별 상세

대규모 MoE 모델은 토큰 라우팅, 전문가별 행렬연산, 가중치 샤딩 및 통신-계산 중첩 등 일반 라이브러리 수준을 넘어서는 인프라 요구를 갖고 있어 확장 시 메모리와 동기화 문제가 병목으로 작용한다. NeMo AutoModel은 HuggingFace Transformers v5의 AutoModelForCausalLM을 서브클래싱해 동일한 from_pretrained() API를 유지하면서도 내부에서 Expert Parallelism, DeepEP fused dispatch, TransformerEngine 커널을 적용해 로딩부터 분산 실행까지 최적화한다. 문서에서 제시한 사용법은 단일 import 변경으로 모델 인스턴스를 얻는 형태이며 실제로 코드 변경 없이 멀티GPU·멀티노드 교육이 가능한 환경을 제공한다. 이는 기존 Transformers 기반 코드를 최소한의 수정으로 대규모 MoE에 확장할 수 있다는 점에서 운영상 장점이 크다.

성능 벤치마크는 single-node 8x H100 환경에서 대표 30B MoE 모델들을 대상으로 v4, v5, NeMo AutoModel 경로를 비교해 NeMo AutoModel이 평균 TPS/GPU를 3.4~3.7배 높이고 피크 메모리를 29~32% 절감함을 보였다. 예를 들어 Qwen3-30B-A3B는 v5의 3,075 TPS/GPU 대비 NeMo AutoModel에서 11,340 TPS/GPU를 기록해 약 3.69배 향상했고 Nemotron-3-Nano-30B-A3B는 v5 4,583 TPS/GPU 대비 15,421 TPS/GPU로 3.36배 향상했다. 대규모 전수 파인튜닝 사례로 Nemotron 3 Ultra 550B A55B는 Transformers v5로는 메모리 초과로 실행되지 않았으나 NeMo AutoModel의 Expert Parallelism(EP=64)으로 16노드(128 GPU) 환경에서 정상적으로 전체 파인튜닝을 수행할 수 있었다. 이러한 수치는 NeMo AutoModel의 EP, DeepEP, TransformerEngine 세 요소가 결합되어 실무적 확장성과 효율을 동시에 끌어올렸음을 의미한다.

이미지 1은 Qwen3-30B-A3B와 Nemotron-3-Nano-30B-A3B의 TPS/GPU와 Peak Memory를 비교한 막대 차트이다. — Chart이미지 1(URL=https://cdn-uploads.huggingface.co/production/uploads/690d0a6c2c5acfe0e1f4777d/rbCVgV6a18c4UcDsiWfZN.png)은 NeMo AutoModel(NVIDIA 경로)과 v5/HF 경로의 평균 TPS/GPU와 피크 메모리 값을 시각적으로 비교해 NeMo AutoModel이 처리량을 약 3.36~3.69배 높이고 피크 메모리를 약 29~32% 낮춘다는 주장을 직관적으로 뒷받침한다. 이 차트는 본문에 제시된 수치(예: Qwen3 3,075 → 11,340 TPS/GPU, Nemotron Nano 4,583 → 15,421 TPS/GPU 등)와 일치하며 벤치마크 결과의 핵심 정량적 근거로 작용한다.

NeMo AutoModel의 성능 우위는 세 가지 기술적 축에서 나온다. 첫째, Expert Parallelism은 전문가 파라미터를 전문가 차원으로 물리적으로 샤딩해 각 GPU가 전체 전문가 중 일부만 보유하도록 하여 per-GPU 메모리 부담을 ep_size 배만큼 줄인다; 예컨대 ep_size=8이면 전문가 가중치가 8분의 1로 줄어들어 약 55 GiB 규모의 전문가 가중치가 각 GPU에 약 6.8 GiB로 감소한다. 둘째, DeepEP는 토큰 디스패치와 결합(combined dispatch)된 GPU 커널로 AllGather/ReduceScatter와 같은 통신 집약적 단계를 연산과 겹치게 실행해 통신 지연을 숨긴다. 셋째, TransformerEngine의 fused attention·linear·RMSNorm 커널은 표준 PyTorch/FlashAttention 구현 대비 각 레이어의 처리 속도를 일관되게 끌어올려 순전파와 역전파 모두에서 이득을 제공한다. 이 세 축의 결합이 단일 모델 경로에서 v5 대비 3.4~3.7배 처리량을 가능케 한 핵심 메커니즘이다.

운영과 호환성 측면에서 NeMo AutoModel은 v5의 WeightConverter 기반 동적 가중치 로딩을 그대로 활용해 체크포인트 텐서를 실행 시점에 병합·변환하고, save_pretrained()는 표준 HF 포맷의 safetensors를 출력해 vLLM이나 SGLang 같은 추론 프레임워크와 바로 연동이 가능하다. 사용자는 분산 설정을 create_distributed_setup_from_config로 구성한 뒤 from_pretrained() 호출만으로 FSDP, EP, TransformerEngine, DeepEP가 적용된 모델을 얻을 수 있어 마이그레이션 비용이 낮다. 다만 Balanced routing(토큰 균형 라우팅) 같은 운영 파라미터와 하드웨어(예: H100 80GB) 요건이 성능에 큰 영향을 미치므로 환경 구성과 라우터 설정을 신중히 맞춰야 기대한 이득을 확보할 수 있다.

python

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

이 파이썬 코드는 FSDP2와 Expert Parallelism(ep_size=8)을 설정해 NeMo AutoModel으로 사전학습된 MoE 모델을 from_pretrained() 한 줄로 로드하여 분산 학습 환경을 구성하는 예시이다.

실무 Takeaway

NeMo AutoModel은 from_pretrained() API 호환을 유지하면서 내부에서 Expert Parallelism, DeepEP, TransformerEngine을 적용해 기존 코드 변경 없이 MoE 학습을 대폭 가속하므로 대규모 MoE 마이그레이션 비용을 낮춘다.
Expert Parallelism으로 전문가 가중치를 전문가 차원으로 샤딩하면 per-GPU 전문가 메모리 요구량이 ep_size 배만큼 줄어들어 단일 GPU 메모리 한계로 불가능했던 대규모 전수 파인튜닝을 가능하게 한다.
DeepEP는 라우팅 통신과 전문가 계산을 융합된 커널로 겹쳐 실행해 통신 대기 시간을 줄이고 TransformerEngine의 fused 연산은 순·역전파 모두에서 일관된 속도 향상을 제공하므로 실운영 처리량과 비용 효율을 동시에 개선한다.