fix(moe): map FSDP expert shards with torch.chunk layout - #4058

Suggestion: Approve

Reviewed 89198045259f0908ce8d68cc1d690236bc9aeaaf against 112da8aa4c1128fec9186f81efa79c858db61a5b, including the shared helper, both Kimi overrides, and their native checkpoint load/save consumers. I found no introduced blocking issue.

The shared helper is appropriate here: it labels the shards FSDP already holds and removes the same incorrect range calculation from three consumers. It does not introduce empty shards or change expert placement. Supported EP configurations require divisibility, so their ranges remain unchanged; the separate balanced EP ownership helper would not describe these uneven FSDP shards correctly.

Validation:

• The three changed native unit-test files: 55 passed, 1 skipped (the CUDA-only unit test was skipped in the CPU run).
• Four real CPU/Gloo ranks with actual DTensors: 112/112 rank-case checks passed, covering shared splitting, both Kimi adapters, exact expert IDs/values/HF keys, empty tails, and replicated/inner-axis/HSDP/divisible-EP controls. The actual baseline produced 25 expected regression failures; all 72 control checks passed on both versions.
• Four H100s, native AutoModel FSDP2 loading with EP/TP/CP=1: a tiny DeepSeek-V3 fixture with 5, 10, and 8 experts exercised local counts [2,2,1,0], [3,3,3,1], and [2,2,2,2]. After complete production setup, every parameter/buffer matched its source value, shape, and dtype exactly on every rank; fixed-input logits had maximum error 0. Native forward/backward, a finite nonzero-gradient optimizer step, and save/consolidation completed. All 15/30/24 expert HF tensors, respectively, matched independently derived post-training values exactly. The actual baseline failed during native loading for both 5 and 10 experts with the expected local-count errors.

One existing limitation remains visible: the full export check failed solely on e_score_correction_bias (maximum difference 0.002). With FSDP optimizations disabled for this topology, that replicated buffer diverges across training ranks and the checkpoint contains rank 1's value rather than the observer's rank 0 value. The unmodified baseline's 8-expert control reproduces the same failure, and every post-training parameter/buffer is identical between baseline and candidate within corresponding ranks. This is not an introduced finding or a full-checkpoint export pass.

Coverage is one-layer FP32, SDPA/torch experts, 16-token native training; Kimi coverage is distributed adapter components, not full Kimi checkpoint loading. The custom observer calls the native recipe and checkpointer and only inspects state; it does not substitute a loader or repair loaded values. Runtime: PyTorch 2.14.0a0+4fdf77b940.nv26.08, Transformers 5.18.0.

Suggestion: Approve

Bot suggestion — second review, on 9588a15839027a90917c04e8620ba5e173f06597 against 3ea1de8d0aa902c3329b585d194c2e5274a85f3f. I found no introduced blocking issue.

The six changed files are unchanged from the first reviewed head. I rechecked the merge with current main, particularly its new upstream safetensors reader. The shared helper remains the right scope: it labels existing FSDP torch.chunk storage at all three consumers. The balanced EP ownership helper cannot describe those uneven shards; supported divisible EP layouts retain their previous ranges.

Fresh validation on this head:

• Three native unit-test files: 55 passed, 1 CUDA-only test skipped in the CPU invocation.
• Four real CPU/Gloo ranks, actual DTensors: 112/112 checks passed, including 5/10-expert uneven shards, empty tails, both Kimi adapters, exact expert IDs/values/HF keys, and replicated/inner-axis/HSDP/divisible-EP controls. Current main reproduces 25 expected failures; all 72 unaffected control checks pass on both versions.
• Two H100s, native AutoModel FSDP2 (dp_shard=[0,1], EP/TP/CP/PP=1), loading the retained native five-expert checkpoint: the actual reader is _CachedMetadataHFReader; all 20 parameters/buffers match source values, shapes and dtypes exactly on each rank. Local expert counts are [3,2]; fixed-input logits have maximum error 0. Native forward/backward, two accumulated microbatches, one AdamW update and save/consolidation complete. All 15 expert HF tensors exactly match independently derived post-training values.

The full export gate still fails solely on the pre-existing replicated e_score_correction_bias: the saved value matches rank 1 and differs from rank 0 by up to 0.002. Current main reproduces the same per-rank bias and export failure. This is not a full-checkpoint export pass or a new PR finding.

The default candidate/main post-training comparison is not bitwise (maximum parameter difference 2.48e-7, equal loss). A separate strict-determinism pair, with identical recorded input hashes, produces bit-identical post-training parameters and buffers on both ranks. The specific source of default runtime variation was not isolated; the original failed exact comparison remains recorded.

Coverage is a tiny one-layer FP32 model, SDPA/torch_mm, 16-token training. Two ranks cannot distinguish the old/new partition formulas: the four-rank CPU tests provide that regression check, and the previous four-H100 native evidence remains separate. Kimi validation is adapter-component coverage, not full Kimi checkpoint loading. The observer uses the production recipe/checkpointer and only records state, reader selection and inputs; it does not replace loading or repair weights. Runtime: PyTorch 2.14.0a0+4fdf77b940.nv26.08, Transformers 5.18.0.