[hubert-marek]

Qwen shared-expert gates can have weight shape [1, hidden_size]. Default FSDP row sharding gives some ranks empty local tensors, which is incompatible with the equal-shard layout needed by the distributed Muon update.

Shard 2D single-row parameters along dimension 1 when their width is divisible by the parameter-sharding degree, obtained directly from parallel_dims.get_mesh("dp_shard_cp").size(). Other parameters retain the existing fused-parameter placement callback or default placement. FSDP still uses the HSDP mesh; the divisibility check excludes replication ranks. The placement rule applies during model setup regardless of optimizer selection.

There is no additional optimizer-setup validation. Non-divisible layouts retain the existing placement and failure behavior in Muon.

Validation performed locally:

• Two-rank CPU/Gloo DTensor reproduction with PyTorch 2.8.0: [1,128] / [0,128] under Shard(0) versus equal [1,64] shards under the placement callback, with exact reconstruction.
• Isolated placement checks across shard degrees 1/2/4/8/16, non-divisible widths, 1D/3D parameters and fallback delegation.
• Ruff lint/format and git diff --check passed. Local diagnostic scripts are not added to the repository. These checks do not exercise a full FSDP + Muon GPU update on current main.

Remaining merge blocker: the NIXL weight sender currently assumes dimension-0 shard offsets. Its handling of column-sharded gates must be fixed or the incompatible transport/layout combination rejected before merge. Checkpoint-layout compatibility also needs review.

---

Note

Medium Risk
Alters FSDP/DTensor layout for a narrow class of MoE gate weights; downstream weight sync or checkpoint code that assumes row-wise sharding may need matching updates.

Overview
FSDP setup now uses a custom shard_placement_fn instead of passing the fused-parameter callback directly. For 2D weights with a single row (e.g. Qwen shared-expert gates shaped [1, hidden_size]), parameters are sharded along column dimension 1 when the hidden width is divisible by the dp_shard_cp mesh size, so every rank holds an equal local shard for Muon’s all-to-all. All other parameters still use the existing fused shard_fused_on_dim1 placement or PyTorch’s default.

The change pulls in Shard from torch.distributed.tensor and reads shard degree from parallel_dims.get_mesh("dp_shard_cp").size() (parameter sharding only, not HSDP replication).

^{Reviewed by Cursor Bugbot for commit 159c63d. Bugbot is set up for automated code reviews on this repo. Configure here.}

[chatgpt-codex-connector]

Codex Review Summary

This comment shows the latest Codex review activity on this pull request.

Review	Status	Commit	Review trigger
📝 Code Review	✅ Completed 2026-09-29T22:42:06.770662Z	159c63d	Draft marked ready

ℹ️ About Codex in GitHub

Your team has set up Codex to review pull requests in this repo. Reviews are triggered when you

• Open a pull request for review
• Mark a draft as ready
• Comment "@codex review" or "@codex security review".

Codex reacts with 👀 while any review is running, comments if it has suggestions, and reacts with 👍 once all reviews finish with no findings.

[chatgpt-codex-connector]

💡 Codex Review

Here are some automated review suggestions for this pull request.

Reviewed commit: aa77830dcd

ℹ️ About Codex in GitHub

Your team has set up Codex to review pull requests in this repo. Reviews are triggered when you

• Open a pull request for review
• Mark a draft as ready
• Comment "@codex review".

If Codex has suggestions, it will comment; otherwise it will react with 👍.

Codex can also answer questions or update the PR. Try commenting "@codex address that feedback".

[S1ro1]

lgtm