perf: reduce .item() calls and fuse optimizer

.gitignore

@@ -92,3 +92,4 @@ wandb/log.txt
 wandb/
 data/
 profiles/
+.tmp/

tests/test_loop_cleanup.py

@@ -76,6 +76,165 @@ def test_run_training_loop_finally_runs_cleanup_on_exception(self):
         )
 
 
+def _run_mock_training_loop(
+    *,
+    prefetch_depth: int,
+    num_steps: int,
+    steps_per_epoch: int,
+    dispatch_results=None,
+):
+    args = SimpleNamespace(
+        training_num_nodes=1,
+        training_num_gpus_per_node=1,
+        num_train_steps=num_steps,
+        draft_accumulation_steps=1,
+        steps_per_epoch=steps_per_epoch,
+        num_epochs=(num_steps + steps_per_epoch - 1) // steps_per_epoch,
+        global_batch_size=1,
+        per_dp_rank_batch_size=1,
+        prefetch_depth=prefetch_depth,
+        enable_perf_metrics=False,
+        save_interval=0,
+        checkpoint_dir=None,
+        save_per_epoch=False,
+        train_with_decode=False,
+    )
+    events = []
+    results = iter(dispatch_results) if dispatch_results is not None else None
+
+    controller = mock.MagicMock()
+    controller.submit_training_dataset.remote.return_value = None
+    controller.reload_dataset.remote.return_value = None
+
+    def dispatch():
+        result = next(results) if results is not None else True
+        events.append(f"dispatch:{result}")
+        return result
+
+    controller.try_dispatch_batch.remote.side_effect = dispatch
+    controller.get_full_status.remote.return_value = {
+        "inference_speed": 0.0,
+        "sample_pool_size": 0,
+        "elapsed_seconds": 0.0,
+        "avg_inference_speed": 0.0,
+        "avg_training_speed": 0.0,
+    }
+
+    actor = mock.MagicMock()
+    actor.get_global_step.remote.return_value = 0
+
+    def train(*, step, num_batches):
+        assert num_batches == 1
+        events.append(f"train:{step}")
+        return {}
+
+    actor.train_from_queue.remote.side_effect = train
+    train_group = mock.MagicMock()
+    train_group._actor_handlers = [actor]
+
+    eval_state = eval_utils.EvalSetupState(
+        eval_interval=0,
+        eval_enabled=False,
+        eval_cache_loaded=False,
+        eval_cache_path=None,
+        best_eval_score=0.0,
+        eval_dispatch_bs=0,
+        eval_dataset_size=0,
+        dp_size=1,
+    )
+
+    with (
+        mock.patch("torchspec.controller.loop.setup_eval", return_value=eval_state),
+        mock.patch("torchspec.controller.loop.ray.get", side_effect=lambda value: value),
+        mock.patch("torchspec.controller.loop.time.sleep") as mock_sleep,
+        mock.patch("torchspec.controller.loop.tqdm"),
+    ):
+        loop.training_loop(
+            args,
+            controller,
+            mock.MagicMock(),
+            train_group,
+            dataset_size=max(num_steps, 1),
+            eval_dataset_size=0,
+        )
+
+    return events, controller, mock_sleep
+
+
+def test_training_loop_prefills_and_refills_controller_queue():
+    events, _, _ = _run_mock_training_loop(
+        prefetch_depth=2,
+        num_steps=3,
+        steps_per_epoch=3,
+    )
+
+    assert events == [
+        "dispatch:True",
+        "dispatch:True",
+        "train:0",
+        "dispatch:True",
+        "train:1",
+        "train:2",
+    ]
+
+
+def test_training_loop_prefetch_zero_preserves_one_step_dispatching():
+    events, _, _ = _run_mock_training_loop(
+        prefetch_depth=0,
+        num_steps=3,
+        steps_per_epoch=3,
+    )
+
+    assert events == [
+        "dispatch:True",
+        "train:0",
+        "dispatch:True",
+        "train:1",
+        "dispatch:True",
+        "train:2",
+    ]
+
+
+def test_training_loop_runs_ready_step_when_lookahead_cannot_be_filled():
+    events, _, mock_sleep = _run_mock_training_loop(
+        prefetch_depth=2,
+        num_steps=3,
+        steps_per_epoch=3,
+        dispatch_results=[True, False, True, True],
+    )
+
+    assert events == [
+        "dispatch:True",
+        "dispatch:False",
+        "train:0",
+        "dispatch:True",
+        "dispatch:True",
+        "train:1",
+        "train:2",
+    ]
+    mock_sleep.assert_not_called()
+
+
+def test_training_loop_does_not_prefetch_across_epoch_boundary():
+    events, controller, _ = _run_mock_training_loop(
+        prefetch_depth=3,
+        num_steps=4,
+        steps_per_epoch=2,
+    )
+
+    assert events == [
+        "dispatch:True",
+        "dispatch:True",
+        "train:0",
+        "train:1",
+        "dispatch:True",
+        "dispatch:True",
+        "train:2",
+        "train:3",
+    ]
+    controller.reload_dataset.remote.assert_called_once_with()
+
+
 def test_generate_eval_cache_dispatches_and_finalizes():
     controller = mock.MagicMock()
     train_group = mock.MagicMock()

