Slower training on A100 than 2080Ti or 3060Ti

Skylar_Wurster: <p>I am training a set of learnable 3D transformation matrices as part of a larger model. I started with just scale/translation and some more efficient operations, which resulted in the A100 vastly outperforming my 2080Ti. However, have since moved to fully learnable transformation matrices, and now the A100 is slower than my 2080Ti.</p>
<p>Here is a minimum executable code snippet assuming you have PyTorch+tensorboard installed.</p>
<pre><code class="lang-auto">from __future__ import absolute_import, division, print_function
import torch
import torch.optim as optim
import torch.nn.functional as F
import time
import os
from torch.utils.tensorboard import SummaryWriter

class benchmark_model(torch.nn.Module):
    def __init__(self):
        super().__init__()
        
        self.transformation_matrices = torch.nn.Parameter(
            torch.zeros(
                [16, 4, 4],
                device = "cuda:0",
                dtype=torch.float32
            ),
            requires_grad=True
        )
        with torch.no_grad():
            self.transformation_matrices[:] = torch.eye(4, device="cuda:0", dtype=torch.float32)
            self.transformation_matrices[:,0:3,:] += torch.rand_like(
                self.transformation_matrices[:,0:3,:],
                device="cuda:0", dtype=torch.float32) * 0.1
            self.transformation_matrices = torch.nn.Parameter(
                self.transformation_matrices @ \
                self.transformation_matrices.transpose(-1, -2),
                requires_grad=True)
            self.transformation_matrices[:,3,0:3] = 0  

    def transform(self, x):
        torch.cuda.synchronize()
        transformation_matrices = self.transformation_matrices
        
        torch.cuda.synchronize()            
        transformed_points = torch.cat(
            [x, torch.ones([x.shape[0], 1], 
            device=x.device,
            dtype=torch.float32)], 
            dim=1).unsqueeze(0).repeat(
                transformation_matrices.shape[0], 1, 1
            )
        
        torch.cuda.synchronize()        
        transformed_points = torch.bmm(transformation_matrices, 
                            transformed_points.transpose(-1, -2)).transpose(-1, -2)
        
        
        torch.cuda.synchronize()        
        transformed_points = transformed_points[...,0:3]
            
        torch.cuda.synchronize()        
        return transformed_points 
    
    def forward(self, x):
        torch.cuda.synchronize()
        transformed_points = self.transform(x)
        
        torch.cuda.synchronize()
        coeffs = torch.linalg.det(self.transformation_matrices[:,0:3,0:3]).unsqueeze(0) / ((2.0*torch.pi)**(3/2))
        
        torch.cuda.synchronize()
        exps = torch.exp(-0.5 * \
            torch.sum(
                transformed_points.transpose(0,1)**20, 
            dim=-1))
        
        torch.cuda.synchronize()
        result = torch.sum(coeffs * exps, dim=-1, keepdim=True)
        
        torch.cuda.synchronize()
        return result
    
if __name__ == '__main__':

    start_time = time.time()
    writer = SummaryWriter(os.path.join('tensorboard','profiletest'))
    
    target_density = torch.randn([10000, 1], device="cuda:0", dtype=torch.float32)
    target_density /= target_density.sum()
    
    model = benchmark_model()
    
    optimizer = optim.Adam([{"params": model.transformation_matrices}], lr=0.0001, 
            betas=[0.99, 0.99]) 
    
    x = torch.rand([10000, 3], device="cuda:0", dtype=torch.float32)*2 - 1 
    
    with torch.profiler.profile(
        activities=[
            torch.profiler.ProfilerActivity.CPU,
            torch.profiler.ProfilerActivity.CUDA,
        ],
        schedule=torch.profiler.schedule(
            wait=2,
            warmup=8,
            active=1,
            repeat=1),
        profile_memory=True,
        with_stack=True,
        with_modules=True,
        on_trace_ready=torch.profiler.tensorboard_trace_handler(
            os.path.join('tensorboard',"profiletest"))) as profiler:
        for iteration in range(10000):
            torch.cuda.synchronize()
            optimizer.zero_grad()
            
            torch.cuda.synchronize() 
            
            torch.cuda.synchronize()          
            density = model(x)
            
            torch.cuda.synchronize()
            density /= density.sum().detach()  
            
            torch.cuda.synchronize()
            density_loss = F.kl_div(
                torch.log(density+1e-16), 
                    torch.log(target_density.detach()+1e-16), 
                    reduction='none', 
                    log_target=True)
            
            torch.cuda.synchronize()
            density_loss.mean().backward()
            
            torch.cuda.synchronize()
            optimizer.step()   
            
            torch.cuda.synchronize()
            profiler.step()    

    end_time = time.time()
    print(f"Took {(end_time-start_time)/60 : 0.02f} min")    

    writer.close()
