Repository: pytorch
Issue: Graph break on wait() of Torchrec's Awaitable instance
Labels: triaged,oncall: pt2,module: dynamo

### 🐛 Describe the bug

When we call the wait() on an Awaitable instance, it causes a graph break. But wait operation should already be supported, we need support this case as well.

### Error logs
```
[rank1]:V1009 23:16:08.229000 3761618 /data/users/shuaiyang/fbsource/fbcode/caffe2/torch/_dynamo/symbolic_convert.py:392] [0/0] [__graph_breaks] Graph break in user code at /data/users/shuaiyang/fbsource/buck-out/v2/gen/fbcode/c58c62e0a86c2716/scripts/jchae/torch_playground/__playground__/playground#link-tree/torchrec/distributed/comm_ops.py:1983
[rank1]:V1009 23:16:08.229000 3761618 /data/users/shuaiyang/fbsource/fbcode/caffe2/torch/_dynamo/symbolic_convert.py:392] [0/0] [__graph_breaks] Reason: Unsupported: call_method UserDefinedObjectVariable(instancemethod) __call__ [] {}
[rank1]:V1009 23:16:08.229000 3761618 /data/users/shuaiyang/fbsource/fbcode/caffe2/torch/_dynamo/symbolic_convert.py:392] [0/0] [__graph_breaks] User code traceback:
[rank1]:V1009 23:16:08.229000 3761618 /data/users/shuaiyang/fbsource/fbcode/caffe2/torch/_dynamo/symbolic_convert.py:392] [0/0] [__graph_breaks]   File "/data/users/shuaiyang/fbsource/buck-out/v2/gen/fbcode/c58c62e0a86c2716/scripts/jchae/torch_playground/__playground__/playground#link-tree/scripts/jchae/torch_playground/playground.py", line 111, in wait_fn
[rank1]:V1009 23:16:08.229000 3761618 /data/users/shuaiyang/fbsource/fbcode/caffe2/torch/_dynamo/symbolic_convert.py:392] [0/0] [__graph_breaks]     res = awaitable.wait()
[rank1]:V1009 23:16:08.229000 3761618 /data/users/shuaiyang/fbsource/fbcode/caffe2/torch/_dynamo/symbolic_convert.py:392] [0/0] [__graph_breaks]   File "/data/users/shuaiyang/fbsource/buck-out/v2/gen/fbcode/c58c62e0a86c2716/scripts/jchae/torch_playground/__playground__/playground#link-tree/torchrec/distributed/types.py", line 335, in wait
[rank1]:V1009 23:16:08.229000 3761618 /data/users/shuaiyang/fbsource/fbcode/caffe2/torch/_dynamo/symbolic_convert.py:392] [0/0] [__graph_breaks]     ret: W = self._wait_impl()
[rank1]:V1009 23:16:08.229000 3761618 /data/users/shuaiyang/fbsource/fbcode/caffe2/torch/_dynamo/symbolic_convert.py:392] [0/0] [__graph_breaks]   File "/data/users/shuaiyang/fbsource/buck-out/v2/gen/fbcode/c58c62e0a86c2716/scripts/jchae/torch_playground/__playground__/playground#link-tree/torchrec/distributed/comm_ops.py", line 121, in _wait_impl
[rank1]:V1009 23:16:08.229000 3761618 /data/users/shuaiyang/fbsource/fbcode/caffe2/torch/_dynamo/symbolic_convert.py:392] [0/0] [__graph_breaks]     ret = self.wait_function.apply(self.pg, self, self.dummy_tensor)
[rank1]:V1009 23:16:08.229000 3761618 /data/users/shuaiyang/fbsource/fbcode/caffe2/torch/_dynamo/symbolic_convert.py:392] [0/0] [__graph_breaks]   File "/data/users/shuaiyang/fbsource/buck-out/v2/gen/fbcode/c58c62e0a86c2716/scripts/jchae/torch_playground/__playground__/playground#link-tree/torchrec/distributed/comm_ops.py", line 1983, in forward
[rank1]:V1009 23:16:08.229000 3761618 /data/users/shuaiyang/fbsource/fbcode/caffe2/torch/_dynamo/symbolic_convert.py:392] [0/0] [__graph_breaks]     myreq.req.wait()
[rank1]:V1009 23:16:08.229000 3761618 /data/users/shuaiyang/fbsource/fbcode/caffe2/torch/_dynamo/symbolic_convert.py:392] [0/0] [__graph_breaks] Traceback (most recent call last):
[rank1]:V1009 23:16:08.229000 3761618 /data/users/shuaiyang/fbsource/fbcode/caffe2/torch/_dynamo/symbolic_convert.py:392] [0/0] [__graph_breaks]   File "/data/users/shuaiyang/fbsource/buck-out/v2/gen/fbcode/c58c62e0a86c2716/scripts/jchae/torch_playground/__playground__/playground#link-tree/torch/_dynamo/symbolic_convert.py", line 616, in wrapper
[rank1]:V1009 23:16:08.229000 3761618 /data/users/shuaiyang/fbsource/fbcode/caffe2/torch/_dynamo/symbolic_convert.py:392] [0/0] [__graph_breaks]     return inner_fn(self, inst)
[rank1]:V1009 23:16:08.229000 3761618 /data/users/shuaiyang/fbsource/fbcode/caffe2/torch/_dynamo/symbolic_convert.py:392] [0/0] [__graph_breaks]   File "/data/users/shuaiyang/fbsource/buck-out/v2/gen/fbcode/c58c62e0a86c2716/scripts/jchae/torch_playground/__playground__/playground#link-tree/torch/_dynamo/symbolic_convert.py", line 1610, in CALL_FUNCTION
[rank1]:V1009 23:16:08.229000 3761618 /data/users/shuaiyang/fbsource/fbcode/caffe2/torch/_dynamo/symbolic_convert.py:392] [0/0] [__graph_breaks]     self.call_function(fn, args, {})
[rank1]:V1009 23:16:08.229000 3761618 /data/users/shuaiyang/fbsource/fbcode/caffe2/torch/_dynamo/symbolic_convert.py:392] [0/0] [__graph_breaks]   File "/data/users/shuaiyang/fbsource/buck-out/v2/gen/fbcode/c58c62e0a86c2716/scripts/jchae/torch_playground/__playground__/playground#link-tree/torch/_dynamo/symbolic_convert.py", line 838, in call_function
[rank1]:V1009 23:16:08.229000 3761618 /data/users/shuaiyang/fbsource/fbcode/caffe2/torch/_dynamo/symbolic_convert.py:392] [0/0] [__graph_breaks]     self.push(fn.call_function(self, args, kwargs))  # type: ignore[arg-type]
[rank1]:V1009 23:16:08.229000 3761618 /data/users/shuaiyang/fbsource/fbcode/caffe2/torch/_dynamo/symbolic_convert.py:392] [0/0] [__graph_breaks]   File "/data/users/shuaiyang/fbsource/buck-out/v2/gen/fbcode/c58c62e0a86c2716/scripts/jchae/torch_playground/__playground__/playground#link-tree/torch/_dynamo/variables/functions.py", line 400, in call_function
[rank1]:V1009 23:16:08.229000 3761618 /data/users/shuaiyang/fbsource/fbcode/caffe2/torch/_dynamo/symbolic_convert.py:392] [0/0] [__graph_breaks]     return super().call_function(tx, args, kwargs)
[rank1]:V1009 23:16:08.229000 3761618 /data/users/shuaiyang/fbsource/fbcode/caffe2/torch/_dynamo/symbolic_convert.py:392] [0/0] [__graph_breaks]   File "/data/users/shuaiyang/fbsource/buck-out/v2/gen/fbcode/c58c62e0a86c2716/scripts/jchae/torch_playground/__playground__/playground#link-tree/torch/_dynamo/variables/functions.py", line 339, in call_function
[rank1]:V1009 23:16:08.229000 3761618 /data/users/shuaiyang/fbsource/fbcode/caffe2/torch/_dynamo/symbolic_convert.py:392] [0/0] [__graph_breaks]     return super().call_function(tx, args, kwargs)
[rank1]:V1009 23:16:08.229000 3761618 /data/users/shuaiyang/fbsource/fbcode/caffe2/torch/_dynamo/symbolic_convert.py:392] [0/0] [__graph_breaks]   File "/data/users/shuaiyang/fbsource/buck-out/v2/gen/fbcode/c58c62e0a86c2716/scripts/jchae/torch_playground/__playground__/playground#link-tree/torch/_dynamo/variables/functions.py", line 111, in call_function
[rank1]:V1009 23:16:08.229000 3761618 /data/users/shuaiyang/fbsource/fbcode/caffe2/torch/_dynamo/symbolic_convert.py:392] [0/0] [__graph_breaks]     return tx.inline_user_function_return(self, [*self.self_args(), *args], kwargs)
[rank1]:V1009 23:16:08.229000 3761618 /data/users/shuaiyang/fbsource/fbcode/caffe2/torch/_dynamo/symbolic_convert.py:392] [0/0] [__graph_breaks]   File "/data/users/shuaiyang/fbsource/buck-out/v2/gen/fbcode/c58c62e0a86c2716/scripts/jchae/torch_playground/__playground__/playground#link-tree/torch/_dynamo/symbolic_convert.py", line 844, in inline_user_function_return
[rank1]:V1009 23:16:08.229000 3761618 /data/users/shuaiyang/fbsource/fbcode/caffe2/torch/_dynamo/symbolic_convert.py:392] [0/0] [__graph_breaks]     return InliningInstructionTranslator.inline_call(self, fn, args, kwargs)
[rank1]:V1009 23:16:08.229000 3761618 /data/users/shuaiyang/fbsource/fbcode/caffe2/torch/_dynamo/symbolic_convert.py:392] [0/0] [__graph_breaks]   File "/data/users/shuaiyang/fbsource/buck-out/v2/gen/fbcode/c58c62e0a86c2716/scripts/jchae/torch_playground/__playground__/playground#link-tree/torch/_dynamo/symbolic_convert.py", line 3000, in inline_call
[rank1]:V1009 23:16:08.229000 3761618 /data/users/shuaiyang/fbsource/fbcode/caffe2/torch/_dynamo/symbolic_convert.py:392] [0/0] [__graph_breaks]     return cls.inline_call_(parent, func, args, kwargs)
[rank1]:V1009 23:16:08.229000 3761618 /data/users/shuaiyang/fbsource/fbcode/caffe2/torch/_dynamo/symbolic_convert.py:392] [0/0] [__graph_breaks]   File "/data/users/shuaiyang/fbsource/buck-out/v2/gen/fbcode/c58c62e0a86c2716/scripts/jchae/torch_playground/__playground__/playground#link-tree/torch/_dynamo/symbolic_convert.py", line 3128, in inline_call_
[rank1]:V1009 23:16:08.229000 3761618 /data/users/shuaiyang/fbsource/fbcode/caffe2/torch/_dynamo/symbolic_convert.py:392] [0/0] [__graph_breaks]     tracer.run()
[rank1]:V1009 23:16:08.229000 3761618 /data/users/shuaiyang/fbsource/fbcode/caffe2/torch/_dynamo/symbolic_convert.py:392] [0/0] [__graph_breaks]   File "/data/users/shuaiyang/fbsource/buck-out/v2/gen/fbcode/c58c62e0a86c2716/scripts/jchae/torch_playground/__playground__/playground#link-tree/torch/_dynamo/symbolic_convert.py", line 991, in run
[rank1]:V1009 23:16:08.229000 3761618 /data/users/shuaiyang/fbsource/fbcode/caffe2/torch/_dynamo/symbolic_convert.py:392] [0/0] [__graph_breaks]     while self.step():
[rank1]:V1009 23:16:08.229000 3761618 /data/users/shuaiyang/fbsource/fbcode/caffe2/torch/_dynamo/symbolic_convert.py:392] [0/0] [__graph_breaks]   File "/data/users/shuaiyang/fbsource/buck-out/v2/gen/fbcode/c58c62e0a86c2716/scripts/jchae/torch_playground/__playground__/playground#link-tree/torch/_dynamo/symbolic_convert.py", line 903, in step
[rank1]:V1009 23:16:08.229000 3761618 /data/users/shuaiyang/fbsource/fbcode/caffe2/torch/_dynamo/symbolic_convert.py:392] [0/0] [__graph_breaks]     self.dispatch_table[inst.opcode](self, inst)
[rank1]:V1009 23:16:08.229000 3761618 /data/users/shuaiyang/fbsource/fbcode/caffe2/torch/_dynamo/symbolic_convert.py:392] [0/0] [__graph_breaks]   File "/data/users/shuaiyang/fbsource/buck-out/v2/gen/fbcode/c58c62e0a86c2716/scripts/jchae/torch_playground/__playground__/playground#link-tree/torch/_dynamo/symbolic_convert.py", line 616, in wrapper
[rank1]:V1009 23:16:08.229000 3761618 /data/users/shuaiyang/fbsource/fbcode/caffe2/torch/_dynamo/symbolic_convert.py:392] [0/0] [__graph_breaks]     return inner_fn(self, inst)
[rank1]:V1009 23:16:08.229000 3761618 /data/users/shuaiyang/fbsource/fbcode/caffe2/torch/_dynamo/symbolic_convert.py:392] [0/0] [__graph_breaks]   File "/data/users/shuaiyang/fbsource/buck-out/v2/gen/fbcode/c58c62e0a86c2716/scripts/jchae/torch_playground/__playground__/playground#link-tree/torch/_dynamo/symbolic_convert.py", line 1610, in CALL_FUNCTION
[rank1]:V1009 23:16:08.229000 3761618 /data/users/shuaiyang/fbsource/fbcode/caffe2/torch/_dynamo/symbolic_convert.py:392] [0/0] [__graph_breaks]     self.call_function(fn, args, {})
[rank1]:V1009 23:16:08.229000 3761618 /data/users/shuaiyang/fbsource/fbcode/caffe2/torch/_dynamo/symbolic_convert.py:392] [0/0] [__graph_breaks]   File "/data/users/shuaiyang/fbsource/buck-out/v2/gen/fbcode/c58c62e0a86c2716/scripts/jchae/torch_playground/__playground__/playground#link-tree/torch/_dynamo/symbolic_convert.py", line 838, in call_function
[rank1]:V1009 23:16:08.229000 3761618 /data/users/shuaiyang/fbsource/fbcode/caffe2/torch/_dynamo/symbolic_convert.py:392] [0/0] [__graph_breaks]     self.push(fn.call_function(self, args, kwargs))  # type: ignore[arg-type]
[rank1]:V1009 23:16:08.229000 3761618 /data/users/shuaiyang/fbsource/fbcode/caffe2/torch/_dynamo/symbolic_convert.py:392] [0/0] [__graph_breaks]   File "/data/users/shuaiyang/fbsource/buck-out/v2/gen/fbcode/c58c62e0a86c2716/scripts/jchae/torch_playground/__playground__/playground#link-tree/torch/_dynamo/variables/functions.py", line 400, in call_function
[rank1]:V1009 23:16:08.229000 3761618 /data/users/shuaiyang/fbsource/fbcode/caffe2/torch/_dynamo/symbolic_convert.py:392] [0/0] [__graph_breaks]     return super().call_function(tx, args, kwargs)
[rank1]:V1009 23:16:08.229000 3761618 /data/users/shuaiyang/fbsource/fbcode/caffe2/torch/_dynamo/symbolic_convert.py:392] [0/0] [__graph_breaks]   File "/data/users/shuaiyang/fbsource/buck-out/v2/gen/fbcode/c58c62e0a86c2716/scripts/jchae/torch_playground/__playground__/playground#link-tree/torch/_dynamo/variables/functions.py", line 339, in call_function
[rank1]:V1009 23:16:08.229000 3761618 /data/users/shuaiyang/fbsource/fbcode/caffe2/torch/_dynamo/symbolic_convert.py:392] [0/0] [__graph_breaks]     return super().call_function(tx, args, kwargs)
[rank1]:V1009 23:16:08.229000 3761618 /data/users/shuaiyang/fbsource/fbcode/caffe2/torch/_dynamo/symbolic_convert.py:392] [0/0] [__graph_breaks]   File "/data/users/shuaiyang/fbsource/buck-out/v2/gen/fbcode/c58c62e0a86c2716/scripts/jchae/torch_playground/__playground__/playground#link-tree/torch/_dynamo/variables/functions.py", line 111, in call_function
[rank1]:V1009 23:16:08.229000 3761618 /data/users/shuaiyang/fbsource/fbcode/caffe2/torch/_dynamo/symbolic_convert.py:392] [0/0] [__graph_breaks]     return tx.inline_user_function_return(self, [*self.self_args(), *args], kwargs)
[rank1]:V1009 23:16:08.229000 3761618 /data/users/shuaiyang/fbsource/fbcode/caffe2/torch/_dynamo/symbolic_convert.py:392] [0/0] [__graph_breaks]   File "/data/users/shuaiyang/fbsource/buck-out/v2/gen/fbcode/c58c62e0a86c2716/scripts/jchae/torch_playground/__playground__/playground#link-tree/torch/_dynamo/symbolic_convert.py", line 844, in inline_user_function_return
[rank1]:V1009 23:16:08.229000 3761618 /data/users/shuaiyang/fbsource/fbcode/caffe2/torch/_dynamo/symbolic_convert.py:392] [0/0] [__graph_breaks]     return InliningInstructionTranslator.inline_call(self, fn, args, kwargs)
[rank1]:V1009 23:16:08.229000 3761618 /data/users/shuaiyang/fbsource/fbcode/caffe2/torch/_dynamo/symbolic_convert.py:392] [0/0] [__graph_breaks]   File "/data/users/shuaiyang/fbsource/buck-out/v2/gen/fbcode/c58c62e0a86c2716/scripts/jchae/torch_playground/__playground__/playground#link-tree/torch/_dynamo/symbolic_convert.py", line 3000, in inline_call
[rank1]:V1009 23:16:08.229000 3761618 /data/users/shuaiyang/fbsource/fbcode/caffe2/torch/_dynamo/symbolic_convert.py:392] [0/0] [__graph_breaks]     return cls.inline_call_(parent, func, args, kwargs)
[rank1]:V1009 23:16:08.229000 3761618 /data/users/shuaiyang/fbsource/fbcode/caffe2/torch/_dynamo/symbolic_convert.py:392] [0/0] [__graph_breaks]   File "/data/users/shuaiyang/fbsource/buck-out/v2/gen/fbcode/c58c62e0a86c2716/scripts/jchae/torch_playground/__playground__/playground#link-tree/torch/_dynamo/symbolic_convert.py", line 3128, in inline_call_
[rank1]:V1009 23:16:08.229000 3761618 /data/users/shuaiyang/fbsource/fbcode/caffe2/torch/_dynamo/symbolic_convert.py:392] [0/0] [__graph_breaks]     tracer.run()
[rank1]:V1009 23:16:08.229000 3761618 /data/users/shuaiyang/fbsource/fbcode/caffe2/torch/_dynamo/symbolic_convert.py:392] [0/0] [__graph_breaks]   File "/data/users/shuaiyang/fbsource/buck-out/v2/gen/fbcode/c58c62e0a86c2716/scripts/jchae/torch_playground/__playground__/playground#link-tree/torch/_dynamo/symbolic_convert.py", line 991, in run
[rank1]:V1009 23:16:08.229000 3761618 /data/users/shuaiyang/fbsource/fbcode/caffe2/torch/_dynamo/symbolic_convert.py:392] [0/0] [__graph_breaks]     while self.step():
[rank1]:V1009 23:16:08.229000 3761618 /data/users/shuaiyang/fbsource/fbcode/caffe2/torch/_dynamo/symbolic_convert.py:392] [0/0] [__graph_breaks]   File "/data/users/shuaiyang/fbsource/buck-out/v2/gen/fbcode/c58c62e0a86c2716/scripts/jchae/torch_playground/__playground__/playground#link-tree/torch/_dynamo/symbolic_convert.py", line 903, in step
[rank1]:V1009 23:16:08.229000 3761618 /data/users/shuaiyang/fbsource/fbcode/caffe2/torch/_dynamo/symbolic_convert.py:392] [0/0] [__graph_breaks]     self.dispatch_table[inst.opcode](self, inst)
[rank1]:V1009 23:16:08.229000 3761618 /data/users/shuaiyang/fbsource/fbcode/caffe2/torch/_dynamo/symbolic_convert.py:392] [0/0] [__graph_breaks]   File "/data/users/shuaiyang/fbsource/buck-out/v2/gen/fbcode/c58c62e0a86c2716/scripts/jchae/torch_playground/__playground__/playground#link-tree/torch/_dynamo/symbolic_convert.py", line 616, in wrapper
[rank1]:V1009 23:16:08.229000 3761618 /data/users/shuaiyang/fbsource/fbcode/caffe2/torch/_dynamo/symbolic_convert.py:392] [0/0] [__graph_breaks]     return inner_fn(self, inst)
[rank1]:V1009 23:16:08.229000 3761618 /data/users/shuaiyang/fbsource/fbcode/caffe2/torch/_dynamo/symbolic_convert.py:392] [0/0] [__graph_breaks]   File "/data/users/shuaiyang/fbsource/buck-out/v2/gen/fbcode/c58c62e0a86c2716/scripts/jchae/torch_playground/__playground__/playground#link-tree/torch/_dynamo/symbolic_convert.py", line 1610, in CALL_FUNCTION
[rank1]:V1009 23:16:08.229000 3761618 /data/users/shuaiyang/fbsource/fbcode/caffe2/torch/_dynamo/symbolic_convert.py:392] [0/0] [__graph_breaks]     self.call_function(fn, args, {})
[rank1]:V1009 23:16:08.229000 3761618 /data/users/shuaiyang/fbsource/fbcode/caffe2/torch/_dynamo/symbolic_convert.py:392] [0/0] [__graph_breaks]   File "/data/users/shuaiyang/fbsource/buck-out/v2/gen/fbcode/c58c62e0a86c2716/scripts/jchae/torch_playground/__playground__/playground#link-tree/torch/_dynamo/symbolic_convert.py", line 838, in call_function
[rank1]:V1009 23:16:08.229000 3761618 /data/users/shuaiyang/fbsource/fbcode/caffe2/torch/_dynamo/symbolic_convert.py:392] [0/0] [__graph_breaks]     self.push(fn.call_function(self, args, kwargs))  # type: ignore[arg-type]
[rank1]:V1009 23:16:08.229000 3761618 /data/users/shuaiyang/fbsource/fbcode/caffe2/torch/_dynamo/symbolic_convert.py:392] [0/0] [__graph_breaks]   File "/data/users/shuaiyang/fbsource/buck-out/v2/gen/fbcode/c58c62e0a86c2716/scripts/jchae/torch_playground/__playground__/playground#link-tree/torch/_dynamo/variables/misc.py", line 1038, in call_function
[rank1]:V1009 23:16:08.229000 3761618 /data/users/shuaiyang/fbsource/fbcode/caffe2/torch/_dynamo/symbolic_convert.py:392] [0/0] [__graph_breaks]     return self.obj.call_method(tx, self.name, args, kwargs)
[rank1]:V1009 23:16:08.229000 3761618 /data/users/shuaiyang/fbsource/fbcode/caffe2/torch/_dynamo/symbolic_convert.py:392] [0/0] [__graph_breaks]   File "/data/users/shuaiyang/fbsource/buck-out/v2/gen/fbcode/c58c62e0a86c2716/scripts/jchae/torch_playground/__playground__/playground#link-tree/torch/_dynamo/variables/misc.py", line 781, in call_method
[rank1]:V1009 23:16:08.229000 3761618 /data/users/shuaiyang/fbsource/fbcode/caffe2/torch/_dynamo/symbolic_convert.py:392] [0/0] [__graph_breaks]     return self.call_apply(tx, args, kwargs)
[rank1]:V1009 23:16:08.229000 3761618 /data/users/shuaiyang/fbsource/fbcode/caffe2/torch/_dynamo/symbolic_convert.py:392] [0/0] [__graph_breaks]   File "/data/users/shuaiyang/fbsource/buck-out/v2/gen/fbcode/c58c62e0a86c2716/scripts/jchae/torch_playground/__playground__/playground#link-tree/torch/_dynamo/variables/misc.py", line 730, in call_apply
[rank1]:V1009 23:16:08.229000 3761618 /data/users/shuaiyang/fbsource/fbcode/caffe2/torch/_dynamo/symbolic_convert.py:392] [0/0] [__graph_breaks]     return variables.UserFunctionVariable(fn, source=source).call_function(
[rank1]:V1009 23:16:08.229000 3761618 /data/users/shuaiyang/fbsource/fbcode/caffe2/torch/_dynamo/symbolic_convert.py:392] [0/0] [__graph_breaks]   File "/data/users/shuaiyang/fbsource/buck-out/v2/gen/fbcode/c58c62e0a86c2716/scripts/jchae/torch_playground/__playground__/playground#link-tree/torch/_dynamo/variables/functions.py", line 339, in call_function
[rank1]:V1009 23:16:08.229000 3761618 /data/users/shuaiyang/fbsource/fbcode/caffe2/torch/_dynamo/symbolic_convert.py:392] [0/0] [__graph_breaks]     return super().call_function(tx, args, kwargs)
[rank1]:V1009 23:16:08.229000 3761618 /data/users/shuaiyang/fbsource/fbcode/caffe2/torch/_dynamo/symbolic_convert.py:392] [0/0] [__graph_breaks]   File "/data/users/shuaiyang/fbsource/buck-out/v2/gen/fbcode/c58c62e0a86c2716/scripts/jchae/torch_playground/__playground__/playground#link-tree/torch/_dynamo/variables/functions.py", line 111, in call_function
[rank1]:V1009 23:16:08.229000 3761618 /data/users/shuaiyang/fbsource/fbcode/caffe2/torch/_dynamo/symbolic_convert.py:392] [0/0] [__graph_breaks]     return tx.inline_user_function_return(self, [*self.self_args(), *args], kwargs)
[rank1]:V1009 23:16:08.229000 3761618 /data/users/shuaiyang/fbsource/fbcode/caffe2/torch/_dynamo/symbolic_convert.py:392] [0/0] [__graph_breaks]   File "/data/users/shuaiyang/fbsource/buck-out/v2/gen/fbcode/c58c62e0a86c2716/scripts/jchae/torch_playground/__playground__/playground#link-tree/torch/_dynamo/symbolic_convert.py", line 844, in inline_user_function_return
[rank1]:V1009 23:16:08.229000 3761618 /data/users/shuaiyang/fbsource/fbcode/caffe2/torch/_dynamo/symbolic_convert.py:392] [0/0] [__graph_breaks]     return InliningInstructionTranslator.inline_call(self, fn, args, kwargs)
[rank1]:V1009 23:16:08.229000 3761618 /data/users/shuaiyang/fbsource/fbcode/caffe2/torch/_dynamo/symbolic_convert.py:392] [0/0] [__graph_breaks]   File "/data/users/shuaiyang/fbsource/buck-out/v2/gen/fbcode/c58c62e0a86c2716/scripts/jchae/torch_playground/__playground__/playground#link-tree/torch/_dynamo/symbolic_convert.py", line 3000, in inline_call
[rank1]:V1009 23:16:08.229000 3761618 /data/users/shuaiyang/fbsource/fbcode/caffe2/torch/_dynamo/symbolic_convert.py:392] [0/0] [__graph_breaks]     return cls.inline_call_(parent, func, args, kwargs)
[rank1]:V1009 23:16:08.229000 3761618 /data/users/shuaiyang/fbsource/fbcode/caffe2/torch/_dynamo/symbolic_convert.py:392] [0/0] [__graph_breaks]   File "/data/users/shuaiyang/fbsource/buck-out/v2/gen/fbcode/c58c62e0a86c2716/scripts/jchae/torch_playground/__playground__/playground#link-tree/torch/_dynamo/symbolic_convert.py", line 3128, in inline_call_
[rank1]:V1009 23:16:08.229000 3761618 /data/users/shuaiyang/fbsource/fbcode/caffe2/torch/_dynamo/symbolic_convert.py:392] [0/0] [__graph_breaks]     tracer.run()
[rank1]:V1009 23:16:08.229000 3761618 /data/users/shuaiyang/fbsource/fbcode/caffe2/torch/_dynamo/symbolic_convert.py:392] [0/0] [__graph_breaks]   File "/data/users/shuaiyang/fbsource/buck-out/v2/gen/fbcode/c58c62e0a86c2716/scripts/jchae/torch_playground/__playground__/playground#link-tree/torch/_dynamo/symbolic_convert.py", line 991, in run
[rank1]:V1009 23:16:08.229000 3761618 /data/users/shuaiyang/fbsource/fbcode/caffe2/torch/_dynamo/symbolic_convert.py:392] [0/0] [__graph_breaks]     while self.step():
[rank1]:V1009 23:16:08.229000 3761618 /data/users/shuaiyang/fbsource/fbcode/caffe2/torch/_dynamo/symbolic_convert.py:392] [0/0] [__graph_breaks]   File "/data/users/shuaiyang/fbsource/buck-out/v2/gen/fbcode/c58c62e0a86c2716/scripts/jchae/torch_playground/__playground__/playground#link-tree/torch/_dynamo/symbolic_convert.py", line 903, in step
[rank1]:V1009 23:16:08.229000 3761618 /data/users/shuaiyang/fbsource/fbcode/caffe2/torch/_dynamo/symbolic_convert.py:392] [0/0] [__graph_breaks]     self.dispatch_table[inst.opcode](self, inst)
[rank1]:V1009 23:16:08.229000 3761618 /data/users/shuaiyang/fbsource/fbcode/caffe2/torch/_dynamo/symbolic_convert.py:392] [0/0] [__graph_breaks]   File "/data/users/shuaiyang/fbsource/buck-out/v2/gen/fbcode/c58c62e0a86c2716/scripts/jchae/torch_playground/__playground__/playground#link-tree/torch/_dynamo/symbolic_convert.py", line 616, in wrapper
[rank1]:V1009 23:16:08.229000 3761618 /data/users/shuaiyang/fbsource/fbcode/caffe2/torch/_dynamo/symbolic_convert.py:392] [0/0] [__graph_breaks]     return inner_fn(self, inst)
[rank1]:V1009 23:16:08.229000 3761618 /data/users/shuaiyang/fbsource/fbcode/caffe2/torch/_dynamo/symbolic_convert.py:392] [0/0] [__graph_breaks]   File "/data/users/shuaiyang/fbsource/buck-out/v2/gen/fbcode/c58c62e0a86c2716/scripts/jchae/torch_playground/__playground__/playground#link-tree/torch/_dynamo/symbolic_convert.py", line 1610, in CALL_FUNCTION
[rank1]:V1009 23:16:08.229000 3761618 /data/users/shuaiyang/fbsource/fbcode/caffe2/torch/_dynamo/symbolic_convert.py:392] [0/0] [__graph_breaks]     self.call_function(fn, args, {})
[rank1]:V1009 23:16:08.229000 3761618 /data/users/shuaiyang/fbsource/fbcode/caffe2/torch/_dynamo/symbolic_convert.py:392] [0/0] [__graph_breaks]   File "/data/users/shuaiyang/fbsource/buck-out/v2/gen/fbcode/c58c62e0a86c2716/scripts/jchae/torch_playground/__playground__/playground#link-tree/torch/_dynamo/symbolic_convert.py", line 838, in call_function
[rank1]:V1009 23:16:08.229000 3761618 /data/users/shuaiyang/fbsource/fbcode/caffe2/torch/_dynamo/symbolic_convert.py:392] [0/0] [__graph_breaks]     self.push(fn.call_function(self, args, kwargs))  # type: ignore[arg-type]
[rank1]:V1009 23:16:08.229000 3761618 /data/users/shuaiyang/fbsource/fbcode/caffe2/torch/_dynamo/symbolic_convert.py:392] [0/0] [__graph_breaks]   File "/data/users/shuaiyang/fbsource/buck-out/v2/gen/fbcode/c58c62e0a86c2716/scripts/jchae/torch_playground/__playground__/playground#link-tree/torch/_dynamo/variables/lazy.py", line 161, in realize_and_forward
[rank1]:V1009 23:16:08.229000 3761618 /data/users/shuaiyang/fbsource/fbcode/caffe2/torch/_dynamo/symbolic_convert.py:392] [0/0] [__graph_breaks]     return getattr(self.realize(), name)(*args, **kwargs)
[rank1]:V1009 23:16:08.229000 3761618 /data/users/shuaiyang/fbsource/fbcode/caffe2/torch/_dynamo/symbolic_convert.py:392] [0/0] [__graph_breaks]   File "/data/users/shuaiyang/fbsource/buck-out/v2/gen/fbcode/c58c62e0a86c2716/scripts/jchae/torch_playground/__playground__/playground#link-tree/torch/_dynamo/variables/user_defined.py", line 926, in call_function
[rank1]:V1009 23:16:08.229000 3761618 /data/users/shuaiyang/fbsource/fbcode/caffe2/torch/_dynamo/symbolic_convert.py:392] [0/0] [__graph_breaks]     return self.call_method(tx, "__call__", args, kwargs)
[rank1]:V1009 23:16:08.229000 3761618 /data/users/shuaiyang/fbsource/fbcode/caffe2/torch/_dynamo/symbolic_convert.py:392] [0/0] [__graph_breaks]   File "/data/users/shuaiyang/fbsource/buck-out/v2/gen/fbcode/c58c62e0a86c2716/scripts/jchae/torch_playground/__playground__/playground#link-tree/torch/_dynamo/variables/user_defined.py", line 794, in call_method
[rank1]:V1009 23:16:08.229000 3761618 /data/users/shuaiyang/fbsource/fbcode/caffe2/torch/_dynamo/symbolic_convert.py:392] [0/0] [__graph_breaks]     return super().call_method(tx, name, args, kwargs)
[rank1]:V1009 23:16:08.229000 3761618 /data/users/shuaiyang/fbsource/fbcode/caffe2/torch/_dynamo/symbolic_convert.py:392] [0/0] [__graph_breaks]   File "/data/users/shuaiyang/fbsource/buck-out/v2/gen/fbcode/c58c62e0a86c2716/scripts/jchae/torch_playground/__playground__/playground#link-tree/torch/_dynamo/variables/base.py", line 343, in call_method
[rank1]:V1009 23:16:08.229000 3761618 /data/users/shuaiyang/fbsource/fbcode/caffe2/torch/_dynamo/symbolic_convert.py:392] [0/0] [__graph_breaks]     unimplemented(f"call_method {self} {name} {args} {kwargs}")
[rank1]:V1009 23:16:08.229000 3761618 /data/users/shuaiyang/fbsource/fbcode/caffe2/torch/_dynamo/symbolic_convert.py:392] [0/0] [__graph_breaks]   File "/data/users/shuaiyang/fbsource/buck-out/v2/gen/fbcode/c58c62e0a86c2716/scripts/jchae/torch_playground/__playground__/playground#link-tree/torch/_dynamo/exc.py", line 304, in unimplemented
[rank1]:V1009 23:16:08.229000 3761618 /data/users/shuaiyang/fbsource/fbcode/caffe2/torch/_dynamo/symbolic_convert.py:392] [0/0] [__graph_breaks]     raise Unsupported(msg, case_name=case_name)
[rank1]:V1009 23:16:08.229000 3761618 /data/users/shuaiyang/fbsource/fbcode/caffe2/torch/_dynamo/symbolic_convert.py:392] [0/0] [__graph_breaks] torch._dynamo.exc.Unsupported: call_method UserDefinedObjectVariable(instancemethod) __call__ [] {}
```

