[diffusion] RL rollout support for the Cosmos3 pipeline#30486
Draft
zhihengy wants to merge 4 commits into
Draft
Conversation
Bind a flow-match Euler scheduler to rollout=True requests (same per-request switch as the Wan rollout PR), reusing the serving scheduler's sigma grid so rollout noise levels match serving exactly. The denoising env ships text_ids/text_mask (cond+uncond) and fps; the dit trajectory now carries the scheduler sigma snapshot so the training side replays exact noise levels. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
- WeightsUpdater: carry the param-mapping merge index through to weight_loader as the shard id; fused params (to_qkv, gate_up_proj) previously narrowed out of bounds. Cosmos3 is the first fused-param model through the tensor-update path. - jit_kernel fa3: drop only_qv=False before calling the kernel so older sgl-kernel builds without the argument keep working. - Cosmos3DecodingStage: propagate rollout_trajectory_data to OutputBatch. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
Contributor
|
Warning You have reached your daily quota limit. Please wait up to 24 hours and I will start processing your requests again! |
…ted grid The checkpoint's karras schedule ignores flow_shift; an explicit override now builds the rollout Euler grid with the requested shift (smoke: karras-16 still scores best, shift grids stay within the UniPC-35 band but degrade visibly at shift 2-3). Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
Imported CUDA-IPC blocks stay cached in the consumer allocator until torch.cuda.ipc_collect(), pinning the sender's exported storage. Large models (Cosmos3-Super, 128GB sync volume) OOM the trainer mid-sync. 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.
Motivation
GRPO-style post-training for Cosmos3 needs rollout log-probs from the serving engine. This wires the existing post-training rollout machinery (SDE-Euler + log-prob, per-request scheduler binding from #30036) into the Cosmos3 pipeline, which uses hand-rolled stages rather than the generic denoising loop.
Modifications
Cosmos3TimestepPreparationStage: rollout=True requests lazily bind a flow-match Euler scheduler that reuses the serving scheduler's sigma grid (set_timesteps(sigmas=...)) — the Cosmos3 checkpoint ships a Karras flow-sigma schedule, so grid reuse keeps rollout noise levels exactly at serving parity regardless of scheduler config.Cosmos3DenoisingStage:RolloutDenoisingMixinhooks (prepare, per-step pre-step latent snapshots, SDE step viaflow_sde_sampling, postprocess). The denoising env shipstext_ids/text_mask(cond+uncond) +fps— KB-scale vs ~150MB/sample for per-layer KV. I2V+rollout is rejected (frame-0 re-injection breaks the Gaussian transition assumption).RolloutDitTrajectorygains asigmassnapshot (serialized in the rollout API) so trainers replay exact noise levels instead of reconstructingtimesteps/1000(fp32 round-trip drifts, and a Karras grid cannot be reconstructed at all).to_qkv,gate_up_proj) narrowed out of bounds — Cosmos3 is the first fused-param model through this path. The merge index now reachesweight_loaderas the shard id.only_qv=Falsebefore the kernel call so older sgl-kernel builds without the argument keep working.Cosmos3DecodingStage: propagaterollout_trajectory_datatoOutputBatch.Includes the two per-request rollout-scheduler infra files from #30036; will rebase once that lands.
Validation
Trained against miles-d (radixark/miles_diffusion#25): rollout↔train log-prob
ratio_abs_minus_1= 1–2.5e-5 across T2I and 17-frame T2V; SDE-Euler at 16 steps matches UniPC-35 output quality (PickScore + visual). Non-rollout serving paths untouched (rollout branch gated onbatch.rollout).CI States
Latest PR Test (Base): ❌ Missing
run-cilabel -- add it to run CI tests.Latest PR Test (Extra): ❌ Blocked --
run-ciis required first.