Generative Recommender Training Efficiency

8/6/2026

This post details the engineering effort to accelerate the training infrastructure for LinkedIn's 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 SLM, and the infrastructure re-architecture to address these. Key contributions include: 1. Distributed training enhancements using PyTorch's FSDP and FSDP2 with HSDP for multi-node setups, improving teacher fine-tuning throughput and hardware FLOPs utilization. 2. A multi-teacher distillation system that treats teachers as pluggable components, using high-throughput inference engines with continuous batching and paged attention for teacher inference, and overlapping teacher inference with student computation. 3. A caching and streaming infrastructure to enable rapid experimentation. These improvements led to an 8X speedup in the multi-teacher distillation pipeline, reducing training time from over two days to under 5 hours.

5/28/2026

Introduced a C++-based fused data loader to reduce I/O bottlenecks and average training step time by 50%. Replaced standard attention kernels with FlashAttention-3 and FlexAttention, achieving up to 25% faster training for Ads GR and 2x faster for 3D mask training. Developed a custom CUDA kernel for fused metrics calculation, reducing update time from ~40ms to ~0.5ms and contributing to 22% GPU hour savings. Enabled fused optimizer operations by turning on the fused flag in Adam, reducing optimizer time by 50% and contributing to 15% GPU hour savings. Implemented fused embedding table lookups to reduce kernel launches and memory traffic.