### Minified repro

```
import functools

import logging
import multiprocessing
import os
import unittest
from dataclasses import dataclass
from typing import Callable, List, Optional, Union

import torch
import torch.distributed as dist
import torchrec.distributed.comm_ops as comm_ops
from torch import nn
from torch.distributed.distributed_c10d import GroupMember
from torchrec.test_utils import get_free_port

torch._logging.set_logs(
    aot_graphs=True,
    graph_code=True,
    graph_breaks=True,
    trace_source=True,
    recompiles=True,
    inductor=logging.DEBUG,
    dynamo=logging.DEBUG,
)


@dataclass
class _CompileConfig:
    # backend is None means no compilation
    backend: Optional[str] = "inductor"
    fullgraph: bool = True
    skip_sync_backward: bool = False
    skip_compile_backward: bool = False
    test_compiled_with_noncompiled_ranks: bool = False


def compile_config_to_fn_transform(
    compile_config: Optional[_CompileConfig],
    # pyre-ignore
) -> Callable:
    if compile_config is None:
        return lambda x: x

    return functools.partial(
        torch.compile,
        backend=compile_config.backend,
        fullgraph=compile_config.fullgraph,
        dynamic=True,
    )


def _test_alltoallv(
    rank: int,
    world_size: int,
    backend: str,
    compile_config: Optional[_CompileConfig] = None,
    specify_pg: bool = False,
) -> None:
    dist.init_process_group(rank=rank, world_size=world_size, backend=backend)
    pg = GroupMember.WORLD
    assert pg is not None

    device = torch.device(f"cuda:{rank}")

    torch.cuda.set_device(device)

    B_global = 2
    D0 = 8
    D1 = 9

    input_embedding0 = torch.rand(
        (B_global, D0),
        device=device,
        requires_grad=True,
    )
    input_embedding0 = torch.tensor(
        [[1, 2, 3, 4, 5, 6, 7, 8], [9, 10, 11, 12, 13, 14, 15, 16]],
        device=device,
        requires_grad=True,
        dtype=torch.float32,
    )
    input_embedding1 = torch.rand(
        (B_global, D1),
        device=device,
        requires_grad=True,
    )
    input_embedding1 = torch.tensor(
        [[11, 12, 13, 14, 15, 16, 17, 18, 19], [20, 21, 22, 23, 24, 25, 26, 27, 28]],
        device=device,
        requires_grad=True,
        dtype=torch.float32,
    )

    input_embeddings = [input_embedding0, input_embedding1]
    out_split = [17, 17]

    # pyre-ignore
    def fn(*args, **kwargs):
        awaitable = comm_ops.alltoallv(
            *args, **kwargs
        )  # the type of awaitable is torchrec.distributed.comm_ops.Request

        return awaitable

    @torch.compile(backend="inductor")
    def wait_fn(awaitable):
        res = awaitable.wait()
        res1 = res[0] + res[1]
        return res1

    # fn_transform = compile_config_to_fn_transform(compile_config)

    with unittest.mock.patch(
        "torch._dynamo.config.skip_torchrec",
        False,
    ):
        v_embs_out_awaitable = fn(
            input_embeddings, out_split=out_split, group=pg if specify_pg else None
        )
        v_embs_out = wait_fn(v_embs_out_awaitable)
        print(f"bbbbbbbbbbbbbbbbbbbb {(v_embs_out)}")
        torch.sum(v_embs_out)

    # res = torch.cat(v_embs_out, dim=1).cpu()
    # assert tuple(res.size()) == (1, 34)
    dist.destroy_process_group()


def test_alltoallv(
    # specify_pg: bool,
    # test_compiled_with_noncompiled_ranks: bool,
) -> None:
    _run_multi_process_test(
        world_size=2,
        backend="nccl",
        # pyre-ignore [6]
        callable=_test_alltoallv,
        # compile_config=_CompileConfig(
        #     test_compiled_with_noncompiled_ranks=test_compiled_with_noncompiled_ranks
        # ),
        compile_config=None,
        specify_pg=None,
    )


def _run_multi_process_test(
    world_size: int,
    backend: str,
    callable: Callable[[], None],
    # pyre-ignore
    *args,
    # pyre-ignore
    **kwargs,
) -> None:
    processes = []
    ctx = multiprocessing.get_context("spawn")
    for rank in range(world_size):
        p = ctx.Process(
            target=callable,
            args=(
                rank,
                world_size,
                backend,
                *args,
            ),
            kwargs=kwargs,
        )
        p.start()
        processes.append(p)

    for p in processes:
        p.join()
        assert p.exitcode == 0


if __name__ == "__main__":
    os.environ["MASTER_ADDR"] = str("localhost")
    os.environ["MASTER_PORT"] = str(get_free_port())
    os.environ["GLOO_DEVICE_TRANSPORT"] = "TCP"
    os.environ["NCCL_SOCKET_IFNAME"] = "lo"
    test_alltoallv()

```

### Versions

nightly

cc @ezyang @chauhang @penguinwu @voznesenskym @EikanWang @jgong5 @Guobing-Chen @XiaobingSuper @zhuhaozhe @blzheng @wenzhe-nrv @jiayisunx @chenyang78 @kadeng @amjames @rec