Knowledge Distillation
— a deep dive

import torch
import torch.nn as nn
import torch.nn.functional as F
from torch.utils.data import DataLoader
from transformers import (
    AutoTokenizer,
    AutoModelForSequenceClassification,
    BertConfig,
    BertForSequenceClassification,
    get_linear_schedule_with_warmup,
)
from datasets import load_dataset

DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
NUM_LABELS = 2        # e.g. SST-2 sentiment classification
TEMPERATURE = 4.0     # Hinton et al. recommend 3–5 for classification
ALPHA = 0.8           # weight on soft (distillation) loss
HARD_LOSS_ALPHA = 0.2 # weight on hard (CE) loss; must sum to 1
EPOCHS = 3
BATCH_SIZE = 32
LR = 3e-5

TEACHER_NAME = "textattack/bert-base-uncased-SST-2"

# ── Teacher: load a fine-tuned BERT-base ──────────────────────────────────────
teacher = AutoModelForSequenceClassification.from_pretrained(TEACHER_NAME)
teacher = teacher.to(DEVICE)
teacher.eval()  # freeze; no gradients needed
for param in teacher.parameters():
    param.requires_grad = False

# ── Student: a tiny BERT with only 4 layers ───────────────────────────────────
student_config = BertConfig(
    vocab_size=teacher.config.vocab_size,
    hidden_size=256,          # 768 in BERT-base → 256 (3× smaller)
    num_hidden_layers=4,       # 12 in BERT-base → 4
    num_attention_heads=4,
    intermediate_size=1024,
    num_labels=NUM_LABELS,
)
student = BertForSequenceClassification(student_config).to(DEVICE)

# ── Shared tokenizer ───────────────────────────────────────────────────────────
tokenizer = AutoTokenizer.from_pretrained("bert-base-uncased")

print(f"Teacher params: {sum(p.numel() for p in teacher.parameters()):,}")
print(f"Student params: {sum(p.numel() for p in student.parameters()):,}")

def distillation_loss(
    student_logits: torch.Tensor,
    teacher_logits: torch.Tensor,
    hard_labels: torch.Tensor,
    temperature: float = TEMPERATURE,
    alpha: float = ALPHA,
) -> torch.Tensor:
    """
    Combined distillation + task loss.

    Args:
        student_logits: raw logits from student  [batch, num_classes]
        teacher_logits: raw logits from teacher  [batch, num_classes]
        hard_labels:    ground-truth class ids   [batch]
        temperature:    softening temperature T
        alpha:          weight on soft (distillation) loss

    Returns:
        Scalar loss tensor.
    """
    # Soft probabilities — apply temperature to BOTH teacher and student
    teacher_probs = F.softmax(teacher_logits / temperature, dim=-1)
    student_log_probs = F.log_softmax(student_logits / temperature, dim=-1)

    # KL divergence: KL(teacher ‖ student)
    # F.kl_div expects (log_input, target); output is mean over batch
    soft_loss = F.kl_div(
        student_log_probs,
        teacher_probs,
        reduction="batchmean",
    )

    # Rescale by T² to preserve gradient magnitude (Hinton et al.)
    soft_loss = soft_loss * (temperature ** 2)

    # Hard cross-entropy loss against ground-truth labels
    hard_loss = F.cross_entropy(student_logits, hard_labels)

    # Combine
    total_loss = alpha * soft_loss + (1.0 - alpha) * hard_loss
    return total_loss

# ── Dataset ────────────────────────────────────────────────────────────────────
dataset = load_dataset("glue", "sst2")

def tokenize(batch):
    return tokenizer(
        batch["sentence"],
        padding="max_length",
        truncation=True,
        max_length=128,
        return_tensors="pt",
    )

train_dataset = dataset["train"].map(tokenize, batched=True)
train_dataset.set_format("torch", columns=["input_ids", "attention_mask", "label"])
train_loader = DataLoader(train_dataset, batch_size=BATCH_SIZE, shuffle=True)

# ── Optimizer & scheduler ──────────────────────────────────────────────────────
optimizer = torch.optim.AdamW(student.parameters(), lr=LR, weight_decay=0.01)
total_steps = len(train_loader) * EPOCHS
scheduler = get_linear_schedule_with_warmup(
    optimizer,
    num_warmup_steps=total_steps // 10,
    num_training_steps=total_steps,
)

