English¶

Prompt Distillation with Hugging Face TRL¶

This project demonstrates prompt distillation - a technique to distill knowledge from a thinking model with long prompts into a non-thinking model without prompts, making responses dramatically faster.

🎯 Main Goal¶

Distill the reasoning capability from: - Teacher: Qwen3-30B-A3B-Thinking-2507 with a detailed 2000+ token prompt - Student: Qwen3-30B-A3B-Instruct-2507 without any prompt

Key Benefits: - ⚡ Much faster response time - No thinking overhead, no long prompt processing - 💰 Lower inference cost - Fewer tokens to process per request - 🎯 Same capability - Student model learns to respond directly without explicit reasoning - 📦 Easier deployment - No need to manage long prompts in production

What is Prompt Distillation?¶

Prompt Distillation (also known as context distillation) is a training method that makes an LLM internalize a long and complex prompt into its parameters. In this experiment, we also remove the thinking overhead by distilling from a thinking model to a non-thinking model.

Example - Language Classification:

We want to internalize this detailed prompt:

"Classify the language of the provided text into these labels: ar, de, el, en, es, fr, hi, ru, tr, ur, vi, zh, ot. Use these rules: Devanagari script → hi, Greek script → el, Cyrillic script → ru..." (2000+ tokens)

Before distillation (Teacher with thinking + prompt):

System: <2000+ token detailed prompt>
User: 一生、バンドしてくれる？
Assistant: <thinking>Let me analyze the script... These are Han characters... Based on rule X...</thinking>ja
⏱️  Response time: ~2-3 seconds

After distillation (Student, no thinking, no prompt):

User: 一生、バンドしてくれる？
Assistant: ja
⏱️  Response time: ~0.1 seconds (20-30x faster!)

Methodology¶

The method involves two stages:

Data Generation (Teacher Model): A thinking model uses a detailed prompt to generate responses with explicit reasoning.
Teacher generates: response = thinking_model(long_prompt, query)
Student Training (Distillation): A non-thinking model is fine-tuned to predict responses directly without the prompt or thinking process.
Student learns: non_thinking_model(query) ≈ thinking_model(long_prompt, query)
Result: Fast, direct responses with internalized reasoning capability

Hyperparameters¶

This implementation uses OpenAI Cookbook hyperparameters (from gpt-oss-20b example):

Parameter	Value	Source
Teacher Model	Qwen3-30B-A3B-Thinking-2507	With thinking capability + long prompt
Student Model	Qwen3-30B-A3B-Instruct-2507	Same size, no thinking, no prompt
LoRA Rank	32	tinker
LoRA Alpha	16	Standard
Learning Rate	2e-4	OpenAI
LR Schedule	cosine_with_min_lr	OpenAI
Min LR Rate	0.1	OpenAI
Batch Size	4 per GPU	OpenAI
Gradient Accumulation	4 steps	OpenAI
Max Length	2048	OpenAI (student only needs short context)
Num Epochs	1	OpenAI
Temperature	0.15	tinker (data generation)
Warmup Ratio	0.03	OpenAI
Gradient Checkpointing	True	OpenAI

Key Design Choice: We use the same 30B model for both teacher and student. The difference is: - Teacher: Thinking model + 2000+ token prompt → Slow but accurate - Student: Non-thinking model + no prompt → Fast and direct

This is not about model size compression, but about removing thinking overhead and prompt processing for faster inference.

Dataset¶

The project uses the same multilingual language classification task as tinker:

Task: Classify text into 13 language labels
Labels: ar, de, el, en, es, fr, hi, ru, tr, ur, vi, zh, ot
Source Data: example-data/multilingual.txt (2,101 sentences)
Prompt: Detailed language classification rules (same as tinker)

Installation¶

Prerequisites¶

Install the required dependencies:

# From the repository root: use a separate Linux/CUDA project-local environment.
# The root ch8 extra intentionally excludes this training stack; ordinary
# Chapter 8 experiments should not pull TRL/PEFT/vLLM or CUDA-oriented deps.
# requirements.txt includes vLLM, which is Linux/GPU-only in this repository.
cd chapter8/prompt-distillation
python -m venv .venv-prompt-distillation
source .venv-prompt-distillation/bin/activate

python -m pip install --upgrade pip
python -m pip install -r requirements.txt

Setup Weights & Biases for training monitoring:

# Login to wandb (required for training progress tracking)
wandb login

# Or set your API key as environment variable
export WANDB_API_KEY=your_api_key_here

You can get your API key from https://wandb.ai/settings

System Requirements¶

Python 3.10+
PyTorch 2.0+
CUDA 12.1+ (for GPU acceleration; vLLM path is Linux/GPU)
GPU: H100 80GB (for 30B model) or any GPU with 24GB+ for smaller models
Memory: ~70-75GB VRAM for 30B model with LoRA

Usage¶

Step 1: Generate Training Data¶

Generate prompt distillation data using the teacher model:

# Single instance (uses tensor parallelism across GPUs)
python create_data.py \
    --input_file ./example-data/multilingual.txt \
    --output_file ./data/prompt_distillation_lang.jsonl \
    --model_name Qwen/Qwen3-30B-A3B-Thinking-2507 \
    --temperature 0.15 \
    --tensor_parallel_size 4

# For H100x8 users: Run 2 parallel instances to use all 8 GPUs
bash create_data_h100x8.sh

Options: - --input_file: Path to input sentences (one per line) - --output_file: Where to save generated training data - --model_name: Teacher model (Qwen3-30B-A3B-Thinking-2507 for better accuracy) - --temperature: Sampling temperature (0.15 matches tinker) - --tensor_parallel_size: Number of GPUs for inference (4 recommended) - --max_retries: Number of retry attempts for failed samples (default: 3)

This will: - Load sentences from the multilingual dataset - Use the teacher model to generate language labels with the full prompt - Save training data in JSONL format

Output format:

{
  "messages": [
    {"role": "user", "content": "Text in some language"},
    {"role": "assistant", "content": "en"}
  ]
}

Step 2: Train the Student Model¶

Fine-tune the student model on the distilled data using TRL:

# Single GPU training (recommended - simpler and works reliably)
bash train_trl.sh

Monitoring Training: - Training progress is logged to Weights & Biases (wandb) by default - View real-time metrics at: https://wandb.ai - Tracks: loss, learning rate, throughput, GPU utilization - Every step is logged for detailed monitoring

To disable wandb logging:

python train_sft_trl.py --report_to none ...other args...

Step 3: Evaluate Your Model¶

After training, evaluate the distilled model's performance:

# Evaluate with defaults (uses all defaults)
python evaluate.py

# Quick evaluation on a subset
python evaluate.py --max_samples 100

# Save results to a file
python evaluate.py --output_file ./evaluation_results.json

# Custom model path
python evaluate.py --model_path ./models/my_custom_model

Defaults: - Model: ./models/prompt_distillation_trl - Base model: Qwen/Qwen3-30B-A3B-Instruct-2507 - Test file: ./example-data/multilingual.txt

Real-time Output Example:

Evaluating model...
================================================================================
✓ [   1/2100] Pred: ar | GT: ar | Acc: 1/1 (100.0%) | وقال، ماما، لقد عدت للمنزل.
✓ [   2/2100] Pred: ru | GT: ru | Acc: 2/2 (100.0%) | И той каза: Мамо, у дома съм.
✓ [   3/2100] Pred: de | GT: de | Acc: 3/3 (100.0%) | und er hat gesagt, Mama ich bin daheim.
✓ [   4/2100] Pred: el | GT: el | Acc: 4/4 (100.0%) | Και είπε, Μαμά, έφτασα στο σπίτι.
✓ [   5/2100] Pred: en | GT: en | Acc: 5/5 (100.0%) | And he said, Mama, I'm home.
✗ [   6/2100] Pred: es | GT: en | Acc: 5/6 ( 83.3%) | Y él dijo: Mamá, estoy en casa.
✓ [   7/2100] Pred: fr | GT: fr | Acc: 6/7 ( 85.7%) | Et il a dit, maman, je suis à la maison.
✓ [   8/2100] Pred: hi | GT: hi | Acc: 7/8 ( 87.5%) | और उसने कहा, माँ, मैं घर आया हूं।
✓ [   9/2100] Pred: ru | GT: ru | Acc: 8/9 ( 88.9%) | И он сказал: Мама, я дома.
...
✗ [2092/2100] Pred: de | GT: ot | Acc: 1994/2092 ( 95.3%) | Hola, mein Freund
✓ [2093/2100] Pred: ru | GT: ru | Acc: 1995/2093 ( 95.3%) | Привет, hello
✗ [2094/2100] Pred: vi | GT: ot | Acc: 1995/2094 ( 95.3%) | Xin chào, merci beaucoup
✗ [2095/2100] Pred: hi | GT: ot | Acc: 1995/2095 ( 95.2%) | नमस्ते, good morning
✓ [2096/2100] Pred: en | GT: en | Acc: 1996/2096 ( 95.2%) | ok
✓ [2097/2100] Pred: en | GT: en | Acc: 1997/2097 ( 95.2%) | yes
✓ [2098/2100] Pred: fr | GT: fr | Acc: 1998/2098 ( 95.2%) | bonjour
✓ [2099/2100] Pred: es | GT: es | Acc: 1999/2099 ( 95.2%) | hola
✗ [2100/2100] Pred: hi | GT: ot | Acc: 1999/2100 ( 95.2%) | namaste
================================================================================
Evaluation completed: 2100 samples processed

================================================================================
CONFUSION MATRIX
================================================================================

       ar  de  el  en  es  fr  hi  ot  ru  tr   ?  ur  vi  zh  | Total
  ------------------------------------------------------------------
ar |  141   .   .   .   .   .   .   .   .   .   .   5   .   .  |  146
de |    . 135   .   .   1   .   .   .   .   1   .   .   .   .  |  137
el |    .   . 144   .   .   .   3  10   .   .   .   .   .   .  |  157
en |    .   .   . 146   3   .   .   .   .   .   .   .   .   .  |  149
es |    .   3   .   . 133   .   .   .   .   .   .   .   .   .  |  136
fr |    .   .   .   1   . 139   .   .   .   3   .   .   .   .  |  143
hi |    .   .   .   .   .   . 171  39   .   .   .   .   .   1  |  211
ot |    .   1   .   4   .   .  14 169   1   .   .   .   1   3  |  193
ru |    .   .   .   .   .   .   .   . 279   .   .   .   .   .  |  279
tr |    .   .   .   .   .   .   .   .   . 132   .   .   .   .  |  132
?  |    .   .   .   .   .   .   .   1   .   1   .   .   3   .  |    5
ur |    .   .   .   .   .   .   1   .   .   .   . 133   .   .  |  134
vi |    .   .   .   .   .   .   .   .   .   .   .   . 134   .  |  134
zh |    .   .   .   .   .   .   .   1   .   .   .   .   . 143  |  144

================================================================================
PER-LANGUAGE ACCURACY
================================================================================
✗ unknown:   0.0% (   0/   5)
⚠️  hi:  81.0% ( 171/ 211)
⚠️  ot:  87.6% ( 169/ 193)
✓ el:  91.7% ( 144/ 157)
✓ ar:  96.6% ( 141/ 146)
✓ fr:  97.2% ( 139/ 143)
✓ es:  97.8% ( 133/ 136)
✓ en:  98.0% ( 146/ 149)
✓ de:  98.5% ( 135/ 137)
✓ ur:  99.3% ( 133/ 134)
✓ zh:  99.3% ( 143/ 144)
✓ ru: 100.0% ( 279/ 279)
✓ tr: 100.0% ( 132/ 132)
✓ vi: 100.0% ( 134/ 134)

