Activation checkpoint not working inside Inductor-compiled submodules

[shino16]

🐛 Bug

When ThunderFX delegates activation checkpointing operations to Inductor-compiled submodules, apparently the checkpointing mechanism fails to execute properly, causing excessive memory usage due to retained intermediate activations. This may apply to all models that containing some Thunder-unsupported ops in activation checkpointed regions, causing e.g. #2501.

Note. When Thunder encounters an unsupported op in a checkpointed region, the entire region will go to Inductor. Huge memory usage in status quo could be a good sign of Thunder's poor coverage.

Code sample

import torch
from torch.utils.checkpoint import checkpoint
from thunder.dynamo import thunderfx

def fn(x):
    return x.sinc().sinc().sinc()

def checkpoint_fn(x):
    return checkpoint(fn, x, use_reentrant=False)

for compiler in [thunderfx, torch.compile]:
    torch.cuda.reset_peak_memory_stats()
    assert torch.cuda.memory_allocated() == 0

    x = torch.randn((1024, 1024, 1024), device="cuda", requires_grad=True)  # 4 GB
    y = compiler(checkpoint_fn)(x)
    del x, y

    peak_mem_usage = torch.cuda.max_memory_allocated()
    print(f"{compiler.__name__}: {peak_mem_usage / 1024 / 1024 / 1024} GB")

thunderfx: 16.0 GB
compile: 8.0 GB

Input for thunder.dynamo.ThunderCompiler:

class GraphModule(torch.nn.Module):
    def forward(self, L_x_: "f32[1024, 1024, 1024]"):
        l_x_ = L_x_
        
         # File: /opt/pytorch/lightning-thunder/tmp/repro.py:9 in checkpoint_fn, code: return checkpoint(fn, x, use_reentrant=False)
        wrap_body_0 = self.wrap_body_0
        tag_activation_checkpoint = torch.ops.higher_order.tag_activation_checkpoint(wrap_body_0, l_x_, use_reentrant = False);  wrap_body_0 = l_x_ = None
        getitem: "f32[1024, 1024, 1024]" = tag_activation_checkpoint[0];  tag_activation_checkpoint = None
        return (getitem,)
        
    class wrap_body_0(torch.nn.Module):
        def forward(self, l_x_: "f32[1024, 1024, 1024]"):
             # File: /opt/pytorch/lightning-thunder/tmp/repro.py:6 in fn, code: return x.sinc().sinc().sinc()
            sinc: "f32[1024, 1024, 1024]" = l_x_.sinc();  l_x_ = None
            sinc_1: "f32[1024, 1024, 1024]" = sinc.sinc();  sinc = None
            sinc_2: "f32[1024, 1024, 1024]" = sinc_1.sinc();  sinc_1 = None
            return (sinc_2,)

Output:

class GraphModule(torch.nn.Module):
    def forward(self, l_x_: "f32[1024, 1024, 1024]"):
        # No stacktrace found for following nodes
        inductor_0 = self.inductor_0(l_x_);  l_x_ = None
        thunder_1 = self.thunder_1(inductor_0);  inductor_0 = None
        return (thunder_1,)
        
    class inductor_0(torch.nn.Module):
        def forward(self, l_x_: "f32[1024, 1024, 1024]"):
             # File: /opt/pytorch/lightning-thunder/tmp/repro.py:9 in checkpoint_fn, code: return checkpoint(fn, x, use_reentrant=False)
            wrap_body_0 = self.wrap_body_0
            tag_activation_checkpoint = torch.ops.higher_order.tag_activation_checkpoint(wrap_body_0, l_x_, use_reentrant = False);  wrap_body_0 = l_x_ = None
            return tag_activation_checkpoint
            
        class _orig_mod(torch.nn.Module):
            def forward(self, l_x_: "f32[1024, 1024, 1024]"):
                 # File: /opt/pytorch/lightning-thunder/tmp/repro.py:9 in checkpoint_fn, code: return checkpoint(fn, x, use_reentrant=False)
                wrap_body_0 = self.wrap_body_0
                tag_activation_checkpoint = torch.ops.higher_order.tag_activation_checkpoint(wrap_body_0, l_x_, use_reentrant = False);  wrap_body_0 = l_x_ = None
                return tag_activation_checkpoint
                
            class wrap_body_0(torch.nn.Module):
                def forward(self, l_x_: "f32[1024, 1024, 1024]"):
                     # File: /opt/pytorch/lightning-thunder/tmp/repro.py:6 in fn, code: return x.sinc().sinc().sinc()
                    sinc: "f32[1024, 1024, 1024]" = l_x_.sinc();  l_x_ = None
                    sinc_1: "f32[1024, 1024, 1024]" = sinc.sinc();  sinc = None
                    sinc_2: "f32[1024, 1024, 1024]" = sinc_1.sinc();  sinc_1 = None
                    return (sinc_2,)
                    
    class thunder_1(torch.nn.Module):
        def forward(self, tag_activation_checkpoint):
             # File: /opt/pytorch/lightning-thunder/tmp/repro.py:9 in checkpoint_fn, code: return checkpoint(fn, x, use_reentrant=False)
            getitem: "f32[1024, 1024, 1024]" = tag_activation_checkpoint[0];  tag_activation_checkpoint = None
            return getitem
            
        class _model(torch.nn.Module):
            def forward(self, tag_activation_checkpoint):
                 # File: /opt/pytorch/lightning-thunder/tmp/repro.py:9 in checkpoint_fn, code: return checkpoint(fn, x, use_reentrant=False)
                getitem: "f32[1024, 1024, 1024]" = tag_activation_checkpoint[0];  tag_activation_checkpoint = None
                return getitem

[shino16]

The GraphModule that ThunderCompiler receives already fails to apply activation checkpointing. So this bug is more on the PyTorch side?

def trivial_backend(gm, sample_args):
    return gm
    return torch.compile(gm)  # Neither works

x = torch.randn((1024, 1024, 1024), device="cuda", requires_grad=True)
y = torch.compile(checkpoint_fn, backend=trivial_backend)(x)
peak_mem_usage = torch.cuda.max_memory_allocated()
print(f"{peak_mem_usage / 1024 / 1024 / 1024} GB")

16.0 GB

[shino16]

So gm with torch.ops.higher_order.tag_activation_checkpoint needs to be handled by PyTorch's backend to perform activation checkpointing. However, when we pass the gm to torch.compile again, it appears to skip the backend entirely.

def compile_with_bad_backend(gm, sample_args):
    def bad_backend(gm, sample_args):
        raise RuntimeError("Backend was called!")
    return torch.compile(gm, backend=bad_backend)

x = torch.randn((1024, 1024, 1024), device="cuda", requires_grad=True)
torch.compile(checkpoint_fn, backend=compile_with_bad_backend)(x)  # No RuntimeError raised

The same issue occurs with ThunderFX. I tested using bad_backend for compiling submodules and observed no runtime error, confirming that the backend is being bypassed.

Since there appears to be no reliable way to ensure Inductor handles torch.ops.higher_order.tag_activation_checkpoint properly in a submodule, we should replace these operators with torch.utils.checkpoint.checkpoint calls that can be executed directly.

[kshitij12345]

Thanks @shino16 for the detailed debug log.

The GraphModule that ThunderCompiler receives already fails to apply activation checkpointing. So this bug is more on the PyTorch side?

It seems that using torch.compile in the backend to torch.compile(fn, backend=backend) doesn't actually use the inductor. This is a surprise as I thought inductor will be used in this case. This affects the thunderfx path in general. Based on the new understanding, for fallback, we don't really use inductor (but probably eager). We should file a separate issue to investigate this behavior and also report to PyTorch if we find this behaviour to be a bug. (cc: @IvanYashchuk for visibility)

This can be verified with the following script and looking at the kernels in the nsys profile.

