Knowledge Distillation

Help the user set up and run distillation from teacher to student models using the Tinker API.

Step 1: Understand the request

Ask the user (if not already specified):

• Student model: Which model to train (e.g., Qwen/Qwen3-8B-Base)
• Teacher model: Which model to distill from (e.g., Qwen/Qwen3-8B, or a checkpoint path)
• Distillation type:
  • On-policy: Student generates, teacher scores via KL — best for reasoning/chat
  • Off-policy reasoning: SFT on teacher-generated reasoning traces (e.g., OpenThoughts3)
  • Multi-teacher: Combine multiple teachers on different datasets

Step 2: Reference existing recipes

Read these files for patterns:

• tinker_cookbook/recipes/distillation/on_policy_distillation.py — On-policy distillation CLI
• tinker_cookbook/recipes/distillation/off_policy_reasoning.py — SFT on OpenThoughts3 traces
• tinker_cookbook/recipes/distillation/on_policy_multi_teacher.py — Multi-teacher setup
• tinker_cookbook/distillation/train_on_policy.py — Core on-policy training loop
• tinker_cookbook/distillation/datasets.py — TeacherConfig, PromptOnlyDatasetBuilder, DistillationDatasetConfig

Step 3: Choose distillation approach

On-Policy Distillation (Recommended)

Student generates samples, teacher provides KL penalty supervision. No correctness rewards needed.

Key config:

• TeacherConfig(base_model="Qwen/Qwen3-8B", load_checkpoint_path=None)
• PromptOnlyDatasetBuilder(dataset_name="deepmath"|"tulu3", ...)
• DistillationDatasetConfig(dataset_builder=..., teacher_config=..., groups_per_batch=...)
• kl_penalty_coef: Weight of KL penalty (default 1.0)
• kl_discount_factor: Discount for future KL (0.0 = no discount)

Off-Policy Reasoning (SFT on Traces)

Standard SFT on pre-generated reasoning traces (e.g., OpenThoughts3). Simpler but less effective than on-policy.

See recipes/distillation/off_policy_reasoning.py — uses the standard SL pipeline from supervised/train.py.

Multi-Teacher Distillation

Combine multiple teacher models on different datasets. Each dataset can have its own teacher.

See recipes/distillation/on_policy_multi_teacher.py — passes multiple DistillationDatasetConfig objects.

Step 4: Write the training script

Follow the on-policy distillation pattern:

import asyncio
import chz
from tinker_cookbook import checkpoint_utils, cli_utils
from tinker_cookbook.distillation import train_on_policy
from tinker_cookbook.distillation.datasets import (
    DistillationDatasetConfig,
    PromptOnlyDatasetBuilder,
    TeacherConfig,
)

@chz.chz
class CLIConfig:
    model_name: str = "Qwen/Qwen3-8B-Base"       # Student
    teacher_model: str = "Qwen/Qwen3-8B"          # Teacher
    dataset: str = "deepmath"                       # deepmath or tulu3
    group_size: int = 4
    groups_per_batch: int = 1024
    learning_rate: float = 1e-4
    max_tokens: int = 4096
    kl_penalty_coef: float = 1.0
    kl_discount_factor: float = 0.0
    lora_rank: int = 128
    loss_fn: str = "importance_sampling"

async def cli_main(cli_config: CLIConfig):
    renderer_name = await checkpoint_utils.resolve_renderer_name_from_checkpoint_or_default_async(
        model_name=cli_config.model_name, ...)

    dataset_builder = PromptOnlyDatasetBuilder(
        dataset_name=cli_config.dataset,
        groups_per_batch=cli_config.groups_per_batch,
        group_size=cli_config.group_size,
        model_name_for_tokenizer=cli_config.model_name,
        renderer_name=renderer_name,
    )
    teacher_config = TeacherConfig(base_model=cli_config.teacher_model)
    dataset_config = DistillationDatasetConfig(
        dataset_builder=dataset_builder,
        teacher_config=teacher_config,
        groups_per_batch=cli_config.groups_per_batch,
    )
    config = train_on_policy.Config(
        dataset_configs=[dataset_config],
        model_name=cli_config.model_name,
        renderer_name=renderer_name,
        learning_rate=cli_config.learning_rate,
        lora_rank=cli_config.lora_rank,
        max_tokens=cli_config.max_tokens,
        kl_penalty_coef=cli_config.kl_penalty_coef,
        kl_discount_factor=cli_config.kl_discount_factor,
        loss_fn=cli_config.loss_fn,
        log_path="/tmp/tinker-examples/distillation/my_run",
    )
    await train_on_policy.main(config)

Step 5: Run

# On-policy distillation (reasoning)
python -m tinker_cookbook.recipes.distillation.on_policy_distillation \
    model_name=Qwen/Qwen3-8B-Base dataset=deepmath learning_rate=1e-4

# Off-policy reasoning (SFT on traces)
python -m tinker_cookbook.recipes.distillation.off_policy_reasoning \
    model_name=Qwen/Qwen3-8B-Base learning_rate=2e-4

# Multi-teacher
python -m tinker_cookbook.recipes.distillation.on_policy_multi_teacher \
    model_name=Qwen/Qwen3-8B-Base learning_rate=1e-4

Step 6: Add tests

If you created a new distillation recipe, add a smoke test:

# tests/recipes/test_recipe_<name>.py
import pytest
from tests.helpers import run_recipe

@pytest.mark.integration
def test_<recipe_name>():
    run_recipe(
        "tinker_cookbook.recipes.<recipe_name>.train",
        ["behavior_if_log_dir_exists=delete", "groups_per_batch=16"],
    )

run_recipe() automatically passes max_steps=2 so the recipe runs 2 training steps and exits. See tests/recipes/test_recipe_on_policy_distillation.py and tests/recipes/test_recipe_on_policy_multi_teacher.py for existing examples.

Step 7: Export weights (optional)

After distillation, export the student model using the tinker_cookbook.weights API:

from tinker_cookbook import weights

adapter_dir = weights.download(tinker_path="tinker://run-id/sampler_weights/final", output_dir="./adapter")
weights.build_hf_model(base_model="Qwen/Qwen3-8B-Base", adapter_path=adapter_dir, output_path="./model")
weights.publish_to_hf_hub(model_path="./model", repo_id="user/my-distilled-model")

Common pitfalls

• Teacher model must be compatible with student's tokenizer/renderer
• On-policy is generally better than off-policy but more compute-intensive
• kl_discount_factor=0.0 means no discounting — increase for longer sequences
• High kl_penalty_coef can make training too conservative
• For multi-teacher, ensure groups_per_batch is balanced across datasets