SpecForge: Training speculative decoding models with SGLang

SpecForge: Training speculative decoding models with SGLang#

Author: Ning Zhang
Knowledge level: Intermediate

Speculative decoding is a powerful technique to improve inter-token latency in memory-bound LLM inference, which needs to run two models in parallel. These are the target model (the main LLM model for the AI application) and the small draft model (a smaller, lightweight LLM that runs alongside it to help speed up inference for the main LLM). The small draft model is important in speculative decoding because its ability to correctly predict the tokens for target model acceptance is the key factor in whether the deployment will be successful or not.

SpecForge, a purpose-built ecosystem for training draft models that integrate natively with SGLang, is an open-source application that is listed as a flagship project by the Large Model Systems Organization (LMSYS). This tutorial demonstrates how to run SpecForge draft model training and SGLang speculative decoding inference with the trained draft model on a system with an AMD Instinct™ MI300X GPU.

Prerequisites#

This tutorial was developed and tested using the following setup.

Operating system#

  • Ubuntu 22.04/24.04: Ensure your system is running Ubuntu 22.04 or 24.04.

Hardware#

  • AMD Instinct GPUs: This tutorial was tested on an AMD Instinct MI300X GPU node with eight GPUs; one MI300X GPU is enough for a reduced run (see NUM_GPUS below). Ensure you are using AMD Instinct GPUs or compatible hardware with ROCm support and that your system meets the official requirements.

  • GPUs required: 1 (set NUM_GPUS = 1 in the run settings). The default settings use eight GPUs.

Software#

  • ROCm 6.3+: This tutorial requires ROCm 6.3 or later. Install and verify ROCm by following the ROCm install guide. After installation, confirm your setup using the amd-smi command. AMD also provides prebuilt ROCm Docker images, for example, ROCm PyTorch, ROCm Ubuntu 22.04, and ROCm Ubuntu 24.04. You can use these prebuilt Docker images to reduce the effort of setting up a ROCm environment.

    Note: For ROCm 6.4 and earlier, use the rocm-smi command instead of amd-smi.

  • Docker: Ensure Docker is installed and configured correctly. Follow the Docker installation guide for your operating system.

    Note: Ensure the Docker permissions are correctly configured. To configure permissions to allow non-root access, run the following commands:

    sudo usermod -aG docker $USER
newgrp docker

    Verify Docker is working correctly with:

    docker run hello-world

Hugging Face API access#

  • Obtain an API token from Hugging Face for downloading models.

  • Ensure the Hugging Face API token has the necessary permissions and approval to access the Meta Llama checkpoints.

Set up the SpecForge environment#

In this tutorial, you will work on the prebuilt ROCm PyTorch image as an example. You can also try the other ROCm images.

Step 1: Launch the Docker image#

Launch the Docker container. Replace /path/to/SpecForge_Project with the full path to the directory on your host machine where the SpecForge code and model files are stored. This tutorial uses the SGLang ROCm Docker images to demonstrate SpecForge draft model training and SGLang speculative decoding inference. For better model training performance, you can also try the ROCm PyTorch Training Docker image.

docker run -it --rm \
  --network=host \
  --device=/dev/kfd \
  --device=/dev/dri \
  --group-add=video \
  --ipc=host \
  --cap-add=SYS_PTRACE \
  --security-opt seccomp=unconfined \
  --shm-size 32G \
  -v /path/to/SpecForge_Project:/SpecForge \
  -w /SpecForge/ \
  lmsysorg/sglang:v0.5.18-rocm720-mi30x

Note: Use SGLang 0.5.18. SpecForge’s online training adds hidden-state capture hooks to SGLang (Step 3), and its patch supports SGLang 0.5.18 only; this is the image SpecForge’s AMD guide uses for MI300X.

Note: This command mounts the current directory to the /SpecForge directory in the container. Ensure the notebook file is either copied to this directory before running the Docker command or uploaded into the Jupyter Notebook environment after it starts. Save the token or URL provided in the terminal output to access the notebook from your web browser. You can download this notebook from the AI Developer Hub GitHub repository.

Step 2: Install and launch Jupyter#

Inside the Docker container, install Jupyter using the following command:

pip install jupyter

Start the Jupyter server:

jupyter-lab --ip=0.0.0.0 --port=8888 --no-browser --allow-root

Note: Ensure port 8888 is not already in use on your system before running the above command. If it is, you can specify a different port by replacing --port=8888 with another port number, for example, --port=8890.

Step 3: Install SpecForge#

Install the official SpecForge code the way its ROCm guide describes. A plain install would replace the container’s ROCm PyTorch and SGLang with the CUDA builds SpecForge pins, so install it with --no-deps and use the container’s packages. Then apply SpecForge’s patch to the container’s SGLang: it adds the --enable-spec-capture server options that online training uses to read the target model’s hidden states. Run the following commands inside the Jupyter notebook running within the Docker container:

%%bash
set -e
git clone https://github.com/sgl-project/SpecForge.git
cd SpecForge
git checkout 03f3682ea852b73488d62dab4cd539033cc60e3e  # sgl-project/SpecForge main as of 2026-10-19 (SpecForge has no release tags)
# --no-deps keeps the container's ROCm PyTorch, transformers and SGLang. SpecForge asks for
# Python 3.11, but it runs on this image's Python 3.10.
pip install --no-deps --ignore-requires-python -e .
# Add the hidden-state capture hooks for online training to the container's SGLang 0.5.18
bash scripts/apply_sglang_spec_capture_patch.sh --target v0.5.18

Step 4: Provide your Hugging Face token#

You’ll require a Hugging Face API token to access Llama-3. Generate your token at Hugging Face Tokens and request access for Llama-3 8B. Tokens typically start with “hf_”.

Run the following interactive block in your Jupyter notebook to set up the token:

Note: Uncheck the “Add token as Git credential” option.

from huggingface_hub import notebook_login

# Prompt the user to log in
notebook_login()

from huggingface_hub import HfApi

try:
    api = HfApi()
    user_info = api.whoami()
    print(f"Token validated successfully! Logged in as: {user_info['name']}")
except Exception as e:
    print(f"Token validation failed. Error: {e}")

Choose the GPUs and the run size#

Training and serving run on NUM_GPUS GPUs. Llama-3 8B and its draft model fit on a single MI300X; more GPUs train faster. To try the tutorial quickly, train on only the first TRAIN_SAMPLES conversations and run each benchmark on BENCHMARK_QUESTIONS questions (None uses everything):

NUM_GPUS = 8  # GPUs for training and for the SGLang server
TRAIN_SAMPLES = None  # ShareGPT conversations to train on (None: all ~120k)
BENCHMARK_QUESTIONS = None  # questions per benchmark (None: each benchmark's default, up to 1024)

Make these run-size settings available to the %%bash cells below, which read them from the environment:

Speculative decoding draft model training#

SpecForge is a framework for training speculative decoding models so you can smoothly port them over to the SGLang serving framework to speed up your inference. It offers two methods of training the draft model: online and offline training. Online training freezes the target model and training draft model at same time, which generates auxiliary hidden states on the fly and needs multiple GPUs to achieve better performance. Offline training generates and saves the hidden states using the target model first and then trains the draft model in a separate process. Offline training requires as little as one GPU because it only needs to accommodate the draft model, but it needs a huge amount of disk space. For example, the UltraChat and ShareGPT datasets would need 12TB of storage. Due to the disk size of the AMD Instinct MI300X system, this tutorial uses online training to demonstrate SpecForge.

SpecForge runs online training as two cooperating parts. An SGLang server runs the target model and captures its hidden states for each training prompt, and the trainer reads them through a Mooncake transfer store and updates the draft model. One specforge train command drives both: its producer sends the prompts to the SGLang server and its consumer trains on the captured states.

Data preparation#

SpecForge provides a script to prepare the UltraChat (200k) and ShareGPT (120k) datasets for draft model training. The commands below process the datasets into jsonl (JSON Lines) files, which are the raw dataset for online training, and place them in the cache/dataset directory as <dataset_name>_train.jsonl. With TRAIN_SAMPLES set, --sample-size keeps only that many conversations. To train the model using your own data, prepare the dataset in jsonl format. The schema should look like this:

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

This tutorial uses the ShareGPT and UltraChat datasets as examples.

%%bash
set -e
cd SpecForge
# ultrachat
python scripts/prepare_data.py --dataset ultrachat $SAMPLE_ARG

# sharegpt
python scripts/prepare_data.py --dataset sharegpt $SAMPLE_ARG
wc -l cache/dataset/sharegpt_train.jsonl