import torch
from thunder.dynamo import thunderfx

def fn(x):
    # `sinc` is unsupported by thunder and uses fallback in thunderfx.
    return x.sinc().sinc().sinc()

x = torch.randn((1024), device="cuda", requires_grad=True)

torch.cuda.nvtx.range_push("Eager")
fn(x)
torch.cuda.nvtx.range_pop()

torch.cuda.nvtx.range_push("thunderfx")
thunderfx(fn)(x)
torch.cuda.nvtx.range_pop()

torch.cuda.nvtx.range_push("torch.compile")
torch.compile(fn)(x)
torch.cuda.nvtx.range_pop()

Eager Kernels

thunderfx Kernels (same as eager)

torch.compile Kernels (triton kernels)

---

For the snippet in comment above, updating the trivial backend to use compile_fx from inductor seems to do the correct thing.

import torch
from torch.utils.checkpoint import checkpoint

def fn(x):
    return x.sinc().sinc().sinc()

def checkpoint_fn(x):
    return checkpoint(fn, x, use_reentrant=False)


torch.cuda.reset_peak_memory_stats()
assert torch.cuda.memory_allocated() == 0

x = torch.randn((1024, 1024, 1024), device="cuda", requires_grad=True)  # 4 GB

def trivial_backend(gm, sample_args):
    from torch._inductor.compile_fx import compile_fx
    return compile_fx(gm, sample_args)

x = torch.randn((1024, 1024, 1024), device="cuda", requires_grad=True)
y = torch.compile(checkpoint_fn, backend=trivial_backend)(x)
peak_mem_usage = torch.cuda.max_memory_allocated()
print(f"{peak_mem_usage / 1024 / 1024 / 1024} GB")  # 8 GB

[shino16]

Thank you @kshitij12345 for your great analysis. I will now make a separate issue quickly.

As a side note, I made a quick workaround specifically for activation checkpointing in #2538. This has fixed #2501, and will perhaps fix the memory spikes we have seen in some models.

[mattteochen]

From some initial investigation, it seems like passing a graphfx to torch.compile has an issue in lowering non-reentrant checkpoint calls. Wrapping the fallback graphfx as a nn.Module (very lightweight) and compiling it with torch.compile will trigger torch dynamo, autograd etc... on a fresh nn.Module. This leads to the expected memory consumption.

I assume passing a graphfx rather than a Python callable to torch.compile may lead to different behaviors (backend).

Currently the torch_inductor Callables we pass to splitter.py are all pointing to torch.compile:

• lightning-thunder/thunder/dynamo/compiler.py

  Line 151 in 5ece7ef

  self._torch_compile,

• lightning-thunder/thunder/dynamo/report.py

  Line 1110 in 5ece7ef

  _, subgraph_info = _splitter(gm, thunder_jit, torch.compile, _unused_sample_args=None)

import gc
import torch
import torch.nn as nn
from torch.utils.checkpoint import checkpoint
from thunder.dynamo import thunderfx
from thunder.dynamo.utils import CompilerType, get_or_create_example_inputs_from_placeholders

def inner(x):
    return x.sinc().sinc().sinc()

def checkpoint_fn(x):
    return checkpoint(inner, x, use_reentrant=False)

def iso_peak(run, *args):
    gc.collect(); torch.cuda.empty_cache(); torch.cuda.synchronize()
    torch.cuda.reset_peak_memory_stats()
    mem_before = torch.cuda.memory_allocated() / (1024**3)
    out = run(*args)
    mem_after = torch.cuda.memory_allocated() / (1024**3)
    gc.collect(); torch.cuda.synchronize()
    print(f"  Memory: before={mem_before:.2f}GB, after={mem_after:.2f}GB, diff={mem_after - mem_before:.2f}GB")

