[DDP PT2] TypeError: convert_frame_assert.<locals>._convert_frame_assert() missing 2 required positional arguments: 'hooks' and 'frame_state'
Labels: triaged, oncall: pt2, module: dynamo, mlperf

### 🐛 Describe the bug

Steps to reproduce:

pytorch version: 796ce672296c9ae8d7387b18403810aa2f1048a1
torchrec version: 005727ef06cdc808aa7d263ab4f2837938a77ed2

1. Get working torchrec install (you need fbgemm-gpu and torchrec, install by source; explicitly uninstall fbgemm-gpu-nightly and torchrec-nightly if they exist)
2. Patch torchrec with

```
diff --git a/examples/golden_training/train_dlrm.py b/examples/golden_training/train_dlrm.py
index 4004fbc..39fb5e4 100644
--- a/examples/golden_training/train_dlrm.py
+++ b/examples/golden_training/train_dlrm.py
@@ -127,6 +127,8 @@ def train(
     )
     sharder = EmbeddingBagCollectionSharder(qcomm_codecs_registry=qcomm_codecs_registry)
 
+    train_model.forward = torch.compile(fullgraph=True)(train_model.forward)
+
     model = DistributedModelParallel(
         module=train_model,
         device=device,
```

In torchrec/examples/golden_training/ run `TORCH_LOGS=dynamo MASTER_ADDR=127.0.0.1 MASTER_PORT=29501 RANK=0 LOCAL_RANK=0 WORLD_SIZE=1 python train_dlrm.py`

Backtrace

