support graph-by-graph benchmarking for PyTorch native checkpointing

[kiya00]

Before submitting

• Was this discussed/approved via a Github issue? (no need for typos and docs improvements)
• Did you read the contributor guideline, Pull Request section?
• Did you make sure to update the docs?
• Did you write any new necessary tests?

What does this PR do?

The converter replaces the Torch operators in the checkpoint function with Thunder operators in-place, and also the compiled thunder/inductor module is replaced in-place, but the ThunderCompilerGraphBenchmarking/saving reproduction script needs the original GraphModule to compile/save.

Previously, the deepcopy of GraphModule is blocked by a pytorch/pytorch#139275. Thanks to @kshitij12345 for helping to fix it, we can use it to support the graph-by-graph benchmarking for PyTorch native checkpointing starting from Torch 2.6

Note:

• The code is also useful for saving the reproduction script ThunderFX: Save the reproducer script into files #1380
• The above mentioned deepcopy bug affects the test case test_thundercompiler_optim_step, so it's skipped
• Torch 2.6 changes the GraphModule structure, the checkpoint function becomes a submodule of the module containing tag_activation_checkpoint

# Torch 2.6
GraphModule(
  (submod_1): GraphModule(
    (wrap_body_0): GraphModule()
  )
)
# before 2.6
GraphModule(
  (wrap_body_0): GraphModule()
  (submod_1): GraphModule()
)

Before 2.6, in order to get the input tensor of submod_1 we need to calculate the wrap_body_0 ourselves (wrap_body_0 is a placeholder node in submod_1 module, and there's no example_value in node.meta); In 2.6, the wrap_body_0 is a get_attr node in submod_1 module, not an input. Since the latest Torch is much cleaner, we don't currently support benchmark checkpointing in Torch<2.6.

Fixes #1381.

[kiya00]

Hi @IvanYashchuk @kshitij12345 , could you review it?

[kshitij12345]

LGTM, thanks @kiya00

[t-vi]

Stamping Thank you @kiya00 @kshitij12345