NeMo Automodel

NeMo Automodel (사전학습)

NeMo Automodel은 NVIDIA의 PyTorch DTensor 기반 학습 라이브러리예요. Expert Parallelism과 FSDP2를 활용해 Transformers 모델을 NeMo Automodel로 사전학습하는 방법을 살펴볼게요.

출처: 문서

본문

NeMo Automodel은 NVIDIA의 오픈소스 PyTorch DTensor 네이티브 학습 라이브러리예요. 연구와 프로덕션 환경에서 빠른 실험을 위해 LLM과 VLM의 대규모·소규모 사전학습과 파인튜닝을 지원하며, FSDP2, 텐서, 파이프라인, expert, 컨텍스트 병렬화 같은 병렬화 전략을 갖췄어요. 높은 처리량을 위해 DeepEP와 TransformerEngine의 커널을 통합해요.

# Instantiating Nemotron V3 Nano with Expert Parallelism, FSDP2, and TransformerEngine + DeepEP kernels.
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.cuda.set_device(int(os.environ.get("LOCAL_RANK", 0)))
torch.manual_seed(1111)

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,
)
print(model)
dist.destroy_process_group()

아래 명령으로 torchrun을 사용해 스크립트를 실행해요.

torchrun --nproc-per-node=8 /path/to/script

Transformers 통합 과정

  • Transformers에서 지원하는 어떤 LLM이나 VLM이든 NeMo Automodel을 통해서도 인스턴스화할 수 있어요. 전체 모델 지원 목록을 참고하세요.
  • AutoModel.from_pretrained()를 기반으로 Hugging Face 모델 위에 구축되며, 동적 고성능 레이어 교체와 Expert Parallelism(EP) 같은 더 섬세한 병렬화를 지원해요.
  • AutoConfig.from_pretrained()의 architectures 필드를 감지해 Nemotron Nano V3 같은 커스텀 구현을 자동으로 불러와요.
  • Transformers API를 밀접하게 따라가서 drop-in 호환성을 제공해요.

리소스 (Resources)

더 알아보기 (Learn more)