def METHOD_NAME(rank, world_size):
    # set up world pg
    os.environ["MASTER_ADDR"] = "localhost"
    os.environ["MASTER_PORT"] = "12355"
    # initialize the process group
    dist.init_process_group("gloo", rank=rank, world_size=world_size)
    # create a device mesh
    mesh = DeviceMesh("cpu", torch.arange(world_size))
    # create and shard the model in tensor parallel fashion
    model_tp = gen_tensor_parallel_model(SimpleMLP(), mesh)
    model_tp(torch.rand(5, 5))
    # print(f"tensor parallel model state_dict: {model_tp.state_dict()}")
    # create a 2-D device mesh for partial replication
    mesh_2d = DeviceMesh("cpu", torch.arange(world_size).reshape(2, 2))
    # replicate the parameters on the first mesh dimension,
    # and shard the parameters on the second mesh dimension
    model_2d = gen_partial_replicate_2d(SimpleMLP(), mesh_2d)
    model_2d(torch.rand(5, 5))
    # create a sub-mesh and shard/replicate params only on submesh
    submesh = DeviceMesh("cpu", [0, 2])
    model_submesh = gen_model_param_in_submesh(SimpleMLP(), submesh)
    model_submesh(torch.rand(5, 5))
    print(f"partial replicate model state_dict: {model_submesh.state_dict()}")
    # checkpoint the model
    # TODO: fully implement checkpoint save/load example
    model = checkpoint(model_2d, mesh)
    # shutting down world pg
    dist.destroy_process_group()