Training loss behaves strangely in mixed-precision training

pepper8362: <p>I have a loss function looks like <code>loss = loss1 + alpha * loss2</code> where <code>loss1</code> and <code>loss2</code> are two loss functions and <code>alpha</code> is a scalar weight. I have the following observations:</p>
<p>(1) If I set <code>loss=loss1</code>, the model trains fine</p>
<p>(2) If I set ‘loss = loss1 + alpha * loss2’ with <code>alpha=0</code>, the traning interrupts with cuda error after a couple of epoches (usually between 10 and 20). Sometimes <code>loss</code> goes up to a high value before the error.</p>
<p>I have done serveal runs to confirm the above observations. This seems strange to me, as the loss in (2) should be equivalent to the one in (1). Could mixed-precision training cause something like this? Thanks!</p>

ptrblck: <p>I don’t see how mixed-precision training is involved in your use case so could you post a minimal and executable code snippet reproducing the issue, please?</p>

pepper8362: <p>Thank you for your reply. The code is part of a complex project, and the error happens when training on a large dataset. I am not sure if the error could be reproduced without these data, even if I extract the relevant code snippet (I can confirm the data is not the source of the problem though). Briefly speaking, it is actually the code I asked <a href="https://discuss.pytorch.org/t/model-distillation-with-mixed-precision-training/189546">here</a>. Sorry for duplicate but I just include the code below for convenience:</p>
<pre><code class="lang-auto">with torch.cuda.amp.autocast(enabled=True):
    outputs_student = model_student(inputs, targets)
    with torch.no_grad():
        outputs_teacher = model_teacher(inputs, targets)
    loss_distillation = distill_loss(outputs_student, outputs_teacher)
    loss_student = some_loss(outputs_student)    
    loss = loss_student + weight*loss_distillation 
optimizer.zero_grad()
scaler.scale(loss).backward()
scaler.step(self.optimizer)
scaler.update()
</code></pre>
<p>The error message is something like</p>
<pre><code class="lang-auto">/opt/conda/conda-bld/pytorch_1639180549130/work/aten/src/ATen/native/cuda/Loss.cu:115: operator(): block: [87,0,0], thread: [59,0,0] Assertion `input_val &gt;= zero &amp;&amp; input_val &lt;= one` failed.
/opt/conda/conda-bld/pytorch_1639180549130/work/aten/src/ATen/native/cuda/Loss.cu:115: operator(): block: [87,0,0], thread: [60,0,0] Assertion `input_val &gt;= zero &amp;&amp; input_val &lt;= one` failed.
/opt/conda/conda-bld/pytorch_1639180549130/work/aten/src/ATen/native/cuda/Loss.cu:115: operator(): block: [87,0,0], thread: [61,0,0] Assertion `input_val &gt;= zero &amp;&amp; input_val &lt;= one` failed.
/opt/conda/conda-bld/pytorch_1639180549130/work/aten/src/ATen/native/cuda/Loss.cu:115: operator(): block: [87,0,0], thread: [62,0,0] Assertion `input_val &gt;= zero &amp;&amp; input_val &lt;= one` failed.
/opt/conda/conda-bld/pytorch_1639180549130/work/aten/src/ATen/native/cuda/Loss.cu:115: operator(): block: [87,0,0], thread: [63,0,0] Assertion `input_val &gt;= zero &amp;&amp; input_val &lt;= one` failed.
</code></pre>
<p>Thanks!</p>

ptrblck: <p>I don’t know which criterion you are using but most likely the model output and/or target tensor contains out-of-bounds values. Check the min/max values of all tensors before passing them to the loss function and make sure their ranges are valid.</p>

pepper8362: <p>Thanks! I tried to save relevant variables when the crash happened with try catch. But it seems the cuda error caused the process to terminate before all data were dumped to the disk (I got EOFError: Ran out of input when I read in the dumped pkl file). Is there anyway to avoid this?</p>
<p>Btw: <code>distill_loss</code> is a sum of three <code>torch.nn.functional.mse_loss</code>. <code>loss_student</code>  is a combination of <code>        torch.nn.BCEWithLogitsLoss</code> and a self-defined iou loss.</p>

ptrblck: <p>The assert could corrupt the CUDA context so try to print the values beforehand or add an <code>assert</code> statement yourself checking the values in <code>[0, 1]</code>.</p>