How Prompt Distillation Enables AI Agent Skill Transfer

Prompt distillation transfers complex capabilities from a large teacher model to a lightweight student by training the student on label-only outputs, eliminating the need for multi-kilobyte prompts during inference.

The bojieli/ai-agent-book repository implements a complete prompt distillation pipeline that compresses elaborate chain-of-thought reasoning into model parameters. This approach allows AI agents to execute specialized skills—such as language classification—without the latency and cost associated with extensive system prompts.

The Two-Stage Distillation Pipeline

Prompt distillation operates as a supervised transfer learning workflow. The process extracts latent knowledge from a prompt-dependent teacher and embeds it directly into a smaller student model's weights.

Teacher Generation with create_data.py

The pipeline begins with teacher generation, where a powerful LLM processes inputs through verbose reasoning chains. In create_data.py, the constant LANGUAGE_CLASSIFICATION_PROMPT defines the full task description spanning thousands of tokens. The function generate_distillation_data() iterates over input sentences, formats each request using a chat template, and runs the teacher model (e.g., Qwen/Qwen3-30B-A3B-Thinking) via vLLM.

For each input, the teacher produces a final label through step-by-step reasoning. The helper parse_final_answer() extracts this label from the teacher's output, discarding the intermediate chain-of-thought. The script writes these input-label pairs to a JSONL training file, creating a distilled dataset free of the original prompt overhead.

Student Fine-Tuning with train_sft_trl.py

The second stage trains a compact student model to predict these labels directly. The script train_sft_trl.py loads the distilled dataset using the Hugging Face datasets library and initializes the TRL SFTTrainer for supervised fine-tuning.

Because the training data contains only raw sentences mapped to labels, the student learns to bypass explicit reasoning steps. The implementation supports LoRA adapters for efficient parameter updates, allowing the student (e.g., Qwen/Qwen3-30B-A3B-Instruct) to internalize the skill without full weight updates. Once trained, the student answers immediately when given a sentence, requiring zero prompt engineering at inference time.

Evaluation and Cost Analysis

After fine-tuning, the repository provides tools to validate skill retention and measure computational savings.

Testing with evaluate.py

The evaluate.py script loads the fine-tuned student model along with any LoRA adapters and runs inference on a held-out test set. It extracts predicted labels using parse_language_label() and computes accuracy metrics, confusion matrices, and per-language statistics. This validates that the distilled model retains the teacher's classification capability despite the removal of the reasoning prompt.

Quantifying Efficiency Gains with compare.py

The compare.py module generates a quantitative report comparing teacher and student inputs. It counts tokens for the teacher (full prompt plus sentence) versus the student (sentence only), revealing input-token reductions often exceeding 90%—eliminating approximately 2,000 tokens per request.

The report also calculates skill fidelity, typically showing the student retains over 95% of the teacher's accuracy. Example-level comparisons illustrate specific sentences where token savings materialize alongside correct predictions, demonstrating that the compressed model maintains expert-level performance.

Complete Implementation Workflow

Execute the full pipeline using these four commands:

# 1. Generate distillation data (teacher → label)

python create_data.py --input_file multilingual.txt --output_file distill.jsonl

# 2. Fine-tune the student with TRL + LoRA

python train_sft_trl.py \
    --train_file distill.jsonl \
    --model_name Qwen/Qwen3-30B-A3B-Instruct-2507 \
    --output_dir models/prompt_distillation_trl \
    --use_lora

# 3. Evaluate the distilled student

python evaluate.py \
    --model_path models/prompt_distillation_trl \
    --test_file multilingual.txt \
    --output_file evaluation_results.json

# 4. Produce a cost/quality comparison report

python compare.py \
    --test_file multilingual.txt \
    --train_data_file distill.jsonl \
    --eval_results evaluation_results.json \
    --output_file compare_report.json

Running compare.py produces output similar to:

输入 token 降低 94.3% → 教师是学生的 17.1 倍。
学生(蒸馏后) : 96.7% (987/1020)

Summary

• Prompt distillation transfers skills from teacher to student through label-only training data, stripping away reasoning prompts after the knowledge extraction phase.
• The create_data.py script handles teacher generation using generate_distillation_data() and parse_final_answer() to produce clean input-label pairs.
• train_sft_trl.py performs student fine-tuning via TRL's SFTTrainer with optional LoRA adapters, teaching the model to respond without prompts.
• Evaluation in evaluate.py confirms the student retains >95% of teacher accuracy while compare.py documents >90% token reduction.

Frequently Asked Questions

What is prompt distillation?

Prompt distillation is a compression technique that trains a small model to reproduce the final outputs of a large model that relies on extensive prompting. According to the bojieli/ai-agent-book source code, this involves generating labels from a teacher using verbose chain-of-thought prompts, then fine-tuning a student on just the input-label pairs to eliminate prompt dependency at runtime.

How much latency and cost does prompt distillation save?

The technique typically reduces input tokens by over 90% by removing multi-kilobyte system prompts. As calculated in compare.py, this translates to the teacher consuming 17x more input tokens than the student, directly reducing API costs and inference latency for high-throughput AI agents.

Which models work best for teacher and student roles?

The repository uses Qwen/Qwen3-30B-A3B-Thinking as the teacher, leveraging its reasoning capabilities through the full LANGUAGE_CLASSIFICATION_PROMPT. For the student, Qwen/Qwen3-30B-A3B-Instruct serves as the base model, fine-tuned to answer directly without reasoning steps while maintaining high accuracy on the transferred skill.

Can this pattern transfer skills beyond language classification?

Yes. While the example implementation in create_data.py focuses on language identification, the same pipeline applies to any task where a detailed prompt guides a teacher model's reasoning. The core pattern—generate labels with prompts, then train a student to predict those labels directly—scales to classification, extraction, and other agent capabilities.

{
  "mcpServers": {
    "instagit": {
      "command": "npx",
      "args": ["-y", "instagit@latest"]
    }
  }
}