I am trying to use torch.cuda.amp 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).
I am using the most recent NGC docker nvcr.io/nvidia/pytorch:21.04-py3 5, which has CUDA 11.3, pytorch 1.9.0.
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.





resnet101
resnet101 amp
swin_base
swin_base amp




V100
49.19
19.62
120.07
67.19


A100
40.9
29.35
47.19
44.97


A100 w/ docker
24.32
13.26
53.68
50.04



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)