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

DisaggKD: Efficient Knowledge Distillation through Teacher–Student Disaggregation and a Fused Loss Ruixuan Tu* The University of Tokyo, RIKEN AIP turx2003@gmail.com

Yuyang Dong SB Intuitions Corp. yuyang.dong@sbintuitions.co.jp

arXiv:2610.17644v1 [cs.LG] 13 Oct 2026

Abstract

on the Qwen3.5 pair at 3 nodes runs out of memory (OOM) where the separated one fits (Figure 1), and our fused loss moves it but not the weights of larger teachers (Table 3). We separate the 2 roles: the teacher runs forwardonly on its own nodes, so no GPU holds the layers of both models. The teacher sends its final hidden states, which are narrower than its vocabulary-wide logits, and the student recovers the logits with a copy of the teacher’s output head. A fused kernel then evaluates the KD loss over vocabulary shards from per-row statistics that combine across tensor-parallel (TP) ranks, without forming a fullvocabulary tensor, so trainer memory stays low. Whether the split pays depends on per-GPU memory and on NVLink and InfiniBand traffic, so we derive closed-form estimates of both. Our contributions are:

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 fullvocabulary 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. Code will be released upon acceptance.

1

• Estimation models (§3) of per-GPU memory (§3.1) and of traffic inside the student, inside the teacher, and between them (§3.2), for both layouts (§3.3); the memory model is fitted on measured peaks and OOM events, the traffic model checked against node counters (Appendices A.1, A.2).

Introduction

Knowledge distillation (KD) trains a student on its teacher’s output distribution (Hinton et al., 2015), and at pretraining scale SlimQwen uses it to recover the quality of compressed students (Tang et al., 2026). The 3 open frameworks we ran place the teacher on the student’s GPUs, keeping its weights beside the student’s weights, gradients, and optimizer state (NVIDIA, 2025a, 2024–2026) or offloading it between phases (NeMo RL, 2025; Zhang et al., 2026), as KDFlow optionally does. Colocated, the 2 models share 1 per-GPU memory budget, and the KD loss adds a logits tensor over the full vocabulary for every token the student processes. That budget is the limit: with ModelOpt’s loss the colocated stack

• A disaggregated KD algorithm, DisaggKD (§4): teacher replicas stream hidden states over NVIDIA Collective Communications Library (NCCL) point-to-point (P2P) links (§4.1), the student evaluates a fused loss over row windows and vocabulary shards, equal to the full-vocabulary loss in exact arithmetic (§4.2, Appendix B.3), and a prefetch ring paces both (§4.3). • A measured evaluation (§5). On the Qwen3.5 pair the language-modeling (LM) losses of the 2 layouts match, the disaggregated run is faster than the colocated runs with ModelOpt’s loss balancer and with the same loss, and a colocated run with our fused loss is as fast or faster on

* Work was done during the internship at SB Intuitions Corp.

1

S κ

A

teacher

other

measured

and NeMo-RL (NeMo RL, 2025). Like KDFlow and Maestro (Yuan et al., 2026), we send hidden states instead of logits, but stream the states from teacher replicas by NCCL P2P to every tensorparallel rank of the student under a prefetch ring, with a vocabulary-sharded loss and memory and traffic models for both placements (§3, §4).

GiB per GPU

160 (q)

Ht

M

79

0

ModelOpt

fused-W

colocated

producer

Disaggregation and fused losses. Separating phases that contend for one device has paid off in LLM serving, with prefill and decode on separate pools (Zhong et al., 2024; Patel et al., 2024; Qin et al., 2025). RLHF systems place the actor, reference, and reward models flexibly (Sheng et al., 2025) and search plans with a cost estimator (Mei et al., 2025); we separate the teacher of KD. On the student side, chunked and fused losses avoid materializing the full logits tensor in cross-entropy (Wijmans et al., 2025; Hsu et al., 2025) and in offline KD (Ryskulov et al., 2026).

trainer

disaggregated

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 κ 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 κ the remainder. Trainer and producer nodes hold separate stacks and exchange (q) hidden-state slots Ht over NCCL P2P.

3

Whether a teacher–student pair can be distilled on a given set of nodes depends on 3 per-GPU quantities: peak memory, bytes moved over the intra-node fabric (NVLink), and bytes moved over the inter-node fabric (InfiniBand). We model these quantities for each layout at its optimal setting.

as few nodes but holds more memory per GPU when the student spans several nodes (§5.2); on 3 larger pairs the disaggregated layout trains where the colocated stack runs OOM, also with a fused loss, and is faster on Llama.

2

Estimation Models

Notation. The teacher has Dt dense and Et routed-expert parameters (Pt = Dt +Et ), ℓt layers, and hidden width dt . The student has Ds , Es , ℓs , and ds . Both share the vocabulary size V . A layout assigns the student tensor (TP), pipeline (PP), context (CP), data (DP), and expert (EP) parallel degrees (tp, pp, cp, dp, ep) on N nodes of 8 GPUs, each with budget M =79.11 GiB. A sequence has L tokens and an iteration processes gbs sequences, so Ntok = gbs · L, each GPU handles τ = Ntok /(dp cp) tokens, and each DP replica runs nmb = gbs/dp microbatches of one sequence. Activations take b=2 bytes per element, so x elements occupy x b bytes, and Vloc = V /tp is the vocabulary shard of one TP rank (Appendix F).

Related Work

KD for language models. At LLM scale the usual objective is the forward KL to the teacher’s next-token distribution, used alone or with hiddenstate losses in compression pipelines and model releases (Muralidharan et al., 2024; Bercovich et al., 2025; Gemma Team, 2024; NVIDIA, 2025c; Mistral Team, 2026) or studied as a mixture with the LM loss (Tang et al., 2026; Peng et al., 2025; Busbridge et al., 2025; Lu and Liu, 2026). Caching the teacher’s top-K logits stores K values and indices per token (Ryskulov et al., 2026); streaming avoids it.

Optimal setting. For a layout, let G be its legal parallelism grids (degrees divide layer, head, expert, and sequence counts and the world size). The optimal setting is the node-minimal grid g ⋆ = arg ming∈G N (g) whose memory constraint holds, ties broken toward the smaller per-GPU peak. In the colocated layout the constraint couples both

KD frameworks. Table 2 compares the open frameworks we ran against DisaggKD; all 3 keep teacher and student on the same GPUs in the modes we ran: Megatron-Bridge ModelOpt (NVIDIA, 2025a, 2024–2026), KDFlow (Zhang et al., 2026), 2

Disaggregated. The 2 roles have separate budgets. The trainer’s peak is on the loss stage, which holds the student stack (1), the streamed-loss retention κstr of §4.2, and backward-pass memory:

models to the student grid; in the disaggregated layout it splits into trainer and producer inequalities, and R follows from the rate constraint. 3.1

Memory

Mtr (g) = S + κstr + Astr + σstr ≤ M,

A student rank stores bf16 weights, fp32 gradients, and a distributed AdamW state, which is sharded over the DP×CP group of dp · cp ranks for dense parameters and over the expert data-parallel group of edp = dp cp tp/ep ranks for routed experts:

with κstr of (12), which holds no full-row logits, the activation and temporary memory Astr (reference-cycle garbage, loss-backward and recompute buffers, and an expert-layer block), and the fitted remainder σstr , the mean of 8 trainers (Table 5). The producer holds the teacher’s checkpoint as a shard (routed experts in fp8 or bf16, the rest bf16 but Nemotron’s shared experts; Appendix B.2), a fitted remainder φp (forward-pass workspace, CUDA context, allocator), and NCCL buffers ωp :   Dt t Mpr (tpp , epp ) = 2 tp bytes + βe tp Eep p p p (5) + φp + ωp ≤ M,

S(g) = (2+4) Ploc bytes   Es s + 12 tp ppDdp + cp pp ep edp bytes, (1) s + ppEsep , Ploc = tpDpp

cp tp edp = dp ep

where Ploc counts the parameters of one rank (for the trainer, of the last pipeline stage, which holds the loss; the formula gives the stage average and Appendix A.1 adds the output layer V ds /tp) and routed experts are sharded over pp · ep only; coefficients are green when fixed by the code, red when fitted to measurements, and black when set as an allowance (Table 5).

where βe is 1 for fp8 routed experts and 2 for bf16 and φp is calibrated (Table 5). The producers keep pace with the trainer when R ρrep ≥ Ntok /Titer (ρrep : replica forward rate, Titer : iteration time).

Colocated. The colocated teacher is built at the student’s grid
in bf16, so its shard is Tcol = t t 2 ppDtp + ppEep bytes per rank. The ModelOpt logits-KL loss runs eagerly at the last pipeline stage. At its high-water mark the loss holds 7 fp32equivalent logits-sized buffers ∈ RL/cp×V /tp (enumerated in Appendix A.1), 1 for the loss-averaging all-reduce, and 1.5 for the backward pass: κcol (g) = (7+1+1.5) ·

L V · · 4 bytes. cp tp

(4)

3.2

Data movement

A GPU moves bytes for 3 reasons: the parallelism of the student, that of the teacher’s forward pass (on the student’s grid when colocated, on the producer’s span when disaggregated), and the teacherto-student transfer. For a model x ∈ {t, s} with hidden width dx , TP all-reduces (Shoeybi et al., 2020), pipeline activation transfers, and the student’s gradient and parameter collectives send per GPU and iteration

(2)

The peak at the last pipeline rank is

ℓx 2(tp−1) Φtp τ dx b bytes, x = νx pp tp

Mcol (g) = S +Tcol +κcol +Acol +σcol ≤ M, (3)

L dx b Φpp x = nmb cp tp bytes,

ℓs L where Acol = 1 pp cp tp ds b bytes holds the stored

(6)

dp cp−1 Φdp s = dp cp Ploc (4+2) bytes,

layer inputs of the 1 microbatch in flight under full activation recomputation (mixture-of-experts (MoE) students add the dispatch term of Table 5), and σcol =12.27 GiB is the device peak minus the torch-allocated peak, the mean of the 2 cp4 runs (Table 5). The peak is the larger of this loss point and the teacher’s forward-pass point. For the Qwen3.5 pair the logits buffers decide the grid: at tp 2 the loss alone needs 144 GiB at cp=1, so cp > 1 is required. The Bal-L+K recipe (ModelOpt’s balancer) holds 1 more.

where νt =2 and νs =6 are 2 all-reduces per layer times the teacher’s 1 forward pass and the student’s 3 passes (forward, recompute, backward), and Φpp x is per stage boundary and direction. CP exchanges and EP all-to-all are in Table 13. Megatron’s rank order (tp, cp, ep, dp, pp) with 8 GPUs per node decides the link: a dimension of degree n whose ranks are ∆ apart spans n∆ ranks, so the dimension stays on NVLink if n∆ ≤ 8 (one node) and otherwise uses InfiniBand at its node boundaries, 3

Memory per GPU (GiB)

Colocated (4 nodes, tp2/pp4/cp4/dp1)

Disaggregated (3 nodes: 2 trainer + 1 producer)

S+Tcol +κcol +Acol +σcol ≤ M 11.5 + 15.8 + 36.0 + 0.25 + 12.3 ≈ predicted 75.8 (fit), measured 76.1

trainer S+κstr +Astr +σstr ≤ M : predicted 55.8 (fit), measured 57.1
Dt Et + tp +φp +ωp ≤ M : predicted producer 2 tp p p 41.8 (fit), measured 39.4

GB per GPU (s at the limit)

Expression

Colocated

Disaggregated

Student, NVLink Teacher, NVLink Student, InfiniBand Teacher, InfiniBand Transfer, InfiniBand NVLink, measured

cp dp Φtp s + Φs + Φs tp cp Φt + Φt + Φep t Φpp s per direction Φpp t per direction Φxfer per last-stage rank all dimensions

701.7 (1.56 s) 191.0 (0.42 s) 4.29 (0.09 s) 2.15 (0.04 s) 0 932.1 (2.07 s)

1179.5 (2.62 s) 2405 (5.34 s), producer 4.29 (0.09 s) 0 4.29 (0.09 s), 34.4 in total 1360 (3.02 s), trainer; 2485 (5.52 s), producer 7.23 calculated, 7.35 measured

Last-stage node, Gbps received

4.56 calculated, 4.61 measured

Table 1: Estimation models per layout at its optimal setting for the Qwen3.5-35B-A3B to 9B pair (L=32768, 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).

so each dimension that spans nodes adds inter-node traffic. The transfer sends each slot once to every last-stage rank of the receiving group, L Φxfer = nmb cp dt b bytes

the traffic accuracy on the Qwen3.5 pair (Appendix A.1). The colocated model classifies 17 of the 18 configurations outside its fit and 5 of the 6 that set constants correctly, and predicts the peaks of 9 runs outside its fit within 3.0 GiB; the 2 errors are the Qwen3.5 Bal-L+K fit row, at the budget, and the Nemotron run on 4 nodes (OOM in the teacher’s forward pass). The trainer estimate is within 4.0 GiB on 13 runs of the fitted families and the producer estimate within 8.4 GiB on its 8 fitting runs. Inter-node traffic is within 2% of the counters (Table 14), NVLink within 15% (Table 6).

(7)

per last-stage rank, or gbs · tp · L dt b bytes in total; the lm_head (dt V b bytes) is written once to a file. 3.3

Predictions at the optimal settings

Table 1 evaluates the models for the Qwen3.5-35BA3B to 9B pair. At 4 nodes 2 grids with tp≤2 fit the budget and the colocated optimum has the lowest predicted peak; at 3 nodes no tp≤2 grid is predicted under the budget for ModelOpt’s loss. The disaggregated configuration uses the (2+1) split of Table 17, the most GPU-efficient split we ran (Table 23), and both roles fit. With ModelOpt’s loss the colocated layout is capacity-bound: its logits buffers force cp ≥ 4 and hence dp=1 on 4 nodes at tp≤2. On the Qwen3.5 pair neither layout is bandwidth-bound: at the link limits the busiest GPU needs under 6 seconds of NVLink and 0.2 seconds of InfiniBand per iteration, and a hidden-state row carries 2dt bytes where a logits row carries 2V , 8× fewer for the Llama pair and 121× fewer for the Qwen3.5 pair.

4

The Disaggregated KD Algorithm

4.1

Roles and hidden-state transfer

DisaggKD is built on Megatron-Bridge (NVIDIA, 2025a) and splits the 2 models across nodes: no GPU or node hosts both models’ layers (Figure 2). The trainer is a standard training run on Nstu nodes with grid (tp, pp, cp, dp, ep). The producer runs R replicas of the teacher on Nprod nodes, forward only, each spanning tpp GPUs (times an EP degree epp ) chosen by (5) independently of the trainer. For every microbatch the producer emits the teacher’s final hidden states instead of its logits. CP divides a sequence’s rows, one per token, over the cp ranks of the student. The slot of cp rank q (q) is the matrix Ht ∈ Rr×dt of its r rows, so a sequence has L = cp · r rows. The tp × cp GPUs of the student’s last pipeline stage in one DP replica,

Validation. Table 9 compares the models with measurements on the Qwen3.5 pair; Table 6 summarizes the memory accuracy on all 5 pairs and 4

tp rank

tp−1

j

0

0 cp rank

(q) Ht

producer

q

slots (columns), windows (bars) 0

(q,j)

Zx

k K−1

cp−1 (a) Window k

(b) Row i of cell (q, j) fused kernel (q,j)

(q)

Ht,k

×

(j)

Wt

=

Zt,k

(j)

zt

MAX all-reduce mx (16)

mx,u , ax,u , cu (15) (j)

mt,u

zt

u=Vloc −1

u=0 (q,j)

(q) Hs,k

×

tp rank j

(j)

Ws

=

Zs,k

(j) zs

teacher

rescale, SUM all-reduce ax , c (16)

(j)

zs

ms,u u=Vloc −1

u=0

running max

row i

student

row losses KDi , CEi (11)

bold cell

Figure 2: The disaggregated algorithm (symbols as in §4.1 and §4.2; · · · : omitted ranks). Top: the producer sends (q) (q,j) slot Ht to every tp rank of its cp rank; cell (q, j) of the receiver group (bold) holds the logits block Zx , and its (q) (j) (q,j) prefetch ring holds up to 4 slots of K windows. (a) Window k: Hx,k Wx gives Zx,k ; 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.

4.2

the last-stage ranks, form a receiver group, and the (q) producer sends Ht to every tp rank of cp rank q. We call the last-stage rank with cp rank q and tp rank j the cell (q, j) (Figure 2, top).

Windowed, fused loss over vocabulary shards

The KD objective is a sum over rows, and within a row the coupling along the vocabulary axis is the reductions of (10). The trainer therefore evaluates a slot in K windows k = 0, . . . , K−1 of C consecutive rows each (Algorithm 2, Appendix B.1), so r = K · C (C does not change the loss in exact (q) arithmetic). Let Hx,k ∈ RC×dx be window k of

The teacher’s lm_head Wt ∈ Rdt ×V is not part of the producer’s forward path. The producer writes this matrix once to a file that the trainers read. The vocabulary is sharded over the tp ranks: tp rank (j) j keeps the column shards Wt ∈ Rdt ×Vloc and (j) Ws ∈ Rds ×Vloc of the teacher’s and the student’s lm_head, with Vloc = V /tp, identical on every cp (q) rank. With Hs ∈ Rr×ds the student’s own final hidden states of the same rows, cell (q, j) therefore (q,j) (q) (j) evaluates the logits block Zx = Hx Wx ∈ r×V R loc for x ∈ {t, s} (teacher, student): the logits Zx ∈ RL×V are partitioned over the receiver group, cp ranks owning different rows and tp ranks different columns. Statistics are reduced over the tp ranks of one cp rank only; cp ranks exchange no row statistics inside the loss. Appendix B.1 gives the row layout, protocol, and schedule.

(q)

Hx . Cell (q, j) computes

(q,j)

(q)

Zx,k = Hx,k Wx(j) ∈ RC×Vloc ,

