Use `torch._logging._internal.trace_structured_artifact` to save `TraceCtx`s when `TORCH_TRACE` is set by crcrpar · Pull Request #2751 · Lightning-AI/lightning-thunder

crcrpar · 2025-11-18T11:55:28Z

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)

Copilot

Pull Request Overview

This PR enables dumping of TraceCtx objects using PyTorch's torch._logging._internal.trace_structured_artifact API when the TORCH_TRACE environment variable is set. This provides enhanced observability for Thunder's compilation process when using thunder.dynamo.compiler.thunderfx.

Adds support for logging compilation traces and graph modules to TORCH_TRACE artifacts
Implements backward compatibility wrapper for trace_structured_artifact to handle different PyTorch versions
Extends test coverage to validate the TORCH_TRACE functionality

Reviewed Changes

Copilot reviewed 4 out of 4 changed files in this pull request and generated 4 comments.

File	Description
thunder/dynamo/utils.py	Adds `log_trace_or_graphmodule_to_torch_trace` helper function and backward compatibility wrapper for `trace_structured_artifact`
thunder/dynamo/compiler.py	Integrates trace logging for graph modules and split reasons in the ThunderCompiler
thunder/init.py	Adds trace logging at various compilation stages (prologue, computation, epilogue, transforms, backward)
thunder/tests/test_dynamo.py	Adds parametrized test to validate functionality with and without TORCH_TRACE enabled

Comments suppressed due to low confidence (1)

thunder/dynamo/compiler.py:144

