[thunderfx] Integrate `trace_structured` and `trace_structured_artifact` into `ThunderCompiler` to use `TORCH_TRACE` and `tlparse` #2182

crcrpar · 2025-06-03T19:24:54Z

What does this PR do?

This enables ThunderFX to save traces such as execution computation and backward traces as artifacts of TORCH_TRACE.
By running scripts with TORCH_TRACE=/path/to/torch_trace_log_dir, without any changes in the scripts, we can check the traces as well as TorchDynamo output torch.fx.GraphModules and Inductor output, if fallback path is used.

This PR inserts trace_structured and trace_structured_artifact into thunderfx optimization path.

Ref:

In the following capture of an example of TORCH_TRACE, thunder_module_execution_computation_trc_12.txt has the execution forward trace.

Example:

$ TORCH_TRACE="torch_trace_test/" python thunder/benchmarks/benchmark_peft.py --model nvidia/Nemotron-Mini-4B-Instruct --compile thunder --max-steps 3 --fixed-num-hidden-layers 2 --trust-remote-code
# This command would open the generated HTML file on your brower
$ tlparse ./torch_trace_test/dedicated_log_torch_trace_<random string>_<random string>.log

crcrpar · 2025-08-12T20:06:34Z

thunder/__init__.py

+                        "encoding": "string",
+                    },
+                    payload_fn=lambda: f"{trace_to_store}\n",
+                    compile_id=compile_options.get("torch_compile_compile_id", None),


This implicit dependency on compile_options should be improved.

crcrpar · 2025-08-12T20:06:41Z

thunder/dynamo/compiler.py