x ∈ {t, s}, (8) (q,j) the window-k rows of its block Zx . Fix a row i of the window; we leave i implicit in zx , mx , ax and c, which all refer to row i. Its log(0) (tp−1) its are zx = [zx , . . . , zx ] ∈ RV , where (j) (q,j) zx ∈ RVloc is the ith row of Zx,k with entries (j)

zx,u , u = 0, . . . , Vloc −1 (Figure 2a), and no rank forms zx . With px = softmax(zx ) over the V 5

ModelOpt

KDFlow

NeMo-RL

DisaggKD (ours)

Teacher and student on same GPUs same GPUs Idle teacher weights kept kept Teacher forward in the student’s step in the student’s step Teacher output logits, V /tp logits, V /tp columns columns KL taken over all V columns all V columns

same GPUs kept (optional offload) phases

same GPUs offloaded 1 phase per step

hidden states

L/cp × V /tp

tokens ×V , chunks optional

logits, V columns top 8192 columns step tokens ×V

separate nodes lm_head only concurrent, ahead hidden states

Largest logits tensor

ModelOpt fused-W

all V columns

L/cp × V /tp

all V columns C × V /tp

Qwen3.5 pair, 4 nodes (ModelOpt: Bal-L+K / KD-only) Tokens per step (103 ) 4194 4194 Tokens/s per GPU 741 / 1450 1303‡

4183 390

1047 645

4129 1846

Qwen3.5 pair, 3 nodes Tokens per step (103 ) OOM Tokens/s per GPU OOM

− −

− −

4194 2352

524 1780

524 2426

4129 2389

Qwen3-32B pair, 1 node (2 for DisaggKD); ModelOpt: Bal-L+K / KD-only Tokens per step (103 ) 524 − 515 Tokens/s per GPU 3500 / 3570 − 6570

Table 2: DisaggKD against ModelOpt (NVIDIA, 2025a, 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.

Nodes DisaggKD (ours)

ModelOpt

Qwen3.5: 35B/3B to 9B 2 284 − 3 74 OOM 4 70 91 to 177

147 72 101‡

Qwen3-235B: 235B/22B to 30B/3B 3 393 OOM 8 109 75

OOM −

Nemotron: 560B/55B to 30B/3B 4 630 OOM 8 316 OOM

OOM −

Llama: 405B to 70B, dense 6 818 − 8 − 484 16 186 301

with KD weight λ. KDi is the soft cross-entropy, which has the gradient of KL(pt ∥ ps ). (j) TP rank j holds only zx , so the rank cannot normalize a softmax alone, and gathering zx would undo the sharding. The rank need not gather. Define, for the full row, X mx = max zx,v , ax = ezx,v −mx ,

ModelOpt fused-W

v

c=

X

v

e

zt,v −mt

v

OOM − −

Then px,v = ezx,v −mx /ax , and the row losses are KDi = lses −

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 30a).

i

KDi = −

X

c , at P

CEi = lses − zs,yi ,

(11)

with lses = log v ezs,v = ms + log as , so 5 numbers (mt , ms , at , as , c) replace the 2 rows of logits. Each is a maximum or a sum over columns, (j) so tp rank j computes its local versions mx , (j) (j) ax , c(j) , that is, (10) over the entries of zx only. A fused kernel evaluates the statistics with 1 thread block per row, walking the entries in tiles with the online-softmax recurrence (Milakov and Gimelshein, 2018); the tp ranks then combine the results by a MAX all-reduce followed, after each rank rescales its sums to the global maximum, by a SUM all-reduce (Appendix B.3, (15), (16)).

columns and label yi , the trainer minimizes X  L= (1 − λ) CEi + λ KDi , CEi = − log ps,yi ,

(10)

zs,v .

pt,v log ps,v ,

v

(9) 6

The local statistics are taken relative to a maximum because ez overflows for large logits while ez−m ≤ 1. The maximum of a row is unknown until all entries have been read, and a separate pass to find the maximum would read the logits 2 times, so the fused kernel (Figure 2b) reads the entries (j) (j) of zt and zs once and after entry u keeps the running maximum mx,u and the sums ax,u and cu taken relative to the maximum; when the maximum rises, the sums are rescaled to the new one by a factor of at most 1. In the forward scan a rank holds at most 1 matrix ∈ RC×Vloc per distribution, never one ∈ Rr×V , and the loss equals the fullvocabulary loss in exact arithmetic (Theorem 1).

Loss (nodes)

Colocated (ModelOpt) 90.5 97.7 KD-only (4) Bal-L+K (4) 176.8 172.1 96.3 – SQ-L+K (4) SQ-L+K fused (4) 100.6† – SQ-L+K fused-W (3) 72.0† – Disaggregated (ours) SQ-L+K (3) SQ-L+K (4)

Slots are received into a ring of up to 4 slots ahead within the iteration (Appendix B.1). Iteration time is the larger of the trainer’s compute time and Ntok /(Rρrep ), so a trainer that outruns its replicas is producer-bound (Table 15).

5.1

Setup

74.3 – 39.2 70.14 69.9 61.9‡ 33.0 68.52

– 70.09

Loss and its position. Offline distillation stores teacher outputs as generated sequences (Kim and Rush, 2016) or truncated top-K logits (Ryskulov et al., 2026). On-policy distillation (OPD) trains the student on its own samples and matches the teacher’s token-level distributions on them (Agarwal et al., 2024; Gu et al., 2024; Qwen Team, 2025). Our objective (9) is neither. Its loss and λ schedule follow SlimQwen (Tang et al., 2026) without its multi-token-prediction (MTP) terms (SQ-L+K, SlimQwen’s LM+KD mixture): the fullvocabulary forward Kullback–Leibler (KL) divergence between the teacher’s and the student’s nexttoken distributions, in its soft cross-entropy form, mixed with the LM loss as (1 − λ) LLM + λ LKD and λ decaying from 1 to 0.75, where λ=1 gives the KD-only variant. ModelOpt and KDFlow minimize the same full-vocabulary forward KL, and NeMo-RL takes the KL over the top 8192 columns. ModelOpt’s loss balancer (Bal-L+K, its LM+KD balancer) rescales the KD term to the LM loss instead of weighting by λ, which gives the KD term a gradient weight of 0.89 (Table 21).

Prefetch and pacing

Experiments

69.22 69.78 – – –

on it, at several node counts; the Qwen3-32B pair fits 1 node and serves the comparison with other stacks.

The proof is in Appendix B.3. With µ = min(pp, 2) microbatches in flight, the trainer retains
κstr = µ r(dt +ds ) b + dt Vloc (b+1 1fp8 ) bytes (12) (hidden-state rows and the teacher’s lm_head shard, with a 1-byte fp8 copy when the teacher-logit product runs in fp8, 1fp8 =1), where a colocated loss holds 9.5 fp32-equivalent buffers of r Vloc logits (Eq. 2).

5

63.4 67.97 67.2 67.51 59.9 – 40.7 – 64.1 –

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.

Theorem 1 (Exactness). For every window size C dividing r, every grouping of the entries read per step, every tp dividing V , and every cp with 2 cp dividing L, the losses obtained from the scanned statistics and their combination over the tp ranks through (11) are the CEi and KDi of (9), and the backward pass returns the gradient of the fullvocabulary softmax. Hence, in exact arithmetic, L and its gradients do not depend on these choices.

4.3

s/it, seed Peak MMLU, seed 42 43 42 43

Recipe and colocated stacks. All runs use nodes of 8 80 GB H100 GPUs (79.11 GiB usable), connected by NVLink and InfiniBand, and train on FineWeb-Edu (Penedo et al., 2024). The ModelOpt runs are Megatron-Bridge runs with the KD-only (native), Bal-L+K, and SQ-L+K losses; the stacks differ in loss weighting, learning-rate floor, teacher precision, and, at 4 nodes, global batch. KDFlow

Pairs. We use 5 open-weight, same-tokenizer pairs (Tables 2 and 3) from the Qwen3.5 (Qwen Team, 2026), Qwen3 (Qwen Team, 2025), Nemotron 3 (NVIDIA, 2025b, 2026), and Llama 3.1 (Llama Team, 2024) families. The Qwen3.5 pair (35B-A3B teacher, 9B student) is primary: every colocated loss and disaggregated split runs 7

2.3 2.2 2.1 0

100

200

iteration

SQ-L+K

KD-only

103 tokens/s

Bal-L+K

KD loss change

LM loss

disagg.

0.2 0.1 0 0

100

200

fused

59

57 56

60 30

fused-W

44 46 42 29

24 15

0 2

iteration

2

3

3

4

4

4

4

4

nodes

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).

and NeMo-RL were run on the primary pair and on Qwen3-32B to Qwen3-4B, on the corpus packed to one sample length (Table 19). In most DisaggKD runs the trainer computes the teacher’s logits with an fp8 matrix product (Appendix B.2). MMLU (Hendrycks et al., 2021) is scored with FlexEval (SB Intuitions, 2026) (Appendix C.1). 5.2

DisaggKD splits 1+1 (Table 3); DisaggKD holds less memory per GPU at 3 nodes. Quality. With seed 42 the 4-node disaggregated student is 1.0 point above the colocated Bal-L+K student and 0.6 above KD-only, and the 3-node student is 1.6 above the 4-node one; with seed 43 each student with a run gains 1.3 to 2.3 points, so layout differences are of the size of run-to-run differences (Table 26).

Results on the Qwen3.5 pair

Loss parity. Figure 3a compares the LM loss of the 4-node disaggregated run with the 4-node colocated runs: their 17-iteration moving averages differ by 0.008 nats on average from the SQ-L+K run and 0.012 from the Bal-L+K run (Table 20); the runs read different document orders (Table 25). The disaggregated runs and the colocated SQ-L+K run log soft cross-entropy and the other ModelOpt runs KL, which differ by the teacher’s entropy, 2.00 nats in the logs, so Figure 3b subtracts the initial level.

5.3

Other pairs

Qwen3-32B. Splitting a pair that fits 1 node colocated doubles the GPU count, which halves the rate per GPU at equal step time. The rest of KDFlow’s 2.7× lead, 1.35×, is the step: our trainer-limited 13.5 s (Table 15) against its 9.8 s and ModelOpt’s 18.4 s (Table 19). Qwen3-235B. Of the colocated grids, the 3 colocated runs of Table 18 that fit memory use ep 16 or more, where the all-to-all dispatcher corrupts the student’s forward pass, so training gives not-anumber (NaN) values or diverges (Table 28). With the all-gather dispatcher the colocated stack trains the pair on 8 nodes faster than DisaggKD (75 s with Bal-L+K, against 109 s; Table 3): the disaggregated pace is the forward time of its 4 producer replicas (duty 1.0, Table 15).

Speedup. At 4 nodes and 32768 tokens the disaggregated run is faster than the colocated Bal-L+K, SQ-L+K, and KD-only runs at seed 42 (Tables 4 and 15); it uses an fp8 teacher and the colocated runs bf16 (Table 16). Bal-L+K adds 86 s per iteration to the KD-only run, 81% of the gap, and the traced loss stage gains about 50 s of event time in collectives, copy operators, and stream synchronizations (Table 24a); SQ-L+K adds 6 s. The colocated stack runs the teacher forward pass (16 s, estimated) and the student step in sequence; the disaggregated stack runs them concurrently (§4.3). A colocated run with the fused loss of DisaggKD (SQ-L+K fused) peaks lower and, with the model’s cross-entropy fusion off (fused-W), trains at a similar speed at 3 nodes and faster at 2 nodes, where

Nemotron. Memory and 2 NaN sources keep the colocated stack from training Nemotron: the teacher runs OOM at 4 nodes, and the 2 8-node grids that fit memory (ep 16) give NaN, in the forward loss with the all-to-all dispatcher and in the gradient norm with either. A Nano-teacher probe gives NaN at ep 16, not at ep 4 (Table 29). 8

Llama. The model predicts every colocated tp≤8 grid over the budget, so the colocated runs use tp 32; 6 nodes have no grid at a global batch of 128 (Table 18). At 16 nodes disaggregated (tp 8) is faster than colocated (tp 32; Table 30a).

6

Distillation scaling laws. In Proceedings of the 42nd International Conference on Machine Learning, volume 267 of Proceedings of Machine Learning Research, pages 5977–6045. PMLR. Gemma Team. 2024. Gemma 2: Improving open language models at a practical size. Preprint, arXiv:2408.00118.

Conclusion

Yuxian Gu, Li Dong, Furu Wei, and Minlie Huang. 2024. MiniLLM: Knowledge distillation of large language models. In International Conference on Learning Representations, volume 2024, pages 32694–32717.

Memory and traffic models predict the colocated memory limits. On Qwen3.5 DisaggKD matches the colocated LM loss and is faster than colocated runs with the same loss, and the fused loss of DisaggKD lets a colocated run train on as few nodes at a similar or faster speed; DisaggKD is memory efficient when the student spans several nodes. On the larger pairs colocated runs with a fused loss still run OOM where DisaggKD trains.

Dan Hendrycks, Collin Burns, Steven Basart, Andy Zou, Mantas Mazeika, Dawn Song, and Jacob Steinhardt. 2021. Measuring massive multitask language understanding. In International Conference on Learning Representations. Geoffrey Hinton, Oriol Vinyals, and Jeff Dean. 2015. Distilling the knowledge in a neural network. Preprint, arXiv:1503.02531.

Limitations DisaggKD needs separate teacher nodes, so where a colocated layout already fits the pair it uses more GPUs. Each framework runs under its own recipe; aligning hyperparameters across frameworks, isolating the cause of the NaN in the colocated Nemotron student, and fully sharded data parallelism (FSDP) at pp 1 are outside this paper. The estimation models are fitted to the parallelism ranges and sequence lengths we ran, and some of their terms rest on few runs. We study where the teacher runs and how its loss is computed, not how to raise downstream quality; the runs are short, with 2 seeds on 1 pair, on 1 node type and 1 English corpus.

Pin-Lun Hsu, Yun Dai, Vignesh Kothapalli, Qingquan Song, Shao Tang, Siyu Zhu, Steven Shimizu, Shivam Sahni, Haowen Ning, and Yanning Chen. 2025. Liger Kernel: Efficient Triton kernels for LLM training. Preprint, arXiv:2410.10989. Yoon Kim and Alexander M. Rush. 2016. Sequencelevel knowledge distillation. In Proceedings of the 2016 Conference on Empirical Methods in Natural Language Processing, pages 1317–1327, Austin, Texas. Association for Computational Linguistics. Vijay Anand Korthikanti, Jared Casper, Sangkug Lym, Lawrence McAfee, Michael Andersch, Mohammad Shoeybi, and Bryan Catanzaro. 2023. Reducing activation recomputation in large transformer models. In Proceedings of Machine Learning and Systems, volume 5, pages 341–353. Curan.

References

Hao Liu, Matei Zaharia, and Pieter Abbeel. 2024. RingAttention with blockwise transformers for nearinfinite context. In International Conference on Learning Representations, volume 2024, pages 3992– 4008.

Rishabh Agarwal, Nino Vieillard, Yongchao Zhou, Piotr Stanczyk, Sabela Ramos Garea, Matthieu Geist, and Olivier Bachem. 2024. On-policy distillation of language models: Learning from self-generated mistakes. In International Conference on Learning Representations, volume 2024, pages 21246–21263.

Llama Team. 2024. The Llama 3 herd of models. Preprint, arXiv:2407.21783.

Akhiad Bercovich, Tomer Ronen, Talor Abramovich, Nir Ailon, Nave Assaf, Mohammed Dabbah, Ido Galil, Amnon Geifman, Yonatan Geifman, Izhak Golan, Netanel Haber, Ehud Dov Karpas, Roi Koren, Itay Levy, Pavlo Molchanov, Shahar Mor, Zach Moshe, Najeeb Nabwani, Omri Puny, and 7 others. 2025. Puzzle: Distillation-based NAS for inferenceoptimized LLMs. In Proceedings of the 42nd International Conference on Machine Learning, volume 267 of Proceedings of Machine Learning Research, pages 3806–3830. PMLR.

Taiming Lu and Zhuang Liu. 2026. Strong teacher not needed? On distillation in LLM pretraining. Preprint, arXiv:2605.23857. Zhiyu Mei, Wei Fu, Kaiwei Li, Guangju Wang, Huanchen Zhang, and Yi Wu. 2025. ReaL: Efficient RLHF training of large language models with parameter reallocation. In Proceedings of Machine Learning and Systems, volume 7. MLSys. Maxim Milakov and Natalia Gimelshein. 2018. Online normalizer calculation for softmax. Preprint, arXiv:1805.02867.

Dan Busbridge, Amitis Shidani, Floris Weers, Jason Ramapuram, Etai Littwin, and Russell Webb. 2025.

9

Mistral Team. 2026. arXiv:2601.08584.

Ministral 3.

Preprint,

Qwen Team. 2025. Qwen3 technical report. Preprint, arXiv:2505.09388.

Saurav Muralidharan, Sharath Turuvekere Sreenivas, Raviraj Joshi, Marcin Chochowski, Mostofa Patwary, Mohammad Shoeybi, Bryan Catanzaro, Jan Kautz, and Pavlo Molchanov. 2024. Compact language models via pruning and knowledge distillation. In Advances in Neural Information Processing Systems, volume 37, pages 41076–41102. Curran Associates, Inc.

Qwen Team. 2026. Qwen3.5: Towards native multimodal agents. Bakbergen Ryskulov, Iker García-Ferrero, David Montero, David Jansen, Ali Hashemi, Jezabel R. Garcia, Antonio Tiene, and Román Orús. 2026. Efficient knowledge distillation for LLMs: Offline TopK logits and a fused chunked KL loss. Preprint, arXiv:2608.03796.

NeMo RL. 2025. NeMo RL: A scalable and efficient post-training library. https://github.com/ NVIDIA-NeMo/RL. GitHub repository.

