[26.08] AutoModel Roadmap

[26.08] AutoModel Roadmap

This post tracks the planned work items for NeMo AutoModel in the 26.08 cycle. Plans may shift as we validate performance, hardware coverage, upstream compatibility, and community priorities. If you have feature requests, use cases, or model families you want us to prioritize, please comment below.

---

Developer and User Experience

• Unified Engine and Recipe APIs
  Introduce a common engine API and consolidate LLM and VLM recipe paths so training, fine-tuning, evaluation, and downstream integrations use fewer modality-specific branches.

• Configuration and Builder Cleanup
  Continue the typed-config and .build() refactor, simplify recipe construction, and make configuration behavior easier to validate and extend.

• Installation and Dependency Compatibility
  Improve CUDA installation paths, FlashAttention packaging, container compatibility, and support for a practical range of Transformers and CUDA toolkit versions.

• Separate Model Implementations from Framework Infrastructure
  Continue splitting the model zoo from reusable infrastructure so model onboarding and core framework changes can evolve independently.

• Consolidate Benchmark and Fine-Tuning Recipes
  Reduce duplicated benchmark and fine-tuning YAML trees and keep performance overrides close to the canonical model recipe.

---

Model Support and Ecosystem Compatibility

• Continuous Day-0 Model Support
  Continue rapid support for new Hugging Face model releases across LLM, VLM, omnimodal, diffusion, discrete diffusion, and retrieval workflows.

• Gemma 4 Coverage and Stabilization
  Expand Gemma 4 coverage across recipes, long-context and context-parallel paths, DGX Spark validation, drafter cleanup, activation checkpointing, and NaN or convergence investigations.

• BAGEL Scale, Performance, and Convergence
  Continue BAGEL multi-node validation, topology-independent initialization, Transformer Engine and fused-projection optimization, and convergence characterization.

• Retrieval Model Coverage
  Expand vision-language retrieval datasets and model support, improve image/data construction performance, and onboard additional retrieval architectures.

• Discrete Diffusion Language Models
  Broaden dLLM recipe and nightly coverage, including DiffusionGemma compatibility as the upstream Transformers implementation evolves.

• Model Loading and Weight-Tying Correctness
  Make tied-embedding support explicit per model family and strengthen checkpoint-loading behavior when model and checkpoint tie semantics differ.

---

Core Infrastructure and Integrations

• Downstream Training Interfaces
  Improve integration hooks for RL, generation, and external training libraries so they can reuse AutoModel models, parallelism, checkpoints, and recipes without depending on private internals.

• Unified Parallelism Architecture
  Reduce separate MoE and FSDP code paths, consolidate device-mesh handling, and make model parallelism easier to compose and extend.

• Decorator-Based Parallelism
  Explore a simpler declarative interface for attaching parallelism policies to model components.

• Flexible Algorithm Hooks
  Add stable extension points for intermediate-layer distillation, per-module precision policy, and dynamic attention masks without requiring users to fork model forward methods.

---

Performance and Parallelism

• DeepEP v2 and MoE Backend Improvements
  Upgrade DeepEP, validate performance gains, improve router kernels, and continue closing MoE correctness and performance gaps.

• Attention Backend Expansion
  Complete MagiAttention integration and packaging, improve FlashAttention 3/4 installation, and evaluate additional optimized attention backends.

• Context Parallelism, THD Packing, and Long Context
  Expand packing plus context parallelism, DeepSeek V4 THD packed paths, Gemma 4 long-context coverage, and end-to-end CP validation on realistic long-sequence fine-tuning workloads.

• CUDA Graph and Compile Paths
  Extend partial CUDA Graph coverage and improve graph-friendly execution for dynamic and sparse workloads.

• Tensor-Parallel Correctness for Custom Modules
  Ensure replicated custom-module weights are synchronized correctly across TP ranks and fix correctness gaps in fused RoPE and MoE context-parallel combinations.

