Repository: pytorch
Issue: torch.compile() fails when an `autograd.Function` gets called and torch.no_grad() is *not* being used
Labels: oncall: distributed,triaged,ezyang's list,oncall: pt2,module: dynamo

### 🐛 Describe the bug

I've been implementing TensorParallel with functional collectives in order to get good performance with torch.compile() on a model at work. My implementation of TensorParallel is based on Megatron-LM/fairscale, using `autograd.Function` as the basis to do the minimum amount of `allReduce` calls both on `forward()` and `backward()`. My TP calls looks like this:

```
import torch
import torch.distributed as dist
import torch.distributed._functional_collectives as distfunc

class _CopyToModelParallelRegion(torch.autograd.Function):
    """Pass the input to the model parallel region."""

    @staticmethod
    def symbolic(graph, input_):
        return input_

    @staticmethod
    def forward(ctx, input_):
        return input_

    @staticmethod
    def backward(ctx, grad_output):
        return _all_reduce(grad_output)

class _ReduceFromModelParallelRegion(torch.autograd.Function):
    """All-reduce the input from the model parallel region."""

    @staticmethod
    def symbolic(graph, input_):
        return _all_reduce(input_)

    @staticmethod
    def forward(ctx, input_):
        return _all_reduce(input_)

    @staticmethod
    def backward(ctx, grad_output):
        return grad_output

def _all_reduce(input_: torch.Tensor) -> torch.Tensor:
    """All-reduce the input tensor across model parallel group."""
    world_size = torch.distributed.get_world_size()

    if world_size == 1:
        return input_

    return distfunc.all_reduce(input_, "sum", list(range(world_size)))

def copy_to_tensor_model_parallel_region(input_):
    return _CopyToModelParallelRegion.apply(input_)


def reduce_from_tensor_model_parallel_region(input_):
    return _ReduceFromModelParallelRegion.apply(input_)
```

When using these in a module that has support for TP, if I try to torch.compile() the module using  `with torch.no_grad()`, it works and there are no issues, achieving great performance. Now, if I don't put the `with torch.no_grad()`, the results are a Dynamo error with the following stack trace (repro at the bottom):

