perf(train): read the step metrics back with one device-to-host copy by orestis-z · Pull Request #1298 · vllm-project/speculators

orestis-z · 2026-10-21T17:19:11Z

Purpose

{k: v.item() for k, v in metrics.items()} blocks on the device once per entry, and a logged step carries a few dozen of them — the loss, the loss total, and an acceptance/accuracy pair per ttt_step. With the default log_freq of 1 that is a few dozen synchronisations every step. This stacks the scalars into one tensor and reads it back with a single tolist(); the values are 0-dim floats either way, so only the number of transfers changes. The validation reduction uses the same helper.

Measured

Single H100, eagle3, synthetic batches, full 151936-token vocab, total_seq_len 2048, 10 warmup + 50 measured steps, 3 reps, median of the per-run means, log_freq=1, scripts/benchmark.py. This PR against #1297:

Qwen3-0.6B	Qwen3-1.7B	Qwen3-8B	Qwen3-32B
0.980x	1.004x	0.997x	0.999x

This is inside the run-to-run spread — it is not a measurable win on this benchmark at any verifier size. The likely reason is that the first .item() already pays for draining the queue, so removing the other ~19 saves little, but that is an explanation after the fact rather than something measured. It is still one sync per step instead of a few dozen, which is what it would cost on a run with more ttt_steps or a slower link, and it puts the readback in one place.

On this evidence the recommendation is to drop it; nothing else depends on it.

Tests

pytest tests/unit/train/test_metrics_readback.py

New file. Covers that the helper returns the same floats as the per-entry .item() path, that it does not call Tensor.item at all (patched to raise), that it handles an empty dict, and that it keeps non-tensor values. Passes.

Checklist

I have filled in:

The purpose of the PR, such as "Fix some issue (link existing issues this PR will resolve)".
The test plan/results, such as providing test command and pasting the results.
(Optional) The necessary documentation update.
I (a human) have written or reviewed the code in this pr to the best of my ability.

🤖 Generated with Claude Code

``{k: v.item() for k, v in metrics.items()}`` blocks on the device once per entry, and a logged step carries a few dozen of them: the loss, the loss total, and an acceptance/accuracy pair per ttt_step. With the default ``log_freq`` of 1 that is a few dozen synchronisations every step, each one draining whatever the step still had queued. Stack the scalars into one tensor and read that back with a single ``tolist()``. The values are 0-dim floats either way, so the result is identical; only the number of transfers changes. The validation reduction reads its accumulated metrics back the same way, so it uses the same helper -- it runs once per epoch rather than per step, but there is no reason to keep two spellings of it. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com> Signed-off-by: Orestis Zambounis <orestis.zambounis@gmail.com>

coderabbitai · 2026-10-21T17:19:20Z

Important

Draft PR not reviewed

Draft PRs are not automatically reviewed by default.

Trigger a manual review

To automatically review draft PRs, update your CodeRabbit configuration:

reviews:
  auto_review:
    drafts: true

Autofix · Keep fixing CodeRabbit findings and required CI, and resolving merge conflicts

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.

❤️ Share

_{Comment @coderabbitai help to get the list of available commands.}

mergify · 2026-10-21T17:19:47Z

Merge Protections

🔴 1 of 1 protections blocking · waiting on 👀 reviews

	Protection	Waiting on
🔴	Require approval from approved reviewers list	👀 reviews

🔴 Require approval from approved reviewers list

Waiting for any of

approved-reviews-by = dsikka
approved-reviews-by = fynnsu
approved-reviews-by = orestis-z
approved-reviews-by = rahul-tuli
approved-reviews-by = shanjiaz

This rule is failing.

All pull requests must have at least one approving review from a member of the approved reviewers list before merging.

