Optimizing Teacher-Student Partitioning for Scalable Knowledge Distillation on HPC Systems
For practitioners deploying KD on HPC clusters, this work provides a practical methodology to accelerate training by exploiting teacher-student asymmetry.
The paper addresses the inefficiency of symmetric model partitioning in Knowledge Distillation (KD) on HPC systems, achieving up to 67% higher samples-per-second by decoupling teacher and student partitioning and selecting optimal split strategies.
Knowledge Distillation (KD) enables training smaller student models under the guidance of larger teacher models, and the widely adopted TRL library implements it. Yet, TRL treats both models symmetrically, missing opportunities to exploit their pronounced asymmetry in memory footprint, and communication requirements. This paper presents an HPC-aware methodology for KD that decouples teacher and student partitioning efficiently. Our approach achieves up to 67% higher samples-per-second than TRL by avoiding unnecessary teacher-model data structures and selecting the best split strategy. We combine vertical and horizontal partitioning of models, deriving an analytical expression that identifies the existence of inflection points between splitting regimes. These results showed that exploiting teacher--student asymmetry through topology-aware parallelism notably accelerated GKD training on production HPC clusters at our company