NeMo Automodel

NeMo Automodel 是 NVIDIA 推出的一个开源 PyTorch DTensor 原生训练库。它支持针对 LLM 和 VLM 的大规模与中小规模预训练及微调，适用于研究和生产环境中的快速实验，并支持包括 FSDP2、张量并行、流水线并行、专家并行和上下文并行在内的多种并行策略。为实现高吞吐量，它集成了来自 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 稠密模型和混合专家（MoE）基准测试
请参阅 NeMo 预训练指南，了解如何使用 NeMo 进行预训练

Update on GitHub