Run slower when using autocast with small batch size

Wu_Haoqian: <p>I meet a problem that it is slower to train with autocast. I experiment on a toy model and find that it keep running slower than normal until a big batch size is set. I tried it on A30, 3090, 2080Ti, 2060 super.<br>
I am confused about the reason and the strategy of batch size setting. A small batch size makes autocast not work but a big batch size would cause crash.</p>
<p>Here is my code:</p>
<pre><code class="lang-auto">import torch
import torch.nn as nn
import time

import argparse
parser = argparse.ArgumentParser()

BATCH_SIZE = 1024
# BATCH_SIZE = 10240

class Model(nn.Module):
    def __init__(self):
        super().__init__()

        layer = []
        for i in range(50):
            layer.append(nn.Linear(512, 512))
        self.layer = nn.Sequential(*layer)

    def forward(self, x):

        y = self.layer(x)

        return y

model = Model().cuda()
scaler = torch.cuda.amp.GradScaler()
optimizer = torch.optim.SGD(model.parameters(), lr=0.001)
loss_fn = nn.MSELoss()

d = []
label = []
for i in range(100):
    d.append(torch.randn(BATCH_SIZE, 512).cuda())
    label.append(torch.randn(BATCH_SIZE, 512).cuda())

print("################ with autocast #######################")

print_flag = False
time0 = time.time()
for i in range(len(d)):
    optimizer.zero_grad()
    with torch.cuda.amp.autocast():
        outputs = model(d[i])
        if not print_flag:
            print("dtype: ", outputs.dtype)
            print_flag = True
        loss = loss_fn(outputs, label[i])

    scaler.scale(loss).backward()
    scaler.step(optimizer)
    scaler.update()
    # loss.backward()
    # optimizer.step()

time1 = time.time()
print("with autocast cost: ", time1 - time0)

#######################################################
print("################ without autocast #######################")

print_flag = False
time0 = time.time()
for i in range(len(d)):
    optimizer.zero_grad()

    outputs = model(d[i])
    if not print_flag:
        print("dtype: ", outputs.dtype)
        print_flag = True
    loss = loss_fn(outputs, label[i])

    loss.backward()
    optimizer.step()

time1 = time.time()
print("without autocast cost: ", time1 - time0)
</code></pre>
<p>when BATCH_SIZE = 1024, the output is (on 3090)</p>
<pre><code class="lang-auto">################ with autocast #######################
dtype:  torch.float16
with autocast cost:  1.9760982990264893
################ without autocast #######################
dtype:  torch.float32
without autocast cost:  1.1353034973144531
</code></pre>
<p>when BATCH_SIZE = 10240, the output is</p>
<pre><code class="lang-auto">################ with autocast #######################
dtype:  torch.float16
with autocast cost:  2.055091619491577
################ without autocast #######################
dtype:  torch.float32
without autocast cost:  3.0732078552246094

</code></pre>

ptrblck: <p>CUDA operations are executed asynchronously, so you would need to synchronize the code before starting and stopping the timers via <code>torch.cuda.synchronize()</code> or use an event-based profiling. Your current profiling code does not measure the runtime and is invalid.</p>

Wu_Haoqian: <p>I have another try but still slower. Here is my code and output. The GPU is 2060 super.</p>
<pre><code class="lang-auto">import torch
import torch.nn as nn
import time

import argparse
parser = argparse.ArgumentParser()


class Model(nn.Module):
    def __init__(self):
        super().__init__()

        layer = []
        for i in range(5):
            layer.append(nn.Linear(128, 128))
        self.layer = nn.Sequential(*layer)

    def forward(self, x):

        y = self.layer(x)

        return y

model = Model().cuda()
scaler = torch.cuda.amp.GradScaler()
optimizer = torch.optim.SGD(model.parameters(), lr=0.001)
loss_fn = nn.MSELoss()

d = []
label = []
for i in range(1000):
    d.append(torch.randn(20, 128).cuda())
    label.append(torch.randn(20, 128).cuda())