```
(/dccstor/aviros_pytorch/conda_envs/pytorch-dev) [avirosmartin@cccxl015 pytorch]$ torchrun --standalone --nnodes=1 --nproc-per-node=2 repro_compile_autograf.py
[2023-08-23 18:04:25,043] torch.distributed.run: [WARNING] master_addr is only used for static rdzv_backend and when rdzv_endpoint is not specified.
[2023-08-23 18:04:25,044] torch.distributed.run: [WARNING] 
[2023-08-23 18:04:25,044] torch.distributed.run: [WARNING] *****************************************
[2023-08-23 18:04:25,044] torch.distributed.run: [WARNING] Setting OMP_NUM_THREADS environment variable for each process to be 1 in default, to avoid your system being overloaded, please further tune the variable for optimal performance in your application as needed. 
[2023-08-23 18:04:25,044] torch.distributed.run: [WARNING] *****************************************
No CUDA runtime is found, using CUDA_HOME='/usr/local/cuda'
No CUDA runtime is found, using CUDA_HOME='/usr/local/cuda'
tensor([[-0.2395,  1.7112, -0.2051,  ..., -3.9962, -2.2627,  1.5794],
        [ 0.6823,  0.9069, -1.0604,  ..., -5.6853, -0.8283,  1.1720],
        [-0.1123,  0.2991, -0.7974,  ..., -4.0582, -1.8296,  1.6378],
        ...,
        [-1.5834,  1.2646, -0.9437,  ..., -2.6262,  1.2763, -2.2860],
        [-1.4105, -2.2290, -2.1080,  ...,  0.4712, -2.1415, -3.8139],
        [ 0.2763,  1.1615,  2.2424,  ..., -2.4191, -2.1367,  1.7405]])tensor([[-0.2395,  1.7112, -0.2051,  ..., -3.9962, -2.2627,  1.5794],
        [ 0.6823,  0.9069, -1.0604,  ..., -5.6853, -0.8283,  1.1720],
        [-0.1123,  0.2991, -0.7974,  ..., -4.0582, -1.8296,  1.6378],
        ...,
        [-1.5834,  1.2646, -0.9437,  ..., -2.6262,  1.2763, -2.2860],
        [-1.4105, -2.2290, -2.1080,  ...,  0.4712, -2.1415, -3.8139],
        [ 0.2763,  1.1615,  2.2424,  ..., -2.4191, -2.1367,  1.7405]])

Traceback (most recent call last):
  File "/dccstor/aviros_pytorch/pytorch/repro_compile_autograf.py", line 114, in <module>
    repro()
  File "/dccstor/aviros_pytorch/pytorch/repro_compile_autograf.py", line 104, in repro
    out = opt_model(input)
  File "/dccstor/aviros_pytorch/pytorch/torch/nn/modules/module.py", line 1518, in _wrapped_call_impl
    return self._call_impl(*args, **kwargs)
  File "/dccstor/aviros_pytorch/pytorch/torch/nn/modules/module.py", line 1527, in _call_impl
    return forward_call(*args, **kwargs)
  File "/dccstor/aviros_pytorch/pytorch/torch/_dynamo/eval_frame.py", line 333, in _fn
    return fn(*args, **kwargs)
  File "/dccstor/aviros_pytorch/pytorch/torch/nn/modules/module.py", line 1518, in _wrapped_call_impl
    return self._call_impl(*args, **kwargs)
  File "/dccstor/aviros_pytorch/pytorch/torch/nn/modules/module.py", line 1527, in _call_impl
    return forward_call(*args, **kwargs)
  File "/dccstor/aviros_pytorch/pytorch/torch/_dynamo/eval_frame.py", line 493, in catch_errors
    return callback(frame, cache_entry, hooks, frame_state)
  File "/dccstor/aviros_pytorch/pytorch/torch/_dynamo/convert_frame.py", line 636, in _convert_frame
    result = inner_convert(frame, cache_size, hooks, frame_state)
  File "/dccstor/aviros_pytorch/pytorch/torch/_dynamo/convert_frame.py", line 133, in _fn
    return fn(*args, **kwargs)
  File "/dccstor/aviros_pytorch/pytorch/torch/_dynamo/convert_frame.py", line 389, in _convert_frame_assert
    return _compile(
  File "/dccstor/aviros_pytorch/pytorch/torch/_dynamo/convert_frame.py", line 581, in _compile
    raise InternalTorchDynamoError(str(e)).with_traceback(
  File "/dccstor/aviros_pytorch/pytorch/torch/_dynamo/convert_frame.py", line 564, in _compile
    guarded_code = compile_inner(code, one_graph, hooks, transform)
  File "/dccstor/aviros_pytorch/pytorch/torch/_dynamo/utils.py", line 189, in time_wrapper
    r = func(*args, **kwargs)
  File "/dccstor/aviros_pytorch/pytorch/torch/_dynamo/convert_frame.py", line 486, in compile_inner
Traceback (most recent call last):
  File "/dccstor/aviros_pytorch/pytorch/repro_compile_autograf.py", line 114, in <module>
    out_code = transform_code_object(code, transform)
  File "/dccstor/aviros_pytorch/pytorch/torch/_dynamo/bytecode_transformation.py", line 1028, in transform_code_object
    repro()
  File "/dccstor/aviros_pytorch/pytorch/repro_compile_autograf.py", line 104, in repro
    transformations(instructions, code_options)
  File "/dccstor/aviros_pytorch/pytorch/torch/_dynamo/convert_frame.py", line 453, in transform
    out = opt_model(input)
  File "/dccstor/aviros_pytorch/pytorch/torch/nn/modules/module.py", line 1518, in _wrapped_call_impl
    tracer.run()
  File "/dccstor/aviros_pytorch/pytorch/torch/_dynamo/symbolic_convert.py", line 2074, in run
    return self._call_impl(*args, **kwargs)
  File "/dccstor/aviros_pytorch/pytorch/torch/nn/modules/module.py", line 1527, in _call_impl
    super().run()
  File "/dccstor/aviros_pytorch/pytorch/torch/_dynamo/symbolic_convert.py", line 724, in run
    return forward_call(*args, **kwargs)
  File "/dccstor/aviros_pytorch/pytorch/torch/_dynamo/eval_frame.py", line 333, in _fn
    and self.step()
  File "/dccstor/aviros_pytorch/pytorch/torch/_dynamo/symbolic_convert.py", line 688, in step
    return fn(*args, **kwargs)
  File "/dccstor/aviros_pytorch/pytorch/torch/nn/modules/module.py", line 1518, in _wrapped_call_impl
    getattr(self, inst.opname)(inst)
  File "/dccstor/aviros_pytorch/pytorch/torch/_dynamo/symbolic_convert.py", line 392, in wrapper
    return inner_fn(self, inst)
  File "/dccstor/aviros_pytorch/pytorch/torch/_dynamo/symbolic_convert.py", line 1115, in CALL_FUNCTION
    return self._call_impl(*args, **kwargs)
  File "/dccstor/aviros_pytorch/pytorch/torch/nn/modules/module.py", line 1527, in _call_impl
    self.call_function(fn, args, {})
  File "/dccstor/aviros_pytorch/pytorch/torch/_dynamo/symbolic_convert.py", line 562, in call_function
    return forward_call(*args, **kwargs)
      File "/dccstor/aviros_pytorch/pytorch/torch/_dynamo/eval_frame.py", line 493, in catch_errors
self.push(fn.call_function(self, args, kwargs))
  File "/dccstor/aviros_pytorch/pytorch/torch/_dynamo/variables/functions.py", line 261, in call_function
    return super().call_function(tx, args, kwargs)
  File "/dccstor/aviros_pytorch/pytorch/torch/_dynamo/variables/functions.py", line 90, in call_function
    return callback(frame, cache_entry, hooks, frame_state)
  File "/dccstor/aviros_pytorch/pytorch/torch/_dynamo/convert_frame.py", line 636, in _convert_frame
    return tx.inline_user_function_return(
  File "/dccstor/aviros_pytorch/pytorch/torch/_dynamo/symbolic_convert.py", line 598, in inline_user_function_return
    result = inner_convert(frame, cache_size, hooks, frame_state)
  File "/dccstor/aviros_pytorch/pytorch/torch/_dynamo/convert_frame.py", line 133, in _fn
    result = InliningInstructionTranslator.inline_call(self, fn, args, kwargs)
  File "/dccstor/aviros_pytorch/pytorch/torch/_dynamo/symbolic_convert.py", line 2179, in inline_call
    return fn(*args, **kwargs)
  File "/dccstor/aviros_pytorch/pytorch/torch/_dynamo/convert_frame.py", line 389, in _convert_frame_assert
    return _compile(
  File "/dccstor/aviros_pytorch/pytorch/torch/_dynamo/convert_frame.py", line 581, in _compile
    raise InternalTorchDynamoError(str(e)).with_traceback(
  File "/dccstor/aviros_pytorch/pytorch/torch/_dynamo/convert_frame.py", line 564, in _compile
    return cls.inline_call_(parent, func, args, kwargs)
  File "/dccstor/aviros_pytorch/pytorch/torch/_dynamo/symbolic_convert.py", line 2286, in inline_call_
    guarded_code = compile_inner(code, one_graph, hooks, transform)
  File "/dccstor/aviros_pytorch/pytorch/torch/_dynamo/utils.py", line 189, in time_wrapper
    r = func(*args, **kwargs)
  File "/dccstor/aviros_pytorch/pytorch/torch/_dynamo/convert_frame.py", line 486, in compile_inner
    tracer.run()
  File "/dccstor/aviros_pytorch/pytorch/torch/_dynamo/symbolic_convert.py", line 724, in run
    out_code = transform_code_object(code, transform)
  File "/dccstor/aviros_pytorch/pytorch/torch/_dynamo/bytecode_transformation.py", line 1028, in transform_code_object
    and self.step()
  File "/dccstor/aviros_pytorch/pytorch/torch/_dynamo/symbolic_convert.py", line 688, in step
    transformations(instructions, code_options)
  File "/dccstor/aviros_pytorch/pytorch/torch/_dynamo/convert_frame.py", line 453, in transform
    getattr(self, inst.opname)(inst)
  File "/dccstor/aviros_pytorch/pytorch/torch/_dynamo/symbolic_convert.py", line 392, in wrapper
    tracer.run()
  File "/dccstor/aviros_pytorch/pytorch/torch/_dynamo/symbolic_convert.py", line 2074, in run
    return inner_fn(self, inst)
  File "/dccstor/aviros_pytorch/pytorch/torch/_dynamo/symbolic_convert.py", line 1115, in CALL_FUNCTION
    self.call_function(fn, args, {})
  File "/dccstor/aviros_pytorch/pytorch/torch/_dynamo/symbolic_convert.py", line 562, in call_function
    super().run()
  File "/dccstor/aviros_pytorch/pytorch/torch/_dynamo/symbolic_convert.py", line 724, in run
    self.push(fn.call_function(self, args, kwargs))
  File "/dccstor/aviros_pytorch/pytorch/torch/_dynamo/variables/misc.py", line 583, in call_function
    and self.step()
  File "/dccstor/aviros_pytorch/pytorch/torch/_dynamo/symbolic_convert.py", line 688, in step
    return self.obj.call_apply(tx, args, kwargs).add_options(self)
  File "/dccstor/aviros_pytorch/pytorch/torch/_dynamo/variables/misc.py", line 353, in call_apply
    getattr(self, inst.opname)(inst)
  File "/dccstor/aviros_pytorch/pytorch/torch/_dynamo/symbolic_convert.py", line 392, in wrapper
    ).call_function(tx, args, kwargs)
  File "/dccstor/aviros_pytorch/pytorch/torch/_dynamo/variables/higher_order_ops.py", line 966, in call_function
    return inner_fn(self, inst)
  File "/dccstor/aviros_pytorch/pytorch/torch/_dynamo/symbolic_convert.py", line 1115, in CALL_FUNCTION
    ) = speculate_subgraph(
  File "/dccstor/aviros_pytorch/pytorch/torch/_dynamo/variables/higher_order_ops.py", line 152, in speculate_subgraph
    self.call_function(fn, args, {})
  File "/dccstor/aviros_pytorch/pytorch/torch/_dynamo/symbolic_convert.py", line 562, in call_function
    output = f.call_function(tx, args, sub_kwargs)
  File "/dccstor/aviros_pytorch/pytorch/torch/_dynamo/variables/torch.py", line 727, in call_function
    self.push(fn.call_function(self, args, kwargs))
  File "/dccstor/aviros_pytorch/pytorch/torch/_dynamo/variables/functions.py", line 261, in call_function
    tensor_variable = wrap_fx_proxy(
  File "/dccstor/aviros_pytorch/pytorch/torch/_dynamo/variables/builder.py", line 1187, in wrap_fx_proxy
    return super().call_function(tx, args, kwargs)
  File "/dccstor/aviros_pytorch/pytorch/torch/_dynamo/variables/functions.py", line 90, in call_function
    return tx.inline_user_function_return(
  File "/dccstor/aviros_pytorch/pytorch/torch/_dynamo/symbolic_convert.py", line 598, in inline_user_function_return
    return wrap_fx_proxy_cls(
  File "/dccstor/aviros_pytorch/pytorch/torch/_dynamo/variables/builder.py", line 1317, in wrap_fx_proxy_cls
    result = InliningInstructionTranslator.inline_call(self, fn, args, kwargs)
  File "/dccstor/aviros_pytorch/pytorch/torch/_dynamo/symbolic_convert.py", line 2179, in inline_call
    example_value = _clone_input(example_value)
  File "/dccstor/aviros_pytorch/pytorch/torch/_dynamo/variables/builder.py", line 1268, in _clone_input
    return cls.inline_call_(parent, func, args, kwargs)
  File "/dccstor/aviros_pytorch/pytorch/torch/_dynamo/symbolic_convert.py", line 2286, in inline_call_
    value = clone_input(value)
  File "/dccstor/aviros_pytorch/pytorch/torch/_dynamo/utils.py", line 592, in clone_input
    result.copy_(x.clone())
  File "/dccstor/aviros_pytorch/pytorch/torch/utils/_stats.py", line 20, in wrapper
    return fn(*args, **kwargs)
  File "/dccstor/aviros_pytorch/pytorch/torch/_subclasses/fake_tensor.py", line 1101, in __torch_dispatch__
    tracer.run()
  File "/dccstor/aviros_pytorch/pytorch/torch/_dynamo/symbolic_convert.py", line 724, in run
    return func(*args, **kwargs)
  File "/dccstor/aviros_pytorch/pytorch/torch/_ops.py", line 435, in __call__
    and self.step()
  File "/dccstor/aviros_pytorch/pytorch/torch/_dynamo/symbolic_convert.py", line 688, in step
    return self._op(*args, **kwargs or {})
  File "/dccstor/aviros_pytorch/pytorch/torch/utils/_stats.py", line 20, in wrapper
    getattr(self, inst.opname)(inst)
  File "/dccstor/aviros_pytorch/pytorch/torch/_dynamo/symbolic_convert.py", line 392, in wrapper
    return fn(*args, **kwargs)
  File "/dccstor/aviros_pytorch/pytorch/torch/_subclasses/fake_tensor.py", line 1238, in __torch_dispatch__
    return inner_fn(self, inst)
  File "/dccstor/aviros_pytorch/pytorch/torch/_dynamo/symbolic_convert.py", line 1115, in CALL_FUNCTION
    return self.dispatch(func, types, args, kwargs)
  File "/dccstor/aviros_pytorch/pytorch/torch/_subclasses/fake_tensor.py", line 1364, in dispatch
    self.call_function(fn, args, {})
  File "/dccstor/aviros_pytorch/pytorch/torch/_dynamo/symbolic_convert.py", line 562, in call_function
    self.push(fn.call_function(self, args, kwargs))
  File "/dccstor/aviros_pytorch/pytorch/torch/_dynamo/variables/misc.py", line 583, in call_function
    ) = self.validate_and_convert_non_fake_tensors(func, converter, args, kwargs)
  File "/dccstor/aviros_pytorch/pytorch/torch/_subclasses/fake_tensor.py", line 1584, in validate_and_convert_non_fake_tensors
    return self.obj.call_apply(tx, args, kwargs).add_options(self)
  File "/dccstor/aviros_pytorch/pytorch/torch/_dynamo/variables/misc.py", line 353, in call_apply
    args, kwargs = tree_map_only(
  File "/dccstor/aviros_pytorch/pytorch/torch/utils/_pytree.py", line 397, in tree_map_only
    ).call_function(tx, args, kwargs)
  File "/dccstor/aviros_pytorch/pytorch/torch/_dynamo/variables/higher_order_ops.py", line 966, in call_function
    return tree_map(map_only(ty)(fn), pytree)
  File "/dccstor/aviros_pytorch/pytorch/torch/utils/_pytree.py", line 327, in tree_map
    return tree_unflatten([fn(i) for i in flat_args], spec)
  File "/dccstor/aviros_pytorch/pytorch/torch/utils/_pytree.py", line 327, in <listcomp>
    ) = speculate_subgraph(
  File "/dccstor/aviros_pytorch/pytorch/torch/_dynamo/variables/higher_order_ops.py", line 152, in speculate_subgraph
    return tree_unflatten([fn(i) for i in flat_args], spec)    
output = f.call_function(tx, args, sub_kwargs)
  File "/dccstor/aviros_pytorch/pytorch/torch/utils/_pytree.py", line 378, in inner
  File "/dccstor/aviros_pytorch/pytorch/torch/_dynamo/variables/torch.py", line 727, in call_function
    return f(x)
  File "/dccstor/aviros_pytorch/pytorch/torch/_subclasses/fake_tensor.py", line 1574, in validate
    tensor_variable = wrap_fx_proxy(
  File "/dccstor/aviros_pytorch/pytorch/torch/_dynamo/variables/builder.py", line 1187, in wrap_fx_proxy
    raise Exception(
torch._dynamo.exc.InternalTorchDynamoError: Please convert all Tensors to FakeTensors first or instantiate FakeTensorMode with 'allow_non_fake_inputs'. Found in aten.copy_.default(tensor([...], size=(50, 100)), FakeTensor(..., size=(50, 100)))

from user code:
   File "/dccstor/aviros_pytorch/pytorch/repro_compile_autograf.py", line 95, in forward
    return reduce_from_tensor_model_parallel_region(out_par)
  File "/dccstor/aviros_pytorch/pytorch/repro_compile_autograf.py", line 51, in reduce_from_tensor_model_parallel_region
    return _ReduceFromModelParallelRegion.apply(input_)

Set TORCH_LOGS="+dynamo" and TORCHDYNAMO_VERBOSE=1 for more information


You can suppress this exception and fall back to eager by setting:
    import torch._dynamo
    torch._dynamo.config.suppress_errors = True

    return wrap_fx_proxy_cls(
  File "/dccstor/aviros_pytorch/pytorch/torch/_dynamo/variables/builder.py", line 1317, in wrap_fx_proxy_cls
    example_value = _clone_input(example_value)
  File "/dccstor/aviros_pytorch/pytorch/torch/_dynamo/variables/builder.py", line 1268, in _clone_input
    value = clone_input(value)
  File "/dccstor/aviros_pytorch/pytorch/torch/_dynamo/utils.py", line 592, in clone_input
    result.copy_(x.clone())
  File "/dccstor/aviros_pytorch/pytorch/torch/utils/_stats.py", line 20, in wrapper
    return fn(*args, **kwargs)
  File "/dccstor/aviros_pytorch/pytorch/torch/_subclasses/fake_tensor.py", line 1101, in __torch_dispatch__
    return func(*args, **kwargs)
  File "/dccstor/aviros_pytorch/pytorch/torch/_ops.py", line 435, in __call__
    return self._op(*args, **kwargs or {})
  File "/dccstor/aviros_pytorch/pytorch/torch/utils/_stats.py", line 20, in wrapper
    return fn(*args, **kwargs)
  File "/dccstor/aviros_pytorch/pytorch/torch/_subclasses/fake_tensor.py", line 1238, in __torch_dispatch__
    return self.dispatch(func, types, args, kwargs)
  File "/dccstor/aviros_pytorch/pytorch/torch/_subclasses/fake_tensor.py", line 1364, in dispatch
    ) = self.validate_and_convert_non_fake_tensors(func, converter, args, kwargs)
  File "/dccstor/aviros_pytorch/pytorch/torch/_subclasses/fake_tensor.py", line 1584, in validate_and_convert_non_fake_tensors
    args, kwargs = tree_map_only(
  File "/dccstor/aviros_pytorch/pytorch/torch/utils/_pytree.py", line 397, in tree_map_only
    return tree_map(map_only(ty)(fn), pytree)
  File "/dccstor/aviros_pytorch/pytorch/torch/utils/_pytree.py", line 327, in tree_map
    return tree_unflatten([fn(i) for i in flat_args], spec)
  File "/dccstor/aviros_pytorch/pytorch/torch/utils/_pytree.py", line 327, in <listcomp>
    return tree_unflatten([fn(i) for i in flat_args], spec)
  File "/dccstor/aviros_pytorch/pytorch/torch/utils/_pytree.py", line 378, in inner
    return f(x)
  File "/dccstor/aviros_pytorch/pytorch/torch/_subclasses/fake_tensor.py", line 1574, in validate
    raise Exception(
torch._dynamo.exc.InternalTorchDynamoError: Please convert all Tensors to FakeTensors first or instantiate FakeTensorMode with 'allow_non_fake_inputs'. Found in aten.copy_.default(tensor([...], size=(50, 100)), FakeTensor(..., size=(50, 100)))

from user code:
   File "/dccstor/aviros_pytorch/pytorch/repro_compile_autograf.py", line 95, in forward
    return reduce_from_tensor_model_parallel_region(out_par)
  File "/dccstor/aviros_pytorch/pytorch/repro_compile_autograf.py", line 51, in reduce_from_tensor_model_parallel_region
    return _ReduceFromModelParallelRegion.apply(input_)

Set TORCH_LOGS="+dynamo" and TORCHDYNAMO_VERBOSE=1 for more information


You can suppress this exception and fall back to eager by setting:
    import torch._dynamo
    torch._dynamo.config.suppress_errors = True

[2023-08-23 18:04:50,118] torch.distributed.elastic.multiprocessing.api: [ERROR] failed (exitcode: 1) local_rank: 0 (pid: 3399902) of binary: /dccstor/aviros_pytorch/conda_envs/pytorch-dev/bin/python
Traceback (most recent call last):
  File "/dccstor/aviros_pytorch/conda_envs/pytorch-dev/bin/torchrun", line 33, in <module>
    sys.exit(load_entry_point('torch', 'console_scripts', 'torchrun')())
  File "/dccstor/aviros_pytorch/pytorch/torch/distributed/elastic/multiprocessing/errors/__init__.py", line 346, in wrapper
    return f(*args, **kwargs)
  File "/dccstor/aviros_pytorch/pytorch/torch/distributed/run.py", line 806, in main
    run(args)
  File "/dccstor/aviros_pytorch/pytorch/torch/distributed/run.py", line 797, in run
    elastic_launch(
  File "/dccstor/aviros_pytorch/pytorch/torch/distributed/launcher/api.py", line 134, in __call__
    return launch_agent(self._config, self._entrypoint, list(args))
  File "/dccstor/aviros_pytorch/pytorch/torch/distributed/launcher/api.py", line 264, in launch_agent
    raise ChildFailedError(
torch.distributed.elastic.multiprocessing.errors.ChildFailedError: 
============================================================
repro_compile_autograf.py FAILED
------------------------------------------------------------
Failures:
[1]:
  time      : 2023-08-23_18:04:50
  host      : cccxl015.pok.ibm.com
  rank      : 1 (local_rank: 1)
  exitcode  : 1 (pid: 3399903)
  error_file: <N/A>
  traceback : To enable traceback see: https://pytorch.org/docs/stable/elastic/errors.html
------------------------------------------------------------
Root Cause (first observed failure):
[0]:
  time      : 2023-08-23_18:04:50
  host      : cccxl015.pok.ibm.com
  rank      : 0 (local_rank: 0)
  exitcode  : 1 (pid: 3399902)
  error_file: <N/A>
  traceback : To enable traceback see: https://pytorch.org/docs/stable/elastic/errors.html
============================================================
```

