NeMo Automodel

NeMo Automodel은 NVIDIA에서 제공하는 오픈소스 PyTorch DTensor 네이티브 학습 라이브러리입니다. FSDP2, 텐서, 파이프라인, 전문가, 컨텍스트 병렬화를 포함한 병렬화 전략을 지원하여 연구 및 프로덕션 환경에서 LLM 및 VLM의 대규모 및 소규모 사전 학습과 파인튜닝을 위한 빠른 실험을 수행할 수 있습니다. 높은 처리량을 위해 DeepEP와 TransformerEngine의 커널을 통합했습니다.

YAML 설정 파일에 학습 실행을 정의합니다(전체 설정 파일 참조).

# Instantiate a Nemotron V3 Nano model
model:
  _target_: nemo_automodel.NeMoAutoModelForCausalLM.from_pretrained
  pretrained_model_name_or_path: nvidia/NVIDIA-Nemotron-3-Nano-30B-A3B-BF16

# Run SFT on HellaSwag
dataset:
  _target_: nemo_automodel.components.datasets.llm.hellaswag.HellaSwag
  path_or_dataset: rowan/hellaswag
  split: train

# Train PEFT adapters
peft:
  _target_: nemo_automodel.components._peft.lora.PeftConfig
  exclude_modules: ["*.out_proj"]  # mamba layers use custom kernels that take in the out_proj.weight directly, thus lora doesn't work here.
  dim: 8
  alpha: 32
  use_triton: True

# Use EP + FSDP2 for training
distributed:
  strategy: fsdp2
  dp_size: none
  tp_size: 1
  cp_size: 1
  ep_size: 4

# ... other parameters

아래 명령을 사용하여 torchrun으로 학습을 시작합니다.

torchrun --nproc-per-node=4 examples/llm_finetune/finetune.py -c /path/to/yaml

Transformers 통합

Transformers에서 지원하는 모든 LLM 또는 VLM은 NeMo Automodel을 통해서도 인스턴스화할 수 있습니다. 전체 모델 지원 범위를 참조하세요.
Hugging Face 모델 위에 AutoModel.from_pretrained()를 기반으로 구축되었으며, 동적 고성능 레이어 교체와 Expert Parallelism(EP)과 같은 더 정교한 병렬화를 지원합니다.
AutoConfig.from_pretrained()의 architecture 필드를 감지하여 Nemotron Nano V3와 같은 커스텀 구현을 자동으로 로드합니다.
Transformers API를 밀접하게 따르며, 즉각적인 호환성을 제공합니다.

리소스

NeMo Automodel
NeMo Transformers API
NeMo Automodel dense 모델과 Mixture-of-Expert(MoE) 벤치마크
NeMo를 사전 학습에 사용하는 방법을 배우려면 NeMo 사전 학습 가이드를 참조하세요.

Update on GitHub