feat(flowgrpo): add bounded-memory chunked replay backward - #465
Open
Yuanjia-Lee wants to merge 1 commit into
Open
Yuanjia-Lee wants to merge 1 commit into
Yuanjia-Lee wants to merge 1 commit into
Conversation
timestep_chunk_size replays and backwards SDE steps in slices of that size, each scaled by len(chunk)/S so the summed gradient equals the default single full-graph backward (the DiffusionNFT iter_scale pattern; the optimizer boundary stays with TrainStack). Default None keeps the current behavior. Bounds the replay graph residency that grows with S even at micro_batch_size=1; costs one extra backward/reduce-scatter per chunk and non-bit-identical bf16 autocast accumulation.
Yuanjia-Lee
requested review from
CjhHa1,
Ideny42 and
haonan3
as code owners
September 14, 2026 15:25
12 tasks
This branch has not been deployed
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
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.
Summary
FlowGRPO replays all S SDE steps and runs one backward over the stacked
[B, S] log-probs, so every step's replay graph stays resident until that
backward.
sd3_trainside.yamlalready runsmicro_batch_size: 1withactivation_checkpointing: false; the resident graph still grows with S.timestep_chunk_sizemakes FlowGRPO replay and backward SDE steps in slicesof that size, each scaled by
len(chunk) / S; the summed gradient matches thedefault single backward (DiffusionNFT's
iter_scalepattern, optimizerboundary left with
TrainStack). DefaultNonekeeps the current behavior.A/B on one rollout and anchor (S=4, micro 1), real FlowGRPO and
SD3DiffusionStage.replay: fp32 gradients match to 7.9e-8 relative withidentical loss and ratio metrics, and the CUDA peak in the compute drops
43.8 → 36.3 MB (chunk 1) and 40.8 MB (chunk 2): a config that cannot hold S
resident graphs can still train all S steps.
Related Issue
N/A
Test Plan
SKIP=no-commit-to-branch pre-commit run --files unirl/algorithms/flowgrpo.py unirl/algorithms/README.md --show-diff-on-failure→ all hooks pass.
SD3DiffusionStagerollout, armstimestep_chunk_sizeNone/1/2 on thesame segment, beta 0 and 0.1. Numbers in the Summary.
Compatibility / Risk
autocast shifts the gradient by 0.46% and the AdamW delta by 4.7e-6 at
lr 3e-4 (fp32 matches).
Reviewer Notes
replay timesteps into one forward for throughput; no overlap.
Checklist