Repository: pytorch
Issue: [compile] KeyError: example_value
Labels: needs reproduction,triaged,oncall: pt2

### 🐛 Describe the bug

`model.compile()` (defaults) gives me:

```
stderr:  - 5597 - torch._dynamo.symbolic_convert - Step 1: torchdynamo start tracing forward
stderr:  - 5597 - torch._dynamo.symbolic_convert - Step 1: torchdynamo start tracing wrapped_fn
stderr:  - 5597 - torch._dynamo.symbolic_convert - Step 1: torchdynamo start tracing range_push
stderr:  - 5597 - torch._dynamo.symbolic_convert - Step 1: torchdynamo start tracing <graph break in range_push>
stderr:  - 5597 - torch._dynamo.symbolic_convert - Step 1: torchdynamo start tracing <graph break in wrapped_fn>
stderr:  - 5597 - torch._dynamo.symbolic_convert - Step 1: torchdynamo start tracing _pre_forward_module_hook
stderr:  - 5597 - torch._dynamo.symbolic_convert - Step 1: torchdynamo start tracing pre_sub_module_forward_function
stderr:  - 5597 - torch._dynamo.symbolic_convert - Step 1: torchdynamo start tracing record_module
stderr:  - 5597 - torch._dynamo.symbolic_convert - Step 1: torchdynamo start tracing <graph break in pre_sub_module_forward_function>
stderr:  - 5597 - torch._dynamo.symbolic_convert - Step 1: torchdynamo start tracing wrapped_fn
stderr:  - 5597 - torch._dynamo.symbolic_convert - Step 1: torchdynamo start tracing <graph break in range_push>
stderr:  - 5597 - torch._dynamo.symbolic_convert - Step 1: torchdynamo start tracing <graph break in wrapped_fn>
stderr:  - 5597 - torch._dynamo.symbolic_convert - Step 1: torchdynamo start tracing fetch_sub_module
stderr:  - 5597 - torch._dynamo.symbolic_convert - Step 1: torchdynamo start tracing iter_params
stderr:  - 5597 - torch._dynamo.symbolic_convert - Step 1: torchdynamo start tracing wrapped_fn
stderr:  - 5597 - torch._dynamo.symbolic_convert - Step 1: torchdynamo start tracing <graph break in range_push>
stderr:  - 5597 - torch._dynamo.symbolic_convert - Step 1: torchdynamo start tracing <graph break in wrapped_fn>
stderr:  - 5597 - torch._dynamo.symbolic_convert - Step 1: torchdynamo start tracing get_all_parameters
stderr:  - 5597 - torch._dynamo.symbolic_convert - Step 1: torchdynamo start tracing external_parameters
stderr:  - 5597 - torch._dynamo.symbolic_convert - Step 1: torchdynamo start tracing range_pop
stderr:  - 5597 - torch._dynamo.symbolic_convert - Step 1: torchdynamo start tracing <graph break in range_pop>
stderr:  - 5597 - torch._dynamo.symbolic_convert - Step 1: torchdynamo start tracing <graph break in fetch_sub_module>
stderr:  - 5597 - torch._dynamo.symbolic_convert - Step 1: torchdynamo start tracing is_initialized
stderr:  - 5597 - torch._dynamo.symbolic_convert - Step 1: torchdynamo start tracing get_rank
stderr:  - 5597 - torch._dynamo.symbolic_convert - Step 1: torchdynamo done tracing get_rank (RETURN_VALUE)
stderr:  - 5597 - torch._dynamo.output_graph - Step 2: calling compiler function debug_wrapper
stderr:  - 5597 - torch._dynamo.output_graph - Step 2: done compiler function debug_wrapper
stderr:  - 5597 - torch._dynamo.symbolic_convert - Step 1: torchdynamo start tracing <graph break in range_push>
stderr:  - 5597 - torch._dynamo.symbolic_convert - Step 1: torchdynamo start tracing current_stream
stderr:  - 5597 - torch._dynamo.symbolic_convert - Step 1: torchdynamo done tracing current_stream (RETURN_VALUE)
stderr:  - 5597 - torch._dynamo.symbolic_convert - Step 1: torchdynamo start tracing <graph break in pre_sub_module_forward_function>
stderr:  - 5597 - torch._dynamo.symbolic_convert - Step 1: torchdynamo start tracing _post_backward_module_hook
stderr:  - 5597 - torch._dynamo.symbolic_convert - Step 1: torchdynamo start tracing _post_backward_module_hook
stderr:  - 5597 - torch._dynamo.symbolic_convert - Step 1: torchdynamo start tracing _apply_to_tensors_only
stderr:  - 5597 - torch._dynamo.symbolic_convert - Step 1: torchdynamo start tracing forward
stderr:  - 5597 - torch._dynamo.symbolic_convert - Step 1: torchdynamo start tracing <graph break in forward>
stderr:  - 5597 - torch._dynamo.symbolic_convert - Step 1: torchdynamo done tracing <graph break in forward> (RETURN_VALUE)
stderr:  - 5597 - torch._dynamo.output_graph - Step 2: calling compiler function debug_wrapper
stderr:  - 5597 - torch._inductor.compile_fx - Step 3: torchinductor compiling FORWARDS graph 1
stderr:  - 5597 - torch._inductor.compile_fx - Step 3: torchinductor done compiling FORWARDS graph 1
stderr:  - 5597 - torch._dynamo.output_graph - Step 2: done compiler function debug_wrapper
stderr:  - 5597 - torch._dynamo.symbolic_convert - Step 1: torchdynamo start tracing is_zero_param
stderr:  - 5597 - torch._dynamo.symbolic_convert - Step 1: torchdynamo start tracing forward
stderr:  - 5597 - torch._dynamo.symbolic_convert - Step 1: torchdynamo start tracing _pre_forward_module_hook
stderr:  - 5597 - torch._dynamo.symbolic_convert - Step 1: torchdynamo start tracing pre_sub_module_forward_function
stderr:  - 5597 - torch._dynamo.symbolic_convert - Step 1: torchdynamo start tracing record_module
stderr:  - 5597 - torch._dynamo.symbolic_convert - Step 1: torchdynamo start tracing <graph break in pre_sub_module_forward_function>
stderr:  - 5597 - torch._dynamo.symbolic_convert - Step 1: torchdynamo start tracing fetch_sub_module
stderr:  - 5597 - torch._dynamo.symbolic_convert - Step 1: torchdynamo start tracing iter_params
stderr:  - 5597 - torch._dynamo.symbolic_convert - Step 1: torchdynamo start tracing get_all_parameters
stderr:  - 5597 - torch._dynamo.symbolic_convert - Step 1: torchdynamo start tracing register_external_parameter
stderr:  - 5597 - torch._dynamo.symbolic_convert - Step 1: torchdynamo start tracing <graph break in register_external_parameter>
stderr:  - 5597 - torch._dynamo.symbolic_convert - Step 1: torchdynamo start tracing all_gather
stderr:  - 5597 - torch._dynamo.symbolic_convert - Step 1: torchdynamo start tracing <graph break in range_push>
stderr:  - 5597 - torch._dynamo.symbolic_convert - Step 1: torchdynamo start tracing _all_gather
stderr:  - 5597 - torch._dynamo.symbolic_convert - Step 1: torchdynamo start tracing _ensure_availability_of_partitioned_params
stderr:  - 5597 - torch._dynamo.symbolic_convert - Step 1: torchdynamo start tracing <graph break in _all_gather>
stderr:  - 5597 - torch._dynamo.symbolic_convert - Step 1: torchdynamo start tracing _allgather_params
stderr:  - 5597 - torch._dynamo.symbolic_convert - Step 1: torchdynamo start tracing <graph break in _allgather_params>
stderr:  - 5597 - torch._dynamo.output_graph - Step 2: calling compiler function debug_wrapper
stderr:  - 5597 - torch._inductor.compile_fx - Step 3: torchinductor compiling FORWARDS graph 2
stderr:  - 5597 - torch._inductor.compile_fx - Step 3: torchinductor done compiling FORWARDS graph 2
stderr:  - 5597 - torch._dynamo.output_graph - Step 2: done compiler function debug_wrapper
stderr:  - 5597 - torch._dynamo.symbolic_convert - Step 1: torchdynamo start tracing <graph break in _allgather_params>
stderr: Traceback (most recent call last):
stderr:   File "/usr/local/lib/python3.8/dist-packages/torch/_dynamo/convert_frame.py", line 324, in _compile
stderr:     out_code = transform_code_object(code, transform)
stderr:   File "/usr/local/lib/python3.8/dist-packages/torch/_dynamo/bytecode_transformation.py", line 445, in transform_code_object
stderr:     transformations(instructions, code_options)
stderr:   File "/usr/local/lib/python3.8/dist-packages/torch/_dynamo/convert_frame.py", line 311, in transform
stderr:     tracer.run()
stderr:   File "/usr/local/lib/python3.8/dist-packages/torch/_dynamo/symbolic_convert.py", line 1726, in run
stderr:     super().run()
stderr:   File "/usr/local/lib/python3.8/dist-packages/torch/_dynamo/symbolic_convert.py", line 576, in run
stderr:     and self.step()
stderr:   File "/usr/local/lib/python3.8/dist-packages/torch/_dynamo/symbolic_convert.py", line 540, in step
stderr:     getattr(self, inst.opname)(inst)
stderr:   File "/usr/local/lib/python3.8/dist-packages/torch/_dynamo/symbolic_convert.py", line 342, in wrapper
stderr:     return inner_fn(self, inst)
stderr:   File "/usr/local/lib/python3.8/dist-packages/torch/_dynamo/symbolic_convert.py", line 965, in CALL_FUNCTION
stderr:     self.call_function(fn, args, {})
stderr:   File "/usr/local/lib/python3.8/dist-packages/torch/_dynamo/symbolic_convert.py", line 474, in call_function
stderr:     self.push(fn.call_function(self, args, kwargs))
stderr:   File "/usr/local/lib/python3.8/dist-packages/torch/_dynamo/variables/misc.py", line 744, in call_function
stderr:     return self.obj.call_method(tx, self.name, args, kwargs).add_options(self)
stderr:   File "/usr/local/lib/python3.8/dist-packages/torch/_dynamo/variables/tensor.py", line 424, in call_method
stderr:     return wrap_fx_proxy(
stderr:   File "/usr/local/lib/python3.8/dist-packages/torch/_dynamo/variables/builder.py", line 754, in wrap_fx_proxy
stderr:     return wrap_fx_proxy_cls(
stderr:   File "/usr/local/lib/python3.8/dist-packages/torch/_dynamo/variables/builder.py", line 789, in wrap_fx_proxy_cls
stderr:     example_value = get_fake_value(proxy.node, tx)
stderr:   File "/usr/local/lib/python3.8/dist-packages/torch/_dynamo/utils.py", line 1134, in get_fake_value
stderr:     args, kwargs = torch.fx.node.map_arg((node.args, node.kwargs), visit)
stderr:   File "/usr/local/lib/python3.8/dist-packages/torch/fx/node.py", line 621, in map_arg
stderr:     return map_aggregate(a, lambda x: fn(x) if isinstance(x, Node) else x)
stderr:   File "/usr/local/lib/python3.8/dist-packages/torch/fx/node.py", line 629, in map_aggregate
stderr:     t = tuple(map_aggregate(elem, fn) for elem in a)
stderr:   File "/usr/local/lib/python3.8/dist-packages/torch/fx/node.py", line 629, in <genexpr>
stderr:     t = tuple(map_aggregate(elem, fn) for elem in a)
stderr:   File "/usr/local/lib/python3.8/dist-packages/torch/fx/node.py", line 629, in map_aggregate
stderr:     t = tuple(map_aggregate(elem, fn) for elem in a)
stderr:   File "/usr/local/lib/python3.8/dist-packages/torch/fx/node.py", line 629, in <genexpr>
stderr:     t = tuple(map_aggregate(elem, fn) for elem in a)
stderr:   File "/usr/local/lib/python3.8/dist-packages/torch/fx/node.py", line 639, in map_aggregate
stderr:     return fn(a)
stderr:   File "/usr/local/lib/python3.8/dist-packages/torch/fx/node.py", line 621, in <lambda>
stderr:     return map_aggregate(a, lambda x: fn(x) if isinstance(x, Node) else x)
stderr:   File "/usr/local/lib/python3.8/dist-packages/torch/_dynamo/utils.py", line 1132, in visit
stderr:     return n.meta["example_value"]
stderr: KeyError: example_value
stderr: 
stderr: from user code:
stderr:    File "/usr/local/lib/python3.8/dist-packages/deepspeed/runtime/zero/partition_parameters.py", line 1333, in <graph break in _allgather_params>
stderr:     partitions[i].narrow(0,
```

### Versions

pt-2.0

@Chillee 

cc @ezyang @msaroufim @wconstab @bdhirsh @anijain2305 @zou3519 @soumith @ngimel