How does torch.fx work with activation checkpointing?

Frank_Leeeee: <p>I am wondering how we can use <code>torch.utils.checkpoint</code> together with <code>torch.fx</code>.</p>
<p>Simple example is given below:</p>
<pre><code class="lang-python">import torch
import torch.nn as nn
import torch.fx as fx
from torch.utils.checkpoint import checkpoint

class MyModule(torch.nn.Module):
    def __init__(self):
        super().__init__()
        self.linear1 = torch.nn.Linear(2, 2)
        self.linear2 = torch.nn.Linear(2, 2)
        self.linear3 = torch.nn.Linear(2, 2)
        
    def _transform(self, x):
        return x.transpose(1, 0)
    
    def forward(self, x):
        x = self.linear1(x)
        x = self.linear2(x) + self.linear3(x)
        trans_x = self._transform(x)
        return trans_x
    
class NestedModule(torch.nn.Module):
    
    def __init__(self):
        super().__init__()
        self.my_mod = MyModule()
        self.linear4 = torch.nn.Linear(2, 2)
    
    def forward(self, x):
        x = checkpoint(self.my_mod, x)
        return self.linear4(x)
    
nested_mod = NestedModule()
gm = fx.symbolic_trace(nested_mod)
gm.recompile()
print(gm)
</code></pre>
<p>The output will ignore the checkpoint logic.</p>
<pre><code class="lang-nohighlight">def forward(self, x):
    my_mod_linear1 = self.my_mod.linear1(x);  x = None
    my_mod_linear2 = self.my_mod.linear2(my_mod_linear1)
    my_mod_linear3 = self.my_mod.linear3(my_mod_linear1);  my_mod_linear1 = None
    add = my_mod_linear2 + my_mod_linear3;  my_mod_linear2 = my_mod_linear3 = None
    transpose = add.transpose(1, 0);  add = None
    linear4 = self.linear4(transpose);  transpose = None
    return linear4
</code></pre>

Radulescu_Petru: <p><a class="mention" href="/u/frank_leeeee">@Frank_Leeeee</a> Did you manage to find something that works?</p>

Frank_Leeeee: <p>Hi, I ended up writing my own codegen.</p>
<aside class="onebox githubblob" data-onebox-src="https://github.com/hpcaitech/ColossalAI/blob/main/colossalai/fx/codegen/activation_checkpoint_codegen.py">
  <header class="source">

      <a href="https://github.com/hpcaitech/ColossalAI/blob/main/colossalai/fx/codegen/activation_checkpoint_codegen.py" target="_blank" rel="noopener nofollow ugc">github.com/hpcaitech/ColossalAI</a>
  </header>

  <article class="onebox-body">
    <h4><a href="https://github.com/hpcaitech/ColossalAI/blob/main/colossalai/fx/codegen/activation_checkpoint_codegen.py" target="_blank" rel="noopener nofollow ugc">colossalai/fx/codegen/activation_checkpoint_codegen.py</a></h4>

<div class="git-blob-info">
  <a href="https://github.com/hpcaitech/ColossalAI/blob/main/colossalai/fx/codegen/activation_checkpoint_codegen.py" rel="noopener nofollow ugc"><code>main</code></a>
</div>


      <pre><code class="lang-py">from typing import Any, Dict, Iterable, List, Tuple

import torch

import colossalai

try:
    from torch.fx.graph import (
        CodeGen,
        PythonCode,
        _custom_builtins,
        _CustomBuiltin,
        _format_target,
        _is_from_torch,
        _Namespace,
        _origin_type_map,
        inplace_methods,
        magic_methods,
    )
    from torch.fx.node import Argument, Node, _get_qualified_name, _type_repr, map_arg
</code></pre>



  This file has been truncated. <a href="https://github.com/hpcaitech/ColossalAI/blob/main/colossalai/fx/codegen/activation_checkpoint_codegen.py" target="_blank" rel="noopener nofollow ugc">show original</a>

  </article>

  <div class="onebox-metadata">
    
    
  </div>

  <div style="clear: both"></div>
</aside>