Skip to content

data_preparation.md

Latest commit

 

History

History
350 lines (271 loc) · 14.1 KB

data_preparation.md

File metadata and controls

350 lines (271 loc) · 14.1 KB
Raw

📝 Data Preparation

📍 Overview

Data is an important aspect of speculative decoding as the quality of the dataset directly affects the acceptance rate of the draft model. In this section, we will introduce how to prepare the dataset for both online and offline training.

☁️ Canonical Dataset Presets

scripts/prepare_data.py --dataset NAME exposes the same 19 text presets used by the checked-in recipes:

Family Presets
General chat ultrachat, sharegpt, eaglechat, perfectblend, perfectblend-llama3.1-8b-instruct, perfectblend-llama3.3-70b-instruct, perfectblend-llama4-scout-instruct, perfectblend-llama4-maverick-instruct, magpie-qwen2.5-pro-1m-v0.1, nebius-llama31-8b-infinity-instruct
Reasoning, math, and code opc, gsm8k, hendrycks_math, math_qa, codealpaca-20k, opencodeinstruct, magicoder-evol-instruct, sciq, camel

VLM data preparation and training, including ALLaVA and ShareGPT4V, are not supported.

Every successful preset writes the same stable id + conversations JSONL contract. Use --split-eval to create a deterministic 95/5 train/eval split; opc additionally accepts --opc-subset.

Run the script from the repository root:

# ultrachat
python scripts/prepare_data.py --dataset ultrachat

# sharegpt
python scripts/prepare_data.py --dataset sharegpt

By default, each command writes <preset>_train.jsonl under cache/dataset. Use --output-path to choose another output directory and --sample-size to cap the number of source rows.

For a local ShareGPT-format dataset, pass a JSON or JSONL file with --data-path:

python scripts/prepare_data.py \
    --dataset sharegpt \
    --data-path ./raw_sharegpt.jsonl \
    --output-path ./cache/dataset

Each raw row must contain an id and a conversations list whose messages use ShareGPT's from and value keys. Custom paths are intentionally limited to the sharegpt preset.

↩️ Regenerate Datasets

When training speculative decoding draft models for a specific target model, instead of using the original dataset, we can regenerate the assistant responses using the target model to better align the draft model with the target model's output distribution. This will improve the acceptance rate of the draft model and the overall performance of the speculative decoding. According to the EAGLE1 paper, the EAGLE method is not very sensitive to the dataset quality, which means the performance is still good even if you use the original dataset. However, if you are looking for optimal performance in the production environment, it is recommended to regenerate the dataset using the target model.

The regeneration utility uses the OpenAI-compatible client to call SGLang's HTTP API. Install it in the source-checkout environment before running the repository script:

pip install -e '.[data]'

We can follow the following steps to regenerate the dataset. In the example below, we will use meta-llama/Llama-3.1-8B-Instruct as an example, you can replace it with your own target model.

  1. Start the SGLang server for the target model.
python3 -m sglang.launch_server \
    --model-path meta-llama/Llama-3.1-8B-Instruct \
    --cuda-graph-max-bs 128 \
    --dtype bfloat16 \
    --mem-fraction-static 0.8 \
    --port 30000
  1. Regenerate the dataset using the regenerate_train_data.py script.
python scripts/regenerate_train_data.py \
    --model meta-llama/Llama-3.1-8B-Instruct \
    --concurrency 128 \
    --max-tokens 98304 \
    --server-address localhost:30000 \
    --temperature 0.8 \
    --input-file-path ./cache/dataset/sharegpt_train.jsonl \
    --output-file-path ./cache/dataset/sharegpt_train_regen.jsonl

For reasoning models, add --reasoning save to store reasoning_content in the regenerated dataset. To use a reasoning model with thinking disabled, add --reasoning disable, which forwards chat_template_kwargs.enable_thinking=false to the SGLang server and does not save reasoning_content.

Multi-turn structured reasoning

A regenerated multi-turn reasoning conversation stores several assistant targets in one row. When training uses last-turn-only loss masking, only the final assistant target is supervised. Convert the validated conversation-level output into one generation-event row per assistant turn so every reasoning target is trained with the visible history available at its serving boundary:

python scripts/expand_reasoning_conversations.py \
    --input-file-path ./cache/dataset/sharegpt_train_regen_reasoning.jsonl \
    --output-file-path ./cache/dataset/sharegpt_train_regen_reasoning_exploded.jsonl

