Source-linked AI summary
DisaggKD: Efficient Knowledge Distillation through Teacher-Student Disaggregation and a Fused Loss
Ruixuan Tu, Yuyang Dong
TL;DR
Colocated knowledge distillation makes teacher weights and full-vocabulary logits compete with the student for per-GPU memory. DisaggKD separates teacher and student nodes and streams hidden states for a fused, vocabulary-sharded loss; on Qwen3.5 it matches colocated LM loss and is faster than same-loss colocated runs, while larger pairs train in the split where fused-loss colocated runs OOM.
Problem
Colocated KD places teacher and student weights under one per-GPU budget and adds a full-vocabulary logits tensor for each student token, limiting which teachers fit.
Source: §1 · p. 1Method
DisaggKD runs teacher replicas forward-only on separate nodes, streams hidden states to the student, and computes a fused KD loss over row windows and vocabulary shards.
Source: §1 · p. 1, §4.1 · p. 4Results
On Qwen3.5, DisaggKD matches colocated LM loss and is faster than same-loss colocated runs; on larger pairs, it trains where fused-loss colocated runs OOM.
Source: §1 · pp. 1–2Takeaways & Limitations
The split is memory efficient when the student spans several nodes, while the fused loss also lets a colocated Qwen3.5 run train on fewer nodes at similar or faster speed.
Source: §1 · p. 2, §5.2 · p. 8The runs are short, use two seeds on one pair, and cover one node type and one English corpus.
Source: Limitations · p. 9
Abstract
from arXiv · showhide
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.
1 Introduction
Colocated knowledge distillation puts teacher and student weights and a vocabulary-wide logits tensor under one per-GPU memory budget, which can cause OOM. DisaggKD separates teacher and student nodes, streams hidden states, and evaluates the KD loss without materializing full-vocabulary logits.
1 Introduction
- Colocated frameworks share one per-GPU budget for both models and add full-vocabulary logits for every student token, limiting which teachers fit.On the Qwen3.5 pair, ModelOpt's colocated loss runs OOM at three nodes where the separated layout fits.Source: §1 · p. 1
- DisaggKD places forward-only teacher replicas on separate nodes and streams their final hidden states to the student, which recovers logits with a copy of the teacher output head.Hidden states are narrower than vocabulary-wide logits, and the student computes the KD loss without forming a full-vocabulary tensor.Source: §1 · p. 1
- The paper models per-GPU memory and intra- and inter-node traffic to estimate where each layout fits and whether the split pays.The memory model is fitted to measured peaks and OOM events, and the traffic model is checked against node counters.Source: §1 · p. 1
- Its fused loss combines per-row statistics across tensor-parallel ranks and computes the KD loss over vocabulary shards, equaling the full-vocabulary loss in exact arithmetic.A prefetch ring paces the teacher producers and student trainer.Source: §1 · p. 1
- The disaggregated Qwen3.5 run matches colocated LM loss and is faster than colocated runs with the same loss; on larger pairs, the split trains where fused-loss colocated runs OOM.The fused loss also lets a colocated Qwen3.5 run train on as few nodes at similar or faster speed.Source: §1 · pp. 1–2
2 Related Work
Prior work studies language-model distillation, colocated KD frameworks, and fused losses; DisaggKD combines teacher-state streaming with a vocabulary-sharded student loss and models for both placements.
2 Related Work
- At LLM scale, forward KL to a teacher's next-token distribution is used alone, with hidden-state losses, or mixed with the LM loss [Muralidharan et al., 2024; Bercovich et al., 2025; Gemma Team, 2024; NVIDIA, 2025c; Mistral Team, 2026; Tang et al., 2026; Peng et al., 2025; Busbridge et al., 2025; Lu and Liu, 2026].Top-K caching stores K logits and indices per token, whereas streaming avoids that cache [Ryskulov et al., 2026].Source: §2 · p. 2
- The three evaluated open KD frameworks keep teacher and student on the same GPUs in the tested modes: Megatron-Bridge ModelOpt, KDFlow, and NeMo-RL [NVIDIA, 2025a, 2024–2026; Zhang et al., 2026; NeMo RL, 2025].KDFlow can optionally offload the teacher between phases.Source: §2 · p. 2
- Like KDFlow and Maestro, DisaggKD sends hidden states rather than logits, but streams them from teacher replicas by NCCL P2P to every student tensor-parallel rank under a prefetch ring [Yuan et al., 2026].It pairs this transfer with a vocabulary-sharded loss and memory and traffic models for both placements.Source: §2 · p. 2
- Prior chunked or fused losses avoid materializing full logits for cross-entropy and offline KD, while DisaggKD separates the KD teacher from the student [Wijmans et al., 2025; Hsu et al., 2025; Ryskulov et al., 2026].The paper situates teacher separation alongside phase separation in LLM serving and flexible placement in RLHF systems [Zhong et al., 2024; Patel et al., 2024; Sheng et al., 2025; Mei et al., 2025].Source: §2 · p. 2
3 Estimation Models
The estimation models select node-minimal legal parallelism grids subject to memory and rate constraints, then estimate per-GPU memory and communication for colocated and disaggregated layouts. On the Qwen3.5 pair, the predictions identify a three-node capacity boundary for the colocated ModelOpt loss while both disaggregated roles fit.
3 Estimation Models
- The models choose the node-minimal legal grid whose memory constraint holds, breaking ties toward lower per-GPU peak; colocated memory couples both models, while disaggregated memory has separate trainer and producer constraints.The models consider per-GPU memory and NVLink and InfiniBand traffic; producer replica count also follows a rate constraint.Source: §3 · pp. 2–3
3.1 Memory
- For the Qwen3.5 pair, logits buffers determine the grid: at tp=2, the loss alone needs 144 GiB at cp=1, so cp>1 is required.In disaggregated training, the trainer’s peak is on the loss stage, which holds the student stack, streamed-loss retention κstr, and backward-pass memory; κstr holds no full-row logits.Source: §3.1 · p. 3
3.2 Data movement
- Traffic estimates account for student and teacher parallelism plus hidden-state transfer; each hidden-state slot goes to every last-stage student rank, and the output head is written once to a file.On Qwen3.5, neither layout is predicted bandwidth-bound; a hidden-state row carries 121× fewer bytes than a logits row, while for Llama the reduction is 8×.Source: §3.2 · pp. 3–4, §3.3 · p. 4
3.3 Predictions at the optimal settings
- At 3 nodes, no colocated grid with tp≤2 is predicted to fit the ModelOpt loss, while both roles of the evaluated disaggregated configuration fit.At 4 nodes, two grids with tp≤2 fit, and the colocated optimum has the lowest predicted peak; the split is the most GPU-efficient disaggregated setting tested.Source: §3.3 · p. 4
- Inter-node traffic estimates are within 2% of counters and NVLink estimates within 15%; the colocated memory model predicts peaks within 3.0 GiB on nine runs outside its fit.The trainer estimate is within 4.0 GiB on 13 runs of fitted families, and the producer estimate within 8.4 GiB on its eight fitting runs.Source: §3.3 · p. 4
4 The Disaggregated KD Algorithm
DisaggKD puts teacher replicas on producer nodes and streams hidden states to student trainers. Its windowed, vocabulary-sharded loss is exact in arithmetic, while a prefetch ring paces production and training.
4.1 Roles and hidden-state transfer
- Teacher replicas run forward-only on producer nodes and send final hidden-state slots to student ranks, which use the teacher’s output head to compute logits.The teacher’s lm_head is read by trainers rather than used in the producer’s forward path.Source: §4.1 · pp. 4–5
4.2 Windowed, fused loss over vocabulary shards
- The fused loss scans row windows over vocabulary shards and combines per-row statistics, yielding the full-vocabulary loss and gradient in exact arithmetic.The loss is independent of window size and how entries are grouped for scanning, within the theorem’s stated divisibility conditions.Source: §4.2 · pp. 5, 7
- Each rank scans logits once using a running maximum and rescaled sums, then TP ranks combine statistics with MAX and SUM all-reduces.This avoids a separate maximum-finding pass and avoids forming a full-vocabulary logits matrix.Source: §4.2 · pp. 6–7
4.3 Prefetch and pacing
- A prefetch ring holds up to four slots ahead, and iteration time is determined by whichever takes longer: trainer compute or producer supply.When the trainer outruns its replicas, the iteration is producer-bound.Source: §4.3 · p. 7
5 Experiments
Across the Qwen3.5 pair, the disaggregated run is faster than colocated runs while matching their LM loss closely, and the fused loss helps colocated runs use fewer nodes. Results on other pairs vary: some colocated configurations run OOM or fail, while an 8-node Qwen3-235B colocated configuration is faster than the split.
5.1 Setup
- The evaluation uses five open-weight, same-tokenizer pairs; Qwen3.5-35B-A3B to 9B is primary, and Qwen3-32B to 4B supports comparison with other stacks.Runs use 8-GPU H100 nodes and FineWeb-Edu; the primary pair is tested at several node counts.Source: §5.1 · p. 7
5.2 Results on the Qwen3.5 pair
- The 4-node disaggregated run’s 17-iteration LM-loss mean differs by 0.008 nats from SQ-L+K and 0.012 from Bal-L+K, while the run is faster than the colocated baselines.Runs used different document orders; their logged KD losses also differ by the teacher’s entropy because they use cross-entropy versus KL.Source: Figure 3 · p. 8, §5.2 · p. 8
- With fused-W, the colocated run is similarly fast at 3 nodes and faster at 2 nodes; the disaggregated split holds less memory per GPU at 3 nodes.At 2 nodes the split is 1+1; fused-W turns off the model’s cross-entropy fusion.Source: §5.2 · p. 8
- With seed 42, the 4-node disaggregated student scores 1.0 point above Bal-L+K and 0.6 above KD-only; seed-43 gains put layout differences within run-to-run differences.With seed 43, each student with a run gains 1.3 to 2.3 points.Source: §5.2 · p. 8
5.3 Other pairs
- On Qwen3-32B, splitting a pair that fits on one node doubles GPU count and halves rate per GPU at equal step time; KDFlow’s remaining 1.35× lead comes from step time.DisaggKD’s trainer-limited step is 13.5 s, versus 9.8 s for KDFlow and 18.4 s for ModelOpt.Source: §5.3 · p. 8
- For Qwen3-235B, the 8-node all-gather colocated run trains faster than DisaggKD: 75 s with Bal-L+K versus 109 s.Three colocated grids that fit memory used an all-to-all dispatcher and produced NaNs or diverged; the all-gather configuration trained.Source: §5.3 · p. 8
- The colocated stack does not train Nemotron: the teacher runs OOM at 4 nodes, while 8-node grids that fit memory produce NaNs.The reported NaNs arise in the forward loss with all-to-all dispatch and in the gradient norm with either dispatcher.Source: §5.3 · p. 8
- At 16 nodes, disaggregated Llama with TP 8 is faster than colocated Llama with TP 32.The memory model predicts every colocated TP≤8 grid exceeds budget.Source: §5.3 · p. 9
6 Conclusion
Memory and traffic models locate the layouts’ memory limits, while the measured results distinguish where disaggregation helps from where a colocated fused loss suffices.
6 Conclusion
- On Qwen3.5, DisaggKD matches colocated LM loss and is faster than colocated runs with the same loss; on larger pairs, it trains where colocated fused-loss runs OOM.The fused loss also lets a colocated run train on as few nodes at similar or faster speed, while disaggregation is memory efficient when the student spans several nodes.Source: §6 · p. 9
Limitations
The evaluation is bounded by its tested hardware, data, and run scope, and by comparisons that do not align every framework’s recipe.
Limitations
- Evidence is limited to short runs, two seeds on one pair, one node type, and one English corpus; framework recipes are not aligned.The models are fitted to tested parallelism ranges and sequence lengths, with some terms based on few runs.Source: Limitations · p. 9
A.1 Memory Model
The memory model accounts for per-module memory in colocated and disaggregated layouts, then checks its estimates against measured peaks and OOM events. Its validation is strong overall, while held-out trainer memory varies across student families.
A.1 Memory Model
- The colocated loss is modeled as 9.5 logits buffers, plus 1 buffer for LM loss; the streamed loss instead retains hidden states rather than full-row logits.For the Qwen3.5 pair at tp=8, streamed loss reduces the held elements by 40×; for dense Llama, by 1.3×.Source: §A.1 · p. 11
- The model accounts for trainer parameters and state, activations and temporary buffers, and a fitted remainder; producer memory includes teacher weights, forward workspace, and NCCL buffers.Table 5 distinguishes coefficients derived from code or framework behavior, fitted to measurements, and treated as allowances.Source: §A.1 · pp. 11, 13, Table 5 · p. 12
- 17 of 18 colocated configurations outside the fit and 5 of 6 that set constants are classified correctly; 9 outside-fit peaks are predicted within 3.0 GiB.The Nemotron 4-node run is an exception: it OOMs during the teacher’s first forward pass, where measured demand is 75.7 GiB against 59.9 GiB predicted.Source: §A.1 · p. 13, Table 6 · p. 13
- For the Llama 405B-to-70B pair, the lowest-peak legal grids at 8, 16, 24, 28, and 32 nodes are predicted at 89 to 102 GiB, over budget; with tp=32, the layout fits on 8 and 16 nodes.Those tp=32 peaks are predicted within 1.2 GiB of torch-allocated peaks, without Llama colocated measurements in the fit.Source: §A.1 · p. 13
- Across 8 trainers from 4 student families, family and model dimensions cannot be separated; held-out Qwen3.5 trainers at cp=4 are among the most underpredicted.They hold more memory outside PyTorch than at cp=2 and are the 3 most underpredicted of 7 held-out trainers.Source: §A.1 · p. 15
A.2 Traffic Model
The traffic model estimates movement within the student and teacher and between them, then compares predictions with hardware counters. The streaming pipeline transfers teacher hidden states and paces producer and trainer, while the teacher-precision tests show a memory–loss trade-off for fp8 experts.
A.2 Traffic Model
- Traffic estimates include student parallelism, the teacher’s forward pass, and teacher-to-student transfer; the Qwen3.5 accounting includes attention, linear-attention, and expert collectives.The producer’s combine reduce-scatter alone carries 57% of calculated collective bytes.Source: §A.2 · pp. 15, 17, Table 13 · p. 17
- Every inter-node counter is within 1.4% of calculation on the colocated run and within 2% on the 3-node split; at 4 nodes, calculation stays within 2% and below every counter.The 4-node DP×CP rings span nodes, and their gradient and parameter traffic is 1.0 to 2.0 times the pipeline’s.Source: §A.2 · pp. 16–17, Table 14 · p. 17
B.1 Streaming Pipeline
- The prefetch ring transfers teacher hidden states to student receiver groups in bf16, with slot headers and FIFO checks preserving ordering.The producer runs at most δ iterations ahead, and the ring holds up to 4 slots per last-stage rank by default.Source: §B.1 · p. 18
- Duty compares producer-only iteration time with measured iteration time: a value of 1 indicates producers limit or match pace, while a value below 1 indicates the trainer limits.The Qwen3.5 (2+1) run is trainer-limited; the Qwen3-235B and Nemotron runs are producer-limited.Source: §B.1 · p. 18, Table 15 · p. 19
B.2 Teacher Precision
- 5.8 GiB less peak allocated memory with fp8 routed experts came with KD loss 0.505 to 0.665, versus 0.249 to 0.290 with bf16 from iterations 4 to 12.The authors note that the two seconds-per-iteration values do not compare precisions.Source: §B.2 · p. 20, Table 16 · p. 20
B.3 Proof of Exactness
The proof merges per-shard and per-tile statistics to recover the full-vocabulary loss and its gradients independently of the sharding and windowing. This is an identity in exact arithmetic, with the loss mask and normalization outside its scope.
B.3 Proof of Exactness
- The merge rule takes the largest part maximum and rescales part sums before combining them into statistics for the union of vocabulary columns.The same computation incorporates the teacher-weighted sum.Source: §B.3 · p. 20
- Scanning columns in tiles yields the statistics of their union, and induction applies this result to the successive tiles.The proof starts from the empty-prefix state and merges each tile with the accumulated state.Source: §B.3 · p. 20
- The all-reduced shard statistics recover the statistics of all V vocabulary columns, establishing the full-vocabulary quantities used by the loss.The proof applies the merge lemma to each tensor-parallel shard and combines the results.Source: §B.3 · pp. 20–21
- Partitioning rows across context-parallel ranks and windows only regroups a finite sum, so the loss and gradients do not depend on C or sharding.This identity is for the loss as a function of hidden states and head weights; the loss mask and normalization lie outside it.Source: §B.3 · p. 21
- The identity is exact-arithmetic only; it does not equate floating-point execution details or training recipes.The remark excludes floating-point summation order, specified fp8/bf16 operations, and differences in loss weights or evaluated columns.Source: §B.3 · p. 21
C Experimental Setup
The experiments search legal parallelism grids under memory constraints for colocated and disaggregated layouts, then evaluate disaggregated runs using specified training and measurement procedures. The search excludes fully sharded data parallelism, and steady-state trainer token rates stay close to producer rates.
C Experimental Setup
- The search selects node-minimal legal grids under memory constraints, with shared parallelism choices for colocated models and separate trainer and producer grids for disaggregated runs.The search spans tensor, pipeline, context, expert, and data parallelism; sequence lengths are 32,768 tokens unless stated otherwise.Source: C Experimental Setup · p. 21
- Fully sharded data parallelism is excluded; the Megatron-FSDP release errors at pipeline parallelism 2 and 4 on MoE students, and pp=1 was not run.The reported FSDP memory comparisons are constrained by those setup and coverage boundaries.Source: C Experimental Setup · pp. 21–22
C.1 Disaggregated Runs
- Runs use FineWeb-Edu, full activation recomputation, and bf16 students with fp32 master weights, with stated pair-specific sequence lengths and batch sizes.The stated recipe also specifies AdamW, learning-rate decay, and, for SQ-L+K runs, a decaying KD weight.Source: §C.1 · p. 22
- Each disaggregated run's trainer token rate is within 4% of the producers' mean rate in steady state.Iteration-time medians use iterations 3 through the last, with stated exceptions for particular runs.Source: §C.1 · p. 22
- MMLU is evaluated with fixed 4-shot prompts, greedy decoding, and exact match after normalization and lowercasing.The evaluation uses the public test set and FlexEval; decoding stops at a newline or 256 new tokens.Source: §C.1 · p. 22
C.2 Compared Stacks
The compared-stack section describes colocated fused-loss variants and measured stack comparisons, alongside configuration and run-reporting details. Its account notes that loss fusion removes a memory limit but does not eliminate all colocated memory constraints.
C.2 Compared Stacks
- Table 17 reports disaggregated node allocation, grid, iterations, median iteration time, throughput, KD loss, and peak memory.Its caption distinguishes trainer torch-allocated memory from producer device memory.Source: Table 17 · p. 23
- Figure 4 plots per-iteration and 17-iteration-average KD-loss trajectories; colors distinguish the smallest disaggregated run from the colocated-node-count run.Students start from released checkpoints, so loss levels differ across pairs.Source: Figure 4 · p. 23
- The fused colocated loss removes the loss-side memory limit, but a 3-node fused run still OOMs on a 15.16 GiB fp32 cross-entropy tile.Unlike DisaggKD's windowed logits, the colocated criterion reads materialized logits; DisaggKD computes LM and KD terms on row windows.Source: §C.2 · pp. 23, 25
- NeMo-RL's Qwen3.5 logits buffer is 30 GiB at 1.05M tokens and 121 GiB at 4.19M; the latter does not fit an 80 GiB GPU beside the models.The 1.05M-token step was the largest NeMo-RL Qwen3.5 run that completed; two 4.19M-token runs failed.Source: §C.2 · p. 25
- Table 19 reports measured stack runs, but its runs differ in loss, sequence length, steps, tokens per step, and teacher precision.The table caption identifies tokens per GPU, step time, and peak memory among its reported measures, so comparisons should account for those differing settings.Source: §C.2 · pp. 25–26, Table 19 · p. 26
D.1 Primary Pair
On the primary Qwen3.5 pair, the disaggregated run’s LM loss closely matches the colocated run with the same loss, while speed and memory trade-offs depend on node count and layout. Fused-W reduces colocated memory use, but colocated runs for the larger pairs still fail during model construction.
Loss parity
- 0.008 nats is the 4-node disaggregated run’s mean absolute LM-loss gap from colocated SQ-L+K over 17-iteration moving averages.The signed mean gap is +0.005 and the largest gap is 0.028; the rescaled run differs in configuration and data order.Source: Table 20 · p. 27, §D.1 · p. 27
Layout comparison
- At 2 nodes, colocated fused-W is faster and uses less memory per GPU than DisaggKD; at 3 nodes, speed is similar and DisaggKD uses less memory per GPU.At 2 nodes, the fused-W run takes 146.8 s per iteration versus 283.7 s for DisaggKD.Source: §D.1 · p. 28, Table 23 · p. 29
Larger pairs
- Colocated fused-W runs for Qwen3-235B, Nemotron, and Llama run out of memory during model construction at the tested node counts and student grids.The reported failures occur before the first forward pass; DisaggKD runs at those configurations.Source: §D.1 · p. 28
Fused loss in the colocated stack
- At 2 nodes, fused-W removes the KD loss’s large fp32 intermediates, while context parallelism of 2 halves the relevant buffer sizes and lets the run fit.With cp1/dp2, a 0.5 GiB request fails; with cp2/dp1, the run trains at a sampled device peak of 65.8 GiB.Source: §D.1 · p. 28
Source of the speedup
- The 4-node speed gap over colocated Bal-L+K includes 2,762 GPU-seconds from its LM term and 657 GPU-seconds of excess over disaggregated KD-only.The colocated LM term adds collectives, copies, conversions, and stream synchronizations on the loss stage; disaggregation overlaps teacher forward with student training.Source: §D.1 · p. 29
Evaluation caveats
- Two seeds do not establish a ranking of layouts, and the parity statistic compares runs on the same corpus rather than the same samples.Document sequences differ from the first iteration, and score shifts across seeds exceed the seed-42 lead of the 4-node disaggregated student.Source: §D.1 · pp. 29, 31
D.2 Other Pairs
The other-pair evaluation reports both training-stability problems in colocated runs and settings where disaggregation trains or runs faster. Results vary by model and grid, with a Llama speedup, Qwen divergence controls, and Nemotron NaNs among the reported findings.
D.2 Other Pairs
- Replacing all-to-all with all-gather at the tp2/pp1/ep16 grid gives the ep-4 cell’s loss of 2.178 and gradient norm of 29, and avoids the wrong forward pass.On the Qwen3-235B pair, the all-gather run has second-iteration loss 2.11 and first-iteration gradient norm 12, versus 1.0 to 1.7 × 10^6 with all-to-all.Source: §D.2 · p. 31
- Disaggregated Qwen3-235B runs complete 200 iterations on both 4+4 and 2+1 nodes without NaNs; a KD-only run also completes 60 iterations at 105 seconds per iteration.The KD-only run skips the LM term and uses a peak learning rate of 4 × 10^-4.Source: §D.2 · p. 32
- All five colocated Nemotron ep-16 cells produce NaN in the first iteration, while the Nano-teacher ep-4 cell remains finite for three iterations.The ep-16 cells use edp 1; Table 29 describes the compared colocated and disaggregated runs.Source: §D.2 · p. 32, Table 29 · p. 32
- At 16 nodes, the Llama disaggregated layout is 1.6× faster than colocated in training runs and 1.3× faster in pace checks.Table 30a reports median iteration time, GPU-seconds per iteration, and peak memory for the Llama runs.Source: §D.2 · p. 33, Table 30 · p. 33
- On Llama, disaggregation runs on 6 nodes where colocated needs 8; at 8 nodes, colocated uses 31.0k GPU-seconds per iteration versus 23.8k for the disaggregated training run.Both trained students score within 1.8 points of the initial student on MMLU.Source: §D.2 · p. 33, Table 30 · p. 33
- On Qwen3-32B to 4B, disaggregation with CR − is 5.1 to 5.2× faster per GPU than with CR all, and remains below ModelOpt and KDFlow.At fixed 0.52M tokens per step, the CR − iteration times are 13.5, 11.9, and 13.5 seconds for 64, 32, and 16 microbatches per replica.Source: §D.2 · p. 33
E Artifacts, Licenses, and Compute
The paper identifies the licenses for its models, corpus, benchmark, and software, and describes an anonymized codebase intended to reproduce measured tables by rerunning configurations. Its glossary organizes notation for model sizes, parallelism, memory, traffic, and the algorithm’s loss and stream.
E Artifacts, Licenses, and Compute
- The listed models, FineWeb-Edu corpus, MMLU benchmark, and software are used under the stated licenses; the trained students are not released.The paper says FineWeb-Edu is a filtered web corpus and MMLU a multiple-choice benchmark, neither annotated by the authors.Source: E Artifacts, Licenses, and Compute · p. 33
- An anonymized codebase includes configurations, model tests, table and figure scripts, and an MMLU protocol, but no run data, so measured tables require rerunning configurations on the stated hardware.Reported runs used about 15,000 H100 80 GB GPU-hours, excluding failed attempts and memory and NaN probes.Source: E Artifacts, Licenses, and Compute · p. 33
F Glossary of Symbols
- The glossary defines model-size and architecture notation, including dense and routed-expert parameters, layers, widths, vocabulary, attention heads, and routed experts.It also names the tensor-, pipeline-, context-, data-, expert-, and expert-tensor-parallel degrees.Source: F Glossary of Symbols · pp. 33–34
- Parallelism and layout symbols specify legal grids, node counts, trainer and producer nodes, memory budget, token counts, microbatches, and interconnect span.The glossary defines the per-GPU memory budget as 79.11 GiB and gives the node size as eight GPUs.Source: F Glossary of Symbols · p. 34
- Memory and traffic notation distinguishes colocated and streamed memory terms, producer and trainer peaks, iteration time, and bytes sent by communication operations.The glossary defines memory terms including student state, teacher weights, activation memory, logits buffers, and streamed-loss memory.Source: F Glossary of Symbols · p. 34
- Algorithm and loss notation covers stream slots and windows, hidden states and logits, softmax and label terms, loss components, and shard-wise scan statistics.It also defines the gradient of the loss with respect to a window’s student logits on a tensor-parallel rank.Source: F Glossary of Symbols · p. 34
Related papers
This paper cites
- HybridFlow: A Flexible and Efficient RLHF Framework2024
- Sequence-Level Knowledge Distillation2016
- Nemotron 3 Ultra: Open, Efficient Mixture-of-Experts Hybrid Mamba-Transformer Model for Agentic Reasoning2026
- Efficient Knowledge Distillation for LLMs: Offline Top-K Logits and a Fused Chunked KL Loss2026