def METHOD_NAME(num_micro_batches):
    """
    This test is to examine the correctness of interleaved 1F1B, compared with torch.
    Be aware it contains some hardcodes.
    """
    world_size = torch.distributed.get_world_size()
    local_rank = torch.distributed.get_rank()
    seed_all(1453)
    NUM_MICRO_BATCHS = num_micro_batches
    BATCH_SIZE = num_micro_batches
    NUM_CHUNKS = 2
    # create model
    torch_model = MlpModel().cuda()
    pp_model = copy.deepcopy(torch_model).cuda()
    DP_DIM, PP_DIM, TP_DIM = 0, 1, 2
    pg_mesh = ProcessGroupMesh(1, world_size, 1)
    stage_manager = PipelineStageManager(pg_mesh, PP_DIM, is_virtual=True)
    schedule = InterleavedSchedule(NUM_MICRO_BATCHS, NUM_CHUNKS, stage_manager)
    sharded_model = torch.nn.ModuleList()
    for idx, (_, sub_model) in enumerate(pp_model.named_children()):
        if idx % (world_size) == local_rank:
            sub_model._forward = sub_model.forward
            sub_model.forward = MethodType(
                partial(pp_linear_fwd,
                        stage_mgr=stage_manager,
                        num_chunks=NUM_CHUNKS,
                        model_chunk_id=len(sharded_model)), sub_model._forward)
            sharded_model.append(sub_model.cuda())
    # create optimizer
    torch_optimizer = torch.optim.SGD(torch_model.parameters(), lr=1)
    pp_optimizer = OptimizerWrapper(torch.optim.SGD(sharded_model.parameters(), lr=1))
    # create
    seed_all(1453)
    if local_rank == 0:
        input_list = [torch.rand(BATCH_SIZE, 4).cuda()]
    else:
        input_list = [torch.zeros(BATCH_SIZE, 4).cuda()]
    torch.distributed.all_reduce(input_list[0])
    criterion = lambda x, y: torch.mean(x)
    # forward and backward
    torch_output = torch_model(input_list[0])
    torch_loss = criterion(torch_output, _)
    torch_loss.backward()
    pp_ret = schedule.forward_backward_step(sharded_model,
                                            iter(input_list),
                                            criterion,
                                            pp_optimizer,
                                            return_loss=True,
                                            return_outputs=True)
    # check loss
    if stage_manager.is_last_stage():
        assert torch.allclose(torch_loss, pp_ret['loss'])
    # check gradients
    torch_grad = []
    for torch_p in torch_model.parameters():
        torch_grad.append(torch_p.grad.data)
    for idx, pp_p in enumerate(sharded_model.parameters()):
        if idx < 2:
            assert torch.allclose(torch_grad[idx + local_rank * 2], pp_p.grad.data)
        else:
            assert torch.allclose(torch_grad[idx + local_rank * 2 + 6], pp_p.grad.data)
    # step
    torch_optimizer.step()
    pp_optimizer.step()
    # check updated param
    torch_param = []
    for torch_p in torch_model.parameters():
        torch_param.append(torch_p.data)
    for idx, pp_p in enumerate(sharded_model.parameters()):
        if idx < 2:
            assert torch.allclose(torch_param[idx + local_rank * 2], pp_p.data)
        else:
            assert torch.allclose(torch_param[idx + local_rank * 2 + 6], pp_p.data)