refactor: Remove separate moe_mesh references

[edjson]

What does this PR do ?

Removes separate moe_mesh parameter form NeMoAutoModel so users only need to pass in device_mesh to the NeMoAutoModel class.

Changelog

• Add specific line by line info of high level changes in this PR.

Before your PR is "Ready for review"

Pre checks:

• Make sure you read and followed Contributor guidelines
• Did you write any new necessary tests?
• Did you add or update any necessary documentation?

If you haven't finished some of the above items you can still open "Draft" PR.

Additional Information

• Follows up on fix: resolve PT 2.11 DeviceMesh deprecation warnings and unify EP mesh #1684
• closes Remove separate moe_mesh references after https://github.com/NVIDIA-NeMo/Automodel/pull/1684 #1699

[copy-pr-bot]

This pull request requires additional validation before any workflows can run on NVIDIA's runners.

Pull request vetters can view their responsibilities here.

Contributors can view more details about this message here.

[edjson]

Hello @adil-a, I am just opening up this draft to talk more about the design changes. My current plans are to:

1. Remove the moe_mesh parameter from auto_model.py on line 468
2. Delete device_mesh.py
3. Update all the device_mesh.py imports to mesh_utils.py

Before I begin, should I initialize moe_mesh in MeshContext.__post_init__, or make moe_mesh from calling it in from_pretrained?

[akoumpa]

/ok to test b3b2ca8

[akoumpa]

/ok to test 136f4cb

[adil-a]

Hi @edjson , thanks for taking a stab at this. I'd like to actually minimize the blast radius from this PR. The main goal is just to ensure that moe_mesh is no longer being passed into NeMoAutoModelForCausalLM (or other NeMoAutoModel model init classes). As it stands, I got Claude to take a first look and it mentioned the following issues:

Validation complete. Most claims hold up — three corrections to flag, one previous claim retracted.

  Confirmed (high-priority, correctness)

  ✅ A1 — _derive_moe_mesh is broken on the recipe path. mesh_utils.py:241-252 builds the moe mesh and registers only DP/DP_SHARD_CP/DP_CP at _flatten_mapping (lines 235-237); never writes _flatten_mapping["ep"].

  ✅ A2 — All 5 recipes drop moe_mesh= from scale_grads_and_clip_grad_norm: recipes/llm/train_ft.py:1483-1498, recipes/llm/kd.py:634-649, recipes/dllm/train_ft.py:335-350, recipes/retrieval/train_bi_encoder.py:366-381, recipes/vlm/finetune.py:1108-1123. EP grad scaling never fires.

  ✅ A3 — 4 recipe-level Checkpointer calls dropped moe_mesh=self.moe_mesh: recipes/llm/train_ft.py:1018-1023, recipes/llm/train_seq_cls.py:101-106, recipes/retrieval/train_bi_encoder.py:190-195, recipes/vlm/finetune.py:783-788. Checkpointer.__init__ still uses moe_mesh for state-dict adapters → silent MoE checkpoint corruption.

  Confirmed (API breaks & stale references)

  ✅ B1, B2 — MeshContext.from_meshes (mesh.py:198-207) and NeMoAutoModel.from_{pretrained,config} (auto_model.py:487-508, 629-650, plus overloads at 870, 1029) all dropped moe_mesh with no shim.

  ✅ C1 — examples/convergence/tulu3/model-verification/extract_nemo_activations.py:107 still passes moe_mesh=dist_setup.moe_mesh to build_model — but see A4 retraction below; the impact is just an unused kwarg, not a TypeError.

  ✅ C2 — recipes/diffusion/train.py:496-502 still passes Checkpointer(..., moe_mesh=None) (cosmetic).

  ✅ C3 — skills/distributed-training/SKILL.md references stale APIs (positional moe_mesh, deleted device_mesh.py, renamed helper).

  ✅ C5 — Three recovered_*.patch files at repo root.

  ✅ D2 — tests/unit_tests/_transformers/test_infrastructure.py:180-181 mock attribute is now dead (PR reads mesh.moe_mesh, not model_wrapper.moe_mesh).

  ✅ E2 — All 5 recipes already obtain a MeshContext from setup_distributed. The "pass mesh: MeshContext" recommendation stands.

  Retractions / corrections

  ❌ A4 (RETRACT) — Coverage agent's claim that recipes/llm/train_ft.py:163-211 build_model retains a moe_mesh=None parameter is wrong. The PR cleanly removed moe_mesh from build_model's signature (params end at sdpa_method) and from its kwargs dict (lines 202-209). So the claimed "TypeError on the MoE path through build_model" doesn't exist. Downgrades severity of C1 to "stale extra kwarg in an example, will silently TypeError there if the example is run, but not in the recipe path."

  ⚠️  C4 (PARTIAL) — The "unused, kept for API compatibility" docstring at components/training/utils.py:202-203 does exist, but it documents clip_grad_norm (line 175) — whose signature doesn't actually have moe_mesh/ep_axis_name either. The sibling scale_grads_and_clip_grad_norm (line 290) is the function that genuinely takes and uses moe_mesh. So the docstring is wrong, just not in the way the agent implied. Pre-existing.

  ⚠️  D1 (PARTIAL) — The cited test lines in test_auto_diffusion_pipeline.py are correct, but the test will pass, not fail. _diffusers/auto_diffusion_pipeline.py:280 still calls FSDP2Manager(..., moe_mesh=moe_mesh) and FSDP2Manager.__init__ (fsdp2.py:81-89) still accepts it. The diffusion pipeline path wasn't refactored. So this is a coverage gap, not a regression.

  ⚠️  E1 (PARTIAL) — instantiate_infrastructure (_transformers/infrastructure.py:249-261) treats only device_mesh and ep_size as deprecated; moe_mesh was fully removed (not deprecated). Doesn't change the recommendation, just tightens the wording.

  Net assessment

  The two critical correctness claims (broken EP derivation, broken EP grad scaling) are real. The new finding (4 broken Checkpointer calls = silent MoE checkpoint corruption) is real and worse than I initially thought. The architectural recommendation (pass MeshContext) holds. The only thing to walk back: the example script extract_nemo_activations.py:107 is stale-but-not-fatal because build_model quietly accepts unknown kwargs only if they aren't threaded — let me note that doesn't actually save
   it; passing an unknown kwarg to a function whose signature doesn't accept it is still a TypeError. So C1 is a will-break, just via a different mechanism than the agent claimed.

As you mentioned, this is just an initial draft and likely a lot of these issues would go away once we have a more complete implementation. However, just to minimize any regressions across the codebase, my proposal is that we keep the internals the same as they are with minimal source code changes to the internal code. Let's keep the goal of this PR to just remove moe_mesh from the user-facing API as mentioned above. Let me know if you have any questions and thank you once again for taking a look :)

[edjson]

Hello @adil-a,

Thank you for the feed back. I understand, and will begin reworking the PR. Thank you.

[edjson]

Opened a new pull request #2123 in approaching a more focused approach as suggested in review.