Distill a LoRA Student from a Frozen Teacher on One GPU with torchtune

Knowledge distillation on a single GPU lets you compress a large frozen teacher into a LoRA-tuned student while keeping the training loop inside one process. This guide walks through PyTorch torchtune's knowledge_distillation_single_device.py, available at https://github.com/pytorch/torchtune under the BSD-3-Clause license. The recipe does not shard the teacher, offload it to disk, or use pipeline parallelism. Both models live on the same accelerator, the teacher runs under torch.no_grad, and only the student's LoRA parameters receive gradients. That makes the code easy to inspect, but total memory is the sum of teacher weights, student weights, and the student optimizer state.

What this guide adds is a section-by-section implementation read: the exact config and loader logic, where loss and gradient scaling happen, and the failure modes a reader hits when VRAM is tight. Recent single-GPU model releases, such as the NVIDIA Nemotron 3.5 Super VL 120B report, show how much can now run on one accelerator, but this KD recipe still needs to be sized deliberately. With a 24 GB card, assume small models, LoRA, bf16 if available, and aggressive gradient accumulation.

Prerequisites

The source recipe is explicit about its boundaries. CPU training is not supported; use a CUDA GPU for the training device. For bf16, the GPU must support bfloat16, otherwise the recipe falls back to fp32. On the software side you need PyTorch, torchtune, torchdata, omegaconf, and tqdm. The optional bitsandbytes library enables 8-bit AdamW and Paged AdamW optimizers, which reduce optimizer memory but do not make this a full quantization-aware training recipe. You also need a teacher checkpoint readable through torchtune's CheckpointClient.

The recipe is driven by an OmegaConf YAML config. Run tune ls to see the available knowledge distillation configs and their example launch commands. You can override any config value from the command line.

Step 1: Lock in precision and memory policy

The recipe rejects fp16 outright because it does not implement gradient scaling. The constructor converts the config's dtype string with torchtune's training helper, then raises if the result is float16:

        self._device = utils.get_device(device=cfg.device)
        # Reduced precision logic
        self._dtype = training.get_dtype(cfg.dtype, device=self._device)
        # fp16 precision is explicitly disabled as it is not supported in this
        # recipe (for example, no gradient scaling).
        if self._dtype == torch.float16:
            raise ValueError(
                "fp16 precision is not supported in this recipe. Please use fp32 or bf16."
            )

This guard matters because the student model, teacher model, and optimizer all use the same dtype. On Ampere or newer GPUs, dtype: bf16 roughly halves activation and gradient memory compared with fp32. If the GPU does not support bf16, training.get_dtype falls back to fp32.

Two other memory controls interact here. Activation checkpointing recomputes activations during the backward pass instead of storing them, saving memory at the cost of slower training. Gradient accumulation simulates a larger batch by summing gradients over several micro-batches before updating weights. The table below summarizes the tradeoffs.

Step 2: Prepare the dataset and stateful loader

The recipe supports a single map-style dataset or a list of datasets concatenated with ConcatDataset. Packed and non-packed datasets use different collation paths: non-packed samples pad to the longest sequence with padded_collate_sft, while packed datasets use padded_collate_packed. The loader is stateful so a resumed run can restore its sampler position:

        if isinstance(cfg_dataset, ListConfig):
            datasets = [
                config.instantiate(single_cfg_dataset, self._tokenizer)
                for single_cfg_dataset in cfg_dataset
            ]
            ds = ConcatDataset(datasets=datasets)
            packed = getattr(ds, "packed", False)
        else:
            ds = config.instantiate(cfg_dataset, self._tokenizer)
            packed = cfg_dataset.get("packed", False)

        sampler = StatefulDistributedSampler(
            ds,
            num_replicas=1,
            rank=0,
            shuffle=shuffle,
            seed=0,
        )
        dataloader = StatefulDataLoader(
            dataset=ds,
            batch_size=batch_size,
            sampler=sampler,
            collate_fn=(
                partial(
                    padded_collate_sft,
                    padding_idx=self._tokenizer.pad_id,
                    ignore_idx=self._loss_fn.ignore_index,
                )
                if not packed
                else padded_collate_packed
            ),
            # dropping last avoids shape issues with compile + flex attention
            drop_last=True,
        )

The collator uses the tokenizer's pad ID and the loss function's ignore index. drop_last=True is deliberate: a partial final batch can break shape assumptions when torch.compile and flex attention are enabled. That also means the number of optimizer steps per epoch is computed from the loader length divided by gradient_accumulation_steps, so a dataset that is not divisible by the effective batch size will silently lose the remainder.

Step 3: Compute the distillation loss

Each batch produces two losses: the standard token-level cross-entropy against labels, and a KD loss against the frozen teacher's logits. The teacher runs under torch.no_grad, and the code deletes both logit tensors after computing losses to prevent a backward memory spike:

    def _loss_step(
        self, batch: dict[str, torch.Tensor]
    ) -> tuple[torch.Tensor, torch.Tensor]:
        # Both are shape [b, s]
        tokens, labels = batch["tokens"], batch["labels"]

        # Get the attention mask and position ids from the dataset if they
        # exist. Currently, only sample packing in PackedDataset returns these
        mask = batch.get("mask", None)  # shape [b, s, s]
        input_pos = batch.get("input_pos", None)  # shape [b, s]

        # run model
        logits = self._model(tokens, mask=mask, input_pos=input_pos)

        # Compute teacher logits
        with torch.no_grad():
            teacher_logits = self._teacher_model(tokens, mask=mask, input_pos=input_pos)

        if not isinstance(logits, list):
            labels = labels.reshape(-1)
            logits = logits.reshape(-1, logits.size(-1))
            teacher_logits = teacher_logits.reshape(-1, teacher_logits.size(-1))

        # Compute kd loss
        kd_loss = self._kd_loss_fn(logits, teacher_logits, labels)

        # Compute loss
        loss = self._loss_fn(logits, labels)

        # free logits otherwise it peaks backward memory
        del logits
        del teacher_logits

        return loss, kd_loss

