[Knowledge Distillation]: NaNs when combining CE + aux MSE in the same backward (BatchNorm state mismatch)

username: Mangoh
cooked: <p>Hi all,</p>
<p>I’m implementing <strong>knowledge-distillation</strong> of a <em>Student</em> ResNet-50 from a <em>Teacher</em> ResNet-101 for speaker-Verification task.<br>
The total loss I want is</p>
<pre><code class="lang-auto">L = α · CE(logits_student, y)        # classification
  + (1-α) · MSE( features_student[layer3,4],
                 features_teacher[layer3,4] )
</code></pre>
<p>Both networks use <strong>SyncBatchNorm</strong> and I train in <strong>distributed mixed precision</strong> (<code>autocast</code> + <code>GradScaler</code>).</p>
<hr>
<h3><a name="p-470371-h-1-what-goes-wrong-if-i-do-the-nave-way-1" class="anchor" href="#p-470371-h-1-what-goes-wrong-if-i-do-the-nave-way-1"></a>1. What goes wrong if I do “the naïve way”</h3>
<pre data-code-wrap="python"><code class="lang-python"># ↓ one forward pass that produces both terms
logits, f3, f4 = student_forward_return_everything(x)

ce  = criterion_CE(logits, y)
with torch.no_grad():
    t3, t4 = teacher_feats(x)        # teacher frozen

mse = 0.5 * F.mse_loss(f3, t3) + 0.5 * F.mse_loss(f4, t4)

loss = α*ce + (1-α)*mse
scaler.scale(loss).backward()        #  &lt;-- NaNs here!
</code></pre>
<p><em>Error</em> (excerpt):</p>
<pre><code class="lang-auto">RuntimeError: Function 'ConvolutionBackward0' returned nan values in its 1th output
</code></pre>
<p><strong>Observation</strong></p>
<ul>
<li>The first term (<code>ce</code>) is computed from the <em>normal</em> forward that ends with the classifier head (uses BN running-stats <strong>V1</strong>).</li>
<li>The second term (<code>mse</code>) relies on <em>intermediate</em> activations.<br>
<strong>When I reuse those same activations</strong> in the joint loss, the backward runs through BNs whose buffers are now <strong>updated again</strong> by the CE pass (stats ≠ saved stats) → NaNs / <code>Saved tensors...</code> errors.</li>
</ul>
<hr>
<h3><a name="p-470371-h-2-workaround-that-does-train-stably-2" class="anchor" href="#p-470371-h-2-workaround-that-does-train-stably-2"></a>2. Workaround that does train stably</h3>
<p>I split the optimisation in <strong>two forwards / two backward calls</strong> so each backward sees its own, consistent BN stats.</p>
<pre data-code-wrap="python"><code class="lang-python">manual_step = True          # tells the main loop not to call step() for me
optimizer.zero_grad()

# ---------- pass 1 : CE ----------
logits = student(x, y)                    # logits WITH margin etc.
ce_loss = criterion_CE(logits, y)
scaler.scale(α * ce_loss).backward(retain_graph=True)

# ---------- pass 2 : feature MSE ----------
with torch.no_grad():
    t_feats = teacher.extract_intermediate_features(x, True)

s_feats = student.module.extract_intermediate_features(x, False)
mse  = 0.5*F.mse_loss(s_feats[3], t_feats[3])
mse += 0.5*F.mse_loss(s_feats[4], t_feats[4])

scaler.scale((1-α) * mse).backward()

# step
scaler.unscale_(optimizer)
torch.nn.utils.clip_grad_norm_(student.parameters(), 1.0)
scaler.step(optimizer)
scaler.update()
</code></pre>
<p><em>No more NaNs – training converges.</em></p>
<p>I applied the same idea to a <strong>Hinton KD</strong> (CE + KL) that needs logits <em>with</em> and <em>without</em> AM-Softmax margin; two passes work fine there as well.</p>
<hr>
<h3><a name="p-470371-h-3-questions-3" class="anchor" href="#p-470371-h-3-questions-3"></a>3. Question(s)</h3>
<ol>
<li><strong>Is the “two-pass / two-backward” pattern the right way</strong> to keep BatchNorm happy in this scenario,<br>
or is there a better recipe to compute a <em>combined</em> loss and still avoid BN-buffer mismatch?</li>
<li>Any other tricks (e.g., hooks that cache intermediate features, gradient checkpointing, gradient scaling gotchas) that would let me simplify this?</li>
</ol>
<p>Any insight appreciated – I just want to know whether my fix is <em>the</em> standard way, or if the community has a cleaner pattern ?</p>
<p>Thanks!</p>