Speculator Model Training Pipeline

This pipeline trains Llama-3.2-1B-Instruct as a speculator (draft) model for Llama-3.1-8B-Instruct (target model).

Architecture Overview

┌─────────────────────────────────────────────────────────────┐
│                     Training Pipeline                       │
├─────────────────────────────────────────────────────────────┤
│                                                             │
│  Target Model                    Draft Model                │
│  ┌──────────────────────┐       ┌──────────────────────┐    │
│  │ Llama-3.1-8B         │       │ Llama-3.2-1B         │    │
│  │ SGLang Engine        │       │ + LoRA Adapters      │    │
│  │ TP=4 on GPUs 0-3     │       │ Training on GPUs 4-7 │    │
│  │                      │       │                      │    │
│  │ Generates:           │       │ Learns from:         │    │
│  │ - Completions        │──────>│ - Hidden States(Optional) │  
│  │ - Hidden States      │       │ - Generated Text     │    │
│  └──────────────────────┘       └──────────────────────┘    │
│                                                             |
│  Loss = α * Distillation(Optional) + β * Cross-Entropy      │
│         (Hidden State MSE)   (Language Modeling)            │
└─────────────────────────────────────────────────────────────┘

Key Features

1. Efficient Inference with SGLang

• Batch Size: 32 (configurable via --batch_size)
• SGLang is optimized for continuous batching and high throughput
• Larger batch sizes improve GPU utilization during teacher inference
• TP=4 enables distributed inference across 4 GPUs

2. LoRA Fine-tuning

• Rank: 64 (configurable via --lora_r)
• Alpha: 128 (configurable via --lora_alpha)
• Target Modules: All attention and MLP layers
  • q_proj, k_proj, v_proj, o_proj (attention)
  • gate_proj, up_proj, down_proj (MLP)
• Memory efficient: Only trains ~1% of parameters

3. Knowledge Distillation (Optional)

• Disabled by default due to different hidden dimensions (8B: 4096, 1B: 2048)
• Distillation Loss: MSE between hidden states (α=0.5) - only if models have matching dims
• Cross-Entropy Loss: Standard language modeling (β=0.5)
• Temperature: 2.0 for softer probability distributions
• Use --use_distillation flag to enable (requires projection layer for different dims)

4. Training Configuration

• Effective Batch Size: 64 (4 GPUs × train_batch_size × gradient_accum)
• Learning Rate: 2e-4 with cosine scheduling
• Warmup: 3% of total steps
• Gradient Clipping: Max norm of 1.0
• Weight Decay: 0.01

5. Data Pipeline

