Practical Project: Distilling a Mathematics AI Model

Chapter Overview

In the previous chapters, we explored the theory behind knowledge distillation and how a smaller AI model can learn from the outputs generated by a larger teacher model. In this chapter, we will build a complete practical project that demonstrates the entire distillation pipeline.

By the end of this project, you will understand how to:

Generate mathematical reasoning using a teacher model.
Build a synthetic training dataset.
Train a smaller student model.
Evaluate the student model.
Improve the student through iterative fine-tuning.

This project uses the Hugging Face ecosystem and is designed to be simple enough for beginners while reflecting real-world workflows used in modern AI development.

Project Goal

Our objective is to train a lightweight language model capable of solving mathematical problems by learning from the responses generated by a much larger language model.

Instead of manually creating thousands of worked-out solutions, we will leverage a powerful teacher model to generate high-quality reasoning and then use those outputs to teach a compact student model.

The complete workflow is shown below:

                  Mathematics Questions
                          │
                          ▼
               Large Teacher Language Model
                          │
             Generate Step-by-Step Solutions
                          │
                          ▼
             Teacher Generated Training Dataset
                          │
                          ▼
               Fine-Tune Student Language Model
                          │
                          ▼
            Efficient Mathematics AI Assistant

Project Requirements

Hardware

Recommended:

NVIDIA GPU (RTX 3060 or better)
16 GB RAM minimum
50 GB free disk space

Alternatively, the project can be executed on:

Google Colab
Kaggle Notebooks
AWS SageMaker
Paperspace
RunPod

Software

Python 3.11+
PyTorch
Hugging Face Transformers
Hugging Face Datasets
PEFT
Accelerate
TRL

Step 1 – Install Required Libraries

Install the required Python packages.

pip install transformers
pip install datasets
pip install accelerate
pip install peft
pip install trl
pip install sentencepiece
pip install torch

Or install everything at once:

pip install transformers datasets accelerate peft trl sentencepiece torch

Step 2 – Organize the Project

Create the following directory structure.

Math-Distillation/

│
├── data/
│      math_questions.json
│      teacher_answers.json
│
├── scripts/
│      generate_teacher.py
│      create_dataset.py
│      train_student.py
│      evaluate.py
│
├── outputs/
│
├── models/
│
├── requirements.txt
│
└── README.md

Keeping a clean project structure makes it easier to scale the project later.

Step 3 – Build the Mathematics Dataset

The teacher model requires mathematical questions.

Create a file named:

data/math_questions.json

Example dataset:

[
    {
        "question":"What is 25 × 14?"
    },
    {
        "question":"Solve x + 5 = 17"
    },
    {
        "question":"Differentiate x²"
    },
    {
        "question":"Integrate 2x"
    },
    {
        "question":"Find the derivative of sin(x)"
    }
]

For real projects, datasets often contain hundreds of thousands of questions collected from educational resources, textbooks, and benchmark datasets.

Step 4 – Choose the Teacher Model

The teacher should be significantly more capable than the student.

Suitable open-source teacher models include:

Model	Parameters
Qwen2.5-7B-Instruct	7 Billion
Llama 3.1 8B	8 Billion
Mistral 7B Instruct	7 Billion
DeepSeek-R1 Distill	8B
Gemma 2 9B	9 Billion

The teacher is responsible for producing high-quality mathematical reasoning.

Step 5 – Generate Teacher Responses

The teacher receives mathematical questions and generates detailed step-by-step solutions.

from transformers import AutoTokenizer
from transformers import AutoModelForCausalLM
import torch
import json

model_name = "Qwen/Qwen2.5-7B-Instruct"

tokenizer = AutoTokenizer.from_pretrained(model_name)

model = AutoModelForCausalLM.from_pretrained(
    model_name,
    torch_dtype=torch.float16,
    device_map="auto"
)

with open("data/math_questions.json") as f:
    questions = json.load(f)

results = []

for item in questions:

    prompt = f"""
Solve the following mathematics problem carefully.

Question:
{item["question"]}

Provide a step-by-step solution.
"""

    inputs = tokenizer(prompt, return_tensors="pt").to(model.device)

    output = model.generate(
        **inputs,
        max_new_tokens=300
    )

    response = tokenizer.decode(
        output[0],
        skip_special_tokens=True
    )

    results.append({
        "question": item["question"],
        "teacher": response
    })

with open("data/teacher_answers.json","w") as f:
    json.dump(results,f,indent=4)

Step 6 – Inspect the Generated Data

Your generated dataset now looks like this.

{
  "question": "25 × 14",

  "teacher": "

Step 1

25 × 10 = 250

Step 2

25 × 4 = 100

Step 3

250 + 100 = 350

Final Answer = 350
"
}

Notice that the teacher does more than produce the correct answer.

It also demonstrates the reasoning process.

This reasoning becomes valuable training data.

Step 7 – Create the Distillation Dataset

Next, combine the question and teacher response into a format suitable for language-model training.

from datasets import Dataset
import json

with open("data/teacher_answers.json") as f:
    teacher = json.load(f)

training_examples = []

