Skip to navigation

Knowledge Distillation Customization

View as Markdown

Learn how to train a smaller student model to mimic a larger teacher model using knowledge distillation (KD).

About

Knowledge Distillation transfers knowledge from a large teacher model to a smaller student model. During training, the student learns to match the teacher’s output probability distribution, producing a compact model that retains much of the teacher’s capability.

What you can achieve with KD:

  • Compress models: Distill an 8B model into a 0.6B model for faster inference and lower deployment costs
  • Reduce latency: Deploy a smaller model that responds faster while preserving quality
  • Lower resource requirements: Serve a distilled model on fewer GPUs

KD vs SFT: Understanding the Trade-offs

AspectFull SFTKnowledge Distillation
Training signalGround-truth labels onlyTeacher’s soft probability distribution + labels
Knowledge sourceDataset examplesTeacher model’s learned representations
Output model sizeSame as input modelTypically a smaller student model
GPU requirementsNeeds to fit one modelNeeds to fit both teacher and student in memory
Best forDomain adaptation, new knowledge injectionModel compression, latency reduction

Key Parameters

ParameterDefaultDescription
teacher_model(required)Teacher model entity URN (e.g., default/qwen3-8b-teacher)
teacher_precisionbf16Precision for the frozen teacher (bf16, fp16, fp32). Lower = less memory
distillation_ratio0.5Balance between CE loss and KD loss. 0.0 = CE only, 1.0 = KD only
distillation_temperature1.0Softmax temperature. Higher = softer distributions, more knowledge transfer

Workflow Overview

This tutorial follows a complete distillation pipeline:

  1. Fine-tune the teacher (SFT on the task dataset) so it learns the domain
  2. Establish a baseline by deploying the base student model and measuring ROUGE scores
  3. Distill into the student using the fine-tuned teacher’s soft targets
  4. Evaluate the distilled student and compare ROUGE scores against the baseline

When to choose KD:

  • You have a high-quality large model and want a smaller, faster version
  • Deployment latency or cost is a constraint
  • The teacher and student share the same vocabulary (e.g., both are Qwen3 models)

When to choose SFT instead: Refer to the Full SFT tutorial when you want to train a model directly on labeled data without a teacher.

Prerequisites

Before starting this tutorial, ensure you have:

  1. Completed the Quickstart to install and deploy NeMo Helix locally
  2. Installed the Python SDK (PyPI wrapper: pip install "nemo-helix[all]"; source checkout: run make bootstrap from the repository root)
  3. Installed evaluation dependencies:
pip install evaluate rouge_score datasets
  1. Two 80 GB GPUs with CUDA 13+ for the teacher’s full SFT. The distillation job and the deployments each use one GPU.

Quick Start

1. Initialize SDK

The SDK connects to the deployment in your current CLI context, using its base URL, default workspace, and credentials. To connect to another deployment, refer to Connect with CLI Contexts. These examples use the default workspace. If your current CLI context uses a different workspace, run nemo config set --workspace default first.

In environments without a CLI configuration file, such as Google Colab, set NHX_BASE_URL (and NHX_ACCESS_TOKEN for an authenticated deployment) before you create the client. NemoClient.from_config() reads both:

import os
os.environ["NHX_BASE_URL"] = "<YOUR_NHX_BASE_URL>"
import json
import os
import time
import uuid
from pathlib import Path
from IPython.display import clear_output
from nemo_helix_plugin.client.client import NemoClient
from nemo_helix_plugin.client.errors import ConflictError
from nemo_helix_plugin.files.storage_config import HuggingfaceStorageConfig
from nemo_helix_plugin.files.types import CreateFilesetRequest
from nemo_helix_plugin.models.types import (
ContainerExecutorConfig,
CreateModelDeploymentConfigRequest,
CreateModelDeploymentRequest,
CreateModelEntityRequest,
Engine,
ModelDeploymentConfigModelSpec,
UpdateModelEntityRequest,
)
from nemo_helix_plugin.secrets.types import HelixSecretCreateRequest
from nemo_customizer.sdk import Customization
from filesets import FilesetFileSystem
from nemo_helix_plugin.files.client import FilesClient
from pydantic import SecretStr
from nemo_helix_plugin.schema import SecretRef
client = NemoClient.from_config()
customization = Customization.from_client(client)

2. Prepare Dataset

Knowledge distillation uses the same dataset formats as SFT. We use the SQuAD dataset for both teacher training and distillation so that the teacher first learns the task, then transfers that knowledge to the student.

We convert each example to the chat messages format. The user turn is the same prompt the evaluation step sends, so the model is trained and queried through the same chat template, including the end-of-turn token that stops generation.

We also hold out a small test split for ROUGE evaluation at the end.

from datasets import load_dataset, DatasetDict
print("Loading dataset rajpurkar/squad")
raw_dataset = load_dataset("rajpurkar/squad")
if not isinstance(raw_dataset, DatasetDict):
raise ValueError("Dataset does not contain expected splits")
print("Loaded dataset")
SEED = 1234
TRAINING_SIZE = 3000
VALIDATION_SIZE = 300
TEST_SIZE = 100
DATASET_PATH = Path("kd-dataset").absolute()
os.makedirs(DATASET_PATH, exist_ok=True)
train_set = raw_dataset.get('train')
split = train_set.train_test_split(test_size=0.05, seed=SEED)
train_ds = split['train'].select(range(min(TRAINING_SIZE, len(split['train']))))
val_ds = split['test'].select(range(min(VALIDATION_SIZE, len(split['test']))))
test_ds = split['test'].select(range(VALIDATION_SIZE, min(VALIDATION_SIZE + TEST_SIZE, len(split['test']))))
def chat_prompt(context, question):
"""User turn shared by the training rows and the evaluation requests."""
return f"Based on the following context, answer the question.\n\nContext: {context}\n\nQuestion: {question}"
def convert_squad(example):
"""Convert SQuAD to a chat `messages` row."""
return {
"messages": [
{"role": "user", "content": chat_prompt(example["context"], example["question"])},
{"role": "assistant", "content": example["answers"]["text"][0]},
]
}
def write_jsonl(dataset, path):
with open(path, "w", encoding="utf-8") as f:
for example in dataset:
f.write(json.dumps(convert_squad(example)) + "\n")
def write_test_jsonl(dataset, path):
"""Save test split with raw context/question for chat-style evaluation."""
with open(path, "w", encoding="utf-8") as f:
for example in dataset:
f.write(json.dumps({
"context": example["context"],
"question": example["question"],
"completion": example["answers"]["text"][0],
}) + "\n")
write_jsonl(train_ds, f"{DATASET_PATH}/training.jsonl")
write_jsonl(val_ds, f"{DATASET_PATH}/validation.jsonl")
write_test_jsonl(test_ds, f"{DATASET_PATH}/testing.jsonl")
print(f"Training: {len(train_ds)} rows")
print(f"Validation: {len(val_ds)} rows")
print(f"Test: {len(test_ds)} rows")
with open(f"{DATASET_PATH}/training.jsonl", 'r') as f:
sample = json.loads(f.readline())
print(f"\nSample user turn: {sample['messages'][0]['content'][:150]}...")
print(f"Sample assistant turn: {sample['messages'][1]['content']}")
DATASET_NAME = "kd-dataset"
try:
client.files.create_fileset(
body=CreateFilesetRequest(
name=DATASET_NAME,
description="Knowledge distillation training data",
),
).data()
print(f"Created fileset: {DATASET_NAME}")
except ConflictError:
print(f"Fileset '{DATASET_NAME}' already exists, continuing...")
fs = FilesetFileSystem(client=FilesClient.from_client(client))
fs.put(
lpath=f"{DATASET_PATH}/",
rpath=f"default/{DATASET_NAME}/",
recursive=True,
)
print("Uploaded files:")
print(json.dumps([f.model_dump() for f in client.files.list_files(name=DATASET_NAME).data().data], indent=2))