Each event ends at its current assistant target and preserves that turn's reasoning_content and visible content. Historical assistant messages keep only visible content; their hidden reasoning is removed from the event context. Train the exploded output with data.train_only_last_turn: true in the typed run config.

The converter accepts only successful rows with non-empty IDs, message content, and assistant reasoning_content. Invalid input is written to a skipped JSONL; if any turn is invalid, the entire source conversation is skipped. Output files must be fresh and distinct from the input.

For maximum performance, we recommend to scale the number of GPUs to regenerate the dataset in data parallel mode. To do this, you can simply add more server addresses to the --server-address argument, e.g. --server-address localhost:30000 localhost:30001 localhost:30002 localhost:30003.

Qwen ShareGPT recipes

The Qwen recipe wraps regeneration, complete-row accounting, and output validation. It expects one or more SGLang servers to already be running; it does not launch or stop the servers itself.

For Qwen3-8B non-reasoning regeneration, start a target server in one terminal:

python -m sglang.launch_server \
    --model-path Qwen/Qwen3-8B \
    --tp-size 1 \
    --dtype bfloat16 \
    --mem-fraction-static 0.8 \
    --host 0.0.0.0 \
    --port 30000

After curl --fail http://127.0.0.1:30000/health succeeds, run the recipe from the SpecForge repository root:

MODEL_PROFILE=qwen3-8b \
INPUT_FILE=./cache/dataset/sharegpt_train.jsonl \
OUTPUT_FILE=./cache/dataset/sharegpt_train_regen_qwen3_8b_temperature0_non_reasoning.jsonl \
SERVER_ADDRESSES="localhost:30000" \
bash examples/data_regeneration/run_qwen_sharegpt_regeneration.sh

This profile sets temperature to zero, disables thinking through chat_template_kwargs.enable_thinking=false, and rejects non-empty structured reasoning in successful rows.

For Qwen3.6-27B structured-reasoning regeneration, start SGLang with the Qwen3 reasoning parser so the OpenAI-compatible response includes reasoning_content:

python -m sglang.launch_server \
    --model-path Qwen/Qwen3.6-27B \
    --tp-size 1 \
    --dtype bfloat16 \
    --mem-fraction-static 0.8 \
    --reasoning-parser qwen3 \
    --host 0.0.0.0 \
    --port 30000

After curl --fail http://127.0.0.1:30000/health succeeds, run:

MODEL_PROFILE=qwen3.6-27b \
INPUT_FILE=./cache/dataset/sharegpt_train.jsonl \
OUTPUT_FILE=./cache/dataset/sharegpt_train_regen_qwen3.6-27b_temperature0_reasoning.jsonl \
SERVER_ADDRESSES="localhost:30000" \
bash examples/data_regeneration/run_qwen_sharegpt_regeneration.sh

The default maximum completion lengths are 4096 tokens for Qwen3-8B and 32768 tokens for Qwen3.6-27B. The server context length must accommodate both the rendered prompt and MAX_TOKENS; override MAX_TOKENS when using a shorter server context.

INPUT_FILE may point to another dataset without changing the recipe as long as it is JSONL in the conversation format documented below, with a non-empty string id and alternating user/assistant messages. Always choose a fresh OUTPUT_FILE: the recipe deliberately refuses to overwrite or resume an existing run.

For an output such as dataset_regen.jsonl, the recipe produces:

  • dataset_regen.jsonl: successful regenerated training rows;
  • dataset_regen_error.jsonl: request or generation errors;
  • dataset_regen_skipped.jsonl: schema-invalid inputs and outputs excluded by the selected reasoning contract.

The run passes accounting only when success + error + skipped == input. It prints the success fraction without a fixed minimum threshold, then strictly validates every successful row, including its conversation structure, reasoning contract, and absence of raw <think> markers.

To distribute regeneration across independently launched servers, provide all addresses as a space-separated list, for example:

SERVER_ADDRESSES="localhost:30000 localhost:30010" \
CONCURRENCY=64 \
bash examples/data_regeneration/run_qwen_sharegpt_regeneration.sh

🤩 Prepare your own dataset

Besides the provided datasets, you can also prepare your own dataset. We support two formats:

Option 1: Conversation Format

You should prepare the dataset in jsonl format and the schema should look like this:

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

Option 2: Pre-formatted Text Format

If you already have conversations formatted with a specific chat template, you can use the pre-formatted text directly:

