Profiler not accounting kernels in the backward pass in record_function

gpoesia: <p>I’m using torch.profiler to profile a full training step (forward + backward + optimizer step). I have each of these 3 stages in a <code>record_function</code> block. However, it seems like the backward pass block is not accounting for CUDA time spent on the relevant kernels called by autograd.</p>
<p>Here is a complete example:</p>
<pre><code class="lang-auto">import torch
import torch.nn as nn
from torch.profiler import profile, record_function

class SimpleNet(nn.Module):
    def __init__(self):
        super(SimpleNet, self).__init__()
        self.fc1 = nn.Linear(10000, 100000)
        self.fc2 = nn.Linear(100000, 10000)

    def forward(self, x):
        x = torch.relu(self.fc1(x))
        x = torch.relu(self.fc2(x))
        return x

def run_step(model, x, backward, optimizer=None):
    with record_function("forward"):
        y = model(x)
    if backward:
        with record_function("backward"):
            loss = torch.mean(y ** 2)
            loss.backward()
        if optimizer is not None:
            with record_function("optimizer"):
                optimizer.step()
                optimizer.zero_grad()
    return y

device = torch.device("cuda")
model = SimpleNet().to(device)
optimizer = torch.optim.AdamW(model.parameters(), lr=0.01)
x = torch.randn(10000, 10000).to(device)
with profile(activities=[torch.profiler.ProfilerActivity.CPU, torch.profiler.ProfilerActivity.CUDA],
             schedule=torch.profiler.schedule(wait=0, warmup=0, active=1, repeat=1),
             ) as prof:
    run_step(model, x, backward=True, optimizer=optimizer)
    prof.step()

print(prof.key_averages().table(sort_by="cuda_time_total", row_limit=10000))
</code></pre>
<p>When running this, we get:</p>
<pre><code class="lang-auto">------------------------------------------------------  ------------  ------------  ------------  ------------  ------------  ------------  ------------  ------------  ------------  ------
------  
                                                   Name    Self CPU %      Self CPU   CPU total %     CPU total  CPU time avg     Self CUDA   Self CUDA %    CUDA total  CUDA time avg    # o
f Calls  
-------------------------------------------------------  ------------  ------------  ------------  ------------  ------------  ------------  ------------  ------------  ------------  ------
------  
                                          ProfilerStep*         0.00%     126.000us        71.34%        2.066s        2.066s       0.000us         0.00%        1.193s        1.193s        
     1  
    autograd::engine::evaluate_function: AddmmBackward0         0.00%      89.000us        27.18%     787.008ms     393.504ms       0.000us         0.00%        1.159s     579.507ms        
     2  
                                         AddmmBackward0         0.00%      67.000us        14.17%     410.129ms     205.065ms       0.000us         0.00%        1.158s     578.788ms        
     2  
                                               aten::mm         0.01%     392.000us        14.16%     410.011ms     136.670ms        1.158s        58.29%        1.158s     385.859ms             3  
                                                forward         0.01%     314.000us        14.42%     417.382ms     417.382ms       0.000us         0.00%        1.130s        1.130s        
     1                     
...
                                                optimizer         0.01%     243.000us        15.15%     438.596ms     438.596ms       0.000us         0.00%      62.893ms      62.893ms         
</code></pre>
<p>So, AddmmBackward0 is supposedly called during loss.backward(). But it is not accounted in the corresponding <code>record_function</code>. The row for the backward pass block comes much later, with trivial CUDA total time:</p>
<pre><code class="lang-auto">                                               backward        27.97%     809.707ms        41.77%        1.209s        1.209s       0.000us         0.00%     409.000us     409.000us        
</code></pre>
<p>So it does account for significant CPU time, but almost nothing in the CUDA total column.</p>
<p>Am I doing something silly in the script above, or is there any current difficulty in aggregating the *Backward kernels in the profiler?</p>