KDFlow: A User-Friendly and Efficient Knowledge Distillation Framework for Large Language Models

# A User-Friendly and Efficient Knowledge Distillation Framework for Large Language Models
Source: [https://arxiv.org/html/2603.01875](https://arxiv.org/html/2603.01875)
Songming Zhang1,2,3, Xue Zhang1,2, Tong Zhang3, Bojie Hu3, Yufeng Chen1,2, andJinan Xu1,2 1School of Computer Science and Technology, Beijing Jiaotong University, Beijing, China 2Key Laboratory of Big Data & Artificial Intelligence in Transportation, \(Beijing Jiaotong University\), Ministry of Education 3WeChat, Tencent Inc, China \{smzhang22,zhang\_xue,chenyf,jaxu\}@bjtu\.edu\.cn

###### Abstract

Knowledge distillation \(KD\) is essential for model compression and post\-training in large language models \(LLMs\), yet existing frameworks often use homogeneous training backends for both teacher and student models, leading to suboptimal efficiency\. In this paper, we proposeKDFlow, a novel framework for LLM distillation that features a decoupled architecture and employs SGLang for teacher inference\. By bridging the training efficiency of FSDP2 and the inference efficiency of SGLang, KDFlow leverages the complementary strengths of both systems within a unified distillation pipeline\. Moreover, our framework achieves efficientfull\-vocabulary distillationby transmitting the teacher’s hidden states using zero\-copy data transfer and recomputing the logits on the student side\. Furthermore, our framework supports both off\-policy and on\-policy distillation and incorporates cross\-tokenizer algorithms through highly extensible and user\-friendly APIs\. Experiments show that KDFlow achieves a1\.44×\\timesto 6\.36×\\timesspeedup over existing KD frameworks, enabling researchers to rapidly prototype and scale LLM distillation with minimal engineering overhead\. The code and documentation are publicly available\.

KDFlow: A User\-Friendly and Efficient Knowledge Distillation Framework for Large Language Models

Songming Zhang1,2,3††thanks:Corresponding author\.††thanks:Work done during internship at Tencent\., Xue Zhang1,2, Tong Zhang3, Bojie Hu3,Yufeng Chen1,2, andJinan Xu1,21School of Computer Science and Technology, Beijing Jiaotong University, Beijing, China2Key Laboratory of Big Data & Artificial Intelligence in Transportation,\(Beijing Jiaotong University\), Ministry of Education3WeChat, Tencent Inc, China\{smzhang22,zhang\_xue,chenyf,jaxu\}@bjtu\.edu\.cn

## 1Introduction

Large Language Models \(LLMs\) have demonstrated remarkable capabilities across diverse tasks, yet their substantial parameter scales pose significant challenges for deployment in resource\-constrained environments\. Knowledge Distillation \(KD\) offers an effective approach to addressing this challenge by transferring knowledge from a large teacher model to a compact student modelHintonet al\.\([2015](https://arxiv.org/html/2603.01875#bib.bib1)\)\.

![Refer to caption](https://arxiv.org/html/2603.01875v3/figures/teacher_fwd_occ.png)Figure 1:Training time per step and the proportion of teacher forward time under different distillation configurations\. The teacher’s MoE architecture poses challenges for FSDP, while being well\-supported by SGLang\.Despite extensive research and application of KD, the infrastructure for LLM distillation remains suboptimal\. In a typical KD process, the teacher and student models play distinct roles: the teacher performs only forward passes \(inference\), while the student requires both forward and backward passes \(training\)\. However, most existing frameworks, such as TRLvon Werraet al\.\([2020](https://arxiv.org/html/2603.01875#bib.bib3)\)or MS\-SWIFTZhaoet al\.\([2025](https://arxiv.org/html/2603.01875#bib.bib2)\), run both models with a unified training backend\. This creates a structural mismatch, as a training engine designed for gradient computation and optimizer state management is ill\-suited to the teacher’s inference\-heavy workload\. As a result, the throughput of LLM distillation is often bounded by inefficient teacher execution, especially for Mixture\-of\-Experts \(MoE\) teachers \(see Figure[1](https://arxiv.org/html/2603.01875#S1.F1)\)\.

FeaturesTRLMS\-SWIFTEasyDistillROLLSlimeverlKDFlow \(Ours\)KD\-Native Design✓✓Decoupled Backends✓✓✓✓Off\-Policy Distillation✓✓✓✓✓On\-Policy Distillation✓✓✓✓✓✓Self Distillation✓✓✓Cross\-Tokenizer Distillation✓✓Multi\-Teacher Distillation✓✓✓✓✓Logits/Logprobsfull vocabfull vocabfull vocabtop\-kkvocabsampled tokenssampled tokenstop\-kkvocabsampled tokensfull vocabDivergence MetricsFKL,RKL,JSDFKL,RKL,JSDFKL,RKLoff\-policy: FKL, RKL, JSD,Skewed F\(R\)KL, AKL,on\-policy: RKLRKLRKLFKL, RKL, JSD,Skewed F\(R\)KL,AKL, TVD

Table 1:Comparisons between KDFlow and existing frameworks\.To address this issue, we presentKDFlow, a high\-performance framework designed specifically for LLM distillation\. KDFlow decouples the teacher and student backends by assigning the student to PyTorch FSDP2 and the teacher to the high\-throughput inference engine SGLangZhenget al\.\([2024](https://arxiv.org/html/2603.01875#bib.bib4)\)\. A key challenge in this architecture is transferring full teacher logits from SGLang processes to FSDP processes: directly transferring logits is infeasible due to their huge size111For 128 sequences with length 4096, full BF16 logits from Qwen3 models occupy128×4096×151936×2​bytes≈160​GB128\\times 4096\\times 151936\\times 2\\text\{bytes\}\\approx 160\\text\{GB\}of memory\., while transferring only top\-k logits breaks the mathematical equivalence of the loss function\. KDFlow solves this by collecting compact teacher hidden states from SGLang and recomputing the full logit distribution on the student side, reducing communication overhead while preserving standard KD\. Compared to prevalent frameworks like TRL and MS\-SWIFT, KDFlow delivers a1\.44×\\timesto 6\.36×\\timestraining speedup in off\-policy distillation, with larger gains for Mixture\-of\-Experts \(MoE\) teachers\.

Despite the decoupled design, KDFlow abstracts away distributed communication and integrates with standard Hugging Face model formats, allowing users to initiate KD with only a few lines of configuration\.

Overall, the contributions of KDFlow include:

- •Efficient Architecture: KDFlow decouples teacher inference from student training by serving the teacher with SGLang and transferring compact hidden states across processes, achieving a 1\.44×\\timesto 6\.36×\\timesspeedup over unified\-backend frameworks such as TRL and MS\-SWIFT\.
- •Comprehensiveness: KDFlow supports comprehensive KD features \(e\.g\., off/on\-policy, cross\-tokenizer, and multi\-teacher distillation\) and provides multiple built\-in divergence metrics and algorithms, making it an out\-of\-the\-box toolkit for LLM distillation\.
- •User\-Friendly Design: KDFlow is a lightweight framework based on FSDP2, and decouples the algorithms from the whole distillation pipeline\.

## 2Related Work

### 2\.1Knowledge Distillation for LLMs

Knowledge Distillation \(KD\) was first proposed byHintonet al\.\([2015](https://arxiv.org/html/2603.01875#bib.bib1)\)to compress large models into smaller ones\. For LLMs, KD is commonly divided into black\-boxKim and Rush \([2016](https://arxiv.org/html/2603.01875#bib.bib5)\)and white\-box distillationZhanget al\.\([2023](https://arxiv.org/html/2603.01875#bib.bib32)\); Guet al\.\([2023](https://arxiv.org/html/2603.01875#bib.bib6)\); Agarwalet al\.\([2024](https://arxiv.org/html/2603.01875#bib.bib7)\); Koet al\.\([2024](https://arxiv.org/html/2603.01875#bib.bib30)\); Wuet al\.\([2025](https://arxiv.org/html/2603.01875#bib.bib31)\)\. White\-box KD aligns the student and teacher output distributions with divergence metrics, providing richer supervision and often better performance\. The paradigm has also expanded from off\-policy distillation on static datasets to on\-policy distillation, where the student learns from its own generated dataGuet al\.\([2023](https://arxiv.org/html/2603.01875#bib.bib6)\); Agarwalet al\.\([2024](https://arxiv.org/html/2603.01875#bib.bib7)\); Xiaoet al\.\([2026](https://arxiv.org/html/2603.01875#bib.bib8)\)\. Recent work further studies cross\-tokenizer KD to handle vocabulary mismatches between teacher and student modelsWanet al\.\([2024](https://arxiv.org/html/2603.01875#bib.bib15)\); Boizardet al\.\([2024](https://arxiv.org/html/2603.01875#bib.bib12)\); Zhanget al\.\([2024](https://arxiv.org/html/2603.01875#bib.bib18),[2025](https://arxiv.org/html/2603.01875#bib.bib17)\); Cuiet al\.\([2025](https://arxiv.org/html/2603.01875#bib.bib16)\); Chenet al\.\([2025](https://arxiv.org/html/2603.01875#bib.bib13)\); Minixhoferet al\.\([2025](https://arxiv.org/html/2603.01875#bib.bib14)\)\. Despite these algorithmic advances, a flexible and efficient framework for KD research remains lacking, which motivates this work\.

![Refer to caption](https://arxiv.org/html/2603.01875v3/figures/architecture.png)Figure 2:Overview of KDFlow\. The framework is built on RayMoritzet al\.\([2018](https://arxiv.org/html/2603.01875#bib.bib27)\)and decouples the distillation pipeline by allocating the teacher model to SGLang and the student model to FSDP2\. Solid and dashed arrows illustrate the data flow for off\-policy and on\-policy distillation, respectively\. Notably, KDFlow transfers compact hidden states from the teacher rather than full logits to reduce communication overhead\.
### 2\.2Existing Frameworks

Existing LLM training and distillation frameworks provide useful support but still have important limitations\. TRLvon Werraet al\.\([2020](https://arxiv.org/html/2603.01875#bib.bib3)\)and EasyDistillWanget al\.\([2025a](https://arxiv.org/html/2603.01875#bib.bib9)\)support KD with homogeneous training engines, which underutilize hardware during teacher forward passes\. SlimeZhuet al\.\([2025](https://arxiv.org/html/2603.01875#bib.bib11)\)and verlShenget al\.\([2025](https://arxiv.org/html/2603.01875#bib.bib10)\)decouple inference and training with engines like vLLM/SGLang, but they are not KD\-specialized and only support on\-policy distillation with incomplete logit information\. MS\-SWIFTZhaoet al\.\([2025](https://arxiv.org/html/2603.01875#bib.bib2)\)supports both off\-policy and on\-policy distillation, but its implementation is relatively heavy\. By contrast,KDFlowis a lightweight and efficient KD framework covering full KD scenarios, as detailed in Table[1](https://arxiv.org/html/2603.01875#S1.T1)\.

## 3The Design of KDFlow

The overall framework of KDFlow is presented in Figure[2](https://arxiv.org/html/2603.01875#S2.F2)\. To address the inefficiency of using a single homogeneous engine for both inference and training, we adopt a top\-down decoupled design\. In this section, we introduce the system architecture, the core communication mechanism, the distillation workflows, and the algorithm abstractions\.

### 3\.1System Architecture

The KD process typically involves multiple stages: teacher inference, student training, and student rollout, which require different backends during training\. Therefore, KDFlow is built upon Ray to manage distributed processes efficiently\. As shown in the middle part of Figure[2](https://arxiv.org/html/2603.01875#S2.F2), the architecture consists of a single controller \(Trainer\) and three functionally independent actor groups:

- •Trainer \(Single Controller\): Trainer is the central coordinator that manages the dataset, controls the training loop, and organizes the data flow among different actor groups\. KDFlow supports both OffPolicyKDTrainer and OnPolicyKDTrainer\.
- •RolloutActorGroup: RolloutActorGroup is used for the rollout process of the student model during on\-policy distillation\. Following frontier RL frameworks like SlimeZhuet al\.\([2025](https://arxiv.org/html/2603.01875#bib.bib11)\), it uses SGLang Router to connect multiple SGLang HTTP servers for load\-balanced inference\. Moreover, KDFlow uses a colocated mode and updates model weights in SGLang via CUDA Interprocess Communication \(IPC\)\.
- •TeacherActorGroup: TeacherActorGroup manages multiple TeacherRayActor instances that execute the forward pass of the teacher\. Specifically, we initialize an SGLang engine222The SGLang engine is more compatible with NumPy ndarray objects than the SGLang HTTP server\.in each TeacherRayActor to obtain the teacher’s hidden states\. This allows KDFlow to leverage SGLang’s high\-throughput and flexible parallel strategies for teacher\-model inference\.
- •StudentActorGroup: StudentActorGroup is deployed with PyTorch FSDP2\. It handles the standard training processes for the student model, including forward pass, backward pass, and optimizer state management\. Then it returns training status \(e\.g\., the loss values\) for logging\.

This decoupled architecture ensures that the inference\-heavy teacher model and the training\-heavy student model run on their respective optimized backends, significantly improving hardware utilization\.

![Refer to caption](https://arxiv.org/html/2603.01875v3/figures/cost.png)Figure 3:Comparison of different decoupled distillation approaches\.
### 3\.2Efficient Communication via Hidden States

In the decoupled architecture, the teacher and student typically run in separate processes, making the transfer of teacher knowledge a key bottleneck\. Directly transmitting full logit distributions becomes prohibitive due to their large data volume, while transferring only top\-k logits reduces bandwidth at the cost of breaking the mathematical equivalence of distillation and degrading performance\. To address this issue, KDFlow adopts hidden\-state transfer and logit recomputation, as shown in Figure[3](https://arxiv.org/html/2603.01875#S3.F3)\. Instead of sending full logits, the TeacherActorGroup outputs only the teacher’s final hidden states, whose dimension \(e\.g\., 4096\) is much smaller than the vocabulary size \(e\.g\., 151936\)\. KDFlow further uses shared memory and the Ray shared object mechanism for zero\-copy transfer across processes\. After receiving the teacher’s hidden states, each StudentRayActor locally recomputes the full logit distributions using the teacher’s language model head\. This design substantially reduces communication volume while preserving the mathematical equivalence to standard logit\-based KD\.

### 3\.3Distillation Workflows

Governed by the Single Controller, KDFlow seamlessly supports two main distillation workflows, illustrating the flexible data routing among the decoupled actors:

Off\-Policy Distillation\(Solid lines in Figure[2](https://arxiv.org/html/2603.01875#S2.F2)\): The student learns from a static dataset\. The Trainer sends the prompts and the corresponding responses directly to the TeacherActorGroup to obtain the teacher’s hidden states\. These hidden states, along with the inputs, are then passed to the StudentActorGroup to compute the distillation loss and update the student’s weights\.

On\-Policy Distillation\(Dashed lines in Figure[2](https://arxiv.org/html/2603.01875#S2.F2)\): The student learns from data generated by the student itself\. First, the Trainer sends prompts to the RolloutActorGroup to generate responses\. Next, these prompt\-response pairs are sent to the TeacherActorGroup to obtain the teacher’s hidden states\. Then, the data and hidden states flow into the StudentActorGroup for gradient updates\. Finally, the updated weights of the student model are synchronized back to the RolloutActorGroup to ensure the next generation step uses the latest policy\.

### 3\.4Comprehensive Abstractions and Algorithms

To provide a user\-friendly and out\-of\-the\-box toolkit, KDFlow strictly separates the underlying system pipeline from the distillation algorithms\. As shown at the bottom of Figure[2](https://arxiv.org/html/2603.01875#S2.F2), KDFlow provides built\-in support for various KD algorithms and divergence metrics, including Forward KL \(FKL\), Reverse KL \(RKL\), Jensen\-Shannon Divergence \(JSD\), and Total Variation Distance \(TVD\)\. Users can also easily implement their custom distillation losses or algorithms with minimal code, eliminating the need to understand the complex distributed communication logic\. Furthermore, KDFlow natively supports cross\-tokenizer distillation\. When the teacher and student models have different vocabularies, directly aligning their full logit distributions is impossible\. Therefore, KDFlow implements the DSKDv2Zhanget al\.\([2025](https://arxiv.org/html/2603.01875#bib.bib17)\)algorithm for cross\-tokenizer knowledge distillation\.

With this decoupled design and efficient communication mechanism, KDFlow significantly improves the overall training throughput, which we will demonstrate in the following experiments\.

![Refer to caption](https://arxiv.org/html/2603.01875v3/figures/loss_curves.png)Figure 4:Loss curves of KDFlow and the pure FSDP implementation when distilling Qwen3\-30B\-A3B to Qwen3\-4B\.ModelFrameworkAlpacaEval 2\.0LC\-Win Rate \(%\)Win Rate \(%\)Original StudentQwen3\-1\.7B26\.0921\.99Distilled StudentMS\-SWIFT28\.4027\.86KDFlow \(Pure FSDP\)28\.1828\.20KDFlow28\.2328\.32KDFlow \(On\-Policy\)28\.2929\.95TeacherQwen3\-30B\-A3B68\.6566\.52

Table 2:Student model performance on AlpacaEval 2\.0 after distillation with different frameworks\.FrameworksTrainingBackendStudent: Qwen3\-4BStudent: Qwen3\-1\.7BQwen3\-14BQwen3\-32BQwen3\-30B\-A3BQwen3\-14BQwen3\-32BQwen3\-30B\-A3BTRLZeRO\-321\.3s/it31\.5s/it\-13\.3s/it23\.4s/it\-MS\-SWIFTZeRO\-316\.6s/it24\.8s/it43\.2s/it11\.5s/it20\.1s/it36\.9s/itROLLFSDP238\.4s/it56\.9s/it67\.9s/it26\.8s/it45\.6s/it53\.8s/itKDFlow \(BF16 Teacher\)FSDP212\.3s/it15\.7s/it11\.3s/it7\.6s/it10\.9s/it5\.9s/itKDFlow \(FP8 Teacher\)FSDP211\.5s/it13\.5s/it11\.1s/it6\.7s/it8\.7s/it5\.8s/itSpeedup\-1\.44×\\times1\.84×\\times3\.89×\\times1\.72×\\times2\.31×\\times6\.36×\\times

Table 3:Off\-policy training efficiency comparison \(seconds per iteration\) across different distillation frameworks\.Speedupis calculated using KDFlow \(FP8\) against the best\-performing baseline \(i\.e\., MS\-SWIFT\)\. All frameworks use identical training settings: global batch size=128, gradient accumulation=8, and max length=4096\.

## 4Experiments

In this section, we comprehensively evaluate the KDFlow framework from the following three perspectives: \(1\)Loss Correctness:We validate the loss correctness of KDFlow by recording and analyzing its loss curves against those of the standard FSDP baseline\. \(2\)KD Performance:We compare the knowledge distillation performance of KDFlow with that of baseline frameworks on a representative downstream task\. \(3\)Training Efficiency:We test the training speed of KDFlow and existing frameworks across multiple teacher\-student setups to show the efficiency of KDFlow\.

### 4\.1Experimental Setup

We evaluate KDFlow on instruction\-following distillation using the Qwen3 model family, covering multiple student sizes and both dense and MoE teacher architectures\. All frameworks are compared under identical training settings on the same 8\-GPU server, and we report loss curves, AlpacaEval 2\.0 performance, and training speed\. Detailed model, dataset, baseline, hardware, and training configurations are provided in Appendix[A](https://arxiv.org/html/2603.01875#A1)\.

### 4\.2Loss Curve Validation

The core technical innovation of KDFlow is to leverage SGLang as the teacher backend and collect the teacher’s hidden states to recompute full logits on the student side\. However, recent studies on RL training have shown that high\-throughput inference engines may introduce numerical discrepancies during model execution, particularly for MoE modelsYaoet al\.\([2025](https://arxiv.org/html/2603.01875#bib.bib19)\); Qiet al\.\([2025](https://arxiv.org/html/2603.01875#bib.bib20)\); Maet al\.\([2025](https://arxiv.org/html/2603.01875#bib.bib21)\); Zhenget al\.\([2025](https://arxiv.org/html/2603.01875#bib.bib22)\)\. This motivates us to investigate whether such discrepancies affect the distillation behavior of KDFlow\. Specifically, we compare the loss curves of KDFlow against the one of the standard KD implementation, where both teacher and student are in the same FSDP process\. As shown in Figure[4](https://arxiv.org/html/2603.01875#S3.F4), the loss curves of KDFlow align well with the baseline throughout the entire training process\. Moreover, the loss curves almost overlap for FP8 and BF16 teacher inference, which suggests that FP8 teacher inference is an efficient alternative for LLM distillation\.

### 4\.3Downstream Task Performance

Beyond training loss, we evaluate the effectiveness of the student models trained via KDFlow on downstream tasks to ensure our decoupled architecture does not compromise distillation performance\. We perform off\-policy distillation from Qwen3\-30B\-A3B to Qwen3\-1\.7B using the Forward KL \(FKL\) divergence across all frameworks\. Table[2](https://arxiv.org/html/2603.01875#S3.T2)presents the performance evaluation on the AlpacaEval 2\.0 benchmark\. The student model distilled via KDFlow achieves an LC\-Win Rate of 28\.23% and a Win Rate of 28\.32%, demonstrating significant improvement over the un\-distilled Qwen3\-1\.7B baseline\. The on\-policy variant further improves the Win Rate to 29\.95%, showing the flexibility of KDFlow across different distillation workflows\. Moreover, when compared to the students distilled using the MS\-SWIFT and pure FSDP baselines, the performance differences remain within a negligible range\. We also report the teacher model performance as a reference, where Qwen3\-30B\-A3B achieves an LC\-Win Rate of 68\.65% and a Win Rate of 66\.52%\. These results indicate that KDFlow safely optimizes the system execution pipeline and reduces communication overhead while strictly preserving the distillation quality\. KDFlow acts as a transparent, high\-performance infrastructure that faithfully executes KD algorithms without sacrificing downstream performance\.

### 4\.4Training Speed and Efficiency

We compare the training speed of KDFlow with existing LLM distillation frameworks in the off\-policy setting, where the rollout cost in on\-policy distillation is excluded to better isolate the efficiency of teacher\-student training\. Specifically, we measure the average training time per step across different teacher\-student setups\. As shown in Table[3](https://arxiv.org/html/2603.01875#S3.T3), KDFlow is consistently the fastest framework among all FSDP2\- and DeepSpeed ZeRO\-3\-based candidates, achieving a 1\.44×\\timesto 6\.36×\\timesspeedup\. The advantage is especially pronounced for the MoE teacher Qwen3\-30B\-A3B, where standard training engines suffer from inefficient routing and expert management\. By serving the teacher with SGLang and transferring compact hidden states instead of full logits, KDFlow reduces the distillation time to 5\.8s/it\.

### 4\.5Cross\-Tokenizer Distillation

We further verify KDFlow in a cross\-tokenizer setting, where the teacher and student use different tokenizers and cannot be directly aligned through vanilla full\-vocabulary KL divergence\. Specifically, we distill knowledge from Qwen3\-14B into Llama3\.2\-3B and find that KDFlow’s DSKD\-based objective further improves over sequence\-level KD\. As shown in Table[4](https://arxiv.org/html/2603.01875#S4.T4), DSKD \+ KL improves the LC\-Win Rate from 29\.72% to 32\.88% and the Win Rate from 39\.32% to 42\.54% compared with sequence\-level KD\.

ModelSettingAlpacaEval 2\.0LC\-Win Rate \(%\)Win Rate \(%\)Original StudentLlama3\.2\-3B15\.2618\.45Distilled StudentSeqKD29\.7239\.32DSKD \+ KL32\.8842\.54TeacherQwen3\-14B73\.3172\.51

Table 4:Cross\-tokenizer distillation performance on AlpacaEval 2\.0\.

## 5Conclusion

In this paper, we introduceKDFlow, an efficient and user\-friendly KD framework for LLMs\. KDFlow decouples the architecture by deploying the teacher model on a high\-throughput inference engine \(SGLang\) and the student on a dedicated training backend \(FSDP2\)\. To eliminate the communication bottleneck caused by this decoupling, we propose a hidden\-state transfer and logit recomputation mechanism, ensuring strict theoretical equivalence to standard full\-vocabulary KD\. Experiments demonstrate that KDFlow achieves a 1\.44×\\timesto 6\.36×\\timestraining speedup over state\-of\-the\-art baselines without compromising downstream task performance\. By natively supporting off\-policy, on\-policy, and cross\-tokenizer distillation, KDFlow serves as a comprehensive infrastructure to accelerate future LLM compression and post\-training research\.

## Limitations

While KDFlow significantly improves the efficiency and flexibility of LLM distillation, it has certain limitations\. First, the current student training backend of KDFlow is built entirely upon PyTorch FSDP2, which still struggles to match the training efficiency and scalability of Megatron\-LMShoeybiet al\.\([2019](https://arxiv.org/html/2603.01875#bib.bib26)\)that supports complex 3D parallelism\. Second, as a research\-oriented framework, KDFlow currently lacks some industrial\-grade optimizations, such as asynchronous training, which are crucial for training on clusters with thousands of GPUs\. However, KDFlow explicitly prioritizes the needs of the research community: user\-friendliness, high flexibility, and rapid prototyping\. By abstracting away complex distributed communication logic, KDFlow allows researchers to easily implement and test novel off\-policy, on\-policy, or cross\-tokenizer distillation algorithms with minimal engineering overhead\. Integrating Megatron\-LM and further optimizations remain important directions for our future work\.

## References

- R\. Agarwal, N\. Vieillard, Y\. Zhou, P\. Stanczyk, S\. R\. Garea, M\. Geist, and O\. Bachem \(2024\)On\-policy distillation of language models: learning from self\-generated mistakes\.InThe twelfth international conference on learning representations,Cited by:[§2\.1](https://arxiv.org/html/2603.01875#S2.SS1.p1.1)\.
- N\. Boizard, K\. E\. Haddad, C\. Hudelot, and P\. Colombo \(2024\)Towards cross\-tokenizer distillation: the universal logit distillation loss for llms\.arXiv preprint arXiv:2402\.12030\.Cited by:[§2\.1](https://arxiv.org/html/2603.01875#S2.SS1.p1.1)\.
- Y\. Chen, Y\. Liu, F\. Meng, Y\. Chen, J\. Xu, and J\. Zhou \(2025\)Enhancing cross\-tokenizer knowledge distillation with contextual dynamical mapping\.InFindings of the Association for Computational Linguistics: ACL 2025,pp\. 8005–8018\.Cited by:[§2\.1](https://arxiv.org/html/2603.01875#S2.SS1.p1.1)\.
- X\. Cui, M\. Zhu, Y\. Qin, L\. Xie, W\. Zhou, and H\. Li \(2025\)Multi\-level optimal transport for universal cross\-tokenizer knowledge distillation on language models\.InProceedings of the AAAI Conference on Artificial Intelligence,Vol\.39,pp\. 23724–23732\.Cited by:[§2\.1](https://arxiv.org/html/2603.01875#S2.SS1.p1.1)\.
- Y\. Dubois, B\. Galambosi, P\. Liang, and T\. B\. Hashimoto \(2024\)Length\-controlled alpacaeval: a simple way to debias automatic evaluators\.arXiv preprint arXiv:2404\.04475\.Cited by:[Appendix A](https://arxiv.org/html/2603.01875#A1.SS0.SSS0.Px1.p1.1)\.
- Y\. Gu, L\. Dong, F\. Wei, and M\. Huang \(2023\)Minillm: knowledge distillation of large language models\.arXiv preprint arXiv:2306\.08543\.Cited by:[§2\.1](https://arxiv.org/html/2603.01875#S2.SS1.p1.1)\.
- G\. Hinton, O\. Vinyals, and J\. Dean \(2015\)Distilling the knowledge in a neural network\.arXiv preprint arXiv:1503\.02531\.Cited by:[§1](https://arxiv.org/html/2603.01875#S1.p1.1),[§2\.1](https://arxiv.org/html/2603.01875#S2.SS1.p1.1)\.
- Y\. Kim and A\. M\. Rush \(2016\)Sequence\-level knowledge distillation\.InProceedings of the 2016 conference on empirical methods in natural language processing,pp\. 1317–1327\.Cited by:[§2\.1](https://arxiv.org/html/2603.01875#S2.SS1.p1.1)\.
- J\. Ko, S\. Kim, T\. Chen, and S\. Y\. Yun \(2024\)DISTILLM: towards streamlined distillation for large language models\.Proceedings of Machine Learning Research235,pp\. 24872–24895\.Cited by:[§2\.1](https://arxiv.org/html/2603.01875#S2.SS1.p1.1)\.
- W\. Ma, H\. Zhang, L\. Zhao, Y\. Song, Y\. Wang, Z\. Sui, and F\. Luo \(2025\)Stabilizing moe reinforcement learning by aligning training and inference routers\.arXiv preprint arXiv:2510\.11370\.Cited by:[§4\.2](https://arxiv.org/html/2603.01875#S4.SS2.p1.1)\.
- B\. Minixhofer, I\. Vulić, and E\. M\. Ponti \(2025\)Cross\-tokenizer distillation via approximate likelihood matching\.Cited by:[§2\.1](https://arxiv.org/html/2603.01875#S2.SS1.p1.1)\.
- P\. Moritz, R\. Nishihara, S\. Wang, A\. Tumanov, R\. Liaw, E\. Liang, M\. Elibol, Z\. Yang, W\. Paul, M\. I\. Jordan,et al\.\(2018\)Ray: a distributed framework for emerging\{\\\{ai\}\\\}applications\.In13th USENIX symposium on operating systems design and implementation \(OSDI 18\),pp\. 561–577\.Cited by:[Figure 2](https://arxiv.org/html/2603.01875#S2.F2)\.
- P\. Qi, Z\. Liu, X\. Zhou, T\. Pang, C\. Du, W\. S\. Lee, and M\. Lin \(2025\)Defeating the training\-inference mismatch via fp16\.arXiv preprint arXiv:2510\.26788\.Cited by:[§4\.2](https://arxiv.org/html/2603.01875#S4.SS2.p1.1)\.
- G\. Sheng, C\. Zhang, Z\. Ye, X\. Wu, W\. Zhang, R\. Zhang, Y\. Peng, H\. Lin, and C\. Wu \(2025\)Hybridflow: a flexible and efficient rlhf framework\.InProceedings of the Twentieth European Conference on Computer Systems,pp\. 1279–1297\.Cited by:[§2\.2](https://arxiv.org/html/2603.01875#S2.SS2.p1.1)\.
- M\. Shoeybi, M\. Patwary, R\. Puri, P\. LeGresley, J\. Casper, and B\. Catanzaro \(2019\)Megatron\-lm: training multi\-billion parameter language models using model parallelism\.arXiv preprint arXiv:1909\.08053\.Cited by:[Limitations](https://arxiv.org/html/2603.01875#Sx1.p1.1)\.
- L\. von Werra, Y\. Belkada, L\. Tunstall, E\. Beeching, T\. Thrush, N\. Lambert, S\. Huang, K\. Rasul, and Q\. Gallouédec \(2020\)TRL: Transformers Reinforcement LearningExternal Links:[Link](https://github.com/huggingface/trl)Cited by:[Appendix A](https://arxiv.org/html/2603.01875#A1.SS0.SSS0.Px2.p1.1),[§1](https://arxiv.org/html/2603.01875#S1.p2.1),[§2\.2](https://arxiv.org/html/2603.01875#S2.SS2.p1.1)\.
- F\. Wan, X\. Huang, D\. Cai, X\. Quan, W\. Bi, and S\. Shi \(2024\)Knowledge fusion of large language models\.arXiv preprint arXiv:2401\.10491\.Cited by:[§2\.1](https://arxiv.org/html/2603.01875#S2.SS1.p1.1)\.
- C\. Wang, J\. Yan, W\. Cai, Y\. Yue, and J\. Huang \(2025a\)Easydistill: a comprehensive toolkit for effective knowledge distillation of large language models\.InProceedings of the 2025 Conference on Empirical Methods in Natural Language Processing: System Demonstrations,pp\. 787–795\.Cited by:[§2\.2](https://arxiv.org/html/2603.01875#S2.SS2.p1.1)\.
- W\. Wang, S\. Xiong, G\. Chen, W\. Gao, S\. Guo, Y\. He, J\. Huang, J\. Liu, Z\. Li, X\. Li,et al\.\(2025b\)Reinforcement learning optimization for large\-scale learning: an efficient and user\-friendly scaling library\.arXiv preprint arXiv:2506\.06122\.Cited by:[Appendix A](https://arxiv.org/html/2603.01875#A1.SS0.SSS0.Px2.p1.1)\.
- T\. Wu, C\. Tao, J\. Wang, R\. Yang, Z\. Zhao, and N\. Wong \(2025\)Rethinking kullback\-leibler divergence in knowledge distillation for large language models\.InProceedings of the 31st International Conference on Computational Linguistics,pp\. 5737–5755\.Cited by:[§2\.1](https://arxiv.org/html/2603.01875#S2.SS1.p1.1)\.
- B\. Xiao, B\. Xia, B\. Yang, B\. Gao, B\. Shen, C\. Zhang, C\. He, C\. Lou, F\. Luo, G\. Wang,et al\.\(2026\)Mimo\-v2\-flash technical report\.arXiv preprint arXiv:2601\.02780\.Cited by:[§2\.1](https://arxiv.org/html/2603.01875#S2.SS1.p1.1)\.
- A\. Yang, A\. Li, B\. Yang, B\. Zhang, B\. Hui, B\. Zheng, B\. Yu, C\. Gao, C\. Huang, C\. Lv,et al\.\(2025\)Qwen3 technical report\.arXiv preprint arXiv:2505\.09388\.Cited by:[Appendix A](https://arxiv.org/html/2603.01875#A1.SS0.SSS0.Px1.p1.1)\.
- F\. Yao, L\. Liu, D\. Zhang, C\. Dong, J\. Shang, and J\. Gao \(2025\)Your efficient rl framework secretly brings you off\-policy rl training\.External Links:[Link](https://fengyao.notion.site/off-policy-rl)Cited by:[§4\.2](https://arxiv.org/html/2603.01875#S4.SS2.p1.1)\.
- S\. Zhang, Y\. Liang, S\. Wang, Y\. Chen, W\. Han, J\. Liu, and J\. Xu \(2023\)Towards understanding and improving knowledge distillation for neural machine translation\.InProceedings of the 61st Annual Meeting of the Association for Computational Linguistics \(Volume 1: Long Papers\),A\. Rogers, J\. Boyd\-Graber, and N\. Okazaki \(Eds\.\),Toronto, Canada,pp\. 8062–8079\.External Links:[Link](https://aclanthology.org/2023.acl-long.448/),[Document](https://dx.doi.org/10.18653/v1/2023.acl-long.448)Cited by:[§2\.1](https://arxiv.org/html/2603.01875#S2.SS1.p1.1)\.
- S\. Zhang, X\. Zhang, Z\. Sun, Y\. Chen, and J\. Xu \(2024\)Dual\-space knowledge distillation for large language models\.InProceedings of the 2024 Conference on Empirical Methods in Natural Language Processing,pp\. 18164–18181\.Cited by:[§2\.1](https://arxiv.org/html/2603.01875#S2.SS1.p1.1)\.
- X\. Zhang, S\. Zhang, Y\. Liang, F\. Meng, Y\. Chen, J\. Xu, and J\. Zhou \(2025\)A dual\-space framework for general knowledge distillation of large language models\.arXiv preprint arXiv:2504\.11426\.Cited by:[§2\.1](https://arxiv.org/html/2603.01875#S2.SS1.p1.1),[§3\.4](https://arxiv.org/html/2603.01875#S3.SS4.p1.1)\.
- Y\. Zhao, J\. Huang, J\. Hu, X\. Wang, Y\. Mao, D\. Zhang, Z\. Jiang, Z\. Wu, B\. Ai, A\. Wang,et al\.\(2025\)Swift: a scalable lightweight infrastructure for fine\-tuning\.InProceedings of the AAAI Conference on Artificial Intelligence,Vol\.39,pp\. 29733–29735\.Cited by:[Appendix A](https://arxiv.org/html/2603.01875#A1.SS0.SSS0.Px2.p1.1),[§1](https://arxiv.org/html/2603.01875#S1.p2.1),[§2\.2](https://arxiv.org/html/2603.01875#S2.SS2.p1.1)\.
- C\. Zheng, K\. Dang, B\. Yu, M\. Li, H\. Jiang, J\. Lin, Y\. Liu, H\. Lin, C\. Wu, F\. Hu,et al\.\(2025\)Stabilizing reinforcement learning with llms: formulation and practices\.arXiv preprint arXiv:2512\.01374\.Cited by:[§4\.2](https://arxiv.org/html/2603.01875#S4.SS2.p1.1)\.
- L\. Zheng, W\. Chiang, Y\. Sheng, T\. Li, S\. Zhuang, Z\. Wu, Y\. Zhuang, Z\. Li, Z\. Lin, Eric\. P\. Xing, J\. E\. Gonzalez, I\. Stoica, and H\. Zhang \(2023\)LMSYS\-chat\-1m: a large\-scale real\-world llm conversation dataset\.External Links:2309\.11998Cited by:[Appendix A](https://arxiv.org/html/2603.01875#A1.SS0.SSS0.Px1.p1.1)\.
- L\. Zheng, L\. Yin, Z\. Xie, C\. L\. Sun, J\. Huang, C\. H\. Yu, S\. Cao, C\. Kozyrakis, I\. Stoica, J\. E\. Gonzalez,et al\.\(2024\)Sglang: efficient execution of structured language model programs\.Advances in neural information processing systems37,pp\. 62557–62583\.Cited by:[§1](https://arxiv.org/html/2603.01875#S1.p3.2)\.
- Z\. Zhu, C\. Xie, X\. Lv, and slime Contributors \(2025\)Slime: an llm post\-training framework for rl scaling\.Note:[https://github\.com/THUDM/slime](https://github.com/THUDM/slime)GitHub repository\. Corresponding author: Xin LvCited by:[§2\.2](https://arxiv.org/html/2603.01875#S2.SS2.p1.1),[2nd item](https://arxiv.org/html/2603.01875#S3.I1.i2.p1.1)\.

## Appendix AExperimental Setup Details

#### Models and Datasets\.

We evaluate the correctness and performance of KDFlow on instruction\-following tasks\. Specifically, we randomly sample 100k prompts from LMSys\-Chat\-1MZhenget al\.\([2023](https://arxiv.org/html/2603.01875#bib.bib23)\)and generate responses with Qwen3\-14B\. Then, we choose the Qwen3 model familyYanget al\.\([2025](https://arxiv.org/html/2603.01875#bib.bib24)\)for distillation since it covers multiple model sizes and architectures\. We use Qwen3\-14B, Qwen3\-32B and Qwen3\-30B\-A3B as the teacher models and Qwen3\-4B and Qwen3\-1\.7B as the student models\. We respectively report the training loss curves, model performance on AlpacaEval 2\.0Duboiset al\.\([2024](https://arxiv.org/html/2603.01875#bib.bib28)\), and training speed in the main experiments\. For AlpacaEval 2\.0, we use Qwen3\-235B\-2507\-Instruct as the evaluator model\.

#### Baselines and Hardware\.

We compare KDFlow against TRLvon Werraet al\.\([2020](https://arxiv.org/html/2603.01875#bib.bib3)\), ROLLWanget al\.\([2025b](https://arxiv.org/html/2603.01875#bib.bib25)\), and MS\-SWIFTZhaoet al\.\([2025](https://arxiv.org/html/2603.01875#bib.bib2)\), three representative frameworks for LLM distillation\. All experiments are conducted on a single server equipped with 8 NVIDIA H20 GPUs with CUDA 12\.9\.