</code></pre>
<p>This has been tested on 5 setups (all python 3.9):<br>
Windows 11/CUDA12/PyTorch+cu117 with 2080Ti<br>
The above machine with WSL2<br>
Ubuntu/CUDA 11.7/PyTorch+cu117 with 3060Ti<br>
Supercomputer cluster (in Illinois) with Linux/CUDA 11.3/PyTorch+cu117 with A100<br>
Supercomputer cluster (in Ohio) with Linux/CUDA 11.7/PyTorch+cu117 with A100</p>
<p>The tests on the 2080Ti and the 3060Ti are both quicker than the A100s.<br>
The profiler (with plenty of GPU sync operations) seems to point to the backward pass being much longer on the A100 in the trace. Seems related to the BMM operation.</p>
<p>I have found if I detach <code>transformed_points</code> on line 66, the A100 again becomes faster.<br>
I also have traces saved for each of these setups to view in tensorboard, here is an image from one of the A100 traces.<br>
<div class="lightbox-wrapper"><a class="lightbox" href="https://discuss.pytorch.org/uploads/default/original/3X/b/e/be1ea8cb1c11d7527663ef68e8593512c556509c.png" data-download-href="https://discuss.pytorch.org/uploads/default/be1ea8cb1c11d7527663ef68e8593512c556509c" title="trace"><img src="https://discuss.pytorch.org/uploads/default/optimized/3X/b/e/be1ea8cb1c11d7527663ef68e8593512c556509c_2_690x195.png" alt="trace" data-base62-sha1="r7Spdv45qPOF1DFO6SGCIpa0RlO" width="690" height="195" srcset="https://discuss.pytorch.org/uploads/default/optimized/3X/b/e/be1ea8cb1c11d7527663ef68e8593512c556509c_2_690x195.png, https://discuss.pytorch.org/uploads/default/optimized/3X/b/e/be1ea8cb1c11d7527663ef68e8593512c556509c_2_1035x292.png 1.5x, https://discuss.pytorch.org/uploads/default/optimized/3X/b/e/be1ea8cb1c11d7527663ef68e8593512c556509c_2_1380x390.png 2x" data-dominant-color="E8E9E9"><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">trace</span><span class="informations">2104×597 78.7 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>Any ideas?</p>

Skylar_Wurster: <p>I found the problem! The backend CUDA settings default to torch.backends.cuda.matmul.allow_tf32=False. After I set torch.backends.cuda.matmul.allow_tf32=True, the A100 machines again have faster performance than the 2080Ti and 3060Ti.</p>

cxv0519: <p>From 1.12 default flag for TF32 has been changed to False!. <a href="https://pytorch.org/docs/stable/notes/cuda.html#tensorfloat-32-tf32-on-ampere-devices" rel="noopener nofollow ugc">link</a></p>
<p>But It seems weird though.<br>
Specification says 2080Ti has 13.4TFLOPS for single precision.<br>
And A100 has 19.5TFLOPS for single precision.<br>
So I think A100 should faster than 2080Ti without using tensor cores…</p>
<p>Thus I think there may be an issue with cloud setup(i/o bottleneck or something…)<br>
Any ideas?</p>

ptrblck: <p>Note that e.g. heuristics are trained for TF32 kernels as the primary target while FP32 kernels are used as a fallback. Depending on your workload, the cuBLAS version etc., a non-optimal kernel might be picked.</p>