Repro (together with the TP code)
```
class FeedForwardBlock(torch.nn.Module):
    def __init__(
        self,
        emb_dim,
        hidden_grow_factor=4,
    ):
        super(FeedForwardBlock, self).__init__()
        self.hidden_grow_factor = hidden_grow_factor
        self.hidden_dim = int(hidden_grow_factor * emb_dim)
        self.w1 = torch.nn.Linear(emb_dim, self.hidden_dim)
        self.a = torch.nn.SiLU()
        self.w2 = torch.nn.Linear(self.hidden_dim, emb_dim)
        self.reset_params()

    def reset_params(self):
        for layer in ["w1", "w2"]:
            torch.nn.init.trunc_normal_(
                getattr(self, layer).weight, mean=0.0, std=(2**0.5 / self.w1.weight.numel() ** 0.5) ** 0.5
            )

    def forward(self, x):
        out = self.a(self.w1(x))
        return self.w2(out)

class TPFeedForwardBlock(FeedForwardBlock):
    def __init__(
        self,
        emb_dim,
        hidden_grow_factor=4,
        world_size=1,
        rank=0,
    ):
        hidden_dim = int(hidden_grow_factor * emb_dim)
        assert hidden_dim % world_size == 0, "Hidden dim must be divisible by world size"
        super(TPFeedForwardBlock, self).__init__(emb_dim, hidden_grow_factor / world_size)
        self.rank = rank
        self.world_size = world_size

    def forward(self, x):
        x_par = copy_to_tensor_model_parallel_region(x)
        out_par = FeedForwardBlock.forward(self, x_par)
        return reduce_from_tensor_model_parallel_region(out_par)
    
def repro():
    model = TPFeedForwardBlock(100, 4, dist.get_world_size(), dist.get_rank())

    opt_model = torch.compile(model)

    input = torch.randn((50, 100))

    out = opt_model(input)
    print(out)

if __name__ == "__main__":
    dist.init_process_group(backend="gloo")
    # This works
    with torch.no_grad():
        repro()

    # This fails
    repro()
```