• Training Precision and Memory Efficiency
  Continue work on FP32 master weights, lower-memory FP8 loading and sharding, and memory-efficient checkpoint or state-dict adaptation.

---

Multimodal, Retrieval, and Knowledge Distillation

• Multimodal Algorithm Flexibility
  Support configurable intermediate-output taps, per-module FP8 opt-in or opt-out, and dynamic attention-mask plumbing for advanced multimodal training algorithms.

• BAGEL Multi-Node Validation
  Validate BAGEL parity, convergence, and performance from single-node tuning through larger multi-node configurations.

• Retriever Throughput and Scale
  Improve Nemotron and Ministral-based vision-language retrieval models, data construction, image processing, activation checkpointing, and distributed tuning.

• Knowledge Distillation Across Tokenizers
  Extend knowledge distillation to teacher and student models that use different tokenizers.

---

Diffusion

• Diffusion Recipe Alignment
  Refactor diffusion recipes toward the shared AutoModel recipe and engine structure while preserving diffusion-specific model and data requirements.

• DiffusionGemma Compatibility and Parity
  Keep DiffusionGemma training and parity coverage aligned with current Transformers releases and multimodal upstream changes.

---

Checkpointing, State, and Robustness

• Async Checkpoint Save/Resume Coverage
  Add stronger CI coverage for asynchronous checkpoint save and resume behavior, including planner and cache regressions.

• Cloud and Object-Storage Validation
  Expand S3 and remote-storage testing for checkpoint and dataset workflows.

• Lower-Memory Loading and Conversion
  Reduce full-model materialization during FP8 loading, state-dict adaptation, checkpoint consolidation, and adapter conversion.

---

Release Quality, CI, and Validation

• Loss and Performance Regression Gates
  Add automated convergence and step-time regression checks so nightly CI catches numerical or performance drift earlier.

• Real-Workload Convergence Validation
  Validate context parallelism and long-context fine-tuning with realistic datasets, then carry representative workloads into recurring convergence coverage.

• Downstream Integration Tests
  Add recurring compatibility, loss-matching, and performance tests for downstream training integrations against current AutoModel main.

• Actionable CI Failure Reporting
  Improve automated regression issues with stronger ownership, pipeline provenance, checkpoint-cache support, and searchable configuration context.

---

Tracker-Backed 26.08 Items

The following public issues and pull requests were used to align this roadmap with the 26.08 plan.

Area	Planned item	Public reference
API and recipes	Unified Engine class and LLM/VLM recipe consolidation	#2269
Architecture	Separate the model zoo from framework infrastructure	#2163
MoE performance	Upgrade to DeepEP v2	#2021
Attention	MagiAttention backend and packaging	#1013, #2384
Long context	Sequence packing with context parallelism	#1473
Long context	End-to-end CP validation on an SFT workload	#2890
Gemma 4	TE context parallelism, recipe stability, and drafter support	#2115, #2199, #2240
BAGEL	Scale, fused kernels, and convergence validation	#2895, #2939
Retrieval	Vision-language data and throughput improvements	#2451, #2452
dLLM	DiffusionGemma training and nightly coverage	#2506, #2783, #2873
Model correctness	Explicit tied-embedding support and checkpoint guards	#2511, #2896
Training precision	FP32 master weights for full-parameter training	#2379
Checkpointing	Async DCP correctness and remote object storage	#2070, #1709, #2234

---

We Want Your Input

Have a feature request or use case that is not covered above? Please comment and include:

1. What you would like to see.
2. Why it matters for your workflow.
3. Any context that helps us prioritize, such as model family, scale, hardware, precision, or deployment target.

We will prioritize based on community feedback, engineering feasibility, upstream dependencies, and release validation risk.

This roadmap names many independent API, model, parallelism, checkpointing, and CI goals but no files or tests; start by selecting one tracker-backed item and reading its linked issue or pull request. Define a concrete scope and acceptance checks with maintainers, since the umbrella issue is complete only when the selected item's stated validation is finished.