[cosmos3] Cosmos3-Nano GRPO support: train pipeline config, VideoAlign reward, T2V recipe - #25
Draft
zhihengy wants to merge 1 commit into
Draft
[cosmos3] Cosmos3-Nano GRPO support: train pipeline config, VideoAlign reward, T2V recipe#25zhihengy wants to merge 1 commit into
zhihengy wants to merge 1 commit into
Conversation
…I recipe Rebased onto the loss-hub/input-dtype-policy refactor as one commit (the pre-refactor history is preserved on backup/feat-cosmos3-pre-rebase; most intermediate commits were add-then-drop pairs already split to other PRs). - Cosmos3TrainPipelineConfig: packed single-sample forward over the MoT DiT (UND/GEN towers), token-level conditioning (text_ids/text_mask/fps on CondKwargs), GEN-only training with UND frozen, timestep-scale handled in-model. - Boundary dtypes moved to the refactor's input_dtype_policy: latents=default, cond=passthrough (the packed forward casts its own inputs; mRoPE position ids ~15000 would be scrambled by a bf16 boundary cast), timestep=fp32 (karras grid is non-integer; bf16 rounds 993.25 -> 992). Replaces the pre-refactor fsdp_cast_forward_inputs / cast_timesteps_to_forward_dtype flags. - use_cfg: guidance <= 1 runs single-branch, matching sglang do_cfg (now in loss_hub.flow_grpo). - --fsdp-frozen-params-dtype: frozen base stays at checkpoint-native dtype; only LoRA params get fp32 master copies. - Engine: multi-GPU engines pin the full contiguous CVD slice; num_gpus owns the rollout allocation with tp/sp from --sglang-* passthrough. - Weight sync: send one copy of the all-gathered payload (size-1 contract; the engine shards internally). Co-authored-by: Cursor <cursoragent@cursor.com>
2 tasks
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
1 participant
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
What
End-to-end GRPO post-training support for nvidia/Cosmos3-Nano (16B MoT: 8B UND + 8B GEN):
Cosmos3TrainPipelineConfig: replicates the diffusers packed-sequence forward (text+vision joint sequence, mRoPE via diffusers helpers), UND tower frozen by param-name fragments (it sits inside the training graph), LoRA targets = GEN attention only.Cosmos3ModelBackendloads transformer+scheduler only.CondKwargsgainstext_ids/text_mask/fps— Cosmos3 has no separate text encoder; shipping tokens verbatim eliminates the text-replay-consistency failure class.runtime_env.py_executable; rolling per-dimension logging (TA collapse is the canonical hacking mode and is invisible in the Overall sum).run-diffusion-grpo-cosmos3-videoalign-4gpu.sh(3 colocate + 1 reward GPU, 16-step SDE, 17f 832x480, LoRA r64) and_smoke_cosmos3_t2i.sh(2-GPU pipeline smoke).Companion sglang-d branch:
feat/cosmos3-rl-rollout(rollout SDE-Euler on the serving sigma grid + trajectory sigmas + fused-param weight-sync fix).Validation
ratio_abs_minus_1stable at 1–2.5e-5 (10x below clip range), cross-engine weight-sync checksums equal.miles-diffusion-grpo/diffusion_grpo_cosmos3_videoalign_*): rollout 768x17f in ~15 min on 3 engines, VideoAlign reward mean -1.1 ± 1.8, first steps healthy.Draft: long-run reward trend still being monitored; batched multi-sample generation per request deliberately deferred (packed forward is single-sample).