DisaggKD: Efficient Knowledge Distillation through Teacher-Student Disaggregation and a Fused Loss

Ruixuan Tu; Yuyang Dong

cs.LGOct 13, 2026

DisaggKD: Efficient Knowledge Distillation through Teacher-Student Disaggregation and a Fused Loss

Authors: Ruixuan Tu, Yuyang Dong

Organizations: The University of Tokyo, RIKEN AIP · SB Intuitions Corp.

Abstract

Knowledge distillation (KD) frameworks place teacher and student on the same GPUs, so per-GPU memory limits the teachers that fit. DisaggKD separates them across nodes: teacher replicas stream hidden states to the student, a KD loss over row windows and vocabulary shards equals the full-vocabulary loss in exact arithmetic without forming the full-vocabulary logits, and a prefetch ring paces producer and trainer. Closed-form memory and traffic models locate the memory limits of each layout. On the Qwen3.5-35B-A3B to 9B pair, the disaggregated layout matches the LM loss of the colocated run with the same loss and is faster than that run; our fused loss also lets a colocated run train on as few nodes at faster speed, while the split is memory efficient when the student spans several nodes. On the larger pairs, colocated runs with a fused loss still run OOM where the split trains, because teacher weights and student states fill the budget.

Figures & tables

Panel of Figure 1 — Figure 1: Per-GPU memory of the Qwen3.5-35B-A3B to 9B pair on 3 nodes: bars stack student states $S$ , activations and buffers $A$ , teacher weights, logits buffers (colocated) or the memory $\kappa$ retained by the streamed loss (trainer), and other memory; diamonds are measured peaks (ModelOpt bar: allocated plus requested bytes at the OOM error). The first colocated bar (ModelOpt’s loss) is the model’s demand; the fused-W bar (Table 3 ; fused loss, model cross-entropy fusion off) splits a measured peak, with $\kappa$ the remainder. Trainer and producer nodes hold separate stacks and exchange hidden-state slots $\bm{H}_{t}^{(q)}$ over NCCL P2P.

	Colocated (4 nodes, tp2/pp4/cp4/dp1)	Disaggregated (3 nodes: 2 trainer + 1 producer)
Memory per GPU (GiB)	$S{+}T_{\mathrm{col}}{+}\kappa_{\mathrm{col}}{+}A_{\mathrm{col}}{+}\sigma_{\mathrm{col}}\leq M$ 11.5 + 15.8 + 36.0 + 0.25 + 12.3 $\approx$ predicted 75.8 (fit), measured 76.1	trainer $S{+}\kappa_{\mathrm{str}}{+}A_{\mathrm{str}}{+}\sigma_{\mathrm{str}}\leq M$ : predicted 55.8 (fit), measured 57.1 producer $\bigl(2\tfrac{D_{t}}{\mathrm{tp}_{p}}{+}\tfrac{E_{t}}{\mathrm{tp}_{p}}\bigr){+}\varphi_{p}{+}\omega_{p}\leq M$ : predicted 41.8 (fit), measured 39.4

Table 1: Estimation models per layout at its optimal setting for the Qwen3.5-35B-A3B to 9B pair (

L{=}32768

\mathrm{gbs}{=}128

; other dimensions in Appendix A.2 ). Memory rows: predictions against measured device peaks (GiB). Traffic rows (( 6 ), ( 7 )): GB sent by a GPU per iteration and, in parentheses, the time at the nominal limits (NVLink 450 GB/s per GPU and direction; InfiniBand 400 Gb/s per GPU); measured NVLink: counter bytes per GPU, iterations 3 to 12 (colocated) and 3 to 10 (disaggregated).

Panel of Figure 2 — Figure 2: The disaggregated algorithm (symbols as in § 4.1 and § 4.2 ; $\cdots$ : omitted ranks). Top : the producer sends slot $\bm{H}_{t}^{(q)}$ to every tp rank of its cp rank; cell $(q,j)$ of the receiver group (bold) holds the logits block $\bm{Z}_{x}^{(q,j)}$ , and its prefetch ring holds up to 4 slots of $K$ windows. (a) Window $k$ : $\bm{H}_{x,k}^{(q)}\bm{W}_{x}^{(j)}$ gives $\bm{Z}_{x,k}^{(q,j)}$ ; row $i$ (white) is what the kernel scans. (b) The kernel walks the row in tiles with a running maximum (red); tp ranks combine the statistics by all-reduce.

	ModelOpt	ModelOpt fused-W	KDFlow	NeMo-RL	DisaggKD (ours)
Teacher and student on	same GPUs	same GPUs	same GPUs	same GPUs	separate nodes
Idle teacher weights	kept	kept	kept (optional offload)	offloaded	lm_head only
Teacher forward	in the student’s step	in the student’s step	phases	1 phase per step	concurrent, ahead
Teacher output	logits, $V/\mathrm{tp}$ columns	logits, $V/\mathrm{tp}$ columns	hidden states	logits, $V$ columns	hidden states
KL taken over	all $V$ columns	all $V$ columns	all $V$ columns	top 8192 columns	all $V$ columns
Largest logits tensor	$L/\mathrm{cp}\times V/\mathrm{tp}$	$L/\mathrm{cp}\times V/\mathrm{tp}$	tokens $\times V$ , chunks optional	step tokens $\times V$	$C\times V/\mathrm{tp}$

Table 2: DisaggKD against ModelOpt ( NVIDIA, 2025a ; NVIDIA, 2024–2026 ) , ModelOpt with the fused-W loss of DisaggKD (Appendix C.2 ), KDFlow ( Zhang et al., 2026 ) , and NeMo-RL ( NeMo RL, 2025 ) . Measured rows are our runs (Table 19 ); NeMo-RL on Qwen3.5 ran OOM at 4.19M tokens per step; ‡ : fused, cross-entropy fusion on;

-

: not run. Bold: best tokens/s per GPU in each pair and node count.

Qwen3.5: 35B/3B to 9B
Nodes	DisaggKD (ours)	ModelOpt	ModelOpt fused-W
2	284	$-$	147
3	74	OOM	72
4	70	91 to 177	101 ‡
Qwen3-235B: 235B/22B to 30B/3B
3	393	OOM	OOM

Table 3: Outcome of each layout by node count: median s/it or OOM;

-

: not run. ModelOpt fused-W (Appendix C.2 ): the fused loss of DisaggKD with the model’s cross-entropy fusion off ( ‡ : on); DisaggKD forms no model cross-entropy. Entries pool the grids and losses run (Tables 18 , 17 , and 30 a).

