Optimizing Teacher-Student Partitioning for Scalable Knowledge Distillation on HPC Systems

2026-06-26

AI Summary

This work addresses the inefficiency of conventional knowledge distillation methods in high-performance computing (HPC) systems, which employ symmetric parallel strategies for teacher and student models despite their markedly different memory and communication demands. To overcome this limitation, we propose an HPC-oriented asymmetric knowledge distillation framework that introduces, for the first time, a topology-aware heterogeneous parallelism mechanism. By combining vertical and horizontal model partitioning, our approach decouples the parallelization strategies of the teacher and student models and derives an analytical expression for the optimal partition-switching inflection point. Integrated with communication optimization and memory-reduction techniques, the proposed method achieves up to a 67% throughput improvement on real-world HPC clusters, substantially accelerating large-scale distillation training.

Abstract

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

Problem

• Knowledge Distillation
• Teacher-Student Asymmetry
• HPC Systems
• Model Partitioning
• Scalable Training

Innovation

• knowledge distillation
• teacher-student asymmetry
• HPC-aware partitioning
• model parallelism
• scalable training