3. Secrets Setup

In this tutorial we use two Qwen3 models from Hugging Face:

Both models share the same tokenizer/vocabulary (required for knowledge distillation) and include a chat template for deployment with /chat/completions.

Hugging Face Authentication:

  • Both models are public, so the Hugging Face token is optional
  • For gated models, you must provide a Hugging Face token via the token_secret parameter and accept the model’s terms on its Hugging Face page
  • Get your token from Hugging Face Settings (requires Read access)
HF_TOKEN = os.getenv("HF_TOKEN")
def create_or_get_secret(name: str, value: str | None, label: str):
if not value:
print(f"{label} is not set - skipping setting secret")
return None
try:
secret = client.secrets.create_secret(
body=HelixSecretCreateRequest(name=name, value=SecretStr(value)),
).data()
print(f"Created secret: {name}")
return secret
except ConflictError:
print(f"Secret '{name}' already exists, continuing...")
return client.secrets.get_secret(name=name).data()
hf_secret = create_or_get_secret("hf-token", HF_TOKEN, "HF_TOKEN")
if hf_secret:
print("HF_TOKEN secret:")
print(hf_secret.model_dump_json(indent=2))

4. Create Model FileSets and Model Entities

Knowledge distillation requires two model entities:

  1. Student model — the smaller model that will be trained (Qwen/Qwen3-0.6B)
  2. Teacher model — the larger model that provides soft targets (Qwen/Qwen3-8B)

Both are post-trained checkpoints, so the output model includes a chat template, which is required for the /chat/completions inference endpoint.

from nemo_helix_plugin.files.storage_config import HuggingfaceStorageConfig
SPEC_TIMEOUT_SECONDS = 120
def create_model(hf_repo: str, model_name: str, description: str):
"""Create a fileset + model entity and wait for ModelSpec."""
try:
client.files.create_fileset(
body=CreateFilesetRequest(
name=model_name,
description=description,
storage=HuggingfaceStorageConfig(
repo_id=hf_repo,
repo_type="model",
token_secret=SecretRef(hf_secret.name) if hf_secret else None,
),
),
)
print(f"Created fileset: {model_name}")
except ConflictError:
print(f"Fileset '{model_name}' already exists.")
try:
model = client.models.create_model(
body=CreateModelEntityRequest(
name=model_name,
fileset=f"default/{model_name}",
),
).data()
print(f"Created Model Entity: {model_name}")
except ConflictError:
print(f"Model '{model_name}' already exists. Updating fileset.")
model = client.models.update_model(
name=model_name,
body=UpdateModelEntityRequest(
fileset=f"default/{model_name}",
),
).data()
print(f"Waiting for ModelSpec on {model_name}...")
spec_start = time.time()
while not model.spec:
if time.time() - spec_start > SPEC_TIMEOUT_SECONDS:
raise TimeoutError(f"ModelSpec not populated within {SPEC_TIMEOUT_SECONDS}s")
time.sleep(2)
model = client.models.get_model(name=model_name).data()
print(f"ModelSpec populated: {model.spec}")
return model
student_model = create_model(
hf_repo="Qwen/Qwen3-0.6B",
model_name="qwen3-0.6b-student",
description="Qwen3 0.6B student model",
)
print()
teacher_model = create_model(
hf_repo="Qwen/Qwen3-8B",
model_name="qwen3-8b-teacher",
description="Qwen3 8B teacher model",
)

Phase 1: Fine-Tune the Teacher

5. Train Teacher with Full SFT

For best distillation results, fine-tune the teacher on the same dataset that will be used for distillation. This ensures the teacher has learned the task-specific knowledge that the student will inherit.

We train the 8B teacher with Full SFT on the SQuAD dataset, sharded across 2 GPUs. Activation checkpointing and a micro batch of 1 keep it within 80 GB per GPU. The global batch and learning rate follow the Automodel 7B full-SFT SQuAD recipe qwen2_5_7b_squad.yaml.

