Skip to content

feat(flowgrpo): add bounded-memory chunked replay backward - #465

Open
Yuanjia-Lee wants to merge 1 commit into
Tencent-Hunyuan:mainfrom
Yuanjia-Lee:feat/grpo-replay-chunks
Open

Yuanjia-Lee wants to merge 1 commit into
Tencent-Hunyuan:mainfrom
Yuanjia-Lee:feat/grpo-replay-chunks

Conversation

@Yuanjia-Lee

Copy link
Copy Markdown
Contributor

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.yaml already runs micro_batch_size: 1 with
activation_checkpointing: false; the resident graph still grows with S.

timestep_chunk_size makes FlowGRPO replay and backward SDE steps in slices
of that size, each scaled by len(chunk) / S; the summed gradient matches the
default single backward (DiffusionNFT's iter_scale pattern, optimizer
boundary left with TrainStack). Default None keeps 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 with
identical 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.
  • Harness on one 4090 (torch 2.11.0+cu130): real FlowGRPO over a real
    SD3DiffusionStage rollout, arms timestep_chunk_size None/1/2 on the
    same segment, beta 0 and 0.1. Numbers in the Summary.

Compatibility / Risk

  • Default off; recipes unchanged. One extra reduce-scatter per chunk; bf16
    autocast shifts the gradient by 0.46% and the AdamW delta by 4.7e-6 at
    lr 3e-4 (fp32 matches).

Reviewer Notes

Checklist

  • I reviewed the changed code and removed unrelated/generated artifacts.
  • I updated tests, docs, and configs where needed, or explained why not.

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.

This branch has not been deployed

No deployments
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

need review Ready and waiting for review

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant