praful-goel/speculative_decoding_models

Model Architecture & Training

Both models are Decoder-only Transformers (GPT-style) implemented in pure PyTorch.

Hyperparameters	Main Model (Target)	Draft Model (Small)	Draft Model (Medium)
Parameters	~150M	~30M	~70M
Layers	12	2	6
Heads	12	4	8
Embedding Dim	768	256	512
Context Length	1024	1024	1024
Vocab Size	50304	50304	50304
Droupout	0.1	0.1	0.1
Dataset	OpenWebText (Sample)	OpenWebText (Sample)	OpenWebText (Sample)

Key Components

1. Token + Position Embeddings

   • Learned token embedding table: (vocab_size, n_embd)
   • Rotary Position Embedding: parameter-free sinusoidal rotations on attention queries and keys

2. Decoder Blocks Each block consists of:

   • Masked Self-Attention (MHA/MQA/GQA)
   • Key-Value (KV) Caching for O(1) generation latency
   • Rotary Positional Embedding to understand relative positions
   • Feed-Forward Network (Linear → GELU → Linear → Dropout)
   • Residual connection
   • Pre-norm RMSNorm

3. Final RMSNorm + Language Modeling Head

   • Linear projection from n_embd → vocab_size to produce logits

Training Strategy

Main Model

The main model was trained in 4 distinct phases to hande hardware constraints and optimize convergence

Phases	Focus	Max Iterations	Warmup Steps	Eval Iterations	Eval Interval	Accumulation Steps	Learning Rate	Weight Decay	Batch Size
Phase 1	Initial Warmup	10_000	200	20	500	32	3e-4	0.1	8
Phase 2	Main Pre-Training	30_000	0	20	1_000	32	3e-4	0.15	8
Phase 3	Large-Batch Scaling	50_000	1_000	10	2_000	16	6e-5	0.1	16
Phase 4	Convergence	40_000	500	10	2_000	16	2e-5	0.08	16

Draft Model (Small)

The draft model (small) was trained only once

max_iters = 40_000
warmup_steps = 1_000
eval_iter = 20
eval_interval = 1_000
accumulation_steps = 16
base_lr = 3e-4
weight_decay=0.1
batch_size = 16

Draft Model (Medium)

The draft model (medium) was trained only once

max_iters = 40_000
warmup_steps = 2_000
eval_iters = 20
eval_interval = 2_000
accumulation_steps = 16
base_lr = 3e-4
weight_decay = 0.1