The Training Infrastructure Behind AI-Powered Job Search: 8X Faster Multi-Teacher Distillation

At LinkedIn, we recently shared how we reimagined our search tech stack with improvements like LLM-powered query understanding, semantic retrieval, and Small Language Model ranking. Changes that power a smarter, faster, and more personalized experience that helps every member find the most relevant opportunities and connections. This blog will go deeper on the training infrastructure that makes it all possible for LinkedIn’s AI-Powered Job Search.

Delivering the most relevant job recommendations requires frequent model retraining. Currently, the speed and efficiency of our training infrastructure act as the primary constraint on how rapidly we can iterate and improve our ranking models. Behind every model improvement - from significant gains in engagement prediction and ranking quality to achieving near-teacher model accuracy from a model one-third the size - there's a training pipeline that has to get dramatically faster and more flexible. Taking our job search as an example, we went from single-node jobs that took more than two days to a distributed multi-teacher distillation pipeline that runs in under 5 hours.

In this blog, you’ll learn about the infrastructure approach behind LinkedIn's multi-teacher distillation for AI-Powered Job Search, including how we went about coordinating multiple specialized teachers and the efficient training pipeline supporting them. These improvements to our infrastructure meant rethinking the pipeline at three levels that we’ll explore further, including: distributed training that keeps pace with growing model complexity, a multi-teacher distillation system that coordinates several specialized models across heterogeneous GPU clusters, and a caching and streaming infrastructure that turns slow, all-or-nothing runs into rapid experimentation at production scale.

The problem: Scaling model training for AI-Powered Job Search

LinkedIn's ranking system for AI-Powered Job Search uses a knowledge distillation approach: large teacher models generate rich training signals, and a smaller student model learns to approximate those signals for low-latency serving. But AI-Powered Job Search isn't a simple single-teacher setup. The production pipeline requires multiple specialized teachers working in concert:

A relevance teacher that scores job-query relevance using structured evaluation prompts.
An engagement teacher that predicts member actions (views, applies, recruiter accepts, profile connections, messaging, and follows) across multiple prediction tasks.
An embedding model that generates contextual representations fed directly into the student's input layer.
A compact student Small Language Model (SLM) that distills knowledge from all of the above into a single model optimized for low-latency serving.

The reason for this multi-teacher setup is a hard serving constraint. Job search has to rank thousands of candidate jobs per query within a tight latency budget, which makes serving multi-billion-parameter teacher models directly a non-starter. Distillation lets us keep their quality while serving a compact 0.6B student at ~22,000 requests/sec/GPU. But that trade pushes cost from serving onto training, so the training pipeline becomes the bottleneck for how quickly the product can iterate.

The baseline pipeline was a chain of loosely coupled stages - fine-tune each teacher, run each teacher over the training data to generate its signals, then train the student on the combined outputs. The data lives on HDFS, the Hadoop Distributed File System, as hundreds of millions of job-query pairs, and because the models span a wide range of sizes and architectures, the stages were stitched together largely by hand. Critically, every iteration re-ran the whole chain end to end, even when only the student had changed.

Training these models takes time. The engagement teacher alone, fine-tuned on massive interaction logs, originally took more than a day per run. The distillation stage, which coordinates teacher inference alongside student training, was even slower. With each full iteration measured in days, the infrastructure, not the modeling, set the ceiling on how fast we could improve the product.

The core challenge for multi-teacher distillation

Why multiple teachers?

Knowledge distillation in its simplest form uses one teacher and one student. But the AI-Powered Job Search’s student SLM must learn from multiple specialized models simultaneously. Each teacher provides a different type of signal: the relevance teacher provides soft labels scored as a logit-based probability over the yes/no tokens, the engagement teacher provides multi-task soft targets across several member action prediction tasks, and the embedding model contributes learned representations that enrich the student's input layer. For the relevance teacher, we compute the "yes" probability by taking the yes score and normalizing it by the sum of the yes and no scores.

Just as important, the set of teachers is not fixed. Over the project's lifetime we have added and swapped teachers as the modeling team explored new signals, including introducing a dedicated click teacher, and the teachers themselves span multiple model families and sizes. From an infrastructure standpoint this rules out hardcoding assumptions about any single teacher's architecture or output format. The system has to treat teachers as pluggable components. Onboarding the click teacher, for example, first required building an internal serving-engine release with sequence-classification support before it could even produce labels.

The student trains on a Knowledge Distillation (KD) loss that combines cross-entropy on hard labels with Kullback-Leibler divergence on the teachers' soft labels. From an infrastructure perspective, this creates two distinct costs, both of which grow with the number of teachers:

The first is compute: every teacher must run a forward pass over each batch, and because the teachers are many times larger than the student, teacher inference can easily dominate step time.
The second is data movement: the student has to gather and merge every teacher's output on each step. The embedding model in particular produces large dense tensors that must be transferred to the student every step; soft labels are smaller, but as the number of teachers and prediction tasks grows, they add up too.

Accelerating teacher training

