feat: add Self-Distillation Fine-Tuning (SDFT) algorithm

[slacki-ai]

Summary

Implements Self-Distillation Fine-Tuning (SDFT) from the paper arxiv 2601.19897 as a new loss="sdft" option in the unsloth training job.

SDFT can achieve higher new-task accuracy while substantially reducing catastrophic forgetting compared to standard SFT, by using the model itself (conditioned on a demonstration) as a teacher signal.

Algorithm

SDFT trains the student (model without demonstrations) to match the token-level distribution of the teacher (same model conditioned on an in-context demonstration) via reverse KL divergence over response tokens:

L(θ) = (1/K) Σ_{t∈response} KL( π_θ(·|y_<t, x)  ‖  π_φ(·|y_<t, x, c) )

• π_θ = student (trainable LoRA model, no demonstration)
• π_φ = teacher (EMA of student weights, sees demonstration)
• x = prompt, c = demonstration, K = number of response tokens

After each optimizer step the teacher EMA is updated: φ ← α·θ + (1−α)·φ

The EMA teacher only tracks the trainable (LoRA adapter) parameters — the frozen base model weights are shared, so no extra GPU memory is needed for a second full model copy.

New files

File	Description
openweights/jobs/unsloth/sdft.py	SDFTTrainer, SDFTDataCollator, EMATeacherCallback, sdft_train()
cookbook/sdft/sdft_qwen3_4b.py	Minimal usage example
cookbook/sdft/data/train.jsonl	Sample training data with demonstration fields

Modified files

File	Change
openweights/jobs/unsloth/validate.py	Add "sdft" to loss Literal; add sdft_ema_alpha (default 0.02) and sdft_demo_template fields
openweights/jobs/unsloth/training.py	Import sdft_train; preserve demonstration field in create_dataset; route loss == "sdft" to sdft_train

Data format

Same conversations JSONL as standard SFT, with an optional demonstration field:

{
  "messages": [
    {"role": "user",      "content": "What is 2+2?"},
    {"role": "assistant", "content": "The answer is 4."}
  ],
  "demonstration": "2+2 equals 4 because addition combines two quantities."
}

When "demonstration" is absent, the trainer automatically falls back to using the last assistant message as the teacher's in-context demo.

Usage

from openweights import OpenWeights
ow = OpenWeights()

training_file = ow.files.upload("train.jsonl", purpose="conversations")["id"]

job = ow.fine_tuning.create(
    model="unsloth/Qwen3-4B",
    training_file=training_file,
    loss="sdft",
    sdft_ema_alpha=0.02,      # EMA rate for teacher (paper recommends 0.01–0.05)
    epochs=1,
    learning_rate=1e-4,
    r=32,
)

Implementation notes

• Weight-swapping: Teacher logits are computed by temporarily replacing the LoRA adapter weights with their EMA values, running a no_grad forward pass, then restoring the student weights before the backward pass. The autograd graph is never contaminated.
• EMA timing: EMATeacherCallback.on_step_end fires after the HuggingFace optimizer step, ensuring the teacher always tracks the updated student.
• Logit alignment: Teacher sequences have a longer prefix (demo context). We align by taking the last T_student positions from the shifted teacher logit tensor, since both sequences share the same suffix (prompt + response).
• Graceful fallback: If teacher_input_ids is absent from a batch (e.g. during evaluation), compute_loss falls back to the standard SFT cross-entropy loss.

Test plan

• Verify validate.py accepts loss="sdft" and rejects loss="sdft" with a preference-* training file
• Verify loss="sft" and loss="dpo" still work unchanged
• Run sdft_qwen3_4b.py on a small model + dataset to confirm a training job completes
• Check that EMA teacher weights diverge from student weights after a few steps (confirming EMA is updating)
• Check that sdft_ema_alpha=0.0 freezes the teacher (loss converges quickly) and sdft_ema_alpha=1.0 sets teacher == student each step (loss → 0)

🤖 Generated with Claude Code