SB Intuitions. 2026. FlexEval: Flexible evaluation tool for language models. https://github.com/ sbintuitions/flexeval. Version 0.17.3, commit 8a07c3c, 23 July 2026.

NVIDIA. 2024–2026. NVIDIA Model Optimizer. https://github.com/NVIDIA/Model-Optimizer. GitHub repository.

Guangming Sheng, Chi Zhang, Zilingfeng Ye, Xibin Wu, Wang Zhang, Ru Zhang, Yanghua Peng, Haibin Lin, and Chuan Wu. 2025. HybridFlow: A flexible and efficient RLHF framework. In Proceedings of the Twentieth European Conference on Computer Systems, EuroSys ’25, page 1279–1297, New York, NY, USA. Association for Computing Machinery.

NVIDIA. 2025a. NeMo Megatron Bridge. https: //github.com/NVIDIA- NeMo/Megatron- Bridge. GitHub repository. NVIDIA. 2025b. NVIDIA Nemotron 3: Efficient and open intelligence. Preprint, arXiv:2512.20856.

Mohammad Shoeybi, Mostofa Patwary, Raul Puri, Patrick LeGresley, Jared Casper, and Bryan Catanzaro. 2020. Megatron-LM: Training multi-billion parameter language models using model parallelism. Preprint, arXiv:1909.08053.

NVIDIA. 2025c. NVIDIA Nemotron Nano 2: An accurate and efficient hybrid Mamba-Transformer reasoning model. Preprint, arXiv:2508.14444. NVIDIA. 2026. Nemotron 3 Ultra: Open, efficient Mixture-of-Experts hybrid MambaTransformer model for agentic reasoning. Preprint, arXiv:2606.15007.

Shengkun Tang, Zekun Wang, Bo Zheng, Liangyu Wang, Rui Men, Siqi Zhang, Xiulong Yuan, Zihan Qiu, Zhiqiang Shen, and Dayiheng Liu. 2026. SlimQwen: Exploring the pruning and distillation in large MoE model pre-training. In Third Conference on Language Modeling.

Pratyush Patel, Esha Choukse, Chaojie Zhang, Aashaka Shah, Íñigo Goiri, Saeed Maleki, and Ricardo Bianchini. 2024. Splitwise: Efficient generative LLM inference using phase splitting. In 2024 ACM/IEEE 51st Annual International Symposium on Computer Architecture (ISCA), pages 118–132.

Erik Wijmans, Brody Huval, Alexander Hertzberg, Vladlen Koltun, and Philipp Krähenbühl. 2025. Cut your losses in large-vocabulary language models. In International Conference on Learning Representations, volume 2025, pages 68174–68193.

Guilherme Penedo, Hynek Kydlíček, Loubna Ben allal, Anton Lozhkov, Margaret Mitchell, Colin Raffel, Leandro Von Werra, and Thomas Wolf. 2024. The FineWeb datasets: Decanting the web for the finest text data at scale. In Advances in Neural Information Processing Systems, volume 37, pages 30811–30849. Curran Associates, Inc.

Xiulong Yuan, Hongqing Chen, Jiaxuan Peng, Fan Zhou, Zhixiang Ruan, Zekun Wang, Bo Zheng, Rui Men, Haiquan Wang, Zhipeng Zhang, Langshi Chen, Man Yuan, Jiaqi Gao, Zhengping Qian, Junyang Lin, Yong Li, Wei Lin, Junhua Wang, and Jingren Zhou. 2026. Accelerating compound LLM training workloads with Maestro. Preprint, arXiv:2605.10501.

Hao Peng, Xin Lv, Yushi Bai, Zijun Yao, Jiajie Zhang, Lei Hou, and Juanzi Li. 2025. Pre-training distillation for large language models: A design space exploration. In Proceedings of the 63rd Annual Meeting of the Association for Computational Linguistics (Volume 1: Long Papers), pages 3603–3618, Vienna, Austria. Association for Computational Linguistics.

Songming Zhang, Xue Zhang, Tong Zhang, Bojie Hu, Yufeng Chen, and Jinan Xu. 2026. KDFlow: A user-friendly and efficient knowledge distillation framework for large language models. Preprint, arXiv:2603.01875.

Ruoyu Qin, Zheming Li, Weiran He, Jialei Cui, Feng Ren, Mingxing Zhang, Yongwei Wu, Weimin Zheng, and Xinran Xu. 2025. Mooncake: Trading more storage for less computation — a KVCache-centric architecture for serving LLM chatbot. In 23rd USENIX Conference on File and Storage Technologies (FAST 25), pages 155–170, Santa Clara, CA. USENIX Association.

Yinmin Zhong, Shengyu Liu, Junda Chen, Jianbo Hu, Yibo Zhu, Xuanzhe Liu, Xin Jin, and Hao Zhang. 2024. DistServe: Disaggregating prefill and decoding for goodput-optimized large language model serving. In 18th USENIX Symposium on Operating Systems Design and Implementation (OSDI 24), pages 193–210, Santa Clara, CA. USENIX Association.

10

A

Estimation Models in Detail

A.1

Memory Model

ranks), device (the peak device memory in use, which adds the non-PyTorch memory), or, at an OOM error, the bytes allocated plus the bytes requested. The trainer’s loss-stage parameter count sums the blocks of the stage, with the layers split evenly over the pipeline stages, the output layer (V ds /tp) on the last stage (a copy of the embedding when the embeddings are tied and pp > 1), and the Gated-DeltaNet, gated-attention, router, and Mamba parameters of each layer type; it reproduces the logged per-rank count of each of 7 trainers within 0.1%. Totals Ds and Es taken from the configuration alone undercount the Qwen3.5 student by 9% (8.2 against 9.0 billion parameters, because a Gated-DeltaNet layer and the gated attention output count as plain attention) and the Qwen3-4B student by 1.1 GiB, because its tied embedding is held on both end stages, so the counts come from the checkpoint tensor index and the block sums. The optimizer term shards over the combined DP×CP group, as Megatron-Core’s distributed optimizer does, so at (dp, cp) = (1, 4) the optimizer state per rank is 4 times smaller than a DP-only form gives. The Adam moments are allocated at the first optimizer step, so a run that stops in its first iteration holds S without them, and its non-PyTorch memory is 3.1 GiB (the mean of 4 OOM events, 2.5, 2.7, 3.2, and 4.0, at the failing rank with the most memory in use; the 4.0 is the Nemotron 4-node run with tp8/pp1/cp4/ep32), against 12.27 GiB at steady state (the device peak minus the torch-allocated peak of the 2 cp4 runs, 12.65 and 11.88 GiB). The student’s routed experts are sharded over pp · ep only because the bridge pins their expert TP to 1, and the Qwen3.5 teacher’s 2 key-value heads limit TP to 2 without head replication.

Per-module accounting. Table 5 lists the terms of (1), (3), (4), and (5) per module. The red constants were fitted to measured OOM events and peak readings (Tables 18, 8, and 7), the green ones follow from the code, and the black ones are allowances. The colocated loss buffers. ModelOpt’s logitsKL loss materializes, at the last pipeline stage, 7 fp32-equivalent logits buffers, each ∈ RL/cp×V /tp : the bf16 student and teacher logits (together 1 fp32equivalent buffer), their 2 fp32 upcasts, the exponentials of the student logits that autograd retains for the backward pass, the student and teacher logprobabilities, and the elementwise KL tensor that is summed over the vocabulary. We fitted the count from the 2 colocated OOM events of the Qwen3.5 pair on 4 nodes (Table 18), whose OOM events fail at the 4th and the 7th buffer. The cp1 event failed on the fourth buffer, as did the 3-node event outside the fit, and the cp2 event on the seventh with 5 to 6 live, which gives the 7 live buffers under the allocation order listed. An eighth buffer stages the loss-averaging all-reduce, set from the Qwen3-235B 8-node cp2 OOM at its first call (Table 18). A run that reaches the backward pass also holds the gradient of the fp32 student logits and its bf16 copy, 1.5 buffers: fitting that count to the peak of each cp4 run gives 1.46 for both runs, and the cp8 run, outside the fit, is predicted within 0.7 GiB (Table 18). Together they are the 9.5 buffers of (2). The LM loss adds 1 more buffer, since the torch-allocated peak of the Bal-L+K run is 3.8 GiB above the KD-only run’s and 1 buffer is 3.8 GiB. The fused loss peaks 19.2 GiB below the SQ-L+K run at cp4, 5.1 buffers, and is outside the count (Table 4). The streamed loss retains κstr of (12) instead, which per microbatch holds r(dt +ds ) hidden elements where the colocated loss holds 2r Vloc logit elements (teacher and student), a reduction of 2Vloc /(dt +ds ): 40× for the Qwen3.5 pair and 1.3× for the dense Llama pair at tp=8.

Producer-side terms. The producer holds the teacher’s whole checkpoint, the vision tower and the multi-token-prediction block where the teacher has them. The shard term of (5) counts dense tensors at 2 bytes and routed experts at 1 byte when fp8 and 2 bytes when bf16, and reproduces the per-rank parameter count in the producer log to 0.15% for the Qwen3.5, Qwen3-235B, Llama, and Qwen3-32B teachers and to 4.2% for the Nemotron teacher, whose Mamba mixers are not fully tensorparallel. The device peak lies 17.7 to 32.1 GiB above the shard for the MoE teachers and 5.9 to 18.2 GiB for the dense ones (Table 7). The forwardpass term φp is the mean of these remainders after

