Duration: 5

This quickstart walks you through a single self-contained Snowflake Notebook that trains a small student LLM using two on-policy distillation methods, evaluates them with paired statistics, and exposes the gradient-geometry intuition behind why one method needs a stability trick.

You will run, on a Snowflake GPU compute pool:

A Snowflake Notebook (inference_time_distillation_kl_clipping.ipynb) shipped with this guide, organised as:

The notebook is included as an asset of this guide:

Download inference_time_distillation_kl_clipping.ipynb

Duration: 10

Create or reuse a GPU compute pool with at least 40 GB VRAM. The notebook prints the detected GPU and decides at runtime whether to use 4-bit + LoRA.

Snowflake Container Runtime ships an older bitsandbytes than current transformers expects. Section 0 of the notebook installs bitsandbytes>=0.43.1 and a few peer pins, then runs a hard preflight that prints the torch CUDA build, the bitsandbytes import status, and transformers' is_bitsandbytes_available() verdict. If any of these fails, the cell raises a RuntimeError with the specific cause (no GPU, broken bnb wheel, stale lru_cache).

After the install completes the first time, restart the kernel so transformers re-imports the new bitsandbytes.

The preflight also surfaces the device and dtype the rest of the notebook will use:

The student is mistralai/Mistral-7B-Instruct-v0.2 loaded in 4-bit + LoRA (~5 GB on the GPU). The optional same-family teacher is mistralai/Mixtral-8x7B-Instruct-v0.1 in 4-bit (~24 GB). Tokenizer match is enforced at runtime — per-token reverse KL is undefined across vocabularies.

Duration: 5

Section 1 loads a small slice of GSM8K, defines gold_answer, extract_pred, correct, and build_prompt. The OPSD privileged-info trick is implemented inside build_prompt: when called with answer=... the gold answer is injected into the teacher's system prompt only.

Section 2 defines all per-token machinery (rollout, logprobs_at, per_token_reverse_kl) plus the eval statistics (wilson_ci, fmt_acc, mcnemar_p, bootstrap_gap_ci, difficulty_bucket, evaluate, pass_at_k). Everything that the slow eval cells call lives here, so re-running this cell is sub-second and safe to inspect.

Duration: 15

The notebook hits the Cortex REST API directly (not the SQL SNOWFLAKE.CORTEX.COMPLETE() function). The single entrypoint is client.complete(messages, model, max_tokens, temperature), which POSTs to /api/v2/cortex/inference:complete with Authorization: Snowflake Token="<session-token>" and parses the Server-Sent Events response into one {content, usage, model} dict.

Section 2.5 evaluates each teacher in CORTEX_MODELS (default: mistral-large2, llama4-maverick) on the same 200-example GSM8K val slice the student is graded on. For every teacher it reports:

This is the test that proves whether a Cortex teacher is worth distilling from. Initial small-n runs typically look like this — note the wide gaps that vanish under proper CIs:

After bumping N_VAL to 200 and adding paired statistics, the table becomes interpretable:

Duration: 20

Section 2.6 implements a calibrated process reward model that scores chain-of-thought reasoning step by step.

split_steps(rollout_text) runs before the judge ever sees the rollout, splitting on numbered-step regex, blank lines, and sentence boundaries. The judge receives a fixed (step_id, step_text) list and returns one 0/1 per step — no segmentation decision left to the model. This stabilises min / mean aggregation across runs.

_prm_cache memoises results by md5(judge_model, rubric, question, steps). Re-running with a different aggregation (min, mean, last) costs zero extra Cortex calls. Iterating on the rubric invalidates the cache automatically.

After definition the cell runs the judge on a single student rollout and prints the segmented step count, all three aggregations, and the cache size — sub-5-second sanity check.

End-to-end behavior of the PRM judge across multiple rollouts:

Section 2.6b is what graduates the PRM from "stub" to "non-stub". It runs prm_judge over every existing cortex_pred[m] rollout from Section 2.5 and reports Spearman rho and AUROC against the verifier ground truth, for each judge model and each aggregation:

Decision rule:

In a typical run, llama4-maverick + mean lands at AUROC ≈ 0.846 — borderline strong, ensemble or rubric tweak likely pushes it past 0.85.

Duration: 10

Section 2.7 computes pass@k for the student under its own distribution at temperature 0.7. This is the upper bound for any rejection-sampling SFT loop that uses student-self rollouts:

Section 2.7b answers the operational question: does the PRM judge approximate the verifier well enough to filter SFT data without gold answers? It generates a fresh sample of student rollouts, computes both verifier ground truth and prm_judge scores, sweeps τ across agg ∈ {min, mean}, and reports precision / recall / F1 with Wilson CIs. If precision > 0.95 with recall > 0.5 at some τ, the PRM filter is operationally usable.

Duration: 20

Section 4 demonstrates OPD with the local Mixtral-8x7B teacher: the student samples a rollout, both the student and the teacher compute logprobs at the generated positions, and the per-token reverse-KL KL(p_T || p_S) is summed over generated tokens as the loss.

Section 5 demonstrates OPSD with the same student model played as its own teacher, but with the gold answer injected into the teacher's system prompt. The student samples without the gold answer; the teacher's distribution sharpens around answer-consistent tokens. The same per-token reverse-KL kernel runs on these logprobs.

Section 6 plots the per-token KL distribution side-by-side. OPD's signal is diffuse — spread across many tokens. OPSD's is concentrated on a small number of pivot tokens that determine the answer. The max / mean ratio quantifies the concentration.

Duration: 10

Section 7 runs OPSD with and without kl.clamp(max=1.0). Without clipping, the pivot-token gradients drive the LoRA update into a regime that destabilises training within ~100 steps — losses blow up and the student collapses. With kl_clip=1.0, only the long tail of the per-token KL distribution is truncated; the bulk of the signal is preserved. The exact same code, just one flag.

Section 9 generalises everything into one function:

Four configurations of (alpha, lambda, pi_T) instantiate four classical algorithms:

This is the single most useful diagram in the notebook: every popular post-training method is one row of this table, distinguished only by who plays the teacher and whether the per-token KL is clipped.

Duration: 5

Mixtral 4-bit ~24 GB + Mistral-7B LoRA ~5 GB + activations + KV cache ≈ 35–45 GB peak. Fits comfortably on an A100-80GB or H100 GPU pool.