Experimenting MoE new sharding

[acisseJZhong]

Stack from ghstack (oldest at bottom):

• fix async tp #4045
• [do not review] inference moe padding #4067
• inference moe expert sp padding #4080
• -> Experimenting MoE new sharding #3996

Summary

This updates the MoE EP/SP sharding contract so token dispatchers return the local sequence shard directly instead of constructing a full global-length sparse contribution.

The main change is that routed expert outputs are represented as the MoE local output layout:

• EP on: sequence-sharded layout
• EP off, experts using pure TP: Partial layout

Shared experts are updated to produce the same layout as routed experts before the routed + shared add. The outer MoE boundary remains responsible for redistributing the final MoE output to sp_layout.

Changes

• Update standard AllToAll combine to scatter into a local-shard output buffer instead of a full sp_size-expanded buffer.
• Remove the SP global-index remapping path from EP token dispatchers.
• Remove local_seq_len_after_padding from the dispatcher combine interface since local-shard combine no longer needs per-rank-to-global SP index mapping.
• Keep num_local_tokens_after_padding for logical sequence padding, so AllToAll can size the local padded output without materializing pad rows.
• Make DeepEP combine stay async for all SP sizes; MoE.forward() syncs before the routed output is read.
• Update MoE sharding annotations so routed and shared expert outputs match before addition:
  • under EP: sequence-sharded
  • without EP: Partial
• Move shared expert TP redistributions onto the linear layers so w1/w3 expose input all-gather and w2 exposes Partial-to-output-layout reduction.
• Update the MoE sharding diagram.

Test plan

loss parity on FSDP4, TP=2, SP on, EP=2, ETP=1: AssertionError: 7.949278831481934 != 7.9492878913879395 : Loss mismatch at step 1: baseline=7.949278831481934, test=7.9492878913879395

NCCL_NVLS_ENABLE=0 python scripts/loss_compare.py . main --baseline-module='deepseek_v3' --baseline-config='deepseek_v3_debugmodel' --baseline-options="--parallelism.pipeline_parallel_degree 1 --parallelism.data_parallel_shard_degree 4 --parallelism.tensor_parallel_degree 2 --parallelism.expert_parallel_degree 2" --baseline-ngpus=8 --test-module='deepseek_v3' --test-config='deepseek_v3_debugmodel' --test-options="--parallelism.pipeline_parallel_degree 1 --parallelism.data_parallel_shard_degree 4 --parallelism.tensor_parallel_degree 2 --parallelism.expert_parallel_degree 2" --test-ngpus=8 --assert-equal --no-seed-checkpoint

loss parity on FSDP4, TP=2, SP off, EP=2, ETP=1: [LOSS_COMPARE] Average difference: 0.000000

NCCL_NVLS_ENABLE=0 python scripts/loss_compare.py . main --baseline-module='deepseek_v3' --baseline-config='deepseek_v3_debugmodel' --baseline-options="--parallelism.spmd_backend spmd_types --parallelism.pipeline_parallel_degree 1 --parallelism.data_parallel_shard_degree 4 --parallelism.tensor_parallel_degree 2 --parallelism.expert_parallel_degree 2 --parallelism.no-enable-sequence-parallel" --baseline-ngpus=8 --test-module='deepseek_v3' --test-config='deepseek_v3_debugmodel' --test-options="--parallelism.spmd_backend spmd_types --parallelism.pipeline_parallel_degree 1 --parallelism.data_parallel_shard_degree 4 --parallelism.tensor_parallel_degree 2 --parallelism.expert_parallel_degree 2 --parallelism.no-enable-sequence-parallel" --test-ngpus=8 --assert-equal --no-seed-checkpoint

loss parity on FSDP4, TP=2, SP on, EP=1, ETP=2: AssertionError: 6.0643310546875 != 6.064278602600098 : Loss mismatch at step 2: baseline=6.0643310546875, test=6.064278602600098

NCCL_NVLS_ENABLE=0 python scripts/loss_compare.py . main --baseline-module='deepseek_v3' --baseline-config='deepseek_v3_debugmodel' --baseline-options="--parallelism.pipeline_parallel_degree 1 --parallelism.data_parallel_shard_degree 4 --parallelism.tensor_parallel_degree 2 --parallelism.expert_parallel_degree 1" --baseline-ngpus=8 --test-module='deepseek_v3' --test-config='deepseek_v3_debugmodel' --test-options="--parallelism.pipeline_parallel_degree 1 --parallelism.data_parallel_shard_degree 4 --parallelism.tensor_parallel_degree 2 --parallelism.expert_parallel_degree 1" --test-ngpus=8 --assert-equal --no-seed-checkpoint

loss parity on FSDP4, TP=2, SP off, EP=1, ETP=2: [LOSS_COMPARE] Average difference: 0.000000

NCCL_NVLS_ENABLE=0 python scripts/loss_compare.py . main --baseline-module='deepseek_v3' --baseline-config='deepseek_v3_debugmodel' --baseline-options="--parallelism.pipeline_parallel_degree 1 --parallelism.data_parallel_shard_degree 4 --parallelism.tensor_parallel_degree 2 --parallelism.expert_parallel_degree 1 --parallelism.no-enable-sequence-parallel" --baseline-ngpus=8 --test-module='deepseek_v3' --test-config='deepseek_v3_debugmodel' --test-options="--parallelism.pipeline_parallel_degree 1 --parallelism.data_parallel_shard_degree 4 --parallelism.tensor_parallel_degree 2 --parallelism.expert_parallel_degree 1 --parallelism.no-enable-sequence-parallel" --test-ngpus=8 --assert-equal --no-seed-checkpoint

passed spmd type checking:

NCCL_NVLS_ENABLE=0 NGPU=8 MODULE=deepseek_v3 CONFIG=deepseek_v3_debugmodel ./run_train.sh \
  --parallelism.pipeline_parallel_degree 1 \
  --parallelism.data_parallel_shard_degree 4 \
  --parallelism.tensor_parallel_degree 2 \
  --parallelism.expert_parallel_degree 2 \
  --parallelism.spmd_backend spmd_types \
  --debug.spmd_typechecking \
  --training.steps 10 \
  activation-checkpoint:none

Also used a temporary script based on TestBitwiseParityMoEEP to run main
and the current commit with identical checkpoint weights and uneven
prompts, then compared all prefill/decode logprobs and generated
token IDs bit-for-bit.

[tianyu-l]

it'd be good to check the numerics doesn't change (for inference) when input sequence length is not always even

[tianyu-l]

add a TODO to deprecate this when DTensor backend is deprecated -- for spmd_types, we can directly access this info from global context