Repository: pytorch
Issue: NotImplementedError Could not run 'c10d::alltoall_' with arguments from the 'Meta' backend. 
Labels: triaged

### 🐛 Describe the bug

I use FakeTensor for Shape_prop_pass and met this error:

Exception has occurred: NotImplementedError
Could not run 'c10d::alltoall_' with arguments from the 'Meta' backend. This could be because the operator doesn't exist for this backend, or was omitted during the selective/custom build process (if using custom build). If you are a Facebook employee using PyTorch on mobile, please visit https://fburl.com/ptmfixes for possible resolutions. 'c10d::alltoall_' is only available for these backends: [CPU, CUDA, BackendSelect, Python, FuncTorchDynamicLayerBackMode, Functionalize, Named, Conjugate, Negative, ZeroTensor, ADInplaceOrView, AutogradOther, AutogradCPU, AutogradCUDA, AutogradXLA, AutogradMPS, AutogradXPU, AutogradHPU, AutogradLazy, AutogradMeta, Tracer, AutocastCPU, AutocastCUDA, FuncTorchBatched, FuncTorchVmapMode, Batched, VmapMode, FuncTorchGradWrapper, PythonTLSSnapshot, FuncTorchDynamicLayerFrontMode, PythonDispatcher].
CPU: registered at ../torch/csrc/distributed/c10d/Ops.cpp:700 [kernel]
CUDA: registered at ../torch/csrc/distributed/c10d/Ops.cpp:704 [kernel]
BackendSelect: fallthrough registered at ../aten/src/ATen/core/BackendSelectFallbackKernel.cpp:3 [backend fallback]
Python: registered at ../aten/src/ATen/core/PythonFallbackKernel.cpp:144 [backend fallback]
FuncTorchDynamicLayerBackMode: registered at ../aten/src/ATen/functorch/DynamicLayer.cpp:491 [backend fallback]
Functionalize: registered at ../aten/src/ATen/FunctionalizeFallbackKernel.cpp:280 [backend fallback]
Named: registered at ../aten/src/ATen/core/NamedRegistrations.cpp:7 [backend fallback]
Conjugate: registered at ../aten/src/ATen/ConjugateFallback.cpp:17 [backend fallback]
Negative: registered at ../aten/src/ATen/native/NegateFallback.cpp:19 [backend fallback]
ZeroTensor: registered at ../aten/src/ATen/ZeroTensorFallback.cpp:86 [backend fallback]
ADInplaceOrView: fallthrough registered at ../aten/src/ATen/core/VariableFallbackKernel.cpp:63 [backend fallback]
AutogradOther: fallthrough registered at ../aten/src/ATen/core/VariableFallbackKernel.cpp:30 [backend fallback]
AutogradCPU: fallthrough registered at ../aten/src/ATen/core/VariableFallbackKernel.cpp:34 [backend fallback]
AutogradCUDA: fallthrough registered at ../aten/src/ATen/core/VariableFallbackKernel.cpp:42 [backend fallback]
AutogradXLA: fallthrough registered at ../aten/src/ATen/core/VariableFallbackKernel.cpp:46 [backend fallback]
AutogradMPS: fallthrough registered at ../aten/src/ATen/core/VariableFallbackKernel.cpp:54 [backend fallback]
AutogradXPU: fallthrough registered at ../aten/src/ATen/core/VariableFallbackKernel.cpp:38 [backend fallback]
AutogradHPU: fallthrough registered at ../aten/src/ATen/core/VariableFallbackKernel.cpp:67 [backend fallback]
AutogradLazy: fallthrough registered at ../aten/src/ATen/core/VariableFallbackKernel.cpp:50 [backend fallback]
AutogradMeta: fallthrough registered at ../aten/src/ATen/core/VariableFallbackKernel.cpp:58 [backend fallback]
Tracer: registered at ../torch/csrc/autograd/TraceTypeManual.cpp:294 [backend fallback]
AutocastCPU: fallthrough registered at ../aten/src/ATen/autocast_mode.cpp:487 [backend fallback]
AutocastCUDA: fallthrough registered at ../aten/src/ATen/autocast_mode.cpp:354 [backend fallback]
FuncTorchBatched: registered at ../aten/src/ATen/functorch/LegacyBatchingRegistrations.cpp:815 [backend fallback]
FuncTorchVmapMode: fallthrough registered at ../aten/src/ATen/functorch/VmapModeRegistrations.cpp:28 [backend fallback]
Batched: registered at ../aten/src/ATen/LegacyBatchingRegistrations.cpp:1073 [backend fallback]
VmapMode: fallthrough registered at ../aten/src/ATen/VmapModeRegistrations.cpp:33 [backend fallback]
FuncTorchGradWrapper: registered at ../aten/src/ATen/functorch/TensorWrapper.cpp:210 [backend fallback]
PythonTLSSnapshot: registered at ../aten/src/ATen/core/PythonFallbackKernel.cpp:152 [backend fallback]
FuncTorchDynamicLayerFrontMode: registered at ../aten/src/ATen/functorch/DynamicLayer.cpp:487 [backend fallback]
PythonDispatcher: registered at ../aten/src/ATen/core/PythonFallbackKernel.cpp:148 [backend fallback]


