[Feature] Add MoE expert tensor parallelism

[jayhenry]

Summary

This PR ports MoE expert tensor parallelism onto the latest upstream/main while keeping checkpoint formats topology-independent. The work is split into four ordered, independently testable commits.

flowchart LR
    A["DTensor-aware HF planning"] --> B["InterleavedShard HF I/O"]
    B --> C["InterleavedShard DCP"]
    C --> D["Expert TP"]

Loading

Changes

• Make HF load/save planning aware of DTensor placements.
• Add InterleavedShard HF checkpoint round trips.
• Add InterleavedShard DCP save/load and reshard planning.
• Add (FSDP, EP, ETP) meshes and Expert TP support for Naive, All2All, and DeepEP dispatchers, including async/Domino and tile-wise FP8 paths.
• Keep HF and DCP checkpoint representations independent of the runtime EP/ETP topology.

Verification

• Pre-commit hooks, Ruff, formatting, mypy, and git diff --check passed.
• HF planning/config tests: 42 passed.
• Dispatcher/model-path tests: 9 passed.
• TrainEngine focused tests: 7 passed.
• DCP model/optimizer round trips: 2 passed.
• InterleavedShard distributed tests passed on both 4 and 8 GPUs.
• Qwen3-30B-A3B 8-GPU end-to-end runs completed 10/10 steps for EP1/ETP1 and EP2/ETP2 DeepEP, in both BF16 and FP8.
• Against this branch's direct upstream/main parent, EP1 throughput changed by +0.12% (BF16) and -0.10% (FP8), with matching loss and peak memory.
• With the historical FA3/non-deterministic runtime flags, EP1 BF16 and FP8 throughput reproduced the previous baseline within 1%, while reducing peak memory by about 4 GB.

Result

The final tree preserves the existing Expert TP behavior and numerical results while integrating it with the current upstream/main checkpoint and training architecture.

[jayhenry]

@claude review

[jayhenry]

@claude review