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)
      data2 = torch.rand(4, 4)
  
      # copy model to cuda
      model = model.to(device="cuda")
      data1 = data1.to(device="cuda")
      data2 = data2.to(device="cuda")
  
      non_fx_out = model(data1, data2)
  
      # trace the module and replace codegen
      tracer = ColoTracer(trace_act_ckpt=True)
      graph = tracer.trace(model)
      codegen = ActivationCheckpointCodeGen()
      graph.set_codegen(codegen)
  
      # check ops are annotated with ckpt
      # also annotate the selected node for offloading
      ckpt_nodes = ["mlp1_linear1", "mlp1_linear2", "relu_relu", "relu"]
      offload_starts = ["mlp1_linear1"]
      for node in graph.nodes:
          if node.name in ckpt_nodes:
              assert "activation_checkpoint" in node.meta
  
              # annotate the selected node for offload
              if node.name in offload_starts:
                  node.meta["activation_offload"] = True
  
      gm = ColoGraphModule(model, graph)
      gm.recompile()
  
      # assert checkpoint function will be generated and
      # the offload option is correct
      code = graph.python_code("self").src
      assert (
          "colossalai.utils.activation_checkpoint.checkpoint(self.checkpoint_0, True, x, use_reentrant=False)"
          in code
          and "colossalai.utils.activation_checkpoint.checkpoint(self.checkpoint_1, False, x, use_reentrant=False)"
          in code
          and "colossalai.utils.activation_checkpoint.checkpoint(self.checkpoint_2, False, y, use_reentrant=False)"
          in code
          and "colossalai.utils.activation_checkpoint.checkpoint(self.checkpoint_3, False, y, relu, use_reentrant=True)"
          in code
      )
  
      # recompile and verify the outputs are consistent
      fx_out = gm(data1, data2)
      assert torch.equal(non_fx_out, fx_out)
  
      gpc.destroy()
  
imports: import pytest
  import torch
  import torch.nn.functional
  from torch.utils.checkpoint import checkpoint
  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