gistfile1.txt

import torch
from torch.utils.checkpoint import checkpoint


class MyModule(torch.nn.Module):
    def __init__(self):
        super().__init__()
        self.linear = torch.nn.Linear(10, 10)

    def _forward_helper(self, x):
        a = torch.sin(x)
        b = self.linear(a)
        torch._dynamo.graph_break()
        c = torch.cos(b)
        return c

    def forward(self, x):
        return checkpoint(self._forward_helper, x, use_reentrant=False)


x = torch.randn(10, 10)
mod = MyModule()
ref = mod(x)

opt_mod = torch.compile(mod, backend="eager")
res = opt_mod(x)

(pytorch)  anijain@devgpu002 ~/local/pytorch (main) $ TORCH_LOGS="dynamo,graph_breaks" python examples/ac.py
[2024-02-06 11:36:07,278] [0/0] torch._dynamo.symbolic_convert: [INFO] Step 1: torchdynamo start tracing forward /data/users/anijain/pytorch/examples/ac.py:17
[2024-02-06 11:36:07,278] [0/0] torch.fx.experimental.symbolic_shapes: [INFO] create_env
[2024-02-06 11:36:07,603] [0/0] torch._dynamo.variables.higher_order_ops: [INFO] speculate_subgraph: while introspecting torch.utils.checkpoint.checkpoint, we were unable to trace function `_forward_helper` into a single graph. This means that Dynamo was unable to prove safety for this API and will fall back to eager-mode PyTorch, which could lead to a slowdown.
[2024-02-06 11:36:07,604] [0/0] torch._dynamo.variables.higher_order_ops: [INFO] 'skip function graph_break in file /data/users/anijain/pytorch/torch/_dynamo/decorators.py'
[2024-02-06 11:36:07,604] [0/0] torch._dynamo.symbolic_convert.__graph_breaks: [DEBUG] Graph break: 'skip function graph_break in file /data/users/anijain/pytorch/torch/_dynamo/decorators.py' from user code at:
[2024-02-06 11:36:07,604] [0/0] torch._dynamo.symbolic_convert.__graph_breaks: [DEBUG]   File "/data/users/anijain/pytorch/examples/ac.py", line 18, in forward
[2024-02-06 11:36:07,604] [0/0] torch._dynamo.symbolic_convert.__graph_breaks: [DEBUG]     return checkpoint(self._forward_helper, x, use_reentrant=False)
[2024-02-06 11:36:07,604] [0/0] torch._dynamo.symbolic_convert.__graph_breaks: [DEBUG]   File "/data/users/anijain/pytorch/examples/ac.py", line 13, in _forward_helper
[2024-02-06 11:36:07,604] [0/0] torch._dynamo.symbolic_convert.__graph_breaks: [DEBUG]     torch._dynamo.graph_break()
[2024-02-06 11:36:07,604] [0/0] torch._dynamo.symbolic_convert.__graph_breaks: [DEBUG]
[2024-02-06 11:36:07,604] [0/0] torch._dynamo.convert_frame: [INFO] Restarting analysis due to _dynamo/symbolic_convert.py:141 in fail_and_restart_analysis
[2024-02-06 11:36:07,605] [0/0_1] torch._dynamo.symbolic_convert: [INFO] Step 1: torchdynamo start tracing forward /data/users/anijain/pytorch/examples/ac.py:17
[2024-02-06 11:36:07,605] [0/0_1] torch.fx.experimental.symbolic_shapes: [INFO] create_env
[2024-02-06 11:36:07,609] [0/0_1] torch.fx.experimental.symbolic_shapes: [INFO] produce_guards
[2024-02-06 11:36:07,617] torch._dynamo.utils: [INFO] TorchDynamo compilation metrics:
[2024-02-06 11:36:07,617] torch._dynamo.utils: [INFO] Function, Runtimes (s)
[2024-02-06 11:36:07,617] torch._dynamo.utils: [INFO] _compile.<locals>.compile_inner, 1.8634
(pytorch)  anijain@devgpu002 ~/local/pytorch (main) $