DOCSEARCHbyAlgolia
Submit Docs

SpecForge: Speculative Decoding Training Framework

repository·main·Indexed 21 days ago

https://github.com/sgl-project/specforge

A specialized training framework for speculative decoding models designed for direct compatibility with the SGLang serving framework to enable high-speed LLM inference. SpecForge supports various training strategies including EAGLE3, P-EAGLE, DFlash, Domino, and DSpark, and provides tools for managing training recipes via YAML, running benchmarks with `bench_eagle3.py`, and handling both online and offline data capture workflows.

Tokens
44.2K
Snippets
117
Records
177
Agent score
70%

What's inside SpecForge

  1. Overview of SpecForge

    main

    SpecForge is a framework developed by the SGLang team for training speculative decoding models. It is designed to allow developers to train models that can be smoothly ported to the SGLang serving framework for accelerated inference.

    Key features include:

    • Direct SGLang Compatibility: No additional porting effort is required to use trained models in SGLang.
    • Unified Runtime: Supports local offline training and server-only online-disaggregated training through a single runtime.
    • Flexible Topologies: Supports various data, tensor, and sequence parallel topologies.
    • Regular Maintenance: The project is actively maintained to ensure code is runnable out-of-the-box.

  2. What is SpecForge?

    main

    SpecForge is a speculative decoding framework designed to speed up inference without losing performance. It is maintained by the SGLang team and is built to be directly compatible with SGLang serving.

    Key capabilities include:

    • SGLang Compatibility: Models are ready for SGLang serving out-of-the-box.
    • Unified Runtime: Supports SGLang-server online training as well as local/disaggregated offline training.
    • Advanced Training Features: Includes consumer DP (Data Parallelism), offline USP (Unsupervised Speculative Pre-training), evaluation, checkpoint selection, and portability across CUDA, ROCm, and Ascend hardware.

  3. What is Speculative Decoding and how does it reduce latency?

    main

    Speculative decoding is a technique used to reduce LLM inference latency. Standard autoregressive decoding is bottlenecked by memory bandwidth because the entire model weight set must be loaded for every single token generated.

    Speculative decoding addresses this by using a small, fast Draft Model to predict several future tokens in advance. These tokens are then verified in parallel by the larger Target Model. Because the verification stage is also memory-bound, verifying multiple tokens takes nearly the same amount of time as generating a single token, resulting in a significant speedup in the overall decoding process.

  4. What is EAGLE3 and how does it work?

    main

    EAGLE3 is a state-of-the-art speculative decoding mechanism integrated into SGLang. Unlike traditional speculative decoding that uses a small language model from the same family as the target model (e.g., using an 8B model to draft for a 70B model), EAGLE3 uses a separate small speculator model that operates in the feature space of the target model.

    Key technical characteristics include:

    1. Feature-based Drafting: Instead of feeding raw tokens to the draft model, EAGLE3 extracts three hidden states from different depths of the target model, concatenates them into a single feature vector, and feeds this vector to the draft model to generate predictions.
    2. Training-time Test: The model is trained by simulating the autoregressive generation process, computing loss between predicted and ground truth sequences to reduce error accumulation and improve acceptance rates.
    3. Dynamic Draft Tree: It utilizes a dynamic draft tree (based on EAGLE2) to store only the candidate tokens most likely to be accepted by the target model, optimizing the acceptance rate.

  5. Choose between Online and Offline training modes

    main

    SpecForge supports two primary data modes based on how features are acquired:

    ModeTarget during trainingDisk useData configDescription
    OnlineExternal/managed SGLang capture serverLowtrain_data_path or prompts_pathCaptures target features while the run is active. Keeps target inference available during training.
    OfflineNot loaded by the trainerHighhidden_states_pathReads precomputed feature checkpoints. Only the draft model must fit on training GPUs.

    Data Source Selection: Set exactly one of the following in the data section:

    • data.train_data_path: Raw conversation or preformatted online data.
    • data.prompts_path: Pre-tokenized online JSONL (contains input_ids and loss_mask).
    • data.hidden_states_path: Precomputed offline feature checkpoints.

  6. Understand DataFlow Contracts and the No-Tensor Boundary

    main

    SpecForge uses a strict separation between the control plane and the data plane via specific data records (contracts).

    To ensure the control plane remains lightweight and testable without heavy dependencies like torch, a strict no-tensor boundary is enforced.

    • Control Plane Records: Must be stdlib-only and carry only metadata. They are defined as @dataclass(frozen=True) to ensure immutability. Examples include PromptTask and SampleRef.
    • Data Plane Records: The only place where actual tensors are permitted. TrainBatch is the primary contract that carries tensors and is deliberately not frozen, as it lives exclusively on the trainer/data-plane side.

    The assert_no_tensors guard is used to enforce this boundary by recursively checking records for any objects that look like tensors (e.g., objects with dtype, shape, and device attributes).

  7. Understand Online vs Offline disaggregated modes

    main

    SpecForge supports two primary modes for disaggregated training:

    Online Mode

    • Data Plane: Uses Mooncake.
    • Producer: Requires a positive deployment.disaggregated.producer_segment_size.
    • Consumer/Roles: The consumer role and online roles use a producer_segment_size of zero because they do not own feature allocations.
    • Prerequisites: Requires an already-running Mooncake deployment and a patched SGLang capture server. These services are typically managed externally and are not started/stopped by specforge train unless using a managed-local recipe.

    Offline Mode

    • Data Plane: May use either a typed shared_dir store or Mooncake.
    • Producer: Requires a positive deployment.disaggregated.producer_segment_size.
    • Consumer: Uses a producer_segment_size of zero.

  8. Understand the SpecForge runtime architecture

    main

    The specforge.runtime package serves as the transport substrate for all specforge train topologies. It is divided into three functional layers:

    1. Contracts (contracts.py): Defines metadata-only records like PromptTask and SampleRef, as well as the TrainBatch record used for carrying tensors across boundaries.
    2. Control Plane (control_plane/): Manages prompt scheduling, online reference (ref) staging, the single online-consumer ledger, and optimizer-boundary DP (Data Parallel) acknowledgement.
    3. Data Plane (data_plane/): Manages feature stores, fixed offline references, consume-once online channels, rank inboxes, and the FeatureDataLoader.

    Note that training and inference compute are handled in separate packages (specforge.training and specforge.inference) and are not part of the runtime package itself.

  9. Prepare custom datasets in Conversation or Pre-formatted Text format

    main

    You can use your own datasets in SpecForge by following one of two supported formats.

    Option 1: Conversation Format

    Use this for standard chat datasets in JSONL format:

    {
    "id": "xxxx",
    "conversations": [
        {
            "role": "user | assistant",
            "content": "The message content"
        }
    ]
}

    Option 2: Pre-formatted Text Format

    Use this if you have conversations already formatted with a specific chat template (e.g., raw model generations). This is useful for matching the exact training distribution of a target model.

    {
    "id": "xxxx",
    "text": "<|im_start|>system\nYou are a helpful assistant.<|im_end|>\n<|im_start|>user\nHello<|im_end|>\n<|im_start|>assistant\nHi there!<|im_end|>\n"
}

    Configuration for Pre-formatted Text: To use the text format, you must update your training configuration:

    1. Set data.is_preformatted: true.
    2. Set data.chat_template to the template used to create the text (e.g., llama3). SpecForge uses this to identify assistant spans for loss masking.
    3. Provide the path via data.train_data_path.
    # Example configuration for pre-formatted text
