Torchvision model vit_b_16 fails to train with AMP

pallgeuer: <p>I am trying to use the <code>vit_b_16</code> torchvision model with AMP (<code>torch.cuda.amp</code> / <code>torch.autocast</code>):<br>
<a href="https://pytorch.org/vision/main/models/generated/torchvision.models.vit_b_16.html" class="onebox" target="_blank" rel="noopener nofollow ugc">https://pytorch.org/vision/main/models/generated/torchvision.models.vit_b_16.html</a></p>
<p>I am encountering the error:</p>
<pre><code class="lang-auto">Traceback (most recent call last):
  File "/home/phil/Code/ReLish/benchmark/train_cls.py", line 95, in main
    train_model(C, train_loader, valid_loader, model, output_layer, criterion, optimizer, scheduler)
  File "/home/phil/Code/ReLish/benchmark/train_cls.py", line 455, in train_model
    output = model(data)
  File "/home/phil/anaconda3/envs/relish/lib/python3.9/site-packages/torch/nn/modules/module.py", line 1130, in _call_impl
    return forward_call(*input, **kwargs)
  File "/home/phil/anaconda3/envs/relish/lib/python3.9/site-packages/torchvision/models/vision_transformer.py", line 298, in forward
    x = self.encoder(x)
  File "/home/phil/anaconda3/envs/relish/lib/python3.9/site-packages/torch/nn/modules/module.py", line 1130, in _call_impl
    return forward_call(*input, **kwargs)
  File "/home/phil/anaconda3/envs/relish/lib/python3.9/site-packages/torchvision/models/vision_transformer.py", line 157, in forward
    return self.ln(self.layers(self.dropout(input)))
  File "/home/phil/anaconda3/envs/relish/lib/python3.9/site-packages/torch/nn/modules/module.py", line 1130, in _call_impl
    return forward_call(*input, **kwargs)
  File "/home/phil/anaconda3/envs/relish/lib/python3.9/site-packages/torch/nn/modules/container.py", line 139, in forward
    input = module(input)
  File "/home/phil/anaconda3/envs/relish/lib/python3.9/site-packages/torch/nn/modules/module.py", line 1130, in _call_impl
    return forward_call(*input, **kwargs)
  File "/home/phil/anaconda3/envs/relish/lib/python3.9/site-packages/torchvision/models/vision_transformer.py", line 113, in forward
    x, _ = self.self_attention(query=x, key=x, value=x, need_weights=False)
  File "/home/phil/anaconda3/envs/relish/lib/python3.9/site-packages/torch/nn/modules/module.py", line 1130, in _call_impl
    return forward_call(*input, **kwargs)
  File "/home/phil/anaconda3/envs/relish/lib/python3.9/site-packages/torch/nn/modules/activation.py", line 1113, in forward
    return torch._native_multi_head_attention(
RuntimeError: expected scalar type Half but found Float
</code></pre>
<p>The error occurs in the first epoch AFTER training is complete, right when the validation should start. Without AMP, no error occurs. With every other model I’ve tried (including Swin transformers) there is no issue.</p>
<p>I have found this, but am not sure how it would apply:</p><aside class="onebox githubissue" data-onebox-src="https://github.com/NVIDIA/apex/issues/121">
  <header class="source">

      <a href="https://github.com/NVIDIA/apex/issues/121" target="_blank" rel="noopener nofollow ugc">github.com/NVIDIA/apex</a>
  </header>

  <article class="onebox-body">
    <div class="github-row">
  <div class="github-icon-container" title="Issue">
	  <svg width="60" height="60" class="github-icon" viewBox="0 0 14 16" aria-hidden="true"><path fill-rule="evenodd" d="M7 2.3c3.14 0 5.7 2.56 5.7 5.7s-2.56 5.7-5.7 5.7A5.71 5.71 0 0 1 1.3 8c0-3.14 2.56-5.7 5.7-5.7zM7 1C3.14 1 0 4.14 0 8s3.14 7 7 7 7-3.14 7-7-3.14-7-7-7zm1 3H6v5h2V4zm0 6H6v2h2v-2z"></path></svg>
  </div>

  <div class="github-info-container">
    <h4>
      <a href="https://github.com/NVIDIA/apex/issues/121" target="_blank" rel="noopener nofollow ugc">RuntimeError: expected scalar type Half but found Float</a>
    </h4>

    <div class="github-info">
      <div class="date">
        opened <span class="discourse-local-date" data-format="ll" data-date="2019-01-10" data-time="17:32:31" data-timezone="UTC">05:32PM - 10 Jan 19 UTC</span>
      </div>


      <div class="user">
        <a href="https://github.com/entavelis" target="_blank" rel="noopener nofollow ugc">
          <img alt="entavelis" src="https://avatars.githubusercontent.com/u/8151641?v=4" class="onebox-avatar-inline" width="20" height="20">
          entavelis
        </a>
      </div>
    </div>

    <div class="labels">
    </div>
  </div>
</div>

  <div class="github-row">
    <p class="github-body-container">Hello, I am trying to use amp to have mixed precision training for my model. I a<span class="show-more-container"><a href="" rel="noopener" class="show-more">…</a></span><span class="excerpt hidden">m implementing a multi-task learning algorithm, so my loss is a summation of loss_a and loss_b, which I then back-propagate. When doing so and following amp usage instructions I get:

_RuntimeError: expected scalar type Half but found Float_

I also tried, back-propagating one loss at the time, while retaining the graph, following your instructions for multiple losses. I get the same error.

In the previous cases, I didn't use model.half(). When I did, my losses started becoming NaN. 

Do you have any suggestions on how to proceed?

Thanks in advance!
VglsD

*Edit: I should also note that my network has Batch Normalization Layers

Environment: Miniconda, Python 3.7, Cuda 10,  PyTorch 1.0, apex master branch</span></p>
  </div>

  </article>

  <div class="onebox-metadata">
    
    
  </div>

  <div style="clear: both"></div>
</aside>
<p>
Maybe it’s something similar to this?</p><aside class="quote" data-post="5" data-topic="91628">
  <div class="title">
    <div class="quote-controls"></div>
    <img loading="lazy" alt="" width="24" height="24" src="https://discuss.pytorch.org/user_avatar/discuss.pytorch.org/mcarilli/48/28517_2.png" class="avatar">
    <a href="https://discuss.pytorch.org/t/runtimeerror-expected-scalar-type-half-but-found-float/91628/5">RuntimeError: expected scalar type Half but found Float</a> <a class="badge-category__wrapper " href="/c/vision/5"><span data-category-id="5" data-drop-close="true" class="badge-category " title="Topics related to either pytorch/vision or vision research related topics"><span class="badge-category__name">vision</span></span></a>
  </div>
  <blockquote>
    It appears LSTMCell and GRUCell have <a href="https://github.com/pytorch/pytorch/blob/db52cd7322ccf3f2121f1a6e592ddb436968aa50/tools/autograd/derivatives.yaml#L1756-L1761" rel="noopener nofollow ugc">dedicated autograd operations</a> that type-check inputs.  I should add those to the FP16 cast list. 
Filed issue <a href="https://github.com/pytorch/pytorch/issues/42605" rel="noopener nofollow ugc">https://github.com/pytorch/pytorch/issues/42605</a>.  Should be a straightforward fix, thanks for reporting this.  I thought the cells were implemented in terms of autograd-exposed primitive ops (eg. gemms) that autocast already covers.
  </blockquote>
</aside>

<p>This is the code I’m using:</p><aside class="onebox githubblob" data-onebox-src="https://github.com/pallgeuer/ReLish/blob/9a31066c3601821f4be8d2c15ace0d2171b8d342/benchmark/train_cls.py#L451">
  <header class="source">

      <a href="https://github.com/pallgeuer/ReLish/blob/9a31066c3601821f4be8d2c15ace0d2171b8d342/benchmark/train_cls.py#L451" target="_blank" rel="noopener nofollow ugc">github.com</a>
  </header>

  <article class="onebox-body">
    <h4><a href="https://github.com/pallgeuer/ReLish/blob/9a31066c3601821f4be8d2c15ace0d2171b8d342/benchmark/train_cls.py#L451" target="_blank" rel="noopener nofollow ugc">pallgeuer/ReLish/blob/9a31066c3601821f4be8d2c15ace0d2171b8d342/benchmark/train_cls.py#L451</a></h4>



    <pre class="onebox"><code class="lang-py">
      <ol class="start lines" start="441" style="counter-reset: li-counter 440 ;">
          <li>		init_detail_stamp = last_detail_stamp = timeit.default_timer()</li>
          <li></li>
          <li>		with torch.inference_mode():</li>
          <li>			for batch_num, (data, target_cpu) in enumerate(valid_loader):</li>
          <li></li>
          <li>				num_in_batch = data.shape[0]</li>
          <li>				data = data.to(device, non_blocking=True)</li>
          <li>				target = target_cpu.to(device, non_blocking=True)</li>
          <li></li>
          <li>				with torch.autocast(device_type=device.type, enabled=amp_enabled):</li>
          <li class="selected">					output = model(data)</li>
          <li>					mean_batch_loss = criterion(output if output_layer is None else output_layer(output), target)</li>
          <li></li>
          <li>				num_valid_samples += num_in_batch</li>
          <li>				output_cpu = output.detach().to(device=cpu_device, dtype=float)</li>
          <li>				output_nans += torch.count_nonzero(output_cpu.isnan()).item()</li>
          <li>				batch_topk_sum = calc_topk_sum(output_cpu, target_cpu, topn=5)</li>
          <li>				for k in range(5):</li>
          <li>					valid_topk[k] += batch_topk_sum[k]</li>
          <li>				valid_loss += mean_batch_loss.item() * num_in_batch</li>
          <li></li>
      </ol>
    </code></pre>



  </article>

  <div class="onebox-metadata">
    
    
  </div>

  <div style="clear: both"></div>
</aside>

<p>For reference I call it as (Python 3.9, PyTorch 1.12.1, CUDA 11.6, cuDNN 8.3.2):</p>
<pre><code class="lang-auto">./train_cls.py --act_func=original --batch_size=32 --dataset=Imagenette --epochs=120 --model=vit_b_16
</code></pre>
<p>Adding <code>--no_amp</code> to the command line turns AMP off.</p>
<p>Can anyone help?</p>

ptrblck: <p>The error seems to be raised by the MHA layer. Could you check if the latest nightly binary still rauses the error, please?</p>

pallgeuer: <p>No, the error is not there under <code>pytorch-nightly/linux-64::pytorch-1.14.0.dev20221011-py3.9_cuda11.6_cudnn8.3.2_0</code>. What is/was the problem, and can I do anything to overcome the issue now, besides moving to a nightly PyTorch release? (the nightly release broke other unrelated parts of my code related to torch.jit.script, which I had to just comment out otherwise it would crash before I have the chance to train)</p>

ptrblck: <p>I don’t know what the fix was but you could check the commit e.g. via <code>git blame</code> to check which issues were fixed and if some of them matched your error.</p>