🛠️ Lab 14 — Distilling a Causal LM with Transformers and TRL's GKDTrainer (50 pts)

Goal: Distill a larger open-source causal language model ("teacher") into a smaller student model using two complementary approaches — a Trainer-API supervised distillation pattern, and TRL's on-policy GKDTrainer — then run a quality-vs-size tradeoff analysis comparing your distilled student to a same-size non-distilled baseline. All work happens in a free Google Colab notebook using free/local, open-source tools only. Steps:

1. Set up your free Colab environment. Open a new Colab notebook, select a free GPU runtime (Runtime → Change runtime type → T4 GPU), and install transformers, trl, datasets, accelerate, and evaluate.
2. Pick a small teacher/student pair. To keep this runnable on a free GPU, choose an open teacher/student pair from the same model family where the student is a clear fraction of the teacher's size (for example, a \~1B-parameter teacher paired with a \~135M–360M-parameter student from the same tokenizer family). Load both with transformers.
3. Prepare a small instruction or task dataset. Pull a small public dataset from the Hugging Face Hub appropriate for causal LM fine-tuning (for example, a small instruction-following or summarization subset). Keep the training subset small (a few thousand examples) so training completes in Colab's session limits.
4. Adapt the Trainer-API distillation pattern (supervised, forward-pass baseline). Following the methodology in the Transformers knowledge-distillation guide (originally written for image classification), adapt the same teacher/student Trainer pattern to your causal LM pair: run the teacher's forward pass to get logits over your dataset, and train the student with a combined loss (standard language-modeling loss on hard labels plus a KL term between student and teacher logits). Train for a small number of epochs.
5. Distill with TRL's GKDTrainer (on-policy). Using the TRL GKD Trainer documentation, configure GKDTrainer with your same teacher/student pair. Train a second student copy this way, letting the student generate its own sequences during training with teacher feedback via the generalized JSD loss.
6. Tune the lambda blending parameter. Run at least two GKDTrainer configurations with different values of the lambda parameter that blends the supervised JSD loss with the on-policy, student-generated-sequence loss (for example, a lower value favoring supervised loss vs. a higher value favoring on-policy loss). Record training loss curves for each.
7. Train a same-size non-distilled baseline. Train a fresh copy of the student architecture directly on the hard labels in your dataset (standard supervised fine-tuning, no teacher involved at all). This is your control condition.
8. Run the quality-vs-size tradeoff analysis. For the teacher, the two distilled students (Trainer-API and best GKDTrainer lambda), and the non-distilled baseline, measure: (a) perplexity or task accuracy on a held-out split, (b) parameter count, and (c) inference latency (average seconds per generated response over a fixed batch of prompts, measured locally in the Colab runtime). Build one table with all four models as rows.
9. Write your analysis. In 300–400 words, answer: Did either distilled student beat the non-distilled baseline of the same size? Which distillation method (Trainer-API vs. GKDTrainer) performed better, and does that match what the reverse-KL / on-policy theory from lecture predicts? What was the accuracy (or perplexity) gained per parameter saved and per unit of latency change, relative to the teacher?

Deliverables: Submit your Colab notebook (.ipynb, shared link or exported file) containing all training runs and code, your tradeoff-analysis table, and your written analysis — either as an uploaded file or pasted into the text entry box.