torchspec/controller/loop.py

@@ -227,6 +227,7 @@ def training_loop(
     )
 
     enable_perf = getattr(args, "enable_perf_metrics", True)
+    prefetch_batches = max(accumulation_steps, getattr(args, "prefetch_depth", 0))
 
     completed_steps = start_step
     current_epoch = completed_steps // steps_per_epoch + 1
@@ -235,23 +236,35 @@ def training_loop(
         logger.info(f"Resuming from step {start_step} (epoch {current_epoch})")
     dispatch_attempts = 0
     consecutive_failures = 0
+    queued_batches = 0
     last_saved_step: int | None = None
     progress = tqdm(total=num_steps, desc="Training", unit="step", initial=start_step)
     while completed_steps < num_steps:
-        # Inner loop: dispatch accumulation_steps batches before training
-        dispatches_done = 0
+        remaining_steps = min(
+            num_steps - completed_steps,
+            steps_per_epoch - steps_in_current_epoch,
+        )
+        target_queued_batches = min(
+            prefetch_batches,
+            remaining_steps * accumulation_steps,
+        )
         if enable_perf:
             t_dispatch = time.time()
         status = None
-        while dispatches_done < accumulation_steps:
+        while queued_batches < target_queued_batches:
             dispatch_attempts += 1
 
             dispatched = ray.get(controller.try_dispatch_batch.remote())
             if dispatched:
-                dispatches_done += 1
+                queued_batches += 1
                 consecutive_failures = 0
             else:
                 consecutive_failures += 1
+                if queued_batches >= accumulation_steps:
+                    # The current step can run while inference refills the pool.
+                    target_queued_batches = queued_batches
+                    consecutive_failures = 0
+                    continue
 
                 # Only fetch status when needed for logging or reload decision
                 if dispatch_attempts % 100 == 0 or consecutive_failures >= 500:
@@ -271,7 +284,7 @@ def training_loop(
                     should_reload = True
                 elif (
                     consecutive_failures >= 500
-                    and (completed_steps > 0 or dispatches_done > 0)
+                    and (completed_steps > 0 or queued_batches > 0)
                     and status is not None
                     and status["sample_pool_size"] < status["dispatch_batch_size"]
                     and status.get("prompt_buffer_size", 0) == 0
@@ -280,7 +293,7 @@ def training_loop(
                         f"Pool insufficient for dispatch "
                         f"(pool_size={status['sample_pool_size']}, "
                         f"need={status['dispatch_batch_size']}, "
-                        f"{dispatches_done}/{accumulation_steps} dispatches done, "
+                        f"{queued_batches}/{accumulation_steps} batches queued, "
                         f"{steps_in_current_epoch}/{steps_per_epoch} steps in epoch). "
                         f"Reloading dataset."
                     )
@@ -299,7 +312,7 @@ def training_loop(
 
                 time.sleep(0.01)
         else:
-            # All accumulation dispatches succeeded — run training
+            # The current optimizer step is fully queued.
             if enable_perf:
                 dispatch_wait = time.time() - t_dispatch
 
@@ -312,6 +325,7 @@ def training_loop(
             ]
 
             train_results = ray.get(train_futures)
+            queued_batches -= accumulation_steps
             completed_steps += 1
 
             # Log metrics from training (use rank 0's metrics - they're already all-reduced)

torchspec/models/eagle3.py

@@ -281,11 +281,7 @@ def forward(
 
             loss = local_sum_loss / local_count.clamp_min(1.0)
             metric_loss = loss.detach()
-            metric_acc = (
-                (local_correct / local_count.clamp_min(1.0)).detach()
-                if float(local_count.detach().float().cpu()) > 0.0
-                else local_correct.detach().float() * 0.0
-            )
+            metric_acc = (local_correct / local_count.clamp_min(1.0)).detach()
 
             if self._usp_sp_group is not None:
                 reduced_stats = torch.stack(
@@ -299,16 +295,9 @@ def forward(
                 reduced_sum_loss, reduced_correct, reduced_count = reduced_stats.unbind()
                 denom = reduced_count.clamp_min(1.0)
                 loss = (local_sum_loss / denom).to(loss.dtype)
-                if reduced_count.item() > 0:
-                    metric_loss = (reduced_sum_loss / denom).detach()
-                    metric_acc = (reduced_correct / denom).to(
-                        device=loss.device, dtype=torch.float32
-                    )
-                    metric_count = reduced_count.to(device=loss.device, dtype=torch.float32)
-                else:
-                    metric_loss = reduced_sum_loss.detach() * 0.0
-                    metric_acc = local_correct.detach().float() * 0.0
-                    metric_count = reduced_count.to(device=loss.device, dtype=torch.float32)
+                metric_loss = (reduced_sum_loss / denom).detach()
+                metric_acc = (reduced_correct / denom).to(device=loss.device, dtype=torch.float32)
+                metric_count = reduced_count.to(device=loss.device, dtype=torch.float32)
             else:
                 metric_count = local_count.detach().float().to(device=loss.device)
 

torchspec/training/eagle3_trainer.py

@@ -59,6 +59,10 @@ class Eagle3Trainer(Trainer):
     def __init__(self, args: Namespace):
         super().__init__(args)
         self.target_lm_head: Optional[torch.nn.Module] = None
+        self._ploss_weights = tuple(
+            _position_decay_weights(args.ttt_length, getattr(args, "ploss_weights", None))
+        )
+        self._ploss_weight_sum = sum(self._ploss_weights)
 
     def init_model(
         self,
@@ -314,10 +318,10 @@ def _forward(self, batch: dict) -> Tuple[List[torch.Tensor], List[torch.Tensor]]
         return plosses, vlosses, acces, acc_counts
 
     def _backward(self, plosses: List[torch.Tensor], accumulation_steps: int = 1) -> torch.Tensor:
-        ploss_weight = _position_decay_weights(
-            len(plosses), getattr(self.args, "ploss_weights", None)
+        ploss = (
+            sum(self._ploss_weights[i] * plosses[i] for i in range(len(plosses)))
+            / accumulation_steps
         )
-        ploss = sum(ploss_weight[i] * plosses[i] for i in range(len(plosses))) / accumulation_steps
         ploss.backward()
         return ploss
 
@@ -461,41 +465,51 @@ def _aggregate_metrics(
         avg_vlosses = torch.stack([m["vlosses"] for m in all_step_metrics]).mean(dim=0)
         avg_acces = torch.stack([m["acces"] for m in all_step_metrics]).mean(dim=0)
 
-        dist.all_reduce(avg_vlosses, op=dist.ReduceOp.AVG)
-        dist.all_reduce(avg_acces, op=dist.ReduceOp.AVG)
+        num_depths = avg_vlosses.shape[0]
+        reduced_metrics = torch.cat((avg_vlosses, avg_acces))
+        dist.all_reduce(reduced_metrics, op=dist.ReduceOp.AVG)
 
-        avg_acc_scalar = avg_acces.mean().item()
+        grad_norm_value = (
+            grad_norm.to(device=avg_vlosses.device, dtype=avg_vlosses.dtype)
+            if grad_norm is not None
+            else avg_vlosses.new_zeros(())
+        )
+        packed_metrics = torch.cat(
+            (
+                reduced_metrics,
+                grad_norm_value.reshape(1),
+            )
+        ).detach()
+        metric_values = packed_metrics.float().cpu().tolist()
 
-        # Simulated acceptance length: acc_0 + acc_0*acc_1 + acc_0*acc_1*acc_2 + ...
-        # Models the expected number of consecutively accepted draft tokens,
-        # which better reflects actual speculative decoding performance.
-        cumulative = 1.0
-        simulated_acc_len = 0.0
-        for i in range(avg_acces.shape[0]):
-            cumulative *= avg_acces[i].item()
-            simulated_acc_len += cumulative
+        ploss_values = metric_values[:num_depths]
+        acc_values = metric_values[num_depths : 2 * num_depths]
+        grad_norm_scalar = metric_values[-1]
 
-        ploss_weights = torch.tensor(
-            _position_decay_weights(
-                avg_vlosses.shape[0], getattr(self.args, "ploss_weights", None)
-            ),
-            device=avg_vlosses.device,
+        weighted_avg_loss_value = (
+            sum(loss * weight for loss, weight in zip(ploss_values, self._ploss_weights))
+            / self._ploss_weight_sum
         )
-        weighted_avg_loss = (avg_vlosses * ploss_weights).sum().item() / ploss_weights.sum().item()
+        avg_acc_scalar = sum(acc_values) / num_depths
+        cumulative = 1.0
+        simulated_acc_len_value = 0.0
+        for acc in acc_values:
+            cumulative *= acc
+            simulated_acc_len_value += cumulative
 
         metrics = {
-            "train/avg_loss": weighted_avg_loss,
+            "train/avg_loss": weighted_avg_loss_value,
             "train/avg_acc": avg_acc_scalar,
-            "train/simulated_acc_len": simulated_acc_len,
-            "train/grad_norm": grad_norm.item() if grad_norm is not None else 0.0,
+            "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,
         }
 
-        for i in range(avg_vlosses.shape[0]):
-            metrics[f"train/ploss_{i}"] = avg_vlosses[i].item()
-            metrics[f"train/acc_{i}"] = avg_acces[i].item()
+        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}")