Panel of Figure 3 — Figure 3: Loss parity and throughput on the Qwen3.5 pair. Colocated runs are ModelOpt with the loss in the legend; fused and fused-W use the fused loss of DisaggKD (hatched: fused-W). (a) LM loss of the 4-node runs (thin: per iteration, thick: 17-iteration mean). (b) KD loss minus its first-3-iteration mean. (c) Median total throughput by node count (2-node, fused, and fused-W bars: 15 iterations).

Colocated (ModelOpt)
	s/it, seed		Peak	MMLU, seed
Loss (nodes)	42	43		42	43
KD-only (4)	90.5	97.7	63.4	67.97	69.22
Bal-L+K (4)	176.8	172.1	67.2	67.51	69.78
SQ-L+K (4)	96.3	–	59.9	–	–
SQ-L+K fused (4)	100.6 †	–	40.7	–	–

Table 4: Primary pair at iteration 200, 1 run per cell; –: not run; † : 15 iterations ( SQ-L+K : 110.3); ‡ : 30 iterations without cache release (Appendix D.1 ; 99.4 with it); W: cross-entropy fusion off, batch 126; peak: torch GiB; best bold.

Appendix figures & tables29 assets

Supplementary material from the paper’s appendix.

Appendix

Module	Term
Student weights, gradients, optimizer	$({\color[rgb]{0,0.45,0}2}{+}{\color[rgb]{0,0.45,0}4})\,P_{\mathrm{loc}}\,\mathrm{bytes}+{\color[rgb]{0,0.45,0}12}(D_{s}/(\mathrm{tp}\,\mathrm{pp}\,\mathrm{dp}\,\mathrm{cp})+E_{s}/(\mathrm{pp}\,\mathrm{ep}\,\mathrm{edp}))\,\mathrm{bytes}$ , experts on $\mathrm{pp}\cdot\mathrm{ep}$ , their optimizer state on $\mathrm{edp}$ ranks
Student activations, colocated	${\color[rgb]{0,0.45,0}1}\tfrac{\ell_{s}}{\mathrm{pp}}\tfrac{L}{\mathrm{cp}\,\mathrm{tp}}d_{s}\,b\,\mathrm{bytes}$ stored layer inputs of the 1 microbatch in flight at the loss stage (full recomputation)
Student attention scores	$(L/\mathrm{cp})^{2}(h/\mathrm{tp})\,b\,\mathrm{bytes}$ when attention is not fused (zero for the students of this paper)
Student MoE dispatch	${\color[rgb]{0.75,0,0}2}\tfrac{L}{\mathrm{cp}}k_{\mathrm{e}}d_{s}\,b/\mathrm{tp}\,\mathrm{bytes}$
Colocated teacher shard	${\color[rgb]{0,0.45,0}2}(D_{t}/(\mathrm{pp}\,\mathrm{tp})+E_{t}/(\mathrm{pp}\,\mathrm{ep}))\,\mathrm{bytes}$ (bf16)
Colocated logits buffers	$({\color[rgb]{0.75,0,0}7}{+}1{+}{\color[rgb]{0.75,0,0}1.5})\,(L/\mathrm{cp})(V/\mathrm{tp})\,{\color[rgb]{0,0.45,0}4}\,\mathrm{bytes}$ (7 buffers, 1 for the loss-averaging all-reduce, and 1.5 for the backward pass; 1 more with the LM loss)

Appendix

Table 5: Per-module terms of the memory model (

k_{\mathrm{e}}

: routed experts per token,

h

: attention heads). Green coefficients follow from the code or the framework; red coefficients are fitted to measurements; black numbers are allowances. Parameter terms are in exact GiB (Appendix A.1 ).

Model	In the fit	Outside the fit
Trainer device peak	8 runs: RMS 1.2, largest 2.5	13 runs: RMS 1.7, largest 4.0
Producer device peak	8 runs: RMS 5.0, largest 8.4	leave one MoE family out (6 runs): RMS 8.6, largest 12.3
Colocated verdict	6 configurations: 5 correct	18 configurations: 17 correct
Colocated peak	4 peaks, 4 OOM points: RMS 0.5, largest 1.2	9 peaks: RMS 1.5, largest 3.0; 2 OOM points, largest $-$ 15.8
Inter-node traffic (Qwen3.5)	$-$	within 2% (3 and 4 nodes)
NVLink traffic (Qwen3.5)	$-$	4.4% (colocated), 15% (trainer), 3% (producer)

Appendix

Table 6: Accuracy of the estimation models: root-mean-square (RMS) and largest error (predicted minus measured, GiB) of the trainer and producer device peaks (Tables 8 and 7 ), where the producer’s constants are fitted on all 8 runs and the second column predicts each MoE family from the others (the 2 dense runs fix their own constants); the colocated rows are verdicts (fits or OOM) and peaks or OOM points (Table 18 ), where the fit rows are those that set a constant; traffic is the deviation of the counters from the calculation (Tables 14 and 1 ). The trainer’s 13 runs outside the fit include the 2 runs with all 3 parts of the cache release (CR all; Appendix D.1 ), of which the (2+1) reading, 79.6 GiB, is at the device limit.

Run (T+P nodes)	Replica	Shard	Pred.	Meas.	Err.	Resid.	Err., family left out
Qwen3.5 (2+1)	TP2, fp8	18.1	41.8	39.4	$+2.4$	21.3	$+3.5$
Qwen3.5 (3+1)	TP2 $\times$ EP2, fp8	10.4	34.1	32.0	$+2.1$	21.6	$+3.2$
Qwen3-235B (2+1)	TP8, bf16	54.7	78.4	72.4	$+6.1$	17.7	$+8.9$
Qwen3-235B (4+4)	TP8, bf16	54.7	78.4	73.3	$+5.2$	18.6	$+8.0$
Nemotron (2+2)	TP8 $\times$ EP2, fp8	38.6	62.3	69.6	$-7.3$	31.0	$-11.2$
Nemotron (4+4)	TP8 $\times$ EP2, fp8	38.6	62.3	70.7	$-8.4$	32.1	$-12.3$

Appendix

Table 7: Producer memory per GPU (GiB) of 8 of the 10 runs of Table 17 . Shard: the whole checkpoint’s per-rank weights, dense tensors in bf16 and routed experts in fp8 or bf16 as listed (Appendix A.1 ); Pred.: shard plus the forward-pass term (22.70 for MoE teachers; 4.4 plus 12.8 bf16 buffers of

L\times d_{t}

for dense ones, at the sequence length of the run: 32768, and 4096 for Qwen3-32B) plus the 1 for NCCL buffers; Meas.: device peak; Err.: predicted minus measured; Resid.: measured minus shard; Err., family left out: the error when the forward-pass term is fitted on the other families (not defined for the 2 dense teachers, whose 2 parameters the 2 dense runs fix). The Qwen3-235B (2+1) row uses the 72.4 of an earlier run of that configuration; the run of Table 17 reads 72.9.

Trainer	Grid	$n_{\mathrm{mb}}$	$S$	$\kappa_{\mathrm{str}}$	$A_{\mathrm{gc}}$	$A_{\mathrm{bwd}}$	$A_{\mathrm{moe}}$	Pred.	Meas.	Err.	Use
Qwen3.5 (2+1)	2/2/2/2	64	18.8	1.1	12.0	8.6	0.0	55.8	57.1	$-$ 1.3	fit
Qwen3.5 (3+1)	2/2/2/3	42	16.7	1.1	7.9	8.6	0.0	49.6	50.3	$-$ 0.7	fit
Qwen3.5 (2+1), tp4	4/2/2/1	128	12.5	0.7	12.0	4.3	0.0	44.9	45.1	$-$ 0.2	fit
Qwen3-32B (1+1), seq 4096	2/2/1/2	64	12.3	1.2	1.9	10.1	0.0	40.8	40.3	+0.5	fit
Llama (8+8)	8/2/2/2	64	37.0	2.0	6.0	3.2	0.0	63.5	61.0	+2.5	fit
Nemotron (2+2)	2/2/1/4/4	32	47.4	2.8	7.9	4.5	2.5	80.4	79.6	+0.8	fit

Appendix

Table 8: Terms of the trainer memory model ( 4 ) and the measured device peak of the trainer runs the model uses (GiB). Grid: tp/pp/cp/dp, with ep appended for MoE students; sequence length 32768 unless stated;

n_{\mathrm{mb}}

: microbatches per DP replica. Pred.: the 5 terms shown plus

\sigma_{\mathrm{str}}{=}15.32

; Meas.: device peak, the maximum over the trainer GPUs; Err.: predicted minus measured. Use: fit (the 8 runs whose mean remainder is

\sigma_{\mathrm{str}}

), held out (outside that mean; the term forms were chosen with their residuals visible), or ceiling (CR all, Table 27 : the device reading is an allocator ceiling). Collected: Python’s garbage collected after every microbatch. The Qwen3-235B pair rows are the 30B student; the Qwen3-235B (2+1), collected row was predicted at 73.8 before its run, 0.3 below this value, which adds the fp8 head copy of its configuration. Teacher precision: the Qwen3-235B (4+4), ep8 and (2+1), collected rows are fp8-teacher runs outside Table 17 ; the 2 ceiling rows are bf16.

Configuration	Pred.	Measured	Use
Coloc., 3 nodes, cp1	86.8	OOM, 72.9+15.2	held out
Coloc., 4 nodes, cp1	86.4	OOM, 72.4+15.2	fit
Coloc., 4 nodes, cp2	78.3	OOM, 71.4+7.6	fit
Coloc., 4 nodes, cp4	75.8	76.1	fit
Coloc., 4 nodes, cp4, Bal-L+K	79.6	79.1	fit
Coloc., 4 nodes, cp8	77.3	77.0	held out

Appendix

Table 9: Predicted (Pred.) memory demand and measured peak per GPU (GiB) on the Qwen3.5 pair (and 1 Qwen3-30B trainer); Use: whether the row set a constant of the model (fit), not (held out), or is outside the model because its loss is not ModelOpt’s (outside). OOM rows compare the predicted demand at the failing allocation with the measured bytes allocated plus requested; all other rows are device peaks (torch-allocated trainer peak 39.2). The colocated cp4 rows fix the non-PyTorch memory and the logits-buffer count, the trainer rows the trainer remainder

\sigma_{\mathrm{str}}

, and the producer rows the forward-pass term

\varphi_{p}

(Appendix A.1 ); “collected” is the (2+1) trainer with Python’s garbage collected after every microbatch. Tables 18 , 8 , and 7 list all calibration points.

Remainder after subtracting	Mean	SD	LOO	LOFO
$S$ , $\kappa_{\mathrm{str}}$	29.91	4.7	5.1	5.7
and $A_{\mathrm{gc}}$	23.09	4.1	4.4	4.5
and $A_{\mathrm{bwd}}$	16.41	2.3	2.4	2.4
and $A_{\mathrm{moe}}$	15.32	1.3	1.4	1.4

Appendix

Table 10: Remainder of the device peaks (GiB) of the 8 fitted trainers after the terms listed are subtracted cumulatively: its mean, SD, and the RMS error of predicting each trainer by the mean of the others (LOO) or of the other model families (LOFO). The last row is

\sigma_{\mathrm{str}}

and its spread.

Trainer	Device	Torch	Reserved	Non-PyTorch	Unallocated	End of iter.	Stage 1	Ramp	$A_{\mathrm{gc}}$
Qwen3.5 (2+1)	57.1	39.2	43.1	13.9	3.8	31.8	+0.3	11.6	12.0
Qwen3.5 (3+1)	50.3	33.0	36.6	13.4	3.6	25.6	$-$	$-$	7.9
Qwen3.5 (2+1), tp4	45.1	29.0	31.1	14.0	2.1	25.1	$-$	11.7	12.0
Qwen3-32B (1+1)	40.3	25.1	28.8	11.5	3.7	15.6	+0.2	1.7	1.9
Llama (8+8)	61.0	50.3	52.7	7.8	2.4	43.8	+0.2	$-$	6.0
Nemotron (2+2)	79.6	67.5	69.8	9.7	3.6	59.3	+2.2	$-$	7.9

Appendix

Table 11: Measured memory of the 8 fitted trainers, the 2 runs with CR all, and the 7 Qwen3.5 trainers outside the

\sigma_{\mathrm{str}}

fit (below the rule; GiB), on the last pipeline stage. Device: peak device memory, the maximum over the trainer GPUs; Torch: peak allocated (the maximum over the trainer GPUs for the 2 Nemotron rows); Reserved: peak reserved; Non-PyTorch: device memory minus reserved on the last-stage rank; Unallocated: reserved minus peak allocated (Device is a maximum over the trainer GPUs and the other columns are logged on 1 last-stage rank, so a row need not add up); End of iter.: allocation after the iteration; Stage 1: steady allocation of the first pipeline stage minus its statics (not logged for 3 runs and the held-out trainers); Ramp: growth of the allocation over the microbatches of an iteration, from the per-microbatch memory log (not logged for 4 runs and the held-out trainers);

A_{\mathrm{gc}}

: the model’s garbage term.

