Validating CrossEntropyLoss Performance

[kevinstephano]

I made this code snippet to show perf of CrossEntropyLoss.

import torch

class MyLoss(torch.nn.Module):
    def __init__(self):
        super(MyLoss, self).__init__()
        self.loss = torch.nn.CrossEntropyLoss()

    def forward(self, inputs, targets):
        out = self.loss(inputs, targets)
        return out

inputs = torch.randn(8192, 32768, device='cuda')
targets = torch.randint(32767, (8192,), device='cuda')

model = torch.compile(MyLoss())

for _ in range(5):
    out = model(inputs, targets)

Test command:

nsys nvprof --print-gpu-trace python my_loss.py

Sample output on A100:

Tensor Sizes:
inputs = [8192, 32768]
targets = [8192[

Kernel1: 1.222ms
Kernel2: 37.1us

4202242471        1222152    1059  8192     1     1   256     1     1       39         0.000         0.001                                                     NVIDIA A100 80GB PCIe (0)    1     7  triton__0d1d2d3d4d                                                                                  
4203465679          37121    1072     1     1     1   256     1     1      184         0.000         0.008                                                     NVIDIA A100 80GB PCIe (0)    1     7  triton__0d1d2d3d4d56d

[csarofeen]

What's the effective bandwidth of the kernels?

[naoyam]

Kernel1 is around 880 GB/s. The other one is just a few GB/s and is negligible.

[kevinstephano]

Closing, old.