data:
  train_data_path: ./your_preformatted_dataset.jsonl
  is_preformatted: true
  chat_template: llama3
    specforge train --config ./my-eagle3-disaggregated.yaml

  10. Supported training strategy combinations

    main

    The SpecForge unified runtime supports text training across various strategies and deployment modes. Note that VLM training (e.g., Qwen2.5-VL) is currently unsupported; the runtime accepts text inputs only.

    Strategy Compatibility Matrix

    StrategySGLang server onlineLocal/dataflow offlineDisaggregated offline
    EAGLE3Yes, consumer DPYes, DP + USPYes, consumer DP
    DFlashYes, consumer DPYes, DPYes, consumer DP
    DominoYes, consumer DPYes, DPYes, consumer DP
    DSparkYes, consumer DPYes, DPYes, consumer DP
    P-EAGLEYes, consumer DP, batch size 1NoNo

    Key Constraints

    • Attention Backends:
      • EAGLE3: sdpa, flex_attention, fa, or offline usp.
      • P-EAGLE: requires flex_attention.
      • DFlash, Domino, DSpark: eager, sdpa, or flex_attention.
    • P-EAGLE: Requires training.batch_size=1.
    • Online Evaluation: Not supported. Evaluation requires precomputed offline features via data.eval_hidden_states_path.
    • EAGLE3 Offline: Derives a deterministic vocabulary mapping from the feature corpus if model.vocab_mapping_path is empty. Disaggregated runs require an explicit shared mapping.

  11. Compare SpecForge training modes (Colocated, Disaggregated, Online)

    main

    SpecForge supports three primary operational modes. Choose the mode based on your data availability and iteration requirements:

    ModeProducer sideConsumer reference sourceFeature storeIteration contract
    Colocated offlinePrecomputed feature filesFixed SampleRef listLocalFeatureStore reads file:// refsRe-iterable; supports epochs and checkpoint resume
    Disaggregated offlineProducer ingests files and writes a static manifestFixed manifest refsShared directory or MooncakeRe-iterable; supports DP/multi-node epochs and checkpoint resume
    OnlinePatched SGLang server writes tensors; producer publishes refsPer-rank StreamingRefQueue inboxMooncakeConsume once; consumer-only recovery supported; no producer resume

    Note on Online Epochs: In online mode, training.num_epochs controls how many prompt passes the producer creates. The consumer always iterates the resulting stream exactly once and cannot replay it for a second epoch.

  12. Configure parallel topologies for training

    main

    The launcher creates process groups based on your typed run configuration. The world size must be divisible by training.sp_ulysses_size * training.sp_ring_size.

    • Online Targets: TP/EP belongs to the external SGLang capture server. Online consumers keep training.tp_size and both SP sizes at 1; each trainer rank receives a disjoint feature stream.
    • Offline Consumers: Keep training.tp_size at 1. Without USP, every trainer rank receives a disjoint reference shard (Data Parallelism).
    • EAGLE3 Offline (USP): Can set training.attention_backend: usp and choose training.sp_ulysses_size and training.sp_ring_size. USP currently uses training.batch_size: 1. SP peers share one sequence while draft-DP groups receive disjoint references.