def METHOD_NAME(shard: bool, model_name: str, size_per_shard: int, test_config: dict):
    (model_fn, data_gen_fn, output_transform_fn, loss_fn,
     _) = next(iter(model_zoo.get_sub_registry(model_name).values()))
    criterion = loss_fn
    plugin = HybridParallelPlugin(**test_config)
    booster = Booster(plugin=plugin)
    def _criterion(outputs, inputs):
        outputs = output_transform_fn(outputs)
        loss = criterion(outputs)
        return loss
    def _preprocess_data(data):
        if booster.plugin.stage_manager is not None:
            for k, v in data.items():
                if torch.is_tensor(v) or 'Tensor' in v.__class__.__name__:
                    new_shape = [1] * v.dim()
                    new_shape[0] = 4
                    data[k] = v.to('cuda').repeat(*new_shape)
            return iter([data])
        else:
            return {k: v.cuda() for k, v in data.items()}
    model = model_fn().cuda()
    optimizer = Adam(model.parameters(), lr=1e-3)
    model, optimizer, criterion, _, _ = booster.boost(model, optimizer, criterion)
    data = data_gen_fn()
    model.train()
    if booster.plugin.stage_manager is not None:
        booster.execute_pipeline(_preprocess_data(data),
                                 model,
                                 _criterion,
                                 optimizer,
                                 return_loss=True,
                                 return_outputs=False)
    else:
        output = model(**_preprocess_data(data))
        loss = criterion(output)
        optimizer.backward(loss)
    optimizer.step()
    with shared_tempdir() as tempdir:
        model_ckpt_path = f"{tempdir}/model"
        optimizer_ckpt_path = f"{tempdir}/optimizer"
        booster.save_model(model, model_ckpt_path, shard=shard, size_per_shard=size_per_shard)
        booster.save_optimizer(optimizer, optimizer_ckpt_path, shard=shard, size_per_shard=size_per_shard)
        dist.barrier()
        new_model = model_fn().cuda()
        new_optimizer = Adam(new_model.parameters(), lr=1e-3)
        new_model, new_optimizer, criterion, _, _ = booster.boost(new_model, new_optimizer, criterion)
        booster.load_model(new_model, model_ckpt_path)
        check_state_dict_equal(model.unwrap().state_dict(), new_model.unwrap().state_dict(), False)
        booster.load_optimizer(new_optimizer, optimizer_ckpt_path)
        check_state_dict_equal(optimizer.unwrap().state_dict(), new_optimizer.unwrap().state_dict(), False)
        dist.barrier()
    # Check whether the loaded model & optimizer works smoothly.
    model.train()
    new_model.train()
    if booster.plugin.stage_manager is not None:
        booster.execute_pipeline(_preprocess_data(data),
                                 model,
                                 _criterion,
                                 optimizer,
                                 return_loss=True,
                                 return_outputs=False)
        booster.execute_pipeline(_preprocess_data(data),
                                 new_model,
                                 _criterion,
                                 new_optimizer,
                                 return_loss=True,
                                 return_outputs=False)
    else:
        old_model_loss = criterion(model(**_preprocess_data(data)))
        optimizer.backward(old_model_loss)
        new_model_loss = criterion(new_model(**_preprocess_data(data)))
        new_optimizer.backward(new_model_loss)
    optimizer.step()
    new_optimizer.step()
    # Check updated weights.
    stage_manager = booster.plugin.stage_manager
    if stage_manager is None or stage_manager.is_first_stage():
        assert_close_loose(model.unwrap().wte.weight.data, new_model.unwrap().wte.weight.data, atol=5e-3, rtol=5e-3)
        assert_close_loose(model.unwrap().h[0].mlp.c_fc.weight.data,
                           new_model.unwrap().h[0].mlp.c_fc.weight.data,
                           atol=5e-3,
                           rtol=5e-3)
    dist.barrier()
    Randomizer.reset_index()
    clear_layout_converter()