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

arXiv: https://arxiv.org/abs/2606.27797 · HTML: https://mlsanity.com/paper/2606.27797/

## Key facts

- No citing papers in the map yet (data through 14 Oct 2026).
- Placed in General knowledge distillation methods, one of 12 directions in Knowledge distillation.
- Builds on 10 of its 13 references that are in the map.
- First posted on arXiv 26 Jun 2026.

## Placed in

- Knowledge distillation / [General knowledge distillation methods](https://mlsanity.com/distill/general-knowledge-distillation-methods/)

## Builds on (most cited first, 10 of 10)

- [Attention Is All You Need](https://mlsanity.com/paper/1706.03762/) (2017)
- [Distilling the Knowledge in a Neural Network](https://mlsanity.com/paper/1503.02531/) (2015)
- [Efficient Memory Management for Large Language Model Serving with PagedAttention](https://mlsanity.com/paper/2309.06180/) (2023)
- [DistilBERT, a distilled version of BERT: smaller, faster, cheaper and lighter](https://mlsanity.com/paper/1910.01108/) (2019)
- [Megatron-LM: Training Multi-Billion Parameter Language Models Using Model Parallelism](https://mlsanity.com/paper/1909.08053/) (2019)
- [ZeRO: Memory Optimizations Toward Training Trillion Parameter Models](https://mlsanity.com/paper/1910.02054/) (2019)
- [On-Policy Distillation of Language Models: Learning from Self-Generated Mistakes](https://mlsanity.com/paper/2306.13649/) (2023)
- [DeepSpeed Inference: Enabling Efficient Inference of Transformer Models at Unprecedented Scale](https://mlsanity.com/paper/2207.00032/) (2022)
- [torchdistill: A Modular, Configuration-Driven Framework for Knowledge Distillation](https://mlsanity.com/paper/2011.12913/) (2020)
- [GKD: A General Knowledge Distillation Framework for Large-scale Pre-trained Language Model](https://mlsanity.com/paper/2306.06629/) (2023)

## Abstract (arXiv)

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