[h-guo18]

What does this PR do?

Type of change: New feature

Stacked on #2990 [2/2] (branch haoguo/dflash-shared-forward, head 28f4f45), whose branch is the base; merge it after that one. There, every DFlash-family training forward returns draft_metrics, a dict of detached 0-dim tensors. Nothing logs them: EagleTrainerWithAccLog records only train_acc and EAGLE's component losses. So DFlash2's selector_coverage, LiLiCorr's lattice diagnostics, DSpark's confidence_loss, Domino's lambda_base and the objective terms never reach TensorBoard, wandb or trainer_state.json.

• Trainer.
  • compute_loss adds each metric to a running sum per name. The sum stays a tensor, so there is no host sync per step.
  • Trainer.log adds each name's mean over the steps that reported it to logs, as draft/<name>. That costs one .tolist() sync per log step.
  • A step with no supervised token takes the forward's early return, which reports no metrics. That step is left out of the mean rather than counted as zero.
  • The trainer names no variant and no metric.
• Why Trainer.log and not EagleTrainingPlot.on_log. HF registers its reporters (TensorBoard, wandb, ...) as callbacks ahead of user callbacks, and Trainer.log copies logs into log_history before any callback runs. A key added in on_log therefore never reaches TensorBoard, and reaches trainer_state.json only by patching the recorded entry.
  • That is how train_acc/parallel_<i>_step_<j> travels on main: since specdec: stop the training regression gate passing unevaluated #2776, EagleTrainingPlot.on_log copies it into log_history[-1] for the regression gate. It reaches the console, trainer_state.json and EagleTrainingPlot's own wandb.log, but not TensorBoard.
  • Trainer.log is the hook HF documents for this ("Subclass and override this method to inject custom behavior").
  • The on_log comment claimed train_acc is "picked up by TensorBoard". It now says where train_acc actually goes. Its handling is unchanged.
