09. Distillation at Scale

Chapter 9 of 18 · 15 min

KEY INSIGHT

Enterprise-scale distillation introduces coordination overhead, computational expense, and distributed training challenges that require systematic infrastructure design. Large-scale distillation trains students against teachers comprising billions of parameters. The computational cost of running the teacher for every training step can exceed the cost of training the student itself. Efficient distillation pipelines must optimize teacher inference alongside student optimization. Batch distillation amortizes teacher computation across multiple student updates. Instead of generating teacher outputs for each student step, generate teacher outputs for many student batches and cache them. Student training proceeds using cached outputs, updating without repeatedly running the teacher. ```python class CachedDistillationPipeline: """ Discretizes teacher inference to amortize computational cost. """ def __init__(self, teacher, cache_size=1000, refresh_fraction=0.1): self.teacher = teacher self.cache = {} # {(input_hash): teacher_output} self.cache_size = cache_size self.refresh_fraction = refresh_fraction self.access_counts = defaultdict(int) def generate_cache(self, dataloader): """Pre-compute teacher outputs for dataset.""" self.teacher.eval() teacher_outputs = [] with torch.no_grad(): for batch_idx, batch in enumerate(dataloader): if len(teacher_outputs) >= self.cache_size: break inputs = self.collate_inputs(batch) with torch.cuda.amp.autocast(): outputs = self.teacher(inputs) teacher_outputs.append({ 'inputs': inputs, 'logits': outputs.logits, 'hidden_states': outputs.hidden_states, 'attention': outputs.attentions, }) # Store in cache for idx, item in enumerate(teacher_outputs[:self.cache_size]): self.cache[idx] = item return self.cache def get_distillation_batch(self, batch_idx): """Retrieve cached teacher output for batch.""" if batch_idx not in self.cache: # Regenerate if cache miss self.regenerate_single(batch_idx) return self.cache[batch_idx] def regenerate_single(self, batch_idx): """Regenerate cache entry for single batch.""" inputs = self.fetch_inputs(batch_idx) with torch.no_grad(), torch.cuda.amp.autocast(): self.cache[batch_idx] = self.teacher(inputs) ``` Distributed distillation splits training across multiple devices. The teacher can reside on one device generating outputs while the student trains on another. Pipeline parallelism overlaps teacher inference with student training stages, hiding latency from teacher computation. A failure mode involves staleness in cached distillation. Training dynamics shift as the student learns—cached teacher outputs reflect the teacher's behavior before student updates. Stale outputs cause the student to chase a moving target with delayed information. Periodic cache refreshes mitigate this issue at computational cost. Multi-teacher distillation trains against multiple teacher models simultaneously. Different teachers capture complementary knowledge—a retrieval teacher and a language understanding teacher, for example. Multi-teacher approaches are more complex but can transfer broader capabilities than single-teacher approaches.

EXERCISE

Compare the training time and final student accuracy between online distillation (teacher runs every step), cached distillation (pre-computed teacher outputs), and mixed distillation (periodic teacher refresh). Identify the conditions where each approach works best.