• Reads from /data/shared/datasets/sft_llama/*/batch_input.jsonl
• Extracts user prompts from the messages field
• Applies Llama chat template automatically
• Supports sampling and shuffling

Usage

Quick Start

chmod +x launch_training.sh
./launch_training.sh

Custom Configuration

python train_speculator.py \
    --target_model "meta-llama/Llama-3.1-8B-Instruct" \
    --draft_model "meta-llama/Llama-3.2-1B-Instruct" \
    --data_dir "/data/shared/datasets/sft_llama" \
    --output_dir "./checkpoints_speculator" \
    --batch_size 32 \
    --train_batch_size 4 \
    --gradient_accumulation_steps 4 \
    --num_epochs 3 \
    --learning_rate 2e-4 \
    --lora_r 64 \
    --lora_alpha 128 \
    --target_gpus "0,1,2,3" \
    --draft_gpus "4,5,6,7" \
    --save_steps 500

Testing with Limited Data

python train_speculator.py \
    --max_samples 10000 \
    --num_epochs 1 \
    --save_steps 100

Arguments Reference

Model Configuration

• --target_model: Path/name of target (teacher) model
• --draft_model: Path/name of draft (student) model
• --max_seq_length: Maximum sequence length (default: 2048)

Training Hyperparameters

• --batch_size: Batch size for SGLang inference (default: 32)
  • Larger values improve throughput but require more VRAM
  • Recommended: 32-64 for 4×A100/H100
• --train_batch_size: Training batch size per GPU (default: 4)
• --gradient_accumulation_steps: Gradient accumulation (default: 4)
• --num_epochs: Number of training epochs (default: 3)
• --learning_rate: Learning rate (default: 2e-4)
• --warmup_ratio: Warmup ratio (default: 0.03)
• --scheduler_type: LR scheduler type: cosine or linear (default: cosine)

Loss Configuration

• --use_distillation: Enable distillation loss (flag, default: disabled)
  • Note: Only use if teacher and student have same hidden dimensions
  • Llama-3.1-8B (4096) vs Llama-3.2-1B (2048) won't work without projection
• --distillation_weight: Weight for distillation loss (default: 0.5)
• --ce_weight: Weight for cross-entropy loss (default: 0.5)
• --temperature: Temperature for distillation (default: 2.0)

LoRA Configuration

• --lora_r: LoRA rank (default: 64)
• --lora_alpha: LoRA alpha (default: 128)
• --lora_dropout: LoRA dropout (default: 0.05)

GPU Configuration

• --target_gpus: GPU IDs for target model (default: "0,1,2,3")
• --draft_gpus: GPU IDs for draft model (default: "4,5,6,7")

Checkpointing

• --output_dir: Output directory for checkpoints
• --save_steps: Save checkpoint every N steps (default: 500)
• --save_total_limit: Max checkpoints to keep (default: 3)
• --logging_steps: Log every N steps (default: 10)

Data Configuration

• --data_dir: Directory containing batch_input.jsonl files
• --max_samples: Limit number of training samples (optional)

Output Structure

checkpoints_speculator/
├── checkpoint-500/
│   ├── adapter_config.json
│   ├── adapter_model.bin
│   ├── training_state.pt
│   └── tokenizer files...
├── checkpoint-1000/
├── checkpoint-epoch-0/
├── checkpoint-epoch-1/
└── training_log.txt

Monitoring Training

Watch the training log in real-time:

tail -f checkpoints_speculator/training_log.txt

Key metrics to monitor:

• Total Loss: Should decrease steadily
• Distillation Loss: Measures hidden state similarity
• CE Loss: Measures language modeling quality
• Learning Rate: Should follow warmup → decay schedule

Loading Trained Model

from transformers import AutoModelForCausalLM, AutoTokenizer
from peft import PeftModel

# Load base model
base_model = AutoModelForCausalLM.from_pretrained(
    "meta-llama/Llama-3.2-1B-Instruct",
    torch_dtype=torch.bfloat16,
    device_map="auto"
)

# Load LoRA weights
model = PeftModel.from_pretrained(
    base_model,
    "./checkpoints_speculator/checkpoint-epoch-2"
)

# Optionally merge for faster inference
model = model.merge_and_unload()

# Load tokenizer
tokenizer = AutoTokenizer.from_pretrained(
    "./checkpoints_speculator/checkpoint-epoch-2"
)

Troubleshooting

Out of Memory (OOM)

1. SGLang OOM (Target Model):

   • Reduce --batch_size
   • Reduce --max_seq_length

2. Training OOM (Draft Model):

   • Reduce --train_batch_size
   • Increase --gradient_accumulation_steps
   • Reduce --lora_r

Slow Training

1. Increase --batch_size for better SGLang throughput
2. Enable gradient checkpointing (requires code modification)
3. Use fewer data samples with --max_samples

Poor Convergence

1. For mismatched models (default setup):

   • Distillation is disabled by default (different hidden dims)
   • Focus on cross-entropy loss only (standard supervised fine-tuning)
   • The model learns from the teacher's generated text, not hidden states

2. Tune learning rate:

   • Try --learning_rate 1e-4 or --learning_rate 5e-4
   • Increase --warmup_ratio to 0.05 or 0.1

3. If using distillation (with --use_distillation):

   • Adjust loss weights:
     • Increase --distillation_weight to focus on mimicking teacher
     • Increase --ce_weight to focus on language modeling
   • Adjust temperature:
     • Higher --temperature (e.g., 3.0) for softer distributions
     • Lower --temperature (e.g., 1.0) for sharper distributions

Requirements

pip install torch transformers peft sglang tqdm

Ensure you have access to:

• Meta Llama models (requires authentication)
• 8 GPUs with sufficient VRAM (recommended: 8×A100)
• SGLang installation and dependencies

Citation

If you use this training pipeline, consider citing:

@software{speculator_training,
  title = {Speculator Model Training with Knowledge Distillation},
  author = {Your Name},
  year = {2024},
}