Distilling frontier reasoning into an edge model — Unsloth + W&B in Colab, Axolotl on Blackwell

The edge deployment story doesn’t start with a smaller model — it starts with a better teacher. Frontier models like Claude Opus 4.6 and Qwen3.5 produce reasoning traces you can’t run on a Jetson, but you can distill their behavior into a 3B–8B student and ship that student as a GGUF. I walked that path recently: mix high-quality reasoning datasets, fine-tune in Google Colab with Unsloth and a Weights & Biases dashboard, then move the serious run to Axolotl on my RTX PRO 6000 Blackwell workstation.

Table of contents

The goal: frontier teacher → edge student

I want a model that thinks in structured reasoning blocks — <think>...</think> wrapped traces followed by a final answer — but runs on hardware I actually deploy: Jetson Orin Nano 8 GB, GeniePod-class edge boards, or a single consumer/workstation GPU. That means:

1. Collect reasoning-heavy SFT data distilled from frontier models (not generic instruction fluff).
2. Fine-tune a small base — unsloth/Llama-3.2-3B in Colab — with wide LoRA (r=64) on all projection layers; graduate to larger students on Blackwell via Axolotl when the data mix is proven.
3. Track every run in W&B so I can compare LR schedules, data mixes, and loss curves without guessing.
4. Export merged weights + Q4_K_M GGUF for sparkinfer / llama.cpp inference on edge.

My Colab prototype is a Llama 3.2 3B run (unsloth/Llama-3.2-3B, Llama-3 chat template, train_on_responses_only) — the notebook export on disk was mislabeled qwopus3_5_27b_colab.py, but nothing inside targets Qwen or a 27B model; the student is 3B Llama end to end. Good for proving the frontier-reasoning data mix. For the run I care about on Blackwell, I graduate to Axolotl. Three reasons, below.

The dataset: ~14k reasoning traces from three teachers

The core idea in that Llama 3.2 3B Colab notebook is mixing frontier reasoning datasets into one normalized conversations format, then applying the Llama-3 chat template and filtering by length. My mix:

~14,233 examples after sampling, shuffle, dedup, and length filter (MAX_CONTEXT_WINDOW = 8192).

Every assistant turn is normalized to:

<think>reasoning trace here</think>
final answer

That format matters: it’s the behavior I want the edge student to learn — show work, then answer. The pipeline loads each dataset, maps to conversations, concatenates, shuffles with a fixed seed, tokenizes through the chat template, and drops sequences over 8K tokens or with malformed thinking blocks.

For Colab prototyping, this is all Python — datasets.map, filter, SFTTrainer. It works. It’s also the kind of script that becomes fragile when you need to reproduce a run six months later on different hardware.

Phase 1: Unsloth + W&B in Colab — Llama 3.2 3B prototype

Unsloth is the right tool for the first pass on unsloth/Llama-3.2-3B: one GPU, minimal boilerplate, optimized LoRA path (use_gradient_checkpointing = "unsloth" saves ~30% VRAM), and train_on_responses_only with the Llama-3 header tokens so you don’t waste gradients on user turns.

The Colab flow:

from unsloth import FastLanguageModel

model, tokenizer = FastLanguageModel.from_pretrained(
    model_name="unsloth/Llama-3.2-3B",
    max_seq_length=2048,
    load_in_4bit=True,
)

model = FastLanguageModel.get_peft_model(
    model,
    r=64,
    lora_alpha=64,
    target_modules=["q_proj", "k_proj", "v_proj", "o_proj",
                    "gate_proj", "up_proj", "down_proj"],
    use_gradient_checkpointing="unsloth",
)

Training hooks into W&B via report_to="wandb" and an explicit wandb.init(project="Llama-3.2-3B-Finetuning", ...). The dashboard screenshot above is what you get: loss per step, hyperparameter panel, run comparison — the difference between “I think epoch 2 helped” and “epoch 2 helped 12% on the eval slice.”

Colab secrets hold WANDB_API_KEY and HF_TOKEN; checkpoints land on Google Drive; export pushes merged 16-bit weights and GGUF (q4_k_m, q8_0, bf16) back to Hugging Face Hub.

What Unsloth is great at: proving the data mix works, finding a sane LR (2e-4 with adamw_8bit, batch 1 × grad-accum 12), and getting a first GGUF you can smoke-test on edge inference the same day.

What it isn’t: a reproducible, version-controlled training system you hand to a colleague or re-run on a Blackwell workstation without re-reading 600 lines of notebook cells.

That’s where Axolotl comes in.

