Torch.cuda.amp cannot speed up on A100

username: thomas0809
cooked: <p>I am trying to use <code>torch.cuda.amp</code> to speed up training. The following code works well on V100 (120s w/o amp and 67s w/ amp), but cannot get a reasonable speedup on A100 (53s w/o amp and 50s w/ amp).</p>
<p>I am using the most recent NGC docker <a href="https://ngc.nvidia.com/catalog/containers/nvidia:pytorch" rel="noopener nofollow ugc">nvcr.io/nvidia/pytorch:21.04-py3</a>, which has CUDA 11.3, pytorch 1.9.0.</p>
<p>The strange thing is if I change the model to resnet101, both V100 and A100 can get speedup, but not for the swin transformer on A100. See the following table for the running time (in sec) on different machines.</p>
<div class="md-table">
<table>
<thead>
<tr>
<th></th>
<th>resnet101</th>
<th>resnet101 amp</th>
<th>swin_base</th>
<th>swin_base amp</th>
</tr>
</thead>
<tbody>
<tr>
<td>V100</td>
<td>49.19</td>
<td>19.62</td>
<td>120.07</td>
<td>67.19</td>
</tr>
<tr>
<td>A100</td>
<td>40.9</td>
<td>29.35</td>
<td>47.19</td>
<td>44.97</td>
</tr>
<tr>
<td>A100 w/ docker</td>
<td>24.32</td>
<td>13.26</td>
<td>53.68</td>
<td>50.04</td>
</tr>
</tbody>
</table>
</div><pre><code class="lang-auto">import torch
import argparse
import os
import time
import timm
from tqdm import trange

parser = argparse.ArgumentParser()
parser.add_argument("--model", type=str, default='swin_base')
parser.add_argument("--fp16", action='store_true')
args = parser.parse_args()

N = 32
input_size = 384

x = torch.randn(N, 3, input_size, input_size).cuda()
y = torch.randint(1000, (N,)).cuda()

if args.model == 'swin_base':
    model = timm.create_model('swin_base_patch4_window12_384').cuda()
elif args.model == 'resnet101':
    model = timm.create_model('resnet101').cuda()

loss_fn = torch.nn.CrossEntropyLoss()
optimizer = torch.optim.SGD(model.parameters(), lr=1e-3)

if args.fp16:
    scaler = torch.cuda.amp.GradScaler()

start = time.time()

for t in trange(100):
#     optimizer.zero_grad()
    with torch.cuda.amp.autocast(enabled=args.fp16):
        y_pred = model(x)
        if args.fp16:
            assert y_pred.dtype is torch.float16
        loss = loss_fn(y_pred, y)
    if args.fp16:
        scaler.scale(loss).backward()
    else:
        loss.backward()
#     optimizer.step()

end = time.time()

print('Time:', end-start)
</code></pre>

username: ptrblck
cooked: <p>I’m not familiar with this model, but note that you are already using TensorCores on the A100, since TF32 is enabled by default. With a proper synchronization I get a runtime of:</p>
<pre><code class="lang-auto">FP32: 102s
TF32: 50s
AMP: 47s
</code></pre>
<p>Based on the profile it also seems that the majority or kernels are (vectorized/unrolled) elementwise kernels and sporadically a cublas kernel is called:<br>
FP32:<br>
<div class="lightbox-wrapper"><a class="lightbox" href="https://discuss.pytorch.org/uploads/default/original/3X/2/3/23a6592eb781ef6c91fc57cbd4a908d5f38b41c9.png" data-download-href="https://discuss.pytorch.org/uploads/default/23a6592eb781ef6c91fc57cbd4a908d5f38b41c9" title="image"><img src="https://discuss.pytorch.org/uploads/default/optimized/3X/2/3/23a6592eb781ef6c91fc57cbd4a908d5f38b41c9_2_690x194.png" alt="image" data-base62-sha1="55n6TsJV7q4j726QZkNUTqjcuZX" width="690" height="194" srcset="https://discuss.pytorch.org/uploads/default/optimized/3X/2/3/23a6592eb781ef6c91fc57cbd4a908d5f38b41c9_2_690x194.png, https://discuss.pytorch.org/uploads/default/optimized/3X/2/3/23a6592eb781ef6c91fc57cbd4a908d5f38b41c9_2_1035x291.png 1.5x, https://discuss.pytorch.org/uploads/default/optimized/3X/2/3/23a6592eb781ef6c91fc57cbd4a908d5f38b41c9_2_1380x388.png 2x" data-dominant-color="E6EBF3"><div class="meta"><svg class="fa d-icon d-icon-far-image svg-icon" aria-hidden="true"><use href="#far-image"></use></svg><span class="filename">image</span><span class="informations">1490×420 28.6 KB</span><svg class="fa d-icon d-icon-discourse-expand svg-icon" aria-hidden="true"><use href="#discourse-expand"></use></svg></div></a></div><br>
FP16:<br>
<div class="lightbox-wrapper"><a class="lightbox" href="https://discuss.pytorch.org/uploads/default/original/3X/d/0/d0f469fe204357a57a50e6f51d84f91fabac07ed.png" data-download-href="https://discuss.pytorch.org/uploads/default/d0f469fe204357a57a50e6f51d84f91fabac07ed" title="image"><img src="https://discuss.pytorch.org/uploads/default/optimized/3X/d/0/d0f469fe204357a57a50e6f51d84f91fabac07ed_2_690x170.png" alt="image" data-base62-sha1="tOuY2W5IFWzDK7t07oYri8Ek7c9" width="690" height="170" srcset="https://discuss.pytorch.org/uploads/default/optimized/3X/d/0/d0f469fe204357a57a50e6f51d84f91fabac07ed_2_690x170.png, https://discuss.pytorch.org/uploads/default/optimized/3X/d/0/d0f469fe204357a57a50e6f51d84f91fabac07ed_2_1035x255.png 1.5x, https://discuss.pytorch.org/uploads/default/optimized/3X/d/0/d0f469fe204357a57a50e6f51d84f91fabac07ed_2_1380x340.png 2x" data-dominant-color="E0E5EF"><div class="meta"><svg class="fa d-icon d-icon-far-image svg-icon" aria-hidden="true"><use href="#far-image"></use></svg><span class="filename">image</span><span class="informations">1692×418 39.3 KB</span><svg class="fa d-icon d-icon-discourse-expand svg-icon" aria-hidden="true"><use href="#discourse-expand"></use></svg></div></a></div></p>
<p>Based on the <a href="https://pytorch.org/docs/stable/amp.html#ops-that-can-autocast-to-float16">Ops that can autocast to float16</a> it seems that no many of these operations are used.</p>

username: omeryasar
cooked: <p>Hi, I want to monitor my gpu activations too so may I ask how do you get with these screen shots? With which app. Thanks in advance.</p>

username: ptrblck
cooked: <p>I used Nsight Systems and you can take a look at <a href="https://dev-discuss.pytorch.org/t/using-nsight-systems-to-profile-gpu-workload/59?u=ptrblck">this post</a> for an example.</p>