Support 3D+ MoE local_map mesh variants

[AlbedoWang]

Summary

This PR extends the DeepSeekV3 MoE local_map implementation so MoE can be represented on 3D+ AutoParallel meshes.

The change makes the MoE communication and compute boundary explicit at the local_map HOP. AutoParallel continues to reason about the outer graph through HOP input/output placements, while the MoE body owns the variant-specific dispatch, expert computation, and collectives.

Supported MoE role axes:

• ep_axis_name: expert parallel axis for expert ownership and token dispatch/combine.
• token_axis_names: token batch sharding axes, defaulting to all mesh axes except ETP.
• expert_tensor_parallel_axis_name: optional ETP axis for expert hidden-dim tensor parallelism.
• expert_fsdp_axis_name: optional expert-weight sharding axis exposed at the HOP boundary.
• expert_fsdp_shard_dim: optional explicit EFSDP shard dimension.

This lets the DSv3 example express MoE layouts such as:

• dp, ep
• ep, etp
• efsdp, ep
• dp_replicate, ep, etp
• dp_replicate, efsdp, ep
• dp_replicate, efsdp, ep, etp
• larger dense+MoE meshes where unrelated dense axes can still exist on the same mesh

Motivation

TorchTitan supports multiple sparse MoE mesh variants by assigning logical roles to mesh dimensions, for example DP replicate, EFSDP, EP, and ETP. AutoParallel needs the same kind of expressiveness, but its integration point is different.

In AutoParallel, local_map is the boundary that exposes MoE communication and computation to tracing as a HOP. DTensor placement reasoning happens at the HOP boundary, not by inferring arbitrary placement transitions inside the local body.

Because of that, this PR makes each MoE variant explicit through:

• boundary placement annotations for tokens and expert weights;
• local body collectives for variant-specific expert computation;
• role arguments that map user-provided mesh axis names onto MoE behavior.

Axis names themselves are only handles for selecting mesh dimensions. The semantics come from which role argument uses that axis.

Changes

3D+ MoE local_map placement construction

Adds a role-based placement helper for DSv3 MoE.

The helper builds token and expert-weight placements from the selected MoE roles:

• token axes shard batch tokens with Shard(0);
• EP shards the expert dimension with Shard(0);
• ETP shards expert hidden dimensions for grouped matmuls;
• EFSDP optionally shards expert weights at the HOP boundary.

This allows the same DSv3 MoE implementation to represent 2D, 3D, 4D, and higher-dimensional mesh variants without changing AutoParallel main graph logic.

Variant-specific MoE body collectives

Extends the local MoE body so the implementation matches the selected boundary placements:

• ETP performs the expert tensor-parallel output all-reduce.
• EFSDP performs expert-weight all-gathers inside the local body before grouped expert matmuls.
• Existing EP dispatch/combine behavior remains the expert-parallel communication path.

This is necessary because AutoParallel does not rely on DTensor inference inside the local_map body.

DeepSeekV3 user-facing role arguments

Extends the DSv3 model/MoE constructor path with role arguments for EP, token axes, ETP, and EFSDP.

Users express the intended layout by constructing a named DeviceMesh and passing the matching role arguments to DeepSeekV3Model. They do not manually write local_map regions for the DSv3 example.

local_map HOP compatibility

Updates placement option handling so the current local_map_hop target name is recognized in addition to older local_map target names.

3D expand propagation regression

Adds an AutoParallel propagation rule for aten.expand.default. This fixes the 3D DSv3 attention failure where DTensor's expand rule raises KeyError: 2 when a candidate strategy shards the singleton head dimension that expand turns into the full head dimension.

Invalid expand mappings are skipped; valid placements continue through normal redistribute-cost generation.

Documentation

Updates docs/local_map_and_moe.md with user-facing guidance for:

• 2D, 3D, 4D, and 5D-style MoE mesh variants;
• how DP replicate, EFSDP, EP, and ETP roles compose;
• why axis names are handles rather than semantics;
• how AutoParallel's HOP-boundary model differs from TorchTitan's external FSDP resharding model.

Validation

Codebase version:

commit: 32466d48ecbac44acccf18ed433a135a41f0f6bc
base: a54c2529b1cb10b6b27a585b4a6fc2f84bc1d232
merge_base: a54c2529b1cb10b6b27a585b4a6fc2f84bc1d232

Results:

python -m py_compile autoparallel/shardings/propagation_rules.py tests/test_moe_local_map_variants.py
exit_status: 0

python -m pytest tests/test_moe_local_map_variants.py -q
8 passed, 14 warnings in 1.94s

python -m pytest tests/test_ops.py -q -k 'local_map or sharding_constraint'
6 passed, 12 deselected, 10483 warnings in 5.04s

python -m pytest tests/test_graph_clustering.py -q -k clustering_high_coverage
1 passed, 1 deselected, 1104559 warnings in 93.60s

Earlier, I also validated that an EFSDP-enabled DSv3 AutoParallel trace contains the expected expert-weight all-gathers inside the local_map body:

all_gather_count=3

3D full-graph DSv3 AutoParallel smoke was also investigated. Before the expand rule, it failed at dense attention aten.expand.default with KeyError: 2. After this fix, that correctness failure is removed, but the local full 3D build did not complete within the diagnostic window because the later full-graph 3D strategy/cost search is still very large. This PR does not claim that full 3D DSv3 AutoParallel optimization completes locally end-to-end.

Notes

This PR does not make AutoParallel synthesize arbitrary MoE implementations. It makes the DSv3 MoE example expose the correct 3D+ local_map boundary and body for supported EP, ETP, EFSDP, and token-sharding role combinations.