[cosmos3] Cosmos3-Nano GRPO support: train pipeline config, VideoAlign reward, T2V recipe#25
Draft
zhihengy wants to merge 30 commits into
Draft
[cosmos3] Cosmos3-Nano GRPO support: train pipeline config, VideoAlign reward, T2V recipe#25zhihengy wants to merge 30 commits into
zhihengy wants to merge 30 commits into
Conversation
Cosmos3 has no separate text encoder: rollout ships text_ids/text_mask/fps, and the training forward replays the diffusers packed-sequence layout (single-sample joint UND+GEN sequence, per-patch timesteps). UND-tower params are frozen before FSDP; LoRA targets the GEN attention projections. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
VideoAlign runs in a pinned conda env (transformers 4.45.x) via Ray runtime_env; the actor lives in its own module so the worker process never imports miles' training deps. Reward = normalized VQ+MQ+TA. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
- cond text_ids/text_mask arrive 1-D after per-sample extraction - cast the fp32 timestep sinusoid to the MLP weight dtype (matches sglang-d's TimestepEmbedder; diffusers crashes under FSDP bf16) - request one sample per rollout request (microgroup-size 1): the Cosmos3 pipeline generates a single sample per request First aligned step: ratio_abs_minus_1=1.0e-5, approx_kl=7e-11. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
Ray's conda runtime_env needs a conda binary on the worker PATH, and cluster python versions must match — pin the videoalign env to py3.11 and hand Ray the interpreter path directly. Adds the adaptation report. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
TA collapse is the canonical VideoAlign hacking mode (DanceGRPO dropped TA; Diffusion-DRF documents full score collapse) and is invisible in the Overall sum. Report gains the same-RM recipe benchmark: success band Δ+0.6~0.8 (SAGE-GRPO parity), +1.5 and above is hacking territory. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
… in the T2V recipe VideoAlign is hacking-prone; the clean SAGE-GRPO precedent trains with a KL leash while the KL-free ablations collapse. Debug mode restores the model_output alignment diagnostics the Wan recipe logs. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
…ebug mode Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
Root cause: MixedPrecisionPolicy(cast_forward_inputs=True) quantized the float32 mRoPE position ids (~15000 after the temporal modality margin, where bf16 spacing is 128) to bf16 at the FSDP boundary, scrambling the rotary phases of every vision token in the train-side forward. Wan never hit this: its DiT forward has no large-magnitude float inputs. - TrainPipelineConfig.fsdp_cast_forward_inputs (default True, unchanged for existing families); Cosmos3 disables it — the packed forward casts its own inputs. - TrainPipelineConfig.cast_timesteps_to_forward_dtype (default True); Cosmos3 disables it: the karras grid has non-integer timesteps and the bf16 round (993.25 -> 992) drifted the cond branch 0.008 -> 0.021. - MILES_ROLLOUT_START_PORT env: shift engine port range so concurrent miles runs on one host cannot collide (free-port probe misses servers bound to a specific interface; collisions produced cross-cluster weight syncs and the earlier 'Invalid device_uuid' failures). - MILES_DIAG_DUMP env: dump compared train/rollout tensors per microbatch (the instrument that isolated this bug). T2V 17f validation: model_output_mean_abs_diff 0.64 -> 0.068, ratio_abs_minus_1 2.5e-5 -> 5.9e-6. Residual 0.068 is fa3-vs-SDPA kernel noise amplified by CFG (g + (1-g) coefficients). Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
…save A 410-rollout run accumulated 41 x 6.9GB snapshots and filled the storage quota mid-run. Default 0 keeps current behavior. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
The Wan-inherited candidate steps 1,2,3 sit on the dense head of the karras flow grid (|dt| 0.0025-0.0054 vs Wan's ~0.1), making the policy gradient ~30-50x weaker per pair — after 30 optim steps the ratio was still at the alignment noise floor (7e-6, clipfrac 0) and reward was flat. Steps 8-14 carry |dt| 0.045-0.16 across sigma 0.91->0.25. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
At the meaningful-dt trained steps the fa3-vs-SDPA velocity diff (~6.9e-2) maps to a 1.4e-4 log-prob deviation — above clip_range 1e-4 — so ~80% of pairs were clipped by pure numerics before any policy motion. The SGLANG_DIFFUSION_ATTENTION_BACKEND env is a dormant path; the live knob is ServerArgs.attention_backend via miles' --sglang-* passthrough. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
Kernel-family alignment (torch_sdpa both sides) left the residual velocity diff unchanged at 6.9e-2, so the log-prob noise floor at the trained steps stays ~1.4e-4 and a 1e-4 clip wiped ~80% of pairs before any policy motion. 1e-3 restores the ~7x floor-to-clip headroom Wan has. Rollout back on fa (sdpa: +30% time, no diff gain). Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
… clip The alignment noise floor is step-dependent: at sigma 0.25-0.4 (steps 13-14) the same velocity diff maps to ~4e-3 in log-prob — above the 1e-3 clip — so rollouts drawing those steps trained at clipfrac=100% (observed rollout 4). Steps 8-11 keep dt 0.045-0.14 with a ~1.3e-4 floor. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
…t, window 2 Per-optim-step policy movement was ~1e-4 against a 1e-3 clip budget (10% utilization) with only 2 steps and 1 trained grid step per ~16 min rollout. Rollout stays the bottleneck; this is ~8x movement per hour. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
dcp set_optimizer_state_dict and lr_scheduler.load_state_dict silently
restore the checkpoint's lr, defeating an lr change passed for the
resumed run. Re-apply args.lr after both loads. Also pass through extra
script args ("$@") for --load/--ckpt-step.
Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
step() recomputes group lr from max_lr each step, so writing base_lrs and param_groups alone is overwritten on the next scheduler step. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
…nt PPO ratio Rollout-stored old log-probs carry ~5e-2 bf16 execution noise vs the trainer forward (unbiased, not attributable to any single op — RoPE and attention verified bit-consistent). Recomputing old with the trainer forward at rollout start cancels that noise from the ratio exactly: first-step ratio_abs_minus_1 becomes 0.0 (was ~1e-4 floor). Matches flow_grpo's trainer-side log-prob convention. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
… CFG); GUIDANCE env in t2i script Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
…ross sigma steps KL = |dmu|^2/2std^2 makes low-sigma steps' spring 25-115x stiffer (karras grid step 16 vs 8), stalling full-grid training at the low-step equilibrium. Clamping std in the denominator aligns their stiffness with the mid-sigma anchor zone. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
Untrack the smoke/diag scripts and the VideoAlign T2V recipe (kept locally, out of the PR). Replace the working pickscore script with a main-style recipe: fixed hyperparameters (the validated CFG-free config: lr 3e-4, adam_beta2 0.95, clip_grad 2e-3, KL beta 1e-3 + std floor, recompute old log-prob, SDE candidate steps 4-15), standard dataset bootstrap, wandb project, and 5gpu naming to match the Wan2.2 recipe convention. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
…t dtype Large frozen bases (Cosmos3-Super 64B) cannot pay the fp32 master cost: load at the frozen dtype and upcast only LoRA params to master dtype. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
rollout_num_gpus_per_engine was forced into tp_size (LLM semantics) and num_gpus stayed 1, breaking multi-GPU diffusion engines (ulysses SP). Single-GPU engines are unaffected (both fields were and remain 1). Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
_pin_to_assigned_gpu pinned exactly one device; ulysses engines with rollout_num_gpus_per_engine>1 crashed with invalid device ordinal on worker rank 1. Single-GPU engines keep the original mapping path. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
Every rank in the engine gather group serializes the full all-gathered weights; ulysses (tp=1) engines accept exactly one payload and replicate internally, so posting one per trainer rank got a 400 rejection. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
Trainer payloads are identical all-gathered copies; size-1 satisfies the engine contract for both replicated (ulysses) and TP engines, which shard internally. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
All ranks IPC-export flattened buckets at serialize time; exported storage parks in CudaIPCSentDataLimbo and is only reclaimed by torch.cuda.ipc_collect(). Without it the trainer accumulates the full sync volume (128GB for Cosmos3-Super) and OOMs mid-sync. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
Serializing a fresh 2GB flattened bucket per send pins each exported storage in the CUDA-IPC ref-count machinery; on some ranks these were never reaped, accumulating the full sync volume (128GB for Cosmos3-Super) and OOMing mid-sync. Copy every bucket into one persistent per-dtype staging tensor and export views of it: one pinned storage per rank for the whole run, bounded regardless of reaping. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
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).
🤖 Generated with Claude Code