Repository: pytorch
Issue:  test/distributed/_tensor/experimental/test_tp_transform.py TensorParallelTest.test_tp_transform_e2e fail with functionalization V2
Labels: oncall: distributed,triaged,module: functionalization,oncall: pt2,module: inductor,module: pt2-dispatcher,module: reinplacing

python test/distributed/_tensor/experimental/test_tp_transform.py TensorParallelTest.test_tp_transform_e2e (With V2 enabled)

**what is going on ?**

we hit this error:
RuntimeError: Attempting to use FunctionalTensor on its own. Instead, please use it with a corresponding FunctionalTensorMode()
the function in that case is 

```
...
    exported_program = torch.export.export(
  File "/home/lsakka/pytorch/torch/export/__init__.py", line 366, in export
    return _export(
  File "/home/lsakka/pytorch/torch/export/_trace.py", line 1021, in wrapper
    raise e
  File "/home/lsakka/pytorch/torch/export/_trace.py", line 994, in wrapper
    ep = fn(*args, **kwargs)
  File "/home/lsakka/pytorch/torch/export/exported_program.py", line 116, in wrapper
    return fn(*args, **kwargs)
  File "/home/lsakka/pytorch/torch/export/_trace.py", line 1940, in _export
    export_artifact = export_func(  # type: ignore[operator]
  File "/home/lsakka/pytorch/torch/export/_trace.py", line 1234, in _strict_export
    return _strict_export_lower_to_aten_ir(
  File "/home/lsakka/pytorch/torch/export/_trace.py", line 1343, in _strict_export_lower_to_aten_ir
    aten_export_artifact = lower_to_aten_callback(
  File "/home/lsakka/pytorch/torch/export/_trace.py", line 641, in _export_to_aten_ir
    gm, graph_signature = transform(aot_export_module)(
  File "/home/lsakka/pytorch/torch/_functorch/aot_autograd.py", line 1262, in aot_export_module
    fx_g, metadata, in_spec, out_spec = _aot_export_function(
  File "/home/lsakka/pytorch/torch/_functorch/aot_autograd.py", line 1497, in _aot_export_function
    fx_g, meta = create_aot_dispatcher_function(
  File "/home/lsakka/pytorch/torch/_functorch/aot_autograd.py", line 524, in create_aot_dispatcher_function
    return _create_aot_dispatcher_function(
  File "/home/lsakka/pytorch/torch/_functorch/aot_autograd.py", line 762, in _create_aot_dispatcher_function
    compiled_fn, fw_metadata = compiler_fn(
  File "/home/lsakka/pytorch/torch/_functorch/_aot_autograd/jit_compile_runtime_wrappers.py", line 111, in aot_dispatch_export
    graph, _, _ = aot_dispatch_base_graph(
  File "/home/lsakka/pytorch/torch/_functorch/_aot_autograd/dispatch_and_compile_graph.py", line 136, in aot_dispatch_base_graph
    fw_module = _create_graph(
  File "/home/lsakka/pytorch/torch/_functorch/_aot_autograd/dispatch_and_compile_graph.py", line 54, in _create_graph
    fx_g = make_fx(
  File "/home/lsakka/pytorch/torch/fx/experimental/proxy_tensor.py", line 2147, in wrapped
    return make_fx_tracer.trace(f, *args)
  File "/home/lsakka/pytorch/torch/fx/experimental/proxy_tensor.py", line 2085, in trace
    return self._trace_inner(f, *args)
  File "/home/lsakka/pytorch/torch/fx/experimental/proxy_tensor.py", line 2056, in _trace_inner
    t = dispatch_trace(
  File "/home/lsakka/pytorch/torch/_compile.py", line 32, in inner
    return disable_fn(*args, **kwargs)
  File "/home/lsakka/pytorch/torch/_dynamo/eval_frame.py", line 654, in _fn
    return fn(*args, **kwargs)
  File "/home/lsakka/pytorch/torch/fx/experimental/proxy_tensor.py", line 1133, in dispatch_trace
    graph = tracer.trace(root, concrete_args)  # type: ignore[arg-type]
  File "/home/lsakka/pytorch/torch/fx/experimental/proxy_tensor.py", line 1652, in trace
    res = super().trace(root, concrete_args)
  File "/home/lsakka/pytorch/torch/_dynamo/eval_frame.py", line 654, in _fn
    return fn(*args, **kwargs)
  File "/home/lsakka/pytorch/torch/fx/_symbolic_trace.py", line 823, in trace
    (self.create_arg(fn(*args)),),
  File "/home/lsakka/pytorch/torch/fx/experimental/proxy_tensor.py", line 1188, in wrapped
    out = f(*tensors)
  File "<string>", line 1, in <lambda>
  File "/home/lsakka/pytorch/torch/_functorch/_aot_autograd/traced_function_transforms.py", line 693, in inner_fn
    outs = fn(*args)
  File "/home/lsakka/pytorch/torch/_functorch/_aot_autograd/traced_function_transforms.py", line 413, in _functionalized_f_helper
    f_outs = fn(*f_args)
  File "/home/lsakka/pytorch/torch/_functorch/_aot_autograd/traced_function_transforms.py", line 78, in inner_fn
    outs = fn(*args)
  File "/home/lsakka/pytorch/torch/_functorch/_aot_autograd/utils.py", line 184, in flat_fn
    tree_out = fn(*args, **kwargs)
  File "/home/lsakka/pytorch/torch/_functorch/_aot_autograd/traced_function_transforms.py", line 863, in functional_call
    out = PropagateUnbackedSymInts(mod).run(
  File "/home/lsakka/pytorch/torch/fx/interpreter.py", line 146, in run
    self.env[node] = self.run_node(node)
  File "/home/lsakka/pytorch/torch/fx/experimental/symbolic_shapes.py", line 6494, in run_node
    result = super().run_node(n)
  File "/home/lsakka/pytorch/torch/fx/interpreter.py", line 203, in run_node
    return getattr(self, n.op)(n.target, args, kwargs)
  File "/home/lsakka/pytorch/torch/fx/interpreter.py", line 275, in call_function
    return target(*args, **kwargs)
  File "/home/lsakka/pytorch/torch/fx/experimental/proxy_tensor.py", line 1236, in __torch_function__
    return func(*args, **kwargs)
  File "/home/lsakka/pytorch/torch/fx/experimental/proxy_tensor.py", line 1274, in __torch_function__
    return func(*args, **kwargs)
  File "/home/lsakka/pytorch/torch/_subclasses/functional_tensor.py", line 225, in __torch_dispatch__
    raise RuntimeError(
RuntimeError: Attempting to use FunctionalTensor on its own. Instead, please use it with a corresponding FunctionalTensorMode()

While executing %ones_like : [num_users=1] = call_function[target=torch.ones_like](args = (%x_2,), kwargs = {})
Original traceback:
  File "/home/lsakka/pytorch/test/distributed/_tensor/experimental/test_tp_transform.py", line 42, in forward
    return x + torch.ones_like(x)

```

