自动模型

简介

Overview

Nemo AutoModel is a Pytorch DTensor‑native SPMD open-source training library under NVIDIA NeMo Framework, designed to streamline and scale training and finetuning for LLMs, VLMs, diffusion models, and retrieval models. Designed for flexibility, reproducibility, and scale, NeMo AutoModel enables both small-scale experiments and massive multi-GPU, multi-node deployments for fast experimentation in research and production environments. <p align="center"> <a href="https://github.com/NVIDIA-NeMo/Automodel"><picture> <source media="(prefers-color-scheme: light)" srcset="https://raw.githubusercontent.com/NVIDIA-NeMo/Automodel/refs/heads/main/docs/automodel_diagram.png"> <img alt="AutoModel Logo" src="https://raw.githubusercontent.com/NVIDIA-NeMo/Automodel/refs/heads/main/docs/automodel_diagram.png"> </picture></a> </p>

What you can expect:

• Hackable with a modular design that allows easy integration, customization, and quick research prototypes.
• Minimal ceremony: YAML-driven recipes; override any field using CLI.
• High performance and flexibility with custom kernels and DTensor support.
• Seamless integration with Hugging Face for day-0 model support, ease of use, and wide range of supported models.
• Efficient resource management using Kubernetes and Slurm, enabling scalable and flexible deployment across configurations.
• Documentation with step-by-step guides and runnable examples.

Feature List

✅ _Available now (v0.4.0 / 26.04 container) | 🔜 Coming next

High-throughput scalable training - ✅ PyTorch DTensor-native SPMD training Same training script can scale from 1 GPU to large multi-node jobs by changing the device mesh/config. - ✅ Composable Parallelism - PyTorch native FSDP2, HSDP, TP, CP, SP and PP for distributed training. - ✅ Optimized kernels - Uses NVIDIA-oriented kernel paths such as Transformer Engine, DeepEP, FlexAttn, TorchSDPA, fused attention, rotary embeddings, Triton, and optional kernel patches. - ✅ MoE acceleration - Includes MoE routing and DeepEP integration, plus expert-parallel configurations used in DeepSeek, Qwen MoE, GPT-OSS, and Nemotron MoE benchmarks. - ✅ FP8 and mixed precision - FP8 support with torchao and Transformer Engine. - ✅ Activation checkpointing - Trades recomputation for lower activation memory, especially useful with FSDP and memory-efficient losses. - ✅ Memory-efficient loss - Linear-Cut / fused linear cross entropy avoids materializing full logits for the loss, reducing output-layer memory pressure. - ✅ Sequence packing - Packs variable-length examples together to reduce padding compute and improve GPU utilization. - ✅ FlashAttention packed-sequence support - Packed masks can feed variable-length FlashAttention paths using per-document cu_seqlens. - ✅ DCP - Supports PyTorch DCP and SafeTensors, sharded and consolidated layouts, merge/reshard utilities, and Hugging Face-compatible outputs. - ✅ Async checkpointing - Can write checkpoints in the background to reduce training stalls caused by I/O. - ✅ Dion optimizer - Distributed Dion optimizer integration. - ✅ Environment Support - SLURM, interactive, SkyPilot, and Kubernetes (via SkyPilot) launchers.

SOTA algorithms - ✅ Pre-training - Support for model pre-training, including DeepSeekV3. - ✅ Learning Algorithms - SFT (Supervised Fine-Tuning), PEFT (LoRA, QLoRA), and QAT (Quantization-Aware Training). - ✅ Knowledge Distillation - Support for knowledge distillation with LLMs.**

Model Coverage and 🤗 Ecosystem compatibility - ✅ Transformers v5 🤗 - Built on latest transformers with device-mesh driven parallelism. - ✅ 🤗 HuggingFace Integration - Works with dense models (e.g., Qwen, Llama3, etc) and large MoEs (e.g., DSv3, DSv4). - ✅ VLM - Finetuning for VLMs (Qwen2.5/3/3.5/3.6 VL, Gemma-3/3n/4 VL, Mistral 3.5/4, LLaVA-OneVision-1.5, Kimi-VL, etc.). - ✅ Omnimodal - Finetuning for omnimodal MoE models (Nemotron-3-Nano-Omni, Qwen3-Omni). - ✅ Diffusion - Pretraining and LoRA finetuning for image/video diffusion models (Qwen-Image, FLUX, Wan2.1, Hunyuan). - ✅ dLLM - Discrete diffusion LM finetuning (LLaDA). - ✅ Retrieval - Bi-encoder and cross-encoder training with in-batch negative sampling. - ✅ Extended MoE support - GPT-OSS, Qwen3 / Qwen3.5 / Qwen3.6 MoE, Qwen-next, MiniMax-M2.x, GLM-4.7 / GLM-5 / GLM-5.1, DeepSeek V3.2 / V4 / V4-Flash, ERNIE 4.5, MiMo-V2-Flash, Ling 2.0, Hy3-preview.

Agentic Development and UX - ✅ Agent-friendly skills - Curated skills/ for common dev tasks (recipe runs, model onboarding, CI).

• 🔜 Muon optimizer - Muon optimizer support.
• 🔜 SonicMoE - Optimized MoE implementation for faster expert computation.
• 🔜 FP8 MoE - FP8 precision training and inference for MoE models.
• 🔜 Cudagraph with MoE - CUDA graph support for MoE layers to reduce kernel launch overhead.
• 🔜 VLM Knowledge Distillation - Extend KD to VLM and omnimodal models.