Units and parameter counts. Every memory term is in GiB (230 bytes): a parameter term is the count in billions times the bytes per parameter times 109 /230 . A measured peak is quoted as the quantity that was logged (Megatron’s decimal-GB readings converted to GiB): torch-allocated (the maximum allocated by PyTorch on the loss-stage 11

Module

Term

(2+4) Ploc bytes + 12(Ds /(tp pp dp cp) + Es /(pp ep edp)) bytes, experts on pp · ep, their optimizer state on edp ranks ℓs L 1 pp d b bytes stored layer inputs of the 1 microbatch in flight at the loss stage cp tp s (full recomputation) Student attention scores (L/cp)2 (h/tp) b bytes when attention is not fused (zero for the students of this paper) L Student MoE dispatch 2 cp ke ds b/tp bytes Colocated teacher shard 2(Dt /(pp tp) + Et /(pp ep)) bytes (bf16) Colocated logits buffers (7+1+1.5) (L/cp)(V /tp) 4 bytes (7 buffers, 1 for the loss-averaging all-reduce, and 1.5 for the backward pass; 1 more with the LM loss) Teacher forward-pass memory, MoE expert-MLP tile (8 fp32 copies of the dispatched token pairs, divided by colocated (the peak is the maximum tp · cp; 0.705 of it for a distributed-checkpoint teacher), Mamba activations (4.99 with this term) bf16 copies of each Mamba layer’s inner states, divided by tp · cp), and 10.9 GiB of staging when the teacher is loaded from a Hugging Face checkpoint Non-PyTorch memory, colocated 12.27 GiB, the mean of the 2 cp4 runs; 3.1 GiB in the first iteration, the mean of 4 σcol OOM events (3 Qwen3.5, 1 Nemotron) Trainer stack, disaggregated (loss S+κstr +Astr +σstr , with Astr = Agc +Abwd +Amoe ; κstr is (12) (1.1 GiB for the stage) Qwen3.5 (2+1) trainer) r Trainer reference-cycle garbage Agc 3 nmb tp ds b bytes (12.0 GiB at 64 microbatches)
r Trainer loss-backward and 2.5 · 4ds Vloc + 6 · 4 CVloc + 34 ℓrc tp ds bytes (8.6 GiB); ℓrc : layers recompute temporary memory Abwd recomputed at once moe r Trainer expert-layer block Amoe 2.5 ℓmoe last tp ds b bytes, ℓlast the expert layers of the loss stage (3.75 GiB for (MoE students) Qwen3-30B-A3B) Trainer remainder σstr 15.32 GiB, the mean remainder of 8 trainers Producer shard (2Dt /tpp + βe Et /(tpp epp )) bytes of the whole checkpoint, βe =1 (fp8 experts) or 2 (bf16) 22.70 GiB (MoE: the mean residual of 6 producers), 4.4 GiB plus 12.8 bf16 buffers Producer remainder φp (forward-pass workspace and other of L × dt (dense: through the 2 dense producers) memory) Producer NCCL buffers ωp 1 GiB, a flat allowance Student weights, gradients, optimizer Student activations, colocated

Table 5: Per-module terms of the memory model (ke : 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).

the 1 GiB of NCCL buffers ωp : 22.70 GiB for the 6 MoE producers, and for the 2 dense producers the line through them, 4.4 GiB plus 12.8 bf16 buffers of L × dt (L=32768 for Llama, 4096 for Qwen332B). The line predicts the Qwen3-32B producers at L=8192 and 16384 (peaks 37.56 and 39.84 GiB) to −0.6 and −1.9 GiB; a single buffer count fitted on Llama underestimates all 3 Qwen3-32B runs by 4.3 to 5.6 GiB. The 2 runs of each MoE family differ by at most 1.1 GiB, and the remainder differs between families (21.4 for Qwen3.5, 18.1 for Qwen3-235B, and 31.6 for Nemotron), which a constant does not follow, so the estimate is too high for the Qwen3-235B producers and too low for the Nemotron producers.

points are the 2 Qwen3.5 4-node events, which fix the 7-buffer count, and the 2 Nemotron 8-node events, which set the teacher’s forward-pass terms and the 10.9 GiB staging of the Hugging Face load. The trainer remainder σstr is the mean remainder of the device peaks (the maximum over a run’s trainer GPUs) of the 8 trainers with CR − after the other terms of (4) are subtracted, and the producer’s φp is the mean remainder of the 6 MoE producer peaks, with the 2 dense producers setting the intercept and the L × dt buffer count (Table 7). Table 9 compares prediction and measurement on the Qwen3.5 pair under both layouts, with the use of each row (fit or held out), and Table 8 lists the terms, prediction, and measurement of each trainer. Outside the trainer fit are the 2 Qwen3-235B trainers with CR all, whose device readings are allocator ceilings, 7 Qwen3.5 trainers on other layouts (Appendix D.1), the Qwen3-4B student at 3 further sequence lengths, and the Qwen3-30B student with the garbage collected after every microbatch,

Fitted and held-out configurations. The colocated constants are set on 4 peak readings and 4 OOM points (Table 18). The readings are the device and torch-allocated peaks of the Qwen3.5 cp4 runs with the KD-only and Bal-L+K losses; they set the non-PyTorch memory σcol . The OOM 12

Model

In the fit

Trainer device peak Producer device peak

8 runs: RMS 1.2, largest 2.5 8 runs: RMS 5.0, largest 8.4

Colocated verdict Colocated peak

The Llama colocated bound. For the 405B-to70B pair the model evaluates every colocated grid from 8 to 32 nodes with TP at most 8, the student’s number of key-value (KV) heads, and PP dividing 2 (the layer counts are 80 and 126). The lowestpeak grid at 8, 16, 24, 28, and 32 nodes is predicted at 89 to 102 GiB, over the budget, and the other node counts have higher minima or no legal grid, so none was run (Table 18). Megatron-Core also allows TP above the KV head count by replicating the heads, as the producer does at 16. With tp=32 the colocated layout fits on 16 nodes and on 8, and the model predicts both within 1.2 GiB of the torchallocated peak without any Llama measurement in the colocated fit (the Llama trainer and producer are in their fits); node counts below 8 were not run.

Outside the fit

13 runs: RMS 1.7, largest 4.0 leave one MoE family out (6 runs): RMS 8.6, largest 12.3 6 configurations: 18 configurations: 5 correct 17 correct 4 peaks, 4 OOM 9 peaks: RMS 1.5, points: RMS 0.5, largest 3.0; 2 OOM largest 1.2 points, largest −15.8 − within 2% (3 and 4 nodes)

Inter-node traffic (Qwen3.5) NVLink traffic − (Qwen3.5)

4.4% (colocated), 15% (trainer), 3% (producer)

Table 6: Accuracy of the estimation models: root-meansquare (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.

Cache-release ceilings. The Qwen3-235B (2+1) trainer with CR all reserves 68.1 GiB and holds 11.5 GiB outside the allocator, with a device peak of 79.6 (Table 11), so its device reading is a ceiling; the (4+4) reading with CR all, 67.8, is a feasibility check for the same reason. Trainer terms. 3 terms account for most of the spread of the trainer’s remainder, the device peak after S and κstr (Table 10); Table 11 gives the measured memory of each trainer and Table 12 the probes of the terms. (1) Reference-cycle garbage. The trainer collects Python garbage at the iteration boundary only, so the allocation of the loss stage grows with every microbatch (Ramp in Table 11). The growth per microbatch is 0.70 to 0.97 of 3 sequence-parallel shards of the bf16 hidden states, 3 (r/tp) ds b bytes, on the 6 trainers where it was logged. On the Qwen3.5 (2+1) trainer, collecting after every microbatch lowers the endof-iteration allocation by 11.8 GiB, as the model gives, and each iteration takes 57% longer (Table 12). (2) Temporary memory of the loss backward and of the recomputed layers. The fp32 head gradient has 4ds Vloc bytes, and the fused backward holds it in 2 to 3 copies (the accumulator, the matmul temporary, and the bf16 cast) next to 5 to 8 fp32 window tensors ∈ RC×Vloc (the forward scan holds 1 per distribution), depending on the order of operations; we use 2.5 and 6, and over that range the standard deviation (SD) of the remainder is 1.1 to 1.6 GiB. The temporary memory measured on the Qwen3-4B student at TP 8 agrees with the terms (Table 12). Recomputing ℓrc layers at once holds 34 (r/tp)ds bytes per layer (Korthikanti et al., 2023), the count of a dense trans-

whose prediction was recorded before its run (Table 8). Colocated verdicts and peaks. The model classifies 17 of the 18 colocated configurations outside its fit correctly and 5 of the 6 that set constants (OOM or fits; Table 18). The 2 exceptions are the Qwen3.5 Bal-L+K run, a fit row that peaks at 79.1 GiB and is predicted 79.6, and the Nemotron run on 4 nodes at tp8/pp1/cp4, which ran OOM in the teacher’s first forward pass although the model predicts a device peak of 63.0 GiB. The measured demand at the failing allocation is 75.7 GiB against 59.9 predicted, and the model has no second tp 8 teacher run to calibrate the teacher’s forward-pass terms. Outside the fit, the model predicts the peaks of 9 runs to within 3.0 GiB (Table 18); the largest errors, −3.0 and −2.8 GiB, are the Qwen3-235B runs on 2 and 4 nodes, which ended with a NaN gradient in their first iteration, before the optimizer step: the model evaluates them without Adam moments and with the 3.1 GiB non-PyTorch memory of that iteration. 13

Run (T+P nodes)

Replica

Shard Pred. Meas.

Err. Resid. Err., family left out

Qwen3.5 (2+1) Qwen3.5 (3+1) Qwen3-235B (2+1) Qwen3-235B (4+4) Nemotron (2+2) Nemotron (4+4) Llama (8+8) Qwen3-32B (1+1)

TP2, fp8 TP2×EP2, fp8 TP8, bf16 TP8, bf16 TP8×EP2, fp8 TP8×EP2, fp8 TP16, bf16 TP2, bf16

18.1 10.4 54.7 54.7 38.6 38.6 47.3 30.5

+2.4 +2.1 +6.1 +5.2 −7.3 −8.4 0.0 0.0

41.8 34.1 78.4 78.4 62.3 62.3 65.4 36.4

39.4 32.0 72.4 73.3 69.6 70.7 65.4 36.4

+3.5 +3.2 +8.9 +8.0 −11.2 −12.3 − −

21.3 21.6 17.7 18.6 31.0 32.1 18.2 5.9

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 × dt 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

nmb

Qwen3.5 (2+1) Qwen3.5 (3+1) Qwen3.5 (2+1), tp4 Qwen3-32B (1+1), seq 4096 Llama (8+8) Nemotron (2+2) Nemotron (4+4) Qwen3-235B (4+4), ep8 Qwen3-235B (2+1) Qwen3-235B (4+4) Qwen3.5 (4+2) Qwen3.5 (4+2), pp4 Qwen3.5 (4+2), cp4 Qwen3.5 (3+1), seq 8192 Qwen3.5 (4+2), pp4/cp4, seq 8192 Qwen3.5 (2+1), collected Qwen3.5 (4+1), pp4/cp4 Qwen3-32B (1+1), seq 8192 Qwen3-32B (1+1), seq 16384 Qwen3-32B (1+1), seq 32768 Qwen3-235B (2+1), collected

2/2/2/2 2/2/2/3 4/2/2/1 2/2/1/2 8/2/2/2 2/2/1/4/4 2/2/1/8/4 2/2/1/8/8 2/2/1/4/4 2/2/1/8/4 2/2/2/4 2/4/2/2 2/2/4/2 2/2/2/3 2/4/4/1 2/2/2/2 2/4/4/1 2/2/1/2 2/2/1/2 2/2/1/2 2/2/1/4/4

64 42 128 64 64 32 16 16 32 16 32 64 64 168 512 64 128 32 16 8 32

S κstr 18.8 16.7 12.5 12.3 37.0 47.4 35.9 23.0 43.8 33.1 15.6 11.5 15.6 16.7 11.5 18.8 11.5 12.3 12.3 12.3 43.8

1.1 1.1 0.7 1.2 2.0 2.8 2.8 1.6 1.3 1.3 1.1 1.1 0.9 0.8 0.8 1.1 0.9 1.3 1.6 2.0 1.6

Agc Abwd Amoe Pred. Meas.

Err. Use

12.0 7.9 12.0 1.9 6.0 7.9 3.9 3.0 6.0 3.0 6.0 12.0 6.0 7.9 12.0 0.2 12.0 1.9 1.9 1.9 0.2

−1.3 −0.7 −0.2 +0.5 +2.5 +0.8 −0.4 −1.3 0.0 −1.8 −0.8 −0.9 −2.4 +0.2 −1.8 −0.8 −4.0 −0.4 0.0 +2.0 −1.2

8.6 8.6 4.3 10.1 3.2 4.5 4.5 9.5 9.5 9.5 8.6 8.6 8.1 7.8 7.7 8.6 8.1 11.4 14.1 19.4 9.5

0.0 0.0 0.0 0.0 0.0 2.5 2.5 3.8 3.8 3.8 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 3.8

55.8 49.6 44.9 40.8 63.5 80.4 65.0 56.1 79.6 66.0 46.7 48.6 46.0 48.5 47.3 44.0 47.8 42.3 45.2 50.9 74.1

57.1 50.3 45.1 40.3 61.0 79.6 65.4 57.4 79.6 67.8 47.5 49.5 48.3 48.3 49.1 44.8 51.8 42.6 45.2 49.0 75.3

fit fit fit fit fit fit fit fit ceiling ceiling held out held out held out held out held out held out held out held out held out held out held out

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; nmb : microbatches per DP replica. Pred.: the 5 terms shown plus σ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 σ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.

former, used for every student (ℓrc =1 except in the 4 Qwen3-32B-pair trainer runs, where ℓrc =8). (3) The expert-layer block of MoE students. Across the optimizer step their first pipeline stage holds more than its statics (Stage 1 in Table 11). An allocation dump of the Qwen3-30B trainer (1 run) attributes this to tensors that stay live across the boundary: the retained autograd graph, the bias-dropout activations, and the dispatch row maps and workspaces

(Table 12), about 2.5 bf16 copies of [r/tp, ds ] per expert layer. The model applies 2.5 such copies to the expert layers of the loss stage (Amoe ); the 2.5 is calibrated on 1 Qwen3-30B dump, at tp 2 and cp 1. The remaining σstr is memory outside PyTorch and unallocated memory that the caching allocator holds reserved (Table 11), and it contains the layer inputs saved for the recomputation, which no other term counts (1.0 GiB on the loss stage 14

Configuration

Pred. Measured

Use

Coloc., 3 nodes, cp1 Coloc., 4 nodes, cp1 Coloc., 4 nodes, cp2 Coloc., 4 nodes, cp4 Coloc., 4 nodes, cp4, Bal-L+K Coloc., 4 nodes, cp8 Coloc., 4 nodes, cp4, SQ-L+K Coloc., 4 nodes, cp4, SQ-L+K fused Coloc., 3 nodes, cp1, SQ-L+K fused Coloc., 3 nodes, cp1, SQ-L+K fused-W Coloc., 2 nodes, cp2, SQ-L+K fused-W Coloc., 2 nodes, cp1, SQ-L+K fused-W Disagg. (2+1), trainer Disagg. (3+1), trainer Disagg. (tp4), trainer Disagg. (4+1, pp4, cp4), trainer Disagg. (2+1, collected), trainer Disagg. Qwen3-30B (2+1, collected), trainer Disagg. (2+1), producer Disagg. (3+1), producer

86.8 OOM, 72.9+15.2 held out 86.4 OOM, 72.4+15.2 fit 78.3 OOM, 71.4+7.6 fit 75.8 76.1 fit 79.6 79.1 fit 77.3 77.0 − 76.8

held out outside

− 58.7

outside

outside

− 65.8

outside

fit fit fit held out

44.0 44.8

held out

74.1 75.3

held out

41.8 39.4

fit

34.1 32.0

fit

S, κstr and Agc and Abwd and Amoe

29.91 23.09 16.41 15.32

4.7 4.1 2.3 1.3

5.1 4.4 2.4 1.4

5.7 4.5 2.4 1.4

3 most under-predicted of the 7 (Table 8).

− OOM, 64.9+0.5 outside 55.8 57.1 49.6 50.3 44.9 45.1 47.8 51.8

Mean SD LOO LOFO

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 σstr and its spread.

− OOM, 57.7+15.2 outside − 78.9

Remainder after subtracting

A.2

Traffic Model

Data-movement derivations. Table 13 lists the volume of each parallelism dimension with its Qwen3.5 values (dt =2048, ds =4096, V =248320, ℓt =40, ℓs =32, τ =220 ), and Table 1 sums them with the transfer. The tables hide 3 points: (1) in the Qwen3.5 architecture every 4th layer is full attention: CP exchanges the keys and values of those layers around a ring (Liu et al., 2024), and the remaining linear-attention layers redistribute their input projection from sequence to head sharding, and their output back, by all-to-all, as Megatron-Bridge implements it; (2) gradients are reduced over the DP×CP group, so the student’s gradient traffic is nonzero at dp=1 when cp > 1; (3) a ring over a group that spans several nodes sends (n−1)/n of its bytes over InfiniBand at each node boundary of the ring, and the transfer multiplies the hiddenstate volume by tp because the sender transmits the rows of each cp rank to its tp last-stage ranks.

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 nonPyTorch memory and the logits-buffer count, the trainer rows the trainer remainder σstr , and the producer rows the forward-pass term φ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.

Traffic validation. Table 14 compares the calculated inter-node bytes of every node with its measured InfiniBand traffic (for the disaggregated runs the node’s total rate over its adapters, averaged over the 5 s samples from iteration 3, checkpoint iterations excluded, times the mean iteration time; for the colocated run the byte counts of the 8 adapters that carry the traffic, read at the ends of iterations 2 and 12; Table 1 divides the bytes by 90.5 s for colocated KD-only, the median iteration time of its 200-iteration run, and by 76.0 s, iterations 3 to 400, for the 400-iteration run of the (2+1) split that carried the counters (not a row of Table 17), whose measured rate is the mean received rate from iteration 3), placing ranks in Megatron’s order (tp, cp, dp, pp) on 8-GPU nodes. On the colocated run each node is 1 pipeline stage, only the pipeline

at (2+1)). On the Qwen3.5 tp-4 trainer, turning NVLink SHARP (NVLS) multicast off changes the device peak by −0.3 GiB and loading kernels eagerly by +1.7 (Table 12); neither probe was run on another family. Memory outside PyTorch differs between families by more than the residuals do; the 8 trainers come from 4 student families, which do not separate family from model dimensions. At cp 4 the held-out Qwen3.5 trainers hold more memory outside PyTorch than at cp 2 (Table 11) and are the

15

Trainer

Device Torch Reserved Non-PyTorch Unallocated End of iter. Stage 1 Ramp

Agc

Qwen3.5 (2+1) Qwen3.5 (3+1) Qwen3.5 (2+1), tp4 Qwen3-32B (1+1) Llama (8+8) Nemotron (2+2) Nemotron (4+4) Qwen3-235B (4+4), ep8 Qwen3-235B (2+1), CR all Qwen3-235B (4+4), CR all

57.1 50.3 45.1 40.3 61.0 79.6 65.4 57.4 79.6 67.8

39.2 33.0 29.0 25.1 50.3 67.5 52.5 40.7 64.3 50.7

43.1 36.6 31.1 28.8 52.7 69.8 54.2 46.0 68.1 56.1

13.9 13.4 14.0 11.5 7.8 9.7 9.5 11.3 11.5 11.2

3.8 3.6 2.1 3.7 2.4 3.6 4.2 5.3 3.8 5.5

31.8 25.6 25.1 15.6 43.8 59.3 43.8 30.5 54.1 40.4

+0.3 − − +0.2 +0.2 +2.2 +2.2 − +3.7 +3.7

11.6 12.0 − 7.9 11.7 12.0 1.7 1.9 − 6.0 − 7.9 − 3.9 2.6 3.0 5.1 6.0 2.1 3.0

Qwen3.5 (4+2) Qwen3.5 (4+2), pp4 Qwen3.5 (4+2), cp4 Qwen3.5 (3+1), seq 8192 Qwen3.5 (4+2), pp4/cp4, seq 8192 Qwen3.5 (2+1), collected Qwen3.5 (4+1), pp4/cp4

47.5 49.5 48.3 48.3 49.1

30.1 31.5 29.5 32.1 30.8

33.6 35.2 32.4 34.9 33.2

13.9 14.3 15.9 13.4 15.8

3.5 3.7 2.9 2.8 2.4

22.7 24.6 22.7 25.6 24.6

− − − − −

− 6.0 − 12.0 − 6.0 − 7.9 − 12.0

44.8 51.8

27.7 31.1

30.8 36.7

13.9 15.2

3.1 5.6

20.0 24.6

− −

− 0.2 − 12.0

Table 11: Measured memory of the 8 fitted trainers, the 2 runs with CR all, and the 7 Qwen3.5 trainers outside the σ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); Agc : the model’s garbage term. Probe

Trainer grid

s/it Torch Device

NCCL Comms Measured; model

Qwen3.5 (2+1), garbage kept Qwen3.5 (2+1), collected

tp2/pp2/cp2/dp2

74.3

39.2

57.1

−

− end of iter. 31.8

tp2/pp2/cp2/dp2

116.9

27.7

44.8

−

Qwen3.5 (2+1), tp4, reference NVLS multicast off eager kernel loading Qwen3-32B (1+1), tp2 Qwen3-32B (1+1), tp8

tp4/pp2/cp2/dp1

97.0

29.0

45.1 59.6 (3.73)

− end of iter. 20.0; −11.8 17 −

tp4/pp2/cp2/dp1 tp4/pp2/cp2/dp1 tp2/pp2/cp1/dp2 tp8/pp1/cp1/dp1

98.2 100.2 − 34.7

29.0 29.0 25.1 11.8

44.8 46.8 − 23.2

59.6 (3.72) 59.6 (3.72) 38.2 (4.78) 23.2 (2.90)

17 17 19 19

Qwen3-235B (2+1), collected

tp2/pp2/cp1/dp4, ep4

402–405

59.6

75.3

−

− − − temporary memory 2.56; 2.52 − retained set 3.57; 3.75; statics 43.75; 43.76

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 Agc of the same grid (calibrated on its ramp). Temporary memory: peak minus end-of-iteration allocation against the term Abwd (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 Amoe (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.

crosses InfiniBand, and every counter is within 1.4% of the calculation; ModelOpt concatenates the teacher’s hidden width to the student’s in the pipeline tensor, so the teacher’s activation-sized bytes cross each boundary in both directions. At

3 nodes the tensor, context, and data groups stay inside the nodes, so only the pipeline activation transfers and the transfer cross InfiniBand, and every counter is within 2% of the calculation. At 4 nodes the ring of each DP×CP group also spans

16

Dimension

Teacher (forward)

ℓt 2(tp−1) Tensor 2 pp τ dt b bytes tp L dt b Pipeline, per nmb cp bytes, both directions tp boundary and colocated direction

Context, attention Context, linear attention Data and context Expert

Student (training)

Qwen3.5 (GB; student; teacher) colocated disaggregated

ℓs 2(tp−1) 6 pp τ ds b bytes tp L ds b nmb cp bytes, both directions tp

412.3; 85.9 4.29; 2.15

824.6; 343.6 4.29; 0

51.5; 8.1

34.4; 0

232.8; 97.0

310.5; 0

5.0; 0

10.1; 0

0; 0

0; 2061.6

h ℓatt t (cp−1) τ 2 kv,t dhead b bytes pp tp

s 4 pp (cp−1) τ 2 kv,s dhead b bytes tp

ℓatt

ℓlin cp−1 w t τ tp b bytes pp cp

w s cp−1 3 pp τ tp b bytes cp

none

dp cp−1 Ploc (4+2) bytes dp cp

 ℓmoe t τ ke dt 2(ep−1) b+ pp  ep etp−1 (b+4) bytes etp

3 spp 2(ep−1) τ ke ds b bytes ep

h

ℓlin

ℓmoe

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 ∆ spans n∆ ranks and uses NVLink if n∆ ≤ 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), tp cp (expert), and tp cp ep 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). ℓatt , ℓlin , and ℓmoe count full-attention, linear-attention, and MoE layers, hkv the KV heads, dhead 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 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 ep=1 and etp=1. Gradients are reduced in fp32 (4 bytes) and parameters gathered in bf16 (2 bytes).

nodes, and its gradient and parameter traffic is 1.0 to 2.0 times the pipeline’s; the calculation stays within 2% and below every counter. NVLink counters (bytes sent per GPU over iterations 3 to 12, and 3 to 10 for the disaggregated trainer and producer) lie above the calculation on the colocated run (4.4%, with the 4 pipeline stages within 3% of each other; the counter snapshots lie within 1 iteration of the stated boundaries), the disaggregated trainer, and the producer (Table 1; Table 6 gives the deviations). The producer’s tensor term, 344 GB, is within 1.4% of the volume of its small all-gather and reduce-scatter pairs (339 GB); the other 86% of the calculated bytes are the expert term, the dispatch all-gather and the fp32 combine reduce-scatter over the expert tensor-parallel group, which move the top-8-replicated tokens of every MoE layer, and the combine reduce-scatter alone carries 57% of the calculated collective bytes. The NVLink counter reproduced a 2-GPU all-reduce, a 4-GPU all-to-all, and a send/receive of known size within 0.4% on 1 node; the InfiniBand counters were not calibrated this way.

Node

Sent calc. meas.

Received Sent: calc. meas. pipe.+ring

3 nodes (dp2) 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) Node 0 44.9 45.5 44.9 45.4 22.5+22.4 Node 1 67.3 68.0 78.6 79.4 22.5+44.8 Node 2 44.9 45.1 67.5 67.9 22.5+22.4 Producer 33.8 34.2 0 0.1 4 nodes colocated (pp4, dp1) Stage 0 51.5 52.1 51.5 52.2 51.5+0 Stage 1 103.1 104.3 103.1 104.2 103.1+0 Stage 2 103.1 104.2 103.1 104.2 103.1+0 Last stage 51.5 52.1 51.5 52.1 51.5+0

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×CP rings. Received bytes at the trainer’s last-stage nodes include the transfer. Colocated: 1 pipeline stage per node, iterations 3 to 12.

B

Method Details

B.1

Streaming Pipeline

Receiver groups and replication. A receiver group is the tp × cp last-stage ranks of one DP 17