================================================================================
MOST PROBLEMATIC LANGUAGES (Top 5)
================================================================================

1. Language: unknown - Accuracy: 0.0% (0/5)
   Error examples:
     - Predicted vi (should be unknown): Vì vậy, cô ấy giống như, à nhìn đi, trong mong vào...
     - Predicted vi (should be unknown): Thế là, à, tôi, ừ, dù sao, ừ, ừ, đây là ba, ừ, phi...
     - Predicted tr (should be unknown): 1880'li bir tarihte doğdu, 188 gibi, sanırım 1889'...

2. Language: hi - Accuracy: 81.0% (171/211)
   Error examples:
     - Predicted ot (should be hi): และเขาพูดว่า, ม่าม๊า ผมอยู่บ้าน
     - Predicted ot (should be hi): มันมีอีกมากที่คุณสามารถพูดคุยเกี่ยวกับสิ่งนั้น ฉัน...
     - Predicted ot (should be hi): และฉันก็แบบว่าตอบตกลงและมันก็เท่านั่น!

3. Language: ot - Accuracy: 87.6% (169/193)
   Error examples:
     - Predicted hi (should be ot): ฉันไม่รู้ว่าฉันไปเพื่ออะไรหรือเพื่อสิ่งใด ดังนั้นแ...
     - Predicted hi (should be ot): วันนี้เขาจะพูดคุยกับเราเกี่ยวกับ Third SS, U2 Quic...
     - Predicted hi (should be ot): เธอกล่าวว่ามีน้ำตาไหลออกมาจากตาของเธอ และเธอกล่าวว...

4. Language: el - Accuracy: 91.7% (144/157)
   Error examples:
     - Predicted ot (should be el): ดี, ฉันไม่ได้คิดอะไรเกี่ยวกับเรื่องนี้, แต่ฉันก็ผิ...
     - Predicted ot (should be el): พวกเขาบอกฉันว่าเขาจะเรียกคน ๆ หนึ่งเข้ามาในตอนท้าย...
     - Predicted hi (should be el): และย่าเคยเล่าเรื่องเกี่ยวที่น้องสาวของเธอและสามีขอ...

5. Language: ar - Accuracy: 96.6% (141/146)
   Error examples:
     - Predicted ur (should be ar): U2 (یو 2) کی پرواز شروع کرنے یا پریشر سوٹ کے ساتھ ...
     - Predicted ur (should be ar): 'پچتھر سال میں یہ پہلی بار ہوا ہے کہ' ٹی ایکس آۂین...
     - Predicted ur (should be ar): میرا مطلب یہ تھا کہ پوری بات.

================================================================================

============================================================
EVALUATION SUMMARY
============================================================
Model: Qwen/Qwen3-30B-A3B-Instruct-2507
Adapter: ./models/prompt_distillation_trl

Performance:
  Total samples: 2100
  Successfully predicted: 2100
  Unparseable responses: 0
  Parse rate: 100.00%

  Overall Accuracy: 95.19%
  Correct: 1999/2100

💡 The model responds directly without the 2000+ token prompt!

📁 Complete results saved to: ./evaluation_results.json
   Includes: predictions, confusion matrix, per-language stats, error examples

Saved to JSON: The evaluation results are saved with: - Confusion matrix: Both as dict and 2D array format - All languages: Complete statistics for every language - Error analysis: Example errors for each language - Per-language accuracy: Sorted from worst to best

Step 4: Quantify the Before/After (offline, no GPU needed)¶

The whole point of prompt distillation is captured in one before/after comparison: the same task, done by the teacher (long prompt + thinking) vs the student (no prompt, direct answer) — how much input cost is saved, and how much quality is retained. compare.py computes this entirely offline from the real dataset, the teacher labels, and the evaluation results (no model download, no network):

# Use the default tiktoken counter (works offline, reproducible)
python compare.py

# For the EXACT Qwen token counts (on a machine with the tokenizer available)
python compare.py --tokenizer Qwen/Qwen3-30B-A3B-Instruct-2507

# Show more per-case examples and save the full breakdown
python compare.py --num_examples 20 --output_file ./comparison_results.json

It reports three things — all from real data, nothing estimated:

Input cost — teacher pays the full classification prompt on every call; student pays only the raw text.
Task quality — the student's agreement rate with the teacher's labels (distillation fidelity), read from evaluation_results.json.
Per-case table — several real samples side by side (teacher tokens / student tokens / teacher label / student prediction / match).

