[cosmos3] Bitwise train/rollout parity: FSDP precision spec + cosmos3_bitwise rollout patch group

[zhihengy]

Stacked on #25 (feat/cosmos3).

Summary

Makes the Cosmos3 GRPO training forward (miles/diffusers/FSDP2) produce bit-identical tensors to the rollout engine (sglang-d), verified across every aligned module output: 103/103 comparable tensor pairs bit-exact (both CFG branches, UND+GEN towers, first and last denoise steps, embed_tokens through proj_out).

Train side

• time_embedder fp32 island: diffusers declares _keep_in_fp32_modules = ["time_embedder"] and sglang-d pins it to fp32 at load; the blanket bf16 gather silently downgraded it on the train side. The family FSDPParallelPlan.param_dtype_patterns (the refactor(fsdp)(5/5): compile per-model precision plans onto the dtype patch #106 mechanism) gathers it at fp32, the wrapper runs it with autocast disabled, and identical sinusoid rows are deduplicated before the MLP (cuBLAS fp32 GEMMs are not bitwise M-invariant, measured on the 4096x4096 MLP; sglang-d runs M=1 per request).
• diffusers RMSNorm → F.rms_norm: eager diffusers RMSNorm rounds to bf16 before the weight mul (two roundings); sglang-d keeps fp32 through the mul. Upgrade the train side (never downgrade) so both run the identical op.

Rollout side (--rollout-patch-group cosmos3_bitwise, per the #97 CLI selection; the recipe sets it)

• Pin TORCH_SDPA attention backend (matches the train-side SDPA path).
• Route sglang RMSNorm through the same F.rms_norm.
• Unfuse MergedColumnParallelLinear (to_qkv, gate_up_proj) into per-slice F.linear GEMMs — fused column blocks are not bitwise equal to the standalone GEMMs diffusers runs. Slices come from output_sizes (with tp=1 init order, output_partition_sizes collapses to one fused slice); the patch fails loud rather than silently falling back fused.
• Eager SiluAndMul and split (unfused) qk-norm/RoPE path.
• Sequential batch-1 CFG: uncond/cond as separate forwards — cuBLAS is not batch-invariant on cosmos3 shapes, and CFG batching also changes text padding.

Validation

• Dump pipeline (per the dumper skills): rollout engine and trainer instrumented with the sglang dumper, records paired by bit-exact latent anchors; pair_metrics shows max_abs_diff = 0 for all 103 aligned pairs; 3 remaining records are benign scalar-vs-broadcast shape mismatches with byte-identical values.
• Originally validated pre-refactor (e707d30) with a ported PrecisionSpec mechanism; now rebased onto main's own precision machinery (refactor(fsdp)(5/5): compile per-model precision plans onto the dtype patch #106 param_dtype_patterns + ParamDtypeMixedPrecisionPolicy) and the refactor(diffusion): move rollout patch group selection to a CLI arg #97 CLI patch-group selection — same fp32 gather island, no parallel implementation. The dump comparison should be re-run on this base before merge.

Test plan

• Re-run the alignment dump pair on this branch (train + rollout) and confirm 103/103 bit-exact
• Cosmos3 T2I GRPO smoke (recipe script in [cosmos3] Cosmos3-Nano GRPO support: train pipeline config, VideoAlign reward, T2V recipe #25) trains with identical loss/log-prob traces

ci-sglang-pr: 0065fbfae1d22626203c158b9e8c0e8f7126c4cd

[Rockdu]

seems like this old bash script is conflicting with the latest one

[Rockdu]

need a careful migration