NeMo Automodel

NeMo Automodel é uma biblioteca de treinamento open-source nativa de PyTorch DTensor da NVIDIA. Ela suporta pré-treinamento e ajuste fino em larga e pequena escala para LLMs e VLMs, permitindo experimentação rápida em ambientes de pesquisa e produção, com estratégias de paralelismo que incluem FSDP2, tensor, pipeline, expert e context parallelism. Para alto throughput, ela integra kernels do DeepEP e do 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()

Execute o script com torchrun usando o comando abaixo.

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

Integração com Transformers

• Qualquer LLM ou VLM suportado no Transformers também pode ser instanciado por meio do NeMo Automodel. Consulte a cobertura completa de modelos.
• Construído sobre modelos do Hugging Face com AutoModel.from_pretrained(), com trocas dinâmicas de camadas de alto desempenho e suporte a paralelismos mais refinados, como Expert Parallelism (EP).
• Detecta o campo de arquitetura em AutoConfig.from_pretrained() para carregar automaticamente implementações personalizadas, como o Nemotron Nano V3.
• Segue de perto a API do Transformers para compatibilidade direta.

Recursos

• NeMo Automodel
• API do NeMo Transformers
• Modelos densos e Mixture-of-Expert (MoE) do NeMo Automodel e benchmarks
• Consulte o guia de fine-tuning do NeMo para aprender a usar o NeMo para fine-tuning