This assignment to 'thunder_options' is unnecessary as it is redefined before this value is used.

        thunder_options = _with_prologue_pruning_transform(

💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.

thunder/dynamo/compiler.py

thunder/dynamo/utils.py

thunder/tests/test_dynamo.py

Copilot

Pull Request Overview

Copilot reviewed 5 out of 5 changed files in this pull request and generated 2 comments.

Comments suppressed due to low confidence (1)

thunder/dynamo/compiler.py:144

This assignment to 'thunder_options' is unnecessary as it is redefined before this value is used.

        thunder_options = _with_prologue_pruning_transform(

💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.

thunder/dynamo/utils.py

Copilot

Pull Request Overview

Copilot reviewed 5 out of 5 changed files in this pull request and generated 1 comment.

💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.

thunder/tests/test_dynamo.py

shino16

It's a very cool feature!

thunder/dynamo/utils.py

thunder/tests/test_dynamo.py

thunder/dynamo/compiler.py

Signed-off-by: Masaki Kozuki <mkozuki@nvidia.com>

so that the fallback path can also correctly work with torch-trace Signed-off-by: Masaki Kozuki <mkozuki@nvidia.com>

Co-authored-by: Copilot <175728472+Copilot@users.noreply.github.com>

Signed-off-by: Masaki Kozuki <mkozuki@nvidia.com>

Co-authored-by: Masato Shinokawa <40788005+shino16@users.noreply.github.com>

Signed-off-by: Masaki Kozuki <mkozuki@nvidia.com>

crcrpar · 2025-11-25T09:35:53Z

@shino16 @kshitij12345 could you give this another look?

shino16

LGTM!

kshitij12345

LGTM, thanks @crcrpar

crcrpar · 2025-11-26T01:38:48Z

@KaelanDt could you review this?

crcrpar requested review from KaelanDt, lantiga and mruberry as code owners November 18, 2025 11:55

crcrpar requested a review from Copilot November 18, 2025 11:56

Copilot started reviewing on behalf of crcrpar November 18, 2025 11:56 View session

Copilot finished reviewing on behalf of crcrpar November 18, 2025 12:00

Copilot AI reviewed Nov 18, 2025

View reviewed changes

thunder/dynamo/compiler.py Show resolved Hide resolved

thunder/dynamo/compiler.py Outdated Show resolved Hide resolved

thunder/dynamo/utils.py Outdated Show resolved Hide resolved

thunder/tests/test_dynamo.py Outdated Show resolved Hide resolved

crcrpar requested a review from Copilot November 18, 2025 12:11

Copilot started reviewing on behalf of crcrpar November 18, 2025 12:11 View session

Copilot finished reviewing on behalf of crcrpar November 18, 2025 12:15

Copilot AI reviewed Nov 18, 2025

View reviewed changes

thunder/dynamo/utils.py Show resolved Hide resolved

thunder/dynamo/utils.py Show resolved Hide resolved

crcrpar requested a review from Copilot November 18, 2025 12:17

Copilot started reviewing on behalf of crcrpar November 18, 2025 12:17 View session

Copilot finished reviewing on behalf of crcrpar November 18, 2025 12:20

Copilot AI reviewed Nov 18, 2025

View reviewed changes

thunder/tests/test_dynamo.py Outdated Show resolved Hide resolved

crcrpar requested review from kshitij12345 and shino16 November 18, 2025 14:48

shino16 reviewed Nov 18, 2025

View reviewed changes

thunder/dynamo/utils.py Show resolved Hide resolved

thunder/dynamo/utils.py Outdated Show resolved Hide resolved

thunder/tests/test_dynamo.py Outdated Show resolved Hide resolved

thunder/tests/test_dynamo.py Outdated Show resolved Hide resolved

shino16 reviewed Nov 18, 2025

View reviewed changes

thunder/dynamo/compiler.py Show resolved Hide resolved

crcrpar requested a review from shino16 November 24, 2025 23:45

crcrpar added 10 commits November 25, 2025 09:54

use trace_structured_artifact

0c14911

Signed-off-by: Masaki Kozuki <mkozuki@nvidia.com>

check compile_id in signature.parameters

c0ed663

Signed-off-by: Masaki Kozuki <mkozuki@nvidia.com>

fix typo of compile_id to torch_compile_compile_id

263e956

Signed-off-by: Masaki Kozuki <mkozuki@nvidia.com>

import GraphModule for isinstance to work

b3875a3

Signed-off-by: Masaki Kozuki <mkozuki@nvidia.com>

fix payload_fn

c01abda

Signed-off-by: Masaki Kozuki <mkozuki@nvidia.com>

fix test: remove pytest fixtures

68fff7b

Signed-off-by: Masaki Kozuki <mkozuki@nvidia.com>

set encoding

01008dd

Signed-off-by: Masaki Kozuki <mkozuki@nvidia.com>

fix cond

10ff345

Signed-off-by: Masaki Kozuki <mkozuki@nvidia.com>

fix key of CompileId

56d1681

Signed-off-by: Masaki Kozuki <mkozuki@nvidia.com>

remove debug print

618ac4d

Signed-off-by: Masaki Kozuki <mkozuki@nvidia.com>

crcrpar and others added 9 commits November 25, 2025 09:54

allow None as m

5da8040

Signed-off-by: Masaki Kozuki <mkozuki@nvidia.com>

use overload for simplicity

ce49291

Signed-off-by: Masaki Kozuki <mkozuki@nvidia.com>

store compiledata

58b0d92

Signed-off-by: Masaki Kozuki <mkozuki@nvidia.com>

propagate compile_id to LazyInductorModule

95bd7b9

so that the fallback path can also correctly work with torch-trace Signed-off-by: Masaki Kozuki <mkozuki@nvidia.com>

Apply suggestions from code review

498486b

Co-authored-by: Copilot <175728472+Copilot@users.noreply.github.com>

import cleanup

3b8fb57

Signed-off-by: Masaki Kozuki <mkozuki@nvidia.com>

Update thunder/tests/test_dynamo.py

802a837

Co-authored-by: Masato Shinokawa <40788005+shino16@users.noreply.github.com>

check torch version instead of signature of trace_structured_artifact

1b3b348

Signed-off-by: Masaki Kozuki <mkozuki@nvidia.com>

prefix

3eeb517

Signed-off-by: Masaki Kozuki <mkozuki@nvidia.com>

crcrpar force-pushed the crpa/trace_structured branch from eb5bf55 to 3eeb517 Compare November 25, 2025 00:55

register compile_id to LazyInductorModule

d93e9ce

Signed-off-by: Masaki Kozuki <mkozuki@nvidia.com>

shino16 approved these changes Nov 25, 2025

View reviewed changes

kshitij12345 approved these changes Nov 25, 2025

View reviewed changes

Conversation

crcrpar commented Nov 18, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

What does this PR do?

Uh oh!

Copilot AI left a comment

Choose a reason for hiding this comment

Pull Request Overview

Reviewed Changes

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Copilot AI left a comment

Choose a reason for hiding this comment

Pull Request Overview

Uh oh!

Uh oh!

Uh oh!

Copilot AI left a comment

Choose a reason for hiding this comment

Pull Request Overview

Uh oh!

Uh oh!

shino16 left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

crcrpar commented Nov 25, 2025

Uh oh!

shino16 left a comment

Choose a reason for hiding this comment

Uh oh!

kshitij12345 left a comment

Choose a reason for hiding this comment

Uh oh!

crcrpar commented Nov 26, 2025

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants

crcrpar commented Nov 18, 2025 •

edited

Loading