NeMo Automodel

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