Fix contrib openfold_triton tests (#1938)

crcrpar · Copilot · web-flow · commit d0e7be181bcc · 2025-10-29T04:58:46.000+09:00
* use `torch.backends.cudnn.flags` instead

Signed-off-by: Masaki Kozuki &lt;mkozuki@nvidia.com&gt;

* import run_tests from common_utils

Signed-off-by: Masaki Kozuki &lt;mkozuki@nvidia.com&gt;

* use TRITON_ALLOW_NON_CONSTEXPR_GLOBALS as a workaround

Signed-off-by: Masaki Kozuki &lt;mkozuki@nvidia.com&gt;

* Update apex/contrib/test/openfold_triton/test_fused_adam_swa.py

Co-authored-by: Copilot &lt;175728472+Copilot@users.noreply.github.com&gt;

---------

Signed-off-by: Masaki Kozuki &lt;mkozuki@nvidia.com&gt;
Co-authored-by: Copilot &lt;175728472+Copilot@users.noreply.github.com&gt;
diff --git a/apex/contrib/test/openfold_triton/test_fused_adam_swa.py b/apex/contrib/test/openfold_triton/test_fused_adam_swa.py
@@ -11,7 +11,7 @@
 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 # See the License for the specific language governing permissions and
 # limitations under the License.
-
+import os
 from itertools import chain
 import random
 import unittest
@@ -88,9 +88,20 @@ def setUp(self):
         self._seed = 19260817
         random.seed(self._seed)
         torch.manual_seed(self._seed)
-        torch.backends.cudnn.deterministic = True
+        # FIXME: correctly fix: """NameError("Cannot access global variable _DTYPE2TRITON from within @jit'ed function.
+        # Triton kernels can only access global variables that are instanstiated as constexpr (`x = triton.language.constexpr(42)`).
+        # Note that this is different from annotating a variable as constexpr (`x: triton.language.constexpr = 42`), which is not supported.
+        # Alternatively, set the envvar TRITON_ALLOW_NON_CONSTEXPR_GLOBALS=1, but we do not promise to support this forever.")"""
+        os.environ["TRITON_ALLOW_NON_CONSTEXPR_GLOBALS"] = "1"
+
+    def tearDown(self):
+        os.environ.pop("TRITON_ALLOW_NON_CONSTEXPR_GLOBALS", None)
 
     def test_fused_update_on_random_data(self):
+        with torch.backends.cudnn.flags(deterministic=True):
+            self._run_fused_update_on_random_data()
+
+    def _run_fused_update_on_random_data(self):
         device = torch.device("cuda:0")
         compute_dtype = torch.float32
         state_dtype = torch.float64
diff --git a/apex/contrib/test/openfold_triton/test_sync_triton_auto_tune_cache_across_gpus.py b/apex/contrib/test/openfold_triton/test_sync_triton_auto_tune_cache_across_gpus.py
@@ -1,18 +1,21 @@
 import os
+
 import torch
 import torch.distributed as dist
+from torch.testing._internal.common_utils import run_tests
 from torch.testing._internal.common_distributed import (
     MultiProcessTestCase,
     requires_nccl,
     skip_if_lt_x_gpu,
-    run_tests,
 )
+
 from apex.contrib.openfold_triton import (
     LayerNormSmallShapeOptImpl,
     sync_triton_auto_tune_cache_across_gpus,
     _tuneable_triton_kernels,
 )
 
+
 class SyncTritonAutoTuneCacheTest(MultiProcessTestCase):
     device_type = "cuda"
     def __init__(self, *args, **kwargs) -> None:
