Repository: hpcaitech/ColossalAI
File: tests/test_checkpoint_io/test_low_level_zero_checkpoint_io.py (Python)

Code before:
# only test 2 is fine
@clear_cache_before_run()
@parameterize("stage", [2])
@parameterize("shard", [True, False])
@parameterize("offload", [False, True])
@parameterize("use_async", [False, True])
def check_low_level_zero_checkpointIO(stage: int, shard: bool, offload: bool, use_async: bool):
    plugin = LowLevelZeroPlugin(stage=stage, max_norm=1.0, initial_scale=32, cpu_offload=offload)
    booster = Booster(plugin=plugin)
    model = resnet18()
    criterion = lambda x: x.mean()
    optimizer = HybridAdam((model.parameters()), lr=0.001)
    model, optimizer, criterion, _, _ = booster.boost(model, optimizer, criterion)

    x = torch.randn(1, 3, 224, 224, device="cuda")
    output = model(x)
    loss = criterion(output)
    booster.backward(loss, optimizer)
    optimizer.step()
    with shared_tempdir() as tempdir:
        model_ckpt_path = f"{tempdir}/model"
        optimizer_ckpt_path = f"{tempdir}/optimizer"
        if use_async:
            model_ckpt_path = f"{model_ckpt_path}.safetensors"
        if not use_async:
            model_ckpt_path = f"{model_ckpt_path}.pt"
        booster.save_model(model, model_ckpt_path, shard=shard, use_async=use_async)

        # lr scheduler is tested in test_torch_ddp_checkpoint_io.py and low level zero does not change it, we can skip it here
        booster.save_optimizer(optimizer, optimizer_ckpt_path, shard=shard)
        booster.checkpoint_io._sync_d2h()
        booster.checkpoint_io._sync_io()
        dist.barrier()

        new_model = resnet18()
        new_optimizer = HybridAdam((new_model.parameters()), lr=0.001)
        new_model, new_optimizer, _, _, _ = booster.boost(new_model, new_optimizer)

        booster.load_model(new_model, model_ckpt_path)
        check_state_dict_equal(model.state_dict(), new_model.state_dict())
        # check master weight
        assert isinstance(new_optimizer, LowLevelZeroOptimizer)
        working_param_id_set = set(id(p) for p in new_model.parameters())
        for p_id, master_param in new_optimizer.working_to_master_param.items():
            assert p_id in working_param_id_set
            working_param = new_optimizer.master_to_working_param[id(master_param)]
            padding = new_optimizer.get_param_padding_size(working_param)
            padded_param = torch.nn.functional.pad(working_param.data.view(-1), (0, padding))
            working_shard = padded_param.chunk(dist.get_world_size())[dist.get_rank()]
            assert torch.equal(
                working_shard, master_param.data.view(-1).to(dtype=padded_param.dtype, device=padded_param.device)

Review comment (suggestion):
If shard, it should be `xxx/model`. If not shard and not async, it should be `xxxx/model.pt`. If not shard and async, it should be `xxx/model.safetensors`

Code after:
# only test 2 is fine
@clear_cache_before_run()
@parameterize("stage", [2])
@parameterize("shard", [False, True])
@parameterize("offload", [False, True])
@parameterize("use_async", [False, True])
def check_low_level_zero_checkpointIO(stage: int, shard: bool, offload: bool, use_async: bool):
    plugin = LowLevelZeroPlugin(stage=stage, max_norm=1.0, initial_scale=32, cpu_offload=offload)
    booster = Booster(plugin=plugin)
    model = resnet18()
    criterion = lambda x: x.mean()
    optimizer = HybridAdam((model.parameters()), lr=0.001)
    model, optimizer, criterion, _, _ = booster.boost(model, optimizer, criterion)

    x = torch.randn(1, 3, 224, 224, device="cuda")
    output = model(x)
    loss = criterion(output)
    booster.backward(loss, optimizer)
    optimizer.step()

    with shared_tempdir() as tempdir:

        model_ckpt_path = f"{tempdir}/model"
        optimizer_ckpt_path = f"{tempdir}/optimizer"
        if not shard and not use_async:
            model_ckpt_path = f"{model_ckpt_path}.pt"
        if not shard and use_async:
            model_ckpt_path = f"{model_ckpt_path}.safetensors"
        booster.save_model(
            model,
            model_ckpt_path,
            shard=shard,
            use_async=use_async,
        )

        # lr scheduler is tested in test_torch_ddp_checkpoint_io.py and low level zero does not change it, we can skip it here
        booster.save_optimizer(optimizer, optimizer_ckpt_path, shard=shard)
        booster.checkpoint_io._sync_d2h()
        booster.checkpoint_io._sync_io()
        dist.barrier()

        new_model = resnet18()
        new_optimizer = HybridAdam((new_model.parameters()), lr=0.001)
        new_model, new_optimizer, _, _, _ = booster.boost(new_model, new_optimizer)

        booster.load_model(new_model, model_ckpt_path)
        check_state_dict_equal(model.state_dict(), new_model.state_dict())
        # check master weight
        assert isinstance(new_optimizer, LowLevelZeroOptimizer)
        working_param_id_set = set(id(p) for p in new_model.parameters())
        for p_id, master_param in new_optimizer.working_to_master_param.items():