Conversation
…he current one trains Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com>
Rockdu
force-pushed
the
feat/sft-rollout-prefetch
branch
from
September 18, 2026 19:48
c674d36 to
e0a3321
Compare
Rockdu
marked this pull request as ready for review
September 18, 2026 21:49
…ct() from the train actor's per-rollout clear_memory Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com>
Rockdu
force-pushed
the
feat/sft-rollout-prefetch
branch
from
September 18, 2026 21:53
e0a3321 to
23ed940
Compare
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.
What
train_sft.py: a train-only entry point that submitsgenerate(i+1)before waiting ontrain(i), so only the first rollout is waited for.train_diffusion.py(the RL loop) is untouched; both SFT recipes now launchtrain_sft.py.RolloutManager.generate(rollout_id, train_in_flight=None): when the SFT encoders sit on the training GPUs, the encode pool waits for the in-flight train step before an encode burst (SftEncodePool.wait_for_train_step, timed asperf/sft_encode_wait_time). With--rollout-num-gpusthe encoders have their own GPUs and encoding overlaps training too.generate(i)returns, before the prefetch moves it.--skip-train-actor-gc-collect: the train actor's per-rolloutclear_memoryskipsgc.collect()and only releases the CUDA cache. Opt-in; the recipes keep the default.Why
An SFT rollout encodes a dataset batch and does not depend on the trained weights, so nothing forces the RL ordering
generate -> train -> update_weights -> generate. With a warm.sft_cachea rollout is a cache load, aray.putand a dp split; that time was fully exposed between train steps.Two traps the loop has to respect:
generate(i+1)already moved it; a resume ati+1would then skip one rollout's samples.Validation
1-GPU H3 LoRA SFT (PR #210's recipe, 16 clips,
--rollout-batch-size 8, colocated encoder),perf/train_wait_timeper rollout, mainebd55fc1vs this branch:--skip-train-actor-gc-collectA 2-GPU run of main and this branch under
--deterministic-modeproduced identical per-step train logs and bitwise-equal DCP checkpoints (model, optimizer, lr_scheduler) at every save.Files
train_sft.py— the train-only loop with rollout prefetchmiles/ray/rollout.py—generatetakestrain_in_flightand hands it to the encode poolmiles/rollout/sft_rollout.py—SftEncodePool.wait_for_train_stepmiles/utils/arguments.py,docs/user-guide/cli-reference.md—--skip-train-actor-gc-collect;--rollout-num-gpushelp mentions the overlapmiles/utils/memory_utils.py,miles/ray/train_actor.py—clear_memory(collect_garbage=...)scripts/run_diffusion_sft_h3_t2va.py,scripts/run_diffusion_sft_wan22.py— launchtrain_sft.pytests/fast/test_train_sft.py— scheduling order, cursor save order,train_in_flightper topologydocs/user-guide/concepts.md,docs/user-guide/launch-script.md,docs/developer/contributor-guide.md,docs/models/h3/lora_sft_guide.md— the new entry point and the overlapChecklist
pre-commit run --all-filespassespytest tests/fast -xis green (284 passed onradixark/miles_diffusion:latestat the branch head)python3 train_diffusion.py --helpstill parses (train_sft.py --helpon the branch head)