```
  0%|                                                                                                                                                                                 | 0/1000 [00:00<?, ?it/s]
Traceback (most recent call last):
  File "/data/users/ezyang/a/torchrec/examples/golden_training/train_dlrm.py", line 163, in <module>
    main()
  File "/data/users/ezyang/a/pytorch/torch/distributed/elastic/multiprocessing/errors/__init__.py", line 346, in wrapper
    return f(*args, **kwargs)
  File "/data/users/ezyang/a/torchrec/examples/golden_training/train_dlrm.py", line 52, in main
    train()
  File "/data/users/ezyang/a/torchrec/examples/golden_training/train_dlrm.py", line 159, in train
    train_pipeline.progress(train_iterator)
  File "/data/users/ezyang/a/torchrec/torchrec/distributed/train_pipeline.py", line 1002, in progress
    losses, output = cast(Tuple[torch.Tensor, Out], self._model(self._batch_i))
  File "/data/users/ezyang/a/pytorch/torch/nn/modules/module.py", line 1518, in _wrapped_call_impl
    return self._call_impl(*args, **kwargs)
  File "/data/users/ezyang/a/pytorch/torch/nn/modules/module.py", line 1527, in _call_impl
    return forward_call(*args, **kwargs)
  File "/data/users/ezyang/a/torchrec/torchrec/distributed/model_parallel.py", line 266, in forward
    return self._dmp_wrapped_module(*args, **kwargs)
  File "/data/users/ezyang/a/pytorch/torch/nn/modules/module.py", line 1518, in _wrapped_call_impl
    return self._call_impl(*args, **kwargs)
  File "/data/users/ezyang/a/pytorch/torch/nn/modules/module.py", line 1527, in _call_impl
    return forward_call(*args, **kwargs)
  File "/data/users/ezyang/a/pytorch/torch/nn/parallel/distributed.py", line 1519, in forward
    else self._run_ddp_forward(*inputs, **kwargs)
  File "/data/users/ezyang/a/pytorch/torch/nn/parallel/distributed.py", line 1355, in _run_ddp_forward
    return self.module(*inputs, **kwargs)  # type: ignore[index]
  File "/data/users/ezyang/a/pytorch/torch/nn/modules/module.py", line 1518, in _wrapped_call_impl
    return self._call_impl(*args, **kwargs)
  File "/data/users/ezyang/a/pytorch/torch/nn/modules/module.py", line 1527, in _call_impl
    return forward_call(*args, **kwargs)
  File "/data/users/ezyang/a/pytorch/torch/_dynamo/eval_frame.py", line 333, in _fn
    return fn(*args, **kwargs)
  File "/data/users/ezyang/a/torchrec/torchrec/models/dlrm.py", line 893, in forward
    logits = self.model(batch.dense_features, batch.sparse_features)
  File "/data/users/ezyang/a/pytorch/torch/nn/modules/module.py", line 1518, in _wrapped_call_impl
    return self._call_impl(*args, **kwargs)
  File "/data/users/ezyang/a/pytorch/torch/nn/modules/module.py", line 1527, in _call_impl
    return forward_call(*args, **kwargs)
  File "/data/users/ezyang/a/pytorch/torch/_dynamo/eval_frame.py", line 490, in catch_errors
    return hijacked_callback(frame, cache_size, hooks, frame_state)
  File "/data/users/ezyang/a/pytorch/torch/_dynamo/convert_frame.py", line 626, in _convert_frame
    result = inner_convert(frame, cache_size, hooks, frame_state)
  File "/data/users/ezyang/a/pytorch/torch/_dynamo/convert_frame.py", line 133, in _fn
    return fn(*args, **kwargs)
  File "/data/users/ezyang/a/pytorch/torch/_dynamo/convert_frame.py", line 379, in _convert_frame_assert
    return _compile(
  File "/data/users/ezyang/a/pytorch/torch/_dynamo/convert_frame.py", line 554, in _compile
    guarded_code = compile_inner(code, one_graph, hooks, transform)
  File "/data/users/ezyang/a/pytorch/torch/_dynamo/utils.py", line 181, in time_wrapper
    r = func(*args, **kwargs)
  File "/data/users/ezyang/a/pytorch/torch/_dynamo/convert_frame.py", line 476, in compile_inner
    out_code = transform_code_object(code, transform)
  File "/data/users/ezyang/a/pytorch/torch/_dynamo/bytecode_transformation.py", line 1028, in transform_code_object
    transformations(instructions, code_options)
  File "/data/users/ezyang/a/pytorch/torch/_dynamo/convert_frame.py", line 443, in transform
    tracer.run()
  File "/data/users/ezyang/a/pytorch/torch/_dynamo/symbolic_convert.py", line 2074, in run
    super().run()
  File "/data/users/ezyang/a/pytorch/torch/_dynamo/symbolic_convert.py", line 724, in run
    and self.step()
  File "/data/users/ezyang/a/pytorch/torch/_dynamo/symbolic_convert.py", line 688, in step
    getattr(self, inst.opname)(inst)
  File "/data/users/ezyang/a/pytorch/torch/_dynamo/symbolic_convert.py", line 439, in wrapper
    self.output.compile_subgraph(self, reason=reason)
  File "/data/users/ezyang/a/pytorch/torch/_dynamo/output_graph.py", line 857, in compile_subgraph
    self.compile_and_call_fx_graph(tx, pass2.graph_output_vars(), root)
  File "/home/ezyang/local/a/pytorch-env/lib/python3.10/contextlib.py", line 79, in inner
    return func(*args, **kwds)
  File "/data/users/ezyang/a/pytorch/torch/_dynamo/output_graph.py", line 957, in compile_and_call_fx_graph
    compiled_fn = self.call_user_compiler(gm)
  File "/data/users/ezyang/a/pytorch/torch/_dynamo/utils.py", line 181, in time_wrapper
    r = func(*args, **kwargs)
  File "/data/users/ezyang/a/pytorch/torch/_dynamo/output_graph.py", line 1024, in call_user_compiler
    raise BackendCompilerFailed(self.compiler_fn, e).with_traceback(
  File "/data/users/ezyang/a/pytorch/torch/_dynamo/output_graph.py", line 1009, in call_user_compiler
    compiled_fn = compiler_fn(gm, self.example_inputs())
  File "/data/users/ezyang/a/pytorch/torch/_dynamo/backends/distributed.py", line 291, in compile_fn
    return self.backend_compile_fn(gm, example_inputs)
torch._dynamo.exc.BackendCompilerFailed: backend='compile_fn' raised:
TypeError: convert_frame_assert.<locals>._convert_frame_assert() missing 2 required positional arguments: 'hooks' and 'frame_state'

Set TORCH_LOGS="+dynamo" and TORCHDYNAMO_VERBOSE=1 for more information
```

I spent some time shuffling around looking for the call to convert_frame_assert that was missing these extra arguments but I couldn't find it. Very mysterious.

Note that this model is using DistributedDataParallel.

Full dynamo debug log: https://gist.github.com/ezyang/e4bf7345326092138969387ba364f3ea

cc @msaroufim @wconstab @bdhirsh @anijain2305 @voznesenskym @penguinwu @EikanWang @jgong5 @Guobing-Chen @XiaobingSuper @zhuhaozhe @blzheng @Xia-Weiwen @wenzhe-nrv @jiayisunx @chenyang78 @aakhundov 

### Versions

main