Squeezing 70B Performance into 8B Weights

The VRAM Constraint

We needed the complex reasoning and nuance capabilities of a 70B parameter model, but our deployment environment only had the VRAM budget to serve an 8B parameter model. Standard fine-tuning of the 8B model on our dataset plateaued early.

Soft Label Distillation

We utilized Knowledge Distillation. Instead of training the 8B "Student" model on hard labels (the actual text), we passed a massive corpus of unlabelled data through the 70B "Teacher" model to generate "soft labels"—the full probability distribution across the entire vocabulary for every single token.

import torch.nn.functional as F

def distillation_loss(student_logits, teacher_logits, temperature=2.0):
    # Soften the probabilities
    soft_targets = F.softmax(teacher_logits / temperature, dim=-1)
    soft_prob = F.log_softmax(student_logits / temperature, dim=-1)
    
    # Calculate Kullback-Leibler divergence
    kl_loss = F.kl_div(soft_prob, soft_targets, reduction='batchmean')
    return kl_loss * (temperature ** 2)

By training the Student model to minimize the KL-divergence between its outputs and the Teacher's soft labels, the Student learned the underlying reasoning patterns and relationships between words, not just the hard text tokens. It achieved 92% of the Teacher's performance at 1/10th the computational cost.