NVIDIA-NeMo / NVIDIA-NeMo/Automodel

Centralize MoE runtime backend capability truth for checkpoint adapters

Open
#3,689 0 comments 0 reactions 1 assignee View on GitHub

@yuhezhang-ai is already working on this.

Since Aug 26, 2026.

Dominant language
Python
Stars
960
Forks
316
Avg merge
3d 20h
Merged PRs (30d)
143

Description

Summary

Centralize the runtime facts used by MoE construction and checkpoint destination selection so the two paths cannot drift.

#3610 mirrors the current MoE constructor locally to avoid treating temporary TE storage as final model memory. The immediate correctness fix is intentionally local because importing from components.distributed would violate the checkpoint component boundary, while importing through components.checkpoint would introduce a MoE/checkpoint cycle.

Scope

  • Move the canonical safe world-size helper to nemo_automodel/shared/.
  • Define the maintained EP-dispatcher set once inside components/moe and use it from MoE construction, checkpoint conversion, and maintained model-specific MoE constructors such as Kimi K3 where applicable.
  • Make BackendConfig reject or deliberately normalize unsupported combinations such as experts="torch" with an EP dispatcher, rather than accepting the configuration and constructing TE experts implicitly.
  • Normalize backend access in the MoE state-dict mixin once its typed contract guarantees self.backend; remove the remaining one-off defensive getattr(backend, "dispatcher", None) next to direct backend dereferences.
  • Audit composite adapters that delegate checkpoint conversion to an inner language adapter. Delegate supports_low_memory_dcp_load only when the outer namespace and temporary-storage behavior preserve the inner capability, and map view_loaded_native_keys back to real outer model FQNs. Start with minimax_m3_vl, step3p7, and kimi_k25_vl; they currently fail safe by leaving the capability disabled. Also cover kimivl, whose shipped Kimi2VL configuration delegates conversion to DeepSeekV3 but does not currently re-prefix the inner view-loaded language keys.
  • Keep component import boundaries acyclic: no checkpoint-to-distributed dependency and no MoE-to-checkpoint dependency.

Acceptance criteria

  • Constructor and checkpoint-capability tests cover deepep, hybridep, and uccl_ep at world size one and greater than one.
  • Every accepted backend/dispatcher combination has an explicit runtime expert implementation.
  • Checkpoint capability tests derive the same model-backed-versus-temporary result as construction.
  • Each audited composite adapter either delegates the inner capability with focused namespace/storage tests or explicitly retains the safe disabled fallback; adapters that delegate conversion also report any view-loaded keys under real outer model FQNs.
  • Import-linting and existing model-construction tests pass.

Parent: #3576

Follow-up to #3610.

Contributor guide

Open the contributing guide

First steps

  1. Read the whole issue, then the project's contributing guide.
  2. Comment on the issue to say you are picking it up — it saves two people doing the same work.
  3. Fork the repository and make your change on a branch.
  4. Open a pull request that references the issue number.

Assessment

This issue has not been assessed yet.

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.