Skip to content

fix(eval): release scored eval samples' generated media to free CPU mem - #84

Merged
Rockdu merged 2 commits into
radixark:mainfrom
Rockdu:fix/eval-release-sample-media
Jul 30, 2026
Merged

fix(eval): release scored eval samples' generated media to free CPU mem#84
Rockdu merged 2 commits into
radixark:mainfrom
Rockdu:fix/eval-release-sample-media

Conversation

@Rockdu

@Rockdu Rockdu commented Jul 29, 2026

Copy link
Copy Markdown
Collaborator

What

  • Release a scored eval sample's generated_output as soon as its reward lands, instead of holding every decoded video until the whole eval dataset finishes.
  • Media is kept for the samples that still read it: the first --diffusion-log-images samples of each dataset (exactly what _log_images uploads, since it takes samples[:max_images] of the index-sorted list) and every sample when --save-debug-rollout-data is set (that path dumps the tensors).

Why

eval_rollout_single_dataset accumulates all scored samples into data and returns them as {"samples": data}, but nothing downstream needs the pixels: _log_eval_rollout_data only reads rewards, and _log_images only uploads the first --diffusion-log-images of them. Meanwhile Sample.generated_output is a float32 [C, F, H, W] tensor (sgl-diffusion's _postprocess_output does output.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 (data dies 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.

Metric Measured
Eval completed 640/640 in 23:31 (2.21 s/it)
eval/pickscore_test 0.7991
generated_output per sample 58.06 MB (float32 [3, 21, 480, 480])
RolloutManager RSS, before eval 1.28 GB (flat)
RolloutManager RSS, during eval 1.70 / 2.66 / 3.80 GB (min / mean / max over 93 samples)
RSS trend during eval +0.32 MB/s (+0.31 GB over the 23.5 min)
RSS right after eval returns 1.97 GB
eval_media/pickscore_test_images in W&B 8 videos, captions carry rewards

Without 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 60 so the known busy-eval health-check false positive could not interrupt the measurement (0 Marking as DEAD this run). That is a run-level knob, not part of this PR.

Also here: raise the declared python floor to 3.12

stage-a-cpu is the one CI job that does not run inside the container image — it sets up Python 3.10 on ubuntu-latest, with a comment saying that matches setup.py (python_requires=">=3.10"). But miles/dashboard/events.py and store.py use enum.StrEnum, which is 3.11+, so collecting this PR's test — it imports miles.rollout.sglang_diffusion_rollout -> miles.dashboard.hooks -> events — failed with ImportError: 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's py_version, and the two workflow jobs that pin a bare-runner Python — instead of adding a backports.strenum dependency or shimming StrEnum in-tree. StrEnum is the only 3.11+ construct anywhere under miles/. stage-a-cpu and the pre-commit job both pass on 3.12, so requirements.txt installs cleanly there.

Files

  • miles/rollout/sglang_diffusion_rollout.pyrelease_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 with register_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_dataset in miles core has the same accumulate-everything shape, where the heavy fields are per-token rollout_log_probs / tokens, the rollout_routed_experts / rollout_indexer_topk arrays under --use-rollout-routing-replay, and deepcopied multimodal_inputs for 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-files passes
  • Added/updated tests for new behaviour
  • pytest -x is green — the new tests pass and pytest tests/fast is 137 passed, 1 pre-existing failure unrelated to this PR: tests/fast/dashboard/test_async_collector.py::test_collector_startup_failure_disables_dashboard monkeypatches miles.dashboard.backend._is_primary, which has never existed in that module (it fails on main too). It survived because that file never calls register_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 merged pytest tests/fast is 138 passed.
  • If launch flags changed, python3 train.py --help still parses — n/a, no flag changes
  • If a public flag was added, it appears in the CLI reference docs — n/a, no new flag
  • If an example was added, it has a real walkthrough — n/a, no new example

@Rockdu
Rockdu force-pushed the fix/eval-release-sample-media branch from 40b9d34 to 6d257b0 Compare July 30, 2026 11:00
@Rockdu
Rockdu marked this pull request as ready for review July 30, 2026 11:01
@Rockdu
Rockdu force-pushed the fix/eval-release-sample-media branch from 99cfdb0 to 14cb5d0 Compare July 30, 2026 19:42
@Rockdu
Rockdu merged commit 1c03d62 into radixark:main Jul 30, 2026
15 checks passed
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