for sample in teacher:

    training_examples.append({

        "text":

f"""Question:
{sample["question"]}

Answer:
{sample["teacher"]}
"""

    })

dataset = Dataset.from_list(training_examples)

dataset.save_to_disk("math_dataset")

Each example now contains both the prompt and the teacher-generated solution.

Step 8 – Select the Student Model

The student should be considerably smaller.

Good choices include:

Student Model	Parameters
Qwen2.5-0.5B	500 Million
Qwen2.5-1.5B	1.5 Billion
SmolLM2	1.7 Billion
Phi-3 Mini	3.8 Billion
Llama 3.2 1B	1 Billion

Smaller models train faster and require significantly less GPU memory.

Step 9 – Fine-Tune the Student

Load the dataset and begin training.

from transformers import AutoTokenizer
from transformers import AutoModelForCausalLM
from transformers import Trainer
from transformers import TrainingArguments

from datasets import load_from_disk

dataset = load_from_disk("math_dataset")

model_name = "Qwen/Qwen2.5-0.5B-Instruct"

tokenizer = AutoTokenizer.from_pretrained(model_name)

model = AutoModelForCausalLM.from_pretrained(model_name)

def tokenize(example):

    return tokenizer(
        example["text"],
        truncation=True,
        padding="max_length",
        max_length=512
    )

tokenized_dataset = dataset.map(tokenize)

training_args = TrainingArguments(

    output_dir="./outputs",

    num_train_epochs=3,

    learning_rate=2e-5,

    per_device_train_batch_size=2,

    logging_steps=20,

    save_steps=200,

    save_total_limit=2

)

trainer = Trainer(

    model=model,

    args=training_args,

    train_dataset=tokenized_dataset

)

trainer.train()

During training, the student repeatedly compares its predictions with the teacher responses and gradually learns to imitate the teacher's reasoning.

Step 10 – Test the Student Model

Once training is complete, evaluate the model.

question = "Solve 36 × 18"

prompt = f"""

Question:
{question}

Answer:

"""

inputs = tokenizer(prompt, return_tensors="pt")

outputs = model.generate(

    **inputs,

    max_new_tokens=200

)

print(tokenizer.decode(outputs[0]))

Example output:

Step 1

36 × 10 = 360

Step 2

36 × 8 = 288

Step 3

360 + 288 = 648

Final Answer = 648

Although much smaller than the teacher, the student can often produce high-quality reasoning for similar mathematical problems.

Step 11 – Evaluate the Student

Evaluation should include both correctness and efficiency.

Common metrics include:

Metric	Description
Accuracy	Percentage of correct answers
BLEU / ROUGE	Text similarity to teacher outputs
Exact Match	Final answer correctness
Inference Speed	Time per response
Memory Usage	GPU or RAM consumption
Model Size	Storage footprint

Testing on unseen mathematical problems is essential to ensure the model has learned to generalize rather than memorize.

Step 12 – Improve the Student

Several strategies can further enhance performance:

Increase the size and diversity of the training dataset.
Remove low-quality or inconsistent teacher responses.
Fine-tune on domain-specific mathematics topics such as algebra, calculus, or geometry.
Use LoRA or QLoRA for parameter-efficient training.
Train for additional epochs while monitoring validation performance.
Evaluate using benchmarks such as GSM8K, MATH, and MMLU-STEM.

Iterative refinement often yields significant improvements without increasing the size of the student model.

Best Practices

To build an effective distilled mathematics model:

Select a high-quality teacher model.
Generate detailed, step-by-step reasoning rather than only final answers.
Use diverse mathematical problems covering multiple topics.
Clean and verify the generated dataset.
Monitor validation metrics during training.
Test on unseen examples before deployment.
Regularly retrain the student as stronger teacher models become available.

Project Summary

In this chapter, we built a complete end-to-end knowledge distillation pipeline for mathematical reasoning.

The process consisted of:

Preparing a mathematics dataset.
Selecting a powerful teacher model.
Generating step-by-step solutions.
Constructing a synthetic training dataset.
Selecting a lightweight student model.
Fine-tuning the student.
Evaluating performance.
Improving the model through iterative refinement.

This workflow closely mirrors the techniques used by organizations developing compact AI models for education, customer support, mobile applications, and edge computing. By mastering this pipeline, you gain practical experience in one of the most important approaches for deploying efficient, scalable, and cost-effective AI systems.

Project Cost Estimation: Distilling a Mathematics AI Model on Google Colab

One of the most common questions when building a knowledge distillation project is:

How much will it cost to train a mathematics AI model using Google Colab?

The answer depends primarily on four factors:

The size of the teacher model.
The size of the student model.
The number of mathematics problems used for training.
The type of GPU available in Google Colab.

For this project, we assume the following configuration:

Component	Selection
Platform	Google Colab
Teacher Model	Llama 3.1 8B
Student Model	Llama 3.2 1B
Fine-Tuning Method	QLoRA (Recommended)

Why Use QLoRA?

Although full fine-tuning is possible, it requires significantly more GPU memory and increases training costs.

QLoRA (Quantized Low-Rank Adaptation) reduces memory consumption by training only a small number of additional parameters while keeping the base model frozen. This approach offers:

Lower GPU memory requirements
Faster training
Reduced cloud computing costs
Performance comparable to full fine-tuning for many tasks

For individual developers, students, and researchers, QLoRA is the recommended approach.

Scenario 1 – Beginner Learning Project

This configuration is ideal for readers who want to understand the complete distillation workflow without spending a large amount of money.

Item	Value
Mathematics Questions	10,000
Teacher Model	Llama 3.1 8B
Student Model	Llama 3.2 1B
Fine-Tuning	QLoRA
GPU	Google Colab L4 or A100
Estimated Training Time	2–4 Hours
Estimated Cost	US$10–25

This setup produces a functional mathematics assistant capable of demonstrating the complete teacher–student distillation pipeline.

Scenario 2 – Intermediate Research Project

For users seeking higher accuracy and broader mathematical coverage, a larger dataset is recommended.

Item	Value
Mathematics Questions	100,000
Teacher Response Generation	8–12 Hours
Student Training	6–10 Hours
GPU	A100
Estimated Cost	US$60–150

The resulting student model generally demonstrates stronger reasoning and better generalization than the beginner configuration.

Scenario 3 – Production-Scale Model

Organizations developing commercial mathematics assistants typically use much larger datasets.

Item	Value
Mathematics Questions	1,000,000
Teacher Inference Time	2–5 Days
Student Training	2–4 Days
Multiple Evaluation Cycles	Yes
Estimated Cost	US$500–2,000+

Production systems usually include additional stages such as data filtering, hyperparameter optimization, benchmark evaluation, and multiple rounds of fine-tuning.

Google Colab GPU Pricing

Google Colab provides access to several GPU types. Pricing varies by region, subscription plan, and availability, but the following estimates are representative.

GPU	Approximate Cost per Hour
NVIDIA T4	US$0.18–0.30
NVIDIA L4	US$0.45–0.80
NVIDIA A100	US$1.20–2.00

For this project, an L4 or A100 GPU is recommended because the teacher model requires substantial computational resources.

GPU Memory Requirements

Teacher Model – Llama 3.1 8B

The amount of GPU memory depends on the numerical precision used.

Precision	Approximate VRAM
FP16	16–20 GB
8-bit Quantization	10–12 GB
4-bit Quantization	6–8 GB

Running the teacher model using 4-bit quantization allows it to fit comfortably on an NVIDIA L4 GPU while maintaining strong performance.

Student Model – Llama 3.2 1B

The student model is considerably smaller and easier to train.

Training Method	Approximate VRAM
FP16	2–3 GB
LoRA	8–12 GB
QLoRA	6–8 GB

Because QLoRA reduces memory usage, it is the preferred method for training on Google Colab.

Dataset Generation Time

Before training begins, the teacher model must generate step-by-step solutions for every mathematics question.

The approximate generation times are shown below.

Number of Questions	Estimated Time
10,000	1–2 Hours
50,000	5–8 Hours
100,000	10–15 Hours
1,000,000	4–6 Days

The actual time depends on the average length of the generated responses and the GPU being used.

Total Cost Breakdown

The following table summarizes the expected costs for each project scale.

Project Size	Questions	Estimated Cost
Beginner	10,000	US$10–25
Intermediate	100,000	US$60–150
Production	1,000,000	US$500–2,000+

As the dataset grows, the cost increases primarily due to the additional time required for teacher inference.

Recommended Configuration for This Book

For readers following this book, the following configuration provides an excellent balance between affordability, performance, and educational value.

Component	Recommendation
Platform	Google Colab Pro
Teacher Model	Llama 3.1 8B (4-bit Quantized)
Student Model	Llama 3.2 1B
Fine-Tuning Method	QLoRA
Mathematics Dataset	20,000 Questions
Estimated Training Time	4–6 Hours
Estimated Total Cost	US$20–40

This configuration allows readers to complete the full knowledge distillation pipeline while keeping cloud computing costs manageable.

Scaling the Project

The project can be expanded gradually as experience and resources grow.

Stage 1 – Beginner

1,000 mathematics questions
Learn the complete workflow
Minimal cloud computing cost

Stage 2 – Intermediate

20,000 mathematics questions
Improved reasoning quality
Suitable for educational applications

Stage 3 – Advanced

100,000 mathematics questions
Strong mathematical reasoning
Appropriate for research projects

Stage 4 – Production

1,000,000+ mathematics questions
Commercial-grade performance
Multiple rounds of evaluation and optimization

By increasing the dataset size over time, developers can progressively improve the capabilities of the student model while controlling infrastructure costs.

Key Takeaways

Google Colab Pro provides an affordable environment for experimenting with AI model distillation.
Llama 3.1 8B serves as a capable teacher model for generating high-quality mathematical reasoning.
Llama 3.2 1B is an efficient student model that can learn much of the teacher's behavior through distillation.
QLoRA significantly reduces GPU memory usage and cloud computing costs while maintaining strong performance.
A dataset of 20,000 mathematics questions offers an excellent balance between cost and educational value for readers following this project.
The complete project can typically be completed for approximately US$20–40, making it accessible to students, researchers, and independent developers.