def METHOD_NAME(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()