def main():
    device = "cuda"

    # 1) Compile with ThunderFX and run
    x_big = torch.randn((1024, 1024, 1024), device=device, requires_grad=True)  # ~4 GB
    cf = thunderfx(checkpoint_fn)
    print('ThunderFX:')
    iso_peak(cf, x_big)

    # 2) Find an FX subgraph that fell back to Inductor (TORCH_INDUCTOR)
    gm_fallback = None
    for sinfo in cf._backend.subgraph_infos:
        for subgm, compiled in sinfo.submodule_to_compiled_functions.items():
            if compiled.compiler == CompilerType.TORCH_INDUCTOR:
                if gm_fallback is None:
                    gm_fallback = subgm
    assert gm_fallback is not None, "No Inductor fallback subgraph found."

    # 3) Compile with GraphFX and run
    c_fx_plain = torch.compile(gm_fallback)  # no AOT lowering
    print('GraphFX')
    iso_peak(c_fx_plain, x_big)

    # 4) Wrap GraphFX in nn.Module and run
    class Wrapped(nn.Module):
        def __init__(self, gm): super().__init__(); self.gm = gm
        def forward(self, *a): return self.gm(*a)

    wrapped = Wrapped(gm_fallback)
    c_fx_wrapped = torch.compile(wrapped)
    print('nn.Module:')
    iso_peak(c_fx_wrapped, x_big)

    # 5) Compile with Torch compile and run
    c_fx_torchcompile = torch.compile(checkpoint_fn)
    print('Torch compile:')
    iso_peak(c_fx_torchcompile, x_big)

if __name__ == "__main__":
    main()

Output:

ThunderFX:
  Memory: before=4.00GB, after=16.00GB, diff=12.00GB
GraphFX
  Memory: before=4.00GB, after=16.00GB, diff=12.00GB
nn.Module:
  Memory: before=4.00GB, after=8.00GB, diff=4.00GB
Torch compile:
  Memory: before=4.00GB, after=8.00GB, diff=4.00GB

[kshitij12345]

Thanks for investigating further @mattteochen

compiling it with torch.compile will trigger torch dynamo, autograd etc... on a fresh nn.Module.

I am surprised as fx.GraphModule inherits from torch.nn.Module (ref). I wonder what does torch.compile do differently when it sees a GraphModule vs when it is wrapped in a torch.nn.Module (but still has to interpret the same GraphModule).

[mattteochen]

Thanks for investigating further @mattteochen

compiling it with torch.compile will trigger torch dynamo, autograd etc... on a fresh nn.Module.

I am surprised as fx.GraphModule inherits from torch.nn.Module (ref). I wonder what does torch.compile do differently when it sees a GraphModule vs when it is wrapped in a torch.nn.Module (but still has to interpret the same GraphModule).

I see, this behaviour is indeed quite strange. Digging further, one thing that I can see is that when we pass a GraphModule to torch compile, the compilation process stops very early. Wrapping with either a nn.Module or a proxy function triggers the compilation process. It's not what we might want".

Proxy function:

gm_fallback = ...
def f(x):
    return gm_fallback(x)

It's still unclear to me why this happened. One reference I've found is here, but I'm still not sure if this is involved or not.

For reference:

(Test case)

def inner(x):
    return x.sinc().sinc().sinc()

def checkpoint_fn(x):
    return checkpoint(inner, x, use_reentrant=False)

(ops not handled by thunder, we fallback to torch.compile pretty much the whole function).

plain gm to torch compile:

gm_fallback with proxy function:

nn.Module:

torch.compile on original function:

All executions run with TORCHINDUCTOR_FORCE_DISABLE_CACHES.

[mattteochen]

Using compile of torch dynamo (compile_fx under the hood) could be a solution, as already shown from peak memory analysis.

import gc
import torch
import torch.nn as nn
from torch.utils.checkpoint import checkpoint
from torch._inductor import compile as inductor_compile
from thunder.dynamo import thunderfx
import operator as _op
from thunder.dynamo.utils import (
    CompilerType,
)