### Versions

Environment (this also fails on latest nightlies):
```
Collecting environment information...
PyTorch version: 2.1.0a0+gitc9f947d
Is debug build: False
CUDA used to build PyTorch: 12.0
ROCM used to build PyTorch: N/A

OS: Red Hat Enterprise Linux release 8.8 (Ootpa) (x86_64)
GCC version: (conda-forge gcc 11.3.0-19) 11.3.0
Clang version: 15.0.7 (Red Hat 15.0.7-1.module+el8.8.0+17939+b58878af)
CMake version: version 3.26.3
Libc version: glibc-2.28

Python version: 3.10.11 | packaged by conda-forge | (main, May 10 2023, 18:58:44) [GCC 11.3.0] (64-bit runtime)
Python platform: Linux-4.18.0-477.15.1.el8_8.x86_64-x86_64-with-glibc2.28
Is CUDA available: False
CUDA runtime version: No CUDA
CUDA_MODULE_LOADING set to: N/A
GPU models and configuration: No CUDA
Nvidia driver version: No CUDA
cuDNN version: No CUDA
HIP runtime version: N/A
MIOpen runtime version: N/A
Is XNNPACK available: True

CPU:
Architecture:        x86_64
CPU op-mode(s):      32-bit, 64-bit
Byte Order:          Little Endian
CPU(s):              16
On-line CPU(s) list: 0-15
Thread(s) per core:  1
Core(s) per socket:  8
Socket(s):           2
NUMA node(s):        2
Vendor ID:           GenuineIntel
CPU family:          6
Model:               62
Model name:          Intel(R) Xeon(R) CPU E5-2667 v2 @ 3.30GHz
Stepping:            4
CPU MHz:             4000.000
CPU max MHz:         4000.0000
CPU min MHz:         1200.0000
BogoMIPS:            6583.80
Virtualization:      VT-x
L1d cache:           32K
L1i cache:           32K
L2 cache:            256K
L3 cache:            25600K
NUMA node0 CPU(s):   0-7
NUMA node1 CPU(s):   8-15
Flags:               fpu vme de pse tsc msr pae mce cx8 apic sep mtrr pge mca cmov pat pse36 clflush dts acpi mmx fxsr sse sse2 ss ht tm pbe syscall nx pdpe1gb rdtscp lm constant_tsc arch_perfmon pebs bts rep_good nopl xtopology nonstop_tsc cpuid aperfmperf pni pclmulqdq dtes64 monitor ds_cpl vmx smx est tm2 ssse3 cx16 xtpr pdcm pcid dca sse4_1 sse4_2 x2apic popcnt tsc_deadline_timer aes xsave avx f16c rdrand lahf_lm cpuid_fault epb pti intel_ppin ssbd ibrs ibpb stibp tpr_shadow vnmi flexpriority ept vpid fsgsbase smep erms xsaveopt dtherm ida arat pln pts md_clear flush_l1d

Versions of relevant libraries:
[pip3] flake8==6.0.0
[pip3] flake8-bugbear==23.3.23
[pip3] flake8-comprehensions==3.12.0
[pip3] flake8-executable==2.1.3
[pip3] flake8-logging-format==0.9.0
[pip3] flake8-pyi==23.3.1
[pip3] flake8-simplify==0.19.3
[pip3] mypy==1.4.1
[pip3] mypy-extensions==1.0.0
[pip3] numpy==1.24.3
[pip3] pytorch-triton==2.1.0+440fd1bf20
[pip3] torch==2.1.0a0+gitc9f947d
[conda] magma-cuda121             2.6.1                         1    pytorch
[conda] mkl                       2023.1.0         h84fe81f_48680    conda-forge
[conda] mkl-include               2023.1.0         h84fe81f_48680    conda-forge
[conda] numpy                     1.24.3          py310ha4c1d20_0    conda-forge
[conda] pytorch-triton            2.1.0+440fd1bf20          pypi_0    pypi
[conda] torch                     2.1.0a0+gitc9f947d           dev_0    <develop>
```

cc @mrshenli @pritamdamania87 @zhaojuanmao @satgera @rohan-varma @gqchen @aazzolini @osalpekar @jiayisuse @H-Huang @kwen2501 @awgu @penguinwu @ezyang @msaroufim @wconstab @bdhirsh @anijain2305 @voznesenskym @EikanWang @jgong5 @Guobing-Chen @XiaobingSuper @zhuhaozhe @blzheng @Xia-Weiwen @wenzhe-nrv @jiayisunx @chenyang78 @aakhundov