print("################ with autocast #######################")

print_flag = False
torch.cuda.synchronize()
time0 = time.time()
for i in range(len(d)):
    optimizer.zero_grad()
    with torch.cuda.amp.autocast():
        outputs = model(d[i])
        if not print_flag:
            print("dtype: ", outputs.dtype)
            print_flag = True
        loss = loss_fn(outputs, label[i])

    scaler.scale(loss).backward()
    scaler.step(optimizer)
    scaler.update()
    # loss.backward()
    # optimizer.step()

torch.cuda.synchronize()
time1 = time.time()
print("with autocast cost: ", time1 - time0)

#######################################################
print("################ without autocast #######################")

print_flag = False
torch.cuda.synchronize()
time0 = time.time()
for i in range(len(d)):
    optimizer.zero_grad()

    outputs = model(d[i])
    if not print_flag:
        print("dtype: ", outputs.dtype)
        print_flag = True
    loss = loss_fn(outputs, label[i])

    loss.backward()
    optimizer.step()

torch.cuda.synchronize()
time1 = time.time()
print("without autocast cost: ", time1 - time0)


</code></pre>
<p>output:</p>
<pre><code class="lang-auto">################ with autocast #######################
dtype:  torch.float16
with autocast cost:  10.757721424102783
################ without autocast #######################
dtype:  torch.float32
without autocast cost:  0.9863979816436768
</code></pre>

Wu_Haoqian: <p>I “warm up” the machine before count the time.</p>
<pre><code class="lang-auto">import torch
import torch.nn as nn
import time

import argparse
parser = argparse.ArgumentParser()


class Model(nn.Module):
    def __init__(self):
        super().__init__()

        layer = []
        for i in range(5):
            layer.append(nn.Linear(128, 128))
        self.layer = nn.Sequential(*layer)

    def forward(self, x):

        y = self.layer(x)

        return y

model = Model().cuda()
scaler = torch.cuda.amp.GradScaler()
optimizer = torch.optim.SGD(model.parameters(), lr=0.001)
loss_fn = nn.MSELoss()

d = []
label = []
for i in range(1000):
    d.append(torch.randn(20, 128).cuda())
    label.append(torch.randn(20, 128).cuda())

#################### warmup #####################
for i in range(len(d)):
    optimizer.zero_grad()
    with torch.cuda.amp.autocast(enabled=True):
        outputs = model(d[i])
        loss = loss_fn(outputs, label[i])

    scaler.scale(loss).backward()
    scaler.step(optimizer)
    scaler.update()

for i in range(len(d)):
    optimizer.zero_grad()

    outputs = model(d[i])
    loss = loss_fn(outputs, label[i])

    loss.backward()
    optimizer.step()
#################### warmup #####################


print("################ with autocast #######################")

print_flag = False
torch.cuda.synchronize()
time0 = time.time()
for i in range(len(d)):
    optimizer.zero_grad()
    with torch.cuda.amp.autocast(enabled=True):
        outputs = model(d[i])
        if not print_flag:
            print("dtype: ", outputs.dtype)
            print_flag = True
        loss = loss_fn(outputs, label[i])

    scaler.scale(loss).backward()
    scaler.step(optimizer)
    scaler.update()
    # loss.backward()
    # optimizer.step()

torch.cuda.synchronize()
time1 = time.time()
print("with autocast cost: ", time1 - time0)

#######################################################
print("################ without autocast #######################")

print_flag = False
torch.cuda.synchronize()
time0 = time.time()
for i in range(len(d)):
    optimizer.zero_grad()

    outputs = model(d[i])
    if not print_flag:
        print("dtype: ", outputs.dtype)
        print_flag = True
    loss = loss_fn(outputs, label[i])

    loss.backward()
    optimizer.step()

torch.cuda.synchronize()
time1 = time.time()
print("without autocast cost: ", time1 - time0)


