fix(eval): release scored eval samples' generated media to free CPU mem - #84
Merged
Merged
Conversation
Rockdu
force-pushed
the
fix/eval-release-sample-media
branch
from
July 30, 2026 11:00
40b9d34 to
6d257b0
Compare
Rockdu
marked this pull request as ready for review
July 30, 2026 11:01
6 tasks
Rockdu
force-pushed
the
fix/eval-release-sample-media
branch
from
July 30, 2026 19:42
99cfdb0 to
14cb5d0
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
generated_outputas soon as its reward lands, instead of holding every decoded video until the whole eval dataset finishes.--diffusion-log-imagessamples of each dataset (exactly what_log_imagesuploads, since it takessamples[:max_images]of the index-sorted list) and every sample when--save-debug-rollout-datais set (that path dumps the tensors).Why
eval_rollout_single_datasetaccumulates all scored samples intodataand returns them as{"samples": data}, but nothing downstream needs the pixels:_log_eval_rollout_dataonly reads rewards, and_log_imagesonly uploads the first--diffusion-log-imagesof them. MeanwhileSample.generated_outputis a float32[C, F, H, W]tensor (sgl-diffusion's_postprocess_outputdoesoutput.cpu().float()), so the entire eval set stays resident in the RolloutManager actor's heap.For the wan2.2-T2V recipe (480x480, 21 frames) that is 58.06 MB per sample: a 2048-prompt eval set pins ~119 GB of host RAM in one actor, and the cost scales with resolution x frames x prompts. It is a per-eval peak rather than unbounded growth (
datadies with the call), but it is large enough to OOM any host that is not a 2 TB node.Validation
16-GPU wan2.2-T2V-A14B LoRA GRPO (2 x 8 H200 + 1 reward GPU, dp_replicate=2 x sp=8, rollout tp=4, colocate), step-0 eval over a 640-prompt subset of the flowgrpo_pickscore test set,
--diffusion-eval-num-steps 28 --diffusion-output-num-frames 21 --diffusion-height 480 --diffusion-width 480 --diffusion-log-images 8. RSS sampled every 15 s from/proc/<RolloutManager pid>/status.eval/pickscore_testgenerated_outputper sample[3, 21, 480, 480])eval_media/pickscore_test_imagesin W&BWithout the release the same eval would have accumulated 640 x 58.06 MB = 36.3 GB of media, i.e. a +26.3 MB/s trend; measured trend is 1.2% of that and shows no upward slope, so the media is being freed as each sample is scored rather than at the end. The 8 W&B video panels confirm the image-logging carve-out still holds. RSS climbing again after eval is the train rollout, which retains DiT trajectories by design.
The probe run also passed
--miles-router-health-check-failure-threshold 40 --rollout-health-check-interval 60so the known busy-eval health-check false positive could not interrupt the measurement (0Marking as DEADthis run). That is a run-level knob, not part of this PR.Also here: raise the declared python floor to 3.12
stage-a-cpuis the one CI job that does not run inside the container image — it sets up Python 3.10 onubuntu-latest, with a comment saying that matchessetup.py(python_requires=">=3.10"). Butmiles/dashboard/events.pyandstore.pyuseenum.StrEnum, which is 3.11+, so collecting this PR's test — it importsmiles.rollout.sglang_diffusion_rollout->miles.dashboard.hooks->events— failed withImportError: cannot import name 'StrEnum'. The declaration and the code have disagreed since the dashboard landed; nothing caught it because no CI-collected test imported that chain.The dev/CI container image is Python 3.12 (Ubuntu 24.04) and every GPU stage runs inside it, so the second commit makes the declaration match reality —
python_requires, the classifiers, isort'spy_version, and the two workflow jobs that pin a bare-runner Python — instead of adding abackports.strenumdependency or shimmingStrEnumin-tree.StrEnumis the only 3.11+ construct anywhere undermiles/.stage-a-cpuand the pre-commit job both pass on 3.12, sorequirements.txtinstalls cleanly there.Files
miles/rollout/sglang_diffusion_rollout.py—release_eval_sample_media()plus its call in the eval completion loop.setup.py,pyproject.toml,.github/workflows/_run-ci.yml,.github/workflows/pre-commit.yml— python floor raised to 3.12.tests/fast/rollout/test_release_eval_sample_media.py— covers release, the image-logging carve-out, and the debug-dump carve-out. Registered withregister_cpu_ci(est_time=15, suite="stage-a-cpu")so the suite runner actually collects it.Not in this PR
miles/rollout/sglang_rollout.py:eval_rollout_single_datasetin miles core has the same accumulate-everything shape, where the heavy fields are per-tokenrollout_log_probs/tokens, therollout_routed_experts/rollout_indexer_topkarrays under--use-rollout-routing-replay, and deepcopiedmultimodal_inputsfor VLM. Worth a matching change there, but its consumers read light fields (response_length,status), so it needs its own patch.Checklist
pre-commit run --all-filespassespytest -xis green — the new tests pass andpytest tests/fastis 137 passed, 1 pre-existing failure unrelated to this PR:tests/fast/dashboard/test_async_collector.py::test_collector_startup_failure_disables_dashboardmonkeypatchesmiles.dashboard.backend._is_primary, which has never existed in that module (it fails onmaintoo). It survived because that file never callsregister_cpu_ci, so the suite runner does not collect it — dropping the stale line makes all 8 dashboard tests pass. Fixed separately in test: collect the unregistered fast CPU tests and update cpu ci python to 3.12 #85; with that mergedpytest tests/fastis 138 passed.python3 train.py --helpstill parses — n/a, no flag changes