Representing DTensor in thunder traces

thunder/core/jit_ext.py

@@ -74,7 +74,12 @@
 from thunder.torch import _torch_to_thunder_function_map
 from thunder.clang import _clang_fn_set
 from thunder.core.pytree import tree_map, tree_iter
-from thunder.torch.experimental.dtensor_torch_and_prims import register_dtensor_torch_and_prims
+from thunder.torch.experimental.dtensor_proxy import is_dtensor_proxy
+from thunder.torch.experimental.dtensor_torch_and_prims import (
+    check_dtensor_spec_repr,
+    handle_check_dtensor_spec_in_prologue,
+    register_dtensor_torch_and_prims,
+)
 
 # TODO: Find a better place to register these ops (mostly in thunder/torch/__init__.py but without cyclical dependency).
 register_dtensor_torch_and_prims()
@@ -278,14 +283,11 @@ def proxify(self, value: WrappedValue) -> Any:
             if p is not uvalue:
                 value.register_proxy(p)
 
-            from thunder.torch.experimental.dtensor_proxy import is_dtensor_proxy
-            from thunder.torch.experimental import dtensor_torch_and_prims
-
             # TODO: other caching modes
             co: CACHE_OPTIONS = get_cache_option()
             if co is CACHE_OPTIONS.CONSTANT_VALUES:
                 if is_dtensor_proxy(p):
-                    self.add_constraint((dtensor_torch_and_prims.check_dtensor_spec_repr, p, uvalue._spec))
+                    self.add_constraint((check_dtensor_spec_repr, p, uvalue._spec))
                 self.add_constraint((clang.check_tensor_shape_and_metadata, p))
             elif co is CACHE_OPTIONS.SYMBOLIC_VALUES:
                 # TODO: establish guarding logic to allow non-broadcast shape change
@@ -1891,8 +1893,6 @@ def is_variableified_tensorproxy(v: Variable | Proxy) -> Proxy:
                         unpack(s)
 
             # Add checks for local tensor, mesh and placment of a DTensor
-            from thunder.torch.experimental.dtensor_torch_and_prims import handle_check_dtensor_spec_in_prologue
-
             if handle_check_dtensor_spec_in_prologue(prim, prologue_trace, args):
                 continue
 

thunder/torch/experimental/dtensor_utils.py

@@ -2,8 +2,6 @@
 from typing import Any
 
 import torch
-from torch._guards import TracingContext
-from functorch.compile import aot_function
 from torch.distributed.tensor import DTensor
 
 from thunder.core.pytree import tree_map
@@ -18,33 +16,28 @@
 from thunder.core.trace import from_trace
 
 from torch._subclasses.fake_tensor import FakeTensorMode
-from torch._guards import TracingContext, tracing
 
 
-def get_fx_graph_and_output(torch_op, *args, **kwargs) -> tuple[torch.fx.GraphModule, Sequence[Any]]:
+def run_with_fake_tensor(torch_op, *args, **kwargs):
     """
-    Generate a Torch FX graph and the corresponding output by tracing a TraceCtx object.
+    Run a torch operation with fake tensors and return the output.
 
     Args:
-        trc (TraceCtx): Trace which will be used to generate the FX Graph
-        args_for_fx (Sequence[Any]): A sequence of input arguments to be passed during the FX tracing.
+        torch_op: The torch operation to execute
+        *args: Arguments to pass to the torch operation
+        **kwargs: Keyword arguments to pass to the torch operation
 
     Returns:
-        tuple[torch.fx.GraphModule, Sequence[Any]]:
-            - A `torch.fx.GraphModule` object representing the traced computation graph in FX.
-            - Output(s) from evaluating the graph with the given inputs.
+        The output of the torch operation executed with fake tensors
     """
 
     def f(*args, **kwargs):
         return torch_op(*args, **kwargs)
 
-    tracing_ctx = TracingContext.try_get()
-    fake_mode = tracing_ctx.fake_mode
-
-    with fake_mode:
+    with FakeTensorMode():
 
         def materialize_fake_tensors(t):
-            # `aot_function` can't handle these proxy types.
+            # Convert proxy types to fake tensors.
             if isinstance(t, NumberProxy):
                 return t.value
 
@@ -63,39 +56,7 @@ def materialize_fake_tensors(t):
 
         args, kwargs = tree_map(materialize_fake_tensors, (args, kwargs))
 
-    def aot_wrapped_fn(*args, **kwargs):
-        return f(*args, **kwargs)
-
-    fwd_graph = None
-
-    def get_graph(name):
-        def f(fx_g: torch.fx.GraphModule, inps):
-            nonlocal fwd_graph
-            assert name != "backward", "aot shouldn't have reached backward as it will be handled by thunder"
-            if name == "forward":
-                fwd_graph = fx_g
-            return fx_g
-
-        return f
-
-    aot_output = aot_function(aot_wrapped_fn, fw_compiler=get_graph("forward"), bw_compiler=get_graph("backward"))(
-        *args, **kwargs
-    )
-    # Example value of `fwd_graph`
-    # def forward(self, arg0_1, arg1_1):
-    #   mul = torch.ops.aten.mul.Tensor(arg0_1, arg1_1);  arg0_1 = arg1_1 = None
-    #   return (mul,)
-
-    # Example value of `aot_output`
-    # DTensor(local_tensor=FakeTensor(..., device='cuda:0', size=(8, 16)), device_mesh=DeviceMesh('cuda', [0, 1]), placements=(Shard(dim=0),))
-
-    return fwd_graph, aot_output
-
-
-def run_with_fake_tensor(torch_op, *args, **kwargs):
-    with tracing(TracingContext(FakeTensorMode())):
-        _, output = get_fx_graph_and_output(torch_op, *args, **kwargs)
-    return output
+    return f(*args, **kwargs)
 
 
 def check_dtensor_cotangent_metadata(dtensor, metadata):