From a3b80eea93eb57d3f2b4b9f94311c87ed8bb0aca Mon Sep 17 00:00:00 2001 From: h-guo18 <67671475+h-guo18@users.noreply.github.com> Date: Mon, 12 Oct 2026 14:58:06 +0000 Subject: [PATCH 1/3] fix(launcher): keep the regression gate's accuracy off the draft/ metrics The next commit has the speculative-decoding example trainer log each DFlash-family variant's training metrics as draft/, which also puts them in trainer_state.json. check_regression.py takes the first key containing "acc" in the last logged entry as the run's accuracy, and TrainerState.save_to_json writes that entry with sorted keys, so the draft/ keys come before train_acc/parallel_0_step_0, the key EagleTrainingPlot copies into the entry for this gate. MIN_FINAL_ACC would then silently start gating DFlash2, DSpark, Domino and LiLiCorr runs on a head diagnostic (selector_accuracy, backbone_accuracy, lilicorr_slot_acc) instead of the training accuracy its thresholds were chosen for: regression_check.json would report that diagnostic as the run's accuracy, and the DFlash2 online example, which runs the gate with REGRESSION_STRICT=1, would pass or fail its training task on it. Skip draft/* so the gate keeps reading train_acc. The new test_check_regression.py case writes such an entry, a draft/selector_accuracy of 0.09 beside a train_acc of 0.17, and fails without the skip, which reads 0.09. The file's 19 tests pass. Co-Authored-By: Claude Opus 5.5 Signed-off-by: h-guo18 <67671475+h-guo18@users.noreply.github.com> --- .../references/algorithms/dflash.md | 2 +- tools/launcher/common/check_regression.py | 11 ++++++----- tools/launcher/tests/test_check_regression.py | 9 +++++++++ 3 files changed, 16 insertions(+), 6 deletions(-) diff --git a/plugins/modelopt/skills/speculative-decoding/references/algorithms/dflash.md b/plugins/modelopt/skills/speculative-decoding/references/algorithms/dflash.md index 90d7f457369..edb1b7142ef 100644 --- a/plugins/modelopt/skills/speculative-decoding/references/algorithms/dflash.md +++ b/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 diff --git a/tools/launcher/common/check_regression.py b/tools/launcher/common/check_regression.py index 47292de44f8..a1d094bcfc6 100644 --- a/tools/launcher/common/check_regression.py +++ b/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 diff --git a/tools/launcher/tests/test_check_regression.py b/tools/launcher/tests/test_check_regression.py index 0832fced283..8e5ad8b7205 100644 --- a/tools/launcher/tests/test_check_regression.py +++ b/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): From 538666e30fb29223658df2fb6da137f540fe0927 Mon Sep 17 00:00:00 2001 From: h-guo18 <67671475+h-guo18@users.noreply.github.com> Date: Mon, 12 Oct 2026 14:58:23 +0000 Subject: [PATCH 2/3] feat(speculative): log the DFlash family's draft_metrics from the example trainer Every DFlash-family training forward returns draft_metrics, but the example trainer never logged them, so DFlash2's selector_coverage, LiLiCorr's lattice diagnostics, DSpark's confidence_loss and Domino's lambda_base never reached TensorBoard or wandb. EagleTrainerWithAccLog now sums each metric as a tensor per micro-batch, so no host sync, and Trainer.log adds each name's mean over the steps that reported it as draft/, with one sync per log step. It has to be Trainer.log, not EagleTrainingPlot.on_log: the HF reporters are callbacks registered ahead of ours, so a key added in on_log never reaches TensorBoard, and log_history is recorded before any callback runs, so such a key reaches trainer_state.json only by patching the recorded entry, as EagleTrainingPlot does for train_acc. train_acc therefore still never reaches TensorBoard; the comment saying it does is corrected, its handling left as is. The trainer names no variant or metric. Each rank averages its own steps, as for train_acc. The LiLiCorr weight description now says where the metrics are logged. Verified: - the new unit test, which fails on the parent with KeyError 'draft/step'; - DFlash, DFlash2, LiLiCorr, DSpark and Domino each trained through the real HF Trainer on GPU with report_to=tensorboard, on transformers 5.5.4, 5.12.1 and 5.18.0: every draft_metrics name lands as train/draft/ and in log_history, equal to the window mean of what the forwards returned, including a window with a step that had no supervised token; train_acc lands in log_history only; - on the trainer_state.json each of those runs saves, check_regression.py reads train_acc/parallel_0_step_0 as the accuracy; without the previous commit it reads draft/selector_accuracy (DFlash2), draft/lilicorr_slot_acc (LiLiCorr) or draft/backbone_accuracy (DSpark, Domino); - torch.cuda.set_sync_debug_mode: no sync in the per-step accumulation, one per log call that has metrics; - 2-rank gloo DDP through the HF Trainer: no hang, each rank logs its own means and rank 0 writes them to TensorBoard; - test_hf_speculative_offline.py, the DFlash-family plugin tests and the launcher's test_check_regression.py on transformers 5.5.4 and 5.18.0 (388 passed on each); - tests/gpu/torch/speculative/plugins (20 passed, 1 skipped); - bitwise A/B of all 23 DFlash-family cases and 15 more, fp32 and bf16, against the parent: 76 MATCH; - FSDP2 probe, dflash draft branch ce/kd: both ranks OK, same losses as the parent. Co-Authored-By: Claude Opus 5.5 Signed-off-by: h-guo18 <67671475+h-guo18@users.noreply.github.com> --- examples/speculative_decoding/eagle_utils.py | 26 +++++++-- modelopt/torch/speculative/config.py | 5 +- .../plugins/test_hf_speculative_offline.py | 53 +++++++++++++++++++ 3 files changed, 79 insertions(+), 5 deletions(-) diff --git a/examples/speculative_decoding/eagle_utils.py b/examples/speculative_decoding/eagle_utils.py index a0f76a02ef4..02378f331b0 100644 --- a/examples/speculative_decoding/eagle_utils.py +++ b/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/``. + + 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 = {} diff --git a/modelopt/torch/speculative/config.py b/modelopt/torch/speculative/config.py index 1a53b21e4d4..e172841f371 100644 --- a/modelopt/torch/speculative/config.py +++ b/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/`, 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'). " diff --git a/tests/unit/torch/speculative/plugins/test_hf_speculative_offline.py b/tests/unit/torch/speculative/plugins/test_hf_speculative_offline.py index d124a2a4de1..e08f287bc1f 100644 --- a/tests/unit/torch/speculative/plugins/test_hf_speculative_offline.py +++ b/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] From 78aee4c98bc8a719b0ce7e80bb959a463bf1852e Mon Sep 17 00:00:00 2001 From: h-guo18 <67671475+h-guo18@users.noreply.github.com> Date: Mon, 12 Oct 2026 14:58:32 +0000 Subject: [PATCH 3/3] docs(speculative): note the logged draft metrics in the changelog Users of examples/speculative_decoding get new draft/ scalars in TensorBoard, wandb and trainer_state.json; say what they are and that they are interval means. Checked with the rst pre-commit hooks. The entry goes under the 0.49.0 heading: release/0.48.0 is branched, so this change first reaches users in 0.49.0. Co-Authored-By: Claude Opus 5.5 Signed-off-by: h-guo18 <67671475+h-guo18@users.noreply.github.com> --- CHANGELOG.rst | 1 + 1 file changed, 1 insertion(+) diff --git a/CHANGELOG.rst b/CHANGELOG.rst index 2907f626226..c73384a325f 100755 --- a/CHANGELOG.rst +++ b/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/``, averaged over each logging interval, to TensorBoard, wandb and ``trainer_state.json``. *Megatron Framework (M-LM / M-Bridge)*