Probe	Trainer grid	s/it	Torch	Device	NCCL	Comms	Measured; model
Qwen3.5 (2+1), garbage kept	tp2/pp2/cp2/dp2	74.3	39.2	57.1	$-$	$-$	end of iter. 31.8
Qwen3.5 (2+1), collected	tp2/pp2/cp2/dp2	116.9	27.7	44.8	$-$	$-$	end of iter. 20.0; $-$ 11.8
Qwen3.5 (2+1), tp4, reference	tp4/pp2/cp2/dp1	97.0	29.0	45.1	59.6 (3.73)	17	$-$
NVLS multicast off	tp4/pp2/cp2/dp1	98.2	29.0	44.8	59.6 (3.72)	17	$-$
eager kernel loading	tp4/pp2/cp2/dp1	100.2	29.0	46.8	59.6 (3.72)	17	$-$
Qwen3-32B (1+1), tp2	tp2/pp2/cp1/dp2	$-$	25.1	$-$	38.2 (4.78)	19	$-$

Appendix

Table 12: Probes of the trainer memory terms (GiB;

-

: not measured), all with CR

-

(Table 27 ). Torch: peak allocated; Device: peak device memory; NCCL: all communicator buffers over the trainer ranks, per rank in parentheses; Comms: communicators per rank. Collected: Python’s garbage collected after every microbatch, which lowers the end-of-iteration allocation by 11.8 GiB against the garbage-kept run, equal to

A_{\mathrm{gc}}

of the same grid (calibrated on its ramp). Temporary memory: peak minus end-of-iteration allocation against the term

A_{\mathrm{bwd}}

(Qwen3-4B student); retained set: tensors alive across the optimizer step (autograd graph 1.50, bias-dropout activations 1.44, dispatch maps and workspaces 0.63) against

A_{\mathrm{moe}}

(the family that calibrates its 2.5 copies); statics: the allocation dump against

S

. Qwen3-235B (2+1), collected is the 30B student with the 235B teacher.

			Qwen3.5 (GB; student; teacher)
Dimension	Teacher (forward)	Student (training)	colocated	disaggregated
Tensor	$2\tfrac{\ell_{t}}{\mathrm{pp}}\tfrac{2(\mathrm{tp}-1)}{\mathrm{tp}}\,\tau\,d_{t}b\,\mathrm{bytes}$	$6\tfrac{\ell_{s}}{\mathrm{pp}}\tfrac{2(\mathrm{tp}-1)}{\mathrm{tp}}\,\tau\,d_{s}b\,\mathrm{bytes}$	412.3; 85.9	824.6; 343.6
Pipeline, per boundary and direction	$n_{\mathrm{mb}}\tfrac{L}{\mathrm{cp}}\tfrac{d_{t}b}{\mathrm{tp}}\,\mathrm{bytes}$ , both directions colocated	$n_{\mathrm{mb}}\tfrac{L}{\mathrm{cp}}\tfrac{d_{s}b}{\mathrm{tp}}\,\mathrm{bytes}$ , both directions	4.29; 2.15	4.29; 0
Context, attention	$\tfrac{\ell^{\mathrm{att}}_{t}}{\mathrm{pp}}(\mathrm{cp}{-}1)\,\tau\,2\tfrac{h_{\mathrm{kv},t}}{\mathrm{tp}}d_{\mathrm{head}}b\,\mathrm{bytes}$	$4\tfrac{\ell^{\mathrm{att}}_{s}}{\mathrm{pp}}(\mathrm{cp}{-}1)\,\tau\,2\tfrac{h_{\mathrm{kv},s}}{\mathrm{tp}}d_{\mathrm{head}}b\,\mathrm{bytes}$	51.5; 8.1	34.4; 0
Context, linear attention	$\tfrac{\ell^{\mathrm{lin}}_{t}}{\mathrm{pp}}\tfrac{\mathrm{cp}-1}{\mathrm{cp}}\,\tau\,\tfrac{w}{\mathrm{tp}}b\,\mathrm{bytes}$	$3\tfrac{\ell^{\mathrm{lin}}_{s}}{\mathrm{pp}}\tfrac{\mathrm{cp}-1}{\mathrm{cp}}\,\tau\,\tfrac{w}{\mathrm{tp}}b\,\mathrm{bytes}$	232.8; 97.0	310.5; 0
Data and context	none	$\tfrac{\mathrm{dp}\,\mathrm{cp}-1}{\mathrm{dp}\,\mathrm{cp}}P_{\mathrm{loc}}\,(4{+}2)\,\mathrm{bytes}$	5.0; 0	10.1; 0

Appendix

Table 13: Bytes sent per GPU and iteration by parallelism dimension, with the values on the Qwen3.5 pair (4 nodes colocated, 3 disaggregated; the disaggregated teacher is the producer). A dimension of degree

n

and rank stride

\Delta

spans

n\Delta

ranks and uses NVLink if

n\Delta\leq 8