{
    "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"
}

This format is useful when you have pre-formatted prompts that were used during training of the target model and have raw generations from the target model.

To use preformatted datasets, set data.is_preformatted: true in the run config. data.chat_template is still required and must match the template used to create the text. SpecForge uses it to identify assistant spans and build the loss mask.

# After copying a disaggregated online example YAML, set these data fields:
# data:
#   train_data_path: ./your_preformatted_dataset.jsonl
#   is_preformatted: true
#   chat_template: llama3
specforge train --config ./my-eagle3-disaggregated.yaml

💾 Prepare offline target features

Offline EAGLE3, DFlash, Domino, and DSpark runs consume target features created by scripts/prepare_hidden_states.py. Select the same strategy and draft configuration that the later training recipe uses; these values determine both the target layers captured and the checkpoint schema. For example, prepare the checked-in Qwen3-8B DFlash recipe with:

torchrun --nproc_per_node=8 \
    scripts/prepare_hidden_states.py \
    --strategy dflash \
    --target-model-path Qwen/Qwen3-8B \
    --draft-model-config configs/qwen3-8b-dflash.json \
    --data-path ./cache/dataset/sharegpt_train.jsonl \
    --output-path ./cache/hidden_states/qwen3-8b-dflash-sharegpt \
    --chat-template qwen \
    --max-length 3072 \
    --tp-size 1 \
    --batch-size 32

Use these strategy/model pairs for the checked-in local offline recipes:

Strategy Target model Draft config Output path used by the recipe
EAGLE3 Qwen/Qwen3-8B configs/qwen3-8b-eagle3.json cache/hidden_states/qwen3-8b-sharegpt
DFlash Qwen/Qwen3-8B configs/qwen3-8b-dflash.json cache/hidden_states/qwen3-8b-dflash-sharegpt
Domino Qwen/Qwen3-8B configs/qwen3-8b-domino.json cache/hidden_states/qwen3-8b-domino-sharegpt
DSpark Qwen/Qwen3-4B configs/qwen3-4b-dspark.json cache/hidden_states/qwen3-4b-dspark-sharegpt

--strategy defaults to eagle3 for compatibility. Pass it explicitly in reproducible jobs. --draft-model-config is required for Domino and DSpark; passing the same explicit config used by training is recommended for every strategy. Capture layers come from the resolved draft config. EAGLE3 retains target-derived defaults for legacy configs that do not define eagle_config.eagle_aux_hidden_state_layer_ids.

Each output record contains the strategy's exact offline feature contract:

Strategy Tensors in each .ckpt or .ckpt.gz record
EAGLE3 input_ids, loss_mask, hidden_state, aux_hidden_state
DFlash and Domino input_ids, loss_mask, hidden_states
DSpark input_ids, loss_mask, hidden_states, target_last_hidden_states

For the DFlash family, hidden_states concatenates the target layers selected by the draft config. DSpark additionally stores the target model's final hidden state for its L1 and confidence objectives. Keep each strategy in a separate output directory; the offline reader validates the contract instead of silently adapting incompatible features.

D-PACE uses training.strategy: dflash and the DFlash feature schema. DTA also uses --strategy dflash, but feature preparation must receive its DTA draft config so the captured layer contract matches the training run.

For preformatted input, add --is-preformatted to the same command and keep --chat-template aligned with the template already applied to the text:

torchrun --nproc_per_node=8 \
    scripts/prepare_hidden_states.py \
    --strategy eagle3 \
    --target-model-path meta-llama/Llama-3.1-8B-Instruct \
    --draft-model-config configs/llama3.1-8b-eagle3.json \
    --data-path ./your_preformatted_dataset.jsonl \
    --output-path ./cache/hidden_states/llama3.1-8b-eagle3 \
    --chat-template llama3 \
    --is-preformatted \
    --max-length 2048

The preparation world size only parallelizes capture. The subsequent offline run still uses the one specforge train entry and self-launches the data-parallel or EAGLE3 USP topology recorded under deployment.trainer. Launch the matching recipe after its data.hidden_states_path points at the generated directory:

specforge train --config examples/configs/qwen3-8b-dflash-offline.yaml

See the Training guide for the complete run schema and supported combinations.

➕ Handling Multiple Datasets

If you have multiple datasets, you can just merge them into the one jsonl file. For example, you can do something like this

cat dataset1.jsonl dataset2.jsonl > merged_dataset.jsonl