The kd_ratio config value, which defaults to 0.5, blends the two objectives in the training loop as (1 - kd_ratio) * class_loss + kd_ratio * kd_loss. If the loss is CEWithChunkedOutputLoss, the code sets the same number of output chunks on both student and teacher and asserts the two loss functions match. Linear projection losses are explicitly rejected for KD in this recipe.

Step 4: Run the accumulation-aware training loop

The training loop accumulates loss weighted by the number of unmasked tokens in each micro-batch, then steps only when the gradient accumulation boundary is reached. Gradients are scaled by 1 / num_tokens so the update approximates a full batch whose examples were summed, not averaged:

                class_loss, kd_loss = self._loss_step(batch)
                running_class_loss += class_loss * current_num_tokens
                running_kd_loss += kd_loss * current_num_tokens
                current_loss = (
                    1 - self._kd_ratio
                ) * class_loss + self._kd_ratio * kd_loss
                current_loss.backward()

                # Step with optimizer
                if (idx + 1) % self._gradient_accumulation_steps == 0:
                    training.scale_grads(self._model, 1 / num_tokens)
                    if self._clip_grad_norm is not None:
                        grad_norm = torch.nn.utils.clip_grad_norm_(
                            self._model.parameters(),
                            max_norm=float(self._clip_grad_norm),
                        )
                    self._optimizer.step()
                    self._optimizer.zero_grad(set_to_none=True)
                    self._lr_scheduler.step()
                    # Update the number of steps when the weights are updated
                    self.global_step += 1

The loop logs combined loss, class_loss, kd_loss, learning rate, tokens per second, and optional peak memory. Mid-epoch checkpoints are saved every save_every_n_steps if configured; otherwise the recipe checkpoints once per epoch. The final checkpoint always saves full tensors.

Step 5: Launch through the parsed entry point

The bottom of the recipe is the OmegaConf-annotated entry point that reads YAML, instantiates the class, runs setup, training, and cleanup:

@config.parse
def recipe_main(cfg: DictConfig) -> None:
    """
    Entry point for the recipe.

    Configurable parameters are read in the following order:
        - Parameters specified in config (see available configs through ``tune ls``)
        - Overwritten by arguments from the command-line
    """
    config.log_config(recipe_name="KDRecipeSingleDevice", cfg=cfg)
    recipe = KDRecipeSingleDevice(cfg=cfg)
    recipe.setup(cfg=cfg)
    recipe.train()
    recipe.cleanup()

Validation is not built into the training loop. After training, run the appropriate torchtune generation or EleutherAI evaluation config from tune ls against the merged checkpoint. The recipe can save adapter-only weights or full merged weights depending on save_adapter_weights_only; for downstream evaluation you usually want a merged checkpoint.

What to watch out for

Teacher dropout does not disappear in eval mode. The _setup_teacher_model comment points to PyTorch issue 124464: SDPA dropout is not disabled by .eval(). The teacher may produce slightly different logits across forward passes. If deterministic distillation matters, compare teacher outputs across two passes before trusting a single run.

Both models share one accelerator. There is no sharding or offload. Add the teacher weights, student weights, activations for both forward passes, and the student optimizer state. The source logs memory after each model initialization; read those numbers before increasing batch size. On memory-constrained parts, treat low_cpu_ram, LoRA rank, activation checkpointing, and bitsandbytes optimizers as first-class controls.

drop_last=True trains on the aligned remainder. If dataset length is not a multiple of the effective batch size, examples in the dropped final batch never participate. Check len(dataloader) // gradient_accumulation_steps against expected steps per epoch, and set max_steps_per_epoch explicitly if needed. A short dataset can otherwise produce very few optimizer steps per epoch.

Resumption is epoch-bound, not arbitrary step. The recipe documentation states that optimizer state and recipe state are saved at the end of an epoch and that mid-epoch checkpointing is not currently supported. save_every_n_steps can create intermediate step checkpoints, but verify the checkpointer includes optimizer, scheduler, and dataloader state before treating them as resume points.

Gradient accumulation changes gradient magnitude. The loop sums token-weighted losses over micro-batches and scales all gradients by 1 / num_tokens. Changing gradient_accumulation_steps changes num_tokens and therefore the scale factor, so keep the learning rate in mind when tuning GA.

kd_ratio is a blunt control. The default of 0.5 blends class loss and KD loss. Watch class_loss and kd_loss separately during tuning; high values can transfer teacher biases, and very low values make the run look like ordinary SFT.

Where to go next

After the final save, inspect the checkpoint directory and run the appropriate torchtune evaluation or generation config against the merged weights. Torchtune's checkpointer deep dive explains how intermediate and final checkpoints are structured. If you need to scale beyond one GPU, check tune ls for a distributed KD config that adds sharding and FSDP. This single-device recipe remains the simplest place to start: lock the dtype, set the dataloader shape, then adjust batch size, kd_ratio, and gradient_accumulation_steps to fit your GPU.