[Fix] GLM-5.3-Flash F2.1 vision training and parallel gradients

[Wuyh11]

This F2.1 follow-up corrects Muon scaling for vision convolution parameters, synchronizes FSDP-ignored FP32 replicas across the complete model mesh, and preserves embedding-input gradients under sequence parallelism.

Target: feat/glm53flash-f2-vision-tower. Related: #2110.

Changes and impact

• Vision convolution updates: Flattened convolution parameters used their trailing dimensions instead of the logical matrix [out_channels, product(input_dimensions)] used for orthogonalization. Derive the Muon learning-rate ratio from each original global parameter shape and the current parameter-group configuration in _build_muon_task, before communication or to_local(). Pass that fresh ratio to the update and write the same value back to optimizer state, so checkpoints record the ratio computed for the latest update task. Saved checkpoint ratios do not influence updates. Preserve momentum/steps and the original clipping-policy load hook, without ratio-specific load/save hooks. Cover 3D/4D/5D convolution/matrix update equivalence, configuration changes, and legacy/DCP checkpoint restoration. The original convolution fix was cherry-picked from JT-Ushio/xtuner-muonsplit-verify, caac8ac8; the ratio-state handling follows the review suggestion.
• FP32 replica synchronization: FP32 parameters retained an EP-only replica mesh while being ignored by FSDP, leaving independent copies across FSDP rows. Expand these replicas to the full world mesh so gradient averaging produces consistent global norms and clipping factors. Detach the local DTensor when creating the replacement leaf parameter; preserve existing full-world replicas and sharded placements. Separate and document the four mesh/layout cases: preserve DTensors already on the world mesh and non-replicated layouts on other meshes; replicate plain parameters on the world mesh; rebuild fully replicated DTensors from other meshes. Explain when this expands the replica set across FSDP rows and when it only normalizes the mesh representation.
• SP embedding gradients: Replace the detached gather in SequenceContext.raw_inputs_embeds with all_gather_tensor_autograd, making inputs_embeds contiguous before the collective to support transposed and strided embedding views. The contiguous copy preserves autograd; backward uses reduce-scatter SUM to return contributions to their owning ranks, including MTP future-token dependencies crossing SP shard boundaries. Keep the cached gathered tensor attached to its graph.
• Vision RoPE dtype documentation: Spell the frequency dtype as torch.float32 and document why the non-persistent buffer must remain real during meta construction.

Validation

Current head a0c799ba: review follow-up (October 12):

• 8 native CPU tests passed, covering preservation of legacy ratio metadata, momentum/steps and caller dictionaries, empty untouched-parameter state, and invalid MuonSplit initialization.
• All 5 committed Muon regression bodies passed on two H200 CUDA/NCCL ranks, using their original CUDA operations and assertions with an already-initialized distributed fixture. These check 3D/4D/5D matrix-update equality, existing scaling policies, legacy checkpoint continuation, group configuration changes, and a real DCP optimizer disk round trip. Loaded stale ratios are preserved until the next update, then overwritten with the ratio actually passed to the task; momentum and steps remain correct. The same checks passed on two CPU/Gloo ranks. These exercise DTensor optimizer updates and DCP APIs, rather than full Trainer/FSDP cold resume.
• Transposed and strided-slice embedding checks passed on two H200 CUDA/NCCL ranks and two CPU/Gloo ranks. The previous gather call reproduced the non-contiguous-input error. The actual updated SequenceContext property and MTP packed-roll code matched the full-sequence forward and backward references exactly, including gradients to the original leaf tensor across an SP boundary. Cache, no-grad, non-SP and None checks also passed. The standalone SP validation code remains outside the commit.
• Both existing clipping-policy load cases passed on CPU, preserving the current policy and caller checkpoint dictionaries, including the default for an added parameter group.
• Ruff lint/format and git diff --check passed for all three files changed in this follow-up. Full pre-commit was unavailable in the runtime. No new collectives were added; contiguous inputs retain the existing storage, while non-contiguous inputs require a copy. No throughput benchmark was run.

Earlier F2.1 validation at 06298610, compared with F2 99266d89 (October 12):

• Native vision CPU suite: 11 passed, 1 skipped, 3 GPU tests deselected. This covers FP32/BF16 HF forward parity, image/video/multi-image cases, backward gradient flow, meta construction, RoPE frequencies, and merge-alignment rejection. The skipped test used an unavailable default checkpoint path; a separate attempt with the real checkpoint is documented below.
• Muon: all 3 then-committed regression bodies passed on two CPU/Gloo ranks, adapting their CUDA device/fixture while retaining the original algorithm and assertions. Extended that CPU probe to 5D Conv3d layouts. F2 reproduced both the convolution-scaling and legacy-checkpoint errors; F2.1 passed them, while the existing-policy case passed on both versions.
• FP32 replicas: four-process Gloo comparison passed. The actual parameter setup, MoE gradient reduction, norm calculation, clipping, and two SGD updates were exercised. F2 reduced [1, 2, 3, 4] to 1.5 / 3.5; F2.1 reduced them to 2.5 on every rank. With nonuniform gradient components, the replica parameter difference after two clipped updates fell from 0.0622052 to 0, matching the global-gradient reference. Native FSDP wrapping was stubbed for this CPU test.
• SP embeddings: four-process Gloo comparison passed. Using the actual packed-roll and SequenceContext methods, F2 lost the future-embedding gradient; F2.1 matched the full-sequence reference exactly, including a dependency crossing the shard boundary. Forward loss was identical between the two versions. Cache, no-grad, integer-ID, non-SP, and None cases also passed. For CPU transport, only the original gather's CUDA device assertion was removed.
• Vision SP: the original padding regression passed on two CPU/Gloo ranks, retaining rtol=1e-5, atol=1e-5. Twenty raw patches were padded to 24, processed through Ulysses and the local projector, then trimmed to five merged tokens. The maximum difference versus non-SP was 0.
• Real 25B checkpoint mapping: The original CPU test failed before loading because the shared profiler calls the unavailable torch.cpu.reset_peak_memory_stats. In a separate CPU-adapted check, disabling only GPU memory profiling and accelerator synchronization allowed all 347 vision/projector parameters to load with 0 missing/unloaded keys and correct analytical FP32 RoPE frequencies. This adapted pass does not erase the original CPU compatibility failure.
• Ruff lint/format checks passed for all five changed files; git diff --check passed. Full pre-commit was not run.

