feat(closed_loop): shard evaluation across all DDP ranks - #337
Draft
HansRobo wants to merge 1 commit into
Draft
Conversation
This was referenced Jul 30, 2026
HansRobo
marked this pull request as draft
July 31, 2026 00:08
HansRobo
force-pushed
the
ky/closed-loop-ddp-fail-fast
branch
from
July 31, 2026 00:15
3d446f4 to
05aab98
Compare
HansRobo
force-pushed
the
ky/closed-loop-pooled-ddp
branch
2 times, most recently
from
July 31, 2026 00:54
1d62ae2 to
acc0be9
Compare
closed_loop_validate ran inside `if global_rank == 0:`, so one GPU evaluated while the rest idled at the next collective -- 48 minutes per checkpoint on a 4-GPU job, 102 on the final save. The existing machinery barriers once per (site, object-mode), which caps at 1.96x on 4 GPUs because a one-route site cannot be split. run_evaluations_ddp() pools all (combo, route) pairs instead, splits once with cost-weighted LPT, runs with no intermediate barriers, then barriers once and merges per combo. 4xH100: media off 2862 -> 1023 s (2.80x), media on 6121 -> 2736 s (2.24x). Summed per-rank time matches the sequential baselines to 0.04% / 0.1%. Output identical apart from order-dependent reductions. 25 tests, where there were 0.
HansRobo
force-pushed
the
ky/closed-loop-ddp-fail-fast
branch
from
July 31, 2026 07:31
05aab98 to
25a2afb
Compare
HansRobo
force-pushed
the
ky/closed-loop-pooled-ddp
branch
from
July 31, 2026 07:31
acc0be9 to
00a1077
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.
Problem
closed_loop_validateruns insideif global_rank == 0:, so one GPU does the evaluation whilethe rest idle at the next collective. On a 4-GPU job that is 48 minutes of 3 idle GPUs per
checkpoint (102 minutes on the final save, which also renders media).
The DDP machinery for this already exists (
shard_items,run_distributed,merge_ddp_shards)but has no caller. It also barriers once per (site, object-mode), so a site with one route runs at
single-GPU speed regardless of GPU count — measured ceiling 1.96x on 4 GPUs.
Fix
run_evaluations_ddp()pools all (combo, route) pairs into one list, splits it once withcost-weighted LPT, runs with no intermediate barriers, then barriers once and merges per combo.
closed_loop_validatemoves out of the rank-0 guard.Result
4xH100, real manifest, 38 segments:
Summed per-rank rollout time matches the sequential baselines to 0.04% / 0.1% — same workload,
negligible contention.
Output is identical. All 38 mp4s match, including 14 routes written concurrently
into one directory.
For reviewers
elapsed_secis now all-reduced per combo, preserving its per-site meaning. This touches anexisting per-site W&B metric deliberately.
run_distributed()now has zero callers (it had none before either). Deleting it is your call.within a site, poorly across sites (0.026-0.079 s/frame). Feeding measured per-route seconds
back would reach 4.00x; that needs rank-0 to plan and broadcast, so it is left as follow-up.
Tests
25 tests, where there were 0. The multi-rank ones spawn real gloo groups on CPU and drive the
production merge path. Failure propagation is mutation-verified.