• Ranks. Each rank averages its own steps, as for train_acc, and nothing is reduced across ranks. So there is no collective for a rank without metrics to hang. The reporters write rank 0's values.
• Regression gate. tools/launcher/common/check_regression.py reads the last logged entry of the latest checkpoint's trainer_state.json and takes its first key containing acc as the run's accuracy; on main that key is train_acc/parallel_0_step_0. TrainerState.save_to_json writes the entry with sorted keys, so this PR's draft/* keys come first.
  • Unguarded, MIN_FINAL_ACC would start gating DFlash2, DSpark, Domino and LiLiCorr runs on selector_accuracy, backbone_accuracy or lilicorr_slot_acc, which its thresholds were not chosen for.
  • Since [1/8] launcher: make the regression gate machine-readable and opt-in strict #2839 that verdict is machine-readable and can fail the training task. regression_check.json would report the diagnostic as the run's accuracy. The Qwen3-8B DFlash2 online example runs the gate with REGRESSION_STRICT=1, so its training task would pass or fail on the selector's accuracy.
  • The gate now skips draft/* and keeps reading train_acc. The MIN_FINAL_ACC row in the speculative-decoding skill's references/algorithms/dflash.md, which the other DFlash-family references point to, notes the exception.
• Docs. The dflash_lilicorr_w_ce description now says where its metrics are logged, and there is a CHANGELOG entry under 0.49.0. DFlash2's comment that "the trainer converts them at the logging boundary" is now true as written.

What each variant logs:

projector_type	draft/<name>
DFlash	loss_ce or loss_kd
DFlash2	its objective term, selector_accuracy, selector_coverage
DSpark	loss_ce, loss_tvd, backbone_accuracy, confidence_loss
Domino	loss_final, loss_backbone, backbone_accuracy, lambda_base
LiLiCorr	26 names: the lattice diagnostics, each term and its weight, loss_ce or loss_kd, origin_loss, origin_accuracy

Usage

Nothing to configure. TensorBoard shows train/draft/<name>, since HF prefixes training keys with train/. trainer_state.json keeps the metrics too:

import json

history = json.load(open(f"{output_dir}/trainer_state.json"))["log_history"]
coverage = [(h["step"], h["draft/selector_coverage"]) for h in history if "draft/selector_coverage" in h]

Testing

• New test_trainer_logs_each_draft_metric_mean_per_window in test_hf_speculative_offline.py, the unit module that already loads eagle_utils.
  • It runs the real HF Trainer loop through EagleTrainerWithAccLog: 4 steps, 2 per log window. The model's second step reports no metrics.
  • It checks log_history, which records the dict the reporters receive. Each window must hold each name's mean, and a float64 metric must read back exactly.
  • It fails on the base with KeyError: 'draft/step', and takes 0.3 s on 4 CPUs (taskset -c 0-3).
• New test_draft_metrics_are_not_the_accuracy in tools/launcher/tests/test_check_regression.py. Its last log entry holds draft/selector_accuracy 0.09 and train_acc/parallel_0_step_0 0.17, written with sorted keys as HF writes them. Against main's check_regression.py it fails, reading 0.09.
• End to end on 1 GPU (scratch script), on transformers 5.5.4, 5.12.1 and 5.18.0: DFlash, DFlash2, LiLiCorr, DSpark (with its confidence head) and Domino (with DominoLambdaCallback).
  • Each trained 6 steps through EagleTrainerWithAccLog and EagleTrainingPlot with report_to=tensorboard. One batch had no supervised token.
  • Every draft_metrics name appears as train/draft/<name> in the TensorBoard event file and as draft/<name> in log_history. Each value equals the mean of what that window's forwards returned. train_acc is in log_history but not in TensorBoard.
  • Under torch.cuda.set_sync_debug_mode, the per-step accumulation does not sync, and log syncs once when it has metrics.
  • On each run's saved trainer_state.json, main's check_regression.py reads draft/selector_accuracy (DFlash2), draft/lilicorr_slot_acc (LiLiCorr) or draft/backbone_accuracy (DSpark, Domino) as the accuracy. With this PR it reads train_acc/parallel_0_step_0 for all five. DFlash logs no draft/* key containing acc.
• 2-rank gloo DDP through the HF Trainer, on the same three versions: it finishes on both ranks, each rank logs its own means, and rank 0's TensorBoard holds rank 0's.
• test_hf_speculative_offline.py, test_hf_{dflash,dflash2,dspark,domino,lilicorr}.py and tools/launcher/tests/test_check_regression.py: 388 passed on each of transformers 5.5.4 and 5.18.0 (CI's tf_min and tf_latest pins), against 386 on the base.
• tests/gpu/torch/speculative/plugins on 1 GPU: 20 passed, 1 skipped. The skipped test needs 2 GPUs.
• Bitwise A/B against the base, comparing loss, gradients, metrics and generated tokens: 23 DFlash-family cases plus 15 more, in fp32 and bf16. The extra cases cover DSpark's gated head and its rnn head online, the KD and D-PACE flags on DSpark and Domino, DFlash2 KD online, more zero-loss steps, and labels that are -100 on some tokens. All 76 runs are identical, since no model code changes.
• FSDP2 probe (2 gloo ranks, the DFlash draft as its own unit, one rank on the zero-loss branch, CE and KD): both ranks OK, with the same losses as on the base.
• pre-commit run --files on the 7 changed files: clean.

Before your PR is "Ready for review"

Make sure you read and follow Contributor guidelines and your commits are signed (git commit -s -S).

Make sure you read and follow the Security Best Practices (e.g. avoiding hardcoded trust_remote_code=True, torch.load(..., weights_only=False), pickle, etc.).

• Is this change backward compatible?: ✅ It only adds log keys, and the regression gate reads the same accuracy key as before.
• If you copied code from any other sources or added a new PIP dependency, did you follow guidance in CONTRIBUTING.md: N/A
• Did you write any new necessary tests?: ✅
• Did you update Changelog?: ✅ Under 0.49.0, Speculative Decoding.
• Did you get Claude approval on this PR?: ❌ Not run yet.

Additional Information

Stack: #2757 ← #2762 ← #2989 [1/2] ← #2990 [2/2] ← this PR. Merge them in that order. Three other PRs sit on #2990 and are independent of this one: DSpark and Domino following dflash_loss_objective (#2991), loss terms in config (#2993), and the shared AR-decode helper (#2994). This branch merges with each of them. The only conflict is in CHANGELOG.rst with the loss-terms PR (#2993), where both add an entry to the same 0.49.0 list; keeping both entries resolves it.

Not changed here: train_acc still never reaches TensorBoard, for EAGLE as for the DFlash family, because EagleTrainingPlot.on_log adds it after the reporters have run. #2776 copied it into trainer_state.json for the gate. Adding it in Trainer.log, as this PR does for draft/*, would also chart it in TensorBoard. That changes how train_acc is handled, so it is left for a follow-up. The same holds for estimated_training_ar, which reaches only the console and wandb.

🤖 Generated with Claude Code

[copy-pr-bot]

Auto-sync is disabled for draft pull requests in this repository. Workflows must be run manually.

Contributors can view more details about this message here.

[coderabbitai]

Important

Draft PR not reviewed

Draft PRs are not automatically reviewed by default.

• Trigger a manual review

To automatically review draft PRs, update your CodeRabbit configuration:

reviews:
  auto_review:
    drafts: true

• Autofix · Keep fixing CodeRabbit findings and required CI, and resolving merge conflicts

_{Comment @coderabbitai help to get the list of available commands.}