[DistMuon] Support tensor-parallel storage layouts

[shuhuayu]

Summary

• Generalize DistMuon's optimizer redistribution schedule from a single storage mesh axis to multiple axes, allowing DP and TP storage layouts to participate in one optimizer transport plan.
• Support Shard and _StridedShard layouts for whole-matrix, per-head block, and expert Muon computation.
• Enable TP layouts in the Kimi K2.5 DistMuon recipe.
• Make QK clipping represent per-head scales as TP-sharded and data-parallel replicated DTensors, preserving the single packed MAX all-reduce with no per-weight collective.

Motivation

DistMuon previously assumed that optimizer redistribution operated over one storage mesh axis. Kimi's FSDP+TP parameters instead use multi-axis layouts such as _StridedShard(0) over the data-parallel axis together with Shard(0) over TP. This prevented the recipe from using TP even though model forward/backward TP was already available.

This change makes the optimizer transport operate over the combined storage axes while retaining the existing storage/compute layout contract. QK clipping is updated separately so TP-local attention heads are scaled against the same distributed storage layout.

Testing

• conda run -n tt12 python -m pytest tests/unit_tests/gpu/flex_shard/test_dist_muon.py -q (3 passed)
• conda run -n tt12 python -m pytest tests/unit_tests/gpu/test_qk_clip.py -q (5 passed)
• Four-GPU Kimi debug training with DP=2, TP=2:
  MODULE=kimi_k2_7 CONFIG=kimi_k2_5_debugmodel NGPU=4 ./run_train.sh --parallelism.tensor_parallel_degree=2 --training.steps=2
  • completed two optimizer steps
  • loss: 8.07139 -> 7.67561
• Formatting, flake8, pydoclint, codespell, and link checks pass for the changed files.