NeMo DFM (Diffusion Foundation Models) is a library under NeMo Framework, focusing on diffusion models for Video, Image, and Text generation. It unifies cutting-edge diffusion-based architectures and training techniques, prioritizing efficiency and performance from research prototyping to production deployment.

Dual-Path Architecture: DFM provides two complementary training paths to maximize flexibility:

• 🌉 Megatron Bridge Path: Built on NeMo Megatron Bridge which leverages Megatron Core for maximum scalability with n-D parallelism (TP, PP, CP, EP, VPP, DP)
• 🚀 AutoModel Path: Built on NeMo AutoModel for PyTorch DTensor-native SPMD training, for easy experimentation and also Day-0 support on 🤗 Hugging Face models.

Choose the path that best fits your workflow—or use both for different stages of development!

# Initialize all submodules (Megatron-Bridge, Automodel, and nested Megatron-LM)
git submodule update --init --recursive

# Build the container
docker build -f docker/Dockerfile.ci -t dfm:dev .

docker run --rm -it --gpus all \
  --entrypoint bash \
  -v $(pwd):/opt/DFM -it dfm:dev

You can find all predefined recipes under recipes directory.

Note: You will have to use uv to run the recipes. Please use --group as megatron-bridge.

uv run --group megatron-bridge python -m torch.distributed.run --nproc-per-node $num_gpus \
  examples/megatron/recipes/wan/pretrain_wan.py \
  --config-file examples/megatron/recipes/wan/conf/wan_1_3B.yaml \
  --training-mode pretrain \
  --mock

Train with PyTorch-native DTensor parallelism and direct 🤗 HF integration:

You can find pre-configured recipes under automodel/finetune and automodel/pretrain directories.

Note: AutoModel examples live under dfm/examples/automodel. Use uv with --group automodel. Configs are YAML-driven; pass -c <path> to override the default.

The fine-tune recipe sets up WAN 2.1 Text-to-Video training with Flow Matching using FSDP2 Hybrid Sharding. It parallelizes heavy transformer blocks while keeping lightweight modules (e.g., VAE) unsharded for efficiency. Adjust batch sizes, LR, and parallel sizes in dfm/examples/automodel/finetune/wan2_1_t2v_flow.yaml. The generation script demonstrates distributed inference with AutoModel DTensor managers, producing an MP4 on rank 0. You can tweak frame size, frames, steps, and CFG in flags.

# Fine-tune WAN 2.1 T2V with FSDP2 (single node, 8 GPUs)
uv run --group automodel torchrun --nproc-per-node=8 \
  dfm/examples/automodel/finetune/finetune.py \
  -c dfm/examples/automodel/finetune/wan2_1_t2v_flow.yaml

# Generate videos with FSDP2 (distributed inference)
uv run --group automodel torchrun --nproc-per-node=8 \
  dfm/examples/automodel/generate/wan_generate.py

Megatron Bridge delivers maximum throughput and scalability with near-linear performance to thousands of nodes. AutoModel provides an easy on-ramp for experimentation and research with PyTorch-native SPMD training.

• 🎥 Multi-Modal Diffusion: Support for video, image, and text generation
• 🔬 Advanced Samplers: EDM, Flow Matching, and custom diffusion schedules
• 🎭 Flexible Architectures: DiT (Diffusion Transformers), WAN (World Action Networks)
• 📊 Efficient Data Loading: Data pipelines with sequence packing
• 💾 Distributed Checkpointing: SafeTensors-based sharded checkpoints
• 🌟 Memory Optimization: Gradient checkpointing, mixed precision, efficient attention
• 🤗 HuggingFace Integration: Seamless integration with the HF ecosystem

DFM provides out-of-the-box support for state-of-the-art diffusion architectures:

For detailed performance benchmarks including throughput metrics across different GPU systems and model configurations, see the (Performance Summary)[https://github.com/NVIDIA-NeMo/DFM/blob/main/docs/performance-summary.md] in our documentation.

DFM/
├── dfm/
│   └── src/
│       ├── megatron/              # Megatron Bridge path
│       │   ├── base/              # Base utilities for Megatron
│       │   ├── data/              # Data loaders and task encoders
│       │   │   ├── common/        # Shared data utilities
│       │   │   ├── <model_name>/  # model-specific data handling
│       │   ├── model/             # Model implementations
│       │   │   ├── common/        # Shared model components
│       │   │   ├── <model_name>/  # model-specific implementations
│       │   └── recipes/           # Training recipes
│       │       ├── <model_name>/  # model-specific training configs
│       ├── automodel              # AutoModel path (DTensor-native)
│       │   ├── _diffusers/        # Diffusion pipeline integrations
│       │   ├── datasets/          # Dataset implementations
│       │   ├── distributed/       # Parallelization strategies
│       │   ├── flow_matching/     # Flow matching implementations
│       │   ├── recipes/           # Training scripts
│       │   └── utils/             # Utilities and validation
│       └── common/                # Shared across both paths
│           ├── data/              # Common data utilities
│           └── utils/             # Batch ops, video utils, etc.
├── examples/                      # Example scripts and configs

We welcome contributions! Please see our Contributing Guide for details on:

• Setting up your development environment
• Code style and testing guidelines
• Submitting pull requests
• Reporting issues

For questions or discussions, please open an issue on GitHub.

NeMo DFM builds upon the excellent work of:

• Megatron-core - Advanced model parallelism
• Megatron Bridge - HuggingFace ↔ Megatron bridge
• NeMo AutoModel - PyTorch-native SPMD training
• PyTorch Distributed - Foundation for distributed training
• Diffusers - Diffusion model implementations