feat(benchmark): bound the measured window and report effective throughput - #970
feat(benchmark): bound the measured window and report effective throughput#970fynnsu wants to merge 7 commits into
Conversation
|
Important Review skippedAuto reviews are disabled on base/target branches other than the default branch. 🗂️ Base branches to auto review (1)
Please check the settings in the CodeRabbit UI or the ⚙️ Run configurationConfiguration used: Path: .coderabbit.yaml Review profile: CHILL Plan: Pro Plus Run ID: You can disable this status message by setting the Use the checkbox below for a quick retry:
Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out. Comment |
Merge Protections🔴 1 of 1 protections blocking · waiting on 👀 reviews
🔴 Require approval from approved reviewers listWaiting for any of
This rule is failing.All pull requests must have at least one approving review from a member of the approved reviewers list before merging.
|
|
The quality checks have failed. Please run |
0081cd5 to
a8ebc1f
Compare
speculatorsbot
left a comment
speculatorsbot
left a comment
There was a problem hiding this comment.
Review Summary
This PR improves the benchmark harness with four well-motivated changes:
-
Strict sample-count validation (
select_measured_profiles): Replaces a silent fallback that could produce misleading results when the dataset was too small. Now raises a clearRuntimeErrorwith an actionable message. -
Time-weighted throughput (
compute_aggregate_throughput): The oldmean(tokens_per_s)equally weighted every step, overweighting short/fast steps. The new metric reconstructs total tokens and divides by total wall time -- a more accurate measure of effective throughput. -
Bounded sampler + worker shutdown (
max_batches,shutdown_dataloader_workers): Caps the DataLoader so persistent workers don't prefetch past the measurement window, and explicitly shuts workers down before distributed teardown to avoid racing Mooncake shutdown. -
Richer provenance: Records
data_path,hidden_states_backend,num_workers, andprefetch_factor-- parameters that affect results but were previously absent.
Tests are well-structured and cover the new functions thoroughly.
CI note: The quality-checks failures are pre-existing on the base branch (generation-recovery). The mypy errors are in unrelated files (test_data_recovery.py, test_data.py, data.py) -- none touched by this PR.
Minor observations (non-blocking)
1. Potential ZeroDivisionError in compute_aggregate_throughput -- If every profile has step_ms == 0 (or profiles is empty), elapsed_s will be zero and the final division raises. In practice this can't happen with real training steps, but a defensive guard or docstring precondition would make the contract clear.
2. shutdown_dataloader_workers relies on DataLoader._iterator -- _iterator is a private attribute of PyTorch's DataLoader. The getattr fallback makes it safe if the attribute disappears, but a brief comment noting the coupling would help future maintainers.
3. max_batches truncation is post-generation -- _generate_batches still computes all batches before truncating. For benchmarks this is fine, but worth noting it doesn't short-circuit the packing algorithm.
Overall the changes are well-scoped, well-tested, and clearly motivated by real pain points in the benchmark harness. Looks good.
| rank0_tokens = sum( | ||
| profile["tokens_per_s"] * profile["step_ms"] / 1000 for profile in profiles | ||
| ) | ||
| return { |
There was a problem hiding this comment.
nit: If profiles is empty, elapsed_s will be 0 and this line will raise ZeroDivisionError. Consider either guarding against empty input or documenting the precondition (e.g. profiles must be non-empty).
if not profiles:
return {"measured_time_s": 0, "rank0_tokens": 0, "effective_rank0_tokens_per_s": 0}| } | ||
|
|
||
|
|
||
| def shutdown_dataloader_workers(loader) -> None: |
There was a problem hiding this comment.
Minor: This accesses loader._iterator, a private CPython implementation detail of torch.utils.data.DataLoader. It works correctly today and the getattr fallback makes it safe if the attribute disappears, but a brief comment noting the coupling would help future readers.
| @@ -244,6 +249,8 @@ def _generate_batches(self, epoch: int) -> list[NDArray]: | |||
| # Translate them so that they are instead relative to the overall unshuffled | |||
| # self.lengths array. | |||
| batches = [indices[batch] for batch in batches] | |||
There was a problem hiding this comment.
Note: The truncation happens after _assign_to_packed_batches has already generated all batches. This means the full packing algorithm runs regardless of max_batches. For benchmark use cases this is fine (the overhead is negligible), but worth being aware of if this parameter is ever used with very large datasets.
a8ebc1f to
07a596a
Compare
07a596a to
9699e6c
Compare
…ifest The completion marker was a bare JSON list of tensor names, so a consumer had no way to tell a correct payload from a truncated or corrupted one, and Mooncake's negative status codes were discarded. A failed put could therefore still be followed by the marker, publishing a sample whose tensors were missing or partially written. Write a versioned manifest recording shape, dtype, and CRC32 for every tensor and verify all three on read, raising MooncakeIntegrityError on mismatch. Check store return codes, remove partially written keys instead of leaving them behind, and reject non-finite tensors at the producer. When a producer fails, publish a small terminal error manifest so the consumer fails fast rather than waiting out its poll timeout. Legacy list manifests are still accepted so in-flight handles keep working. Also correct the install hint: the package is mooncake-transfer-engine-cuda13, not mooncake-transfer-engine-cuda-13. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com> Signed-off-by: Fynn Schmitt-Ulms <fschmitt@redhat.com>
Signed-off-by: Fynn Schmitt-Ulms <fschmitt@redhat.com>
…ad sizing MooncakeStoreConfig already carried global_segment_size, local_buffer_size, and num_writer_threads, but the vLLM-side CLI had no way to set them, so every deployment was pinned to the dataclass defaults regardless of sequence length or concurrency. Add --mooncake-global-segment-gib, --mooncake-local-buffer-gib, and --mooncake-writer-threads, and pass them through both client constructors. Sizes are taken in GiB because the underlying fields are byte counts that are awkward to write on a command line. Lower the default writer-thread count from 16 to 4. Sixteen writers oversubscribe the transfer engine when several clients share a node without improving throughput. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com> Signed-off-by: Fynn Schmitt-Ulms <fschmitt@redhat.com>
Signed-off-by: Fynn Schmitt-Ulms <fschmitt@redhat.com>
A failed generate/load/validate round trip returned None, the collator silently dropped the sample, and a rank whose whole batch failed produced an empty batch. Under DDP that rank skipped the step's collectives while its peers waited, so an isolated data fault stalled the job instead of degrading it. Nothing bounded a persistently failing worker either: it retried forever against a broken endpoint. Retry the complete round trip (generate, load, validate) rather than only the HTTP request, controlled by --generation-validation-retries, and always delete the handle of a failed attempt so a corrupt payload is not left in the store. Replace the None sentinel with a picklable GenerationFailure the collator can count, and trip a circuit breaker after --max-consecutive-generation-failures consecutive failures in one worker. A rank that loses every sample now builds a padded batch with an empty loss mask and runs the normal forward/backward, contributing zero gradients while joining every collective, so peers are unaffected. Only a tripped breaker is all-reduced, so all ranks stop at the same step rather than one vanishing first. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com> Signed-off-by: Fynn Schmitt-Ulms <fschmitt@redhat.com>
Signed-off-by: Fynn Schmitt-Ulms <fschmitt@redhat.com>
…ghput Three problems made benchmark numbers hard to trust. A short dataset silently produced fewer profiles than requested and the harness fell back to measuring warmup steps instead of failing. Throughput was reported as mean(tokens_per_s), which weights every step equally and so overweights short fast steps. And persistent DataLoader workers kept prefetching past the measured window, so generation work outside the measurement competed with it and teardown raced Mooncake shutdown. Replace the silent fallback with select_measured_profiles(), which raises when the dataset cannot supply warmup + measured steps. Add compute_aggregate_throughput(), reconstructing token counts per profile and dividing by total elapsed time for the window rank 0 actually observed; report it in the summary and in compare output. Cap the sampler with max_batches so a bounded benchmark stops generating past its last measured step, and shut workers down before distributed teardown. Benchmarks also set raise_on_generate_error so a generation fault aborts the run instead of quietly measuring degraded batches, and record data_path, hidden_states_backend, num_workers, and prefetch_factor, which change results but were absent from provenance. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com> Signed-off-by: Fynn Schmitt-Ulms <fschmitt@redhat.com>
9699e6c to
c954dbf
Compare
|
This pull request has merge conflicts that must be resolved before it can be |
2 participants
Purpose
Improve benchmarking scripts to compute metrics (like throughput) more accurately, and with greater control over the benchmarking setup (e.g. number of samples).
Tests
Used to tune training performance for a large distributed model run.
Checklist
I have filled in:
Stack created with GitHub Stacks CLI • Give Feedback 💬