def flat_gm_outputs(graph_module):
    for n_out in graph_module.graph.nodes:
        if n_out.op != "output":
            continue
        out_arg = n_out.args[0]
        if isinstance(out_arg, tuple) and len(out_arg) == 1:
            out_arg = out_arg[0]
        if (
            isinstance(out_arg, torch.fx.node.Node)
            and out_arg.op == "call_function"
            and out_arg.target == torch.ops.higher_order.tag_activation_checkpoint
        ):
            with graph_module.graph.inserting_before(n_out):
                get0 = graph_module.graph.call_function(_op.getitem, args=(out_arg, 0))
            n_out.args = (get0,)
            continue
        def _flatten_outputs(arg, out_list):
            if isinstance(arg, (tuple, list)):
                for item in arg:
                    _flatten_outputs(item, out_list)
            else:
                out_list.append(arg)

        flat_outs = []
        _flatten_outputs(out_arg, flat_outs)
        if len(flat_outs) == 1:
            n_out.args = (flat_outs[0],)
        else:
            n_out.args = (tuple(flat_outs),)
    graph_module.recompile()

def inner(x):
    return x.sinc().sinc().sinc()

def checkpoint_fn(x):
    return checkpoint(inner, x, use_reentrant=False)

def iso_peak(run, *args):
    gc.collect(); torch.cuda.empty_cache(); torch.cuda.synchronize()
    torch.cuda.reset_peak_memory_stats()
    mem_before = torch.cuda.memory_allocated() / (1024**3)
    out = run(*args)
    mem_after = torch.cuda.memory_allocated() / (1024**3)
    gc.collect(); torch.cuda.synchronize()
    print(f"  Memory: before={mem_before:.2f}GB, after={mem_after:.2f}GB, diff={mem_after - mem_before:.2f}GB, max={torch.cuda.max_memory_allocated()/1e9:.2f}GB")

def inductor_backend(gm, example_inputs):
    return inductor_compile(gm, example_inputs)

def main():
    device = "cuda"

    x_big = torch.randn((512, 512, 512), device=device, requires_grad=True) 
    cf = thunderfx(checkpoint_fn)
    print('ThunderFX:')
    iso_peak(cf, x_big)

    gm_fallback = None
    for sinfo in cf._backend.subgraph_infos:
        for subgm, compiled in sinfo.submodule_to_compiled_functions.items():
            if compiled.compiler == CompilerType.TORCH_INDUCTOR:
                if gm_fallback is None:
                    gm_fallback = subgm
    assert gm_fallback is not None, "No Inductor fallback subgraph found."

    def f(x):
        return gm_fallback(x)

    c_fx_plain = torch.compile(f)
    print('GraphFX with proxy')
    iso_peak(c_fx_plain, x_big)

    c_fx_plain = torch.compile(gm_fallback)
    print('GraphFX with no proxy')
    iso_peak(c_fx_plain, x_big)

    flat_gm_outputs(gm_fallback)
    c_inductor = inductor_compile(gm_fallback, [x_big.detach().clone()])
    print('Inductor compile')
    iso_peak(c_inductor, x_big)
    

if __name__ == "__main__":
    main()

Outputs:

ThunderFX:
  Memory: before=0.50GB, after=2.00GB, diff=1.50GB, max=2.15GB
GraphFX with proxy
  Memory: before=0.50GB, after=1.00GB, diff=0.50GB, max=1.14GB
GraphFX with no proxy
  Memory: before=0.50GB, after=2.00GB, diff=1.50GB, max=2.15GB
Inductor compile
  Memory: before=0.50GB, after=1.00GB, diff=0.50GB, max=1.14GB

Tracing:

The first compilation region should match:

    c_fx_plain = torch.compile(f)
    print('GraphFX with proxy')
    iso_peak(c_fx_plain, x_big)

while the latter should match (with no reference to the front end part):

    flat_gm_outputs(gm_fallback)
    c_inductor = inductor_compile(gm_fallback, [x_big.detach().clone()])
    print('Inductor compile')
    iso_peak(c_inductor, x_big)

cc @shino16