FP16 and BF16 way slower than FP32 and TF32

Robin_Lobel: <p>I don’t know what I’m doing wrong, but my FP16 and BF16 bench are way slower than FP32 and TF32 modes. Here are my results with the 2 GPUs at my disposal (RTX 2060 Mobile, RTX 3090 Desktop):</p>
<blockquote>
<p>Benching precision speed on a NVIDIA GeForce RTX 2060</p>
<p>benching FP32…<br>
epoch 0 took 13.9146514s<br>
epoch 1 took 11.6350846s<br>
epoch 2 took 11.867831299999999s</p>
<p>benching FP16…<br>
epoch 0 took 15.745933399999998s<br>
epoch 1 took 16.212830699999998s<br>
epoch 2 took 16.495791399999987s</p>
</blockquote>
<blockquote>
<p>Benching precision speed on a NVIDIA GeForce RTX 3090</p>
<p>benching FP32…<br>
epoch 0 took 5.7641565s<br>
epoch 1 took 4.0729165s<br>
epoch 2 took 4.0790243s</p>
<p>benching TF32…<br>
epoch 0 took 4.042242200000002s<br>
epoch 1 took 4.0321663s<br>
epoch 2 took 4.080792600000002s</p>
<p>benching FP16…<br>
epoch 0 took 5.053079000000004s<br>
epoch 1 took 5.029029299999998s<br>
epoch 2 took 4.973819899999995s</p>
<p>benching BF16…<br>
epoch 0 took 11.721234800000005s<br>
epoch 1 took 11.542296499999999s<br>
epoch 2 took 11.566654600000007s</p>
</blockquote>
<p>And here’s the file I made to generate the bench, a simple MNIST classifier that you can easily run on your own computer (I run it with PyTorch 1.12 on Windows and the latest NVIDIA drivers):</p>
<pre><code class="lang-auto">#MNIST example inspired by https://github.com/pytorch/examples/blob/main/mnist/main.py
import timeit
import torch
import torch.nn as nn
import torch.nn.functional as F
import torchvision
dataset = torchvision.datasets.MNIST(root = 'data', train = True, transform = torchvision.transforms.ToTensor(), download = True)
loader = torch.utils.data.DataLoader(dataset, batch_size=64, pin_memory=True, shuffle=True)

class Net(nn.Module):
    def __init__(self):
        super(Net, self).__init__()
        self.conv1 = nn.Conv2d(1, 32, 3, 1)
        self.conv2 = nn.Conv2d(32, 64, 3, 1)
        self.dropout1 = nn.Dropout(0.25)
        self.dropout2 = nn.Dropout(0.5)
        self.fc1 = nn.Linear(9216, 128)
        self.fc2 = nn.Linear(128, 10)

    def forward(self, x):
        x = self.conv1(x)
        x = F.relu(x)
        x = self.conv2(x)
        x = F.relu(x)
        x = F.max_pool2d(x, 2)
        x = self.dropout1(x)
        x = torch.flatten(x, 1)
        x = self.fc1(x)
        x = F.relu(x)
        x = self.dropout2(x)
        x = self.fc2(x)
        output = F.log_softmax(x, dim=1)
        return output

device = torch.device("cuda")

def bench(mode, epochs=3):
    print("benching "+mode+"...")
    torch.backends.cuda.matmul.allow_tf32 = True if mode == 'TF32' else False
    scaler = torch.cuda.amp.GradScaler(enabled=True if mode == 'FP16' else False)

    model = Net().to(device)
    optimizer = torch.optim.Adadelta(model.parameters(), lr=1.0)
    scheduler = torch.optim.lr_scheduler.StepLR(optimizer, step_size=1, gamma=0.7)
    model.train()
    for epoch in range(epochs):
        start_time = timeit.default_timer()
        model.train()
        for batch_idx, (data, target) in enumerate(loader):
            data, target = data.to(device), target.to(device)
            optimizer.zero_grad()
            with torch.autocast(device_type='cuda', dtype=torch.bfloat16 if mode == 'BF16' else torch.float16, enabled=True if '16' in mode else False):
                output = model(data)
                loss = F.nll_loss(output, target)
            scaler.scale(loss).backward()
            scaler.step(optimizer)
            scaler.update()
        scheduler.step()
        torch.cuda.synchronize()
        end_time = timeit.default_timer()
        print("epoch "+str(epoch)+" took "+str(end_time-start_time)+"s")
    print("")

