[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

• precision.py: PrecisionSpec/compile_precision — module-granular FSDP2 gather-dtype overrides, compiled on clean FQNs (pre-LoRA/pre-FSDP) and lowered into the apply_fsdp2 wrap plan.
• 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 spec 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", applied at engine startup)

• 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.
• Validated on the pre-refactor base (e707d30). This PR ports the patches onto the loss-hub/autocast refactor (notably: autocast disabled inside the fp32 time-embedder island so the fp32 gather survives the bf16 autocast boundary); 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

Made with Cursor