While executing %runtime_apply_1 : [#users=1] = call_function[target=colossalai.auto_parallel.passes.runtime_apply_pass.runtime_apply](args = (%transformer_wte, %origin_node_sharding_spec_dict, %sharding_spec_convert_dict, 11, 0), kwargs = {})
Original traceback:
None
  File "/usr/local/lib/python3.9/site-packages/torch/_ops.py", line 287, in __call__
    return self._op(*args, **kwargs or {})
  File "/usr/local/lib/python3.9/site-packages/torch/_subclasses/fake_tensor.py", line 1175, in dispatch
    raise not_implemented_error
  File "/usr/local/lib/python3.9/site-packages/torch/_subclasses/fake_tensor.py", line 1175, in dispatch
    raise not_implemented_error
  File "/usr/local/lib/python3.9/site-packages/torch/_subclasses/fake_tensor.py", line 988, in __torch_dispatch__
    return self.dispatch(func, types, args, kwargs)
  File "/usr/local/lib/python3.9/site-packages/torch/utils/_stats.py", line 20, in wrapper
    return fn(*args, **kwargs)
  File "/usr/local/lib/python3.9/site-packages/torch/distributed/distributed_c10d.py", line 3281, in all_to_all
    work = group.alltoall(output_tensor_list, input_tensor_list, opts)
  File "/usr/local/lib/python3.9/site-packages/torch/distributed/distributed_c10d.py", line 1451, in wrapper
    return func(*args, **kwargs)
  File "/usr/local/lib/python3.9/site-packages/colossalai/tensor/comm_spec.py", line 66, in _all_to_all
    dist.all_to_all(output_tensor_list, input_tensor_list, group)
  File "/usr/local/lib/python3.9/site-packages/colossalai/tensor/comm_spec.py", line 321, in forward
    output = _all_to_all(input_, comm_spec)
  File "/usr/local/lib/python3.9/site-packages/torch/autograd/function.py", line 506, in apply
    return super().apply(*args, **kwargs)  # type: ignore[misc]
  File "/usr/local/lib/python3.9/site-packages/colossalai/tensor/comm_spec.py", line 368, in all_to_all
    return _AllToAll.apply(input_, comm_spec)
  File "/usr/local/lib/python3.9/site-packages/colossalai/tensor/comm_spec.py", line 514, in covert_spec_to_action
    tensor = pattern_to_func_dict[self.comm_pattern](tensor, self)
  File "/usr/local/lib/python3.9/site-packages/colossalai/tensor/shape_consistency.py", line 742, in apply_for_autoparallel_runtime
    tensor = comm_spec.covert_spec_to_action(tensor)
  File "/usr/local/lib/python3.9/site-packages/colossalai/auto_parallel/passes/runtime_apply_pass.py", line 30, in runtime_apply
    return shape_consistency_manager.apply_for_autoparallel_runtime(node, origin_sharding_spec, target_sharding_spec)
  File "/usr/local/lib/python3.9/site-packages/torch/fx/interpreter.py", line 252, in call_function
    return target(*args, **kwargs)
  File "/usr/local/lib/python3.9/site-packages/torch/fx/interpreter.py", line 180, in run_node
    return getattr(self, n.op)(n.target, args, kwargs)
  File "/usr/local/lib/python3.9/site-packages/torch/fx/passes/fake_tensor_prop.py", line 31, in run_node
    result = super().run_node(n)
  File "/usr/local/lib/python3.9/site-packages/torch/fx/interpreter.py", line 139, in run
    self.env[node] = self.run_node(node)
  File "/usr/local/lib/python3.9/site-packages/torch/fx/passes/fake_tensor_prop.py", line 38, in propagate
    return super().run(*fake_args)
  File "/usr/local/lib/python3.9/site-packages/colossalai/_analyzer/fx/passes/shape_prop.py", line 287, in shape_prop_pass
    FakeTensorProp(module, mode=fake_mode).propagate(*args)
  File "/workspace/workfile/nanoGPT_colossalai/nanoGPT/ColossalAI/colossalai/auto_parallel/tensor_shard/initialize.py", line 150, in transform_to_sharded_model
    shape_prop_pass(gm, *meta_args.values(), sharding_spec_dict, origin_spec_dict, comm_actions_dict)
  File "/workspace/workfile/nanoGPT_colossalai/nanoGPT/ColossalAI/colossalai/auto_parallel/tensor_shard/initialize.py", line 274, in initialize_model
    gm, sharding_spec_dicts = transform_to_sharded_model(gm, meta_args, solution, device_mesh, strategies_constructor,
  File "/workspace/workfile/nanoGPT_colossalai/nanoGPT/ColossalAI/colossalai/auto_parallel/tensor_shard/initialize.py", line 342, in autoparallelize
    rst_to_unpack = initialize_model(model,
  File "/workspace/workfile/nanoGPT_colossalai/nanoGPT/train.py", line 299, in train
    gm, solution = autoparallelize(model, meta_input_sample, return_solution=True)
  File "/workspace/workfile/nanoGPT_colossalai/nanoGPT/train.py", line 405, in <module>
    train()

### Versions

pytorch 2.0