Input: teacher replica, corpus order, look-ahead δ iterations, start slot η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 η = η0 , η0 +1, . . . do 3 if η mod R ̸= my replica then continue 4 Ht ← forward of microbatch η (final hidden states) 5 wait until η < (iteration + δ) · slots per iteration 6 slice rows per last-stage rank; send header and payload to the receiver group 7 end for

replica, so a trainer has dp receiver groups. Rows are divided over the cp ranks in Megatron’s loadbalanced layout for causal attention (Llama Team, 2024): cp rank q owns pieces q and 2 cp − 1 − q of 2 cp equal pieces of the sequence (the layout that ring-attention implementations call zigzag), and slots are sent in bf16. The producer serves receiver groups in 2 layouts. In the dedicated layout (replicas ≥ groups, a multiple of dp) each replica serves exactly 1 group and each group is served by R/dp replicas; in the shared layout (replicas < groups) every replica joins every group and sends 1 message per slot per destination rank. Slot η goes to group η mod ngroups from replica η mod R at microbatch size 1. The Qwen3.5 3-node run is dedicated and its 4-node run shared.

Algorithm 1: Producer loop of one replica.

all runs and does not show which side limits. We therefore also measure the forward-only rate ρrep of one replica and compare the producer-only time Ntok /(Rρrep ) with the measured iteration time (Table 15). For the Qwen3.5 runs we measured ρrep with the producer writing to a sink on idle nodes with a warm page cache (the steady rate of slots after the first); for the Qwen3-235B runs we took the steady forward time per slot that each replica logs, and for the other pairs the mean forward time of the first slots of the run. The duty, the ratio of the 2 times, is 1 when the producers limit and below 1 when the trainer does, because the transfer of the Qwen3.5 pair takes 0.09 s per iteration at the link limits (Table 1). The producers limit the Qwen3-235B and Nemotron runs and the Llama (8+8) training run; for Nemotron and Llama ρrep is the first-slot forward rate, so their duty of 1 is an upper bound, and the Llama training run’s ρrep was read on a shared network (the pace check’s 6.1 gives it a duty of 0.92). At constant ρrep the iteration time falls with the replica count R. The trainer limits every run with a duty below 1, among them the Qwen3.5 (2+1) run, the first (3+1) run (duty 0.90), the Qwen3-32B run, the Llama (4+2) run, and the Llama (8+8) pace check, whose replicas consumed less than their forward-only rate; the (3+1) rerun has duty 0.97 and the seed-43 run with CR − duty 1, so the (3+1) split is balanced within 10%. The forward-only rate of the TP2×EP2 fp8 replica of the Qwen3.5 producers is the same on all 8 nodes measured (32.6 to 32.9k tokens/s). The seed-43 run with CR all is limited by the trainer, and with CR − by the producers (Table 15).

Communication protocol. Each receiver group has its own NCCL communicator containing the group’s last-stage ranks and the producer ranks mapped to it, so each communicator contains only the ranks of 1 receiver group and its producer ranks. A slot is an int64 header carrying the slot id followed by the bf16 payload ∈ Rr×dt , already sliced to the cp rank’s share of the rows; per-pair FIFO order and the header check make ordering structural, and the trainer raises on any slot-id mismatch. The producer runs at most δ iterations ahead of the trainer’s consumption counter, receive buffers are posted at request creation, and the prefetch ring holds up to 4 slots per last-stage rank by default. Only a 128-byte communicator id per receiver group crosses the rendezvous store, and every hidden-state byte flows over InfiniBand P2P. The lm_head weights are written once to a file in the stream directory, and V must be divisible by tp. The trainer’s step for 1 microbatch is Algorithm 2, and the producer loop is Algorithm 1. Resume. The producer takes a start slot so that it neither recomputes nor re-sends the prefix the trainer has consumed, and logs the start slot; unit tests check the slot arithmetic of a resume. The window width C does not change the loss in exact arithmetic; it sets the size of the per-window buffers, which scale with C Vloc . Pace-limiting side. Each producer replica logs its cumulative token count; the mean rate between the 5th and the last report is the consumed rate. Because a windowed pipeline equalizes the mean rates of producer and trainer, the consumed rate agrees with Ntok /(R · iteration time) within 4% in 18

Pair (T+P nodes)

R Consumed ρrep Ntok /(Rρrep ) Meas. Duty Limit

CR

Qwen3.5 (1+1), 15 it. Qwen3.5 (2+1) Qwen3.5 (3+1) Qwen3.5 (3+1), rerun Qwen3.5 (3+1), seed 43 Qwen3.5 (3+1), seed 43, 30 it. Qwen3.5 (3+1), seed 43, 30 it. Qwen3.5 (2+1), bf16 teacher Qwen3-235B (2+1) Qwen3-235B (4+4) Nemotron (2+2) Nemotron (4+4) Llama (8+8), training run Llama (8+8), pace check Llama (4+2) Qwen3-32B (1+1)

4 4 2 2 2 2 2 2 1 4 1 2 4 4 1 4

− − − − all − all all all all − − − all all −

3.7 13.8 29.5 31.4 20.1 33.2 − 13.6 10.6 9.3 6.7 6.6 5.6 4.5 5.1 9.6

28.1 28.1 32.8 32.8 32.8 32.8 32.8 34.7 10.7 9.6 6.6 6.6 5.6 6.1 6.7 10.4

37 37 63 63 63 63 63 60 393 109 631 316 188 173 627 13

283.7 74.3 69.9 65.0 99.4 61.9 94.4 153.0 392.9 108.5 630.0 315.9 186.1 234.8 818 13.5

0.13 0.50 0.90 0.97 0.63 1.0 0.67 0.40 1.0 1.0 1.0 1.0 1.0 0.74 0.77 0.93

