code: def _run_offload_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 input
      model = MyNet().cuda()
      data = torch.rand(4, 4).cuda()
  
      # trace the module and replace codegen
      tracer = ColoTracer(trace_act_ckpt=True)
      graph = tracer.trace(model)
      codegen = ActivationCheckpointCodeGen()
      graph.set_codegen(codegen)
  
      # annotate the activation offload part
      # also annotate the activation_checkpoint so we could test both types
      # of input offload
      for node in graph.nodes:
          if node.name == "linear0":
              node.meta["activation_offload"] = [0, True, False]
          if node.name == "linear1":
              node.meta["activation_offload"] = [0, True, False]
          if node.name == "linear2":
              node.meta["activation_offload"] = [1, True, True]
          if node.name == "linear4":
              node.meta["activation_offload"] = [2, False, True]
          if node.name == "linear5":
              node.meta["activation_checkpoint"] = [0]
              node.meta["activation_offload"] = True
  
      gm = ColoGraphModule(copy.deepcopy(model), graph)
      gm.recompile()
  
      # assert we have all the components
      code = graph.python_code("self").src
      assert (
          "def pack_hook_input(self, x):" in code
          and "def unpack_hook(self, packed):" in code
          and "def pack_hook_no_input(self, x):" in code
          and "setattr(x, 'offload', True)" in code
          and "setattr(linear3, 'offload', False)" in code
          and "with torch.autograd.graph.saved_tensors_hooks(self.pack_hook_input, self.unpack_hook):"
          in code
          and "with torch.autograd.graph.save_on_cpu(pin_memory=True):"
          in code
          and "with torch.autograd.graph.saved_tensors_hooks(self.pack_hook_no_input, self.unpack_hook):"
          in code
          and "colossalai.utils.activation_checkpoint.checkpoint(self.checkpoint_0, True, linear4, use_reentrant=False)"
          in code
      )
  
      _test_fwd_and_bwd(model, gm, data)
      gpc.destroy()
  
file_path: ColossalAI/tests/test_fx/test_codegen/test_offload_codegen.py
imports: import copy
  import pytest
  import torch
  from torch.fx import GraphModule
  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 3 for calling a global function with example usage: _run_offload_codegen(rank, world_size, port) without return types