lightning-thunder icon indicating copy to clipboard operation
lightning-thunder copied to clipboard

Use `torch._logging._internal.trace_structured_artifact` to save `TraceCtx`s when `TORCH_TRACE` is set

Open crcrpar opened this issue 2 weeks ago • 2 comments

What does this PR do?

Taking over #2182.

This enables us to dump TraceCtx by setting TORCH_TRACE env var if we use thunder.dynamo.compiler.thunderfx.

What are stored?

  • prologue, computation, and optional epilogue/backward traces from initial to execution.
  • compile data
  • split reasons

Ref: https://docs.pytorch.org/docs/stable/compile/programming_model.observability.html


I got the following as an artifact of TORCH_TRACE="./trace-tests/003" python thunder/benchmarks/benchmark_litgpt.py --compile "thunder_dynamo" --n_layers 2

Build products below:

[0/0]
-_0_0_0/dynamo_output_graph_0.txt (0)
-_0_0_0/graphmodule_after_splitter_1.txt (1)
-_0_0_0/dynamo_cpp_guards_str_2.txt (2)
-_0_0_0/compilation_metrics_3.html (3)
-_0_0_-/prologue_4.txt (4)
-_0_0_-/computation_5.txt (5)
-_0_0_-/epilogue_6.txt (6)
-_0_0_-/prologue_after_ExtractionOnlyPrologueTransform_7.txt (7)
-_0_0_-/computation_after_ExtractionOnlyPrologueTransform_8.txt (8)
-_0_0_-/epilogue_after_ExtractionOnlyPrologueTransform_9.txt (9)
-_0_0_-/computation_after_fwd_bwd_split_10.txt (10)
-_0_0_-/initial_backward_11.txt (11)
-_0_0_-/ex_backward_12.txt (12)
-_0_0_-/ex_computation_13.txt (13)

For TORCH_TRACE="./trace-tests/012" python thunder/benchmarks/benchmark_inference.py --mode thunder,

...
[4/0]
-_4_0_0/dynamo_output_graph_22.txt (22)
-_4_0_0/graphmodule_after_splitter_23.txt (23)
-_4_0_0/graph_split_reasons_24.json (24)
-_4_0_0/dynamo_cpp_guards_str_25.txt (25)
-_4_0_0/compilation_metrics_26.html (26)
-_4_0_-/prologue_27.txt (27)
-_4_0_-/computation_28.txt (28)
-_4_0_-/epilogue_29.txt (29)
-_4_0_-/prologue_after_ExtractionOnlyPrologueTransform_30.txt (30)
-_4_0_-/computation_after_ExtractionOnlyPrologueTransform_31.txt (31)
-_4_0_-/epilogue_after_ExtractionOnlyPrologueTransform_32.txt (32)
-_4_0_-/prologue_after_ExtractionOnlyPrologueTransform_33.txt (33)
-_4_0_-/computation_after_ExtractionOnlyPrologueTransform_34.txt (34)
-_4_0_-/epilogue_after_ExtractionOnlyPrologueTransform_35.txt (35)
-_4_0_-/ex_computation_36.txt (36)
-_4_0_-/compile_data_37.json (37)
-_4_0_-/prologue_46.txt (46)
-_4_0_-/computation_47.txt (47)
-_4_0_-/epilogue_48.txt (48)
-_4_0_-/prologue_after_ExtractionOnlyPrologueTransform_49.txt (49)
-_4_0_-/computation_after_ExtractionOnlyPrologueTransform_50.txt (50)
-_4_0_-/epilogue_after_ExtractionOnlyPrologueTransform_51.txt (51)
-_4_0_-/prologue_after_ExtractionOnlyPrologueTransform_52.txt (52)
-_4_0_-/computation_after_ExtractionOnlyPrologueTransform_53.txt (53)
-_4_0_-/epilogue_after_ExtractionOnlyPrologueTransform_54.txt (54)
-_4_0_-/ex_computation_55.txt (55)
-_4_0_-/compile_data_56.json (56)
-_4_0_-/prologue_68.txt (68)
-_4_0_-/computation_69.txt (69)
-_4_0_-/epilogue_70.txt (70)
-_4_0_-/prologue_after_ExtractionOnlyPrologueTransform_71.txt (71)
-_4_0_-/computation_after_ExtractionOnlyPrologueTransform_72.txt (72)
-_4_0_-/epilogue_after_ExtractionOnlyPrologueTransform_73.txt (73)
-_4_0_-/prologue_after_ExtractionOnlyPrologueTransform_74.txt (74)
-_4_0_-/computation_after_ExtractionOnlyPrologueTransform_75.txt (75)
-_4_0_-/epilogue_after_ExtractionOnlyPrologueTransform_76.txt (76)
-_4_0_-/ex_computation_77.txt (77)
-_4_0_-/compile_data_78.json (78)

crcrpar avatar Nov 18 '25 11:11 crcrpar