print("Benching precision speed on a "+torch.cuda.get_device_name(0))
print("")
bench(mode="FP32")
if torch.cuda.is_bf16_supported():
    bench(mode="TF32")
bench(mode="FP16")
if torch.cuda.is_bf16_supported():
    bench(mode="BF16")
</code></pre>

ptrblck: <p>I cannot reproduce a slowdown using a recent PyTorch build with CUDA 11.6 and see:</p>
<pre><code class="lang-python">benching FP32...
epoch 0 took 5.307978553988505s
epoch 1 took 3.9439379789982922s
epoch 2 took 3.9074656259908807s

benching TF32...
epoch 0 took 4.101563502015779s
epoch 1 took 3.946826447005151s
epoch 2 took 4.008894422004232s

benching FP16...
epoch 0 took 4.420351507025771s
epoch 1 took 4.3506982740073s
epoch 2 took 4.440277656976832s

benching BF16...
epoch 0 took 4.17264149300172s
epoch 1 took 4.073878110008081s
epoch 2 took 4.178816699975869s
</code></pre>
<p>A few things for your profiling:</p>
<ul>
<li>You would need to synchronize the code also before starting the timer, not only before stopping it in case some kernels are still running (e.g. transferring the weights to the GPU).</li>
<li>You are profiling a full training run including the <code>DataLoader</code>, transferring the data to the GPU, and are using a tiny model. Depending on your system the actual model runtime might be tiny and you might see a large overhead from the actual data loading etc. so that actual model speedups won’t be directly visible. If you want to compare different numerical precisions and their speed, I would recommend to profile the model in isolation first.</li>
<li>You model workload is small so even if lower precision <code>dtype</code>s give a speedup, the actual kernel launches, the dispatching etc. might be visible. If this small model is your real workload, you might want to try <a href="https://pytorch.org/docs/stable/notes/cuda.html#cuda-graphs">CUDA graphs</a>.</li>
<li>To enable <code>bfloat16</code> calculations in conv layers, set <code>os.environ["TORCH_CUDNN_V8_API_ENABLED"] = "1"</code> or <code>export</code> this env variable. However, <code>bfloat16</code> should be enabled for Ampere+, so your Turing GPU might not see any benefits.</li>
</ul>

Robin_Lobel: <p>Thanks for the advices, indeed that’s a small model/dataset just for testing so I’m not surprised to not see any dramatic speed-up, yet I’m puzzled by the visible slowdown as soon as 16 bits floats are being used. that doesn’t sound right. Is autocasting adding a significant overhead that could be noticed on small models ?</p>
<p>The bench were done on 2 different systems, both running Windows 11 21H2, latest NV drivers, Python 3.9, PyTorch 1.12.1 with CUDA 11.6, with NVME SSD and fast CPU/RAM.<br>
I’m a bit surprised you did not see similar results (although it’s FP16/BF16 seems also slightly slower on average on your system too), did you run the bench on Windows or Linux ? Maybe it’s Windows-specific ?</p>
<p>I’ve added your suggestions (additional sync+cudnn flag), which gives me different numbers, but still slower 16 bits training:</p>
<pre><code class="lang-auto">Benching precision speed on a NVIDIA GeForce RTX 2060

benching FP32...
epoch 0 took 15.4749659s
epoch 1 took 11.017322s
epoch 2 took 13.260371799999998s