trainer trainer trainer trainer trainer producer trainer trainer producer producer producer producer producer trainer trainer trainer

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; ρrep : forward-only rate of one replica, both in 103 tokens/s; Ntok /(Rρrep ): the iteration time if the trainer were instantaneous; Meas.: measured median iteration time, both in seconds; Duty: Ntok /(Rρ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. (q)

Input: student hidden states Hs , prefetch ring, shards (j) (j) Wt , Ws , labels y, window size C (q) 1 Ht ← pop the next slot from the ring; post receives for up to 4 following slots of the same iteration 2 for k = 0, . . . , K−1 do (q,j) (q) (j) 3 Zx,k ← Hx,k Wx for x ∈ {t, s} 4 kernel: for each row, one pass over the entries u = 0, . . . , Vloc −1 in tiles gives (j) (j) mx , ax (x ∈ {t, s}) and c(j) (j) MAX all-reduce of mx ; rescale; SUM 5 (j) (j) all-reduce of ax , c over the tp ranks; the label logit of CEi comes from a separate vocabulary-parallel pass P 6 L ← L + i [(1−λ)CEi + λ KDi ] (q,j) over the rows of window k; free Zx,k 7 end for 8 backward: recompute each window; accumulate (q) (j) ∂L/∂Hs and ∂L/∂Ws

routed experts are most of an MoE teacher’s parameters (546.7B of 560.4B for the Nemotron teacher), so the shard is 51% of its bf16 size. For the 560B teacher the model predicts that bf16 everywhere does not fit a 16-GPU replica: about 65 GiB per GPU for weights alone, plus the 31 to 32 GiB remainder measured with fp8 (no bf16 run of this teacher exists). The experts are stored as fp8 primary weights (e4m3 with tensor-wise scales), their grouped GEMMs run under Transformer Engine’s current-scaling fp8 autocast, and activations, including the hidden states sent to the trainer, stay bf16. The Qwen3-235B producers and those of the dense Llama and Qwen3-32B teachers are bf16 throughout, which is why the shard term of (5) counts 1 byte per routed-expert parameter of an fp8 teacher and 2 otherwise. On the trainer, the teacher-logit matrix product can also run in fp8 (e4m3 inputs with per-tensor scales), which leaves the hidden states sent to the trainer in bf16; it was on in the Qwen3.5 runs with an fp8 teacher, the Nemotron and Qwen3-32B runs, and the 2 Qwen3235B trainers of Table 8 with κstr of 1.6, and off in the Qwen3.5 run with a bf16 teacher, the other Qwen3-235B runs, and the Llama runs. The fp8 build changes the stock provider in 4 ways: (1) the quantization recipe is set on the teacher’s model provider and matched to the expert modules by name; (2) the mixed-precision wrapper is skipped, because it casts the fp8 parameters back

Algorithm 2: Streamed KD step for one microbatch at cell (q, j) (§4.2).

B.2

Teacher Precision

In the producer. The producer runs the teacher forward only, without gradients, optimizer state, or backward activations, so its weights and a fitted remainder φp (forward-pass workspace, CUDA context, allocator) fill its memory. It holds the routed experts of the Qwen3.5 and Nemotron teachers in fp8 and every other module (attention, mamba mixers, latent projections, norms, embeddings) in bf16; the shared experts are bf16 on GPT-path teachers and fp8 on the hybrid Nemotron teacher. The 19

Experts

with (−∞, 0, 0) for the empty P set. Then (10) is the case of all V columns, log v∈S ezx,v = mx,S + log ax,S , and 1 ≤ ax,S ≤ |S| for nonempty S.

s/it Torch KD loss, it. 1, 2, Grad norm, it. 3 1, 2, 3

bf16

100.3

fp8

119.1

61.7 0.215, 0.207, 0.207 55.9 0.215, 0.207, 0.795

0.70, 0.70, 1.04 0.70, 0.70, 14.5

Lemma 1 (Merge). Let P0 , . . . , Pn−1 be n disjoint sets of columns with union S, and let l index the parts. Then

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.

mx,S = max mx,Pl , l X ax,S = ax,Pl emx,Pl −mx,S , l

cS =

l

=

cS =

emx,Pl −mx,S

X

ezx,v −mx,Pl ,

v∈Pl

Lemma 2 (Scan). Let the columns of S be read in consecutive tiles U0 , . . . , UU −1 , and let the state after tile u be the merge (14) of the state after tile u−1 with the statistics (13) of Uu , starting from (−∞, 0, 0). Then the state after tile u equals the statistics of U0 ∪ · · · ∪ Uu . For tiles of 1 entry, (j) write zx,u for zx,u , u = 0, . . . , Vloc −1, and let mx,u , ax,u and cu be the state after entries 0, . . . , u, starting from mx,−1 = −∞, ax,−1 = 0 and c−1 = 0; this update is the recurrence mx,u = max(mx,u−1 , zx,u ), ax,u = ax,u−1 emx,u−1 −mx,u + ezx,u −mx,u , cu = cu−1 emt,u−1 −mt,u

(15)

+ ezt,u −mt,u zs,u . Whenever the maximum rises, the exponential factor on the previous sum rescales it to the new maximum, and this factor is at most 1. Proof. We use induction on u. The empty prefix has statistics (−∞, 0, 0). If the state after tile u−1 is the statistics of U0 ∪ · · · ∪ Uu−1 , Lemma 1 with n = 2 shows that merging it with Uu gives the statistics of U0 ∪ · · · ∪ Uu . A column v has statistics (zx,v , 1, zs,v ), and inserting them into (14) with n = 2 gives (15).

v∈S

e

v∈Pl

and the inner sum is ax,Pl . The same computation with the weight zs,v and the teacher’s maximum gives cS .

We prove Theorem 1 for C dividing r, 2 cp dividing L, and tp dividing V . A row has logits zt , zs ∈ RV (§4.2). For a set S of its columns define the statistics X mx,S = max zx,v , ax,S = ezx,v −mx,S , zt,v −mt,S

X l

Proof of Exactness

X

cPl emt,Pl −mt,S .

Proof. The maximum of a union is the largest of the maxima. Splitting the sum by part and inserting e±mx,Pl gives XX ax,S = ezx,v −mx,S

In the colocated stack. The colocated runs use none of this engineering. Megatron-Bridge’s distillation provider builds their teacher in bf16, and loading a native-fp8 release through it drops the scale tensors in the families we ran. Giving the colocated stack an fp8 teacher means doing the same on the teacher’s provider: setting the recipe and skipping the wrapper, naming the experts, and casting the pre-converted bf16 checkpoint into fp8 parameters on load. We implemented it for the routed experts: the build engages and the peak allocated memory falls by 5.8 GiB, but the colocated run’s KD loss leaves the bf16 run’s at iteration 3 (Table 16). The recipe covers routed experts only, so the dense Llama and Qwen3-32B teachers stay in bf16.

v∈S

(14)

l

to bf16; (3) the expert modules of GPT-path models are given module names, because the recipe matches modules by name and an unnamed build constructs every expert in bf16; (4) the checkpoint is cast into the fp8 parameters on load, with the scale tensors of a native-fp8 release applied first, and the CPU copies of the initial values that the fp8 parameters keep are released after the load. No run isolates a single change.

B.3

X

zs,v ,

v∈S

(13) 20

Proof of Theorem 1. Statistics. Tp rank j scans the columns of its shard in any tiling, so by Lemma 2 (j) (j) its final state (mx , ax , c(j) ) is the statistics of the shard. Lemma 1 applied to the tp shards yields the combination

the loss mask, and the normalization lie outside the identity. Remark. Theorem 1 is an identity in exact arithmetic. It does not cover the floating-point summation order, the trainer’s fp8 teacher-logit product, the producer’s fp8 expert weights, or the bf16 hidden-state transfer. ModelOpt and KDFlow minimize the full-vocabulary forward KL, which equals KDi up to the teacher’s entropy, under different loss weights, and NeMo-RL takes it over 8192 columns (§5.1), so the identity equates per-row loss functions and not recipes.

mx = max m(j) x , j

ax =

X

(j)

mx −mx a(j) , x e

j

c=

X

(j)

c(j) emt −mt ,

j

(16) so the all-reduced (mx , ax , c) are the statistics of all V columns, that is, (10). zt,v −mt /a we have t PLoss values. With pt,v = e p z = c/a , and lse = m + log as , so t s s v t,v s,v X X − pt,v log ps,v = lses − pt,v zs,v v

C

Search space. Both layouts are searched as the optimal setting of §3 defines, over the same axes: tp over the sizes that the heads, key-value heads, hidden size, and Mamba groups admit and that we ran, pp over the divisors of the layer count, cp over 1 to 16 in powers of 2, ep over the divisors of the expert count, and dp the remaining factor, which divides the global batch (128, or 126 where dp is 3 or 7: the 3-, 6-, 24-, and 28-node colocated cells and the 4-node disaggregated run). A colocated cell (a pair on a number of nodes) uses 1 grid for both models, because ModelOpt requires the teacher at the student’s tp, pp, and cp and has no interleaved pipeline; the search also gives both models one ep, the largest common divisor of the expert counts, with ep·pp dividing the GPUs. tp and pp are common to the 2 models (Qwen3.5 tp 1 to 2, Qwen3-235B 1 to 4, Nemotron 1 to 8, Llama 1 to 32). Across the 12 searched colocated cells of Table 18, a cell has 12 (Qwen3.5 on 3 nodes) to 305 (Nemotron on 8 nodes) grids. A disaggregated cell has a trainer grid over the student alone, replicas of the teacher spanning tpp × epp GPUs (tpp over the divisors of its heads and Mamba groups, up to 16, and 8 for Nemotron), and a split of the nodes between the 2 roles. The wire requires the number of replicas R to be a multiple of dp, or below dp at microbatch 1; the trainer and producer inequalities of §3.1 bound each role, and R ρrep ≥ Ntok /Titer bounds R below. The Megatron-Bridge runs use microbatch 1, full activation recomputation, and 32768-token sequences unless the text or a table states otherwise.

v

c = lses − , at − log ps,yi = lses − zs,yi . These are (11) and the definitions in (9). Exactly 1 tp rank holds zs,yi and the others contribute 0. (The implementation evaluates the LM term in a separate vocabulary-parallel log-softmax pass over the same window, which yields the same normalizer.) Gradients. For a row, ∂KDi /∂zs,v = ps,v − pt,v and ∂CEi /∂zs,v = ps,v − [v=yi ], hence ∂L = (1 − λ) (ps − eyi ) + λ (ps − pt ). ∂zs This depends on the row only through the statistics (ms , as ) and (mt , at ), which the backward pass regenerates (a MAX and a SUM all-reduce per distribution and window), so it obtains ps and pt from recomputed logits and the same values as a full-row evaluation. For window k, tp rank j holds the block G(j) ∈ RC×Vloc of this gradient, and the chain rule through (8) gives ∂L

= (j)

∂Ws ∂L

(q) ∂Hs,k

XX q

=

X

(q)⊤

(q,j)

Hs,k Gk

Experimental Setup

,

k

G(j) Ws(j)⊤ ,

j

the latter being an all-reduce over the tp ranks. Windows and cp ranks. L is a sum of per-row terms (9). Partitioning the rows first over the cp ranks and then into windows only regroups a finite sum, so the loss as a function of the hidden states and head weights, and its gradients, depend on neither C nor the sharding; the student’s forward pass,

Sharded data parallelism. The search of both layouts excludes fully sharded data parallelism (FSDP). In the Megatron-Bridge release we use, the

21

Megatron-FSDP device mesh reshapes the world as (dp·cp, ep, tp) with no pipeline axis, and setup raised an error at pp 2 and 4 on MoE students. Every trainer grid of Table 17 has pp 2, and the only colocated grids of the primary pair on 4 nodes that the model predicts under the budget (2 of 31) have pp 4 and pp 2 (Table 18); these counts come from the memory model without a sharded-state term, and FSDP at pp 1 was not run. The producer holds weights alone, split over tpp × epp ranks. KDFlow trains the student with FSDP2 and NeMo-RL with its DTensor policy, which wraps it with FSDP2, and on the primary pair KDFlow’s device reading is at the limit and NeMo-RL’s 2.1 GiB below it (79.1 and 77.0 GiB, sampled device memory; Table 19): KDFlow with the idle teacher offloaded and the loss chunked, NeMo-RL at a quarter of the batch, because its full-vocabulary logits buffer, 121 GiB at the full batch, does not fit. At cp 1 the logits buffers are 144 of the 186 GiB of the colocated peak on 3 nodes (Figure 1), memory that sharding the states does not reduce (the buffers shrink with cp). C.1

colocated runs hold the teacher in bf16; one disaggregated run matches the learning-rate floor and the teacher precision and multiplies its total loss by 2 in the same way (below). MMLU is scored 4shot by generative exact match with FlexEval over 14,042 items. The window width is C=1024 to 4096, and DisaggKD is built on Megatron-Bridge. Runs. Table 17 lists the disaggregated runs of the main configurations with their node allocation and grid, and Figure 4 plots their KD loss; Appendices D.1 and D.2 give the primary-pair and other-pair measurements. For each run of Table 17 the trainer’s token rate is within 4% of the producers’ mean rate in the steady state (Table 15). Median iteration times are over iterations 3 to the last, except 8 to the last for the Qwen3.5 (2+1) run and 6 to the last for the seed-43 colocated runs; the captions of Tables 27, 24, 30, and 16 state the windows of their runs. MMLU evaluation. MMLU (Hendrycks et al., 2021) is scored with the public FlexEval release (SB Intuitions, 2026) on the public test set: a fixed 4-shot prompt for each of the 57 subjects, greedy decoding of at most 256 new tokens that stops at a newline, and exact match after NFKC normalization, lowercasing, and stripping.

Disaggregated Runs

Training recipe. Model sizes in Table 3 are config-derived (the Nemotron release name reads 550B). Each GPU reports 79.65 GiB to the driver, of which 79.11 GiB is allocatable by CUDA (the capacity PyTorch reports, the budget M ); in-run device readings include the remainder, so a full card reads 79.6 GiB and the readings of 79.6 are at the limit. All runs train on FineWeb-Edu at sequence length 32,768 (4096 for the Qwen3-32B pair unless a table states otherwise) with full activation recomputation, disaggregated teachers with fp8 routed-expert weights and bf16 elsewhere (the Qwen3-235B, Llama, and Qwen3-32B teachers are bf16 throughout), and bf16 students with fp32 master weights, at global batch 128 (4.19M tokens per iteration, 0.52M for the Qwen3-32B pair; 126 for the 4-node Qwen3.5 run so that dp=3 divides it). The recipe is AdamW at learning rate 4 × 10−4 (5 × 10−5 for the Llama pair) with cosine decay, 32 warmup iterations, and, in the SQ-L+K runs, KD weight λ decaying from 1.0 to 0.75 (λ=1 in the KD-only runs). On the Qwen3.5, Qwen3-235B, Nemotron, and Qwen3-32B pairs the colocated runs decay the learning rate to 10−5 where the disaggregated runs decay to 3 × 10−5 (both layouts decay to 3.75 × 10−6 on Llama), and in the Bal-L+K runs ModelOpt’s rescaling makes its total loss 2× its LM loss at every iteration. The

C.2

Compared Stacks

Colocated configurations. Table 18 lists the colocated configurations behind Table 3. The rows that ended with NaN in their first iteration (Qwen3235B on 2 and 4 nodes, Nemotron with cp8 and with tp4) stopped before the first optimizer step and are evaluated without Adam moments and with the 3.1 GiB non-PyTorch memory of that iteration. The rows that set constants of the model are the 2 Qwen3.5 cp4 rows, the OOM points of the 2 Qwen3.5 4-node rows, and those of the 2 Nemotron 8-node rows with cp4 (Appendix A.1); the other rows are outside the fit, except that the 3.1 GiB firstiteration constant is the mean of the 3 Qwen3.5 OOM points (including the 3-node one) and the Nemotron 4-node OOM point. Fused and fused-W losses. All colocated SQ-L+K runs of the primary pair evaluate LSQ of (19) on the logits that ModelOpt’s forward pass materializes; they differ in how KDi and CEi are computed. The fused and fused-W runs install the loss module of DisaggKD in ModelOpt, the kernel of Section 4.2 and the label pick of the trainer, so 22

s/it 103 tok/s

Run (T+P nodes) Trainer grid

Producer replicas

Iter.

KD loss

Peak

CR

Qwen3.5 (1+1) Qwen3.5 (2+1) Qwen3.5 (3+1) Qwen3-235B (2+1) Qwen3-235B (4+4) Qwen3-235B (4+4), KD-only Nemotron (2+2) Nemotron (4+4) Llama (8+8) Qwen3-32B (1+1)

tp2/pp2/cp2/dp1 tp2/pp2/cp2/dp2 tp2/pp2/cp2/dp3 tp2/pp2/ep4

4×TP2 4×TP2 2×TP2×EP2 1×TP8 (bf16)

15 283.7 200 74.3 200 69.9 200 392.9

14.8 56.4 59.1 10.7

2.228 → 2.488 2.228 → 2.227 2.226 → 2.225 2.020 → 1.962

57.5; 39.4 39.2; 39.4 33.0; 32.0 64.3; 72.9

− − − all

tp2/pp2/ep4, dp8

4×TP8 (bf16)

200 108.5

38.7 2.020 → 1.962 50.7; 73.3

all

tp2/pp2/ep4, dp8

4×TP8

60 105.0

40.0 2.066 → 2.029 51.0; 72.9

all

tp2/pp2/ep4 tp2/pp2/ep4, dp8 tp8/pp2/cp2/dp2 tp2/pp2/dp2

1×TP8×EP2 2×TP8×EP2 4×TP16 (bf16) 4×TP2

200 630.0 200 315.9 400 186.1 400 13.5

6.7 2.19 → 1.87 13.3 2.19 → 1.87 22.5 1.91 → 1.81 38.8 2.139 → 2.099

− − − −

67.7; 69.6 52.5; 70.7 50.3; 65.4 25.1; 36.4

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 (103 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).

smaller node set

equal node count

Qwen3.5-35B → 9B

Qwen3-235B → 30B

2.5 200 iter

2.2 KD loss

KD loss

2.4 2.3

200 iter

2

2.2

200 iter

200 iter

2.1

0

50

100 iteration

150

1.8

200

0

50

100 iteration

150

200

Llama-3.1-405B → 70B

Nemotron-Ultra → Nano 2

400 iter

KD loss

KD loss

2.2 200 iter

2 1.8

1.9 1.8

200 iter

0

50

100 iteration

150

200

1.7

0

100

200 iteration

300

400

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.

their results are a contribution of this work and not a competing baseline: the module removes the loss-side memory limit of the colocated stack, and the layout comparison is between 2 stacks

that share it. (1) SQ-L+K: the criterion forms the fp32 log-probabilities of the whole vocabulary and accumulates KDi in tiles of 16384 columns, and the model forward computes CEi with Mega23

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 3 nodes, tp2/pp4/cp1/dp3

133.1

− 120.8

−

185.9

− 173.6

− 86.8

4 nodes, cp1 184.6 4 nodes, cp2 112.1 4 nodes, cp4/dp1 75.8 4 nodes, cp4/dp1, Bal-L+K 79.6 4 nodes, cp4/dp1, SQ-L+K − 4 nodes, cp4/dp1, SQ-L+K − fused 3 nodes, tp2/pp4/cp1/dp3, − SQ-L+K fused 3 nodes, tp2/pp4/cp1/dp3, − SQ-L+K fused-W 2 nodes, tp2/pp4/cp2/dp1, − SQ-L+K fused-W 2 nodes, tp2/pp4/cp1/dp2, − SQ-L+K fused-W 4 nodes, tp2/pp2/cp8/dp1 77.3 4 nodes, cp4/dp1, 8192 tokens 48.7 4 nodes, tp2/pp2/cp2/dp4, 8192 77.3 tokens Qwen3-235B 3 nodes 210.4 3 nodes, ep4, SQ-L+K fused-W − 6 nodes, cp2 107.2 6 nodes, tp4/pp1/cp4/ep16/dp3 73.5 2 nodes, tp4/pp1/cp4/ep16 64.3

− 172.3 − 99.8 76.1 63.6 79.1 67.4 76.8 − 58.7 −

− 86.4 − 78.3 63.4 − 67.2 − 59.9 − 40.7 −

Nemotron

Qwen3-32B Llama

−

− not run 88.0 OOM at the 4th buffer (student divide) 87.6 OOM at the 4th buffer 79.0 OOM at the 7th buffer − trains 200 iterations − trains 200 iterations − trains 200 iterations − trains 15 iterations

−

−

−

−

78.9

−

64.1

−

− trains 15 iterations, batch 126

65.8

−

50.3

−

− trains 15 iterations

−

−

−

−

77.0 65.0 − 36.4 − 65.0

65.7 37.4 65.7

− − −

− 198.1 − − − 94.9 − 61.2 67.3 61.2

− − − − −

− − − − −

72.9 OOM in the first forward

65.4 OOM after iteration 1 (windowed pick) − fits, trains 3 iterations − fits, 12 iterations − fits, 12 iterations − OOM in model construction 85.8 OOM in model construction − OOM in the loss computation − not run − NaN in the gradient norm of the first backward pass − NaN in the gradient norm of the first backward pass − OOM in the loss computation (74.4 in use, 4.6 requested) − fits; diverges with the all-to-all dispatcher − trains 60 iterations, 75 s/it

4 nodes, tp4/pp1/cp8/ep32

36.7

39.5 33.6

−

−

8 nodes, cp2

87.1

− 74.8

−

−

8 nodes, cp4

64.5

− 52.2

−

−

8 nodes, cp4, all-gather dispatcher, Bal-L+K 4 nodes, tp2/pp4/cp1

64.5

66.3 52.2

54.4

−

334.8

− 322.5

−

−

4 nodes, ep4, SQ-L+K fused-W − 4 nodes, tp8/pp1/cp4/ep32 63.0 8 nodes, cp4, Hugging Face − teacher 8 nodes, cp4, 93.8 distributed-checkpoint teacher 8 nodes, tp2/pp4/cp8 71.0

− − − 59.9 − −

− − − 59.9 − 96.7

− OOM during model construction (78.7 in use, 7.1 requested) 77.0 OOM in model construction 75.7 OOM in the teacher forward 96.6 OOM in the teacher forward

− 81.5

− 77.6

77.7 OOM in the teacher forward

− 67.9

−

−

8 nodes, tp4/pp4/cp4

− 48.9

−

−

8 nodes, tp4/pp4/cp4, all-gather 52.0 72.4 48.9 − dispatcher 1 node, tp2/pp2/dp2, 4096 51.0 − 38.8 38.9 tokens 1 node, same grid, Bal-L+K 52.2 − 39.9 40.0 6 nodes, tp16/pp1/cp1/dp3, 112.4 − 100.1 − global batch 126 6 nodes, tp16/dp3, SQ-L+K − − − − fused-W 8, 16, 24, 28, 32 nodes, tp≤8 89 to 102 76 to 90 (lowest-peak grid) 8 nodes, tp32/pp2/cp1/dp1 47.8 − 35.6 36.7 16 nodes, tp32/pp2/cp2/dp1 39.0 − 26.8 27.5

−

52.0

−

− memory fits; NaN in the gradient norm of the first backward pass − memory fits; NaN in the forward loss of the first iteration − NaN in the gradient norm of the first backward pass − trains 12 iterations

− −

− trains 12 iterations − not run with Bal-L+K

−

80.1 OOM in model construction

−

− not run

− −

− trains 5 iterations − trains 200 iterations

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.

24

tron’s fused vocabulary-parallel cross-entropy. (2) SQ-L+K fused: one criterion computes both terms per token and returns (1P − λι ) CEi + λι KDi . It evaluates KDi = lses − v pt,v zs,v with the one-pass kernel of DisaggKD (an online running maximum and sum of exponentials over tiles, and the gradient ps − pt of (18) re-derived from the saved per-token statistics, so no fp32 tensor of size tokens×V /tp is saved for the backward pass; the kernel saves the logits that ModelOpt already holds and per-token statistics), and CEi = lses − zs,yi by the vocabulary-parallel label pick, in windows of 1024 rows under activation checkpointing. The cross-entropy of the model forward is still computed and logged, with zero gradient weight. Unlike DisaggKD, which forms the logits window by window from hidden states, this criterion reads the full materialized logits. (3) SQ-L+K fused-W (W: the model’s cross-entropy fusion off): (2) with that fusion turned off, so that the logged-only cross-entropy takes the unfused vocabulary-parallel path, which normalizes in place on its single fp32 copy of the logits. The 3-node fused run ran OOM allocating an fp32 tile of tokens×V /tp values (15.16 GiB at cp1 with 32768 rows per rank) inside the fused cross-entropy (Appendix D.1). DisaggKD evaluates no model cross-entropy: its forward step returns the pre-head hidden states, and the loss module computes the LM and KD terms on windows of 1024 to 4096 rows (Appendix C.1), so the fusion option does not apply to it.

cept NeMo-RL on Qwen3.5, which ran 4096-token samples), and counted tokens by re-tokenizing the packed samples. KDFlow receives the packed text as a response to a fixed prompt without a chat template, because its 2 chat-template paths disagreed by 2 tokens per sample; its samples render to 32,682 tokens on average on Qwen3.5 (at most 32,723, under its limit of 32,768) and 4026 on Qwen3-32B. NeMo-RL’s samples average 4091 tokens. The tokens per step are in Table 19. NeMo-RL’s x-token recipe (its off-policy distillation) exports the teacher’s full-vocabulary logits into a persistent per-rank buffer whose size is proportional to the tokens per step on the Qwen3.5 pair: 30 GiB at 1.05M tokens and 121 GiB at 4.19M, which does not fit an 80 GiB GPU beside the models, as 2 failed 4.19M-token runs showed. The step of 1.05M tokens in Table 19 was the largest that ran on Qwen3.5; 4.19M tokens ran OOM and 2.1M was not attempted. KDFlow at 4.18M tokens ran OOM without sleep mode (the student’s forward, with the teacher’s weights resident) and without the chunked loss (a teacher-head chunk wanted 15 GiB with 8 GiB free); with sleep mode and a chunk of 4096 tokens its device reading is at the limit (Table 19). Systems comparison. Table 19 lists the measured runs behind Table 2. The global batch is 128 for every stack, except the 4-node DisaggKD run and the 3-node fused-W run on Qwen3.5, whose global batch is 126 so that dp=3 divides it, and NeMo-RL on Qwen3.5, whose batch is 256 sequences of 4096 tokens. DisaggKD trains on Qwen3-32B with SQ-L+K on 1 trainer node and 1 producer node of 4 TP2 replicas, so 2 nodes against 1. KDFlow and NeMo-RL are summarized by the median step (NeMo-RL on Qwen3-32B in 2 runs with medians 36.9 and 36.7 s), except KDFlow on Qwen3-32B, whose 1 s step log gives the mean of its 9 steady steps (1 step of 36 s excluded). KDFlow’s optional sleep mode, which offloads the idle model, was off on Qwen3-32B, as in its example script, so its teacher’s weights stay in GPU memory beside the student’s, whereas NeMo-RL offloads the teacher after its forward pass. On Qwen3.5, NeMo-RL’s expert layers ran without the groupedGEMM kernel. In the rows on the teacher’s output and the loss in Table 2, columns are vocabulary entries held per TP rank, tokens are the microbatch’s tokens on a GPU, and the KL is the forward KL between teacher and student next-token distributions;

Colocated-side enablers. Keeping the colocated runs on Megatron-Bridge used 3 enablers: a textextracted copy of the Qwen3.5 teacher (the release’s vision wrapper cannot ride PP), a preconverted sharded checkpoint for the Qwen3-235B teacher, and communicator pre-warming at training start; the converted and extracted teachers were not compared with the originals. Token accounting. Our first KDFlow and NeMoRL inputs held 1 document per sample, and FineWeb-Edu documents average 840 to 900 tokens, so a global batch of 128 would hold 0.11M tokens against 4.19M (Qwen3.5) or 0.52M (Qwen332B) for ModelOpt and DisaggKD, which concatenate documents into full-length sequences. We therefore ran KDFlow and NeMo-RL on the corpus packed into samples of the same length (documents concatenated, the remainder dropped; 32,768 tokens for Qwen3.5 and 4096 for Qwen3-32B, ex25

Nodes Steps Tokens (103 ) s/step Tok/s/GPU Peak Setting

Stack

Pair

ModelOpt, KD-only ModelOpt, Bal-L+K ModelOpt, SQ-L+K ModelOpt, SQ-L+K fused ModelOpt, SQ-L+K fused-W ModelOpt, SQ-L+K fused-W DisaggKD (ours) KDFlow NeMo-RL DisaggKD (ours)

Qwen3.5

4

200

4194

90.5

1450

63.4 tp2/pp4/cp4/dp1

Qwen3.5

4

200

4194 176.8

741

67.2 tp2/pp4/cp4/dp1

Qwen3.5

4

200

4194

96.3

1361

59.9 tp2/pp4/cp4/dp1

Qwen3.5

4

15

4194 100.6

1303

40.7 tp2/pp4/cp4/dp1

Qwen3.5

3

15

4129

72.0

2389

64.1 tp2/pp4/cp1/dp3

Qwen3.5

2

15

4194 146.8

1786

50.3 tp2/pp4/cp2/dp1

Qwen3.5 Qwen3.5 Qwen3.5 Qwen3.5

2 4 4 3

15 12 12 200

4194 283.7 4183 335 1047 50.7 4194 74.3

924 390 645 2352

57.5 79.1 77.0 39.2

DisaggKD (ours) Qwen3.5

4

200

4129

69.9

1846

ModelOpt, KD-only ModelOpt, Bal-L+K KDFlow NeMo-RL DisaggKD (ours)

Qwen3-32B

1

12

524

18.4

3570

38.9 tp2/pp2/dp2

Qwen3-32B

1

12

524

18.7

3500

40.0 tp2/pp2/dp2

Qwen3-32B Qwen3-32B Qwen3-32B

1 1 2

12 11 400

515 524 524

9.8 36.8 13.5

6570 1780 2426

78.9 sleep mode off 77.5 batch 128 × 4096 25.1 1+1 nodes, tp2/pp2/dp2, CR −

1+1 nodes, tp2/pp2/cp2/dp1 sleep mode, chunked loss batch 256 × 4096 2+1 nodes, tp2/pp2/cp2/dp2, CR − 33.0 3+1 nodes, tp2/pp2/cp2/dp3, CR −

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.

NeMo-RL adds a CE term (scale 0.1) and KDFlow weights KL and CE 0.5 each at learning rate 2 × 10−5 . NeMo-RL’s same-tokenizer path takes it over the 8192 teacher columns with the largest logits in the microbatch (its default, of V ≈152,000 for Qwen3-32B and 248,000 for Qwen3.5), renormalized within that subset, whereas the other 3 use all V columns.

D

Additional Results

D.1

Primary Pair

isfy CEi = − log ps,yi , X KDi = − pt,v log ps,v , v

Hi = −

X

pt,v log pt,v ,

(17)

v

KLi =

X

p

t,v pt,v log ps,v = KDi − Hi ,

v

P with LKL = i KLi . The teacher entropy Hi has no student parameters, so ∇LKL = ∇LKD , and the gradients with respect to the student logits zs of row i are

The SQ-L+K and Bal-L+K losses. For row i with label yi and one-hot vector eyi , the label loss, the soft cross-entropy, and the KL divergence sat-

∂CEi ∂zs = ps − eyi ,

∂KDi ∂KLi ∂zs = ∂zs = ps − pt .

(18)

26

With ι ∈ [0, 1] the completed fraction of the iterations, the losses of the primary-pair runs are

(a) LM loss gap to the colocated run Bal-L+K

Disagg., 4 nodes Disagg., rescaled Disagg., 4 nodes

LSQ = (1 − λι ) LLM + λι LKD , λι = 1 − 0.25 ι, LBal = LLM + γ LKL , Lres = LLM + γres LKD ,

γ = LLLM , KL

(19)

LM γres = L LKD .

LSQ is our objective (9) (SQ-L+K): the disaggregated runs, except the rescaled run, and the colocated run of Table 4 that installs it in ModelOpt with the cross-entropy form of the KD term. LBal is ModelOpt’s loss balancer (Bal-L+K) of the colocated runs, and Lres is the rescaled disaggregated run of Table 20, which applies the rescaling to our cross-entropy value; γ and γres are held constant in the backward pass. The rescaled losses are constant in value, LBal = Lres = 2 LLM , so their totals carry no KD value, and the KD term acts through the gradient. By (18), the gradients with respect to zs of row i are ∂LSQ ∂zs = ps − (1 − λι ) eyi − λι pt , ∂LBal ∂zs = (1 + γ) ps − eyi − γ pt ,

SQ-L+K

w

|gap|

signed |gap|

signed |gap|

signed

1 5 9 17 33 Max

0.020 0.014 0.012 0.012 0.011 0.056

+0.001 0.021 +0.001 0.014 +0.001 0.012 +0.001 0.010 +0.001 0.009 0.043

+0.005 0.017 +0.005 0.010 +0.006 0.009 +0.007 0.008 +0.008 0.006 0.028

+0.005 +0.005 +0.005 +0.005 +0.006

(b) KD loss, first 3 iterations subtracted Run Rise End |gap| signed Disagg., 4 nodes Coloc., KD-only Coloc., Bal-L+K Coloc., SQ-L+K

0.16 0.00 – – 0.15 +0.01 0.011 −0.010 0.14 +0.01 0.019 −0.015 0.16 +0.01 0.011 −0.010

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.

(20)

∂Lres ∂zs = (1 + γres ) ps − eyi − γres pt .

Each gradient moves ps toward the target (1 − ω) eyi + ω pt with KD weight ω = λι , γ/(1+γ), and γres /(1+γres ), at scale 1, 1+γ, and 1+γres . Per token, LKL is 0.2 to 0.35 against LLM of 2.1 to 2.2, so γ is 6.5 to 9.6 at the iterations of Table 21, the weight of Bal-L+K averages 0.887 over the run against 0.875 for λι , and the logged gradient norm at iteration 100 is 6.2 for Bal-L+K against 0.29 for the colocated SQ-L+K run and 0.63 for the disaggregated run. The logged KD values differ by Hi , about 2 nats. For the rescaled run γres = 0.96 (2.2 to 2.3 against 2.2), a constant weight of 0.49 (Table 21); its mixture is not ModelOpt’s. At λι = 1 the KD-only losses LKD and LKL have the same gradient.

the 4-node gap at w=17 is 0.008 (signed +0.005, largest 0.028). The rescaled run applies the colocated teacher precision (bf16), learning-rate floor (10−5 ), and global batch (128), and rescales the KD term to the LM loss so that its total loss is 2× the LM loss, as in ModelOpt. It runs on 3 nodes (2+1) against 4 for the colocated run, at 153.0 s per iteration against 176.8 (Tables 15 and 4), and differs from the colocated run in the grid, the kernels, the KD share (Table 21), the document order, and CR all. Table 20b gives the KD losses: the logged values differ by 2.00 nats on average (SD 0.02) from the ModelOpt runs and by 0.02 from the SQ-L+K run, and after subtracting each run’s first 3 iterations they rise by 0.14 to 0.16 and end within 0.01 of their start.

Loss parity. Table 20a gives the gap between the LM loss of a disaggregated run and that of the colocated Bal-L+K run, over moving averages of w iterations. We smooth over 17 iterations (±8) because the loss changes by 0.014 to 0.018 nats between consecutive iterations, the size of the gap at w=1; the mean absolute gap shrinks with the window, and the signed mean stays at +0.001 for the 4-node run and rises from +0.005 to +0.008 for the rescaled run. Against the colocated SQ-L+K run

Fused loss in the colocated stack. The colocated runs SQ-L+K fused and fused-W (Appendix C.2; 15 iterations each, otherwise the SQ-L+K recipe) differ from DisaggKD in reading the materialized logits. Its torch peak is 40.7 GiB and its sampled device peak 58.7 GiB, against 59.9 and 76.8 for SQ-L+K, and its median iteration takes 100.6 s over iterations 3 to 15, against 110.3 s for SQ-L+K 27

Iteration Colocated Bal-L+K SQ-L+K (disagg. and coloc.) Disaggregated, rescaled Colocated KD-only

20

than the colocated runs with ModelOpt’s loss. On the 3 larger pairs the fused-W cells at the node counts and student grids of DisaggKD run OOM during model construction, before the first forward pass: Qwen3-235B on 3 nodes (tp2/pp2/cp1/ep4, global batch 126; 74.9 GiB in use and a request of 13.5 GiB), Nemotron on 4 nodes (tp2/pp2/cp1/ep4; 79.09 GiB in use while the teacher’s expert layers are built), and Llama on 6 nodes (tp16/pp1/cp1/dp3, global batch 126; 66.0 GiB in use and a request of 16.4 GiB for the gradient buffer of the student). The colocated memory model of Appendix A.1 describes ModelOpt’s loss and does not predict these runs.

50 100 200

0.88 0.87 0.89 0.91 0.98 0.94 0.88 0.75 0.49 0.49 0.49 0.49 1 1 1 1

Table 21: Weight of the KD term in the loss of the primary-pair runs: λ for the mixture, γ/(1+γ) with γ = LM/KL for Bal-L+K and γres = LM/KD for the rescaled run.

over the same iterations (96.3 s over iterations 3 to 200). At iteration 1, with the same weights and data, the total loss of a rank is 18207.279 against 18207.285 (sums over the rank’s tokens), and the model-forward LM sum agrees with the fused label pick to 2 × 10−7 (the printed precision). The 15-iteration runs differ from the 200-iteration ones in λι and in the learning rate, whose warmup covers iterations 1 to 14 of the 15, so only iteration 1 is compared, and the 15-iteration KD losses of different layouts are not comparable. At 3 nodes (tp2/pp4/cp1/dp3, global batch 126) the fused run runs OOM in the first forward, with 66.23 GiB in use and a request of 15.16 GiB: the fp32 tile that the model’s own fused cross-entropy forms at 32768 rows per rank, to which the fused criterion gives no gradient weight. With that fusion off (fusedW) the 3-node run trains 15 iterations with the iteration-1 values of the 4-node runs (LM 2.0945, KD 2.1907), at 72.0 s per iteration (4.13M tokens, 2389 tokens/s per GPU, against 74.3 s and 2352 for DisaggKD on 3 nodes), with a torch peak of 64.1 GiB and a sampled device peak of 78.9 GiB against the 79.11 GiB limit. At 2 nodes the fusedW run (tp2/pp4/cp2/dp1, global batch 128) trains 15 iterations at 146.8 s per iteration, with a torch peak of 50.3 GiB and a sampled device peak of 65.8 GiB; with cp1 and dp2 it ran OOM after the first iteration, with 78.9 GiB in use and a request of 0.5 GiB at the windowed label pick. DisaggKD on 2 nodes (1 trainer node with tp2/pp2/cp2/dp1, 1 producer node with 4 TP2 replicas) trains 15 iterations at 283.7 s per iteration with a trainer torch peak of 57.5 GiB; the student is the bottleneck at dp1 and the producers are idle most of the time (duty 0.13, Table 15). At 2 nodes the colocated fused-W run is therefore faster and holds less memory per GPU than DisaggKD, at 3 nodes the 2 runs have a similar speed and DisaggKD holds less memory per GPU, and DisaggKD is faster

Fused-W on 2 nodes. 2 fp32 buffers with tokens×V /tp values each decide the peak of a colocated run of the primary pair, 15.16 GiB each at 32768 rows per rank (V /tp=124160) and 7.58 GiB at cp2. (1) The intermediates of ModelOpt’s logitsKL loss, on the KD side: the stock criterion computes fp32 log-probabilities of this size and saves the student-side tensors for the backward pass. The fused criterion removes them: the kernel of DisaggKD saves no fp32 tensor of this size (only the logits that ModelOpt already holds and pertoken scalars), and the label pick materializes 1 window of 1024 rows at a time. (2) The model’s own cross-entropy, on the student side: Megatron’s fused vocabulary-parallel cross-entropy in the model forward forms 1 such tensor, although the fused criterion gives that loss zero gradient weight. Turning the fusion off (W) routes it through the unfused path, which normalizes in place on its single fp32 copy of the logits. Table 22 lists the 3-node cells, which are 3 different runs (the KDonly run has no LM term, hence no buffer (2)): the failing allocation moves from the loss to the model forward, and with the fusion off the fused run fits. At 2 nodes the same student grid with cp1 and dp2 holds 32768 rows per rank and ran OOM after the first iteration at a request of 0.5 GiB (64.9 GiB allocated, 78.9 GiB in use), whereas cp2 and dp1 halves every buffer of this size, fits at 50.3 GiB allocated (65.8 GiB sampled device peak), and trains at 146.8 s per iteration. DisaggKD on the same 2 nodes trains the student on 8 GPUs while the producer node waits for the trainer, and takes 283.7 s (Table 15). Loss parity under the offset. Figure 3a plots the LM loss, the student’s next-token cross-entropy, which every run defines in the same way whatever 28

Loss

with the same shuffle seed, so a reconstruction of both data pipelines from the corpus index gives different document sequences from the first iteration (Table 25, seed-42 runs). The parity statistic therefore compares runs on the same corpus, not on the same samples.

Alloc. Req. In use Failed in

KD-only, ModelOpt SQ-L+K fused

72.9 15.2

75.5 the loss

57.7 15.2

SQ-L+K fused-W

64.1

66.2 model forward 78.9 trains

−

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 fusedW). Run Colocated, cp4/dp1 Colocated fused, cp4/dp1 Colocated fused-W, cp1/dp3 Colocated fused-W, cp2/dp1 Disaggregated, cp2/dp1 Disaggregated, cp2/dp3 Disaggregated, cp2/dp2

Nodes

Source of the speedup. On the 4-node Qwen3.5 pair the colocated Bal-L+K, colocated SQ-L+K, colocated KD-only, and disaggregated runs cost 5656, 3082, 2894, and 2237 GPU-seconds per iteration (median iteration time times the 8 GPUs of each of 4 nodes, producer node included; the colocated SQ-L+K fused run costs 3219 and the 3node fused-W run 1728, Table 23). At seed 42 the 3419 GPU-second gap to the colocated Bal-L+K run is the colocated LM term, 2762 GPU-seconds (81%; 68% with the seed-43 colocated runs and the 61.9 s pace), plus the 657 (19%) by which the colocated KD-only run still exceeds the disaggregated run, which also computes the LM term. With SQ-L+K the colocated LM term costs 188 GPUseconds, and the gap to the disaggregated run is 845, of which 657 is the KD-only excess. Table 24a locates the added events: on the last pipeline stage, which evaluates the loss, the LM term adds NCCL collective events and a chain of small copy and conversion operators with their stream synchronizations, and the GPU-kernel share of traced time falls (85.9% to 77.6%), while stage 1 is unchanged (98.0% and 98.5%). The KD-only gap has a different source. The colocated stack executes the teacher forward pass and the student training step sequentially on the same 4 nodes, so the iteration time is their sum. The streamed SQ-L+K trainer, which executes only the student step on the same grid, takes 70.9 s against the colocated 90.5 s (seed 42) and 97.7 s (seed 43, different nodes), 19.6 and 26.8 s less, against a teacher-forward estimate of 16 s from the producer rate (Table 24d); at 8192 tokens on the same grid the gap is 59.7 s (Table 24c). In the 3+1 layout the producer node executes the teacher forward pass while the trainer nodes execute the student step. The producers’ forward time is within 10% of the iteration time (Table 15), so the iteration time is the longer of the producers’ forward time and the trainers’ step time, not their sum. Grid. With ModelOpt’s loss, 3 trainer nodes cannot host the colocated grid, which needs 4 nodes per replica; with the fused-W loss the colocated stack runs the narrow grid tp2/pp4/cp1/dp3 on 3

s/it GPU-s Peak

4 96.3 4 100.6† 3 72.0†

3082 59.9 3219 40.7 1728 64.1

2 146.8†

2349 50.3

1+1 283.7† 3+1 69.9 2+1 74.3

4539 57.5 2237 33.0 1783 39.2

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.

its training weights, so the curves measure one quantity. The weights change the trajectory, not the definition, so we measured how much: the 4node disaggregated run (mean KD weight 0.875) and the rescaled 2+1 run (0.49, which differs from the 4-node disaggregated run in node split, global batch, learning-rate floor, teacher precision, and CR) lie 0.012 and 0.010 nats from the colocated run. Figure 3b plots KD losses that are defined differently, cross-entropy against KL; their offset is the teacher’s entropy (the logged difference is 2.00 nats; no run logged the entropy), which we remove by subtracting each run’s first-3-iteration mean. The plotted disaggregated run is the mixture, whose weight is the closer to ModelOpt’s. A comparison at equal share would apply the balancer to the KL value, which needs the teacher’s entropy per token and is not part of these runs. The runs also differ in their data order: the colocated runs hold out 0.01% of the documents for validation (a 9999:1:0 split) before shuffling the document order, and the disaggregated runs read the whole corpus

29

(a) Profile of the colocated runs

(b) Streamed trainer on other grids KD-only Bal-L+K

Loss stage, GPU-kernel share (%) Loss stage, NCCL collectives (s) Loss stage, copy operators (s) Loss stage, stream syncs (s) Stage 1, GPU-kernel share (%)

Trainer grid

Nodes

s/it GPU-s Relative CR

85.9

77.6

tp2/pp2/cp2/dp4

4+2 39.3

1886

1.00

−

2.3

18.3

tp2/pp4/cp2/dp2

4+2 50.8

2438

1.29

−

7.1 13.2 98.0

23.4 30.4 98.5

tp2/pp2/cp4/dp2 tp2/pp4/cp4/dp1 tp2/pp2/cp2/dp3

4+2 50.0 4+2 70.9 3+1 65.0

2400 3403 2080

1.27 1.80 −

− − −

(c) KD-only at 8192-token sequences Stack

Grid

Nodes

Colocated Colocated Streamed Streamed

tp2/pp4/cp4/dp1 tp2/pp2/cp2/dp4 tp2/pp2/cp2/dp3 tp2/pp4/cp4/dp1

s/it GPU-s Peak

4 137.2 4 67.5 3+1 114.7 4+2 77.5

4389 2161 3672 3720

37.4 65.7 32.1 30.8

(d) Colocated KD-only iteration at 32768 tokens (tp2/pp4/cp4/dp1) Run

Part of the iteration

s/it GPU-s

Colocated, seed 42 Colocated, seed 43 forward-only difference Streamed trainer, 4+2 nodes Producers, estimated

all all both forward passes and the loss training step minus evaluation step all teacher forward

90.5 97.7 77.9 19.8 70.9 16

2894 3126 − − 3403 −

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 the whole corpus, same shuffle seed

0 of 128

Disaggregated

and needs fewer GPU-seconds than the split. This does not extend to 32768 tokens, where the model predicts ModelOpt’s loss over the budget on the narrow grid (the measured cp2 OOM of Table 18 is on pp 4), and the same 3+1 layout takes 65.0 to 69.9 s for 4.13M tokens at 32768 tokens per sequence against 114.7 s at 8192 (Tables 24c and 15). The colocated runs hold the teacher in bf16 (Appendix B.2). Holding the routed experts in fp8 lowers the colocated stack’s torch peak at its cp4 grid, but the KD loss and gradient norm of that run leave those of the bf16 run at iteration 3 (Table 16), and the model predicts the cp2 grid over the budget with fp8 experts as well (no fp8 run at cp2).

−

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).

nodes (Table 23). The grid with the smaller cp is faster for the streamed trainer at fixed student GPUs (Table 24b) and keeps the trainers’ step time within 10% of the producers’ forward time in the 3+1 runs with CR − (Table 15). At 8192 tokens the colocated stack fits the 4-node grids of Table 24c

