wehub-resource-sync
256 lines
12 KiB
Python
256 lines
12 KiB
Python
# Copyright (c) Microsoft Corporation.
|
|
# SPDX-License-Identifier: Apache-2.0
|
|
|
|
# DeepSpeed Team
|
|
|
|
import os
|
|
import torch
|
|
import numbers
|
|
|
|
import deepspeed
|
|
from deepspeed.runtime.zero.stage_1_and_2 import DeepSpeedZeroOptimizer
|
|
from deepspeed.runtime.fp16.fused_optimizer import FP16_Optimizer
|
|
from deepspeed.runtime.fp16.unfused_optimizer import FP16_UnfusedOptimizer
|
|
from deepspeed.runtime.zero.stage3 import DeepSpeedZeroOptimizer_Stage3
|
|
from deepspeed.runtime.zero.partition_parameters import ZeroParamStatus
|
|
|
|
from unit.common import preferred_dtype
|
|
from unit.simple_model import *
|
|
from unittest.mock import MagicMock, patch
|
|
|
|
|
|
def compare_deepspeed_states(saved_model, loaded_model):
|
|
# These are compared in more depth in other places
|
|
assert hasattr(loaded_model, 'module')
|
|
|
|
assert saved_model.sparse_tensor_module_names == loaded_model.sparse_tensor_module_names
|
|
assert saved_model.skipped_steps == loaded_model.skipped_steps
|
|
assert saved_model.global_steps == loaded_model.global_steps
|
|
|
|
|
|
def zero3_params_to_fetch(param_list):
|
|
return [p for p in param_list if hasattr(p, 'ds_id') and p.ds_status == ZeroParamStatus.NOT_AVAILABLE]
|
|
|
|
|
|
def compare_model_states(saved_model, loaded_model, compare_optimizer=True, load_module_only=False):
|
|
if not load_module_only:
|
|
compare_deepspeed_states(saved_model, loaded_model)
|
|
|
|
params_to_fetch = zero3_params_to_fetch(
|
|
list(saved_model.module.named_parameters()) + list(loaded_model.module.named_parameters()))
|
|
enable_gather = len(params_to_fetch) > 0
|
|
with deepspeed.zero.GatheredParameters(params_to_fetch, enabled=enable_gather):
|
|
for p0, p1 in zip(saved_model.module.named_parameters(), loaded_model.module.named_parameters()):
|
|
np0, p0 = p0
|
|
np1, p1 = p1
|
|
if 'deepspeed_moe.gate.wg' in np0:
|
|
# these params are converted to float at runtime, cast to half for comparison
|
|
p1 = p1.half()
|
|
p0 = p0.half()
|
|
assert id(p0) != id(p1), f'Comparing fp16 model state tensor against itself : {id(p0)} <====> {id(p1)}'
|
|
try:
|
|
assert torch.allclose(p0, p1,
|
|
atol=1e-07), f"FP16 model state {p0} is not equal to {p1}, names:{np0}, {np1}"
|
|
except RuntimeError as err:
|
|
print(f"FP16 model state {p0} is not equal to {p1}, names:{np0}, {np1}")
|
|
raise err
|
|
|
|
if not compare_optimizer:
|
|
return
|
|
|
|
if DeepSpeedZeroOptimizer_Stage3 is not None and isinstance(saved_model.optimizer, DeepSpeedZeroOptimizer_Stage3):
|
|
for p0, p1 in zip(saved_model.optimizer.fp32_partitioned_groups_flat,
|
|
loaded_model.optimizer.fp32_partitioned_groups_flat):
|
|
assert torch.allclose(p0, p1, atol=1e-07), f"Fp32 model states {p0} is not equal to {p1}"
|
|
|
|
elif isinstance(saved_model.optimizer, DeepSpeedZeroOptimizer):
|
|
for p0, p1 in zip(saved_model.optimizer.single_partition_of_fp32_groups,
|
|
loaded_model.optimizer.single_partition_of_fp32_groups):
|
|
assert id(p0) != id(p1), f'Comparing fp32 model state tensor against itself: {id(p0)} <====> {id(p1)}'
|
|
assert torch.allclose(p0, p1, atol=1e-07), f"Fp32 model states {p0} is not equal to {p1}"
|
|
|
|
elif isinstance(saved_model.optimizer, FP16_Optimizer):
|
|
for p0, p1 in zip(saved_model.optimizer.fp32_groups_flat, loaded_model.optimizer.fp32_groups_flat):
|
|
assert id(p0) != id(p1), f'Comparing fp32 model state tensor against itself: {id(p0)} <====> {id(p1)}'
|
|
assert torch.allclose(p0, p1, atol=1e-07), f"FP32 model states {p0} is not equal to {p1}"
|
|
|
|
elif isinstance(saved_model.optimizer, FP16_UnfusedOptimizer):
|
|
for params0, params1 in zip(saved_model.optimizer.fp32_groups, loaded_model.optimizer.fp32_groups):
|
|
for p0, p1 in zip(params0, params1):
|
|
assert id(p0) != id(p1), f'Comparing fp32 model state tensor against itself: {id(p0)} <====> {id(p1)}'
|
|
assert torch.allclose(p0, p1, atol=1e-07), f"FP32 model states {p0} is not equal to {p1}"
|
|
elif isinstance(saved_model.optimizer, torch.optim.Optimizer):
|
|
pass
|
|
else:
|
|
assert False, f'Unexpected Optimizer Type: {saved_model.optimizer}'
|
|
|
|
|
|
def compare_state_dicts(state0, state1, expected_mismatch_keys=[]):
|
|
key_set0 = set(k for k in state0.keys() if k not in expected_mismatch_keys)
|
|
key_set1 = set(k for k in state1.keys() if k not in expected_mismatch_keys)
|
|
assert key_set0 == key_set1, f'failure due to key mismatch {key_set0} != {key_set1}'
|
|
|
|
for k in key_set0:
|
|
s0 = state0[k]
|
|
s1 = state1[k]
|
|
if k in expected_mismatch_keys:
|
|
continue
|
|
if isinstance(s0, torch.Tensor) and isinstance(s1, torch.Tensor):
|
|
assert id(s0) != id(s1), f'Comparing optimizer state tensor against itself: {id(s0)} <====> {id(s1)}'
|
|
assert torch.equal(s0.to('cpu'), s1.to('cpu'))
|
|
else:
|
|
assert s0 == s1, f'failures with keys = {k}, {k}, values = {s0} and {s1}'
|
|
|
|
|
|
def compare_opt_state_dicts(state0, state1, expected_mismatch_keys=[]):
|
|
for param_group0, saved_param_group1 in zip(state0['param_groups'], state1['param_groups']):
|
|
compare_state_dicts(param_group0, saved_param_group1, expected_mismatch_keys)
|
|
|
|
assert "state" in state0
|
|
assert "state" in state1
|
|
assert len([state0["state"].keys()]) == len([state1["state"].keys()])
|
|
|
|
for (k0, s0), (k1, s1) in zip(state0["state"].items(), state1["state"].items()):
|
|
assert k0 == k1, f'failure due to key mismatch {k0} != {k1}'
|
|
compare_state_dicts(s0, s1, expected_mismatch_keys)
|
|
|
|
|
|
def compare_optimizer_states(saved_model, loaded_model, hidden_dim, fp16=True):
|
|
saved_optimizer = saved_model.optimizer.optimizer if fp16 else saved_model.optimizer
|
|
loaded_optimizer = loaded_model.optimizer.optimizer if fp16 else loaded_model.optimizer
|
|
|
|
for state0, state1 in zip(saved_optimizer.state.values(), loaded_optimizer.state.values()):
|
|
compare_state_dicts(state0, state1)
|
|
|
|
|
|
def compare_lr_scheduler_states(saved_model, loaded_model):
|
|
assert hasattr(saved_model, 'lr_scheduler')
|
|
assert hasattr(loaded_model, 'lr_scheduler')
|
|
|
|
saved_scheduler = saved_model.lr_scheduler
|
|
loaded_scheduler = loaded_model.lr_scheduler
|
|
|
|
assert hasattr(saved_scheduler, 'state_dict')
|
|
assert hasattr(loaded_scheduler, 'state_dict')
|
|
|
|
saved_sd = saved_scheduler.state_dict()
|
|
loaded_sd = loaded_scheduler.state_dict()
|
|
|
|
print(f"saved_sd = {saved_sd}")
|
|
print(f"loaded_sd = {loaded_sd}")
|
|
|
|
assert saved_sd.keys() == loaded_sd.keys()
|
|
|
|
for state0, state1 in zip(saved_sd.values(), loaded_sd.values()):
|
|
if isinstance(state0, numbers.Number) and isinstance(state1, numbers.Number):
|
|
assert state0 == state1
|
|
|
|
|
|
# following mixture-of-experts.md
|
|
def create_moe_param_groups(model):
|
|
from deepspeed.moe.utils import split_params_into_different_moe_groups_for_optimizer
|
|
|
|
parameters = {'params': [p for p in model.parameters()], 'name': 'parameters'}
|
|
return split_params_into_different_moe_groups_for_optimizer(parameters)
|
|
|
|
|
|
def create_deepspeed_model(config_dict, model, base_optimizer):
|
|
ds_model, _, _, _ = deepspeed.initialize(config=config_dict,
|
|
model=model,
|
|
model_parameters=create_moe_param_groups(model),
|
|
optimizer=base_optimizer)
|
|
ds_model.empty_partition_cache()
|
|
return ds_model
|
|
|
|
|
|
def checkpoint_correctness_verification(config_dict,
|
|
models,
|
|
hidden_dim,
|
|
tmpdir,
|
|
load_optimizer_states=False,
|
|
load_lr_scheduler_states=False,
|
|
train_batch=False,
|
|
base_optimizers=[None, None],
|
|
empty_tag=False,
|
|
seq_dataloader=False,
|
|
load_module_only=False,
|
|
dtype=None):
|
|
if dtype is None:
|
|
dtype = preferred_dtype()
|
|
|
|
ds_model = create_deepspeed_model(config_dict=config_dict, model=models[0], base_optimizer=base_optimizers[0])
|
|
|
|
if seq_dataloader:
|
|
data_loader = sequence_dataloader(model=ds_model,
|
|
total_samples=50,
|
|
hidden_dim=hidden_dim,
|
|
device=ds_model.device,
|
|
dtype=dtype)
|
|
else:
|
|
data_loader = random_dataloader(model=ds_model,
|
|
total_samples=50,
|
|
hidden_dim=hidden_dim,
|
|
device=ds_model.device,
|
|
dtype=dtype)
|
|
|
|
if train_batch:
|
|
ds_model.set_dataloader(data_loader)
|
|
for _, batch in enumerate(data_loader):
|
|
loss = ds_model.train_batch()
|
|
else:
|
|
for _, batch in enumerate(data_loader):
|
|
loss = ds_model(batch[0], batch[1])
|
|
ds_model.backward(loss)
|
|
ds_model.step()
|
|
|
|
# Flush zero stage 3 cache
|
|
ds_model.empty_partition_cache()
|
|
|
|
trained_model = ds_model
|
|
|
|
save_folder = os.path.join(tmpdir, 'saved_checkpoint')
|
|
save_tag = None if empty_tag else '1'
|
|
|
|
trained_model.save_checkpoint(save_folder, tag=save_tag)
|
|
|
|
dist.barrier()
|
|
|
|
for root, _, files in os.walk(save_folder):
|
|
for f in files:
|
|
if "_expert_" in f and "_model_states" in f:
|
|
expert = torch.load(os.path.join(root, f), weights_only=False)
|
|
needed, storages = 0, {}
|
|
for name, tensor in expert.items():
|
|
needed += tensor.size().numel()
|
|
storage = tensor.storage()
|
|
# some storage can be shared within an expert's checkpoint
|
|
storages[storage.data_ptr()] = storage.size()
|
|
stored = sum(v for _, v in storages.items())
|
|
assert needed == stored, f"MoE expert checkpoint uses more storage than required: {f}"
|
|
|
|
loaded_model = create_deepspeed_model(config_dict=config_dict, model=models[1], base_optimizer=base_optimizers[1])
|
|
assert list(trained_model.parameters())[0].dtype == list(loaded_model.parameters())[0].dtype
|
|
|
|
context = patch.object(loaded_model, "_get_optimizer_ckpt_name",
|
|
wraps=loaded_model._get_optimizer_ckpt_name) if not load_optimizer_states else MagicMock()
|
|
with context as optim_load_state_dict_mock:
|
|
loaded_model.load_checkpoint(save_folder,
|
|
tag=save_tag,
|
|
load_optimizer_states=load_optimizer_states,
|
|
load_lr_scheduler_states=load_lr_scheduler_states,
|
|
load_module_only=load_module_only)
|
|
if not load_optimizer_states:
|
|
# should not attempt to get the file name to load it
|
|
optim_load_state_dict_mock.assert_not_called()
|
|
|
|
compare_model_states(trained_model,
|
|
loaded_model,
|
|
compare_optimizer=load_optimizer_states,
|
|
load_module_only=load_module_only)
|
|
|
|
if load_optimizer_states:
|
|
compare_optimizer_states(trained_model, loaded_model, hidden_dim, dtype == torch.float16)
|
|
|
|
if load_lr_scheduler_states:
|
|
compare_lr_scheduler_states(trained_model, loaded_model)
|