Draft vocabulary mapping#

The Llama-3 EAGLE-3 draft model predicts tokens from a reduced vocabulary of 32,000 tokens (draft_vocab_size in configs/llama3-8B-eagle3.json) instead of the target model’s 128,256. The mapping between the two vocabularies keeps the tokens that appear most often in the assistant replies of the training data. Online training needs this mapping as a file, so that every part of the run uses the same one. The cell below builds it from the ShareGPT data with SpecForge’s dataset functions:

%%bash
set -e
cd SpecForge
python - <<'EOF'
from datasets import load_dataset
from transformers import AutoConfig, AutoTokenizer

from specforge.data.preprocessing import build_eagle3_dataset, generate_vocab_mapping_file

tokenizer = AutoTokenizer.from_pretrained("meta-llama/Meta-Llama-3-8B-Instruct")
conversations = load_dataset("json", data_files="cache/dataset/sharegpt_train.jsonl")["train"]
dataset = build_eagle3_dataset(conversations, tokenizer, chat_template="llama3", max_length=2048)
draft_config = AutoConfig.from_pretrained("configs/llama3-8B-eagle3.json")
generate_vocab_mapping_file(
    dataset,
    target_vocab_size=draft_config.vocab_size,
    draft_vocab_size=draft_config.draft_vocab_size,
    cache_dir="cache/vocab_mapping",
    cache_key="llama3-8b-eagle3",
)
EOF

The tutorial runs its servers from the notebook. These helpers start a server in the background, wait until it answers (or stop with its log if it fails), and stop it again:

Hide code cell source

 
import os
import signal
import subprocess
import time
import urllib.request


def server_is_up(url, headers=None):
    try:
        with urllib.request.urlopen(urllib.request.Request(url, headers=headers or {}),
                                    timeout=5):
            return True
    except Exception:
        return False


def ensure_server(command, url, log_path, timeout=1200, headers=None):
    """Start `command` in the background (its output goes to log_path) unless
    url already answers, then wait until it does. Returns the server process,
    or None if a server was already running. Raises, with the end of the log,
    if the server exits or is not ready within timeout seconds."""
    if server_is_up(url, headers):
        print(f"Already running: {url}")
        return None
    log = open(log_path, "w")
    server = subprocess.Popen(command, shell=True, stdout=log, stderr=subprocess.STDOUT,
                              start_new_session=True)
    deadline = time.time() + timeout
    while time.time() < deadline and server.poll() is None:
        if server_is_up(url, headers):
            print(f"Ready after {timeout - int(deadline - time.time())}s: {url}")
            return server
        time.sleep(5)
    tail = "".join(open(log_path).readlines()[-40:])
    state = f"exited with code {server.returncode}" if server.poll() is not None else "timed out"
    raise RuntimeError(f"Server {state} before {url} answered. Last log lines:\n{tail}")


def stop_server(server):
    """Stop a server started by ensure_server, and everything it started."""
    if server is not None and server.poll() is None:
        os.killpg(server.pid, signal.SIGTERM)
        try:
            server.wait(timeout=60)
        except subprocess.TimeoutExpired:
            os.killpg(server.pid, signal.SIGKILL)

Draft model online training#

