DisaggKD moves the teacher onto separate nodes and speeds up distillation of Qwen3.5-35B-A3B into 9B

The team behind DisaggKD proposes removing a familiar constraint from knowledge distillation: the teacher and the student no longer have to live on the same GPUs. On the Qwen3.5-35B-A3B (teacher) and 9B (student) pair, the disaggregated setup reached the same LM loss as a colocated run using the same loss, and finished faster. On larger pairs, colocated runs run out of memory, while disaggregated ones keep training.

Why colocation hits a memory wall

In most KD frameworks, the teacher and the student share the same accelerators. Each card has to hold the teacher's weights, the student's weights, gradients, and optimizer states. The stronger the teacher, the less room is left for everything else, and the teacher's size becomes bounded by the memory of a single GPU.

The second problem is hidden in the loss itself. Classic KD compares the teacher and student distributions over the entire vocabulary. That requires logits of size "tokens × vocabulary", and modern models have large vocabularies. This matrix becomes one of the biggest consumers of memory at peak.

How DisaggKD works

The authors break the design into three parts.

First: splitting across nodes. Teacher replicas run separately and pass hidden states to the student. Logits do not travel over the network.

Second: fused loss. The KD loss is computed over windows of rows and vocabulary shards. In exact arithmetic, the result equals the full-vocabulary loss, but the full logits matrix is never formed. The equivalence is claimed specifically for exact arithmetic. In practice, the summation order in floating point can produce small discrepancies, so the authors compare real training curves separately.

Third: prefetch ring. This is a ring buffer that matches the pace of the producer (the teacher) and the consumer (the trainer). While the student computes a step, the next hidden states are already on their way to it, so neither side sits idle.

The authors also provide closed-form models of memory and network traffic. With them, you can estimate in advance where the memory limit lies for each layout (colocated or disaggregated).

What the measurements showed

• Qwen3.5-35B-A3B → 9B pair: the disaggregated setup reaches the same LM loss as the colocated one with the same loss, and runs faster than it.
• Fused loss is useful on its own. A colocated run with it can use fewer nodes and run faster.
• If the student spans several nodes, the disaggregated setup is more memory-efficient.
• On larger pairs, colocated runs fail with OOM even with fused loss: the teacher's weights and the student's states eat the entire budget. The disaggregated setup trains on these pairs.

The abstract gives no absolute speedup figures, GPU counts, or sizes for the larger pairs. You will have to check the size of the gain in the paper itself and verify it against your own hardware.

Who this is useful for

The authors' findings point to a clear fork in the road. If the teacher is small and the student fits on one or two nodes, a colocated setup with fused loss is enough: it already saves memory and speeds up training. A disaggregated layout is worth considering when the teacher has tens of billions of parameters (like 35B-A3B with its experts) or the student is spread across several nodes. In that case, the memory balance decides whether the run is possible at all.

There is a cost, too. The disaggregated setup requires fast links between nodes: hidden states must be streamed without pauses. That is why the paper's traffic models are practical: before renting a cluster, you can plug your numbers into them and check whether the network will become a bottleneck.

The fused loss idea is useful in its own right. It can be ported into existing KD pipelines without changing the cluster topology, and it immediately removes the logits memory peak. For teams distilling large open models into compact ones, it is the cheapest way to test the approach on their own data.