# ── Training ───────────────────────────────────────────────────────────────────
student.train()
for epoch in range(EPOCHS):
    running_loss = 0.0
    for step, batch in enumerate(train_loader):
        input_ids      = batch["input_ids"].to(DEVICE)
        attention_mask = batch["attention_mask"].to(DEVICE)
        labels         = batch["label"].to(DEVICE)

        # Get teacher logits (no grad)
        with torch.no_grad():
            teacher_out = teacher(
                input_ids=input_ids,
                attention_mask=attention_mask,
            )
            teacher_logits = teacher_out.logits

        # Get student logits
        student_out = student(
            input_ids=input_ids,
            attention_mask=attention_mask,
        )
        student_logits = student_out.logits

        # Compute combined loss
        loss = distillation_loss(
            student_logits=student_logits,
            teacher_logits=teacher_logits,
            hard_labels=labels,
            temperature=TEMPERATURE,
            alpha=ALPHA,
        )

        optimizer.zero_grad()
        loss.backward()
        torch.nn.utils.clip_grad_norm_(student.parameters(), 1.0)
        optimizer.step()
        scheduler.step()

        running_loss += loss.item()
        if step % 100 == 0:
            print(f"Epoch {epoch+1} | Step {step} | Loss: {loss.item():.4f}")

    avg = running_loss / len(train_loader)
    print(f"\n── Epoch {epoch+1} complete. Avg loss: {avg:.4f} ──\n")

# ── Save ───────────────────────────────────────────────────────────────────────
student.save_pretrained("./distilled-tiny-bert")
tokenizer.save_pretrained("./distilled-tiny-bert")
print("✓ Distilled student saved to ./distilled-tiny-bert")

class HintLayer(nn.Module):
    """Projects student hidden states → teacher hidden dim."""
    def __init__(self, student_dim: int, teacher_dim: int):
        super().__init__()
        self.proj = nn.Linear(student_dim, teacher_dim, bias=False)

    def forward(self, x):
        return self.proj(x)

# One hint layer per matched pair: student layer 2 → teacher layer 6
hint = HintLayer(
    student_dim=student_config.hidden_size,   # 256
    teacher_dim=teacher.config.hidden_size,   # 768
).to(DEVICE)

def feature_distillation_loss(
    student_hidden: torch.Tensor,  # [batch, seq_len, 256]
    teacher_hidden: torch.Tensor,  # [batch, seq_len, 768]
) -> torch.Tensor:
    """MSE loss between projected student states and teacher states."""
    projected = hint(student_hidden)  # [batch, seq_len, 768]
    return F.mse_loss(projected, teacher_hidden.detach())

# ── In your training loop, add to total loss: ──────────────────────────────────
# student_out = student(..., output_hidden_states=True)
# teacher_out = teacher(..., output_hidden_states=True)
#
# feat_loss = feature_distillation_loss(
#     student_hidden=student_out.hidden_states[2],   # student layer 2
#     teacher_hidden=teacher_out.hidden_states[6],   # teacher layer 6
# )
# loss = response_loss + 0.1 * feat_loss

from trl import GKDConfig, GKDTrainer
from transformers import AutoModelForCausalLM, AutoTokenizer
from datasets import load_dataset

# ── Load teacher and student ───────────────────────────────────────────────────
teacher = AutoModelForCausalLM.from_pretrained(
    "meta-llama/Meta-Llama-3-8B-Instruct",
    torch_dtype=torch.bfloat16,
    device_map="auto",
)
student = AutoModelForCausalLM.from_pretrained(
    "meta-llama/Llama-3.2-1B-Instruct",  # student is 1B
    torch_dtype=torch.bfloat16,
    device_map="auto",
)
tokenizer = AutoTokenizer.from_pretrained("meta-llama/Llama-3.2-1B-Instruct")

# ── Dataset ────────────────────────────────────────────────────────────────────
dataset = load_dataset("HuggingFaceH4/ultrachat_200k", split="train_sft[:10000]")

# ── GKD config ────────────────────────────────────────────────────────────────
training_args = GKDConfig(
    output_dir="./llama-1b-distilled",
    num_train_epochs=1,
    per_device_train_batch_size=2,
    gradient_accumulation_steps=8,
    learning_rate=2e-5,
    bf16=True,
    temperature=2.0,           # T for soft label computation
    lmbda=0.5,                 # λ = weight on distillation loss (vs. SFT loss)
    seq_kd=False,              # True = use sequence-level KD (teacher generates)
    max_seq_length=1024,
    logging_steps=10,
    save_strategy="epoch",
)

# ── Trainer ───────────────────────────────────────────────────────────────────
trainer = GKDTrainer(
    model=student,
    teacher_model=teacher,
    args=training_args,
    tokenizer=tokenizer,
    train_dataset=dataset,
)

trainer.train()
trainer.save_model("./llama-1b-distilled")