Harden the test with dynamic shapes (#3807)

tests/python/test_multidevice.py

@@ -115,19 +115,14 @@ def multidevice_schedule(self):
 
 @pytest.mark.mpi
 def test_linear_loop_split(multidevice_test):
-    class Model(FusionDefinition):
-        def __init__(self, num_devices, batch, sequence, hidden):
-            super().__init__()
-            self._num_devices = num_devices
-            self._batch = batch
-            self._sequence = sequence
-            self._hidden = hidden
+    d = multidevice_test.size
+    mesh = nvfuser.DeviceMesh(range(d))
 
+    class Model(FusionDefinition):
         def definition(self):
-            d, b, s, e = self._num_devices, self._batch, self._sequence, self._hidden
-            self.inp = self.define_tensor([b, s, e])
-            self.weight = self.define_tensor([d * e, e])
-            self.bias = self.define_tensor([d * e])
+            self.inp = self.define_tensor([-1, -1, -1])
+            self.weight = self.define_tensor([-1, -1])
+            self.bias = self.define_tensor([-1])
             self.out = self.ops.linear(self.inp, self.weight, self.bias)
             self.add_output(self.out)
 
@@ -147,9 +142,6 @@ def multidevice_schedule(self):
             self.sched.parallelize(self.out, -3, nvfuser.ParallelType.mesh_x)
             self.sched.set_allocation_as_loop(self.out)
 
-    d = multidevice_test.size
-    mesh = nvfuser.DeviceMesh(range(d))
-
     torch.cuda.set_device(multidevice_test.local_rank)
 
     b, s, e = 2, 1024, 768
@@ -161,7 +153,7 @@ def multidevice_schedule(self):
     unsharded_bias_tensor = torch.randn(d * e)
     sharded_bias_tensor = multidevice_test.shard_tensor(unsharded_bias_tensor, 0, mesh)
 
-    fd = Model(d, b, s, e)
+    fd = Model()
     (out_tensor,) = fd.execute([inp_tensor, sharded_weight_tensor, sharded_bias_tensor])
 
     # [b, s, d*e]
@@ -229,18 +221,13 @@ def multidevice_schedule(self) -> None:
 
 @pytest.mark.mpi
 def test_matmul_loop_split(multidevice_test):
-    class Model(FusionDefinition):
-        def __init__(self, num_devices, batch, sequence, hidden):
-            super().__init__()
-            self._num_devices = num_devices
-            self._batch = batch
-            self._sequence = sequence
-            self._hidden = hidden
+    d = multidevice_test.size
+    mesh = nvfuser.DeviceMesh(range(d))
 
+    class Model(FusionDefinition):
         def definition(self):
-            d, b, s, e = self._num_devices, self._batch, self._sequence, self._hidden
-            self.inp = self.define_tensor([b, s, e])
-            self.weight = self.define_tensor([e, d * e])
+            self.inp = self.define_tensor([-1, -1, -1])
+            self.weight = self.define_tensor([-1, -1])
             self.out = self.ops.matmul(self.inp, self.weight)
             self.add_output(self.out)
 
@@ -259,10 +246,6 @@ def multidevice_schedule(self):
             self.sched.parallelize(self.out, -3, nvfuser.ParallelType.mesh_x)
             self.sched.set_allocation_as_loop(self.out)
 
-    d = multidevice_test.size
-    mesh = nvfuser.DeviceMesh(range(d))
-    rank = multidevice_test.rank
-
     torch.cuda.set_device(multidevice_test.local_rank)
 
     b, s, e = 2, 1024, 768
@@ -272,7 +255,7 @@ def multidevice_schedule(self):
         unsharded_weight_tensor, -1, mesh
     )
 
-    fd = Model(d, b, s, e)
+    fd = Model()
     (out_tensor,) = fd.execute([inp_tensor, sharded_weight_tensor])
 
     # [b, s, d*e]
@@ -286,16 +269,16 @@ def multidevice_schedule(self):
 
 @pytest.mark.mpi
 def test_matmul_allreduce_loop_split(multidevice_test):
-    d, b, s, e = multidevice_test.size, 1, 4, 8
+    d = multidevice_test.size
     mesh = nvfuser.DeviceMesh(range(d))
 
     class Model(FusionDefinition):
         def definition(self) -> None:
             self.inp = self.define_tensor(
-                [b * s, d * e], contiguity=True, dtype=DataType.Half
+                [-1, -1], contiguity=True, dtype=DataType.Half
             )
             self.weight = self.define_tensor(
-                [d * e, e], contiguity=True, dtype=DataType.Half
+                [-1, -1], contiguity=True, dtype=DataType.Half
             )
             self.out = self.ops.matmul(self.inp, self.weight)
             self.add_output(self.out)
@@ -323,10 +306,9 @@ def multidevice_schedule(self) -> None:
             self.sched._set_device_mesh(self.local_out, mesh)
             self.sched.parallelize(self.local_out, -2, nvfuser.ParallelType.mesh_x)
 
-    rank = multidevice_test.rank
-
     torch.cuda.set_device(multidevice_test.local_rank)
 
+    b, s, e = 1, 4, 8
     unsharded_inp = torch.randn(b * s, d * e, dtype=torch.half)
     unsharded_weight = torch.randn(d * e, e, dtype=torch.half)
     sharded_inp = multidevice_test.shard_tensor(unsharded_inp, -1, mesh)