Tutorial deep_atlas shows a huge negative loss in PyTorch 22.09-23.03 on some GPUs

[mingxin-zheng]

Describe the bug

On V100/A100, the tutorial deep_atlas is reporting correct number on the reg_loss in the 42nd cell when printing the losses

print(f"\n\nBest reg_net validation loss: {best_reg_validation_loss}")
print(f"Best seg_net validation loss: {best_seg_validation_loss}")

On Geforce 3080Ti/3090Ti and L40, the reg loss is a large negative number, for example:

To Reproduce
Steps to reproduce the behavior:

1. Install MONAI 1.0.1 or 1.2rc4 on Nvidia PyTorch 22.09 or 23.03 container

[binliunls]

Hi @wyli ,
When I ran this tutorial on the 3090Ti, I got some negative values in tensor t2_sum - t_avg*t_sum in this line as shown below. I think this should not be correct and the values should always be positive. What do you think about this?

Thanks,
Bin

[wyli]

agreed, looks like some precision issue, I think this should be reported back to the framework team if we can narrow it down to a minimal reproducible example... perhaps you can try to save the inputs to the lncc_loss with torch.save

[binliunls]

I think it is a NGC pytorch docker related precision problem. Here is the code to show the bug.

import torch
import torch.nn.functional as F

torch.backends.cudnn.benchmark=False
torch.manual_seed(0)
torch.set_printoptions(precision=5)

x = (1000 * torch.rand((2, 1, 96, 96, 96)) + 2000).to("cuda")
kernel = torch.tensor([[[[[1.0]], [[1.0]], [[1.0]]]]]).to(x)

print(x[1, 0, 68 : 68 + 3, 57 : 57 + 3, 43 : 43 + 3])
torch_y = F.conv3d(x, weight=kernel, groups=1)
print(torch_y[1, 0, 68 : 68 + 3, 57 : 57 + 3, 43 : 43 + 3])

Here are three experiments.

1. On a 3090Ti with NGC pytorch 23.03 container, the output is:

tensor([[[2886.76953, 2181.87695, 2682.72485],
         [2972.62549, 2509.76709, 2155.61938],
         [2899.60498, 2226.87720, 2881.76416]],

        [[2009.33215, 2248.57861, 2991.17578],
         [2323.16748, 2200.64209, 2854.13525],
         [2977.41748, 2440.96704, 2035.22681]],

        [[2787.95068, 2324.95898, 2909.67480],
         [2139.43262, 2556.39453, 2748.65894],
         [2487.08765, 2212.69141, 2001.22424]]], device='cuda:0')


tensor([[[7683., 6754., 8584.],
         [7436., 7266., 7758.],
         [8366., 6878., 6918.]],

        [[7459., 6926., 8040.],
         [7132., 6886., 8296.],
         [7704., 6784., 6508.]],

        [[8122., 6876., 7858.],
         [6894., 6720., 7742.],
         [7150., 6880., 7417.]]], device='cuda:0')

2. On an A100 with NGC pytorch 23.03 container, the output is:

tensor([[[2886.76953, 2181.87695, 2682.72485],
         [2972.62549, 2509.76709, 2155.61938],
         [2899.60498, 2226.87720, 2881.76416]],

        [[2009.33215, 2248.57861, 2991.17578],
         [2323.16748, 2200.64209, 2854.13525],
         [2977.41748, 2440.96704, 2035.22681]],

        [[2787.95068, 2324.95898, 2909.67480],
         [2139.43262, 2556.39453, 2748.65894],
         [2487.08765, 2212.69141, 2001.22424]]], device='cuda:0')


tensor([[[7684.05225, 6755.41455, 8583.57520],
         [7435.22559, 7266.80371, 7758.41406],
         [8364.11035, 6880.53564, 6918.21533]],

        [[7458.84961, 6926.97754, 8038.90674],
         [7130.95898, 6887.61133, 8296.36719],
         [7701.62305, 6786.12939, 6507.63770]],

        [[8122.47559, 6876.62500, 7858.00000],
         [6893.88086, 6720.79150, 7741.64746],
         [7147.68799, 6880.23145, 7416.61182]]], device='cuda:0')

3. On a 3090Ti with pytorch official docker 2.0.1-cuda11.7-cudnn8-devel, the output is:

tensor([[[2886.76953, 2181.87695, 2682.72485],
         [2972.62549, 2509.76709, 2155.61938],
         [2899.60498, 2226.87720, 2881.76416]],

        [[2009.33215, 2248.57861, 2991.17578],
         [2323.16748, 2200.64209, 2854.13525],
         [2977.41748, 2440.96704, 2035.22681]],

        [[2787.95068, 2324.95898, 2909.67480],
         [2139.43262, 2556.39453, 2748.65894],
         [2487.08765, 2212.69141, 2001.22424]]], device='cuda:0')


tensor([[[7684.05225, 6755.41455, 8583.57520],
         [7435.22559, 7266.80371, 7758.41406],
         [8364.11035, 6880.53564, 6918.21533]],

        [[7458.84961, 6926.97754, 8038.90674],
         [7130.95898, 6887.61133, 8296.36719],
         [7701.62305, 6786.12939, 6507.63770]],

        [[8122.47559, 6876.62500, 7858.00000],
         [6893.88086, 6720.79150, 7741.64746],
         [7147.68799, 6880.23145, 7416.61182]]], device='cuda:0')

It seems that the NGC pytorch 23.03 container on the 3090Ti just rounds the output number somehow. And this leads to a large round-off error, since this logic is recursively used in the separable_filtering function of the loss function.

[wyli]

thanks, maybe you should fix the random seed or try to use exactly the same input to make it easier to see the differences? and perhaps set the datatypes and the print options to make sure they are not just printing differences https://pytorch.org/docs/stable/generated/torch.set_printoptions.html

[binliunls]

thanks, maybe you should fix the random seed or try to use exactly the same input to make it easier to see the differences? and perhaps set the datatypes and the print options to make sure they are not just printing differences https://pytorch.org/docs/stable/generated/torch.set_printoptions.html

Hi @wyli ,
Thank you for the advices! I have updated the comment accordingly.

Thanks,
Bin

[binliunls]

One more thing. I think this relates to the tf32 precision of the cuda computation. If I set the tf32 precision off in the first experiment mentioned above, the output becomes normal and the tutorial could run successfully.

These two lines need to be added to the program to turn off the tf32 precision.

torch.backends.cuda.matmul.allow_tf32 = False
torch.backends.cudnn.allow_tf32 = False

[wyli]

so do you think setting allow_tf32=False will address the original issue? sorry I didn't read your comments properly

[binliunls]

Sorry, I forgot to add the tutorial result. I ran the deep atlas tutorial on the 3090Ti NGC pytorch docker with allow_tf32=False. The output turned out normally to me. Haven't got a chance to run it on L40 and 3080Ti.

[wyli]

ok thanks, we can add allow_tf32=False to the notebook as a solution for this ticket, what do you think?

[binliunls]

ok thanks, we can add allow_tf32=False to the notebook as a solution for this ticket, what do you think?

Agreed. Should we add allow_tf32=False to the beginning of the whole tutorial or just add it to the loss cell and reset it to the previous value at the end of that cell?

[wyli]

created a ticket to track the issue Project-MONAI/MONAI#6525, feel free to create a PR if you have a solution for it...