Fix AutoTP metadata updates for unsharded modules (#8299)
## Summary
- defer module metadata updates until recursive AutoTP replacement has
completed
- update logical dimensions only when a subtree contains parameters
carrying AutoTP TP metadata
- preserve metadata for replicated vision modules and high-dimensional
weights while retaining existing sharded Linear behavior
- add regression coverage for PatchEmbed-style Conv3d wrappers,
unsharded attention subtrees, and genuinely sharded parameters
Fixes #8285.
## Problem
`update_mp_params()` previously ran before Linear replacement and scaled
logical attributes for every traversed module. In multimodal models this
changed replicated vision metadata, such as Qwen2-VL PatchEmbed
`embed_dim`, even though the corresponding Conv3d weight was not
partitioned. The metadata then disagreed with the full weight shape and
broke vision forward.
## Validation
Current master baseline `e70b90352e93e07862c45ebb946d75b813f025a2`
reproduces on 2-GPU NCCL with `patch_embed.embed_dim=60` while
`proj.out_channels=120`.
- CPU targeted: `7 passed, 20 deselected`
- 2-process Gloo metadata check: unsharded `num_heads=16`, sharded
`hidden_size=6` on both ranks
- 2-GPU custom-pattern suite: `27 passed`
- 4-GPU AutoTP/ZeRO adjacent suite: `33 passed, 2 skipped, 24
deselected`
- synthetic training on 1/2/3/5 RTX 3090 GPUs: forward, backward, and
optimizer step all complete with zero cross-rank output, PatchEmbed
gradient, and PatchEmbed parameter difference
- 5-GPU synthetic, 3 runs with 3 warmup and 20 measured steps each:
`5.008 ms` mean step latency, `409,216 tokens/s`, `28,170,752` peak
allocated bytes, `52,428,800` peak reserved bytes
- Qwen2-VL-2B AutoTP=2 vision forward: PatchEmbed `1280 -> 1280`, vision
heads `16 -> 16`, language q_proj `[1536,1536] -> [768,1536]`, zero
cross-rank output difference
- changed-file pre-commit hooks, `DS_BUILD_OPS=0 python setup.py check`,
`compileall`, and `git diff --check` pass
---------
Signed-off-by: Wang Zupeng <zupenwang@gmail.com>
Co-authored-by: Ma, Guokai <guokai.ma@gmail.com>
