def METHOD_NAME(world_size, dtype: torch.dtype):
    """
    In this test, two pairs of model and optimizers are created.
    1. zero: use sharded optimizer and fp16 parameters
    2. torch: use torch DDP and fp32 parameters
    We feed these two sets of models with the same input and check if the
    differences in model output and updated parameters are within tolerance.
    """
    local_rank = torch.distributed.get_rank()
    seed_all(1453)
    # create models
    torch_model = MlpModel().cuda()
    zero_model = copy.deepcopy(torch_model).to(dtype)
    torch_model = DDP(torch_model.cuda(), static_graph=True).cuda()
    # create optimizer
    zero_optimizer = torch.optim.SGD(zero_model.parameters(), lr=1)
    # we only test stage 1 here
    # in `check_sharded_param_consistency.py`, we will test whether
    # level 1 and 2 will produce exactly the same results
    zero_optimizer = LowLevelZeroOptimizer(zero_optimizer,
                                           overlap_communication=True,
                                           initial_scale=1,
                                           reduce_bucket_size=1024 * 1024)
    torch_optimizer = torch.optim.SGD(torch_model.parameters(), lr=1)
    seed_all(1453 + local_rank)
    # create
    input_data = torch.rand(32, 123).cuda()
    # zero-dp forward
    zero_output = zero_model(input_data.to(dtype))
    # torch-ddp forward
    torch_output = torch_model(input_data)
    loose_close(zero_output, torch_output, dtype=dtype)
    # zero-dp backward
    zero_optimizer.backward(zero_output.mean().float())
    # torch-ddp backward
    torch_output.mean().backward()
    # check grad
    for (n, p), z1p in zip(torch_model.named_parameters(), zero_model.parameters()):
        if p.grad is not None:
            zero_grad_list = zero_optimizer._grad_store.get_partitioned_gradients_by_param_id(0, id(z1p))
            torch_grad_list = split_ddp_grad(p.grad, world_size)
            for zero_grad, torch_grad in zip(zero_grad_list, torch_grad_list):
                loose_close(zero_grad, torch_grad, dtype=dtype)
    # zero-dp step
    zero_optimizer.step()
    # torch ddp step
    torch_optimizer.step()
    # check updated param
    for (n, p), z1p in zip(torch_model.named_parameters(), zero_model.parameters()):
        loose_close(p.data, z1p.data, dtype=dtype)