OUTLETS: Learning Output Length Prediction Alongside Speculative Decoding

Introduction

This repository contains the official implementation of our paper OUTLETS: Output-Length Prediction from Speculative Decoding Backbones.

Datasets

• sharegpt
• alpaca
• lmsys

Models

• meta-llama/Llama-3.2-1B-Instruct
• meta-llama/Llama-3.1-8B-Instruct
• Qwen/Qwen3-30B-A3B

Default Path to put the huggingface models and datasets

• model: /data/models/
• dataset: /data/datasets/

Installation

conda create -n outlets python=3.11
conda activate outlets
pip install -r requirements.txt
# Install SpecForge from source: https://docs.sglang.io/SpecForge/get_started/installation.html

Data Preparation

1. Start VLLM Server

CUDA_VISIBLE_DEVICES="0,1,2,3" vllm serve /data/models/Qwen3-30B-A3B/ \
  --host 0.0.0.0 \
  --port 8000 \
  --max-model-len 5000 \
  --data-parallel-size 4

2. Regenerate Datasets

cd scripts
python gen_sharegpt.py \
  --vllm_url http://127.0.0.1:8000/v1/chat/completions \
  --tokenizer_path /data/models/Qwen3-30B-A3B
# For other datasets: gen_alpaca.py, gen_lmsys.py

3. Split and Clean Data

cd scripts
python split_data.py
python preprocess.py

Training

1. Build Vocabulary Cache

cd outlets
python build_vocab_cache.py \
  --dataset-name sharegpt \
  --data-path ../mydata/sharegpt_train_2048.jsonl \
  --tokenizer-path /data/models/Qwen3-30B-A3B \
  --draft-vocab-size 32000 \
  --vocab-size 151936 \
  --num-proc 16

2. Train OUTLETS Model

deepspeed --master_port=29000 --include localhost:0 main.py \
  --deepspeed_config ds_config.json \
  --dataset sharegpt \
  --task join \
  --ols

3. Train MLP Baseline (Optional)

deepspeed --master_port=29000 --include localhost:0 main.py \
  --deepspeed_config ds_config.json \
  --dataset sharegpt \
  --task join \
  --mlp

4. Monitor Training with TensorBoard

tensorboard --logdir=runs --port=6006

Baseline Methods

Proxy-based Methods

cd baseline
python proprecess.py # need update the args in scripts

1. BERT Predictor

python bert_predictor.py \
  --dataset sharegpt \
  --output_dir bert_sharegpt \
  --epoch 10 \
  --device "cuda:0"

2. OPT Predictor

python opt_predictor.py \
  --dataset sharegpt \
  --output_dir opt_sharegpt \
  --epochs 10 \
  --device "cuda:0" \
  --num_classes 50

PIA

3. Start VLLM Server for PIA

CUDA_VISIBLE_DEVICES=0 vllm serve /data/models/Qwen3-30B-A3B/ \
  --host 0.0.0.0 \
  --port 8000 \
  --max-model-len 2048

4. Run PIA Prediction (in another terminal)

cd baseline
python pia_prediction.py \
  --dataset sharegpt \
  --data_split test \
  --vllm_url http://127.0.0.1:8000/v1/chat/completions \
  --max_tokens 40 \
  --temperature 0.0 \
  --num_threads 40 \
  --prompt_template suffix \
  --output_dir sharegpt_pia_results