wehub-resource-sync
273 lines
10 KiB
Python
273 lines
10 KiB
Python
# Copyright (c) Microsoft Corporation.
|
|
# SPDX-License-Identifier: Apache-2.0
|
|
|
|
# DeepSpeed Team
|
|
|
|
import os
|
|
import math
|
|
|
|
import deepspeed
|
|
from types import SimpleNamespace
|
|
from torch.utils._pytree import tree_map
|
|
|
|
from deepspeed.utils.torch import required_torch_version
|
|
from deepspeed.checkpoint import UNIVERSAL_CHECKPOINT_INFO
|
|
from deepspeed.checkpoint.ds_to_universal import main as convert_to_universal
|
|
|
|
from unit.common import DistributedTest, DistributedFixture
|
|
from unit.simple_model import *
|
|
from unit.util import bf16_required_version_check
|
|
|
|
from unit.checkpoint.common import compare_opt_state_dicts, compare_state_dicts
|
|
|
|
import pytest
|
|
import deepspeed.comm as dist
|
|
|
|
|
|
def get_expected_mismatch_keys():
|
|
# torch 1.2.* stores raw tensor id numbers in checkpoint state which leads to
|
|
# false positive mismatches in checkpoint state comparisons.
|
|
# Newer torch versions store tensor ids as 0, 1, 2, ...
|
|
return [] if required_torch_version(min_version=1.4) else ['params']
|
|
|
|
|
|
def maybe_step(t):
|
|
return not torch.is_tensor(t) or (t.device.type == 'cpu' and t.numel() == 1)
|
|
|
|
|
|
def gather_opt_state(optimizer_state):
|
|
|
|
def gather_tensor(t):
|
|
|
|
if maybe_step(t):
|
|
return t
|
|
else:
|
|
buffer = [torch.zeros_like(t.flatten()) for _ in range(dist.get_world_size())]
|
|
dist.all_gather(buffer, t.flatten())
|
|
return torch.cat(buffer)
|
|
|
|
return tree_map(gather_tensor, optimizer_state)
|
|
|
|
|
|
def remove_pad_in_opt_state(optimizer_state, num_params):
|
|
|
|
def remove_pad(t):
|
|
if maybe_step(t):
|
|
return t
|
|
else:
|
|
return t[:num_params]
|
|
|
|
return tree_map(remove_pad, optimizer_state)
|
|
|
|
|
|
CP_TAG = "test_tag"
|
|
|
|
|
|
def init_ds_engine(model, ds_config, use_torch_adam):
|
|
|
|
if use_torch_adam:
|
|
ds_optimizer = torch.optim.Adam(model.parameters(), lr=0.1)
|
|
del ds_config["optimizer"]
|
|
model, _, _, _ = deepspeed.initialize(config=ds_config, model=model, optimizer=ds_optimizer)
|
|
else:
|
|
model, _, _, _ = deepspeed.initialize(config=ds_config, model=model, model_parameters=model.parameters())
|
|
|
|
return model
|
|
|
|
|
|
def train_save_convert(ds_config, hidden_dim, load_optim, use_torch_adam, dtype, tmpdir, world_size):
|
|
if dtype == torch.bfloat16 and not bf16_required_version_check():
|
|
return
|
|
|
|
test_step = 8
|
|
|
|
model = SimpleModel(hidden_dim, nlayers=2)
|
|
model = init_ds_engine(model, ds_config, use_torch_adam)
|
|
data_loader = random_dataloader(model=model,
|
|
total_samples=test_step,
|
|
hidden_dim=hidden_dim,
|
|
device=model.device,
|
|
dtype=dtype)
|
|
for batch in data_loader:
|
|
loss = model(batch[0], batch[1])
|
|
model.backward(loss)
|
|
model.step()
|
|
|
|
if ds_config["zero_optimization"]["stage"] == 3:
|
|
model.optimizer._set_fp32_optimizer_param_groups()
|
|
sd = model.optimizer.optimizer.state_dict() if load_optim else None
|
|
model.optimizer._clear_fp32_optimizer_param_groups()
|
|
else:
|
|
sd = model.optimizer.optimizer.state_dict() if load_optim else None
|
|
|
|
client_state = {}
|
|
client_state[UNIVERSAL_CHECKPOINT_INFO] = {}
|
|
client_state['iteration'] = test_step
|
|
model.save_checkpoint(tmpdir, tag=CP_TAG, client_state=client_state)
|
|
|
|
cp_dir = os.path.join(tmpdir, CP_TAG)
|
|
univ_cp_dir = f"{cp_dir}_universal"
|
|
|
|
args = SimpleNamespace(input_folder=cp_dir,
|
|
output_folder=univ_cp_dir,
|
|
num_extract_workers=1,
|
|
num_merge_workers=1,
|
|
keep_temp_folder=False,
|
|
strict=True,
|
|
inject_missing_state=False)
|
|
|
|
dist.barrier()
|
|
if dist.get_rank() == 0:
|
|
convert_to_universal(args)
|
|
|
|
model_state = model.state_dict()
|
|
optimizer_state = None
|
|
if load_optim:
|
|
if ds_config["zero_optimization"]["stage"] == 3:
|
|
model.optimizer._set_fp32_optimizer_param_groups()
|
|
optimizer_state = gather_opt_state(model.optimizer.optimizer.state_dict())
|
|
model.optimizer._clear_fp32_optimizer_param_groups()
|
|
update_gathered_stage3_optimizer(optimizer_state, model._get_zero_param_shapes(), world_size)
|
|
else:
|
|
optimizer_state = gather_opt_state(model.optimizer.optimizer.state_dict())
|
|
|
|
if dist.get_rank() == 0:
|
|
torch.save((model_state, optimizer_state), os.path.join(tmpdir, "baseline_state.pt"))
|
|
|
|
dist.barrier()
|
|
model.destroy()
|
|
|
|
|
|
@pytest.fixture
|
|
def ds_config(zero_stage, dtype, sub_group_size):
|
|
ds_config = {
|
|
"train_batch_size": 8,
|
|
"optimizer": {
|
|
"type": 'Adam'
|
|
},
|
|
"zero_optimization": {
|
|
"stage": zero_stage,
|
|
}
|
|
}
|
|
if dtype == torch.float16:
|
|
ds_config["fp16"] = {"enabled": True, "initial_scale_power": 8}
|
|
elif dtype == torch.bfloat16:
|
|
ds_config["bf16"] = {"enabled": True}
|
|
if sub_group_size > 0:
|
|
ds_config["zero_optimization"]["sub_group_size"] = sub_group_size
|
|
return ds_config
|
|
|
|
|
|
class _baseline(DistributedFixture):
|
|
world_size = None
|
|
|
|
def run(self, tmpdir, ds_config, zero_stage, dtype, load_optim, use_torch_adam):
|
|
hidden_dim = 10
|
|
train_save_convert(ds_config, hidden_dim, load_optim, use_torch_adam, dtype, tmpdir, self.world_size)
|
|
|
|
|
|
class baseline_ws2(_baseline):
|
|
world_size = 2
|
|
|
|
|
|
class baseline_ws4(_baseline):
|
|
world_size = 4
|
|
|
|
|
|
# Stage3 use shard parameter, need to reorganize the optimizer parameters.
|
|
def update_gathered_stage3_optimizer(optimizer_state, param_shapes, world_size):
|
|
for sub_group_id, group in enumerate(optimizer_state["param_groups"]):
|
|
group["params"] = None
|
|
|
|
new_state = {}
|
|
for sub_group_id, sub_group_param_shape in enumerate(param_shapes):
|
|
total_numel = optimizer_state['state'][sub_group_id]['exp_avg'].numel()
|
|
assert total_numel % world_size == 0
|
|
numel_per_rank = total_numel // world_size
|
|
param_offset_in_current_rank = 0
|
|
for param_name, param_shape in sub_group_param_shape.items():
|
|
param_numel = param_shape.numel()
|
|
param_partition_numel = math.ceil(param_numel / world_size)
|
|
param_optimizer_tensor = {
|
|
"exp_avg": torch.zeros(param_numel),
|
|
"exp_avg_sq": torch.zeros(param_numel),
|
|
"step": optimizer_state['state'][sub_group_id]['step'],
|
|
}
|
|
for key in ["exp_avg", "exp_avg_sq"]:
|
|
write_offset = 0
|
|
for rank in range(world_size):
|
|
offset = param_offset_in_current_rank + rank * numel_per_rank
|
|
length = min(param_partition_numel, param_numel - rank * param_partition_numel)
|
|
tmp = optimizer_state['state'][sub_group_id][key].narrow(0, offset, length)
|
|
param_optimizer_tensor[key].narrow(0, write_offset, length).copy_(tmp)
|
|
write_offset += length
|
|
param_offset_in_current_rank += param_partition_numel
|
|
new_state[param_name] = param_optimizer_tensor
|
|
optimizer_state["state"] = new_state
|
|
|
|
|
|
@pytest.mark.parametrize('dtype', [torch.bfloat16, torch.float16, torch.float32])
|
|
@pytest.mark.parametrize("zero_stage", [1, 3])
|
|
@pytest.mark.parametrize("use_torch_adam", [False, True])
|
|
@pytest.mark.parametrize("load_optim", [False, True])
|
|
@pytest.mark.parametrize("sub_group_size", [-1, 100])
|
|
class TestZeROUniversalCheckpointDP(DistributedTest):
|
|
|
|
def _run_test(self, tmpdir, dtype, ds_config, load_optim, use_torch_adam, world_size):
|
|
if dtype == torch.bfloat16 and not bf16_required_version_check():
|
|
pytest.skip(
|
|
" DeepSpeed BFloat16 tests need torch >= 1.10, NCCL >= 2.10.3, CUDA > =11.0 and HW support for BFloat16 to run correctly"
|
|
)
|
|
|
|
hidden_dim = 10
|
|
loaded_model_state, loaded_optimizer_state = torch.load(f"{tmpdir}/baseline_state.pt", weights_only=False)
|
|
|
|
ds_config["checkpoint"] = {"load_universal": True}
|
|
univ_model = SimpleModel(hidden_dim, nlayers=2)
|
|
univ_model = init_ds_engine(univ_model, ds_config, use_torch_adam)
|
|
univ_model.load_checkpoint(tmpdir, tag=f"{CP_TAG}_universal", load_optimizer_states=load_optim)
|
|
|
|
model_state = univ_model.state_dict()
|
|
compare_state_dicts(model_state, loaded_model_state)
|
|
|
|
if load_optim:
|
|
if ds_config["zero_optimization"]["stage"] == 3:
|
|
univ_model.optimizer._set_fp32_optimizer_param_groups()
|
|
optimizer_state = gather_opt_state(univ_model.optimizer.optimizer.state_dict())
|
|
univ_model.optimizer._clear_fp32_optimizer_param_groups()
|
|
update_gathered_stage3_optimizer(optimizer_state, univ_model._get_zero_param_shapes(), world_size)
|
|
else:
|
|
optimizer_state = gather_opt_state(univ_model.optimizer.optimizer.state_dict())
|
|
# padding sizes may differ when dp sizes are different
|
|
param_count = sum(p.numel() for p in univ_model.parameters())
|
|
optimizer_state = remove_pad_in_opt_state(optimizer_state, param_count)
|
|
loaded_optimizer_state = remove_pad_in_opt_state(loaded_optimizer_state, param_count)
|
|
|
|
compare_opt_state_dicts(optimizer_state, loaded_optimizer_state, get_expected_mismatch_keys())
|
|
|
|
# Run training again to verify that the optimizer has necessary states
|
|
test_step = 8
|
|
data_loader = random_dataloader(model=univ_model,
|
|
total_samples=test_step,
|
|
hidden_dim=hidden_dim,
|
|
device=univ_model.device,
|
|
dtype=dtype)
|
|
for batch in data_loader:
|
|
loss = univ_model(batch[0], batch[1])
|
|
univ_model.backward(loss)
|
|
univ_model.step()
|
|
|
|
univ_model.destroy()
|
|
|
|
@pytest.mark.world_size(2)
|
|
def test_dp_world_size_2to2(self, baseline_ws2, tmpdir, dtype, ds_config, load_optim, use_torch_adam):
|
|
self._run_test(tmpdir, dtype, ds_config, load_optim, use_torch_adam, 2)
|
|
|
|
@pytest.mark.world_size(2)
|
|
def test_dp_world_size_4to2(self, baseline_ws4, tmpdir, dtype, ds_config, load_optim, use_torch_adam):
|
|
self._run_test(tmpdir, dtype, ds_config, load_optim, use_torch_adam, 2)
|
|
|
|
@pytest.mark.world_size(4)
|
|
def test_dp_world_size_2to4(self, baseline_ws2, tmpdir, dtype, ds_config, load_optim, use_torch_adam):
|
|
self._run_test(tmpdir, dtype, ds_config, load_optim, use_torch_adam, 4)
|