Before tackling distillation itself, we needed teacher fine-tuning to keep pace, since a slow teacher pipeline bottlenecks everything downstream. Our first step was scaling out training across multiple GPUs and nodes using PyTorch's Fully Sharded Data Parallel (FSDP) with a configurable sharding policy that could adapt to each model in our pipeline. We needed this flexibility because the pipeline includes models of varying sizes and architectures, and a single sharding strategy doesn't work well across all of them. On top of FSDP, we layered a series of complementary optimizations: Liger kernels for more efficient GPU utilization, an upgraded runtime stack on the latest CUDA and PyTorch versions to fully utilize the GPUs, and a migration to H200 GPUs with their larger memory and bandwidth. None of these techniques is novel on its own; the engineering was in composing them so the gains actually stacked rather than cancelling out.

We later migrated to FSDP2 and added Hybrid Sharded Data Parallel (HSDP) support. FSDP2 improved on the original FSDP with better memory efficiency and more flexible composability with other parallelism strategies. HSDP was particularly valuable for our multi-node setup. It applies full sharding within each node (so GPUs on the same machine split the model to save memory) while replicating across nodes (so each node holds a complete copy and only needs to sync gradients). This means inter-node communication is limited to gradient all-reduces rather than the more expensive parameter all-gathers that full cross-node sharding requires. For our workloads, this was a meaningful improvement since intra-node NVLink bandwidth is much higher than inter-node network bandwidth, and HSDP lets us take advantage of that topology.

Together, these optimizations roughly doubled teacher fine-tuning throughput (from ~25 hours to under 12 hours on 20M records), increased usable batch size by 75%, and doubled hardware FLOPs utilization (HFU), all while preserving model quality. With teacher training no longer the bottleneck, we could turn to the harder problem of coordinating multiple specialized teachers during distillation.

Online distillation: Real-time, teacher-student coordination

In online distillation, the teacher models run in inference mode alongside the student during training. Each training step, the student sends its current batch to each teacher. The teachers process the batch and return their outputs. The relevance and engagement teachers return soft label distributions over their respective prediction tasks, while the embedding model returns dense vector representations. The student then combines these signals, computing its KD loss against the teachers' soft labels and fusing the embedding outputs into its own input representation, before running backpropagation.

We started with the simplest arrangement: a single node, with all teachers co-located on the same machine as the student. Teachers could serve their outputs locally with no network transfer, which kept the system easy to reason about. But it had a low ceiling: teachers and the student competed for the same GPUs, so adding a teacher or growing the batch quickly ran out of memory, and the whole job was capped at what one machine could hold.

The single biggest lever was how we ran teacher inference. Our first implementation generated teacher outputs with an ordinary model forward pass, which left the GPUs badly underutilized. We switched to running each teacher on a high-throughput inference engine, the same one that serves our production ranking, with continuous batching and paged attention. That change alone made teacher inference roughly 2X faster and let us push 4X larger batches (or 5X more augmented samples) through the same hardware, which is what made full-scale online distillation practical in the first place. On top of it, we applied tensor parallelism to the larger teachers and prefetched the next batch's teacher inference so it overlapped with the student's own computation.

As models and data outgrew a single machine, we moved to a multi-node setup with replicated teachers. The student is sharded across nodes with data parallelism, and each teacher is replicated on every node so the student always has a local copy to query. This keeps every teacher output on-box and avoids paying the per-step network cost of shipping soft labels and embeddings between machines (the bandwidth bottleneck described earlier) at the expense of dedicating extra GPUs to teacher replicas.

As the system matured, we made the whole thing configurable. Teacher combinations, parallelism strategies, and hyperparameters became parameters rather than code changes, so the broader team could try different distillation recipes without touching core infrastructure. Together these brought online multi-teacher distillation from almost two days (~45 hours) down to roughly 10 hours. But one structural inefficiency remained. Every run re-ran all teacher inference from scratch, even when neither the teachers nor the data had changed. In practice, teams iterate on the student far more often than on the teachers or the data, so re-paying for teacher inference on every run was pure waste. This is what drove the move to offline distillation.

Offline distillation: Decoupling teacher inference from student training

The offline distillation pipeline separates the process into two stages.

In the first stage, each teacher runs inference over the training data once and writes its outputs (soft labels or embeddings) to HDFS. Each teacher's cache is keyed on model version and data fingerprint, so updating one teacher doesn't invalidate another's cached outputs.
In the second stage, the student trains entirely from these pre-computed outputs, reading them from storage rather than querying live teacher fleets. No teacher GPUs are needed during student training at all, which is where most of the per-run cost goes.

A key design decision was making the cache operate at per-shard granularity rather than all-or-nothing per teacher. When the training data is updated, only the new or modified shards trigger teacher inference, while unchanged shards reuse cached outputs. This makes repeated runs dramatically cheaper, which is exactly the common case of iterating on student architecture or loss while the teachers stay fixed.

The remaining challenge was data volume. The offline datasets are large - tens of millions of records - each now carrying the original training example plus dense embeddings and multi-task soft labels from several teachers, and they come in a mix of formats (e.g. Avro, JSONL, and pre-formatted datasets).