## Purpose Split out of the measurements in #1259 / #1262 / #1263, one change per PR. Every training step built its metric scalars with `torch.tensor(<python float>, device=cuda)`, which copies from pageable host memory and so cannot be issued asynchronously — the host blocks until the work already queued on the device drains. The same step ran `torch.masked_select` on the accuracy mask, whose output size depends on the mask's values, forcing another device-to-host sync before the count could be read. Building those scalars with `torch.full`/`torch.zeros`/`torch.ones` and counting masked positions with a reduction gives identical numbers without either sync. `full_total` and `cond_total` stay separate tensors so a caller reducing the metrics dict in place cannot count one of them twice. ### Measured Single H100, eagle3, synthetic batches, full 151936-token vocab, `total_seq_len` 2048, 10 warmup + 50 measured steps, 3 reps, median of the per-run means, `log_freq=1` (the default), `scripts/benchmark.py`. This PR against main: | Qwen3-0.6B | Qwen3-1.7B | Qwen3-8B | Qwen3-32B | |---|---|---|---| | **1.111x** | **1.110x** | **1.030x** | **1.013x** | <details> <summary>Where the gain comes from, change by change</summary> Same harness and settings, each change measured against its parent commit, so every cell is one PR's marginal contribution at that verifier size. Run-to-run spread within a cell is under 1% everywhere except main at 0.6B. | PR | | Qwen3-0.6B | Qwen3-1.7B | Qwen3-8B | Qwen3-32B | |---|---|---|---|---|---| | #1294 | metrics off host | **1.111x** | **1.110x** | **1.030x** | **1.013x** | | #1295 | reuse the argmax | **1.076x** | **1.026x** | **1.013x** | 1.007x | | #1296 | single-pass fused KL | **1.047x** | **1.017x** | **1.012x** | 1.003x | | #1297 | logs off the thread | **1.029x** | **1.042x** | **1.015x** | 1.006x | | #1298 | one metric readback | 0.980x | 1.004x | 0.997x | 0.999x | | #1299 | H2D prefetch † | 1.024x | 1.013x | 1.029x | 1.021x | | #1300 | skip dropped positions | 0.997x | 0.988x | 0.980x | 0.988x | | #1301 | precast frozen heads | **1.055x** | **1.096x** | **1.059x** | **1.036x** | | | **cumulative** | **1.358x** | **1.329x** | **1.141x** | **1.077x** | † #1299's row is not the prefetch. That commit also pins the synthetic batch in `scripts/benchmark.py`, and the real `DataLoader` already sets `pin_memory=True`, so it is a benchmark-fidelity fix rather than a training speedup. Applying only the pinning hunk on top of #1298 gives 1.017x / 1.017x / 1.030x / 1.022x — the whole row. Discounting it, the cumulative row is **1.335x / 1.307x / 1.108x / 1.054x**. Reading it: - **#1294 and #1301 carry most of the gain.** They are the two largest at every size, and #1301 is the only one that also moves peak memory (−1.5 GB at 0.6B to −7.4 GB at 32B). - **#1295, #1296 and #1297 fade with verifier size** but stay positive — 1–2% at 8B, 0.3–0.7% at 32B. - **#1298 is a tie at every size.** One sync per step instead of a few dozen is the right shape, but it does not show up here. - **#1299's prefetch is inside noise** in this regime once the harness fix is separated out, though it does pay once the metric log emit is deferred — see #1299. - **#1300 is negative at `seq_len` 2048 at every size** and costs 0.3–1.5 GB of peak; it needs both long sequences and a small hidden size to pay. Only #1300 and #1301 change peak allocated memory at all; the other six are flat to the megabyte. Most of these changes remove fixed per-step host and metric overhead, so the relative win shrinks as the verifier grows even though the absolute saving does not. At 8B and 32B the step is dominated by the optimizer (42% and 58%, all Muon Newton–Schulz), which nothing here touches — the same point #1111's side note made. </details> ## Tests ```bash pytest tests/unit/models/test_metrics.py ``` New cases cover that the metric tensors are built without a host-to-device copy and that the masked counts match `masked_select`. Passes. ## Checklist I have filled in: - [x] The purpose of the PR, such as "Fix some issue (link existing issues this PR will resolve)". - [x] The test plan/results, such as providing test command and pasting the results. - [ ] (Optional) The necessary documentation update. - [x] I (a human) have written or reviewed the code in this pr to the best of my ability. 🤖 Generated with [Claude Code](https://claude.com/claude-code) --------- Signed-off-by: Orestis Zambounis <orestis.zambounis@gmail.com> Signed-off-by: Orestis Zambounis <23146389+orestis-z@users.noreply.github.com> Co-authored-by: Claude Opus 5 <noreply@anthropic.com>

orestis-z added this pull request to stack #1302 October 21, 2026 17:19

This was referenced Oct 21, 2026

perf(train): keep metric construction off the host #1294

Merged

perf(train): copy the next batch to the device on a side stream #1299

Closed

orestis-z closed this Oct 22, 2026

orestis-z removed this pull request from stack #1302 October 22, 2026 14:32

Navigation Menu

Repository navigation

perf(train): read the step metrics back with one device-to-host copy - #1298

perf(train): read the step metrics back with one device-to-host copy#1298
orestis-z wants to merge 1 commit into
perf/log-off-training-threadfrom
perf/stacked-metric-readback

orestis-z commented Oct 21, 2026 •

edited

Loading

Uh oh!

coderabbitai Bot commented Oct 21, 2026

Draft PR not reviewed

Uh oh!

mergify Bot commented Oct 21, 2026

Uh oh!

Reviewers

Uh oh!

Assignees

Labels

Projects

Milestone

Development

Uh oh!

1 participant

Conversation

orestis-z commented Oct 21, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Purpose

Measured

Tests

Checklist

Uh oh!

coderabbitai Bot commented Oct 21, 2026

Draft PR not reviewed

Uh oh!

mergify Bot commented Oct 21, 2026

Merge Protections

🔴 Require approval from approved reviewers list

Uh oh!

Reviewers

Uh oh!

Assignees

Labels

Projects

Milestone

Development

Uh oh!

1 participant

orestis-z commented Oct 21, 2026 •

edited

Loading