mirror of
https://github.com/pytorch/pytorch.git
synced 2025-10-20 21:14:14 +08:00
Add structured logging for tensor fakeification (#126879)
This adds dumps of MetaTensorDesc and MetaStorageDesc to structured logs
when they are triggered from Dynamo. The logs look like this:
```
V0522 08:13:25.267000 140224882566144 torch/_subclasses/meta_utils.py:195] {"describe_storage": {"id": 0, "describer_id": 0, "size": 32}, "frame_id": 0, "frame_compile_id": 0, "attempt": 0}
V0522 08:13:25.267000 140224882566144 torch/_subclasses/meta_utils.py:220] {"describe_tensor": {"id": 0, "ndim": 1, "dtype": "torch.float32", "device": "device(type='cpu')", "size": [8], "is_leaf": true, "stride": [1], "storage": 0, "view_func": "<built-in method _view_func_unsafe of Tensor object at 0x7f882959e840>", "describer_id": 0}, "frame_id": 0, "frame_compile_id": 0, "attempt": 0}
V0522 08:13:25.268000 140224882566144 torch/_subclasses/meta_utils.py:1594] {"describe_source": {"describer_id": 0, "id": 0, "source": "L['x']"}, "frame_id": 0, "frame_compile_id": 0, "attempt": 0}
```
The `describer_id` is used to disambiguate ids. We expect it to be
unique per frame id, but if there is a bug it possibly is not. Note you will get
redundant dumps when evaluation restarts.
tlparse can use this to give a visualization of input tensors to a
model, you could also use this to generate example inputs to run graphs
on.
Some care is taken to avoid redumping the tensor metadata multiple
times, which would happen ordinarily because AOTAutograd refakifies
everything after Dynamo, to deal with metadata mutation.
Partially fixes https://github.com/pytorch/pytorch/issues/126644
Signed-off-by: Edward Z. Yang <ezyang@meta.com>
Pull Request resolved: https://github.com/pytorch/pytorch/pull/126879
Approved by: https://github.com/jamesjwu
This commit is contained in:
Edward Z. Yang
committed by
PyTorch MergeBot
PyTorch MergeBot
parent
b1792a622d
commit
0aaac68c57
@ -77,6 +77,14 @@ class StructuredTraceTestingFormatter(logging.Formatter):
|
||||
metadata["stack"] = "STACK"
|
||||
if "compilation_metrics" in metadata:
|
||||
metadata["compilation_metrics"] = "METRICS"
|
||||
if "describe_storage" in metadata:
|
||||
metadata["describe_storage"]["describer_id"] = "ID"
|
||||
if "describe_tensor" in metadata:
|
||||
metadata["describe_tensor"]["describer_id"] = "ID"
|
||||
if "view_func" in metadata["describe_tensor"]:
|
||||
metadata["describe_tensor"]["view_func"] = "VIEW_FUNC"
|
||||
if "describe_source" in metadata:
|
||||
metadata["describe_source"]["describer_id"] = "ID"
|
||||
|
||||
return json.dumps(metadata)
|
||||
|
||||
@ -136,6 +144,9 @@ class StructuredTraceTest(TestCase):
|
||||
self.buffer.getvalue(),
|
||||
"""\
|
||||
{"dynamo_start": {"stack": "STACK"}, "frame_id": 0, "frame_compile_id": 0, "attempt": 0}
|
||||
{"describe_storage": {"id": 0, "describer_id": "ID", "size": 4000000}, "frame_id": 0, "frame_compile_id": 0, "attempt": 0}
|
||||
{"describe_tensor": {"id": 0, "ndim": 2, "dtype": "torch.float32", "device": "device(type='cuda', index=0)", "size": [1000, 1000], "is_leaf": true, "stride": [1000, 1], "storage": 0, "view_func": "VIEW_FUNC", "describer_id": "ID"}, "frame_id": 0, "frame_compile_id": 0, "attempt": 0}
|
||||
{"describe_source": {"describer_id": "ID", "id": 0, "source": "L['a']"}, "frame_id": 0, "frame_compile_id": 0, "attempt": 0}
|
||||
{"dynamo_output_graph": {"sizes": {"l_a_": [1000, 1000], "ones": [1000, 1000], "output": [1000, 1000]}}, "frame_id": 0, "frame_compile_id": 0, "attempt": 0, "has_payload": "HASH"}
|
||||
{"aot_forward_graph": {}, "frame_id": 0, "frame_compile_id": 0, "attempt": 0, "has_payload": "HASH"}
|
||||
{"artifact": {"name": "fx_graph_cache_hash", "encoding": "json"}, "frame_id": 0, "frame_compile_id": 0, "attempt": 0, "has_payload": "HASH"}
|
||||
@ -157,6 +168,9 @@ class StructuredTraceTest(TestCase):
|
||||
self.buffer.getvalue(),
|
||||
"""\
|
||||
{"dynamo_start": {"stack": "STACK"}, "frame_id": 0, "frame_compile_id": 0, "attempt": 0}
|
||||
{"describe_storage": {"id": 0, "describer_id": "ID", "size": 4000000}, "frame_id": 0, "frame_compile_id": 0, "attempt": 0}
|
||||
{"describe_tensor": {"id": 0, "ndim": 2, "dtype": "torch.float32", "device": "device(type='cuda', index=0)", "size": [1000, 1000], "is_leaf": true, "stride": [1000, 1], "storage": 0, "view_func": "VIEW_FUNC", "describer_id": "ID"}, "frame_id": 0, "frame_compile_id": 0, "attempt": 0}
|
||||
{"describe_source": {"describer_id": "ID", "id": 0, "source": "L['a']"}, "frame_id": 0, "frame_compile_id": 0, "attempt": 0}
|
||||
{"dynamo_output_graph": {"sizes": {"l_a_": [1000, 1000], "ones": [1000, 1000], "output": [1000, 1000]}}, "frame_id": 0, "frame_compile_id": 0, "attempt": 0, "has_payload": "HASH"}
|
||||
{"aot_forward_graph": {}, "frame_id": 0, "frame_compile_id": 0, "attempt": 0, "has_payload": "HASH"}
|
||||
{"artifact": {"name": "fx_graph_cache_hash", "encoding": "json"}, "frame_id": 0, "frame_compile_id": 0, "attempt": 0, "has_payload": "HASH"}
|
||||
@ -182,6 +196,12 @@ class StructuredTraceTest(TestCase):
|
||||
self.buffer.getvalue(),
|
||||
"""\
|
||||
{"dynamo_start": {"stack": "STACK"}, "frame_id": 0, "frame_compile_id": 0, "attempt": 0}
|
||||
{"describe_storage": {"id": 0, "describer_id": "ID", "size": 4000000}, "frame_id": 0, "frame_compile_id": 0, "attempt": 0}
|
||||
{"describe_tensor": {"id": 0, "ndim": 2, "dtype": "torch.float32", "device": "device(type='cpu')", "size": [1000, 1000], "is_leaf": true, "stride": [1000, 1], "storage": 0, "view_func": "VIEW_FUNC", "describer_id": "ID"}, "frame_id": 0, "frame_compile_id": 0, "attempt": 0}
|
||||
{"describe_source": {"describer_id": "ID", "id": 0, "source": "L['x']"}, "frame_id": 0, "frame_compile_id": 0, "attempt": 0}
|
||||
{"describe_storage": {"id": 1, "describer_id": "ID", "size": 4000000}, "frame_id": 0, "frame_compile_id": 0, "attempt": 0}
|
||||
{"describe_tensor": {"id": 1, "ndim": 2, "dtype": "torch.float32", "device": "device(type='cpu')", "size": [1000, 1000], "is_leaf": true, "stride": [1000, 1], "storage": 1, "view_func": "VIEW_FUNC", "describer_id": "ID"}, "frame_id": 0, "frame_compile_id": 0, "attempt": 0}
|
||||
{"describe_source": {"describer_id": "ID", "id": 1, "source": "L['y']"}, "frame_id": 0, "frame_compile_id": 0, "attempt": 0}
|
||||
{"dynamo_output_graph": {"sizes": {"l_x_": [1000, 1000], "l_y_": [1000, 1000], "add": [1000, 1000]}}, "frame_id": 0, "frame_compile_id": 0, "attempt": 0, "has_payload": "HASH"}
|
||||
{"aot_forward_graph": {}, "frame_id": 0, "frame_compile_id": 0, "attempt": 0, "has_payload": "HASH"}
|
||||
{"artifact": {"name": "fx_graph_cache_hash", "encoding": "json"}, "frame_id": 0, "frame_compile_id": 0, "attempt": 0, "has_payload": "HASH"}
|
||||
@ -191,6 +211,9 @@ class StructuredTraceTest(TestCase):
|
||||
{"dynamo_cpp_guards_str": {}, "frame_id": 0, "frame_compile_id": 0, "attempt": 0, "has_payload": "HASH"}
|
||||
{"compilation_metrics": "METRICS", "frame_id": 0, "frame_compile_id": 0, "attempt": 0}
|
||||
{"dynamo_start": {"stack": "STACK"}, "frame_id": 0, "frame_compile_id": 1, "attempt": 0}
|
||||
{"describe_storage": {"id": 0, "describer_id": "ID", "size": 4000000}, "frame_id": 0, "frame_compile_id": 1, "attempt": 0}
|
||||
{"describe_tensor": {"id": 0, "ndim": 2, "dtype": "torch.float32", "device": "device(type='cpu')", "size": [1000, 1000], "is_leaf": true, "stride": [1000, 1], "storage": 0, "view_func": "VIEW_FUNC", "describer_id": "ID"}, "frame_id": 0, "frame_compile_id": 1, "attempt": 0}
|
||||
{"describe_source": {"describer_id": "ID", "id": 0, "source": "L['x']"}, "frame_id": 0, "frame_compile_id": 1, "attempt": 0}
|
||||
{"dynamo_output_graph": {"sizes": {"l_x_": [1000, 1000], "add": [1000, 1000]}}, "frame_id": 0, "frame_compile_id": 1, "attempt": 0, "has_payload": "HASH"}
|
||||
{"aot_forward_graph": {}, "frame_id": 0, "frame_compile_id": 1, "attempt": 0, "has_payload": "HASH"}
|
||||
{"artifact": {"name": "fx_graph_cache_hash", "encoding": "json"}, "frame_id": 0, "frame_compile_id": 1, "attempt": 0, "has_payload": "HASH"}
|
||||
@ -211,6 +234,9 @@ class StructuredTraceTest(TestCase):
|
||||
self.buffer.getvalue(),
|
||||
"""\
|
||||
{"dynamo_start": {"stack": "STACK"}, "frame_id": 0, "frame_compile_id": 0, "attempt": 0}
|
||||
{"describe_storage": {"id": 0, "describer_id": "ID", "size": 4000000}, "frame_id": 0, "frame_compile_id": 0, "attempt": 0}
|
||||
{"describe_tensor": {"id": 0, "ndim": 2, "dtype": "torch.float32", "device": "device(type='cpu')", "size": [1000, 1000], "is_leaf": true, "stride": [1000, 1], "storage": 0, "view_func": "VIEW_FUNC", "describer_id": "ID"}, "frame_id": 0, "frame_compile_id": 0, "attempt": 0}
|
||||
{"describe_source": {"describer_id": "ID", "id": 0, "source": "L['a']"}, "frame_id": 0, "frame_compile_id": 0, "attempt": 0}
|
||||
{"dynamo_output_graph": {"sizes": {"l_a_": [1000, 1000], "ones": [1000, 1000], "output": [1000, 1000], "ones_1": [1000, 1000], "output_1": [1000, 1000]}}, "frame_id": 0, "frame_compile_id": 0, "attempt": 0, "has_payload": "HASH"}
|
||||
{"aot_forward_graph": {}, "frame_id": 0, "frame_compile_id": 0, "attempt": 0, "has_payload": "HASH"}
|
||||
{"artifact": {"name": "fx_graph_cache_hash", "encoding": "json"}, "frame_id": 0, "frame_compile_id": 0, "attempt": 0, "has_payload": "HASH"}
|
||||
@ -234,6 +260,9 @@ class StructuredTraceTest(TestCase):
|
||||
self.buffer.getvalue(),
|
||||
"""\
|
||||
{"dynamo_start": {"stack": "STACK"}, "frame_id": 0, "frame_compile_id": 0, "attempt": 0}
|
||||
{"describe_storage": {"id": 0, "describer_id": "ID", "size": 4000000}, "frame_id": 0, "frame_compile_id": 0, "attempt": 0}
|
||||
{"describe_tensor": {"id": 0, "ndim": 2, "dtype": "torch.float32", "device": "device(type='cpu')", "size": [1000, 1000], "is_leaf": true, "requires_grad": true, "stride": [1000, 1], "storage": 0, "view_func": "VIEW_FUNC", "describer_id": "ID"}, "frame_id": 0, "frame_compile_id": 0, "attempt": 0}
|
||||
{"describe_source": {"describer_id": "ID", "id": 0, "source": "L['a']"}, "frame_id": 0, "frame_compile_id": 0, "attempt": 0}
|
||||
{"compilation_metrics": "METRICS", "frame_id": 0, "frame_compile_id": 0, "attempt": 0}
|
||||
""", # noqa: B950
|
||||
)
|
||||
@ -263,6 +292,9 @@ class StructuredTraceTest(TestCase):
|
||||
self.buffer.getvalue(),
|
||||
"""\
|
||||
{"dynamo_start": {"stack": "STACK"}, "frame_id": 0, "frame_compile_id": 0, "attempt": 0}
|
||||
{"describe_storage": {"id": 0, "describer_id": "ID", "size": 4000000}, "frame_id": 0, "frame_compile_id": 0, "attempt": 0}
|
||||
{"describe_tensor": {"id": 0, "ndim": 2, "dtype": "torch.float32", "device": "device(type='cpu')", "size": [1000, 1000], "is_leaf": true, "requires_grad": true, "stride": [1000, 1], "storage": 0, "view_func": "VIEW_FUNC", "describer_id": "ID"}, "frame_id": 0, "frame_compile_id": 0, "attempt": 0}
|
||||
{"describe_source": {"describer_id": "ID", "id": 0, "source": "L['a']"}, "frame_id": 0, "frame_compile_id": 0, "attempt": 0}
|
||||
{"dynamo_output_graph": {"sizes": {"l_a_": [1000, 1000], "output": [1000, 1000]}}, "frame_id": 0, "frame_compile_id": 0, "attempt": 0, "has_payload": "HASH"}
|
||||
{"aot_joint_graph": {}, "frame_id": 0, "frame_compile_id": 0, "attempt": 0, "has_payload": "HASH"}
|
||||
{"aot_forward_graph": {}, "frame_id": 0, "frame_compile_id": 0, "attempt": 0, "has_payload": "HASH"}
|
||||
@ -310,10 +342,16 @@ class StructuredTraceTest(TestCase):
|
||||
{"dynamo_cpp_guards_str": {}, "rank": 0, "frame_id": 0, "frame_compile_id": 0, "attempt": 1, "has_payload": "HASH"}
|
||||
{"compilation_metrics": "METRICS", "rank": 0, "frame_id": 0, "frame_compile_id": 0, "attempt": 1}
|
||||
{"dynamo_start": {"stack": "STACK"}, "rank": 0, "frame_id": 1, "frame_compile_id": 0, "attempt": 0}
|
||||
{"describe_storage": {"id": 0, "describer_id": "ID", "size": 4194304}, "rank": 0, "frame_id": 1, "frame_compile_id": 0, "attempt": 0}
|
||||
{"describe_tensor": {"id": 0, "ndim": 2, "dtype": "torch.float32", "device": "device(type='cuda', index=0)", "size": [1024, 1024], "is_leaf": true, "stride": [1024, 1], "storage": 0, "view_func": "VIEW_FUNC", "describer_id": "ID"}, "rank": 0, "frame_id": 1, "frame_compile_id": 0, "attempt": 0}
|
||||
{"describe_source": {"describer_id": "ID", "id": 0, "source": "L['x']"}, "rank": 0, "frame_id": 1, "frame_compile_id": 0, "attempt": 0}
|
||||
{"dynamo_output_graph": {"sizes": {"l_x_": [1024, 1024], "l__self___layers_0": [1024, 1024], "l__self___layers_1": [1024, 1024]}}, "rank": 0, "frame_id": 1, "frame_compile_id": 0, "attempt": 0, "has_payload": "HASH"}
|
||||
{"optimize_ddp_split_graph": {}, "rank": 0, "frame_id": 1, "frame_compile_id": 0, "attempt": 0, "has_payload": "HASH"}
|
||||
{"optimize_ddp_split_child": {"name": "submod_0"}, "rank": 0, "frame_id": 1, "frame_compile_id": 0, "attempt": 0, "has_payload": "HASH"}
|
||||
{"optimize_ddp_split_child": {"name": "submod_1"}, "rank": 0, "frame_id": 1, "frame_compile_id": 0, "attempt": 0, "has_payload": "HASH"}
|
||||
{"describe_storage": {"id": 0, "describer_id": "ID", "size": 4194304}, "rank": 0, "frame_id": 1, "frame_compile_id": 0, "attempt": 0}
|
||||
{"describe_tensor": {"id": 0, "ndim": 2, "dtype": "torch.float32", "device": "device(type='cuda', index=0)", "size": [1024, 1024], "is_leaf": true, "stride": [1024, 1], "storage": 0, "view_func": "VIEW_FUNC", "describer_id": "ID"}, "rank": 0, "frame_id": 1, "frame_compile_id": 0, "attempt": 0}
|
||||
{"describe_source": {"describer_id": "ID", "id": 0, "source": "L['x']"}, "rank": 0, "frame_id": 1, "frame_compile_id": 0, "attempt": 0}
|
||||
{"aot_joint_graph": {}, "rank": 0, "frame_id": 1, "frame_compile_id": 0, "attempt": 0, "has_payload": "HASH"}
|
||||
{"aot_forward_graph": {}, "rank": 0, "frame_id": 1, "frame_compile_id": 0, "attempt": 0, "has_payload": "HASH"}
|
||||
{"aot_backward_graph": {}, "rank": 0, "frame_id": 1, "frame_compile_id": 0, "attempt": 0, "has_payload": "HASH"}
|
||||
@ -350,6 +388,9 @@ class StructuredTraceTest(TestCase):
|
||||
{"dynamo_cpp_guards_str": {}, "frame_id": 0, "frame_compile_id": 0, "attempt": 1, "has_payload": "HASH"}
|
||||
{"compilation_metrics": "METRICS", "frame_id": 0, "frame_compile_id": 0, "attempt": 1}
|
||||
{"dynamo_start": {"stack": "STACK"}, "frame_id": 1, "frame_compile_id": 0, "attempt": 0}
|
||||
{"describe_storage": {"id": 0, "describer_id": "ID", "size": 4}, "frame_id": 1, "frame_compile_id": 0, "attempt": 0}
|
||||
{"describe_tensor": {"id": 0, "ndim": 1, "dtype": "torch.float32", "device": "device(type='cpu')", "size": [1], "is_leaf": true, "stride": [1], "storage": 0, "view_func": "VIEW_FUNC", "describer_id": "ID"}, "frame_id": 1, "frame_compile_id": 0, "attempt": 0}
|
||||
{"describe_source": {"describer_id": "ID", "id": 0, "source": "L['x']"}, "frame_id": 1, "frame_compile_id": 0, "attempt": 0}
|
||||
{"dynamo_output_graph": {"sizes": {"l_x_": [1], "add": [1]}}, "frame_id": 1, "frame_compile_id": 0, "attempt": 0, "has_payload": "HASH"}
|
||||
{"aot_forward_graph": {}, "frame_id": 1, "frame_compile_id": 0, "attempt": 0, "has_payload": "HASH"}
|
||||
{"artifact": {"name": "fx_graph_cache_hash", "encoding": "json"}, "frame_id": 1, "frame_compile_id": 0, "attempt": 0, "has_payload": "HASH"}
|
||||
@ -379,11 +420,23 @@ class StructuredTraceTest(TestCase):
|
||||
self.buffer.getvalue(),
|
||||
"""\
|
||||
{"dynamo_start": {"stack": "STACK"}, "frame_id": 0, "frame_compile_id": 0, "attempt": 0}
|
||||
{"describe_storage": {"id": 0, "describer_id": "ID", "size": 800}, "frame_id": 0, "frame_compile_id": 0, "attempt": 0}
|
||||
{"describe_tensor": {"id": 0, "ndim": 2, "dtype": "torch.float32", "device": "device(type='cpu')", "size": [10, 20], "is_leaf": true, "stride": [20, 1], "storage": 0, "view_func": "VIEW_FUNC", "describer_id": "ID"}, "frame_id": 0, "frame_compile_id": 0, "attempt": 0}
|
||||
{"describe_source": {"describer_id": "ID", "id": 0, "source": "L['a']"}, "frame_id": 0, "frame_compile_id": 0, "attempt": 0}
|
||||
{"describe_storage": {"id": 1, "describer_id": "ID", "size": 2400}, "frame_id": 0, "frame_compile_id": 0, "attempt": 0}
|
||||
{"describe_tensor": {"id": 1, "ndim": 2, "dtype": "torch.float32", "device": "device(type='cpu')", "size": [20, 30], "is_leaf": true, "stride": [30, 1], "storage": 1, "view_func": "VIEW_FUNC", "describer_id": "ID"}, "frame_id": 0, "frame_compile_id": 0, "attempt": 0}
|
||||
{"describe_source": {"describer_id": "ID", "id": 1, "source": "L['b']"}, "frame_id": 0, "frame_compile_id": 0, "attempt": 0}
|
||||
{"dynamo_output_graph": {"sizes": {"l_a_": [10, 20], "l_b_": [20, 30], "matmul": [10, 30]}}, "frame_id": 0, "frame_compile_id": 0, "attempt": 0, "has_payload": "HASH"}
|
||||
{"dynamo_guards": {}, "frame_id": 0, "frame_compile_id": 0, "attempt": 0, "has_payload": "HASH"}
|
||||
{"dynamo_cpp_guards_str": {}, "frame_id": 0, "frame_compile_id": 0, "attempt": 0, "has_payload": "HASH"}
|
||||
{"compilation_metrics": "METRICS", "frame_id": 0, "frame_compile_id": 0, "attempt": 0}
|
||||
{"dynamo_start": {"stack": "STACK"}, "frame_id": 0, "frame_compile_id": 1, "attempt": 0}
|
||||
{"describe_storage": {"id": 0, "describer_id": "ID", "size": 200}, "frame_id": 0, "frame_compile_id": 1, "attempt": 0}
|
||||
{"describe_tensor": {"id": 0, "ndim": 2, "dtype": "torch.float32", "device": "device(type='cpu')", "size": [5, 10], "is_leaf": true, "stride": [10, 1], "storage": 0, "view_func": "VIEW_FUNC", "describer_id": "ID"}, "frame_id": 0, "frame_compile_id": 1, "attempt": 0}
|
||||
{"describe_source": {"describer_id": "ID", "id": 0, "source": "L['a']"}, "frame_id": 0, "frame_compile_id": 1, "attempt": 0}
|
||||
{"describe_storage": {"id": 1, "describer_id": "ID", "size": 600}, "frame_id": 0, "frame_compile_id": 1, "attempt": 0}
|
||||
{"describe_tensor": {"id": 1, "ndim": 2, "dtype": "torch.float32", "device": "device(type='cpu')", "size": [10, 15], "is_leaf": true, "stride": [15, 1], "storage": 1, "view_func": "VIEW_FUNC", "describer_id": "ID"}, "frame_id": 0, "frame_compile_id": 1, "attempt": 0}
|
||||
{"describe_source": {"describer_id": "ID", "id": 1, "source": "L['b']"}, "frame_id": 0, "frame_compile_id": 1, "attempt": 0}
|
||||
{"dynamo_output_graph": {"sizes": {"l_a_": ["s0", "s1"], "l_b_": ["s1", "s3"], "matmul": ["s0", "s3"]}}, "frame_id": 0, "frame_compile_id": 1, "attempt": 0, "has_payload": "HASH"}
|
||||
{"dynamo_guards": {}, "frame_id": 0, "frame_compile_id": 1, "attempt": 0, "has_payload": "HASH"}
|
||||
{"dynamo_cpp_guards_str": {}, "frame_id": 0, "frame_compile_id": 1, "attempt": 0, "has_payload": "HASH"}
|
||||
@ -414,11 +467,17 @@ class StructuredTraceTest(TestCase):
|
||||
self.buffer.getvalue(),
|
||||
"""\
|
||||
{"dynamo_start": {"stack": "STACK"}, "frame_id": 0, "frame_compile_id": 0, "attempt": 0}
|
||||
{"describe_storage": {"id": 0, "describer_id": "ID", "size": 4}, "frame_id": 0, "frame_compile_id": 0, "attempt": 0}
|
||||
{"describe_tensor": {"id": 0, "ndim": 1, "dtype": "torch.float32", "device": "device(type='cpu')", "size": [1], "is_leaf": true, "stride": [1], "storage": 0, "view_func": "VIEW_FUNC", "describer_id": "ID"}, "frame_id": 0, "frame_compile_id": 0, "attempt": 0}
|
||||
{"describe_source": {"describer_id": "ID", "id": 0, "source": "L['x']"}, "frame_id": 0, "frame_compile_id": 0, "attempt": 0}
|
||||
{"dynamo_output_graph": {"sizes": {"l_x_": [1], "x": [1]}}, "frame_id": 0, "frame_compile_id": 0, "attempt": 0, "has_payload": "HASH"}
|
||||
{"dynamo_guards": {}, "frame_id": 0, "frame_compile_id": 0, "attempt": 0, "has_payload": "HASH"}
|
||||
{"dynamo_cpp_guards_str": {}, "frame_id": 0, "frame_compile_id": 0, "attempt": 0, "has_payload": "HASH"}
|
||||
{"compilation_metrics": "METRICS", "frame_id": 0, "frame_compile_id": 0, "attempt": 0}
|
||||
{"dynamo_start": {"stack": "STACK"}, "frame_id": 0, "frame_compile_id": 1, "attempt": 0}
|
||||
{"describe_storage": {"id": 0, "describer_id": "ID", "size": 4}, "frame_id": 0, "frame_compile_id": 1, "attempt": 0}
|
||||
{"describe_tensor": {"id": 0, "ndim": 1, "dtype": "torch.float32", "device": "device(type='cpu')", "size": [1], "is_leaf": true, "stride": [1], "storage": 0, "view_func": "VIEW_FUNC", "describer_id": "ID"}, "frame_id": 0, "frame_compile_id": 1, "attempt": 0}
|
||||
{"describe_source": {"describer_id": "ID", "id": 0, "source": "L['x']"}, "frame_id": 0, "frame_compile_id": 1, "attempt": 0}
|
||||
{"dynamo_output_graph": {"sizes": {"l_x_": [1], "x": [1]}}, "frame_id": 0, "frame_compile_id": 1, "attempt": 0, "has_payload": "HASH"}
|
||||
{"dynamo_guards": {}, "frame_id": 0, "frame_compile_id": 1, "attempt": 0, "has_payload": "HASH"}
|
||||
{"dynamo_cpp_guards_str": {}, "frame_id": 0, "frame_compile_id": 1, "attempt": 0, "has_payload": "HASH"}
|
||||
|
||||
@ -509,10 +509,14 @@ def create_aot_dispatcher_function(
|
||||
# see note [Tensor Fakification and Symbol Caching]
|
||||
symbolic_context = None
|
||||
source = None
|
||||
trace = True
|
||||
if tracing_context := torch._guards.TracingContext.try_get():
|
||||
if x in tracing_context.tensor_to_context:
|
||||
symbolic_context = tracing_context.tensor_to_context[x]
|
||||
source = symbolic_context.tensor_source
|
||||
# We already fakeified this tensor in Dynamo, don't
|
||||
# dump the trace for it again
|
||||
trace = False
|
||||
if (
|
||||
idx < aot_config.num_params_buffers
|
||||
and config.static_weight_shapes
|
||||
@ -527,6 +531,7 @@ def create_aot_dispatcher_function(
|
||||
static_shapes=False,
|
||||
symbolic_context=symbolic_context,
|
||||
source=source,
|
||||
trace=trace,
|
||||
)
|
||||
|
||||
return [convert(idx, x) for idx, x in enumerate(flat_args)]
|
||||
|
||||
@ -795,7 +795,11 @@ class TorchLogsFormatter(logging.Formatter):
|
||||
)
|
||||
if self._is_trace:
|
||||
assert s == ""
|
||||
r = f"{prefix} {json.dumps(record.metadata)}"
|
||||
try:
|
||||
r = f"{prefix} {json.dumps(record.metadata)}"
|
||||
except TypeError:
|
||||
log.warning("failing metadata: %r", record.metadata)
|
||||
raise
|
||||
if record.payload is not None:
|
||||
r += "".join(f"\n\t{l}" for l in record.payload.split("\n"))
|
||||
return r
|
||||
|
||||
@ -291,6 +291,7 @@ class FakeTensorConverter:
|
||||
*,
|
||||
source=None,
|
||||
symbolic_context=None,
|
||||
trace=True,
|
||||
):
|
||||
# see note [Tensor Fakification and Symbol Caching]
|
||||
if not symbolic_context and not source and shape_env:
|
||||
@ -333,6 +334,7 @@ class FakeTensorConverter:
|
||||
callback=mk_fake_tensor,
|
||||
source=source,
|
||||
symbolic_context=symbolic_context,
|
||||
trace=trace,
|
||||
)
|
||||
if out is NotImplemented:
|
||||
raise UnsupportedFakeTensorException("meta converter nyi")
|
||||
@ -1925,6 +1927,7 @@ class FakeTensorMode(TorchDispatchMode):
|
||||
static_shapes=None,
|
||||
source: Optional[Source] = None,
|
||||
symbolic_context=None,
|
||||
trace=True,
|
||||
):
|
||||
shape_env: Optional[ShapeEnv] = self.shape_env
|
||||
if static_shapes is None:
|
||||
@ -1940,6 +1943,7 @@ class FakeTensorMode(TorchDispatchMode):
|
||||
shape_env=shape_env,
|
||||
source=source,
|
||||
symbolic_context=symbolic_context,
|
||||
trace=trace,
|
||||
)
|
||||
|
||||
|
||||
|
||||
@ -1,13 +1,15 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import contextlib
|
||||
|
||||
import dataclasses
|
||||
import warnings
|
||||
import weakref
|
||||
|
||||
from dataclasses import dataclass
|
||||
from typing import (
|
||||
Any,
|
||||
Callable,
|
||||
ClassVar,
|
||||
ContextManager,
|
||||
Dict,
|
||||
List,
|
||||
@ -20,6 +22,7 @@ from typing import (
|
||||
from typing_extensions import TypeAlias
|
||||
|
||||
import torch
|
||||
from torch._C._autograd import CreationMeta
|
||||
from torch._C._functorch import (
|
||||
_add_batch_dim,
|
||||
_unwrap_functional_tensor,
|
||||
@ -33,13 +36,13 @@ from torch._C._functorch import (
|
||||
maybe_get_level,
|
||||
peek_interpreter_stack,
|
||||
)
|
||||
from torch._logging import trace_structured
|
||||
from torch.utils._mode_utils import no_dispatch
|
||||
|
||||
from torch.utils._python_dispatch import is_traceable_wrapper_subclass
|
||||
from torch.utils.weak import WeakIdKeyDictionary
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from torch._C._autograd import CreationMeta
|
||||
from torch._C._functorch import CInterpreter
|
||||
from torch._guards import Source
|
||||
|
||||
@ -142,6 +145,9 @@ MetaStorageId: TypeAlias = int
|
||||
MetaTensorId: TypeAlias = int
|
||||
|
||||
|
||||
DESCRIBER_NEXT_ID = 0
|
||||
|
||||
|
||||
class MetaTensorDescriber:
|
||||
"""
|
||||
Given a Tensor/Storage, generate a MetaTensorDesc/MetaStorageDesc
|
||||
@ -154,6 +160,9 @@ class MetaTensorDescriber:
|
||||
"""
|
||||
|
||||
def __init__(self, *, copy_data=False):
|
||||
global DESCRIBER_NEXT_ID
|
||||
self.id = DESCRIBER_NEXT_ID
|
||||
DESCRIBER_NEXT_ID += 1
|
||||
self.next_tensor_id: MetaTensorId = 0
|
||||
self.next_storage_id: MetaStorageId = 0
|
||||
# Tensor -> int
|
||||
@ -161,6 +170,8 @@ class MetaTensorDescriber:
|
||||
# Storage -> int
|
||||
self.lookup_storage = WeakIdKeyDictionary()
|
||||
self.copy_data = copy_data
|
||||
self.traced_tensors = set()
|
||||
self.traced_storages = set()
|
||||
|
||||
def get_tensor_id(self, t: torch.Tensor):
|
||||
if t not in self.lookup_tensor:
|
||||
@ -174,19 +185,25 @@ class MetaTensorDescriber:
|
||||
self.next_storage_id += 1
|
||||
return self.lookup_storage[s]
|
||||
|
||||
# NB: the describe functions NOT maintain a cache and will happily regen the
|
||||
# description
|
||||
|
||||
def describe_storage(self, s: torch.UntypedStorage):
|
||||
return MetaStorageDesc(
|
||||
def describe_storage(self, s: torch.UntypedStorage, *, trace: bool = False):
|
||||
r = MetaStorageDesc(
|
||||
id=self.get_storage_id(s),
|
||||
size=s.size(),
|
||||
# NB: We don't do the copy yet; copy happens when we start
|
||||
# creating the new storages
|
||||
data=s if self.copy_data else None,
|
||||
)
|
||||
if trace and r.id not in self.traced_storages:
|
||||
trace_structured(
|
||||
"describe_storage",
|
||||
metadata_fn=lambda: r.as_json(self.id),
|
||||
)
|
||||
self.traced_storages.add(r.id)
|
||||
return r
|
||||
|
||||
def describe_tensor(self, t: torch.Tensor, recurse: bool = True):
|
||||
def describe_tensor(
|
||||
self, t: torch.Tensor, *, recurse: bool = True, trace: bool = False
|
||||
):
|
||||
is_leaf = safe_is_leaf(t)
|
||||
is_view = t._is_view()
|
||||
is_sparse = t.is_sparse
|
||||
@ -218,7 +235,7 @@ class MetaTensorDescriber:
|
||||
):
|
||||
# NB: We actually don't use storage to do views, but might as well
|
||||
# put it in for accuracy
|
||||
storage = self.describe_storage(t.untyped_storage())
|
||||
storage = self.describe_storage(t.untyped_storage(), trace=trace)
|
||||
storage_offset = t.storage_offset()
|
||||
|
||||
stride = None
|
||||
@ -239,7 +256,7 @@ class MetaTensorDescriber:
|
||||
autograd_meta_from = None
|
||||
current_level = None
|
||||
if is_batchedtensor_v or is_gradtrackingtensor_v:
|
||||
unwrapped = self.describe_tensor(get_unwrapped(t))
|
||||
unwrapped = self.describe_tensor(get_unwrapped(t), trace=trace)
|
||||
# xla and lazy tensors present as functional tensors, but we want them
|
||||
# to be handled specially
|
||||
elif is_functional and t.device.type not in ("xla", "lazy"):
|
||||
@ -249,13 +266,15 @@ class MetaTensorDescriber:
|
||||
)
|
||||
if not is_functorch_wrapped:
|
||||
torch._sync(t)
|
||||
unwrapped = self.describe_tensor(torch._from_functional_tensor(t))
|
||||
unwrapped = self.describe_tensor(
|
||||
torch._from_functional_tensor(t), trace=trace
|
||||
)
|
||||
autograd_meta_from = t
|
||||
else:
|
||||
reapply_views = torch._C._functionalization_reapply_views_tls()
|
||||
# NB: has side effects!
|
||||
unwrapped = self.describe_tensor(
|
||||
_unwrap_functional_tensor(t, reapply_views)
|
||||
_unwrap_functional_tensor(t, reapply_views), trace=trace
|
||||
)
|
||||
# TODO: It's pretty suspicious that functional tensors don't have
|
||||
# valid level and thus we just grab whatever the current level
|
||||
@ -273,12 +292,15 @@ class MetaTensorDescriber:
|
||||
if is_traceable_wrapper_subclass_v:
|
||||
assert hasattr(t, "__tensor_flatten__")
|
||||
raw_attrs, ctx = t.__tensor_flatten__()
|
||||
attrs = {attr: self.describe_tensor(getattr(t, attr)) for attr in raw_attrs}
|
||||
attrs = {
|
||||
attr: self.describe_tensor(getattr(t, attr), trace=trace)
|
||||
for attr in raw_attrs
|
||||
}
|
||||
type_v = type(t)
|
||||
|
||||
# TODO: Is it important to enable torch.inference_mode before querying
|
||||
# these values?
|
||||
return MetaTensorDesc(
|
||||
r = MetaTensorDesc(
|
||||
id=self.get_tensor_id(t),
|
||||
storage=storage,
|
||||
is_inference=t.is_inference(),
|
||||
@ -318,22 +340,30 @@ class MetaTensorDescriber:
|
||||
# TODO: I actually think recursing here is correct, but we have at
|
||||
# least an infinite cycle from base -> values -> base
|
||||
# https://github.com/pytorch/pytorch/issues/122089
|
||||
crow_indices=self.describe_tensor(t.crow_indices(), recurse=False)
|
||||
crow_indices=self.describe_tensor(
|
||||
t.crow_indices(), recurse=False, trace=trace
|
||||
)
|
||||
if recurse and t.layout in {torch.sparse_csr, torch.sparse_bsr}
|
||||
else None,
|
||||
col_indices=self.describe_tensor(t.col_indices(), recurse=False)
|
||||
col_indices=self.describe_tensor(
|
||||
t.col_indices(), recurse=False, trace=trace
|
||||
)
|
||||
if recurse and t.layout in {torch.sparse_csr, torch.sparse_bsr}
|
||||
else None,
|
||||
ccol_indices=self.describe_tensor(t.ccol_indices(), recurse=False)
|
||||
ccol_indices=self.describe_tensor(
|
||||
t.ccol_indices(), recurse=False, trace=trace
|
||||
)
|
||||
if recurse and t.layout in {torch.sparse_csc, torch.sparse_bsc}
|
||||
else None,
|
||||
row_indices=self.describe_tensor(t.row_indices(), recurse=False)
|
||||
row_indices=self.describe_tensor(
|
||||
t.row_indices(), recurse=False, trace=trace
|
||||
)
|
||||
if recurse and t.layout in {torch.sparse_csc, torch.sparse_bsc}
|
||||
else None,
|
||||
values=self.describe_tensor(t.values(), recurse=False)
|
||||
values=self.describe_tensor(t.values(), recurse=False, trace=trace)
|
||||
if recurse and is_sparse_compressed(t)
|
||||
else None,
|
||||
grad=self.describe_tensor(safe_grad(t))
|
||||
grad=self.describe_tensor(safe_grad(t), trace=trace)
|
||||
if safe_grad(t) is not None
|
||||
else None,
|
||||
creation_meta=torch._C._autograd._get_creation_meta(t)
|
||||
@ -344,7 +374,7 @@ class MetaTensorDescriber:
|
||||
if is_batchedtensor_v or is_gradtrackingtensor_v
|
||||
else None,
|
||||
bdim=maybe_get_bdim(t) if is_batchedtensor_v else None,
|
||||
base=self.describe_tensor(t._base)
|
||||
base=self.describe_tensor(t._base, trace=trace)
|
||||
if recurse and t._is_view() and t._base is not None
|
||||
else None,
|
||||
fake_mode=torch._subclasses.fake_tensor.maybe_get_fake_mode(t),
|
||||
@ -360,6 +390,13 @@ class MetaTensorDescriber:
|
||||
current_level=current_level,
|
||||
data=t if self.copy_data else None,
|
||||
)
|
||||
if trace and r.id not in self.traced_tensors:
|
||||
trace_structured(
|
||||
"describe_tensor",
|
||||
metadata_fn=lambda: r.as_json(self.id),
|
||||
)
|
||||
self.traced_tensors.add(r.id)
|
||||
return r
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
@ -370,43 +407,58 @@ class MetaStorageDesc:
|
||||
# serializable in JSON, you want to do something special here anyway
|
||||
data: Optional[torch.UntypedStorage]
|
||||
|
||||
def as_json(self, describer_id):
|
||||
return {
|
||||
"id": self.id,
|
||||
"describer_id": describer_id,
|
||||
"size": self.size if isinstance(self.size, int) else repr(self.size),
|
||||
}
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class MetaTensorDesc:
|
||||
id: MetaTensorId
|
||||
is_inference: bool
|
||||
is_leaf: bool
|
||||
requires_grad: bool
|
||||
ndim: int
|
||||
dtype: torch.dtype
|
||||
is_sparse: bool
|
||||
is_mkldnn: bool
|
||||
is_functorch_wrapped: bool
|
||||
is_batchedtensor: bool
|
||||
is_legacy_batchedtensor: bool
|
||||
is_gradtrackingtensor: bool
|
||||
is_view: bool
|
||||
is_nested: bool
|
||||
is_traceable_wrapper_subclass: bool
|
||||
is_functional: bool
|
||||
is_conj: bool
|
||||
is_neg: bool
|
||||
device: torch.device
|
||||
layout: torch.layout
|
||||
|
||||
# NB: Sometimes, size, stride and storage_offset contain SymInt, in which
|
||||
# case this is NOT serializable. That only happens when you're
|
||||
# re-fakeifying a fake tensor with an existing ShapeEnv... maybe we
|
||||
# can get rid of this use case entirely. Notably, even if we are
|
||||
# fakeifying a real tensor into a fake tensor with symbolic shapes, the
|
||||
# size here is NOT dynamic
|
||||
# NB: These also contain SymInt because wrap_meta_outputs_with_default_device_logic
|
||||
# goes through this codepath. But it really should not LOL.
|
||||
# NB: size could potentially be None as you can override it and make it
|
||||
# throw an error, but we don't currently have any subclasses that do this
|
||||
# except C++ nested tensor but we're going to have nested int to make this
|
||||
# defined on NJT
|
||||
size: Tuple[int, ...]
|
||||
dynamo_dynamic_indices: List[int]
|
||||
|
||||
layout: torch.layout = torch.strided
|
||||
is_inference: bool = False
|
||||
is_leaf: bool = False
|
||||
requires_grad: bool = False
|
||||
is_sparse: bool = False
|
||||
is_mkldnn: bool = False
|
||||
is_functorch_wrapped: bool = False
|
||||
is_batchedtensor: bool = False
|
||||
is_legacy_batchedtensor: bool = False
|
||||
is_gradtrackingtensor: bool = False
|
||||
is_view: bool = False
|
||||
is_nested: bool = False
|
||||
is_traceable_wrapper_subclass: bool = False
|
||||
is_functional: bool = False
|
||||
is_conj: bool = False
|
||||
is_neg: bool = False
|
||||
stride: Optional[Tuple[int, ...]] = None
|
||||
storage_offset: int = 0
|
||||
# NB: We have a choice whether or not to store the id or a direct pointer
|
||||
# to the data structure. For ease of use, we store the data structure,
|
||||
# but this means that when we serialize, we have to swizzle these pointers
|
||||
# back into ids (so we have accurate aliasing relationships)
|
||||
storage: Optional[MetaStorageDesc] = None
|
||||
sparse_dim: Optional[int] = None # is_sparse, is_sparse_compressed
|
||||
dense_dim: Optional[int] = None # is_sparse, is_sparse_compressed
|
||||
@ -424,6 +476,19 @@ class MetaTensorDesc:
|
||||
grad: Optional[MetaTensorDesc] = None
|
||||
|
||||
# Everything below is NOT serializable, need some more work
|
||||
|
||||
_UNSERIALIZABLE: ClassVar[List[str]] = [
|
||||
"ctx",
|
||||
"type",
|
||||
"fake_mode",
|
||||
"view_func",
|
||||
"level",
|
||||
"current_level",
|
||||
"functorch_stack",
|
||||
"autograd_meta_from",
|
||||
"data",
|
||||
]
|
||||
|
||||
ctx: Optional[object] = None # is_traceable_wrapper_subclass
|
||||
type: Optional[Type] = None # is_traceable_wrapper_subclass
|
||||
fake_mode: Optional[FakeTensorMode] = None
|
||||
@ -459,6 +524,44 @@ class MetaTensorDesc:
|
||||
# entirely clear how to make it all lexical again, so we haven't done
|
||||
# it for now.
|
||||
|
||||
# NB: This will reference numeric IDs, and it is assumed that you've
|
||||
# already serialized everything this recursively references
|
||||
def as_json(self, describer_id):
|
||||
def json(k, v):
|
||||
# Some best-effort debugging serialization for unserializable
|
||||
# fields (feel free to add other special cases as appropriate)
|
||||
if k in ["data", "autograd_meta_from"]:
|
||||
return None # never repr these
|
||||
if k in set(self._UNSERIALIZABLE):
|
||||
return repr(v)
|
||||
if isinstance(v, (torch.device, torch.dtype, torch.layout)):
|
||||
return repr(v)
|
||||
if isinstance(v, torch.SymInt):
|
||||
return repr(v)
|
||||
if isinstance(v, (tuple, list)):
|
||||
return [json(k, v1) for v1 in v]
|
||||
if isinstance(v, (MetaStorageDesc, MetaTensorDesc)):
|
||||
return v.id
|
||||
if isinstance(v, CreationMeta):
|
||||
return str(v)
|
||||
if k == "attrs" and isinstance(v, dict):
|
||||
return {k1: v1.id for k1, v1 in v.items()}
|
||||
return v
|
||||
|
||||
r = {
|
||||
field.name: json(field.name, getattr(self, field.name))
|
||||
for field in dataclasses.fields(self)
|
||||
if not (
|
||||
getattr(self, field.name) is field.default
|
||||
or (
|
||||
field.name == "dynamo_dynamic_indices"
|
||||
and not getattr(self, field.name)
|
||||
)
|
||||
)
|
||||
}
|
||||
r.update({"describer_id": describer_id})
|
||||
return r
|
||||
|
||||
@property
|
||||
def shape(self):
|
||||
return self.size
|
||||
@ -887,9 +990,10 @@ class MetaConverter:
|
||||
|
||||
def tensor_visitor_fn(
|
||||
visited_t: torch.Tensor,
|
||||
# These arguments are never passed, we just use them to close
|
||||
# over these relevant values
|
||||
shape_env=shape_env,
|
||||
callback=callback,
|
||||
source=source,
|
||||
):
|
||||
# It's possible to close over an undefined tensor (e.g. NJT's lengths).
|
||||
if visited_t is None:
|
||||
@ -1443,6 +1547,10 @@ class MetaConverter:
|
||||
callback=lambda t: t(),
|
||||
source=None,
|
||||
symbolic_context=None,
|
||||
# Controls whether or not we should dump the tensor metadata to structured logs
|
||||
# when source is not None. Because we refakify after Dynamo is done,
|
||||
# we don't want to dump info again from AOTAutograd, it is redundant.
|
||||
trace=True,
|
||||
):
|
||||
# TODO: zero tensors? We appear to have eliminated them by
|
||||
# excluding complex for now
|
||||
@ -1475,9 +1583,22 @@ class MetaConverter:
|
||||
# non-Tensor types don't count as hit or miss
|
||||
return t
|
||||
|
||||
if source is None:
|
||||
trace = False
|
||||
|
||||
# Describe the tensor. NB: do NOT disable ambient modes, we may need
|
||||
# to query them when figuring out what to put in here
|
||||
t_desc = self.describer.describe_tensor(t)
|
||||
t_desc = self.describer.describe_tensor(t, trace=trace)
|
||||
|
||||
if trace:
|
||||
trace_structured(
|
||||
"describe_source",
|
||||
metadata_fn=lambda: {
|
||||
"describer_id": self.describer.id,
|
||||
"id": t_desc.id,
|
||||
"source": source.name(),
|
||||
},
|
||||
)
|
||||
|
||||
# Do the meta-fication. Here, we disable all the ambient modes, to
|
||||
# better simulate what would be like to re-fakeify from a fresh
|
||||
|
||||
Reference in New Issue
Block a user