Previously, a training job had to stage the entire dataset to a local disk before it could start. We replaced that with a streaming data pipeline that reads directly from HDFS during training, eliminating the staging step and cutting both startup time and local-storage requirements. To produce the caches efficiently in the first place, teacher inference runs as a concurrent, high-throughput generation job rather than one teacher at a time. Profiling this path also surfaced a dataloader bug whose fix both sped up training and improved model quality.

Offline distillation does add an upfront cost. The first stage has to generate the cache, but that generation runs in under 2 hours on resources comparable to an online run. It is paid only once per teacher-and-data version, then amortized across every subsequent student iteration, which needs no teacher inference at all. With teacher inference fully decoupled and cached, end-to-end student training dropped to under 5 hours, more than 8X faster than the original ~45-hour pipeline and about half the time of the optimized online path, with no loss in model quality.

Bringing it together: A unified distillation framework

Before this, online and offline distillation shared no common machinery, so the two paths duplicated a lot of work, the same teacher setup, data plumbing, and orchestration rebuilt twice and kept in sync by hand. We consolidated them into a single platform that supports both modes under one architecture. Running the whole thing as a Ray job on a heterogeneous GPU cluster buys three things that the old pipeline couldn't:

Teachers scale independently of each other and of the student
Teachers and the trainer can run on different hardware matched to each one's needs, rather than forcing everything onto a single GPU type
The topology itself is flexible, the number of teachers, and the data- and tensor-parallel degree of each, are configuration rather than a rewrite

This last point lifts a hard limit from the online path, where teachers were co-located on the same machine as the trainer. Everything, the student plus every teacher, had to fit within a single node's GPUs (eight, on our hardware), which capped how large a teacher could be and how much data- or tensor-parallelism it could use. With teachers running as their own Ray actors on their own nodes, that ceiling is gone, a teacher can be sized and parallelized for its own needs without competing for the trainer's GPUs. Within a single run, each teacher can operate in live inference mode or cached mode depending on what's available.

The platform has four types of components working together. Inference workers run teacher models on dedicated GPU nodes using a high-throughput serving engine (the same one that powers our production ranking). Each teacher can be independently scaled with data parallelism, tensor parallelism, or distributed across multiple nodes depending on model size. When running live, these workers also write their outputs to a cache in the background, so subsequent runs with the same teacher and data can skip inference entirely.

Cache readers handle the opposite case. When a teacher's pre-computed outputs are available in shared storage, they provide them to the trainer directly, replacing the need for a live inference fleet. A collector component sits on the trainer node and merges outputs from all teachers, regardless of whether they came from live inference or cache, aggregating soft labels via averaging or learned weighting. Finally, a trainer orchestrator manages the full training loop: it checks cache availability for each teacher, routes requests to live inference or cache accordingly, fuses embedding outputs into the student input, computes the distillation loss, and runs backpropagation.

The result is that online and offline are not separate pipelines but per-teacher decisions within the same run. If one teacher's cache is warm while another has been updated, the system reads from cache for the first and runs live inference for the second. Over successive runs, the system naturally converges toward fully cached operation as more teachers stabilize, progressively reducing GPU costs without manual intervention. When data volumes grow and direct transfer becomes impractical, the caching layer absorbs the cost by letting teams pre-compute once and retrain many times.

The goal was not just to make training faster. The real value was that the modeling team could move from multi-day iteration cycles to same-day experiments, which made it possible to test ideas quickly and ship the SLM launch. The improvements compounded across the pipeline. Teacher fine-tuning dropped from about 25 hours to under 12, with 75% larger batches and roughly 2x better hardware FLOPs utilization. Online multi-teacher distillation went from about 45 hours to around 10. Offline distillation brought end-to-end student training under 5 hours, making the full student-training loop about 8x faster than where we started.

Two things mattered as much as the raw numbers:

First, every speedup preserved model quality: the distilled student reaches near-teacher ranking quality with significant engagement-prediction gains from multi-task distillation, and the offline path is quality-neutral versus online, so none of these wins were paid for in accuracy.
Second, most of the savings come from eliminating GPU waste rather than adding GPUs. Caching removes redundant teacher inference across runs (cached student training needs no teacher GPUs at all), and because online and offline are per-teacher decisions within one Ray platform, repeated runs converge toward fully cached operation, automatically shedding GPU cost as teachers stabilize.

What's next

We are now turning this platform into a standalone distillation framework that is not tied to any specific backend. The core design is simple: teacher inference workers, a student trainer, and a cache between them. Because of that, very little is specific to job search. Different serving engines and trainers can plug in behind the same interface, and the framework can support many teacher-student setups, including single-teacher, multi-teacher, multi-task, and embedding-based distillation. The goal is for teams across LinkedIn to adopt distillation without rebuilding this infrastructure, and it is already being validated beyond job search across multiple teams and use cases.

Looking further out, the same architecture, a trainer coordinating with auxiliary models served for inference, is not specific to distillation. Reinforcement-learning post-training has the same shape, with reward and reference models standing in for teachers, so unifying distillation and RL on a single framework is a direction we are actively exploring.