Phase 2: why Axolotl (three reasons)

Axolotl has become the industry-standard choice for serious post-training — Nous Research, OpenPipe, and most teams I respect use it — because it solves the biggest headache in fine-tuning: reproducibility and execution complexity. Instead of gluing together Hugging Face transformers, peft, accelerate, and distributed backends in fragile Python, you manage the entire run in a single declarative YAML file.

On an RTX PRO 6000 Blackwell, that’s the framework I graduate to. Three reasons:

1. Declarative, version-controlled configs

Everything lives in YAML: base model, dataset format, tokenization, LoRA ranks, hardware knobs. If a run fails at 3 a.m., you don’t debug a notebook — you diff two 50-line config files. You hand a colleague the YAML, not a Colab link with hardcoded paths. Fine-tuning becomes an engineering practice, not an ad-hoc art form.

2. Multipack tokenization

Short instructions are common in SFT. A 400-token sample in a 4096-token window means ~90% of your context block is padding — wasted compute and memory bandwidth on empty tokens.

Axolotl’s sample packing (sample_packing: true) packs multiple distinct sequences into one full context block with an attention mask so samples don’t leak into each other. On reasoning datasets with variable-length traces, this alone can cut wall-clock training time dramatically. Unsloth can train; Axolotl trains efficiently at scale.

3. Cutting-edge architectures + Blackwell FP8

The Axolotl community merges support for new techniques fast — GRPO, DPO, reward modeling — and on Blackwell it integrates with PyTorch TorchAO for experimental FP8 mixed precision via torch.compile, mapping to 5th-gen Tensor Cores on the PRO 6000. That’s the hardware-level win Unsloth’s single-GPU Colab path doesn’t expose cleanly.

A minimal Axolotl config

A wide-rank LoRA on an 8B-class student (swap base_model for your edge target):

base_model: meta-llama/Meta-Llama-3-8B-Instruct
model_type: LlamaForCausalLM
tokenizer_type: PreTrainedTokenizerFast

datasets:
  - path: frontier_reasoning_mix.jsonl
    type: sharegpt
    conversation: llama3

adapter: lora
lora_r: 64
lora_alpha: 128
lora_dropout: 0.05
lora_target_modules:
  - q_proj
  - v_proj
  - k_proj
  - o_proj
  - gate_proj
  - up_proj
  - down_proj

sequence_len: 4096
sample_packing: true
pad_to_sequence_len: true

bf16: true
fp8: false              # true + torch_compile when using TorchAO on Blackwell

gradient_accumulation_steps: 2
micro_batch_size: 4
optimizer: adamw_torch
torch_compile: false

One command:

accelerate launch -m axolotl.cli.train config.yaml

Same W&B integration (wandb in the Axolotl config’s wandb_project field) — the dashboard stays; the execution layer upgrades.

The pipeline end-to-end

  Frontier teachers          ~14k reasoning         Colab prototype
  (Opus 4.6 / Qwen3.5)  ──►  SFT mix          ──►   (Unsloth + W&B)
                                                            │
                                                     good enough?
                                    ┌───────────────────────┴───────────────────────┐
                                    ▼ smoke test                        serious run  ▼
                          GGUF Q4_K_M                                Axolotl YAML
                          (edge inference)                  (RTX PRO 6000 Blackwell)
                                    ▲                                          │
                                    └───────────────── Merged weights + GGUF ──┘

What I’d do differently next time

1. Start the Axolotl YAML on day one for anything I intend to re-run — use Unsloth only for the first 500-step smoke test, not the full 2-epoch run.
2. Log the dataset hash in W&B alongside loss — when the teacher mix changes, the run name should too.
3. Eval on reasoning benchmarks, not just training loss — IFEval and a held-out reasoning slice tell you if the student actually learned to think, not just to lower CE loss on memorized traces.
4. Export early, test on edge hardware often — a GGUF that benchmarks well on a 5090 can still OOM or quality-collapse on 8 GB Jetson; test where you ship.

The takeaway

Distilling frontier reasoning into an edge model is a data problem first, a tracking problem second, and a framework hygiene problem third. Unsloth + W&B in Colab gets you to a first GGUF fast. Axolotl on Blackwell is what makes the run reproducible, packed, and hardware-aware — the difference between a notebook experiment and something you’d actually ship on a GeniePod or Jetson agent stack.

The teacher is cloud-scale. The student has to fit in 8 GB. The path between them is ~14k well-formed reasoning traces, a W&B dashboard you trust, and a YAML file you can re-run six months from now without opening Colab.