Skip to content

perf(rollout): six independent rollout-path fixes and knobs - #198

Open
Rockdu wants to merge 6 commits into
radixark:mainfrom
Rockdu:h3-rollout-perf
Open

perf(rollout): six independent rollout-path fixes and knobs#198
Rockdu wants to merge 6 commits into
radixark:mainfrom
Rockdu:h3-rollout-perf

Conversation

@Rockdu

@Rockdu Rockdu commented Aug 21, 2026

Copy link
Copy Markdown
Collaborator

What

Six independent rollout-path improvements, one commit each:

Commit Change
perf(rm_hub) submit reward chunks to Ray inside the worker thread — .remote() serialises the payload into plasma in the calling thread, so submitting on the event loop stalled every concurrent request
perf(rm_hub) select the reward frames before converting a clip to uint8 (8 of 107 frames survive; converting all first cost several full-size float32 copies)
fix(router) register /remove_worker; engines call it on shutdown and it fell through the catch-all, proxied to a worker as a generate request
fix(fsdp) report checkpoint key mismatches — the LoRA path loads with strict=False and discarded set_state_dict's result, so a checkpoint whose keys match nothing "loaded" silently
perf(rollout) --rollout-request-stagger: one-shot jitter before slot acquisition; co-admitted requests stay in phase for the whole rollout (measured: median inter-arrival gap 0.00 s, 16 arrivals in the first 5% of each rollout)
feat(rollout) --rollout-video-dtype {keep,uint8}: opt-in engine-side quantisation of the rollout video, default unchanged; needs the sglang counterpart to take effect and is harmless without it

Validation

  • Frame selection is bit-identical to the old path on unit-scale, byte-scale and bf16 inputs, 9-11x faster (measured on H200).
  • Full stack ran a 17-GPU H3 t2va run resumed from a LoRA checkpoint: log_prob_mean_abs_diff 3.0e-05 (in the verified band), 0 faults; with uint8 on, resp_bytes 1384 -> 436 MiB and train_wait 1010-1140 s -> 486-499 s.
  • pre-commit run clean on all touched files; pytest tests/fast: 215 passed.

Checklist

  • pre-commit run --all-files passes — run on touched files
  • Added/updated tests for new behaviour — fixture updated for the new arg; no new tests
  • pytest -x is green — tests/fast, 215 passed
  • If launch flags changed, python3 train.py --help still parses
  • If a public flag was added, it appears in the CLI reference docs — not updated
  • If an example was added, it has a real walkthrough — n/a

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

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant