fix: metric deferral + optimizer optimization

torchspec/controller/loop.py

@@ -39,6 +39,18 @@
 from torchspec.utils.logging import get_tb_writer, logger
 
 
+def _write_training_metrics(metrics: dict, train_step: int, inference_step: int) -> None:
+    if getattr(wandb, "run", None) is not None:
+        wandb.log(metrics)
+
+    tb_writer = get_tb_writer()
+    if tb_writer is not None:
+        for key, value in metrics.items():
+            if isinstance(value, (int, float)):
+                scalar_step = inference_step if key.startswith("inference/") else train_step
+                tb_writer.add_scalar(key, value, scalar_step)
+
+
 def _maybe_sync_draft_weights(args, completed_steps, train_group, inference_engines):
     """Sync draft model weights to inference engines (decode mode only)."""
     weight_sync_enabled = getattr(args, "decode_weight_sync_enabled", False)
@@ -238,6 +250,7 @@ def training_loop(
     consecutive_failures = 0
     queued_batches = 0
     last_saved_step: int | None = None
+    previous_dispatch_wait: float | None = None
     progress = tqdm(total=num_steps, desc="Training", unit="step", initial=start_step)
     while completed_steps < num_steps:
         remaining_steps = min(
@@ -315,6 +328,8 @@ def training_loop(
             # The current optimizer step is fully queued.
             if enable_perf:
                 dispatch_wait = time.time() - t_dispatch
+            else:
+                dispatch_wait = 0.0
 
             train_futures = [
                 actor.train_from_queue.remote(
@@ -330,44 +345,44 @@ def training_loop(
 
             # Log metrics from training (use rank 0's metrics - they're already all-reduced)
             metrics = train_results[0] if train_results and train_results[0] else {}
+            metric_step = int(metrics.pop("_metrics_step", completed_steps))
+            metric_dispatch_wait = (
+                previous_dispatch_wait
+                if metric_step != completed_steps and previous_dispatch_wait is not None
+                else dispatch_wait
+            )
+            previous_dispatch_wait = dispatch_wait
             if metrics:
                 # Add step counters for wandb x-axis (required in shared mode)
-                metrics["train/step"] = completed_steps
+                metrics["train/step"] = metric_step
                 metrics["inference/step"] = completed_steps
 
                 # Add inference metrics (e2e_latency, spec metrics, etc.)
                 inference_metrics = ray.get(inference_manager.flush_metrics.remote())
                 metrics.update(inference_metrics)
 
                 if enable_perf:
-                    metrics["perf/dispatch_wait"] = dispatch_wait
+                    metrics["perf/dispatch_wait"] = metric_dispatch_wait
                     step_time = metrics.get("perf/step_time", 0)
                     if step_time > 0:
                         metrics["perf/train_capacity"] = args.global_batch_size / step_time
 
-                    if completed_steps % 5 == 0 or completed_steps <= 5:
+                    if metric_step % 5 == 0 or metric_step <= 5:
                         data_t = metrics.get("perf/data_time", 0)
                         compute_t = metrics.get("perf/compute_time", 0)
                         fwd_t = metrics.get("perf/forward_time", 0)
                         bwd_t = metrics.get("perf/backward_time", 0)
                         opt_t = metrics.get("perf/optimizer_time", 0)
                         logger.info(
-                            f"TIMING step={completed_steps}: "
+                            f"TIMING step={metric_step}: "
                             f"step={step_time:.3f}s "
                             f"data={data_t:.3f}s "
                             f"compute={compute_t:.3f}s "
                             f"[fwd={fwd_t:.3f}s bwd={bwd_t:.3f}s opt={opt_t:.3f}s] "
-                            f"dispatch={dispatch_wait:.3f}s"
+                            f"dispatch={metric_dispatch_wait:.3f}s"
                         )
 
-                if getattr(wandb, "run", None) is not None:
-                    wandb.log(metrics)
-
-                tb_writer = get_tb_writer()
-                if tb_writer is not None:
-                    for key, value in metrics.items():
-                        if isinstance(value, (int, float)):
-                            tb_writer.add_scalar(key, value, completed_steps)
+                _write_training_metrics(metrics, metric_step, completed_steps)
 
             # ── Eval at explicit interval (if configured) ─────────
             # Skip if a checkpoint save is about to run (it will eval anyway)
@@ -448,6 +463,21 @@ def training_loop(
         # Inner while broke (max steps reached during reload), break outer loop
         break
 
+    final_train_results = train_group.flush_pending_train_metrics()
+    if not isinstance(final_train_results, (list, tuple)):
+        final_train_results = []
+    final_metrics = final_train_results[0] if final_train_results and final_train_results[0] else {}
+    if final_metrics:
+        final_metric_step = int(final_metrics.pop("_metrics_step", completed_steps))
+        final_metrics["train/step"] = final_metric_step
+        final_metrics["inference/step"] = completed_steps
+        if enable_perf and previous_dispatch_wait is not None:
+            final_metrics["perf/dispatch_wait"] = previous_dispatch_wait
+            step_time = final_metrics.get("perf/step_time", 0)
+            if step_time > 0:
+                final_metrics["perf/train_capacity"] = args.global_batch_size / step_time
+        _write_training_metrics(final_metrics, final_metric_step, completed_steps)
+
     progress.close()
 
     # Always save a final checkpoint unless saved.

torchspec/models/ops/flex_attention.py

@@ -18,6 +18,8 @@
 # OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
 # SOFTWARE.
 
+from functools import lru_cache
+
 import torch
 import torch._dynamo as dynamo
 import torch._inductor.config as inductor_config
@@ -342,7 +344,7 @@ def mask_mod(b, h, q, kv):
     )
 
 
-def eagle3_block_mask(
+def _make_eagle3_block_mask(
     Q_LEN: int,
     KV_LEN: int,
     *,
@@ -404,3 +406,63 @@ def eagle3_block_mask(
         device=device,
         BLOCK_SIZE=(Q_BS, KV_BS),
     )
+
+
+@lru_cache(maxsize=128)
+def _cached_eagle3_block_mask(
+    Q_LEN: int,
+    KV_LEN: int,
+    B: int,
+    H: int,
+    device: str,
+    Q_BS: int,
+    KV_BS: int,
+    lck: int,
+) -> "BlockMask":
+    return _make_eagle3_block_mask(
+        Q_LEN=Q_LEN,
+        KV_LEN=KV_LEN,
+        B=B,
+        H=H,
+        device=torch.device(device),
+        BLOCK_SIZE=(Q_BS, KV_BS),
+        lck=lck,
+    )
+
+
+def eagle3_block_mask(
+    Q_LEN: int,
+    KV_LEN: int,
+    *,
+    B: int = 1,
+    H: int = 1,
+    device: torch.device = "cuda",
+    BLOCK_SIZE: "int | tuple[int, int]" = 128,
+    lck: int = 0,
+) -> "BlockMask":
+    """Return a cached, read-only Eagle3 block mask for a sequence shape."""
+    Q_BS, KV_BS = _normalize_block_size(BLOCK_SIZE)
+    if is_torchdynamo_compiling():
+        return _make_eagle3_block_mask(
+            Q_LEN=Q_LEN,
+            KV_LEN=KV_LEN,
+            B=B,
+            H=H,
+            device=device,
+            BLOCK_SIZE=(Q_BS, KV_BS),
+            lck=lck,
+        )
+
+    resolved_device = torch.device(device)
+    if resolved_device.type == "cuda" and resolved_device.index is None:
+        resolved_device = torch.device("cuda", torch.cuda.current_device())
+    return _cached_eagle3_block_mask(
+        Q_LEN,
+        KV_LEN,
+        B,
+        H,
+        str(resolved_device),
+        Q_BS,
+        KV_BS,
+        lck,
+    )

torchspec/ray/train_group.py

@@ -135,6 +135,12 @@ def train_from_queue(self, step: int, num_batches: int):
             [actor.train_from_queue.remote(step, num_batches) for actor in self._actor_handlers]
         )
 
+    def flush_pending_train_metrics(self):
+        """Materialize the final one-step-delayed metrics on every trainer."""
+        return ray.get(
+            [actor.flush_pending_train_metrics.remote() for actor in self._actor_handlers]
+        )
+
     def save_model(self, step, force_sync=False):
         """Save training model"""
         return ray.get(

torchspec/training/eagle3_trainer.py

@@ -29,7 +29,7 @@
 from torchspec.training import checkpoint
 from torchspec.training.fsdp import apply_fsdp2, fsdp2_load_full_state_dict
 from torchspec.training.optimizer import BF16Optimizer
-from torchspec.training.trainer import Trainer
+from torchspec.training.trainer import DeferredTrainMetrics, Trainer
 from torchspec.utils.distributed import get_gloo_group
 from torchspec.utils.logging import logger
 from torchspec.utils.tensor import padding
@@ -63,6 +63,7 @@ def __init__(self, args: Namespace):
             _position_decay_weights(args.ttt_length, getattr(args, "ploss_weights", None))
         )
         self._ploss_weight_sum = sum(self._ploss_weights)
+        self._metric_buffers: list[torch.Tensor] = []
 
     def init_model(
         self,
@@ -458,7 +459,7 @@ def _save_dump_data(
 
     def _aggregate_metrics(
         self, all_step_metrics: list[dict], step: int, *, grad_norm: torch.Tensor = None
-    ) -> dict:
+    ) -> dict | DeferredTrainMetrics:
         if not all_step_metrics:
             return {}
 
@@ -474,14 +475,60 @@ def _aggregate_metrics(
             if grad_norm is not None
             else avg_vlosses.new_zeros(())
         )
-        packed_metrics = torch.cat(
-            (
-                reduced_metrics,
-                grad_norm_value.reshape(1),
+        packed_metrics = (
+            torch.cat(
+                (
+                    reduced_metrics,
+                    grad_norm_value.reshape(1),
+                )
+            )
+            .detach()
+            .float()
+        )
+        metric_step = self.global_step
+        learning_rate = self.optimizer.get_learning_rate()
+
+        def finalize_values(metric_values: list[float]) -> dict:
+            return self._build_train_metrics(
+                metric_values,
+                num_depths=num_depths,
+                metric_step=metric_step,
+                learning_rate=learning_rate,
             )
-        ).detach()
-        metric_values = packed_metrics.float().cpu().tolist()
 
+        if packed_metrics.device.type != "cuda":
+            return finalize_values(packed_metrics.tolist())
+
+        if not self._metric_buffers:
+            self._metric_buffers = [
+                torch.empty(
+                    packed_metrics.numel(),
+                    dtype=torch.float32,
+                    device="cpu",
+                    pin_memory=True,
+                )
+                for _ in range(2)
+            ]
+        cpu_values = self._metric_buffers[metric_step % len(self._metric_buffers)]
+        cpu_values.copy_(packed_metrics, non_blocking=True)
+        ready_event = torch.cuda.Event()
+        ready_event.record()
+        return DeferredTrainMetrics(
+            cpu_values=cpu_values,
+            source_values=packed_metrics,
+            ready_event=ready_event,
+            metric_step=metric_step,
+            finalize_values=finalize_values,
+        )
+
+    def _build_train_metrics(
+        self,
+        metric_values: list[float],
+        *,
+        num_depths: int,
+        metric_step: int,
+        learning_rate: float,
+    ) -> dict:
         ploss_values = metric_values[:num_depths]
         acc_values = metric_values[num_depths : 2 * num_depths]
         grad_norm_scalar = metric_values[-1]
@@ -502,16 +549,16 @@ def _aggregate_metrics(
             "train/avg_acc": avg_acc_scalar,
             "train/simulated_acc_len": simulated_acc_len_value,
             "train/grad_norm": grad_norm_scalar,
-            "train/global_step": self.global_step,
-            "train/lr": self.optimizer.get_learning_rate(),
-            "train/step": step,
+            "train/global_step": metric_step,
+            "train/lr": learning_rate,
+            "train/step": metric_step,
         }
 
         for i in range(num_depths):
             metrics[f"train/ploss_{i}"] = ploss_values[i]
             metrics[f"train/acc_{i}"] = acc_values[i]
 
         if dist.get_rank() == 0:
-            logger.debug(f"step {step}: {metrics}")
+            logger.debug(f"step {metric_step}: {metrics}")
 
         return metrics

torchspec/training/optimizer.py

@@ -49,6 +49,7 @@ def __init__(
             self.fp32_params,
             lr=lr,
             weight_decay=weight_decay,
+            fused=True,
         )
         self.scheduler = LRSchedulerWithWarmup(
             self.optimizer,
@@ -71,21 +72,26 @@ def step(self, closure=None):
             grad_norm: The gradient norm before clipping (for logging).
         """
         with torch.no_grad():
+            grad_destinations = []
+            grad_sources = []
             for p, mp, g in zip(self.model_params, self.fp32_params, self.fp32_grads):
                 if p.grad is not None:
-                    g.copy_(p.grad)
+                    grad_destinations.append(g)
+                    grad_sources.append(p.grad)
                     mp.grad = g
                 else:
                     mp.grad = None
+            if grad_destinations:
+                torch._foreach_copy_(grad_destinations, grad_sources)
 
         grad_norm = torch.nn.utils.clip_grad_norm_(self.fp32_params, self.max_grad_norm)
         self.optimizer.step()
 
         self.optimizer.zero_grad()
         self.scheduler.step()
         with torch.no_grad():
-            for p, mp in zip(self.model_params, self.fp32_params):
-                p.data.copy_(mp.data.to(p.dtype))
+            torch._foreach_copy_(self.model_params, self.fp32_params)
+            for p in self.model_params:
                 p.grad = None
 
         return grad_norm
@@ -105,8 +111,7 @@ def load_state_dict(self, state_dict):
     def sync_fp32_params_from_model(self):
         """Reinitialize fp32_params from model params. Call after loading model checkpoint."""
         with torch.no_grad():
-            for mp, p in zip(self.fp32_params, self.model_params):
-                mp.data.copy_(p.data.to(torch.float32))
+            torch._foreach_copy_(self.fp32_params, self.model_params)
 
     def state_dict(self):
         return self.optimizer.state_dict()