DTensor: skip nvfuser test if NVFUSER_DISABLE=multidevice is set (#2724)

Co-authored-by: beverlylytle <[email protected]>

Author: kshitij12345

thunder/tests/distributed/test_dtensor.py

@@ -2,6 +2,7 @@
 from itertools import product
 from collections.abc import Sequence
 from looseversion import LooseVersion
+import os
 
 import pytest
 import torch
@@ -341,6 +342,9 @@ def test_dtensor_columnwise_parallel(self, jit_fn):
         ],
     )
     def test_dtensor_grouped_mm(self, executor, input_shardings):
+        if executor == "nvfuser" and "multidevice" in os.environ.get("NVFUSER_DISABLE", ""):
+            raise unittest.SkipTest("test_dtensor_grouped_mm: nvfuser multidevice is disabled")
+
         if LooseVersion(torch.__version__) < "2.8":
             raise unittest.SkipTest("test_dtensor_grouped_mm: torch._grouped_mm is not available in torch < 2.8")
 

thunder/tests/distributed/test_moe.py

@@ -1,5 +1,7 @@
 from functools import partial
 import copy
+import os
+import unittest
 
 import torch
 from torch.distributed.tensor.placement_types import Placement, Shard, Replicate
@@ -143,6 +145,9 @@ def parallelize_moe_model(model: llama4_moe.Llama4MoE, device_mesh: torch.distri
 
 class TestLlama4MoEDistributed(DistributedParallelTestCase):
     def test_llama4_moe_distributed(self):
+        if "multidevice" in os.environ.get("NVFUSER_DISABLE", ""):
+            raise unittest.SkipTest("test_llama4_moe_distributed: nvfuser multidevice is disabled")
+
         # Get world size
         world_size = self.world_size
         device = f"cuda:{self.rank}"