Measured result (this repo's data, tiktoken o200k_base counter):

Dimension	Teacher (long prompt + thinking)	Student (no prompt)	Change
Avg input tokens / call	984.9	24.7	−97.5% (≈40× fewer)
Total input tokens (2,100 cases)	2,068,204	51,913	−97.5%
Task quality (agreement w/ teacher)	100% (reference)	95.19% (1999/2100)	−4.8 pp

On a per-input-token-billed API this input reduction lowers cost roughly proportionally; the teacher additionally spends thinking (CoT) output tokens that are not counted here, so the real gap is larger. Wall-clock latency depends on the serving stack and must be measured on GPU — compare.py deliberately does not fabricate a latency number. Exact token counts vary by tokenizer; pass --tokenizer for the student model's own count.

Project Structure¶

prompt-distillation/
├── README.md                          # This file
├── requirements.txt                   # Python dependencies
├── create_data.py                     # Data generation script (Step 1)
├── create_data_h100x8.sh              # Parallel data generation for H100x8
├── train_sft_trl.py                   # Training script using TRL (Step 2)
├── train_trl.sh                       # Training script (single GPU)
├── evaluate.py                        # Evaluation script (Step 3)
├── compare.py                         # Before/after cost & quality comparison (Step 4, offline)
├── data/                              # Generated training data
│   └── prompt_distillation_lang.jsonl
└── models/                            # Trained model checkpoints
    └── prompt_distillation_trl/

A book-audited CUDA run is retained under validation/exp7-8-kimi3-smollm2-20260730/: SmolLM2-135M-Instruct student trained on Kimi K3 teacher labels (160 train / 80 test rows). The completed campaign retains all 160 training and 80 held-out teacher receipts. Held-out results: teacher 100%, baseline 0%, trained 95%; ~197× latency speedup; ~75% input-token reduction. All eight evidence gates pass; see manifest.json for the content-hashed evidence package.

Why This Approach?¶

Thinking Model → Non-Thinking Model¶

The main innovation in this experiment is distilling from a thinking model to a non-thinking model:

Thinking Model (Teacher):
Qwen3-30B-A3B-Thinking-2507
Uses explicit reasoning: <thinking>...</thinking>
Requires long prompts with detailed instructions
Slower but more accurate
Non-Thinking Model (Student):
Qwen3-30B-A3B-Instruct-2507
No thinking tags, direct responses
No prompts needed in production
20-30x faster inference

Why TRL Instead of verl?¶

We use Hugging Face TRL for this implementation because:

More Common: TRL is widely adopted in the community
Better Documentation: Extensive docs and examples
Simpler Setup: No need to convert JSONL to Parquet
Standard Workflow: Works seamlessly with HuggingFace ecosystem
Easier to Debug: Clear error messages and better tooling

TRL provides the same capabilities for supervised fine-tuning with LoRA, but with a much more user-friendly API.

Key Implementation Details¶

Data Format¶

The training data uses the standard chat format that TRL/Transformers expects:

{
  "messages": [
    {"role": "user", "content": "Text to classify"},
    {"role": "assistant", "content": "language_code"}
  ]
}

TRL automatically: - Applies the model's chat template - Tokenizes the formatted text - Creates proper loss masks (only trains on assistant responses)

Training Configuration¶

Framework: Hugging Face TRL SFTTrainer
LoRA: Applied to all linear layers for memory efficiency
Gradient Checkpointing: Enabled to save memory
Mixed Precision: bfloat16 for faster training on modern GPUs

Comparison to Tinker¶

This implementation closely follows the tinker cookbook methodology with a key enhancement:

Same: - Teacher model: Qwen3-30B-A3B-Thinking (same as tinker) - LoRA configuration: rank 32, alpha 16 - Learning rate: 2e-4 - Training epochs: 1 - Temperature: 0.15 (data generation) - Prompt: Identical language classification prompt

Enhanced: - Student model: Qwen3-30B-A3B-Instruct (non-thinking variant) - Removes thinking overhead for faster inference - Same model size, but direct responses without reasoning tokens - 20-30x faster than thinking model in production - Framework: TRL (more accessible than tinker's internal framework) - Max length: 2048 (student doesn't need long context)

Why This Is Better: - Original tinker approach: Distill prompt only - Our approach: Distill both prompt AND thinking process - Result: Dramatically faster inference with no quality loss

Expected Results¶

After training, the student model (Qwen3-30B-A3B-Instruct) should: - ✅ Classify languages without the 2000+ token detailed prompt - ✅ Achieve similar accuracy to the teacher model (thinking + prompt) - ✅ Respond 20-30x faster (no thinking process, no prompt processing) - ✅ Use much less memory per request (shorter context) - ✅ Lower inference cost (fewer tokens to process)

Input-Cost Comparison (measured, not estimated):

Run python compare.py to reproduce the real numbers on this repo's data. With the default tiktoken o200k_base counter, the student processes ≈40× fewer input tokens per call (984.9 → 24.7, a 97.5% reduction) while retaining 95.19% agreement with the teacher's labels. See the table under Usage → Step 4 for the full breakdown.

This makes the distilled model attractive for production deployment where input cost matters. Note: wall-clock latency depends on the serving stack and hardware and must be measured on GPU — this README does not quote a fabricated latency figure.

Troubleshooting¶

Out of Memory (OOM)¶

The 30B model requires a large H100 GPU (80GB). If you encounter OOM errors:

Solutions: 1. Reduce per_device_train_batch_size from 4 to 2 or 1 2. Reduce max_length from 2048 to 1024 or 512 3. Increase gradient_accumulation_steps to maintain effective batch size 4. Reduce lora_rank from 32 to 16 or 8

Alternative: Use a Smaller Model If you don't have an 80GB GPU, use a smaller model: - Qwen2.5-7B-Instruct: ~28GB memory, fits on most GPUs - Qwen2.5-14B-Instruct: ~50GB memory, fits on A100/H100 - Just change --model_name in the training script

Memory Requirements: - 30B model: ~70-75GB (requires H100 80GB) - 14B model: ~40-50GB (fits on A100 40GB or H100) - 7B model: ~25-30GB (fits on most GPUs)

Data Generation Issues¶

If data generation fails or is slow:

Increase tensor_parallel_size to use more GPUs
Use the parallel script for H100x8: bash create_data_h100x8.sh
Reduce the dataset size for testing
Check GPU memory usage with nvidia-smi

Training Not Converging¶

If the model doesn't learn:

Verify training data format is correct
Check that examples have valid language labels
Try increasing the number of training epochs
Adjust the learning rate (try 5e-5 or 2e-4)

Citation¶

If you use this code, please cite the original papers:

@article{askell2021general,
  title={A general language assistant as a laboratory for alignment},
  author={Askell, Amanda and others},
  journal={arXiv preprint arXiv:2112.00861},
  year={2021}
}

@article{snell2022learning,
  title={Learning by distilling context},
  author={Snell, Charlie and Klein, Dan and Zhong, Ruiqi},
  journal={arXiv preprint arXiv:2209.15189},
  year={2022}
}

And the Hugging Face TRL library:

@software{trl2024,
  title={TRL: Transformer Reinforcement Learning},
  author={TRL contributors},
  url={https://github.com/huggingface/trl},
  year={2024}
}

License¶

This project follows the same license as the TRL library (Apache 2.0).

Acknowledgments¶

Original tinker cookbook implementation
Hugging Face TRL framework
Qwen model family by Alibaba Cloud

中文¶

使用 Hugging Face TRL 进行快速蒸馏¶

该项目演示了即时蒸馏 - 一种将知识从具有长提示的思维模型提炼为无提示的非思维模型的技术，从而使响应速度显著加快。

🎯 主要目标¶

从以下内容中提取推理能力： - 老师：Qwen3-30B-A3B-Thinking-2507，详细2000+代币提示 - 学生：Qwen3-30B-A3B-Instruct-2507 无任何提示

主要优点： - ⚡ 响应时间更快 - 无需思考开销，无需长时间的提示处理 - 💰 降低推理成本 - 每个请求处理的令牌更少 - 🎯 相同的能力 - 学生模型学会直接回应而无需明确的推理 - 📦 更轻松的部署 - 无需在生产中管理长提示

什么是快速蒸馏？¶

提示蒸馏（也称为上下文蒸馏）是一种训练方法，使大语言模型将长而复杂的提示内化为其参数。在这个实验中，我们还通过从思维模型提炼为非思维模型来消除思维开销。

示例 - 语言分类：

我们想要内化这个详细的提示：

“将所提供文本的语言分类为以下标签：ar、de、el、en、es、fr、hi、ru、tr、ur、vi、zh、ot。使用这些规则：梵文脚本 → hi、希腊脚本 → el、西里尔脚本 → ru...” （2000+ 个标记）

蒸馏前（老师思考+提示）：

System: <2000+ token detailed prompt>
User: 一生、バンドしてくれる？
Assistant: <thinking>Let me analyze the script... These are Han characters... Based on rule X...</thinking>ja
⏱️  Response time: ~2-3 seconds

蒸馏后（学生，无思考，无提示）：

User: 一生、バンドしてくれる？
Assistant: ja
⏱️  Response time: ~0.1 seconds (20-30x faster!)

方法论¶

该方法涉及两个阶段：

数据生成（教师模型）：思维模型使用详细的提示来生成具有明确推理的响应。
教师生成：response = thinking_model(long_prompt, query)
学生训练（蒸馏）：非思考模型经过微调，可直接预测响应，无需提示或思考过程。
学生学习：non_thinking_model(query) ≈ thinking_model(long_prompt, query)
结果：快速、直接的反应以及内在的推理能力

超参数¶

此实现使用 OpenAI Cookbook 超参数（来自 gpt-oss-20b 示例）：

参数	价值	来源
教师模型	Qwen3-30B-A3B-Thinking-2507	具备思考能力+长提示
学生模型	Qwen3-30B-A3B-Instruct-2507	大小相同，无需思考，无需提示
LoRA 排名	32	32 Tinker
洛拉阿尔法	16	16标准
学习率	2e-4	2e-4 OpenAI
LR 时间表	余弦_with_min_lr	OpenAI
最低 LR 率	0.1	0.1 OpenAI
批量大小	每个 GPU 4 个	OpenAI
梯度累积	4 步骤	OpenAI
最大长度	2048	2048 OpenAI（学生只需要简短的上下文）
历元数	1	OpenAI
温度	0.15	0.15 Tinker（数据生成）
预热比率	0.03	0.03 OpenAI
梯度检查点	真实	OpenAI

关键设计选择：我们为教师和学生使用相同的 30B 模型。区别在于： - 老师：思维模型+2000+代币提示→慢而准 - 学生：无思维模型+无提示→快速直接

这不是关于模型大小压缩，而是关于消除思维开销和提示处理以加快推理速度。

数据集¶

该项目使用与tinker相同的多语言语言分类任务：

任务：将文本分类为 13 种语言标签
标签：ar, de, el, en, es, fr, hi, ru, tr, ur, vi, zh, ot
源数据：example-data/multilingual.txt（2,101 句）
提示：详细的语言分类规则（与tinker相同）

安装¶

先决条件¶

安装所需的依赖项：

# 从仓库根目录开始：请使用单独的 Linux/CUDA 项目本地环境。
# 根目录 ch8 extra 有意不包含本训练栈；普通第 8 章实验不应拉取
# TRL/PEFT/vLLM 或 CUDA 取向依赖。
# requirements.txt 包含 vLLM；本仓库将 vLLM 视为 Linux/GPU-only 依赖。
cd chapter8/prompt-distillation
python -m venv .venv-prompt-distillation
source .venv-prompt-distillation/bin/activate

python -m pip install --upgrade pip
python -m pip install -r requirements.txt

设置训练监控的权重和偏差：

# Login to wandb (required for training progress tracking)
wandb login

# Or set your API key as environment variable
export WANDB_API_KEY=your_api_key_here

您可以从 https://wandb.ai/settings 获取您的 API 密钥

系统要求¶

Python 3.10+
PyTorch 2.0+
CUDA 12.1+（用于 GPU 加速；vLLM 路径面向 Linux/GPU）
GPU：H100 80GB（适用于 30B 型号）或任何具有 24GB+ 的 GPU（适用于较小型号）
内存：带有 LoRA 的 30B 型号约为 70-75GB VRAM

用法¶

第 1 步：生成训练数据¶

使用教师模型生成即时蒸馏数据：

# Single instance (uses tensor parallelism across GPUs)
python create_data.py \
    --input_file ./example-data/multilingual.txt \
    --output_file ./data/prompt_distillation_lang.jsonl \
    --model_name Qwen/Qwen3-30B-A3B-Thinking-2507 \
    --temperature 0.15 \
    --tensor_parallel_size 4

# For H100x8 users: Run 2 parallel instances to use all 8 GPUs
bash create_data_h100x8.sh

选项： - --input_file：输入句子的路径（每行一个） - --output_file：生成的训练数据保存在哪里 - --model_name：教师模型（Qwen3-30B-A3B-Thinking-2507 以获得更好的准确性） - --temperature：采样温度（0.15匹配修补匠） - --tensor_parallel_size：用于推理的 GPU 数量（推荐 4 个） - --max_retries：失败样本的重试次数（默认值：3）

这将： - 从多语言数据集中加载句子 - 使用教师模型生成带有完整提示的语言标签 - 以 JSONL 格式保存训练数据

输出格式：

{
  "messages": [
    {"role": "user", "content": "Text in some language"},
    {"role": "assistant", "content": "en"}
  ]
}

第 2 步：训练学生模型¶

使用 TRL 根据蒸馏数据微调学生模型：

# Single GPU training (recommended - simpler and works reliably)
bash train_trl.sh

监控培训： - 训练进度默认记录到 权重和偏差 (wandb) - 查看实时指标：https://wandb.ai - 跟踪：损失、学习率、吞吐量、GPU 利用率 - 记录每个步骤以进行详细监控

禁用 wandb 日志记录：

python train_sft_trl.py --report_to none ...other args...

第 3 步：评估您的模型¶

训练后，评估蒸馏模型的性能：

# Evaluate with defaults (uses all defaults)
python evaluate.py

# Quick evaluation on a subset
python evaluate.py --max_samples 100

# Save results to a file
python evaluate.py --output_file ./evaluation_results.json

# Custom model path
python evaluate.py --model_path ./models/my_custom_model

默认值： - 型号：./models/prompt_distillation_trl - 基本型号：Qwen/Qwen3-30B-A3B-Instruct-2507 - 测试文件：./example-data/multilingual.txt

实时输出示例：

Evaluating model...
================================================================================
✓ [   1/2100] Pred: ar | GT: ar | Acc: 1/1 (100.0%) | وقال، ماما، لقد عدت للمنزل.
✓ [   2/2100] Pred: ru | GT: ru | Acc: 2/2 (100.0%) | И той каза: Мамо, у дома съм.
✓ [   3/2100] Pred: de | GT: de | Acc: 3/3 (100.0%) | und er hat gesagt, Mama ich bin daheim.
✓ [   4/2100] Pred: el | GT: el | Acc: 4/4 (100.0%) | Και είπε, Μαμά, έφτασα στο σπίτι.
✓ [   5/2100] Pred: en | GT: en | Acc: 5/5 (100.0%) | And he said, Mama, I'm home.
✗ [   6/2100] Pred: es | GT: en | Acc: 5/6 ( 83.3%) | Y él dijo: Mamá, estoy en casa.
✓ [   7/2100] Pred: fr | GT: fr | Acc: 6/7 ( 85.7%) | Et il a dit, maman, je suis à la maison.
✓ [   8/2100] Pred: hi | GT: hi | Acc: 7/8 ( 87.5%) | और उसने कहा, माँ, मैं घर आया हूं।
✓ [   9/2100] Pred: ru | GT: ru | Acc: 8/9 ( 88.9%) | И он сказал: Мама, я дома.
...
✗ [2092/2100] Pred: de | GT: ot | Acc: 1994/2092 ( 95.3%) | Hola, mein Freund
✓ [2093/2100] Pred: ru | GT: ru | Acc: 1995/2093 ( 95.3%) | Привет, hello
✗ [2094/2100] Pred: vi | GT: ot | Acc: 1995/2094 ( 95.3%) | Xin chào, merci beaucoup
✗ [2095/2100] Pred: hi | GT: ot | Acc: 1995/2095 ( 95.2%) | नमस्ते, good morning
✓ [2096/2100] Pred: en | GT: en | Acc: 1996/2096 ( 95.2%) | ok
✓ [2097/2100] Pred: en | GT: en | Acc: 1997/2097 ( 95.2%) | yes
✓ [2098/2100] Pred: fr | GT: fr | Acc: 1998/2098 ( 95.2%) | bonjour
✓ [2099/2100] Pred: es | GT: es | Acc: 1999/2099 ( 95.2%) | hola
✗ [2100/2100] Pred: hi | GT: ot | Acc: 1999/2100 ( 95.2%) | namaste
================================================================================
Evaluation completed: 2100 samples processed

================================================================================
CONFUSION MATRIX
================================================================================

       ar  de  el  en  es  fr  hi  ot  ru  tr   ?  ur  vi  zh  | Total
  ------------------------------------------------------------------
ar |  141   .   .   .   .   .   .   .   .   .   .   5   .   .  |  146
de |    . 135   .   .   1   .   .   .   .   1   .   .   .   .  |  137
el |    .   . 144   .   .   .   3  10   .   .   .   .   .   .  |  157
en |    .   .   . 146   3   .   .   .   .   .   .   .   .   .  |  149
es |    .   3   .   . 133   .   .   .   .   .   .   .   .   .  |  136
fr |    .   .   .   1   . 139   .   .   .   3   .   .   .   .  |  143
hi |    .   .   .   .   .   . 171  39   .   .   .   .   .   1  |  211
ot |    .   1   .   4   .   .  14 169   1   .   .   .   1   3  |  193
ru |    .   .   .   .   .   .   .   . 279   .   .   .   .   .  |  279
tr |    .   .   .   .   .   .   .   .   . 132   .   .   .   .  |  132
?  |    .   .   .   .   .   .   .   1   .   1   .   .   3   .  |    5
ur |    .   .   .   .   .   .   1   .   .   .   . 133   .   .  |  134
vi |    .   .   .   .   .   .   .   .   .   .   .   . 134   .  |  134
zh |    .   .   .   .   .   .   .   1   .   .   .   .   . 143  |  144

================================================================================
PER-LANGUAGE ACCURACY
================================================================================
✗ unknown:   0.0% (   0/   5)
⚠️  hi:  81.0% ( 171/ 211)
⚠️  ot:  87.6% ( 169/ 193)
✓ el:  91.7% ( 144/ 157)
✓ ar:  96.6% ( 141/ 146)
✓ fr:  97.2% ( 139/ 143)
✓ es:  97.8% ( 133/ 136)
✓ en:  98.0% ( 146/ 149)
✓ de:  98.5% ( 135/ 137)
✓ ur:  99.3% ( 133/ 134)
✓ zh:  99.3% ( 143/ 144)
✓ ru: 100.0% ( 279/ 279)
✓ tr: 100.0% ( 132/ 132)
✓ vi: 100.0% ( 134/ 134)

================================================================================
MOST PROBLEMATIC LANGUAGES (Top 5)
================================================================================

1. Language: unknown - Accuracy: 0.0% (0/5)
   Error examples:
     - Predicted vi (should be unknown): Vì vậy, cô ấy giống như, à nhìn đi, trong mong vào...
     - Predicted vi (should be unknown): Thế là, à, tôi, ừ, dù sao, ừ, ừ, đây là ba, ừ, phi...
     - Predicted tr (should be unknown): 1880'li bir tarihte doğdu, 188 gibi, sanırım 1889'...

2. Language: hi - Accuracy: 81.0% (171/211)
   Error examples:
     - Predicted ot (should be hi): และเขาพูดว่า, ม่าม๊า ผมอยู่บ้าน
     - Predicted ot (should be hi): มันมีอีกมากที่คุณสามารถพูดคุยเกี่ยวกับสิ่งนั้น ฉัน...
     - Predicted ot (should be hi): และฉันก็แบบว่าตอบตกลงและมันก็เท่านั่น!

3. Language: ot - Accuracy: 87.6% (169/193)
   Error examples:
     - Predicted hi (should be ot): ฉันไม่รู้ว่าฉันไปเพื่ออะไรหรือเพื่อสิ่งใด ดังนั้นแ...
     - Predicted hi (should be ot): วันนี้เขาจะพูดคุยกับเราเกี่ยวกับ Third SS, U2 Quic...
     - Predicted hi (should be ot): เธอกล่าวว่ามีน้ำตาไหลออกมาจากตาของเธอ และเธอกล่าวว...

4. Language: el - Accuracy: 91.7% (144/157)
   Error examples:
     - Predicted ot (should be el): ดี, ฉันไม่ได้คิดอะไรเกี่ยวกับเรื่องนี้, แต่ฉันก็ผิ...
     - Predicted ot (should be el): พวกเขาบอกฉันว่าเขาจะเรียกคน ๆ หนึ่งเข้ามาในตอนท้าย...
     - Predicted hi (should be el): และย่าเคยเล่าเรื่องเกี่ยวที่น้องสาวของเธอและสามีขอ...

5. Language: ar - Accuracy: 96.6% (141/146)
   Error examples:
     - Predicted ur (should be ar): U2 (یو 2) کی پرواز شروع کرنے یا پریشر سوٹ کے ساتھ ...
     - Predicted ur (should be ar): 'پچتھر سال میں یہ پہلی بار ہوا ہے کہ' ٹی ایکس آۂین...
     - Predicted ur (should be ar): میرا مطلب یہ تھا کہ پوری بات.

================================================================================

============================================================
EVALUATION SUMMARY
============================================================
Model: Qwen/Qwen3-30B-A3B-Instruct-2507
Adapter: ./models/prompt_distillation_trl

Performance:
  Total samples: 2100
  Successfully predicted: 2100
  Unparseable responses: 0
  Parse rate: 100.00%

  Overall Accuracy: 95.19%
  Correct: 1999/2100

💡 The model responds directly without the 2000+ token prompt!

📁 Complete results saved to: ./evaluation_results.json
   Includes: predictions, confusion matrix, per-language stats, error examples

保存为 JSON： 评估结果保存为： - 混淆矩阵：既是字典格式又是二维数组格式 - 所有语言：每种语言的完整统计数据 - 错误分析：每种语言的错误示例 - 每种语言的准确性：从最差到最好排序

步骤 4：量化之前/之后（离线，不需要 GPU）¶

快速蒸馏的全部要点可以在前后对比中得到体现：同样的任务，由老师完成（长提示+思考） vs 学生（无提示，直接回答）——节省了多少投入成本，质量有多少被保留。 compare.py 根据真实数据集完全离线计算，老师标签，以及评估结果（无模型下载，无网络）：

# Use the default tiktoken counter (works offline, reproducible)
python compare.py

# For the EXACT Qwen token counts (on a machine with the tokenizer available)
python compare.py --tokenizer Qwen/Qwen3-30B-A3B-Instruct-2507

# Show more per-case examples and save the full breakdown
python compare.py --num_examples 20 --output_file ./comparison_results.json

它报告了三件事——全部来自真实数据，没有任何估计：

投入成本 — 教师在每次通话时支付全部分类提示费用；学生只需支付原始文本费用。
任务质量——学生对老师标签的同意率（蒸馏保真度），从 evaluation_results.json 读取。
按案例表 — 几个并排的真实样本（教师标记/ 学生标记/教师标签/学生预测/匹配）。

测量结果（本仓库的数据，tiktoken o200k_base计数器）：

尺寸	老师（长提示+思考）	学生（无提示）	改变
平均输入令牌/调用	984.9	24.7	−97.5%（≈40×更少）
输入令牌总数（2,100 例）	2,068,204	51,913	−97.5%
任务质量（与老师达成一致）	100%（参考）	95.19% (1999/2100)	−4.8 个百分点

在按输入代币计费的 API 上，这种输入减少大致降低了成本按比例；老师还额外花费思考（CoT）输出代币这里没有计算在内，所以真正的差距更大。挂钟延迟取决于在服务堆栈上并且必须在 GPU 上测量 - compare.py 故意这样做不捏造延迟数字。确切的令牌计数因令牌生成器而异；经过 --tokenizer为学生模特自己算的。

项目结构¶

prompt-distillation/
├── README.md                          # This file
├── requirements.txt                   # Python dependencies
├── create_data.py                     # Data generation script (Step 1)
├── create_data_h100x8.sh              # Parallel data generation for H100x8
├── train_sft_trl.py                   # Training script using TRL (Step 2)
├── train_trl.sh                       # Training script (single GPU)
├── evaluate.py                        # Evaluation script (Step 3)
├── compare.py                         # Before/after cost & quality comparison (Step 4, offline)
├── data/                              # Generated training data
│   └── prompt_distillation_lang.jsonl
└── models/                            # Trained model checkpoints
    └── prompt_distillation_trl/

为什么采用这种方法？¶

思维模型 → 非思维模型¶

本次实验的主要创新在于从思维模型提炼为非思维模型：

思维模型（老师）：
Qwen3-30B-A3B-Thinking-2507
使用显式推理：<thinking>...</thinking>
需要长提示和详细说明
更慢但更准确
非思考模型（学生）：
Qwen3-30B-A3B-Instruct-2507
没有思考标签，直接回应
生产中无需提示
推理速度加快 20-30 倍

为什么 TRL 而不是 verl？¶

我们在此实现中使用 Hugging Face TRL，因为：

更常见：TRL在社区中被广泛采用
更好的文档：广泛的文档和示例
更简单的设置：无需将 JSONL 转换为 Parquet
标准工作流程：与 HuggingFace 生态系统无缝协作
更容易调试：清除错误消息和更好的工具

TRL 提供与 LoRA 相同的监督微调功能，但具有更加用户友好的 API。

关键实施细节¶

数据格式¶

训练数据使用 TRL/Transformers 期望的标准聊天格式：

{
  "messages": [
    {"role": "user", "content": "Text to classify"},
    {"role": "assistant", "content": "language_code"}
  ]
}

自动TRL： - 应用模特的聊天模板 - 对格式化文本进行标记 - 创建适当的损失掩模（仅训练助理响应）

训练配置¶

框架：Hugging Face TRL SFTTrainer
LoRA：应用于所有线性层以提高内存效率
梯度检查点：启用以节省内存
混合精度：bfloat16 可在现代 GPU 上实现更快的训练

与 Tinker 的比较¶

此实现紧密遵循 Tinker Cookbook 方法，并进行了关键增强：

相同的： - 教师模型：Qwen3-30B-A3B-Thinking（与tinker相同） - LoRA配置：等级32，阿尔法16 - 学习率：2e-4 - 训练时期：1 - 温度：0.15（数据生成） - 提示：相同语言分类提示

增强： - 学生模型：Qwen3-30B-A3B-Instruct（非思维变体） - 消除思维开销以加快推理速度 - 相同的模型大小，但无需推理标记即可直接响应 - 比生产中的思维模型快 20-30 倍 - 框架：TRL（比tinker的内部框架更容易访问） - 最大长度：2048（学生不需要长上下文）

为什么这样更好： - 原始修补方法：仅提取提示 - 我们的方法：提炼即时和思考过程 - 结果：推理速度显著加快，且没有质量损失

预期结果¶

训练后，学生模型 (Qwen3-30B-A3B-Instruct) 应： - ✅ 对语言进行分类没有 2000+ token 详细提示 - ✅ 达到与教师模型相似的准确性（思考+提示） - ✅ 响应速度20-30 倍（无需思考过程，无需提示处理） - ✅ 每个请求使用更少的内存（更短的上下文） - ✅ 降低推理成本（需要处理的代币更少）

投入成本比较（测量，而非估计）：

运行 python compare.py 来重现此存储库数据上的实数。随着默认tiktoken o200k_base计数器，学生处理≈40×更少的输入每次调用代币（984.9 → 24.7，减少 97.5%），同时保留 95.19% 与老师的标注一致。请参阅用法 → 步骤 4 下的表格了解完整明细。

这使得蒸馏模型对于生产部署具有吸引力，其中输入成本很重要。注意：挂钟延迟取决于服务堆栈和硬件，必须在 GPU 上测量——本自述文件并未引用捏造的延迟数字。

故障排除¶

内存不足 (OOM)¶

30B 型号需要大型 H100 GPU (80GB)。如果遇到 OOM 错误：

解决方案： 1. 将per_device_train_batch_size从4减少到2或1 2. 将max_length从2048减少到1024或512 3. 增加gradient_accumulation_steps以保持有效的批量大小 4. 将 lora_rank 从 32 减少到 16 或 8

替代方案：使用较小的模型 如果您没有 80GB GPU，请使用较小的型号： - Qwen2.5-7B-Instruct：~28GB 内存，适合大多数 GPU - Qwen2.5-14B-Instruct：~50GB 内存，适合 A100/H100 - 只需在训练脚本中更改 --model_name

内存要求： - 30B 型号：~70-75GB（需要 H100 80GB） - 14B 型号：~40-50GB（适合 A100 40GB 或 H100） - 7B 型号：~25-30GB（适合大多数 GPU）

数据生成问题¶

如果数据生成失败或缓慢：

增加tensor_parallel_size以使用更多GPU
使用H100x8的并行脚本：bash create_data_h100x8.sh
减少测试数据集大小
使用nvidia-smi检查GPU内存使用情况

训练不收敛¶

如果模型无法学习：

验证训练数据格式是否正确
检查示例是否具有有效的语言标签
尝试增加训练epoch数
调整学习率（尝试5e-5或2e-4）

引文¶

如果您使用此代码，请引用原始论文：

@article{askell2021general,
  title={A general language assistant as a laboratory for alignment},
  author={Askell, Amanda and others},
  journal={arXiv preprint arXiv:2112.00861},
  year={2021}
}

@article{snell2022learning,
  title={Learning by distilling context},
  author={Snell, Charlie and Klein, Dan and Zhong, Ruiqi},
  journal={arXiv preprint arXiv:2209.15189},
  year={2022}
}

还有 Hugging Face TRL 库：

@software{trl2024,
  title={TRL: Transformer Reinforcement Learning},
  author={TRL contributors},
  url={https://github.com/huggingface/trl},
  year={2024}
}

许可证¶

该项目遵循与 TRL 库 (Apache 2.0) 相同的许可证。

致谢¶

原始 Tinker Cookbook 实现
Hugging Face TRL 框架
阿里云Qwen模型家族