The earlier 06298610 results above used CPU/Gloo. The current review follow-up adds the CUDA/NCCL Muon and embedding-gather checks listed first; full 25B training, full Trainer/FSDP resume, compile, and throughput were not rerun. The current F2 baseline already has real CPU RoPE initialization; this PR's RoPE change explicitly names the FP32 dtype and updates its comments.

Earlier GPU evidence for the FP32/SP fixes (September 29–30):

• Four-GPU FSDP2 × EP2 regression: 1/1 passed, with replica-mean gradients, matching norms, and clipping results. An eight-GPU probe confirmed complete FSDP × EP replication for 28 affected FP32 parameters.
• MTP GPU regressions: 2/2 passed, including elementwise agreement with the packed full-sequence gradient reference and the no-grad gather case. In the causal comparison over all 347 vision/projector parameters, the SP2-versus-SP1 gradient relative L2 difference decreased from approximately 14.40% to 3.49%.
• 8×H200 continuous training: BF16, EP4, global batch 8, 16K packs, AdamW, MTP coefficient 0.1, balance alpha 0.001, activation offload and recomputation. SP1/2/4/8 each completed 20 optimizer steps on all eight ranks, with finite gradients, advancing optimizer state, and zero cross-rank gradient-norm spread at all 80 steps.

SP	Optimizer steps per rank	Step-20 loss	First-step visual gradient relative L2 vs SP1
1	20	10.071869	Reference
2	20	10.069990	3.468954%
4	20	10.079679	3.283411%
8	20	10.100826	1.267570%

The earlier GPU runs used an older GLM53 snapshot and AdamW; they provide evidence for the FP32/SP fixes, rather than full-model training validation of the current combined head. Current Muon CUDA/NCCL checks are listed separately above. Strict full-model numerical equivalence across SP sizes remains unresolved. No before/after throughput benchmark was run.

[YanhuiDua]

LGTM

[YanhuiDua]

SP gather 缺少 .contiguous()，非连续 embedding 会报错，原代码中 gather_for_sequence_parallel() 会先执行 input = input.contiguous()

建议改为：

gathered = all_gather_tensor_autograd(
    self.inputs_embeds.contiguous(),
    gather_dim=1,
    group=self.sequence_parallel_mesh,
)

[YanhuiDua]

建议不要修改原有的 _restore_clip_grad_policy hook，也不需要在加载时额外清理 lr_ratio。每次更新都会重新计算比例，不读取 checkpoint 中保存的旧值，因此旧值不会影响更新；clip hook 继续只负责恢复 clipping 策略即可。

为了保证后续保存的 state_dict 中，lr_ratio 与更新采用的值一致，建议在 _build_muon_task() 中，根据原始参数的全局 shape 和当前参数组配置重新计算，并将同一个值写回 state、传给更新任务：

lr_ratios = [self._param_lr_ratio(p, group) for p in params]
assert len(set(lr_ratios)) == 1, (
    f"Found different lr_ratios: {set(lr_ratios)}"
)

for state, lr_ratio in zip(states, lr_ratios):
    # Keep the saved ratio aligned with the value used for this update.
    state["lr_ratio"] = lr_ratio

后续 muon_update_batch_async() 继续传入 lr_ratio=lr_ratios[0]。这样更新使用重新计算的正确比例，保存 checkpoint 时也自然保存同一个值，无需额外的加载清理或保存时重算逻辑。

[jayhenry]

直接报错？

[YanhuiDua]

Follow-up on the FP32 ignored-param mesh fix: #2162.

#2157 correctly identified that EP-only FP32 replicas desync FSDP-rank grads for norms/clipping, but the four-way case split in base.py is heavier
than needed.

The invariant is simpler: every FP32 FSDP-ignored parameter should end up as a full-world replica — DTensor(world_mesh, [Replicate()] * ndim).
#2162 always rematerializes that way via full_tensor() + distribute_tensor, instead of preserving alternate meshes/placements. Suggest reviewing
#2162 as a cleaner follow-up to this part of #2157.

[YanhuiDua]

Follow-up to the Muon conv lr_ratio fix here: #2161 refactors matrix layout so Newton-Schulz and lr_ratio share one _muon_matrix_shape contract (2D/3D), instead of inferring shapes separately.