code: def _run_act_ckpt_codegen(rank, world_size, port):
      # launch colossalai to make sure we could execute colossalai.utils.checkpoint currently
      colossalai.launch(
          config={},
          rank=rank,
          world_size=world_size,
          host="localhost",
          port=port,
          backend="nccl",
      )
  
      # build model and run forward
      model = MyModule()
      data1 = torch.rand(4, 4)
  
      # copy model to cuda
      model = model.to(device="cuda")
      data1 = data1.to(device="cuda")
  
      non_fx_out = model(data1)
  
      # trace the module and replace codegen
      tracer = ColoTracer(trace_act_ckpt=True)
      graph = tracer.trace(model)
      codegen = ActivationCheckpointCodeGen()
      graph.set_codegen(codegen)
  
      # annotate nested checkpoint
      for node in graph.nodes:
          if node.name == "linear1":
              node.meta["activation_checkpoint"] = [0, 0, 0]
              continue
          if node.name == "linear2":
              node.meta["activation_checkpoint"] = [0, 0, None]
          if node.name == "linear3":
              node.meta["activation_checkpoint"] = [0, 0, 1]
          if node.name == "linear4":
              node.meta["activation_checkpoint"] = [0, 1, None]
          if node.name == "linear5":
              node.meta["activation_checkpoint"] = 1
      gm = ColoGraphModule(model, graph)
      gm.recompile()
  
      # assert checkpoint function will be generated and
      code = graph.python_code("self").src
      assert (
          "colossalai.utils.activation_checkpoint.checkpoint(self.checkpoint_0_0, False, x, use_reentrant=False)"
          in code
          and "colossalai.utils.activation_checkpoint.checkpoint(self.checkpoint_0_1, False, linear3, use_reentrant=False)"
          in code
          and "colossalai.utils.activation_checkpoint.checkpoint(self.checkpoint_0_0_0, False, x, use_reentrant=False)"
          in code
          and "colossalai.utils.activation_checkpoint.checkpoint(self.checkpoint_0_0_1, False, linear2, use_reentrant=False)"
          in code
          and "colossalai.utils.activation_checkpoint.checkpoint(self.checkpoint_0, False, x, use_reentrant=False)"
          in code
          and "colossalai.utils.activation_checkpoint.checkpoint(self.checkpoint_1, False, linear4, use_reentrant=False)"
          in code
      )
  
      # recompile and verify the outputs are consistent
      fx_out = gm(data1)
      assert torch.equal(non_fx_out, fx_out)
  
      gpc.destroy()
  
file_path: colossalai/tests/test_fx/test_codegen/test_nested_activation_checkpoint_codegen.py
imports: import pytest
  import torch
  import colossalai
  from colossalai.fx import ColoTracer
  from colossalai.fx.graph_module import ColoGraphModule
  from colossalai.legacy.core import global_context
  from colossalai.testing import rerun_if_address_is_in_use, spawn
label_desc: Use image node_id 1 for calling a global function with example usage: _run_act_ckpt_codegen(rank, world_size, port) without return types