feat(parallelism): expert parallelism for the Nemotron-3 Nano MoE layers

[le1nux]

What

Adds expert parallelism (EP) for the mixture-of-experts layers of the Nemotron-3 Nano architecture, plus a few improvements that came out of running the full 30B-A3B model.

Stacked PR. Base branch is nano_integration (#459), not main. Please merge #459 first; this PR then contains exactly one commit on top of it.

Expert parallelism

Under EP the routed experts of each MoE layer are partitioned across an ep mesh dimension, so each rank stores and evaluates only num_experts / ep_degree of them. Every rank still routes its own tokens, so tokens travel to the rank owning their expert (dispatch all-to-all) and results travel back (combine all-to-all).

The alternative — what modalities does today — is to let FSDP2 all-gather the full expert stack on every rank. For the 16-layer Nemotron config that is 2.38 GiB of bf16 expert weights per MoE layer per forward, and activation checkpointing makes the forward run twice.

Design

• src/modalities/models/parallelism/expert_parallelism.py — ExpertParallelGroupedExperts (dispatch/combine wrapper, drop-in for GroupedExperts, so MoE needs no knowledge of EP) and shard_experts_over_ep_mesh (re-creates w1/w2 as Shard(0) DTensors on the meta-device model, so checkpoints still see the full global expert stack).
• The dispatch/combine structure follows TorchTitan (torchtitan/models/common/moe.py, torchtitan/distributed/expert_parallel.py, BSD 3-Clause), attributed in the module docstring. Two things are done deliberately differently:
  1. The post-all-to-all permutation is built entirely on device. The obvious implementation loops over (sender, local_expert) pairs and calls .item(), costing 2 * ep_degree * num_local_experts device-to-host syncs per MoE layer per forward — several thousand per step for the 52-layer model. _build_permute_indices constructs the index tensor with vectorized ops instead.
  2. Exactly one device-to-host sync per dispatch, for the split lists all_to_all_single needs on the host; both split vectors are copied in a single transfer.
• Autograd flows through both all-to-alls via all_to_all_single_autograd and through the index_select/index_copy permutation pair.

Device mesh

EP is carved out of data_parallel_shard_degree rather than multiplying into the world size, so the number of distinct data shards (and therefore the dataloader's dp_degree) is unchanged. With expert_parallel_degree > 1, dp_shard is materialized as dp_shard_mod_ep × ep and then re-exposed as a flattened dp_shard alias, so callers that only care about data parallelism need no changes. ep is the inner (fastest-varying) dimension so an all-to-all group spans consecutive global ranks, i.e. stays inside a node for ep_degree <= devices per node.

Because the flattened dp_shard is addressable but absent from mesh_dim_names, mesh lookups now resolve through device_mesh[name] (_resolve_sub_mesh) instead of scanning mesh_dim_names. Config validation rejects ep_degree > dp_shard_degree, non-divisible degrees, and (for now) any combination with TP/PP/CP.

FSDP2 interaction

Expert-parallel stacks must not be sharded on the full dp_shard dimension — ranks differing in their ep coordinate hold different experts, so all-gathering across them would mix unrelated weights. Each expert stack therefore gets its own FSDP unit on dp_shard_mod_ep, the part of the data-parallel dimension EP did not consume.

Gradient clipping

Routed expert gradients are sharded over (dp_shard_mod_ep, ep) while all other gradients are sharded over dp_shard. Both the norm reduction (torch.stack) and the in-place rescaling (aten._foreach_mul_) are batched, and neither has a sharding rule for operands from different meshes. Gradients are now grouped by mesh, reduced per group, and the group norms combined — (Σ_g norm_g^p)^(1/p), or max_g for the infinity norm. With a single mesh this is bit-identical to the previous computation.

Measured, 4× A100-SXM4-80GB (NVSwitch), torch 2.9.1+cu128

config_fineweb_nemotron_nano_ep_fsdp2.yaml vs. config_fineweb_nemotron_nano_fsdp2.yaml (byte-identical apart from the EP degree, the new component and the header), median over steps 4+ of a ~5 minute run:

expert_parallel_degree	samples/s	tokens/s/GPU	MFU	peak mem	speedup
1 (non-EP baseline)	5.40	5530	0.170	48.97 GiB	1.000×
2	5.90	6042	0.190	52.42 GiB	1.093×
4	7.10	7270	0.230	48.02 GiB	1.315×

Per-step ranges were [5.30, 5.60], [5.80, 6.00] and [7.00, 7.30] — non-overlapping.

Two things worth reading off this table. Degree 2 is the worst of both worlds on memory: the experts are still FSDP-sharded over the 2 remaining data-parallel ranks so the all-gather buffer is still allocated, and the all-to-all buffers are added on top. And the speedup is not the MoE math getting faster — benchmarked in isolation without FSDP, one EP MoE layer is 0.89× on the forward and 1.06× on forward+backward versus a replicated one, roughly break-even. The entire end-to-end win comes from no longer all-gathering expert weights.

Other improvements

• ChunkedCLMCrossEntropyLoss — moves the lm_head out of the model's forward and applies it chunk-by-chunk inside torch.utils.checkpoint, so the [batch, seq_len, vocab_size] logits (the single largest activation at a 131k vocabulary, plus the fp32 up-cast inside CE) are never fully materialized. Same goal as TorchTitan's ChunkedLossWrapper, but torch-native via checkpoint recomputation rather than a manual per-chunk backward. Enabled by GPT2LLM.lm_head / set_skip_lm_head and by separate_lm_head_fsdp_unit on FSDP2WrappedModelConfig (the head needs its own gather hook when it is called outside the model forward). Overlaps with feat: Add torchtitan-style chunked lm_head cross-entropy loss (ChunkedLMHeadCrossEntropyLoss) #458 — happy to drop this part if that PR lands first.
• Compilable losses — the pure-tensor CE core is now a free function (clm_cross_entropy_loss, clm_cross_entropy_loss_sum) so it can be handed to torch.compile without self in the graph, plus a LossFactory.get_compiled_loss / loss: compiled component mirroring ModelFactory.get_compiled_model. Only the numeric core is compiled, never the batch unpacking.
• GroupedExperts accepts DTensor weights (_local_weights), and its tokens_per_expert is now documented as local counts under EP. No-op without EP.
• Two full-depth 30B-A3B FineWeb configs, plain and torch.compile'd. Both headers document the measured parameter counts (31.563B total / 3.565B active) and the deviations from the lorem-ipsum config. The compiled variant's header states plainly that compiling this architecture is close to a no-op — the two block types holding 97% of the parameters fragment into 12 and 4 graphs — and that the uncompiled sibling is the recommended default.

New components

type	variant	config
model	expert_parallelized	ExpertParallelizedModelConfig
loss	chunked_clm_cross_entropy_loss	ChunkedCLMCrossEntropyLossConfig
loss	compiled	CompiledLossConfig

Tests

• tests/config/test_device_mesh_config.py — the EP degree constraints, including that EP does not consume world size, pinned down without a process group.
• tests/models/nemotron/test_expert_parallelism.py — the dispatch permutation against a loop reference, that it is a permutation, that it groups by expert, and (via monkeypatch) that it performs no device-to-host readback.
• tests/fsdp2_parallelization/test_expert_parallelism_fsdp2.py (2 GPUs) — three workers: expert weights end up as DTensors on (dp_shard_mod_ep, ep) holding only this rank's share while router and shared experts stay data-parallel; a full forward/backward/optimizer step including the cross-mesh gradient clipping; and an EP MoE layer matching a replicated one bit-for-bit.

Reviewer notes

• docs/components/nemotron.md still says "Expert parallelism — Not supported. The device mesh has no expert dimension." (line 197) and describes the full 30B config as needing EP for real throughput (line 11). Both are stale as of this PR; I can push the doc update on request, or fold it into the feat(model): add Nemotron-3 Nano hybrid Mamba-Transformer MoE architecture #459 doc.
• The .gitignore change adds data/experiments/* and soofi/*. The latter is local scratch and can be dropped if you'd rather not carry it upstream.