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)* 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/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/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] 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):