benching FP16...
epoch 0 took 15.714570200000004s
epoch 1 took 15.929676400000005s
epoch 2 took 16.421480399999993s
</code></pre>
<pre><code class="lang-auto">Benching precision speed on a NVIDIA GeForce RTX 3090

benching FP32...
epoch 0 took 7.703148500000001s
epoch 1 took 3.7395274s
epoch 2 took 3.7481109999999997s

benching TF32...
epoch 0 took 3.752561700000001s
epoch 1 took 3.6327957000000026s
epoch 2 took 3.7228427999999987s

benching FP16...
epoch 0 took 4.7007553999999985s
epoch 1 took 4.666421100000001s
epoch 2 took 4.604346900000003s

benching BF16...
epoch 0 took 4.327719200000004s
epoch 1 took 4.267118599999996s
epoch 2 took 4.287737299999996s
</code></pre>
<p>full updated code below:</p>
<pre><code class="lang-auto">#MNIST example inspired by https://github.com/pytorch/examples/blob/main/mnist/main.py
import os
import timeit
import torch
import torch.nn as nn
import torch.nn.functional as F
import torchvision
dataset = torchvision.datasets.MNIST(root = 'data', train = True, transform = torchvision.transforms.ToTensor(), download = True)
loader = torch.utils.data.DataLoader(dataset, batch_size=64, pin_memory=True, shuffle=True)

class Net(nn.Module):
    def __init__(self):
        super(Net, self).__init__()
        self.conv1 = nn.Conv2d(1, 32, 3, 1)
        self.conv2 = nn.Conv2d(32, 64, 3, 1)
        self.dropout1 = nn.Dropout(0.25)
        self.dropout2 = nn.Dropout(0.5)
        self.fc1 = nn.Linear(9216, 128)
        self.fc2 = nn.Linear(128, 10)

    def forward(self, x):
        x = self.conv1(x)
        x = F.relu(x)
        x = self.conv2(x)
        x = F.relu(x)
        x = F.max_pool2d(x, 2)
        x = self.dropout1(x)
        x = torch.flatten(x, 1)
        x = self.fc1(x)
        x = F.relu(x)
        x = self.dropout2(x)
        x = self.fc2(x)
        output = F.log_softmax(x, dim=1)
        return output

device = torch.device("cuda")

def bench(mode, epochs=3):
    print("benching "+mode+"...")
    os.environ["TORCH_CUDNN_V8_API_ENABLED"] = "1"
    torch.backends.cuda.matmul.allow_tf32 = True if mode == 'TF32' else False
    scaler = torch.cuda.amp.GradScaler(enabled=True if mode == 'FP16' else False)

    model = Net().to(device)
    optimizer = torch.optim.Adadelta(model.parameters(), lr=1.0)
    scheduler = torch.optim.lr_scheduler.StepLR(optimizer, step_size=1, gamma=0.7)
    model.train()
    for epoch in range(epochs):
        torch.cuda.synchronize()
        start_time = timeit.default_timer()
        model.train()
        for batch_idx, (data, target) in enumerate(loader):
            data, target = data.to(device), target.to(device)
            optimizer.zero_grad()
            with torch.autocast(device_type='cuda', dtype=torch.bfloat16 if mode == 'BF16' else torch.float16, enabled=True if '16' in mode else False):
                output = model(data)
                loss = F.nll_loss(output, target)
            scaler.scale(loss).backward()
            scaler.step(optimizer)
            scaler.update()
        scheduler.step()
        torch.cuda.synchronize()
        end_time = timeit.default_timer()
        print("epoch "+str(epoch)+" took "+str(end_time-start_time)+"s")
    print("")

print("Benching precision speed on a "+torch.cuda.get_device_name(0))
print("")
bench(mode="FP32")
if torch.cuda.is_bf16_supported():
    bench(mode="TF32")
bench(mode="FP16")
if torch.cuda.is_bf16_supported():
    bench(mode="BF16")

</code></pre>