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

src/speculators/train/trainer.py

@@ -56,6 +56,20 @@ def _all_reduce_metrics(metrics: dict[str, torch.Tensor]) -> dict[str, torch.Ten
     return dict(zip(keys, stacked, strict=True))
 
 
+def _metrics_to_host(metrics: dict[str, torch.Tensor]) -> dict[str, float]:
+    """Read *metrics* back to Python floats with one device-to-host copy.
+
+    ``v.item()`` per entry is one synchronous copy each, and there are a few
+    dozen per logged step. Stacking first pays for a single transfer; the
+    values are 0-dim floats either way, so the results are identical.
+    """
+    if not metrics:
+        return {}
+    keys = list(metrics)
+    stacked = torch.stack([metrics[k].float().reshape(()) for k in keys])
+    return dict(zip(keys, stacked.tolist(), strict=True))
+
+
 class _StepTimer:
     # Each mark()/now() forces an accelerator.synchronize to capture true GPU time.
     # This serialises the CUDA pipeline, so profiled steps are slower; keep
@@ -581,7 +595,7 @@ def train_epoch(
                 if self.is_distributed:
                     metrics = _all_reduce_metrics(metrics)
 
-                metrics = {k: v.item() for k, v in metrics.items()}
+                metrics = _metrics_to_host(metrics)
                 world_size = dist.get_world_size() if self.is_distributed else 1
                 metrics = normalize_counted_metrics(metrics, world_size)
                 lr_info = (
@@ -666,7 +680,7 @@ def val_epoch(self, epoch: int) -> dict[str, float] | None:
         if accumulated:
             if self.is_distributed:
                 accumulated = _all_reduce_metrics(accumulated)
-            val_metrics = {k: v.item() for k, v in accumulated.items()}
+            val_metrics = _metrics_to_host(accumulated)
 
         world_size = dist.get_world_size() if self.is_distributed else 1
         val_metrics = {k: v / num_batches for k, v in val_metrics.items()}

tests/unit/train/test_metrics_readback.py

@@ -0,0 +1,48 @@
+"""Device-to-host readback of the per-step metric scalars."""
+
+from unittest.mock import patch
+
+import torch
+
+from speculators.train.trainer import _metrics_to_host
+
+
+def _metrics() -> dict[str, torch.Tensor]:
+    return {
+        "loss": torch.tensor(1.25),
+        "loss_total": torch.tensor(1.0),
+        "acc_0": torch.tensor(0.1),
+        "count": torch.tensor(7, dtype=torch.int64),
+        "halved": torch.tensor([0.5]),
+    }
+
+
+def test_matches_per_entry_item():
+    metrics = _metrics()
+
+    assert _metrics_to_host(metrics) == {
+        k: v.float().reshape(()).item() for k, v in metrics.items()
+    }
+
+
+def test_preserves_key_order():
+    metrics = _metrics()
+
+    assert list(_metrics_to_host(metrics)) == list(metrics)
+
+
+def test_empty_metrics_read_back_as_empty():
+    assert _metrics_to_host({}) == {}
+
+
+def test_reads_back_without_a_sync_per_entry():
+    """One transfer for the whole dict, not one per metric.
+
+    ``.item()`` blocks on the device each time it is called, and a logged step
+    carries a few dozen metrics; this guards the regression back to a
+    per-entry comprehension.
+    """
+    with patch.object(
+        torch.Tensor, "item", side_effect=AssertionError("per-entry sync")
+    ):
+        assert _metrics_to_host(_metrics())["loss"] == 1.25