DTensor: support linear (#2422)

Author: kshitij12345

thunder/executors/nvfuserex_impl.py

@@ -2434,6 +2434,7 @@ def linear(
 
 
 register_supported(PrimIDs.LINEAR, linear, _linear_check)
+register_supported(DTensorPrimIDs.LINEAR, linear, _linear_check)
 
 
 def _matmul_check(

thunder/tests/distributed/test_dtensor.py

@@ -34,10 +34,45 @@
 }
 
 
+# NOTE: OpInfo may use `clang` or `ltorch` ops to be jitted with thunder.jit.
+#       However, for the current DTensor implementation, we add a dispatch in the `torch` operation lookaside
+#       to choose between DTensor supported symbol (from `dtensor_torch_and_prims.py`) or the usual `ltorch` symbol.
+#       This is why we need to make sure that the OpInfo uses PyTorch native op as `op` which is passed to thunder.jit.
+class DTensorOpInfo:
+    def __init__(self, *, name, op, torch_reference, supports_grad, sample_inputs):
+        self.name = name
+        assert "torch" in op.__module__, "OpInfo must use PyTorch native op as `op` which is passed to thunder.jit"
+        self.op = op
+        self.torch_reference = torch_reference
+        # NOTE: Not all DTensor ops support grad initially, use this to disable grad tests for them
+        self.supports_grad = supports_grad
+        # NOTE: This should generally reuse the sample_inputs from the OpInfo
+        self.sample_inputs = sample_inputs
+
+
 # DTensor supported ops
-dtensor_supported_ops = ("reshape",)
+dtensor_supported_opinfos = (
+    DTensorOpInfo(
+        name="reshape",
+        op=torch.reshape,
+        torch_reference=torch.reshape,
+        supports_grad=True,
+        sample_inputs=get_opinfo("reshape").sample_inputs,
+    ),
+    DTensorOpInfo(
+        name="linear",
+        op=torch.nn.functional.linear,
+        torch_reference=torch.nn.functional.linear,
+        supports_grad=False,
+        sample_inputs=get_opinfo("linear").sample_inputs,
+    ),
+)
 
-dtensor_supported_opinfos = [get_opinfo(op) for op in dtensor_supported_ops]
+skip_opinfos = (
+    # RuntimeError: Metadata (placement and mesh) has changed for cotangent between tracing and runtimeduring tracing
+    # it was Spec(S(1) on (1, 2, 1)) but at runtime it is Spec(S(1) on (1, 2, 1)).
+    "reshape",
+)
 
 
 @unittest.skipUnless(
@@ -189,15 +224,20 @@ def fn(x):
         lambda op, executor: op.name + "_" + executor,
     )
     def test_dtensor_opinfo(self, op: OpInfo, executor):
+        if op.name in skip_opinfos:
+            raise unittest.SkipTest(f"test_dtensor_opinfo: Skipping {op.name} as it is in skip_opinfos")
+
         # NOTE: This test only tests for dtype=torch.float32 and requires_grad=True
         #       not for all dtype which are supported by the operation.
         num_devices = self.world_size
         mesh = DeviceMesh("cuda", list(range(num_devices)))
 
-        thunder_op = thunder.jit(op.op, executors=executors_map[executor].executors_list())
+        thunder_op = thunder.jit(op.op, executors=executors_map[executor].executors_list(), nv_enable_linear=True)
+        torch_op = op.torch_reference
 
         tested_sample_count = 0
-        for sample in op.sample_inputs("cpu", dtypes.float32, requires_grad=True):
+
+        for sample in op.sample_inputs("cpu", dtypes.float32, requires_grad=op.supports_grad):
             # DTensorConverter converts inputs tensors to DTensor and creates DTensor
             # with possible placements based on the input shapes.
             # See - https://github.com/pytorch/pytorch/blob/eaa5d9d3d3dc642832b269b184f0c3ab8c990274/torch/testing/_internal/distributed/_tensor/common_dtensor.py#L521
@@ -206,8 +246,6 @@ def test_dtensor_opinfo(self, op: OpInfo, executor):
                 if not dtensor_converter.successful():
                     continue
 
-                torch_op = op.torch_reference
-
                 # Computes PyTorch result
                 try:
                     torch_result = torch_op(*dtensor_args, **dtensor_kwargs)
@@ -220,34 +258,38 @@ def test_dtensor_opinfo(self, op: OpInfo, executor):
                 thunder_result = thunder_op(*dtensor_args, **dtensor_kwargs)
                 torch.testing.assert_close(thunder_result, torch_result)
 
-                torch_flats, _ = tree_flatten((dtensor_args, dtensor_kwargs))
-                torch_result = filter_differentiable_outputs(torch_result)
-                if torch_result == []:
-                    raise RuntimeError("test_dtensor_opinfo: Expected atleast 1 differentiable output.")
-
-                grads = []
-                assert isinstance(torch_result, torch.Tensor) or isinstance(torch_result, Sequence), (
-                    "test_dtensor_opinfo:Expected a single torch tensor or a sequence of torch tensors"
-                )
-                if isinstance(torch_result, Sequence):
-                    for x in torch_result:
-                        assert isinstance(x, torch.Tensor), (
-                            "test_dtensor_opinfo: Expected a single torch tensor or a sequence of torch tensors"
-                        )
-                        if is_output_differentiable(x):
-                            grads.append(torch.ones_like(x))
-                else:
-                    if is_output_differentiable(torch_result):
-                        grads = [torch.ones_like(torch_result)]
-
-                torch_tensors_requiring_grad = tuple(
-                    f for f in torch_flats if isinstance(f, torch.Tensor) and f.requires_grad
-                )
-                torch_grad_result = torch.autograd.grad(torch_result, torch_tensors_requiring_grad, grads)
-
-                thunder_result = filter_differentiable_outputs(thunder_result)
-                thunder_grad_result = torch.autograd.grad(thunder_result, torch_tensors_requiring_grad, grads)
-                torch.testing.assert_close(thunder_grad_result, torch_grad_result)
+                trace = thunder.last_traces(thunder_op)[0]
+                assert any("dtensor" in bsym.sym.name for bsym in trace.bound_symbols)
+
+                if op.supports_grad:
+                    torch_flats, _ = tree_flatten((dtensor_args, dtensor_kwargs))
+                    torch_result = filter_differentiable_outputs(torch_result)
+                    if torch_result == []:
+                        raise RuntimeError("test_dtensor_opinfo: Expected atleast 1 differentiable output.")
+
+                    grads = []
+                    assert isinstance(torch_result, torch.Tensor) or isinstance(torch_result, Sequence), (
+                        "test_dtensor_opinfo:Expected a single torch tensor or a sequence of torch tensors"
+                    )
+                    if isinstance(torch_result, Sequence):
+                        for x in torch_result:
+                            assert isinstance(x, torch.Tensor), (
+                                "test_dtensor_opinfo: Expected a single torch tensor or a sequence of torch tensors"
+                            )
+                            if is_output_differentiable(x):
+                                grads.append(torch.ones_like(x))
+                    else:
+                        if is_output_differentiable(torch_result):
+                            grads = [torch.ones_like(torch_result)]
+
+                    torch_tensors_requiring_grad = tuple(
+                        f for f in torch_flats if isinstance(f, torch.Tensor) and f.requires_grad
+                    )
+                    torch_grad_result = torch.autograd.grad(torch_result, torch_tensors_requiring_grad, grads)
+
+                    thunder_result = filter_differentiable_outputs(thunder_result)
+                    thunder_grad_result = torch.autograd.grad(thunder_result, torch_tensors_requiring_grad, grads)
+                    torch.testing.assert_close(thunder_grad_result, torch_grad_result)
 
                 # Increment tested sample count
                 tested_sample_count += 1

thunder/tests/test_dynamo.py

@@ -841,7 +841,7 @@ def find_target_module(model, target_module_name):
     assert submodule is not None
     for n in submodule.graph.nodes:
         if n.op == "call_function":
-            assert isinstance(n.target, Symbol)
+            assert isinstance(n.target, Symbol) or callable(n.target)
 
 
 @instantiate(

thunder/torch/experimental/dtensor_torch_and_prims.py

@@ -35,6 +35,7 @@ class DTensorPrimIDs(Enum):
     RESHAPE = auto()
     CONVERT_ELEMENT_TYPE = auto()
     BROADCAST_IN_DIM = auto()
+    LINEAR = auto()
 
 
 dtensor_torchsymbol = partial(torchsymbol, allow_tensor_subclass_proxy=True)
@@ -241,6 +242,33 @@ def dtensor_broadcast_in_dim_meta(a, shape, broadcast_dimensions):
 pytorchex.register_implementation(dtensor_broadcast_in_dim_prim, dtensor_broadcast_in_dim_prim_impl)
 
 
+def dtensor_linear_meta(a, w, bias):
+    output = run_with_fake_tensor(torch.nn.functional.linear, a, w, bias)
+    local_tensor_proxy = TensorProxy(like=a.local_tensor)
+    local_tensor_proxy = TensorProxy(
+        like=a.local_tensor, shape=output._local_tensor.shape, dtype=dtypes.to_dtype(output._local_tensor.dtype)
+    )
+    spec = output._spec
+    spec_proxy = AnyProxy(spec, history=a.history)
+    return create_dtensor_proxy_from_proxies(local_tensor_proxy, spec_proxy, False)
+
+
+# TODO: Add grad rule once the prims used for linear grad-rule are available.
+dtensor_linear_prim = make_prim(DTensorPrimIDs.LINEAR, "dtensor_linear_prim", meta=dtensor_linear_meta)
+
+dtensor_linear_prim_impl = pytorchex.register_operator(
+    "dtensor_linear_prim", like=dtensor_linear_prim, fn=torch.nn.functional.linear
+)
+
+pytorchex.register_implementation(dtensor_linear_prim, dtensor_linear_prim_impl)
+
+
+@dtensor_torchsymbol(torch.nn.functional.linear, id="dtensor.torch.nn.functional.linear")
+def dtensor_linear(a: TensorLike, w: TensorLike, bias: None | TensorLike = None) -> TensorLike:
+    return dtensor_linear_prim(a, w, bias)
+
+
 def register_dtensor_torch_and_prims():
     register_function_for_dtensor(torch.mul, ltorch.mul, dtensor_mul, is_method=True)
     register_function_for_dtensor(torch.reshape, ltorch.reshape, dtensor_reshape, is_method=True)
+    register_function_for_dtensor(torch.nn.functional.linear, ltorch.linear, dtensor_linear, is_method=False)