**the code that throws:**
```
 def __torch_dispatch__(self, func, types, args=(), kwargs=None):
        unrecognized_types = [
            t
            for t in types
            if t not in [torch.Tensor, torch._subclasses.FakeTensor, FunctionalTensor]
        ]
        if unrecognized_types:
            not_implemented_log.debug(
                "FunctionalTensor unrecognized subclass(es): %s", unrecognized_types
            )
            return NotImplemented

        if kwargs is None:
            kwargs = {}
        # assert False
        print(func) 
        # FunctionalTensor needs to plumb all metadata requests to the inner tensor.
        # In theory we don't have to do this - but if we want to service metadata requests here,
        # we need to carefully make sure all metadata is accurate (including metadata mutations)
        if func in FunctionalTensor.metadata_fns:
            # All metadata accesses should be plumbed to the inner tensor, that way we don't have to worry
            # about the problem of keeping metadata in sync between the wrapper and inner tensor.
            # This also alleviates us from having to manually handle metadata mutations on the wrapper.
            assert len(kwargs) == 0
            if func in [
                torch.ops.aten.is_strides_like_format.default,
                torch.ops.aten.is_contiguous.memory_format,
            ]:
                assert len(args) == 2 and isinstance(args[0], FunctionalTensor)
                return func(torch._from_functional_tensor(args[0].elem), args[1])
            assert len(args) == 1 and isinstance(args[0], FunctionalTensor)

            return func(torch._from_functional_tensor(args[0].elem))
    

        # Originally I tried to implement my subclass without giving it a torch_dispatch, but I gave up:
        # - _make_wrapper_subclass requires a __torch_dispatch__
        # - If we want to use _make_subclass(), we have a problem: the subclass will share a TensorImpl with the inner tensor,
        #   which is of type FunctionalTensorWrapper! We explicitly do not want our wrapper to be a FunctionalTensorWrapper.
        # - If we use the default tensor.__new__(), we have another problem: it returns inner_tensor.alias(),
        #   which causes every subclass created above autograd to have autograd view metadata
        #   (in addition to also being a FunctionalTensorWrapper).
        raise RuntimeError(
            "Attempting to use FunctionalTensor on its own. Instead, please use it with a corresponding FunctionalTensorMode()"
        )

```
The input to that function is aten.detach which is not in  FunctionalTensor.metadata_fns.

**Why does it happen only with V2?**
 aten.detach is not part of the original graph. but the computation of the base tensor in inference mode end up making us tracing through this detach function? that is if i remove this code from the functional tensor constructor the problem is solved. seem like calling untyped_storage end up generating this detach?

 ```
    if out.is_base_tensor():
                out._inference_mode_base = None
                # This assumes that the FunctionalTensor.elem does not change its storage after this point.
                # Otherwise this would be invalid.
                mode._storage_to_base[out.elem.untyped_storage()] = out
            else:
                out._inference_mode_base = mode._storage_to_base[
                    out.elem.untyped_storage()
                ]
                assert out._inference_mode_base is not None
        return out
  ```

I tried two things that solve the issue:
1. do not compute base tensor in inference mode if we are in export mode (the case for this test, **But** **idk if there could be a more general issue that can be hit even in non export mode and if that just avoids it**) 
https://github.com/pytorch/pytorch/pull/137760

2. add detach to  FunctionalTensor.metadata_fns. (**I have no idea if that is safe**) lol

3. both if 2 is safe? otherwise is it possible to run into this issue in non export mode!?



cc @XilunWu @H-Huang @awgu @kwen2501 @wanchaol @fegin @fduwjj @wz337 @wconstab @d4l3k @c-p-i-o @bdhirsh @ezyang @voznesenskym @penguinwu @EikanWang @jgong5 @Guobing-Chen @XiaobingSuper @zhuhaozhe @blzheng @wenzhe-nrv @jiayisunx @ipiszy @yf225 @chenyang78 @kadeng @muchulee8 @ColinPeppler @amjames @desertfire @chauhang @zou3519