from nemo_automodel_plugin.schema import AutomodelJobInput
job_suffix = uuid.uuid4().hex[:4]
TEACHER_JOB_NAME = f"teacher-sft-job-{job_suffix}"
TEACHER_OUTPUT_NAME = f"teacher-model-{job_suffix}"
teacher_spec = AutomodelJobInput(
model=f"default/{teacher_model.name}",
dataset={"training": f"default/{DATASET_NAME}"},
training={
"training_type": "sft",
"finetuning_type": "all_weights",
"max_seq_length": 2048,
"activation_checkpointing": True,
},
schedule={"epochs": 1},
batch={"global_batch_size": 64, "micro_batch_size": 1},
optimizer={"learning_rate": 1e-5, "weight_decay": 0.0},
parallelism={"num_gpus_per_node": 2},
output={"name": TEACHER_OUTPUT_NAME},
)
teacher_job = customization.automodel.jobs.create(
spec=teacher_spec, workspace="default", name=TEACHER_JOB_NAME
)
TRAINED_TEACHER_NAME = TEACHER_OUTPUT_NAME
print(f"Teacher training job: {teacher_job.job.name}")
print(f"Output teacher model: {TRAINED_TEACHER_NAME}")
import time
teacher_status = None
while True:
teacher_status = client.jobs.get_status(name=TEACHER_JOB_NAME, workspace="default")
print(f"Status: {teacher_status.status}")
if teacher_status.status.lower() in ("completed", "failed", "cancelled", "error"):
break
time.sleep(15)
assert teacher_status.status.lower() == "completed"

Phase 2: Establish Baseline (Base Student)

6. Deploy the Base Student Model

Before distillation, deploy the base student model (Qwen3-0.6B, without any fine-tuning) to establish a baseline ROUGE score. After distillation, we compare the distilled student against this baseline to measure improvement.

baseline_suffix = uuid.uuid4().hex[:4]
BASELINE_DEPLOYMENT_CONFIG = f"baseline-student-cfg-{baseline_suffix}"
BASELINE_DEPLOYMENT_NAME = f"baseline-student-{baseline_suffix}"
baseline_deployment_config = client.models.create_deployment_config(
body=CreateModelDeploymentConfigRequest(
name=BASELINE_DEPLOYMENT_CONFIG,
engine=Engine.VLLM,
model_spec=ModelDeploymentConfigModelSpec(
model_namespace="default",
model_name=student_model.name,
),
executor_config=ContainerExecutorConfig(
gpu=1,
image_name="docker.io/vllm/vllm-openai",
image_tag="v0.22.1",
),
),
).data()
baseline_deployment = client.models.create_deployment(
body=CreateModelDeploymentRequest(
name=BASELINE_DEPLOYMENT_NAME,
config=baseline_deployment_config.name,
),
).data()
print(f"Baseline student deployment: {baseline_deployment.name}")
def wait_for_deployment(deployment_name: str, timeout_minutes: int = 30):
"""Poll deployment until ready."""
start = time.time()
timeout = timeout_minutes * 60
while True:
dep = client.models.get_deployment(name=deployment_name).data()
elapsed = time.time() - start
clear_output(wait=True)
print(f"Deployment: {deployment_name}")
print(f"Status: {dep.status}")
print(f"Elapsed: {int(elapsed // 60)}m {int(elapsed % 60)}s")
if dep.status == "READY":
print("\nDeployment is ready!")
remaining = int(timeout - elapsed)
if remaining <= 0:
raise TimeoutError(f"Deployment timeout after {timeout_minutes} minutes")
if not client.models.wait_for_deployment_status(
deployment_name=deployment_name,
desired_status="READY",
workspace="default",
timeout=remaining,
):
raise TimeoutError("Deployment did not become ready")
if not client.models.wait_for_gateway(deployment_name, workspace="default", timeout=remaining):
raise TimeoutError("Inference gateway did not become ready")
return dep
if dep.status in ("FAILED", "ERROR", "TERMINATED", "LOST"):
raise RuntimeError(f"Deployment failed with status: {dep.status}")
if elapsed > timeout:
raise TimeoutError(f"Deployment timeout after {timeout_minutes} minutes")
time.sleep(15)
try:
dep_status = wait_for_deployment(BASELINE_DEPLOYMENT_NAME)
assert dep_status.status == "READY"
except Exception:
# Free GPUs if readiness fails before the later baseline-cleanup cell runs.
try:
client.models.delete_deployment(name=BASELINE_DEPLOYMENT_NAME)
if not client.models.wait_for_deployment_status(
deployment_name=BASELINE_DEPLOYMENT_NAME,
desired_status="DELETED",
workspace="default",
timeout=600,
):
raise TimeoutError(
f"Deployment {BASELINE_DEPLOYMENT_NAME} was not deleted within timeout"
)
client.models.delete_deployment_config(name=BASELINE_DEPLOYMENT_CONFIG)
except Exception as cleanup_error:
print(f"Baseline cleanup after readiness failure also failed: {cleanup_error}")
raise