@@ -114,8 +127,29 @@ def __call__(self, gm: torch.fx.GraphModule, sample_args: list[torch.SymInt, tor

        # The whole graph may not be supported by `thunder`, so we split it in `thunder` supported sections
        # and unsupported sections which are passed to `torch.compile(backend='inductor')`
-        split_module, subgraph_info = _splitter(gm, self._thunder_jit, self._torch_compile, sample_args)
+        split_module, subgraph_info = _splitter(gm, wrapped_thunder_jit, self._torch_compile, sample_args)


thunder/dynamo/splitter.py

crcrpar · 2025-08-12T20:18:49Z

This gist might be helpful https://gist.github.com/crcrpar/cfdc7dbb499cbe6b327f04be5856f078 about CompileID

Copilot

Pull Request Overview

This PR integrates structured tracing functionality from PyTorch's torch._logging._internal into Thunder's compiler and splitter components to enable better debugging and analysis through tlparse. The changes add trace artifacts at key compilation and splitting stages to provide visibility into Thunder's internal operations.

Key changes:

Added structured tracing to capture graph modules at various stages (original, post-checkpoint conversion, split graphs)
Added tracing for split reasons and unsupported nodes/contexts
Integrated Thunder execution traces (computation, prologue, epilogue) with torch compile ID tracking

Reviewed Changes

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

File	Description
thunder/dynamo/splitter.py	Adds trace artifacts for unsupported nodes/contexts and original/processed graph modules during splitting
thunder/dynamo/compiler.py	Adds tracing for original and split graphs, split reasons, and torch compile ID integration
thunder/init.py	Adds structured tracing for Thunder execution traces with compile ID support

_{Tip: Customize your code reviews with copilot-instructions.md. Create the file or learn how to get started.}

thunder/dynamo/splitter.py

thunder/__init__.py

thunder/dynamo/splitter.py

kshitij12345

It would be nice to have an example of end-to-end workflow in the PR description of how to use it with tlparse. It would be helpful for those (like me) who are not familiar with it.

kshitij12345 · 2025-08-13T13:54:46Z

thunder/__init__.py

@@ -428,6 +438,22 @@ def acquire_initial_trace(fn, args, kwargs, cd, cs, ad_hoc_executor):
            last_interpreter_log = jit_results.interpreter_log
            cs.last_interpreter_log = last_interpreter_log
            cs.last_interpreted_instructions = (i for i in last_interpreter_log if isinstance(i, dis.Instruction))
+
+            for name_in_artifact, trace_to_store in (


I think we should only invoke this logging function (_trace_structured) if logging is specified by the user. This will also prevent failures on main path if this internal API changes.

I agree. _log_to_torch_trace could be _maybe_log_to_torch_trace which incorporates the check on the user specifications.

Another thought:

def _log_to_torch_trace( string_format:str, trace_tuples: list[tuple], compile_id, ): for trc, *format_args in trace_tuples: if trc is None: continue name = string_format.format(*format_args) helper(name, trc, compile_id)

where helper is _log_to_torch_trace as defined below would allow these lines to be replaced with

trace_tuples = [ (computation_trc, "computation"), ... ] _log_to_torch_trace( "thunder_module_initial_{}_trc", trace_tuples, compile_id, )

Maybe this isn't so much shorter, but it is less indentation. But up to you @crcrpar, whichever you find more readable.

I'm not that inclined to update the helper to take a list of tuples of trace and name. That indentation feels fine to me

kshitij12345 · 2025-08-13T13:56:07Z

thunder/__init__.py

@@ -524,7 +550,23 @@ def apply_transforms_and_build_cache_entry(cd, cs, cache_info, prologue_trc, com
            if requires_grad:
                from thunder.transforms.autodiff import grad_transform_on_trace

+                _trace_structured(


It would be nice to have a helper function log_thunder_trace (or something like that). It would be easier to read and also to use.

kshitij12345 · 2025-08-13T13:56:56Z

thunder/dynamo/splitter.py

@@ -183,9 +207,30 @@ def is_thunder_supported_partition(node: torch.fx.Node) -> bool:
                partial(_get_example_inputs_from_placeholder, only_metadata=True), placeholders
            )
            example_input_metadatas.append(list(example_input_metadata))
+
+            trace_structured_artifact(


It would be nice to have a helper called log_fx_graph for readability and usage.

crcrpar · 2025-08-13T17:01:26Z

Oops, I noticed the current implementation only saves the last set of execution traces of a GraphModule. So if a graph module is chunked into thunder_0 -> inductor_1 -> thunder_2, then only the traces for thunder_2 seem to be saved.

t-vi · 2025-08-27T08:15:50Z

@kshitij12345 @kiya00 @beverlylytle @riccardofelluga @IvanYashchuk any of you wanting to review this? (or someone else)

so that we can use `TORCH_TRACE` and tlparse even for thunderfx Reference: - https://docs.pytorch.org/docs/stable/torch.compiler_troubleshooting.html#tlparse-torch-trace - [torch._logging._internal.trace_structured](https://github.com/pytorch/pytorch/blob/6f7694f18f6cbfc684fb44cc2f4f3a06218d919e/torch/_logging/_internal.py#L1209) - [torch._logging._internal.trace_structured_artifact](https://github.com/pytorch/pytorch/blob/6f7694f18f6cbfc684fb44cc2f4f3a06218d919e/torch/_logging/_internal.py#L1194) Signed-off-by: Masaki Kozuki <[email protected]>

as the former would be more Python-like. Signed-off-by: Masaki Kozuki <[email protected]>

…`tlparse` Signed-off-by: Masaki Kozuki <[email protected]>

also, propagating `CompileID` of torch.compile from `ThunderCompiler.__call__` to `thunder.jit` as compile id does not seem to be available when saving those artifacts Signed-off-by: Masaki Kozuki <[email protected]>

for more information, see https://pre-commit.ci

Signed-off-by: Masaki Kozuki <[email protected]>

Co-authored-by: Copilot <[email protected]>

for more information, see https://pre-commit.ci

Signed-off-by: Masaki Kozuki <[email protected]>

crcrpar requested review from mruberry, lantiga and t-vi as code owners June 3, 2025 19:24

crcrpar added the thunderfx for things that could be applicable to the dynamo+thunder frontend label Jun 3, 2025

This comment was marked as outdated.

Sign in to view

crcrpar force-pushed the logging-for-torch_trace-tlparse branch from 84fb0db to 05108cc Compare June 4, 2025 12:40

crcrpar marked this pull request as draft June 13, 2025 09:26

crcrpar force-pushed the logging-for-torch_trace-tlparse branch from 05108cc to 6d8b13c Compare August 12, 2025 18:39

This comment was marked as resolved.

Sign in to view

crcrpar commented Aug 12, 2025

View reviewed changes

crcrpar requested a review from Copilot August 12, 2025 20:20

crcrpar changed the title ~~Integrate trace_structured and trace_structured_artifact into ThunderCompiler~~ Integrate trace_structured and trace_structured_artifact into ThunderCompiler to use TORCH_TRACE and tlparse Aug 12, 2025

Copilot AI reviewed Aug 12, 2025

View reviewed changes

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

thunder/__init__.py Outdated Show resolved Hide resolved

thunder/__init__.py Outdated Show resolved Hide resolved

crcrpar commented Aug 12, 2025

View reviewed changes

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

crcrpar marked this pull request as ready for review August 13, 2025 07:05

kshitij12345 reviewed Aug 13, 2025

View reviewed changes

crcrpar changed the title ~~Integrate trace_structured and trace_structured_artifact into ThunderCompiler to use TORCH_TRACE and tlparse~~ [thunderfx] Integrate trace_structured and trace_structured_artifact into ThunderCompiler to use TORCH_TRACE and tlparse Aug 13, 2025

crcrpar force-pushed the logging-for-torch_trace-tlparse branch from 2bd6c36 to ec7fb0c Compare August 17, 2025 16:24

crcrpar and others added 9 commits September 3, 2025 13:17

use GraphModule.print_readable not str(GraphModule.graph)

c34db96

as the former would be more Python-like. Signed-off-by: Masaki Kozuki <[email protected]>

(ab)use trace_structured_artifact to make split reasons handled by …

8f11a06

…`tlparse` Signed-off-by: Masaki Kozuki <[email protected]>

Store execution prologue, computation, and epilogue traces

04f64b8

also, propagating `CompileID` of torch.compile from `ThunderCompiler.__call__` to `thunder.jit` as compile id does not seem to be available when saving those artifacts Signed-off-by: Masaki Kozuki <[email protected]>

[pre-commit.ci] auto fixes from pre-commit.com hooks

7441d67

for more information, see https://pre-commit.ci

use trace_structured to specify compile_id

a140a17

Signed-off-by: Masaki Kozuki <[email protected]>

removing trace_structured as it requires tlparse customization

ddced44

Signed-off-by: Masaki Kozuki <[email protected]>

remove trace_structured

8e26859

Signed-off-by: Masaki Kozuki <[email protected]>

clean up following trace_structured removal

032e109

crcrpar and others added 16 commits September 3, 2025 13:17

fix name of GraphModules inside splitter

2c98d8a

Signed-off-by: Masaki Kozuki <[email protected]>

Apply suggestions from code review of copilot

d3b5378

Co-authored-by: Copilot <[email protected]>

[pre-commit.ci] auto fixes from pre-commit.com hooks

e6209e3

for more information, see https://pre-commit.ci

deduplicate metadata_fn

5a9f1e5

Signed-off-by: Masaki Kozuki <[email protected]>

avoid redefinition of trace

c2eb191

Signed-off-by: Masaki Kozuki <[email protected]>

store backward_trc if available

b722c0d

Signed-off-by: Masaki Kozuki <[email protected]>

store trace before/after grad_transform

23dd2bb

Signed-off-by: Masaki Kozuki <[email protected]>

check compile_id arg

2482974

Signed-off-by: Masaki Kozuki <[email protected]>

add todo comment

5280335

Signed-off-by: Masaki Kozuki <[email protected]>

remove compile_id from _trace_structured callsite

6e1689a

Signed-off-by: Masaki Kozuki <[email protected]>

fix

2bb659e

Signed-off-by: Masaki Kozuki <[email protected]>

[no ci] todo: use node index

c64b83e

Signed-off-by: Masaki Kozuki <[email protected]>

propagate chunk index of GraphModule

ac7b851

Signed-off-by: Masaki Kozuki <[email protected]>

wrap trace_structured for trace/graph-module

aa87d1e

Signed-off-by: Masaki Kozuki <[email protected]>

fix typos

59765a6

Signed-off-by: Masaki Kozuki <[email protected]>

clean up

b48e9fe

Signed-off-by: Masaki Kozuki <[email protected]>

crcrpar force-pushed the logging-for-torch_trace-tlparse branch from acf6999 to b48e9fe Compare September 3, 2025 04:18

[thunderfx] Integrate trace_structured and trace_structured_artifact into ThunderCompiler to use TORCH_TRACE and tlparse #2182

Are you sure you want to change the base?

[thunderfx] Integrate trace_structured and trace_structured_artifact into ThunderCompiler to use TORCH_TRACE and tlparse #2182

Conversation

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

Uh oh!

What does this PR do?

Uh oh!

This comment was marked as outdated.

This comment was marked as resolved.

Uh oh!

crcrpar Aug 12, 2025

Choose a reason for hiding this comment

Uh oh!

crcrpar Aug 12, 2025

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Uh oh!

crcrpar commented Aug 12, 2025

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!

kshitij12345 left a comment • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

kshitij12345 Aug 13, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

beverlylytle Aug 27, 2025

Choose a reason for hiding this comment

Uh oh!

crcrpar Sep 2, 2025

Choose a reason for hiding this comment

Uh oh!

kshitij12345 Aug 13, 2025

Choose a reason for hiding this comment

Uh oh!

kshitij12345 Aug 13, 2025

Choose a reason for hiding this comment

Uh oh!

crcrpar commented Aug 13, 2025

Uh oh!

t-vi commented Aug 27, 2025

Uh oh!

Uh oh!

[thunderfx] Integrate `trace_structured` and `trace_structured_artifact` into `ThunderCompiler` to use `TORCH_TRACE` and `tlparse` #2182

[thunderfx] Integrate `trace_structured` and `trace_structured_artifact` into `ThunderCompiler` to use `TORCH_TRACE` and `tlparse` #2182

crcrpar commented Jun 3, 2025 •

edited

Loading

kshitij12345 left a comment •

edited

Loading

kshitij12345 Aug 13, 2025 •

edited

Loading