SpecForge provides training recipes for Llama3/4, Qwen3, and other popular LLM models in examples/configs. This tutorial uses the Llama-3 8B model to demonstrate how to run online training using SpecForge. Online training uses three processes:

  1. A Mooncake master, the store through which the captured hidden states reach the trainer.

  2. The SGLang capture server, which runs the target model. --enable-spec-capture --spec-capture-method eagle3 makes it save the hidden states EAGLE-3 trains on: the final layer and the auxiliary layers in --spec-capture-aux-layer-ids. For EAGLE-3 these are layers [1, num_layers // 2 - 1, num_layers - 4], which is 1 15 28 for the 32 layers of Llama-3 8B. Capture needs --chunked-prefill-size -1 and --disable-radix-cache, and --context-length must exceed the training sequence length (2048).

  3. The trainer, specforge train.

The capture server runs on GPU 0 and the trainer on the other NUM_GPUS - 1 GPUs. With a single GPU, both share GPU 0: the capture server takes a quarter of its memory (--mem-fraction-static 0.25). The MOONCAKE_* variables tell the capture server and the trainer where the Mooncake master is. The cell below starts the Mooncake master and the capture server:

# Where the capture server and the trainer find the Mooncake master
os.environ.update(
    MOONCAKE_LOCAL_HOSTNAME="127.0.0.1",
    MOONCAKE_METADATA_SERVER="http://127.0.0.1:35880/metadata",
    MOONCAKE_MASTER_SERVER_ADDR="127.0.0.1:35551",
    MOONCAKE_PROTOCOL="tcp",
    MOONCAKE_GLOBAL_SEGMENT_SIZE=str(32 << 30),
)
# The capture server uses GPU 0; the trainer uses the other GPUs, or shares GPU 0
os.environ["TRAINER_GPUS"] = ",".join(str(gpu) for gpu in range(1, NUM_GPUS)) or "0"

mooncake_master = ensure_server(
    "mooncake_master --enable_http_metadata_server=true --rpc_port=35551"
    " --http_metadata_server_port=35880 --metrics_port=35903"
    " --enable_metric_reporting=false --default_kv_lease_ttl=500",
    "http://127.0.0.1:35880/health", "mooncake_master.log", timeout=120,
)
capture_server = ensure_server(
    "HIP_VISIBLE_DEVICES=0 python3 -m sglang.launch_server"
    " --model-path meta-llama/Meta-Llama-3-8B-Instruct --skip-tokenizer-init"
    " --tp-size 1 --context-length 4096 --mem-fraction-static 0.25"
    " --chunked-prefill-size -1 --disable-radix-cache"
    " --enable-spec-capture --spec-capture-method eagle3 --spec-capture-aux-layer-ids 1 15 28"
    " --host 127.0.0.1 --port 30000",
    "http://127.0.0.1:30000/health", "capture_server.log", timeout=1800,
)

The training settings come from SpecForge’s Llama-3.1 8B online recipe, examples/configs/online/disaggregated/external/llama3.1-8b-eagle3.yaml: batch size 1 per GPU, learning rate 1e-4 and the sdpa attention backend, which runs on ROCm. It already points at the ShareGPT file and at a capture server on port 30000. The dotted section.field=value arguments override the recipe: the Llama-3 8B target model, the vocabulary mapping, a 2048-token sequence length and the number of trainer GPUs. Because model training takes a long time, you will only run one epoch of training, which might not have the best accuracy. The recipe also stops after 10,000 steps, so training.max_steps is set to the number of steps in one epoch. The trainer’s output goes to llama3-8b-eagle3-online.log; the cell prints the producer’s summary and the loss and accuracy of the last logged step:

%%bash
set -eo pipefail
cd SpecForge
# Llama-3 8B draft model training on the trainer GPUs, for one epoch
TRAIN_GPU_COUNT=$(echo "$TRAINER_GPUS" | tr ',' '\n' | wc -l)
SAMPLES=$(wc -l < cache/dataset/sharegpt_train.jsonl)
HIP_VISIBLE_DEVICES=$TRAINER_GPUS specforge train \
    -c examples/configs/online/disaggregated/external/llama3.1-8b-eagle3.yaml \
    model.target_model_path=meta-llama/Meta-Llama-3-8B-Instruct \
    model.vocab_mapping_path=cache/vocab_mapping/llama3-8b-eagle3.pt \
    data.max_length=2048 \
    training.num_epochs=1 \
    training.max_steps=$(( (SAMPLES + TRAIN_GPU_COUNT - 1) / TRAIN_GPU_COUNT )) \
    deployment.trainer.nproc_per_node=$TRAIN_GPU_COUNT \
    run_id=llama3-8b-eagle3-online \
    output_dir=outputs/llama3-8b-eagle3-online \
    deployment.disaggregated.control_dir=outputs/llama3-8b-eagle3-online/control \
    deployment.disaggregated.consumer_state_dir=outputs/llama3-8b-eagle3-online/consumer-state \
    > llama3-8b-eagle3-online.log 2>&1 || { tail -n 40 llama3-8b-eagle3-online.log; exit 1; }

grep "drive_producer returning" llama3-8b-eagle3-online.log
grep "^step " llama3-8b-eagle3-online.log | tail -n 1 | grep -oE "^step [0-9]+|'(loss|acc)': [0-9.]+" | paste -sd ' '
ls outputs/llama3-8b-eagle3-online/llama3-8b-eagle3-online-latest/

Training is done, so stop the capture server and the Mooncake master to free the GPU:

stop_server(capture_server)
stop_server(mooncake_master)

After the training is done, the checkpoint is in outputs/llama3-8b-eagle3-online/llama3-8b-eagle3-online-latest under the output_dir directory. It holds the training state, which the SGLang speculative decoder cannot load directly, so export the draft model to the model directory format SGLang reads:

%%bash
set -e
cd SpecForge
specforge export --to sglang \
    --checkpoint outputs/llama3-8b-eagle3-online/llama3-8b-eagle3-online-latest \
    --draft-config configs/llama3-8B-eagle3.json \
    --output-dir exports/llama3-8b-eagle3-sglang
ls exports/llama3-8b-eagle3-sglang

(Optional) SGLang speculative decoding inference#

This section shows how to test the trained draft model based on the SGLang framework. If you are already familiar with SGLang speculative decoding inference, you can skip this section.

SGLang speculative decoding server#

The SGLang inference server supports EAGLE speculative decoding through the following options (the command below also sets --cuda-graph-max-bs-decode, the largest decode batch size SGLang captures in a graph):

  1. speculative-draft-model-path: Specifies the draft model. This parameter is required.

  2. speculative-num-steps: Specifies the depth of auto-regressive drafting. This variable increases the speculation range but risks rejection cascades. The default is 5.

  3. speculative-eagle-topk: Specifies the branching factor per step. This improves candidate diversity. It leads to a higher acceptance rate but also causes higher memory/compute consumption. The default is 4.

  4. speculative-num-draft-tokens: Specifies the maximum parallel verification capacity. It causes deeper tree evaluation but leads to higher GPU memory usage. The default is 8.

The settings for these options in the following SGLang server command are intended to demonstrate the speculative decoding functional test and are not about maximizing performance. You can find the best combinations of these parameters in bench_speculative.py.

sglang_server = ensure_server(
    f"python3 -m sglang.launch_server --model meta-llama/Meta-Llama-3-8B-Instruct --speculative-algorithm EAGLE3 --speculative-draft-model-path SpecForge/exports/llama3-8b-eagle3-sglang --speculative-num-steps 3 --speculative-eagle-topk 1 --speculative-num-draft-tokens 4 --mem-fraction-static 0.75 --cuda-graph-max-bs-decode 2 --tp {NUM_GPUS} --context-length 8192 --trust-remote-code --host 0.0.0.0 --port 30000 --dtype bfloat16 --attention-backend triton",
    "http://localhost:30000/health", "sglang_server.log", timeout=1800,
)

SGLang speculative decoding benchmarking#

SpecForge’s specforge benchmark command tests speculative decoding performance with different test datasets. After launching the SGLang inference server, as shown above, it sends the GSM8K, MATH-500, MTBench, and HumanEval prompts to the server and reports the output throughput and the average acceptance length: the number of tokens produced per verification step of the target model. It is 1 without speculative decoding and grows as the target model accepts more of the draft model’s tokens. A draft trained for only a few steps stays close to 1.

%%bash
set -e
# GSM8K
specforge benchmark --model meta-llama/Meta-Llama-3-8B-Instruct --dataset gsm8k $QUESTIONS_ARG

# MATH-500
specforge benchmark --model meta-llama/Meta-Llama-3-8B-Instruct --dataset math500 $QUESTIONS_ARG

# MTBench
specforge benchmark --model meta-llama/Meta-Llama-3-8B-Instruct --dataset mt-bench $QUESTIONS_ARG

# HumanEval
specforge benchmark --model meta-llama/Meta-Llama-3-8B-Instruct --dataset humaneval $QUESTIONS_ARG

Stop the server when you are done benchmarking:

stop_server(sglang_server)

Summary#

Speculative decoding has emerged as a breakthrough for accelerating LLM inference. SpecForge is designed for speculative decoding draft model training and is tightly integrated with the SGLang inference engine. This tutorial demonstrated how to run SpecForge draft model training and test the model using SGLang on AMD GPUs. SpecForge itself is constantly evolving, and AMD is working closely with the open-source community to enhance this work for AMD platforms. This tutorial will hopefully encourage you to try, test, and contribute to SpecForge on AMD GPUs and help shape the future of AI acceleration.

Contents