7. Generate Baseline Predictions on Test Set

with open(f"{DATASET_PATH}/testing.jsonl", "r", encoding="utf-8") as f:
test_data = [json.loads(line) for line in f]
contexts = [row["context"] for row in test_data]
questions = [row["question"] for row in test_data]
reference_completions = [row["completion"] for row in test_data]
print(f"Test samples: {len(contexts)}")
print(f"Sample context: {contexts[0]}")
print(f"Sample question: {questions[0]}")
print(f"Sample reference: {reference_completions[0]}")
from nemo_helix_plugin.inference_gateway.client import InferenceGatewayClient
from nemo_helix_plugin.inference_gateway.types import JsonBody
def generate_completions(
deployment_name: str,
output_model_name: str,
contexts: list[str],
questions: list[str],
) -> list[str]:
"""Generate completions for a list of context/question pairs using a deployed model."""
gateway = InferenceGatewayClient.from_client(client)
completions = []
for context, question in zip(contexts, questions):
messages = [{"role": "user", "content": chat_prompt(context, question)}]
response = gateway.provider_post(
name=deployment_name,
workspace="default",
trailing_uri="v1/chat/completions",
body=JsonBody(
{
"model": f"default/{output_model_name}",
"messages": messages,
"temperature": 0,
"max_tokens": 128,
# Qwen3 reasons by default; turn it off for a short answer.
"chat_template_kwargs": {"enable_thinking": False},
}
),
).data()
completions.append(response["choices"][0]["message"]["content"])
return completions
print("Generating baseline (base student) predictions...")
baseline_completions = generate_completions(BASELINE_DEPLOYMENT_NAME, student_model.name, contexts, questions)
print(f"Generated {len(baseline_completions)} baseline predictions")
print(f"\nSample baseline output: {baseline_completions[0]}")

8. Delete Baseline Deployment

Delete the baseline student deployment to free GPU resources for the distillation training job and subsequent distilled model deployment.

client.models.delete_deployment(name=BASELINE_DEPLOYMENT_NAME)
print(f"Deleted baseline deployment: {BASELINE_DEPLOYMENT_NAME}")
if not client.models.wait_for_deployment_status(
deployment_name=BASELINE_DEPLOYMENT_NAME,
desired_status="DELETED",
workspace="default",
timeout=600,
):
raise TimeoutError(
f"Deployment {BASELINE_DEPLOYMENT_NAME} was not deleted within timeout"
)
client.models.delete_deployment_config(name=BASELINE_DEPLOYMENT_CONFIG)
print(f"Deleted baseline deployment config: {BASELINE_DEPLOYMENT_CONFIG}")

Phase 3: Distill into Student