Seeds. Table 26 gives the primary-pair runs with seed 43 for the data order and random state. The 30

Disagg. (ours), 4 nodes Disagg. (ours), 3 nodes Coloc., Bal-L+K Coloc., KD-only

Seed 42 Seed 43 Change

p −6

68.52

70.09

+1.57

70.14

–

–

67.51

69.78

+2.26 9×10−13

67.97

69.22

+1.25

10−4

–

+1.01

0.001

–

+0.56

0.09

–

+1.62

10−8

Between layouts, seed 42 4 nodes − – coloc. Bal-L+K 4 nodes − – coloc. KD-only 3 nodes − 4 – nodes

6×10

evaluates the loss, releases after every microbatch; iter, once at each iteration boundary; and producer, where each producer replica releases after each of its first 400 slots. 6 of the 8 cells over the 3 parts on the Qwen3-32B pair at 4096 tokens (Table 27, 1 run each) put 54.6 of the 55.3 s on the per-microbatch part, 0.85 s per microbatch; the iteration part adds 1.4 s, and the producer part, active only in the first 12 iterations, adds 0.0 s over iterations 8 to 30. The CR column of the run tables records which parts were on: − for none, mb, iter, producer, combinations such as mb+iter, and all for the 3. It changes the pace; the mean KD loss over iterations 8 to 30 differs by at most 0.02 nats across the rows of Table 27. On the seed-43 configuration the release costs at least 0.77 s per microbatch (42 per DP rank; the CR − run is producer-limited; Table 15). The 200-iteration seed-43 run has it on (Table 15), and the seed-43 pace in Table 4 is the run with CR −.

–

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

− mb iter mb+iter producer all

13.5 – 68.1 +54.6 14.9 +1.4 69.3 +55.8 13.5 +0.0 68.8 +55.3

Effect KD loss

D.2

2.42 2.42 2.41 2.42 2.40 2.41

Other Pairs

Divergence control on Qwen3-235B. The 3 preserved colocated Qwen3-235B runs on 8 nodes (cp 4, ep 16, all-to-all) ended in Inf or NaN after 18, 44, and 58 completed iterations (Table 28). The table compares them with runs of the same student on other grids, with the dense Qwen3-4B teacher at 4096 tokens in the 4B cells. The student’s LM loss at the second iteration is within 0.03 of the Hugging Face loss at ep 4 with tp 2 and at ep 8 with tp 1 and tp 2, and above it at ep 16 with tp 1 and tp 2; the gradient norm of the first iteration rises from 29 to 105 and 106 . Replacing only the allto-all token dispatcher by the all-gather dispatcher at the tp2/pp1/ep16 grid gives the loss (2.178) and gradient norm (29) of the ep-4 cell, whereas sequential experts and an fp32 router do not; on the 235B pair it gives a second-iteration loss of 2.11 and a first-iteration gradient norm of 12, against 1.0 to 1.7 × 106 with all-to-all. The checked weights match the Hugging Face tensors, and the output of the first layer is 43% above that of the healthy run at ep 16 in the tp2 run (6% below at tp 1; layer 0 was not measured in the all-gather cell). At ep 16 the all-to-all dispatcher gives a wrong forward pass and the all-gather dispatcher on the same grid does not, and the tp1/pp2 (edp 1) and tp2/pp1 (edp 2) grids are both above the reference, as is the cp-4 grid with the 235B teacher. With the all-gather dispatcher the colocated 8-node run differs from the learning-rate-4 × 10−4 run that diverged after 18 iterations in the dispatcher and the learning-rate

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).

seed-43 runs score 1.25 to 2.26 points above the seed-42 runs, more than the seed-42 lead of the 4-node disaggregated student (+1.01, +0.56); the 3- and 4-node students, which differ in node split, global batch, and trainer grid, differ by 1.62 points, and 2 seeds do not establish a ranking of the layouts. The seed-43 colocated runs take 97.7 s (KD-only) and 172.1 s (Bal-L+K) per iteration, on different nodes. Cache release. The allocator’s cached blocks can be returned to the driver (the cache release). The release has 3 parts, each with its own switch: mb, where the last pipeline stage of the trainer, which 31

Run (30B student)

tp/pp/cp/ep edp Dispatcher

LM, it. 2

Grad norm, it. 1 Layer 0 Outcome

Torch forward Disaggregated, 235B Disaggregated, 235B Colocated, 4B Colocated, 4B Colocated, 4B Colocated, 4B Colocated, 4B Colocated, 4B Colocated, 4B, seq. experts Colocated, 4B, fp32 router

− 2/2/1/8 2/2/1/4 2/2/1/4 1/2/1/8 2/2/1/8 1/2/1/16 2/1/1/16 2/1/1/16 2/1/1/16

− − 2 flex 2 flex 4 all-to-all 2 all-to-all 2 all-to-all 1 all-to-all 2 all-to-all 2 all-gather 2 all-to-all

2/1/1/16

2 all-to-all

−

−

Colocated, 235B, 3 runs Colocated, 235B

2/2/4/16

2 all-to-all

2.63 to 2.68

1.0 to 1.7 × 106

2/2/4/16

2 all-gather

2.11

12

2.16 − 2.11 2.2 2.11 1.7 2.18 29 2.18 29 2.18 29 2.70 1.2 × 105 3.15 to 3.22 1 × 105 to 1.4 × 106 2.18 29 3.11 7.8 × 105

− reference − stable, 20 it. − stable, 200 it. 77.4 finite, 2 it. − finite, 2 it. − finite, 2 it. 72.9 above reference 110.9 losses rise − finite, 2 it. − above reference − NaN in the forward loss, first iteration − Inf/NaN, 18 to 58 it. − finite, 60 it.

Table 28: The Qwen3-30B-A3B student on different grids, 3 to 8 nodes (edp: expert data-parallel size, GPUs divided by ep·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 × 10−8 relative.

schedule (the 2 runs share their first 18 iterations, after which the 60-iteration schedule decays the learning rate sooner), and completes its 60 planned iterations with finite losses (the longest all-to-all run failed after 58) (Table 28; memory and pace in Table 18). Disaggregated on 4+4 nodes with the KD-only loss (λ=1, the LM term skipped), peak learning rate 4 × 10−4 , and the 200-iteration schedule, the pair completes the 60 iterations of a run stopped there with no NaN and no skipped iteration, at 105 s per iteration (Table 17). The disaggregated SQ-L+K runs of Table 17 on 4+4 and on 2+1 nodes run 200 iterations each without a NaN.

Layout Teacher Grid

edp Dispatcher Outcome

Disagg. Ultra Coloc. Ultra

2/2/1/4 2/4/8/16

4 flex 1 all-to-all

Coloc. Ultra

4/4/4/16

1 all-to-all

Coloc. Ultra

4/4/4/16

1 all-gather

Coloc. Nano

4/4/4/16

1 all-gather

Coloc. Nano

4/4/1/16

1 all-gather

Coloc. Nano

4/4/4/4

4 all-gather

finite, 200 it. NaN, gradient norm NaN, forward loss NaN, gradient norm NaN, gradient norm NaN, gradient norm finite, 3 it.

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 × 10−5 , except the 2/4/8/16 row (4 × 10−4 ).

NaN on the Nemotron pair. Table 29 lists the colocated Nemotron runs next to the disaggregated one. The student gives NaN in its first iteration in each of the 5 colocated ep-16 cells (tp 2 or 4, pp 4, cp 1, 4, or 8), with the all-to-all and all-gather dispatchers and the Ultra and Nano teachers, and stays finite for 3 iterations at ep 4 (edp 4) with the Nano teacher, the all-gather dispatcher, and cp 4; every ep-16 cell has edp 1. The all-to-all dispatcher gives NaN already in the forward loss at tp4/cp4. The Nano-teacher rows run the Nano student with its own initial checkpoint as teacher, so ModelOpt’s KL stays in the loss and only the Ultra teacher is

removed. Llama pair. Table 30a lists the Llama runs. At 16 nodes each layout has a training run and a pace check. The pace of each layout was measured again in a 20-iteration pace check of the same layout. The colocated pace moves by 1.3% between the 2. The 32

(a) Llama Layout (nodes)

(b) Qwen3-32B to 4B Iter. s/it GPU-s Peak CR MMLU Seq. × batch

Colocated (8) 5 484 Colocated (16), training 200 301 pace check 20 305 Disagg. (ours, 4+2) 5 818 Disagg. (ours, 8+8), 400 186 training pace check 20 235 Teacher Initial student

36.7 – 27.5 – 27.4 – 68.6 all 50.3 −

− 76.98 − − 76.41

30.1k 50.3 all

− 85.21 78.17

31.0k 38.5k 39.0k 39.3k 23.8k

4096 × 128 4096 × 128 8192 × 64 16384 × 32

Microbatches

s/it Tokens/s/GPU CR

64 13.5 64 69.6 32 11.9 16 13.5

2426 471 2755 2427

− all − −

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. × 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.

pace check ran with CR all and is 48.7 s slower, 0.76 s per microbatch over 64 microbatches per DP rank; the release costs 0.77 s per microbatch on the primary pair, and no Llama run isolates it. The speedup at 16 nodes is 1.6× from the training runs, with CR −, and 1.3× from the pace checks. At 8 nodes the colocated layout needs 31.0k GPUseconds per iteration, 3% more than the disaggregated pace check at 16 nodes (30.1k) and 30% more than its training run (23.8k); the disaggregated layout runs on 6 nodes (39.3k GPU-seconds, 5 iterations) where the colocated layout needs 8. Both students score within 1.8 points of the initial student on MMLU (MMLU in Table 30a).

(Llama 3.1 Community License). The corpus is FineWeb-Edu (ODC-By 1.0) and the benchmark is MMLU (MIT License). The software is MegatronBridge, Megatron-Core, Model Optimizer, NeMoRL, Transformer Engine, and the evaluation tool (Apache License 2.0), and KDFlow and Triton (MIT License). All were used for research distillation, as their licenses permit, and the trained students are not released. FineWeb-Edu is a filtered web corpus and MMLU a multiple-choice benchmark; neither is annotated by us. An anonymized codebase accompanies the submission: it holds the code, the run configuration behind each table, the memory and traffic models with their tests, the scripts that produce the tables and figure data, and a standalone MMLU protocol; it holds no run data, so the measured tables are reproduced by rerunning the configurations on the stated hardware. The runs reported in this paper used about 15,000 GPU-hours on NVIDIA H100 80 GB GPUs, including the baselines and the MMLU evaluations but not failed attempts or memory and NaN probes.

Qwen3-32B pair. On Qwen3-32B to 4B the disaggregated run with CR − is 5.1 to 5.2× faster per GPU than with CR all (Tables 27 and 30b), and stays below ModelOpt and KDFlow (Table 2). Table 30b holds the tokens per step fixed at 0.52M and varies the sequence length: with CR − the iteration time is 13.5, 11.9, and 13.5 s at 64, 32, and 16 microbatches per replica, whereas with CR all the 4096-token configuration pays 0.86 s per microbatch (0.85 s for the per-microbatch part alone; Table 27). In the 4096-token CR − run the producers’ duty is 0.93 (Table 15).

E

F

Glossary of Symbols

Subscripts t and s denote the teacher and the student, and x ∈ {t, s} stands for either. Superscripts (q), (j), and (q, j) denote the cp rank, the tp rank, and the cell of both. Indices start at 0. Quantities in bytes carry the unit explicitly, and GiB is 230 bytes.

Artifacts, Licenses, and Compute

The models are Qwen3.5, Qwen3, and Qwen34B-Base (Apache License 2.0), Nemotron 3 Nano (NVIDIA Nemotron Open Model License) and Ultra (OpenMDW License 1.1), and Llama 3.1

Models and sizes (§3). Dx , Ex Dense and routed-expert parameters of model x; Pt = Dt +Et ℓx Layers of model x

33

lin moe ℓatt x , ℓx , ℓx

dx (dt , ds ) V , Vloc h, hkv dhead ke w

Agc , Abwd , Amoe

Full-attention, linear-attention, and MoE layers of model x Hidden width of model x Vocabulary size and its share V /tp per tp rank Attention heads and key-value heads Width of one attention head Routed experts per token Per-token width of the linear-attention context exchange (Table 13); in Table 20, the width in iterations of the moving average

ℓrc φ p , ω p , βe

Mcol , Mtr , Mpr ρrep

Parallelism and layout (§3). tp, pp, cp, dp, ep, Tensor-, pipeline-, context-, data-, etp expert-, and expert-tensor-parallel degrees edp Data-parallel size of the routed experts, dp cp tp/ep at etp=1 g, G, g ⋆ A parallelism grid, the set of legal grids, and the node-minimal grid that fits N (g), N Nodes of grid g; number of nodes, each of 8 GPUs Nstu , Nprod Trainer nodes and producer nodes R Producer replicas of the teacher tpp , epp Tensor and expert span of one producer replica M Per-GPU memory budget, 79.11 GiB L Tokens (rows) per sequence; L = cp · r gbs Global batch size in sequences Ntok Tokens per iteration, gbs · L τ Tokens handled by one GPU per iteration, Ntok /(dp cp) nmb Microbatches per DP replica, gbs/dp µ Microbatches in flight across the pipeline b Bytes per activation element, 2 n Degree of a parallel dimension (Table 13); number of parts in Lemma 1 ∆, n∆ Distance between the ranks of a dimension in Megatron’s rank order and the number of ranks it spans; the dimension uses NVLink if n∆ ≤ 8, the GPUs of a node, and InfiniBand otherwise

Titer

Trainer reference-cycle garbage, loss-backward and recompute temporary memory, and expert-layer block Layers recomputed at once Producer forward-pass activation and workspace memory, its NCCL buffers, and bytes per routed-expert parameter (1 for fp8, 2 for bf16) Predicted peak memory per GPU of the colocated stack, the trainer, and the producer Forward-only token rate of one producer replica (measurement per pair in Appendix B.1) Iteration time

Traffic model (§3.2). pp cp Φtp Bytes sent per GPU and iteration by TP x , Φx , Φx , ep Φx collectives, pipeline activation transfers (per stage boundary and direction), CP exchanges, and expert dispatch and combine collectives of model x Φdp Bytes sent per GPU by the student’s s gradient reduction and parameter gather Φxfer Bytes per last-stage rank and iteration of the teacher-to-student transfer νt , νs All-reduces per layer: 2 for the teacher’s forward pass, 6 for the student’s forward, recompute, and backward passes Algorithm (§4). q, j, (q, j) cp rank, tp rank, and the cell (last-stage rank) with both i Row (token) index r Rows per slot, one cp rank’s share of a sequence K, C, k Windows per slot, rows per window (r = K · C), and window index; K also denotes top-K in Sections 2 and 5 u, v Entry index within a tp rank’s vocabulary shard and column index over the full vocabulary η, η0 Slot index in the producer’s stream and the start slot of a resume ngroups Receiver groups, one per DP replica (dp) δ Iterations the producer may run ahead of the trainer (q) (q) Hx , Hx,k , Hx Final hidden states of model x for the rows of cp rank q (the slot), its window k, and all L rows (j) Wx , Wx lm_head of model x and its column shard on tp rank j

Memory model (§3.1). S Student weights, gradients, and optimizer state per GPU Ploc Student parameters per GPU Acol , Astr Activation memory in the colocated stack and activation and temporary memory of the streamed trainer (both drawn as A in Figure 1) Tcol Teacher weight shard per GPU in the colocated layout κcol Colocated logits buffers κstr Memory retained by the streamed loss, the windowed fused loss on streamed hidden states σcol , σstr Non-PyTorch memory (CUDA context, NCCL, allocator) of the colocated stack and the trainer

(q,j)

Zx

(q,j)

, Zx,k

Block of the logits of model x held by cell (q, j), and its window k

(j)

Logits of one row, their shard on tp rank j, and the entry of column v

zx , zx , zx,v

Loss and statistics (§4.2). px , px,v Softmax of zx over the V columns and its entry for column v yi , eyi Label of row i and its one-hot vector L, λ Training loss and the weight of the KD term

34

γ, γres LSQ , LBal , Lres

LLM , LKD KLi , eyi , ω LKL , Hi CEi , KDi

lses mx , ax , c

(j)

(j)

mx , ax , c(j) mx,u , ax,u , cu mx,S , ax,S , cS S, Pl , Uu |S| G(j) R, ⊤, e [v=yi ] p

Ratio of the LM loss to the KL loss (ModelOpt’s balancer) or to the KD loss (rescaled run) in a batch The SQ-L+K, Bal-L+K, and rescaled-run losses (19); λι : KD weight at the completed fraction ι of the iterations Language-modeling and P P KD terms of the loss, i CEi and i KDi KL divergence of row i, the one-hot label vector, and the KD weight of the target (1 − ω)eyi + ωpt P i KL(pt ∥ ps ) and the teacher’s entropy of row i Label loss and distillation loss (soft cross-entropy) of row i; KL(pt ∥ ps ) differs from KDi by the teacher’s entropy Log-sum-exp of the student’s logits of a row Row maximum of zx , sum of ez−m , and teacher-weighted sum of student logits The same statistics over the shard of tp rank j Scan state after entry u of a shard The statistics over a set S of columns A set of columns, its n disjoint parts (index l), and the tiles read in one scan (U tiles) Number of columns in S Gradient of L with respect to the window’s student logits on tp rank j Real numbers, matrix transpose, and the base of the natural logarithm 1 if v = yi and 0 otherwise p-value of the paired McNemar test (Table 26), not the softmax

35

Related documents

• DisaggKD: Efficient Knowledge Distillation through Teacher-Student Disaggregation and a Fused Loss#1342856
• Knowledge Distillation in Wide Neural Networks: Risk Bound, Data Efficiency and Imperfect Teacher#974116
• Distillation as Probability Transport: Routed On-Policy Distillation#985698
• Persistent Teacher Anchoring for Tool-Using Agents#971442
• Clean Data Can Still Carry Backdoors: Support-Persistent Backdoors for Model Reuse#1356544