The training infrastructure behind AI-Powered Job Search: 8X faster multi-teacher distillation

This post details the engineering effort behind optimizing the training infrastructure for AI-Powered Job Search, specifically focusing on multi-teacher distillation. It describes the challenges of coordinating multiple specialized teacher models (relevance, engagement, embedding) with a student Small Language Model (SLM) for low-latency serving. Key contributions include: accelerating teacher fine-tuning through PyTorch FSDP2, HSDP, Liger kernels, and H200 GPUs, which doubled throughput and doubled hardware FLOPs utilization; and enabling online distillation by switching to high-throughput inference engines with continuous batching and paged attention for teacher inference, which made teacher inference 2X faster and allowed for 4X larger batches. The post also discusses tensor parallelism for larger teachers and prefetching to overlap computation.