Skip to content

Speculators v0.7.0

Latest
Latest

Choose a tag to compare

View all tags
@dsikka dsikka released this 30 Jul 21:43
v0.7.0
1c2895f
This commit was created on GitHub.com and signed with GitHub’s verified signature.
GPG key ID: B5690EEEBB952194
Verified
Learn about vigilant mode.
spec_simple_v07

Speculators v0.7.0 introduces the DSpark algorithm, a config-file-first training CLI, an expanded loss function library, DDP training with updated mixed precision, and broad improvements to training infrastructure and performance. This release also extends sliding window attention to Eagle3 and P-EAGLE, adds tool-call and multi-turn support to on-policy response regeneration, adds support for more loss functions, and makes Muon the default optimizer.

Key Features:

  • DSpark algorithm support: DFlash + Markov logit-bias head + per-position confidence head
  • Config-file-first training CLI with typed YAML configuration
  • Expanded loss library: D-PACE, TV, JSD, RKL, Negative Log-Acceptance, and Hybrid LK
  • DDP as the default distributed backend with torch.autocast mixed precision
  • Sliding window attention extended to Eagle3 and P-EAGLE
  • Tool-call and multi-turn on-policy response regeneration
  • Muon optimizer as the default
  • NVIDIA SPEED-Bench evaluation support
  • Mistral training support
  • Partial MRoPE support for draft model training
  • Multi-turn and tool call response regeneration

DSpark Algorithm Support

Speculators v0.7.0 adds support for DSpark, a new speculator type that extends DFlash with a low-rank Markov logit-bias head for intra-block token dependency, a per-position confidence head, and a TV-dominant training loss. Running DFlash with the heads disabled (--markov-rank 0) reproduces standard DFlash behavior.
Key Components

  1. Markov Logit-Bias Head: Models sequential token dependencies within a draft block
--markov-head-type vanilla
--markov-head-type gated
--markov-head-type rnn
  1. Per-Position Confidence Head: Predicts acceptance probability at each draft position.
    Optionally incorporates the Markov embedding with:
--confidence-head-with-markov
  1. Sample-from-Anchor: A new --sample-from-anchor mode where every block position, including slot, produces predictions. This matches the DSpark author’s implementation, while disabling this feature aligns with the DFlash implementation (and can be useful if finetuning a DSpark model from a DFlash checkpoint).

Example:

scripts/train.py \
    --speculator-type dspark \
    --markov-rank 256 \
    --markov-head-type vanilla \
    --enable-confidence-head \
    --verifier-name-or-path Qwen/Qwen3-8B \
    --data-path ./data \
    --save-path ./ckpt

With this, we released RedHatAI/GLM-5.2-speculator.dspark and RedHatAI/gemma-4-31B-it-speculator.dspark, they provide significant performance gains when compared to autoregressive drafting algorithms such as EAGLE-3.

Screenshot 2026-07-30 at 11 45 28 AM

Config-File-First Training CLI

The handwritten argparse interface in scripts/train.py has been replaced with a typed, layered configuration system built on pydantic-settings.

Features

  • --config PATH
    • Load a YAML configuration file.
    • Precedence: command-line flag > YAML > default.
  • --dump-config
    • Prints the fully resolved configuration as a round-trippable run.yaml and exits.
  • Reproducibility
    • Every run writes run.yaml and train_command.txt alongside checkpoints.
  • Grouped Help
    • Help is organized into General, Verifier, Draft, Data, Loss, Optimizer, Scheduler, Trainer, and algorithm-specific groups (DFlash, DSpark, P-EAGLE, MTP).

All existing flag-based training recipes continue to work unchanged.

Expanded Loss Function Library

D-PACE (Dynamic Position-Aware Cross-Entropy)

https://arxiv.org/abs/2605.18810

  • Dynamically shifts per-position loss weights toward later positions as earlier positions stabilize.
  • Achieves approximately 5% higher mean accepted length compared to fixed-decay cross-entropy.
  • Enable with --per-position-loss-weight dpace.

Additional Losses

  • Total Variation (TV) with a Triton-fused kernel for long-context efficiency
  • Jensen-Shannon Divergence (JSD)
  • Reverse KL (RKL)
  • Negative Log-Acceptance (LK)
  • Hybrid LK

