Skip to content

feat(sft): add train_sft.py for SFT data prefetch - #248

Open
Rockdu wants to merge 2 commits into
radixark:mainfrom
Rockdu:feat/sft-rollout-prefetch
Open

Rockdu wants to merge 2 commits into
radixark:mainfrom
Rockdu:feat/sft-rollout-prefetch

Conversation

@Rockdu

@Rockdu Rockdu commented Sep 17, 2026

Copy link
Copy Markdown
Collaborator

What

  • train_sft.py: a train-only entry point that submits generate(i+1) before waiting on train(i), so only the first rollout is waited for. train_diffusion.py (the RL loop) is untouched; both SFT recipes now launch train_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 as perf/sft_encode_wait_time). With --rollout-num-gpus the encoders have their own GPUs and encoding overlaps training too.
  • The data-source cursor is saved right after generate(i) returns, before the prefetch moves it.
  • --skip-train-actor-gc-collect: the train actor's per-rollout clear_memory skips gc.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_cache a rollout is a cache load, a ray.put and a dp split; that time was fully exposed between train steps.

Two traps the loop has to respect:

  • Colocated encoders (the recipe default) share GPUs with training; the H3 encoder is ~70 GB, so an overlapped encode burst must wait for the train step.
  • Saving the data-source cursor after training, as the RL loop does, would save it after generate(i+1) already moved it; a resume at i+1 would 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_time per rollout, main ebd55fc1 vs this branch:

rollout cache main branch, --skip-train-actor-gc-collect
0 cold, first rollout 278.2 s 237.0 s
1 cold, encode waits for train 0 150.1 s 146.8 s
2 warm, epoch-boundary save 6.2 s 5.7 s
3 warm 0.55 s 0.023 s

A 2-GPU run of main and this branch under --deterministic-mode produced 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 prefetch
  • miles/ray/rollout.pygenerate takes train_in_flight and hands it to the encode pool
  • miles/rollout/sft_rollout.pySftEncodePool.wait_for_train_step
  • miles/utils/arguments.py, docs/user-guide/cli-reference.md--skip-train-actor-gc-collect; --rollout-num-gpus help mentions the overlap
  • miles/utils/memory_utils.py, miles/ray/train_actor.pyclear_memory(collect_garbage=...)
  • scripts/run_diffusion_sft_h3_t2va.py, scripts/run_diffusion_sft_wan22.py — launch train_sft.py
  • tests/fast/test_train_sft.py — scheduling order, cursor save order, train_in_flight per topology
  • docs/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 overlap

Checklist

  • pre-commit run --all-files passes
  • Added/updated tests for new behaviour
  • pytest tests/fast -x is green (284 passed on radixark/miles_diffusion:latest at the branch head)
  • If launch flags changed, python3 train_diffusion.py --help still parses (train_sft.py --help on the branch head)
  • If a public flag was added, it appears in the CLI reference docs
  • If an example was added, it has a real walkthrough (n/a)

…he current one trains

Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com>
@Rockdu Rockdu added the run-ci-basic Authorize and run only basic unlabeled CI tests label Sep 18, 2026
@Rockdu Rockdu changed the title feat(sft): add train_sft.py, which generates the next rollout while the current one trains feat(sft): add train_sft.py for SFT data prefetch Sep 18, 2026
@Rockdu
Rockdu force-pushed the feat/sft-rollout-prefetch branch from c674d36 to e0a3321 Compare September 18, 2026 19:48
@Rockdu
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
Rockdu force-pushed the feat/sft-rollout-prefetch branch from e0a3321 to 23ed940 Compare September 18, 2026 21:53
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

run-ci-basic Authorize and run only basic unlabeled CI tests

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant