feat(speculative): log the DFlash family's draft_metrics from the example trainer

CHANGELOG.rst

@@ -18,6 +18,7 @@ Changelog
 *Speculative Decoding*
 
 - Add the **external draft** speculative-decoding mode: train a pretrained standalone draft model against a cached base model, instead of grafting a draft head onto the target. Select it with ``mtsp.convert(draft, [("external", cfg)])`` or the ``speculative_external`` recipe; the draft keeps its own embeddings and lm_head and exports as an ordinary HF checkpoint.
+- ``examples/speculative_decoding`` now logs the DFlash family's training metrics -- e.g. DFlash2's ``selector_coverage``, LiLiCorr's lattice diagnostics and Domino's ``lambda_base`` -- as ``draft/<name>``, averaged over each logging interval, to TensorBoard, wandb and ``trainer_state.json``.
 
 *Megatron Framework (M-LM / M-Bridge)*
 

examples/speculative_decoding/eagle_utils.py

@@ -225,7 +225,7 @@ def make_speculative_data_module(
 
 
 class EagleTrainerWithAccLog(Trainer):
-    """Wrapper around Trainer that logs training accuracy."""
+    """Wrapper around Trainer that logs training accuracy and the model's ``draft_metrics``."""
 
     def __init__(
         self,
@@ -235,6 +235,8 @@ def __init__(
     ):
         super().__init__(*args, **kwargs)
         self.lora_lr_multiplier = lora_lr_multiplier
+        # name -> (sum, count) of each draft metric since the last log
+        self._draft_metric_sums = {}
 
     def create_optimizer(self, model=None):
         """Override to give LoRA parameters a higher learning rate.
@@ -316,7 +318,7 @@ def create_optimizer(self, model=None):
         return self.optimizer
 
     def compute_loss(self, *args, **kwargs):
-        """Override compute_loss to save train accs and per-component losses in trainer state."""
+        """Override compute_loss to record train accs, per-component losses and draft metrics."""
         if not hasattr(self.state, "training_accs"):
             self.state.training_accs = []
         if not hasattr(self.state, "component_losses"):
@@ -333,8 +335,25 @@ def compute_loss(self, *args, **kwargs):
             val = getattr(outputs, attr, None)
             if val is not None:
                 self.state.component_losses[key].append(val.item())
+        # Summed as tensors: reading the values here would sync the host on every step.
+        for name, value in (getattr(outputs, "draft_metrics", None) or {}).items():
+            total, count = self._draft_metric_sums.get(name, (0.0, 0))
+            self._draft_metric_sums[name] = (total + value, count + 1)
         return loss
 
+    def log(self, logs, *args, **kwargs):
+        """Log ``logs`` with each draft metric's mean since the last log, as ``draft/<name>``.
+
+        Each rank averages its own steps, as for ``train_acc``; nothing is reduced across ranks.
+        """
+        # Not in EagleTrainingPlot.on_log: the HF reporters are callbacks that run before it.
+        if self._draft_metric_sums:
+            sums = self._draft_metric_sums
+            means = torch.stack([total / count for total, count in sums.values()]).tolist()
+            logs.update({f"draft/{name}": mean for name, mean in zip(sums, means)})
+            self._draft_metric_sums = {}
+        super().log(logs, *args, **kwargs)
+
 
 class LoRAWarmupCallback(TrainerCallback):
     """Manages LoRA warmup: freezes LoRA during warmup, unfreezes after."""
@@ -484,7 +503,8 @@ def on_log(self, args, state, control, **kwargs):
             print_rank_0(f"Step {state.global_step} Training Acc: [{acc_str}]")
         except Exception:
             print_rank_0(f"Step {state.global_step} Training Acc: {average_acc}")
-        # Log accuracy to HF Trainer's logs dict (picked up by TensorBoard)
+        # The HF reporters (TensorBoard) have already run: these reach only the console,
+        # wandb below and, through the log_history update, trainer_state.json.
         logs = kwargs.get("logs")
         if logs is None:
             logs = {}

modelopt/torch/speculative/config.py

@@ -341,8 +341,9 @@ class DFlashConfig(ModeloptBaseConfig):
             "each is the coefficient with which its term enters the total, and "
             "`loss == origin_loss + lilicorr_loss` holds exactly. Both halves and all three "
             "weights are reported in the `draft_metrics` dict the model attaches to its "
-            "forward output, so a consumer of those metrics can check the identity per step "
-            "and read which composition produced a checkpoint. The three weights are validated "
+            "forward output, which examples/speculative_decoding logs as `draft/<name>`, so a "
+            "consumer of those metrics can check the identity per step and read which "
+            "composition produced a checkpoint. The three weights are validated "
             "all-or-nothing (a negative value means unset): a config that sets some but not "
             "others is rejected rather than inheriting a default composition. Shipped "
             "variants: 0.25 ('base') and 0.125 ('margin'). "

plugins/modelopt/skills/speculative-decoding/references/algorithms/dflash.md

@@ -130,7 +130,7 @@ DFlash gates in three places rather than on one acceptance-rate number.
 | Env var | Meaning |
 | --- | --- |
 | `MAX_FINAL_LOSS` | Final loss must be below this |
-| `MIN_FINAL_ACC` | Final accuracy must be above this (any log key containing `acc`) |
+| `MIN_FINAL_ACC` | Final accuracy must be above this (any log key containing `acc`, except `draft/*`) |
 
 Qwen3-8B online reference uses `MAX_FINAL_LOSS=5.0`, `MIN_FINAL_ACC=0.15`. Its
 convergence baseline (8×B200, bs=1, seq_len=4096, 5-layer draft, block_size=16, 100K

tests/unit/torch/speculative/plugins/test_hf_speculative_offline.py

@@ -31,6 +31,8 @@
 import pytest
 import torch
 from _test_utils.torch.transformers_models import get_tiny_llama, get_tiny_tokenizer
+from transformers import TrainingArguments
+from transformers.utils import ModelOutput
 
 import modelopt.torch.speculative as mtsp
 from modelopt.torch.speculative.eagle.default_config import default_eagle_config
@@ -765,3 +767,54 @@ def test_collator_batches_multiple(tmp_path):
         train_len,
         HIDDEN_SIZE,
     )
+
+
+# ---------------------------------------------------------------------------
+# EagleTrainerWithAccLog: draft metrics
+# ---------------------------------------------------------------------------
+
+
+class _DraftMetricsModel(torch.nn.Module):
+    """A DFlash-family training forward's outputs; step 2 has nothing to train, so no metrics."""
+
+    def __init__(self):
+        super().__init__()
+        self.weight = torch.nn.Parameter(torch.ones(()))
+        self.calls = 0
+
+    def forward(self, x):
+        self.calls += 1
+        out = ModelOutput(loss=(self.weight * x).sum(), train_acc=[[0.5]])
+        if self.calls != 2:
+            out["draft_metrics"] = {
+                "step": torch.tensor(float(self.calls)),
+                "scheduled": torch.tensor(0.1, dtype=torch.float64),
+            }
+        return out
+
+
+def test_trainer_logs_each_draft_metric_mean_per_window(tmp_path):
+    """Each window logs the mean over the steps that reported it, before the reporters run.
+
+    ``log_history`` records the logs dict exactly as the HF reporters (TensorBoard, wandb)
+    receive it, which a callback's ``on_log`` cannot add to in time.
+    """
+    args = TrainingArguments(
+        output_dir=str(tmp_path),
+        max_steps=4,
+        logging_steps=2,
+        per_device_train_batch_size=1,
+        report_to="none",
+        save_strategy="no",
+        use_cpu=True,
+        disable_tqdm=True,
+    )
+    trainer = _eagle_utils.EagleTrainerWithAccLog(
+        model=_DraftMetricsModel(), args=args, train_dataset=[{"x": torch.tensor(1.0)}] * 4
+    )
+    trainer.train()
+
+    logged = [h for h in trainer.state.log_history if "loss" in h]
+    assert [h["draft/step"] for h in logged] == [1.0, 3.5]
+    # float64 survives to the log: a scheduled weight reads back exactly.
+    assert [h["draft/scheduled"] for h in logged] == [0.1, 0.1]

tools/launcher/common/check_regression.py

@@ -21,9 +21,9 @@
 
 Environment variables (all optional — no check if unset):
     MAX_FINAL_LOSS:  Final loss must be below this value
-    MIN_FINAL_ACC:   Final accuracy must be above this value (any key containing 'acc').
-                     A threshold whose metric is absent fails the check rather than
-                     passing unevaluated.
+    MIN_FINAL_ACC:   Final accuracy must be above this value (any key containing 'acc',
+                     other than the draft/* metrics). A threshold whose metric is absent
+                     fails the check rather than passing unevaluated.
     MAX_FINAL_PERPLEXITY: Final perplexity must be below this value
 
 An empty variable counts as unset, and one that isn't a finite number fails the check.
@@ -92,9 +92,10 @@ def get_final_metrics(trainer_state_path):
     last = logs[-1]
     metrics = {"loss": float(last["loss"])}
 
-    # Find any accuracy key (train_acc/parallel_0_step_0, eval_accuracy, etc.)
+    # Find any accuracy key (train_acc/parallel_0_step_0, eval_accuracy, etc.). draft/* are a
+    # draft variant's diagnostics, e.g. Domino's backbone_accuracy, not the run's accuracy.
     for key, value in last.items():
-        if "acc" in key.lower():
+        if "acc" in key.lower() and not key.startswith("draft/"):
             metrics["accuracy"] = float(value)
             break
 

tools/launcher/tests/test_check_regression.py

@@ -82,6 +82,15 @@ def test_metrics_pick_up_any_accuracy_key(tmp_path):
     assert metrics["accuracy"] == pytest.approx(0.5153)
 
 
+def test_draft_metrics_are_not_the_accuracy(tmp_path):
+    """A draft/* diagnostic, e.g. DFlash2's selector_accuracy, is not the run's train_acc."""
+    entry = {"loss": 6.5, "draft/selector_accuracy": 0.09, "train_acc/parallel_0_step_0": 0.17}
+    state = tmp_path / "trainer_state.json"
+    # Sorted, as TrainerState.save_to_json writes it: draft/* comes before train_acc.
+    state.write_text(json.dumps({"log_history": [entry]}, sort_keys=True))
+    assert check_regression_mod.get_final_metrics(state)["accuracy"] == pytest.approx(0.17)
+
+
 def test_latest_checkpoint_is_chosen_by_step_number(tmp_path):
     """checkpoint-12000 is newer than checkpoint-8000, though it sorts first as a string."""
     for step in (8000, 12000):