LocalNormalizedCrossCorrelationLoss and TF32 numerical stability #6525

Activity

1. wyli

   added

   Contribution wanted

   Feature request

   on May 17, 2023
2. qingpeng9802 commented on Jul 19, 2023

   

   qingpeng9802

   on Jul 19, 2023

   Contributor

   More actions

   After some investigations, I believe this is probably not a numerical stability issue. See https://dev-discuss.pytorch.org/t/pytorch-and-tensorfloat32/504 for more similar issues. The issue is NGC container sets the default value torch.backends.cuda.matmul.allow_tf32 = False to True somewhere (since tf32 is a trick for significantly accelerating training).
   Therefore, the issue only occurs in the NGC container but not in the pytorch official docker.

   Maybe we should mention the tf32 issue to the users somewhere in the documentation?
3. wyli commented on Jul 19, 2023

   

   wyli

   on Jul 19, 2023

   ContributorAuthor

   More actions

   thanks, before we have a workaround, perhaps adding a warning message in the constructor when the flag is True?

   MONAI/monai/losses/image_dissimilarity.py

   Line 90 in 42e3674

   self.ndim = spatial_dims
4. qingpeng9802 commented on Jul 20, 2023

   

   qingpeng9802

   on Jul 20, 2023

   Last edited by qingpeng9802

   Contributor

   More actions

   It is okay to add a warning for the loss. However, my larger concern is that other operations in monai will be also affected by the tf32 issue (since all operations uses cuda.matmul are affected). This may lead to significant reproducibility issues.

   My proposal is adding something like
   https://github.com/Lightning-AI/lightning/pull/16037/files#diff-909e246d6c36514f952ae5023bd9fbcc3e8f2c6a0837ebf81d7dc96790b5f938R190-R210
   to related classes/functions in monai. Then, monai will print warnings when the flag is True. Not sure when it is better to print warnings, maybe during import? Maybe warnings can be suppressed when the flage is explicitly set by users, but it seems technically challenging.
   &
   adding a part in the documentation to educate users how to use tf32 properly.
5. qingpeng9802 commented on Jul 21, 2023

   

   qingpeng9802

   on Jul 21, 2023

   Last edited by qingpeng9802

   Contributor

   More actions

   By checking https://catalog.ngc.nvidia.com/orgs/nvidia/containers/pytorch/layers Tag 23.06-py3, LAYERS: 08aa16a90c shows that ENV TORCH_ALLOW_TF32_CUBLAS_OVERRIDE=1, which modifies the flag in the NGC container. Therefore, the scope of the warning can be narrowed down to the users who have this environment variable, which is more feasible.
6. wyli commented on Jul 21, 2023

   

   wyli

   on Jul 21, 2023

   ContributorAuthor

   More actions

   thanks @qingpeng9802, I think that's a very good point and I create a separate ticket to track the features #6754.

   👍React with 👍1Reacted by Qingpeng Li
7. vikashg

   closed this as completedon Dec 20, 2023