</code></pre>
<p>output:</p>
<pre><code class="lang-auto">################ with autocast #######################
dtype:  torch.float16
with autocast cost:  1.8655202388763428
################ without autocast #######################
dtype:  torch.float32
without autocast cost:  0.9314119815826416
</code></pre>
<p>the machine is 2060 super</p>

ptrblck: <p>Thanks for the follow-up. On my 2080Ti I see a runtime of:</p>
<pre><code class="lang-auto">################ with autocast #######################
dtype:  torch.float16
with autocast cost:  0.34975481033325195
################ without autocast #######################
dtype:  torch.float32
without autocast cost:  0.21597671508789062
</code></pre>
<p>and after checking the profile my guess that this tiny workload is CPU-limited and is thus not seeing a speedup seems to be valid.<br>
With autocast:<br>
<div class="lightbox-wrapper"><a class="lightbox" href="https://discuss.pytorch.org/uploads/default/original/3X/2/2/22add136ee8eaa5f02721df389fe020b40859c8e.png" data-download-href="https://discuss.pytorch.org/uploads/default/22add136ee8eaa5f02721df389fe020b40859c8e" title="image"><img src="https://discuss.pytorch.org/uploads/default/optimized/3X/2/2/22add136ee8eaa5f02721df389fe020b40859c8e_2_690x197.png" alt="image" data-base62-sha1="4WMDpIPht04Nc4UWGqzu4qo1Ink" width="690" height="197" srcset="https://discuss.pytorch.org/uploads/default/optimized/3X/2/2/22add136ee8eaa5f02721df389fe020b40859c8e_2_690x197.png, https://discuss.pytorch.org/uploads/default/optimized/3X/2/2/22add136ee8eaa5f02721df389fe020b40859c8e_2_1035x295.png 1.5x, https://discuss.pytorch.org/uploads/default/optimized/3X/2/2/22add136ee8eaa5f02721df389fe020b40859c8e_2_1380x394.png 2x" data-dominant-color="F8F8F9"><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">1542×441 46.8 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>Without autocast:<br>
<div class="lightbox-wrapper"><a class="lightbox" href="https://discuss.pytorch.org/uploads/default/original/3X/a/c/ace7353664b656071cb7a8fc51c1bea4add6696f.png" data-download-href="https://discuss.pytorch.org/uploads/default/ace7353664b656071cb7a8fc51c1bea4add6696f" title="image"><img src="https://discuss.pytorch.org/uploads/default/optimized/3X/a/c/ace7353664b656071cb7a8fc51c1bea4add6696f_2_690x189.png" alt="image" data-base62-sha1="oFzu6amITY5gibesn41oyZbB6J9" width="690" height="189" srcset="https://discuss.pytorch.org/uploads/default/optimized/3X/a/c/ace7353664b656071cb7a8fc51c1bea4add6696f_2_690x189.png, https://discuss.pytorch.org/uploads/default/optimized/3X/a/c/ace7353664b656071cb7a8fc51c1bea4add6696f_2_1035x283.png 1.5x, https://discuss.pytorch.org/uploads/default/optimized/3X/a/c/ace7353664b656071cb7a8fc51c1bea4add6696f_2_1380x378.png 2x" data-dominant-color="F8F8F9"><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">1583×434 47.8 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>As you can see, both use cases suffer from huge white spaces between the actual kernel runtimes.<br>
While the kernel execution times are faster using <code>autocast</code>, the overhead of loss scaling etc. is visible since the actual model computation is too small to hide it and you are seeing a slowdown in the end.<br>
To see a speedup, you would have to scale up the actual workload by using a larger model instead of a toy model or by increasing the batch size.<br>
You can recreate the same profiles via Nsight Systems to analyze them further.</p>
<p>PS: if you want to speedup these tiny workloads, you might want to check <a href="https://pytorch.org/docs/stable/notes/cuda.html#cuda-graphs">CUDA Graphs</a>.</p>

Wu_Haoqian: <p>I see. Thanks a lot!</p>