9. Create Knowledge Distillation Job

Now create a distillation job that trains the 0.6B student using the fine-tuned 8B teacher’s output distribution. The KD loss settings and learning rate follow the Automodel reference recipe llama3_2_1b_kd.yaml. The model field specifies the student, and teacher_model references the trained teacher model entity from Phase 1.

GPU Requirements:

KD requires loading both student and teacher models, so plan GPU memory accordingly:

  • 0.6B student + 8B teacher: 1 GPU (80 GB)
  • 3B student + 8B teacher: 4 GPUs
  • 8B student + 70B teacher: 8+ GPUs

Use teacher_precision="bf16" (default) to reduce teacher memory footprint.

from nemo_automodel_plugin.schema import AutomodelJobInput
KD_JOB_NAME = f"my-kd-job-{job_suffix}"
KD_OUTPUT_NAME = f"kd-student-{job_suffix}"
kd_spec = AutomodelJobInput(
model=f"default/{student_model.name}",
dataset={"training": f"default/{DATASET_NAME}"},
training={
"training_type": "distillation",
"finetuning_type": "all_weights",
"teacher_model": f"default/{TRAINED_TEACHER_NAME}",
"teacher_precision": "bf16",
"distillation_ratio": 0.5,
"distillation_temperature": 1.0,
"max_seq_length": 2048,
},
schedule={"epochs": 2},
batch={"global_batch_size": 32, "micro_batch_size": 1},
optimizer={"learning_rate": 1e-5, "weight_decay": 0.0},
parallelism={"num_gpus_per_node": 1},
output={"name": KD_OUTPUT_NAME},
)
kd_job = customization.automodel.jobs.create(
spec=kd_spec, workspace="default", name=KD_JOB_NAME
)
DISTILLED_STUDENT_NAME = KD_OUTPUT_NAME
print(f"Distillation job: {kd_job.job.name}")
print(f"Output student model: {DISTILLED_STUDENT_NAME}")

10. Track Distillation Progress

The same progress display used for teacher training prints completion percentage and the latest loss, and charts loss by step once history is available.

import time
kd_status = None
while True:
kd_status = client.jobs.get_status(name=KD_JOB_NAME, workspace="default")
print(f"Status: {kd_status.status}")
if kd_status.status.lower() in ("completed", "failed", "cancelled", "error"):
break
time.sleep(15)
assert kd_status.status.lower() == "completed"

Phase 4: Evaluate the Distilled Student Model

11. Deploy the Distilled Student Model

The output model has the same architecture as the 0.6B student—only its weights have been updated via distillation. It requires just 1 GPU to deploy.

deploy_suffix_2 = uuid.uuid4().hex[:4]
STUDENT_DEPLOYMENT_CONFIG = f"kd-student-deploy-cfg-{deploy_suffix_2}"
STUDENT_DEPLOYMENT_NAME = f"kd-student-deploy-{deploy_suffix_2}"
student_deployment_config = client.models.create_deployment_config(
body=CreateModelDeploymentConfigRequest(
name=STUDENT_DEPLOYMENT_CONFIG,
engine=Engine.VLLM,
model_spec=ModelDeploymentConfigModelSpec(
model_namespace="default",
model_name=DISTILLED_STUDENT_NAME,
),
executor_config=ContainerExecutorConfig(
gpu=1,
image_name="docker.io/vllm/vllm-openai",
image_tag="v0.22.1",
),
),
).data()
student_deployment = client.models.create_deployment(
body=CreateModelDeploymentRequest(
name=STUDENT_DEPLOYMENT_NAME,
config=student_deployment_config.name,
),
).data()
print(f"Student deployment: {student_deployment.name}")
wait_for_deployment(STUDENT_DEPLOYMENT_NAME)

12. Generate Student Predictions on Test Set

print("Generating distilled student predictions...")
student_completions = generate_completions(STUDENT_DEPLOYMENT_NAME, DISTILLED_STUDENT_NAME, contexts, questions)
print(f"Generated {len(student_completions)} student predictions")
print(f"\nSample student output: {student_completions[0]}")

