AI distillation reduces its cost with a new technique

How do you get a massive artificial intelligence model to teach what it knows without needing a room full of GPUs? Research from Multiverse Computing proposes a practical answer: save only the large model’s most important predictions and calculate the training function in parts, without building enormous matrices in memory.

The result is a considerable reduction in VRAM usage, especially when working with long contexts. In a demonstration with a 20-billion-parameter GPT-OSS model, the system went from needing four GPU nodes to running on one, while the time per step dropped from 57 to 12.23 seconds.

Why model distillation is so expensive

Knowledge distillation involves training a smaller model, called the student, to imitate the responses of another, larger model, known as the teacher. The idea is simple: instead of always running the giant model, you transfer some of its capabilities to a cheaper, faster version.

But the traditional process has a problem. During every training step, both the teacher and student models must remain loaded. In addition, the teacher produces a probability distribution over its entire vocabulary for every input token.

Does that sound excessive? Consider this case: the gpt-oss-120b model uses a vocabulary of 201,088 tokens. With sequences of 32,768 tokens and a batch of four examples, a single probability matrix from the teacher can take up nearly 50 GB in bfloat16.

When you add the weights, activations, gradients, and optimizer states, one iteration can reach approximately 250 GB of VRAM. That exceeds the capacity of an H200 GPU with 141 GB and forces the training to be distributed across multiple devices.

The challenge is not only loading the large model. It is also storing all the predictions it produces for every position in a long sequence.

The proposal: separate the teacher from training

The first modification consists of replacing online distillation with an offline strategy known as offline distillation.

Instead of keeping the teacher active throughout training, the system runs it only once. For each token, it saves only its top-k predictions—that is, the options with the highest probability. In the main experiment, the 100 most likely options are stored for each position.

The student is then trained using that prediction file. The teacher no longer needs to occupy memory or recalculate its results, and the same cache can be reused across multiple experiments.

Saving only the 100 best options may seem like a major simplification compared with a vocabulary of more than 200,000 tokens. However, those predictions contain most of the useful information the student needs to learn the teacher’s behavior.

A KL function that does not build the entire matrix

The second modification affects the loss function based on Kullback–Leibler divergence, or KL loss. This function measures how far the student’s prediction distribution differs from the teacher’s.

The traditional method creates a complete matrix with one dimension for every token in the vocabulary and another for every position in the sequence. With contexts of 32K, 64K, or even 256K tokens, that matrix grows quickly and can cause an unmanageable memory spike.

The team proposes an implementation that processes the sequence in chunks. Instead of creating and keeping the entire matrix, it calculates one part, adds its result to the accumulated loss, and discards that chunk before continuing.

Three ways to calculate the loss

The paper compares three alternatives for offline distillation:

• Dense KL: reconstructs a complete teacher distribution from the top-100 and compares it with the student’s full probabilities. It is the closest reference to the traditional method, but consumes a lot of memory.
• Forward chunked KL: keeps the teacher’s predictions in sparse format and processes the sequence in parts. It reduces consumption, although it still keeps the student’s complete logits in memory.
• Fused chunked KL: integrates the student’s output projection directly into the loss calculation. It never materializes all of its logits at the same time and recalculates each chunk during backpropagation to avoid storing it.

This last version performs part of the work twice—once during the forward pass and again during backpropagation. In exchange, it dramatically reduces the memory peak as the context length increases.

Less memory without losing training quality

In a comparison conducted with an H200 GPU, a Llama 3.1 8B Instruct model as the teacher, and a 3.2-billion-parameter Llama student, all four configurations reached virtually identical loss curves with an 8K-token context.

With 8K contexts, the fused variant is not the fastest. Its advantage appears when sequences become much longer.

In an isolated test with an output projection network, consumption at 32K tokens dropped from 85.2 GiB with dense loss to 5.45 GiB with the fully chunked version. That represents a 15.6-fold reduction. In addition, the dense method stopped working beyond 64K tokens because it ran out of memory.

At 256K tokens, the fused version used 11.6 GiB, compared with 134.2 GiB for the next most efficient alternative. At that length, it was also approximately 3.3 times faster per iteration.

The difference is noticeable in large models

The most concrete impact appeared when distilling a 20-billion-parameter GPT-OSS model with a context of 32,768 tokens. The reduction in memory made it possible to go from four GPU nodes to just one.

The time per step fell from 57 to 12.23 seconds, nearly five times faster. At the same time, performance per GPU increased from 74.2 to 345.7 TFLOP/s.

These numbers matter because distillation is rarely a one-time experiment. Teams need to test different data, loss functions, context lengths, and training configurations. If every attempt requires hundreds of gigabytes of VRAM, research becomes slow and expensive.

With a more efficient implementation, compressing models can become an iterative process accessible to more teams—not just laboratories with large clusters.

Small models with many of the same capabilities

In the recovery experiment, a student with approximately 3.2 billion parameters was distilled from a Llama 3.1 8B Instruct model. The smaller model retained most of the teacher’s accuracy on BoolQ and HellaSwag.

On MMLU, it remained about nine points behind the original model, despite having less than half as many parameters. This does not mean compression is free: the student loses some capability, especially on complex tasks. But it does show that a significant reduction can preserve useful performance.

The technique may also help with what the team calls healing, or the recovery of capabilities after compressing a model. What is the goal? To obtain smaller systems that are cheaper to run and able to preserve abilities that would normally be lost during reduction.

Open source code to try the technique

Multiverse Computing published the implementation of fused chunked KL loss on GitHub: Full-Chunked-KL-Loss.

The research also analyzes other factors, such as the loss function used and how sequences are packed during training. For those working with large models, the main contribution is not a new architecture, but a smarter way to manage memory.

Generative AI often looks like a race to build increasingly large models. However, making them useful in the real world also requires learning how to reduce them. If a technique makes it possible to transfer capabilities to smaller models using fewer GPUs and longer contexts, innovation becomes less dependent on the size of the laboratory developing it.

Original source

https://huggingface.co/blog/MultiverseComputingCAI/efficient-knowledge-distillation