(the GPUs of a node) and InfiniBand at its node boundaries otherwise; the strides are 1 (tensor), tp (context, and data at ep

{=}1

\mathrm{tp}\,\mathrm{cp}

(expert), and

\mathrm{tp}\,\mathrm{cp}\,\mathrm{ep}\,\mathrm{dp}

(pipeline). The factors 2 and 6 are 2 all-reduces per layer times the teacher’s 1 forward pass or the student’s 3 passes (forward, recompute, backward); 3 and 4 count passes (the ring attention ( Liu et al., 2024 ) backward exchanges keys and values and their gradients, hence 4).

\ell^{\mathrm{att}}

\ell^{\mathrm{lin}}

, and

\ell^{\mathrm{moe}}

count full-attention, linear-attention, and MoE layers,

h_{\mathrm{kv}}

the KV heads,

d_{\mathrm{head}}

the head width, and

w{=}16{,}448

the per-token width of the linear-attention exchange of the Qwen3.5 layers (12,352 for the input projection of

q,k,v

, gate, and 2 scalars per head, 4,096 for the output). The expert term counts every routed copy of a token: the all-to-all over the EP group and, at

\mathrm{etp}>1

, the dispatch all-gather in bf16 and the combine reduce-scatter in fp32 over the expert tensor-parallel group; it is 0 at

\mathrm{ep}{=}1

and

\mathrm{etp}{=}1

. Gradients are reduced in fp32 (4 bytes) and parameters gathered in bf16 (2 bytes).

3 nodes (dp2)
	Sent		Received		Sent:
Node	calc.	meas.	calc.	meas.	pipe.+ring
Stage 0	34.4	34.8	34.4	34.8	34.4+0
Last stage	34.4	34.9	68.7	69.6	34.4+0
Producer	34.4	34.8	0	0.1
4 nodes (dp3)

Appendix

Table 14: Inter-node traffic per node and iteration (GB): calculated and measured, and the calculated split of the sent bytes into pipeline activation transfers and DP

\times

CP rings. Received bytes at the trainer’s last-stage nodes include the transfer. Colocated: 1 pipeline stage per node, iterations 3 to 12.

Pair (T+P nodes)	$R$	Consumed	$\rho_{\mathrm{rep}}$	$N_{\mathrm{tok}}/(R\rho_{\mathrm{rep}})$	Meas.	Duty	Limit	CR
Qwen3.5 (1+1), 15 it.	4	3.7	28.1	37	283.7	0.13	trainer	$-$
Qwen3.5 (2+1)	4	13.8	28.1	37	74.3	0.50	trainer	$-$
Qwen3.5 (3+1)	2	29.5	32.8	63	69.9	0.90	trainer	$-$
Qwen3.5 (3+1), rerun	2	31.4	32.8	63	65.0	0.97	trainer	$-$
Qwen3.5 (3+1), seed 43	2	20.1	32.8	63	99.4	0.63	trainer	all
Qwen3.5 (3+1), seed 43, 30 it.	2	33.2	32.8	63	61.9	1.0	producer	$-$

Appendix

Table 15: Which side limits the iteration time (T+P: trainer and producer nodes;

R

: producer replicas). Consumed: mean per-replica rate between the 5th and the last report of the producer logs;

\rho_{\mathrm{rep}}

: forward-only rate of one replica, both in

10^{3}

tokens/s;

N_{\mathrm{tok}}/(R\rho_{\mathrm{rep}})

: the iteration time if the trainer were instantaneous; Meas.: measured median iteration time, both in seconds; Duty:

N_{\mathrm{tok}}/(R\rho_{\mathrm{rep}})

over Meas., capped at 1; CR: the parts of the cache release that were on (Table 27 ). The producers limit a run when the duty is 1 (the trainer is then at least as fast) and the trainer limits it when the duty is below 1. Rerun: the 3+1 run of Table 4 repeated for 12 iterations (65.0 s against 69.9 s, 7%). Each row is 1 run.

Input: teacher replica, corpus order, look-ahead $\delta$ iterations, start slot $\eta_{0}$
1	load the teacher shard (routed experts in fp8 for the Qwen3.5 and Nemotron teachers, bf16 otherwise); write the lm_head shards once
2	for $\eta=\eta_{0},\eta_{0}{+}1,\dots$ do
3	if $\eta\bmod R\neq$ my replica then continue
4	$\bm{H}_{t}\leftarrow$ forward of microbatch $\eta$ (final hidden states)
5	wait until $\eta<(\text{iteration}+\delta)\cdot\text{slots per iteration}$
6	slice rows per last-stage rank; send header and payload to the receiver group

Appendix

Algorithm 1: Producer loop of one replica.

Input: student hidden states $\bm{H}_{s}^{(q)}$ , prefetch ring, shards $\bm{W}_{t}^{(j)},\bm{W}_{s}^{(j)}$ , labels $y$ , window size $C$
1	$\bm{H}_{t}^{(q)}\leftarrow$ pop the next slot from the ring; post receives for up to 4 following slots of the same iteration
2	for $k=0,\dots,K{-}1$ do
3	$\bm{Z}_{x,k}^{(q,j)}\leftarrow\bm{H}_{x,k}^{(q)}\bm{W}_{x}^{(j)}$ for $x\in\{t,s\}$
4	kernel: for each row, one pass over the entries $u=0,\dots,V_{\mathrm{loc}}{-}1$ in tiles gives $m_{x}^{(j)},a_{x}^{(j)}$ ( $x\in\{t,s\}$ ) and $c^{(j)}$
5	MAX all-reduce of $m_{x}^{(j)}$ ; rescale; SUM all-reduce of $a_{x}^{(j)},c^{(j)}$ over the tp ranks; the label logit of $\mathrm{CE}_{i}$ comes from a separate vocabulary-parallel pass
6	$\mathcal{L}\leftarrow\mathcal{L}+\sum_{i}[(1{-}\lambda)\mathrm{CE}_{i}+\lambda\,\mathrm{KD}_{i}]$ over the rows of window $k$ ; free $\bm{Z}_{x,k}^{(q,j)}$

Appendix

Algorithm 2: Streamed KD step for one microbatch at cell

(q,j)

(§ 4.2 ).

Experts	s/it	Torch	KD loss, it. 1, 2, 3	Grad norm, it. 1, 2, 3
bf16	100.3	61.7	0.215, 0.207, 0.207	0.70, 0.70, 1.04
fp8	119.1	55.9	0.215, 0.207, 0.795	0.70, 0.70, 14.5

Appendix

Table 16: Colocated KD-only runs of the Qwen3.5 pair on 4 nodes (tp2/pp4/cp4/dp1, 32768 tokens, 12 iterations) with the teacher’s routed experts in bf16 or fp8 (Appendix B.2 ). s/it: median over iterations 3 to 12; Torch: peak allocated in GiB after iteration 1. From iteration 4 to 12 the KD loss is 0.505 to 0.665 with fp8 against 0.249 to 0.290 with bf16, so the 2 s/it values do not compare precisions.

Run (T+P nodes)	Trainer grid	Producer replicas	Iter.	s/it	$10^{3}$ tok/s	KD loss	Peak	CR
Qwen3.5 (1+1)	tp2/pp2/cp2/dp1	4 $\times$ TP2	15	283.7	14.8	2.228 $\to$ 2.488	57.5; 39.4	$-$
Qwen3.5 (2+1)	tp2/pp2/cp2/dp2	4 $\times$ TP2	200	74.3	56.4	2.228 $\to$ 2.227	39.2; 39.4	$-$
Qwen3.5 (3+1)	tp2/pp2/cp2/dp3	2 $\times$ TP2 $\times$ EP2	200	69.9	59.1	2.226 $\to$ 2.225	33.0; 32.0	$-$
Qwen3-235B (2+1)	tp2/pp2/ep4	1 $\times$ TP8 (bf16)	200	392.9	10.7	2.020 $\to$ 1.962	64.3; 72.9	all
Qwen3-235B (4+4)	tp2/pp2/ep4, dp8	4 $\times$ TP8 (bf16)	200	108.5	38.7	2.020 $\to$ 1.962	50.7; 73.3	all
Qwen3-235B (4+4), KD-only	tp2/pp2/ep4, dp8	4 $\times$ TP8	60	105.0	40.0	2.066 $\to$ 2.029	51.0; 72.9	all

Appendix

Table 17: Disaggregated runs (T+P: trainer and producer nodes): trainer grid, producer replicas (fp8 experts; the Qwen3-235B, Llama, and Qwen3-32B teachers are bf16), iterations completed (Iter.), median iteration time (s/it), throughput (

10^{3}

tokens/s), KD loss at the first and last iteration, and peak memory in GiB (trainer torch-allocated on the last pipeline stage; producers’ device). CR: the parts of the cache release that were on (Table 27 ).

Panel of Figure 4 — Figure 4: KD-loss trajectories of the disaggregated runs on the 4 pairs of Table 3 (per-iteration values, thin; 17-iteration moving averages, thick). Students start from released checkpoints, so levels differ across pairs. Blue: the smallest disaggregated run of the pair; orange: the run on the node count of the colocated cell (4 nodes for Qwen3.5, 8 for Qwen3-235B and Nemotron); the Llama panel shows its 16-node run (blue). The labels give the number of iterations of each run, in the color of its curve.

		Device		Torch		OOM point
Pair	Configuration	Pred.	Meas.	Pred.	Meas.	Pred.	Meas.	Outcome
Qwen3.5	3 nodes, tp2/dp3, best grid tp2/pp2/cp2, global batch 126	133.1	$-$	120.8	$-$	$-$	$-$	not run
	3 nodes, tp2/pp4/cp1/dp3	185.9	$-$	173.6	$-$	86.8	88.0	OOM at the 4th buffer (student divide)
	4 nodes, cp1	184.6	$-$	172.3	$-$	86.4	87.6	OOM at the 4th buffer
	4 nodes, cp2	112.1	$-$	99.8	$-$	78.3	79.0	OOM at the 7th buffer
	4 nodes, cp4/dp1	75.8	76.1	63.6	63.4	$-$	$-$	trains 200 iterations

Appendix

Table 18: Colocated configurations for each pair (Megatron-Bridge ModelOpt; KD-only unless a row says Bal-L+K , SQ-L+K , SQ-L+K fused, or all-gather; W: the model’s cross-entropy fusion off), 32768-token sequences unless stated. Pred. and Meas. compare the model with the measurement of the same quantity in GiB: the device peak, the peak allocated by PyTorch (torch), or, at an OOM error, the bytes allocated plus the bytes requested (OOM point);

-

: not predicted or not measured. The model’s MoE dispatch term is that of the all-to-all dispatcher, and the all-gather rows and the SQ-L+K rows, whose loss is not the modeled one, are not counted in Table 6 .

Stack	Pair	Nodes	Steps	Tokens ( $10^{3}$ )	s/step	Tok/s/GPU	Peak	Setting
ModelOpt, KD-only	Qwen3.5	4	200	4194	90.5	1450	63.4	tp2/pp4/cp4/dp1
ModelOpt, Bal-L+K	Qwen3.5	4	200	4194	176.8	741	67.2	tp2/pp4/cp4/dp1
ModelOpt, SQ-L+K	Qwen3.5	4	200	4194	96.3	1361	59.9	tp2/pp4/cp4/dp1
ModelOpt, SQ-L+K fused	Qwen3.5	4	15	4194	100.6	1303	40.7	tp2/pp4/cp4/dp1
ModelOpt, SQ-L+K fused-W	Qwen3.5	3	15	4129	72.0	2389	64.1	tp2/pp4/cp1/dp3
ModelOpt, SQ-L+K fused-W	Qwen3.5	2	15	4194	146.8	1786	50.3	tp2/pp4/cp2/dp1

Appendix

Table 19: Measured runs of the compared stacks: nodes of 8 GPUs, steps run, tokens per step as counted, median step time (a 9-step mean for KDFlow on Qwen3-32B; steps 3 to 15 for the 15-step run), tokens/s per GPU, and peak memory per GPU in GiB (torch-allocated for ModelOpt and DisaggKD , device memory for KDFlow and NeMo-RL, whose torch peaks were not logged; the producers’ device peaks of DisaggKD are in Table 17 ). KDFlow on Qwen3.5 used sleep mode, and its step includes the teacher offload cycle; NeMo-RL on Qwen3.5 ran at a quarter of the batch without the grouped-GEMM kernel. CR (in the Setting column): the parts of the cache release that were on (Table 27 ). Bold: highest tokens/s per GPU in each pair and node count; the runs differ in loss, sequence length, steps, tokens per step, and teacher precision.

(a) LM loss gap to the colocated run
	Bal-L+K				SQ-L+K
	Disagg., 4 nodes		Disagg., rescaled		Disagg., 4 nodes
$w$	$\|$ gap $\|$	signed	$\|$ gap $\|$	signed	$\|$ gap $\|$	signed
1	0.020	$+$ 0.001	0.021	$+$ 0.005	0.017	$+$ 0.005
5	0.014	$+$ 0.001	0.014	$+$ 0.005	0.010	$+$ 0.005
9	0.012	$+$ 0.001	0.012	$+$ 0.006	0.009	$+$ 0.005

Appendix

Table 20: Loss parity on the primary pair. (a) Mean absolute and mean signed gap between the LM loss of a disaggregated run and a colocated run ( Bal-L+K or SQ-L+K ), over moving averages of

w

iterations, at the iterations at which both have a full window (183 at

w{=}17

for the 4-node run, 184 for the rescaled 2+1 run against the Bal-L+K run). (b) Largest rise of each run’s KD loss, its value at the last full window (iteration 192), and the gap to the disaggregated curve (disaggregated minus colocated; mean absolute and mean signed over iterations 9 to 192). Max: the maximum absolute gap at

w{=}17

Iteration	20	50	100	200
Colocated Bal-L+K	0.88	0.87	0.89	0.91
SQ-L+K (disagg. and coloc.)	0.98	0.94	0.88	0.75
Disaggregated, rescaled	0.49	0.49	0.49	0.49
Colocated KD-only	1	1	1	1

Appendix

Table 21: Weight of the KD term in the loss of the primary-pair runs:

\lambda

for the mixture,

\gamma/(1{+}\gamma)

with

\gamma=\mathrm{LM}/\mathrm{KL}

for Bal-L+K and

\gamma_{\mathrm{res}}=\mathrm{LM}/\mathrm{KD}

for the rescaled run.

Loss	Alloc.	Req.	In use	Failed in
KD-only, ModelOpt	72.9	15.2	75.5	the loss
SQ-L+K fused	57.7	15.2	66.2	model forward
SQ-L+K fused-W	64.1	$-$	78.9	trains

Appendix

Table 22: Colocated 3-node cell of the primary pair (tp2/pp4/cp1/dp3) at its first failing allocation or, for the run that trains, at its peak (GiB): PyTorch-allocated bytes (Alloc.), bytes requested by the failing allocation (Req.), and device memory in use (device peak for fused-W).

Run	Nodes	s/it	GPU-s	Peak
Colocated, cp4/dp1	4	96.3	3082	59.9
Colocated fused, cp4/dp1	4	100.6 †	3219	40.7
Colocated fused-W, cp1/dp3	3	72.0 †	1728	64.1
Colocated fused-W, cp2/dp1	2	146.8 †	2349	50.3
Disaggregated, cp2/dp1	1+1	283.7 †	4539	57.5
Disaggregated, cp2/dp3	3+1	69.9	2237	33.0

Appendix

Table 23: SQ-L+K runs of the primary pair at 32768 tokens (tp2; pp4 colocated, pp2 disaggregated): median s/it, GPU-seconds per iteration over all nodes (producers included), and torch-allocated peak in GiB (the trainer’s for the disaggregated runs). † : 15 iterations; fused-W: the model’s cross-entropy fusion off, global batch 126 at 3 nodes and 128 at 2 nodes.

(a) Profile of the colocated runs			(b) Streamed trainer on other grids
	KD-only	Bal-L+K	Trainer grid	Nodes	s/it	GPU-s	Relative	CR
Loss stage, GPU-kernel share (%)	85.9	77.6	tp2/pp2/cp2/dp4	4+2	39.3	1886	1.00	$-$
Loss stage, NCCL collectives (s)	2.3	18.3	tp2/pp4/cp2/dp2	4+2	50.8	2438	1.29	$-$
Loss stage, copy operators (s)	7.1	23.4	tp2/pp2/cp4/dp2	4+2	50.0	2400	1.27	$-$
Loss stage, stream syncs (s)	13.2	30.4	tp2/pp4/cp4/dp1	4+2	70.9	3403	1.80	$-$
Stage 1, GPU-kernel share (%)	98.0	98.5	tp2/pp2/cp2/dp3	3+1	65.0	2080	$-$	$-$

Appendix

Table 24: (a) Profiles of iteration 3 of the colocated KD-only and Bal-L+K runs (tp2/pp4/cp4/dp1), on 1 rank per pipeline stage. Times are seconds of traced time in the iteration, whose wall time on the loss stage is 111.3 s (KD-only) and 120.0 s ( Bal-L+K ); iteration 3 is not steady (the median iterations take 90.5 s and 176.8 s, Table 4 ); the rows are overlapping event times, so differences and shares, not absolute times, are the result. (b) Streamed trainer ( SQ-L+K ) on the Qwen3.5 pair at 4.19M tokens per iteration (4.13M in the 3+1 row), on 4 trainer and 2 producer nodes (Nodes), 12 iterations each; GPU-s: GPU-seconds per iteration over all nodes, producers included; Relative: s/it over the first row’s; no row of the 4+2 nodes is limited by its producers, and the last row is a rerun of the 3+1 run of Table 4 (12 iterations). (c) KD-only runs at 8192-token sequences and 4.19M tokens per iteration (4.13M for the streamed 3+1 run), 12 iterations each, of the colocated stack (ModelOpt) and the streamed trainer with CR

-

; GPU-s: GPU-seconds per iteration over all nodes, producers included; Peak: torch-allocated GiB. The trainers of the streamed runs wait at most 0.004 s per iteration for hidden states. (d) Forward-only: an evaluation-mode run of the seed-42 configuration (median of steps 3 to 14), on different nodes, subtracted from the seed-43 training run; the teacher forward is estimated from the forward-only rate of one producer replica of the 3+1 run (32.8k tokens/s, Table 15 ) over 8 replicas, the 4 nodes. CR: the parts of the cache release that were on (Table 27 ).

Runs	Documents read	Shared
Colocated	99.99% of the corpus (9999:1:0), shuffled	0 of 128
Disaggregated	the whole corpus, same shuffle seed	$-$

Appendix

Table 25: Document order of the primary-pair runs. Shared: sequences of iteration 1 whose documents are the same in the colocated and the disaggregated run, reconstructed from the corpus index (seed-42 runs).

	Seed 42	Seed 43	Change	$p$
Disagg. (ours), 4 nodes	68.52	70.09	$+1.57$	$6{\times}10^{-6}$
Disagg. (ours), 3 nodes	70.14	–	–	–
Coloc., Bal-L+K	67.51	69.78	$+2.26$	$9{\times}10^{-13}$
Coloc., KD-only	67.97	69.22	$+1.25$	$10^{-4}$
Between layouts, seed 42
4 nodes $-$ coloc. Bal-L+K	–	–	$+1.01$	0.001

Appendix

Table 26: MMLU (%) of the primary-pair students at iteration 200 for 2 seeds (the 3-node run only for seed 42), with the change over all 14,042 questions and its exact 2-sided McNemar

p

; the last block gives the differences between layouts at seed 42.

CR	s/it	Effect	KD loss
$-$	13.5	–	2.42
mb	68.1	$+54.6$	2.42
iter	14.9	$+1.4$	2.41
mb+iter	69.3	$+55.8$	2.42
producer	13.5	$+0.0$	2.40
all	68.8	$+55.3$	2.41

Appendix

Table 27: Parts of the cache release on Qwen3-32B to 4B (1+1 nodes, 4096-token sequences, 64 microbatches per DP rank), 30 iterations per row; s/it: median of iterations 8 to 30; Effect: s/it minus that of the row with no part on; KD loss: mean over iterations 8 to 30 (2.139 at iteration 1 in every row). CR: the parts of the cache release that were on: mb (the last pipeline stage releases after every microbatch), iter (once at each iteration boundary), producer (each producer replica after each of its first 400 slots), all for the 3, and

-

for none; the release changes the pace, and the KD loss agrees within 0.02 nats (Appendix D.1 , Cache release).

Run (30B student)	tp/pp/cp/ep	edp	Dispatcher	LM, it. 2	Grad norm, it. 1	Layer 0	Outcome
Torch forward	$-$	$-$	$-$	2.16	$-$	$-$	reference
Disaggregated, 235B	2/2/1/8	2	flex	2.11	2.2	$-$	stable, 20 it.
Disaggregated, 235B	2/2/1/4	2	flex	2.11	1.7	$-$	stable, 200 it.
Colocated, 4B	2/2/1/4	4	all-to-all	2.18	29	77.4	finite, 2 it.
Colocated, 4B	1/2/1/8	2	all-to-all	2.18	29	$-$	finite, 2 it.
Colocated, 4B	2/2/1/8	2	all-to-all	2.18	29	$-$	finite, 2 it.

Appendix

Table 28: The Qwen3-30B-A3B student on different grids, 3 to 8 nodes (edp: expert data-parallel size, GPUs divided by ep

\cdot

pp; flex: the flex dispatcher with the DeepEP backend; seq. experts: sequential experts instead of grouped GEMM). Disaggregated: DisaggKD with the Qwen3-235B teacher; colocated: ModelOpt with the Qwen3-4B teacher (dense, same vocabulary; 4096-token sequences, learning rate

10^{-8}

for the 2-iteration cells) or the Qwen3-235B teacher (32768 tokens). LM, it. 2: the student’s LM loss at the second iteration, after 1 optimizer step (at learning rate

10^{-8}

in the 4B-teacher cells); the Torch forward is the Hugging Face loss on the first iteration’s 4096-token batch (the Megatron runs read 2.155 there; the second batch gives 2.178 in the healthy cells), and the 32768-token rows have none. Grad norm, it. 1: gradient norm of the first iteration; Layer 0: L2 norm of the first layer’s output on the first batch (normalized by the data-parallel size); it.: iterations completed. The checked weight sums of the 3 runs with a layer-0 norm match the Hugging Face checkpoint to

5\times 10^{-8}

relative.

Layout	Teacher	Grid	edp	Dispatcher	Outcome
Disagg.	Ultra	2/2/1/4	4	flex	finite, 200 it.
Coloc.	Ultra	2/4/8/16	1	all-to-all	NaN, gradient norm
Coloc.	Ultra	4/4/4/16	1	all-to-all	NaN, forward loss
Coloc.	Ultra	4/4/4/16	1	all-gather	NaN, gradient norm
Coloc.	Nano	4/4/4/16	1	all-gather	NaN, gradient norm
Coloc.	Nano	4/4/1/16	1	all-gather	NaN, gradient norm

Appendix

Table 29: Nemotron-3 Nano student on 8 nodes (Disagg.: DisaggKD on 4+4 nodes) with the Ultra teacher or, for the student alone, the Nano checkpoint as its own teacher. Grid: tp/pp/cp/ep; edp: expert data-parallel size. NaN: the first value that is not finite, in the first iteration; finite: iterations completed. The colocated runs share data, global batch 128, 32768 tokens, and learning rate

2\times 10^{-5}

, except the 2/4/8/16 row (

4\times 10^{-4}

(a) Llama							(b) Qwen3-32B to 4B
Layout (nodes)	Iter.	s/it	GPU-s	Peak	CR	MMLU	Seq. $\times$ batch	Microbatches	s/it	Tokens/s/GPU	CR
Colocated (8)	5	484	31.0k	36.7	–	$-$	$4096\times 128$	64	13.5	2426	$-$
Colocated (16), training	200	301	38.5k	27.5	–	76.98	$4096\times 128$	64	69.6	471	all
pace check	20	305	39.0k	27.4	–	$-$	$8192\times 64$	32	11.9	2755	$-$
Disagg. (ours, 4+2)	5	818	39.3k	68.6	all	$-$	$16384\times 32$	16	13.5	2427	$-$
Disagg. (ours, 8+8), training	400	186	23.8k	50.3	$-$	76.41

Appendix

Table 30: (a) Llama runs at 4.19M tokens per iteration: iterations, median iteration time, GPU-seconds per iteration, and peak memory in GiB (colocated: per GPU; disaggregated: the trainer; torch-allocated). Colocated grids are tp32/pp2/cp1/dp1 (8 nodes) and tp32/pp2/cp2/dp1 (16); disaggregated trainers are tp8/pp2/cp2 with dp1 (4+2) and dp2 (8+8); MMLU (%) is of the trained students, the initial student, and the teacher. The 4+2 median is over iterations 2 to 5. A training run is the full-length run; a pace check is 20 iterations of the same layout, run separately from the training run. (b) Disaggregated Qwen3-32B to 4B at 0.52M tokens per iteration on 2 nodes (tp2/pp2/dp2, 4 producer replicas); Seq.

\times

batch: sequence length times global batch; Microbatches: per DP replica. CR: the parts of the cache release that were on (Table 27 ). A dash marks the colocated runs, which do not use it.

Language Model Distillation

Language model distillation transfers a larger model’s capabilities to smaller or specialized models using outputs, probabilities, or internal representations.

See all 522 papers

Aug 4, 2026

Efficient Knowledge Distillation for LLMs: Offline Top-K Logits and a Fused Chunked KL Loss

Small language models are often the only option for deployment under tight latency, cost, and on-premises constraints, but they are rarely trained from scratch: a compressed model is usually recovered through knowledge distillation (KD). This recovery step largely decides the final quality, yet it is expensive. We present a practitioner's study of how to make distillation training efficient, organised around two systems contributions. First, we show that offline KD (caching the teacher's top-

K

logits once and training the student against the cache) matches online distillation at near-identical training loss while removing the teacher from memory, running about 29% faster per iteration, and reaching up to 41% higher throughput on a single H200 GPU. Second, we introduce a fused, chunked KL loss that never materialises the full vocabulary-sized logit tensor, making peak memory linear in the sequence length. This removes the memory spike that otherwise caps context length and lets us train at four times the context (32{,}768 tokens) on a single GPU. A separate output-head-only toy benchmark isolates the loss kernel and confirms its memory and iteration-rate scaling from 4K to 256K tokens. Together these make large-scale healing and hundreds of ablations affordable. We also report supporting ablations on loss design and sequence packing. We release our chunked-loss implementation: https://github.com/CompactifAI/Full-Chunked-KL-Loss.

Bakbergen Ryskulov, Iker García-Ferrero, David Montero +5

Multiverse Computing

Online Knowledge Distillation Efficient Language Model Training Language Model Distillation Knowledge Distillation

Ask PDF

Jun 26, 2026

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

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

Adrian P. Dieguez, Victor Conchello Vendrell, Alex Batlle +3

Qualcomm AI Research

High-Performance Computing Distributed Training Knowledge Distillation

Ask PDF

Mar 2, 2026

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

Knowledge distillation (KD) is widely used to compress and post-train large language models (LLMs), yet many existing frameworks execute teacher inference with the same training-oriented backend as student optimization, leading to suboptimal efficiency. In this paper, we propose KDFlow, a novel framework for LLM distillation that features a decoupled architecture and employs SGLang for teacher inference. KDFlow combines SGLang for teacher inference with PyTorch FSDP2 for student optimization, allowing each model to run on a backend tailored to its workload. To enable efficient full-vocabulary distillation in this decoupled architecture, KDFlow transfers the teacher's final hidden states via Ray's object store and recomputes teacher logits on each student worker using a frozen copy of the teacher's output head. 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 a 1.44

\times

to 6.36

\times

speedup over MS-SWIFT in off-policy distillation and a 1.43

\times

to 1.75

\times

speedup over verl in on-policy distillation. KDFlow further scales to 64 GPUs, achieving 3.68

\times

and 2.52

\times

strong-scaling speedups in two representative model configurations. The code and documentation are publicly available.

Songming Zhang, Xue Zhang, Tong Zhang +3

School of Computer Science and Technology, Beijing Jiaotong University, Beijing, China · Key Laboratory of Big Data & Artificial Intelligence in Transportation, (Beijing Jiaotong University), Ministry of Education · Weixin AI

Cross-Tokenizer Knowledge Distillation Language Model Distillation Knowledge Distillation

Ask PDF

Abstract

Language Model Distillation

Explore similar work

Efficient Knowledge Distillation for LLMs: Offline Top-K Logits and a Fused Chunked KL Loss

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

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