EAL Metric

Expected Accepted Length (EAL) is now available as a training evaluation metric.

Composable Losses

--loss-fn '{"ce": 0.1, "tv": 0.9}'

DDP Training + Updated AMP

The default distributed backend has changed from FSDP to DDP.

  • --fsdp-shard remains available as an opt-in option.
  • Master weights remain in FP32.
  • torch.autocast manages reduced precision during forward and backward passes.
  • Precision-sensitive layers are handled automatically.

Sliding Window Attention for Eagle3 and P-EAGLE

Sliding window attention now extends beyond DFlash to Eagle3 and P-EAGLE.

  • Prevents sharp acceptance-rate drops on long-context requests.
  • Default sliding window size is 2048 tokens.
  • Enabled by default for all draft layers in DFlash and DSpark.
  • --full-attention-indices allows selective overrides.

Tool-Call and Multi-Turn Response Regeneration

Tool-Call Regeneration

  • Tool-call tokens are regenerated on-policy.
  • Cached tool results are spliced back positionally.

Multi-Turn Regeneration

  • Every user turn is regenerated using the newly generated conversation prefix.
  • System prompts remain unchanged.
  • Each assistant turn receives its own reasoning_content.

New Dataset Presets

  • Hermes function-calling dataset
  • Open-PerfectBlend dataset

Muon Optimizer as Default

Muon is now the default optimizer.
Validation shows improved performance over Adam for decoder-only (Qwen3-8B) and MoE models.

  • Muon applies to 2D weight matrices.
  • AdamW automatically handles norms, biases, embeddings, and lm_head.

Training Infrastructure Improvements

  • Sub-epoch checkpointing (--checkpoint-freq 0.5)
  • Scheduler warmup ratio (--scheduler-warmup-ratio)
  • prepare-data safety and correctness improvements (RFC #583)
  • Step-level performance profiling
  • Configurable train/validation split (--train-data-ratio)
  • Restored --max-steps
  • Batch validation metric reduction (one synchronization per epoch)
  • Fused hidden-state dtype casting during collate copy
  • MLflow metric handler
  • Refactored hidden-state transfer with abstract backend
  • On-policy (regenerated) data recommended as the default training strategy
  • Dry run(--dryrun) produce a mock checkpoint with expected config.json and randomized weights, useful to test if the checkpoint runs correctly in vLLM

DFlash Performance Improvements

  • Sorted sampled anchors for contiguous FlexAttention blocks (~3× training speedup)
  • Compiled create_block_mask to avoid dense mask materialization
  • Compute verifier targets only at anchored positions
  • Divergence losses computed in float32 during BF16 training for improved numerical stability
  • Evaluation Improvements
  • NVIDIA SPEED-Bench support in evaluate.py
  • GuideLLM upgraded from 0.6.0 to 0.7.1
  • Simple benchmarking script for quick acceptance-rate evaluation
  • Other Notable Changes
  • Mistral model training support
  • Partial MRoPE support for draft model training
  • Device-agnostic synchronization via torch.accelerator.synchronize()
  • Default drafter architecture changed from Llama to Qwen3
  • Reduced torch.compile graph breaks and recompilations
  • Added CODEOWNERS and required reviewer approvals

Bug Fixes

  • fix(mtp): Pin static shapes for compiled MTP forward to prevent torch.compile crashes across PyTorch versions.
  • fix(train): Fixed DSpark + FSDP2 + Muon crash.
  • fix(train): Prevent double reduction of aliased metric tensors in distributed training.
  • fix(train): Balance sampler sample counts and bound validation rank skew.
  • fix(dspark): Correct position-0 decay weight and Markov bigram alignment for sample_from_anchor.
  • fix(loss): Compute divergence losses in float32 under BF16 training.
  • fix(eagle3): Always teacher-force TTT tokens.
  • fix(regen): Stable --resume identity and request retries.
  • fix: Apply --draft-attn-impl when loading with --from-pretrained.

New Contributors

Full Changelog: v0.6.0...v0.7.0

Contributors

mgoin, imargulis, and 18 other contributors
Assets 4
Loading