13. Compute ROUGE Scores

Compare the base student (before distillation) and the distilled student against the ground-truth reference completions using ROUGE metrics.

import evaluate
rouge = evaluate.load("rouge")
baseline_scores = rouge.compute(predictions=baseline_completions, references=reference_completions)
student_scores = rouge.compute(predictions=student_completions, references=reference_completions)
metrics = list(baseline_scores.keys())
header = f"{'Model':<35} " + " ".join(f"{m:>10}" for m in metrics)
separator = "-" * len(header)
print("=" * 60)
print("ROUGE SCORE COMPARISON")
print("=" * 60)
print(header)
print(separator)
print(f"{'Base Student (0.6B, no training)':<35} " + " ".join(f"{baseline_scores[m]:>10.4f}" for m in metrics))
print(f"{'Distilled Student (0.6B, KD)':<35} " + " ".join(f"{student_scores[m]:>10.4f}" for m in metrics))
print("Sample predictions (first 3):\n")
for i in range(min(3, len(contexts))):
print(f"--- Sample {i + 1} ---")
print(f"Question: {questions[i]}")
print(f"Reference: {reference_completions[i]}")
print(f"Baseline: {baseline_completions[i][:200]}")
print(f"Distilled: {student_completions[i][:200]}")
print()

Interpreting ROUGE Scores:

MetricMeasures
ROUGE-1Unigram overlap between prediction and reference
ROUGE-2Bigram overlap (captures phrase-level similarity)
ROUGE-LLongest common subsequence (captures sentence structure)
ROUGE-LsumROUGE-L computed over full summaries

What to expect:

  • The base student (0.6B, no training) provides a lower bound since it has not seen the task data
  • The distilled student (0.6B, KD) should significantly outperform the base student, demonstrating the knowledge transferred from the 8B teacher
  • If the distilled student scores are not much higher than the baseline, try increasing distillation_temperature, adjusting distillation_ratio, or training for more epochs

Hyperparameters

For detailed information on all available hyperparameters, recommended values, and tuning guidance, refer to the Hyperparameter Reference.


Troubleshooting

Job fails during model download:

  • Verify authentication secrets are configured (refer to Managing Secrets)
  • For gated Hugging Face models, accept the license on the model page
  • Check both model (student) and teacher_model URNs are correct
  • Ensure both model entities exist: client.models.get_model(name=...)

Job fails with OOM (Out of Memory) error:

KD loads both models, so OOM is more likely than with SFT:

  1. First try: Use teacher_precision="bf16" to reduce teacher memory
  2. Still OOM: Reduce micro_batch_size to 1
  3. Still OOM: Reduce global_batch_size and max_seq_length
  4. Last resort: Increase num_gpus_per_node

No chat template / /chat/completions fails:

  • Use post-trained checkpoints (e.g., Qwen3-0.6B) instead of base checkpoints (Qwen3-0.6B-Base). Base checkpoints may not include a chat template in their tokenizer, so the output model would also lack one.

Distilled model quality is poor:

  • Increase distillation_temperature (try 2.0–5.0) to transfer more nuanced knowledge
  • Adjust distillation_ratio—if dataset labels are high-quality, lower the ratio; if the teacher is strong, raise it
  • Increase epochs or max_steps for more training
  • Verify teacher and student share the same vocabulary

Vocabulary mismatch error:

  • Teacher and student must use the same tokenizer. Use models from the same family (e.g., Qwen3-0.6B + Qwen3-8B)

Deployment fails:

  • Verify output model exists: client.models.get_model(name=DISTILLED_STUDENT_NAME)
  • Check job logs: client.jobs.get_logs(name=kd_job.job.name)
  • Check student deployment status: client.models.get_deployment(name=student_deployment.name)
  • The distilled model has the same size as the student, so GPU requirements match the student model

Next Steps