wehub-resource-sync
777 lines
32 KiB
Python
777 lines
32 KiB
Python
# Copyright (c) Microsoft Corporation.
|
|
# SPDX-License-Identifier: Apache-2.0
|
|
|
|
# DeepSpeed Team
|
|
|
|
import deepspeed
|
|
from types import SimpleNamespace
|
|
from deepspeed.ops.op_builder import CPUAdamBuilder
|
|
from deepspeed.checkpoint.utils import clone_tensors_for_torch_save, get_model_ckpt_name_for_rank
|
|
from deepspeed.accelerator import get_accelerator
|
|
from deepspeed.runtime.zero import ZeroParamStatus
|
|
from deepspeed.runtime.zero.config import DeepSpeedZeroConfig
|
|
from deepspeed.utils.torch import required_torch_version
|
|
|
|
from unit.common import DistributedTest, DistributedFixture
|
|
from unit.simple_model import *
|
|
|
|
from unit.checkpoint.common import *
|
|
|
|
import pytest
|
|
|
|
|
|
class TestZeROCheckpoint(DistributedTest):
|
|
world_size = 2
|
|
|
|
@pytest.mark.parametrize('zero_stage', [3])
|
|
def test_pipeline_checkpoint_loading(self, tmpdir, zero_stage):
|
|
config_dict = {
|
|
"train_batch_size": 2,
|
|
"optimizer": {
|
|
"type": 'Adam'
|
|
},
|
|
"zero_optimization": {
|
|
"stage": zero_stage,
|
|
"pipeline_loading_checkpoint": True,
|
|
}
|
|
}
|
|
if get_accelerator().is_bf16_supported():
|
|
config_dict["bf16"] = {"enabled": True}
|
|
elif get_accelerator().is_fp16_supported():
|
|
config_dict["fp16"] = {"enabled": True, "initial_scale_power": 8}
|
|
hidden_dim = 10
|
|
|
|
with deepspeed.zero.Init(config_dict_or_path=config_dict):
|
|
models = [SimpleModel(hidden_dim, empty_grad=False) for _ in range(2)]
|
|
|
|
checkpoint_correctness_verification(config_dict, models, hidden_dim, tmpdir, load_module_only=True)
|
|
|
|
@pytest.mark.parametrize('zero_stage, use_cpu_offload, adam_optimizer', [(0, False, 'Adam'), (1, False, 'Adam'),
|
|
(2, False, 'Adam'),
|
|
(2, True, 'deepspeed_adam'),
|
|
(3, False, 'Adam'),
|
|
(3, True, 'deepspeed_adam')])
|
|
def test_load_optimizer_state(self, tmpdir, zero_stage, use_cpu_offload, adam_optimizer):
|
|
if use_cpu_offload and not deepspeed.ops.__compatible_ops__[CPUAdamBuilder.NAME]:
|
|
pytest.skip("cpu-adam is not compatible")
|
|
|
|
config_dict = {
|
|
"train_batch_size": 2,
|
|
"steps_per_print": 1,
|
|
"optimizer": {
|
|
"type": 'Adam',
|
|
"params": {
|
|
"lr": 0.00015,
|
|
"betas": [0.8, 0.999],
|
|
"eps": 1e-8,
|
|
"weight_decay": 3e-7
|
|
}
|
|
},
|
|
"wall_clock_breakdown": True,
|
|
"zero_optimization": {
|
|
"stage": zero_stage,
|
|
"cpu_offload": use_cpu_offload
|
|
}
|
|
}
|
|
if get_accelerator().is_bf16_supported():
|
|
config_dict["bf16"] = {"enabled": True}
|
|
elif get_accelerator().is_fp16_supported():
|
|
config_dict["fp16"] = {"enabled": True, "initial_scale_power": 8}
|
|
hidden_dim = 10
|
|
|
|
if zero_stage == 3:
|
|
with deepspeed.zero.Init(config_dict_or_path=config_dict):
|
|
models = [SimpleModel(hidden_dim, empty_grad=False) for _ in range(2)]
|
|
else:
|
|
models = [SimpleModel(hidden_dim, empty_grad=False) for _ in range(2)]
|
|
|
|
checkpoint_correctness_verification(config_dict, models, hidden_dim, tmpdir, load_optimizer_states=True)
|
|
|
|
@pytest.mark.parametrize('zero_stage, use_cpu_offload, adam_optimizer', [(1, False, "Adam"), (2, False, "Adam"),
|
|
(2, True, 'deepspeed_adam'),
|
|
(3, False, 'Adam'),
|
|
(3, True, 'deepspeed_adam')])
|
|
def test_not_load_optimizer_state(self, tmpdir, zero_stage, use_cpu_offload, adam_optimizer):
|
|
if use_cpu_offload and not deepspeed.ops.__compatible_ops__[CPUAdamBuilder.NAME]:
|
|
pytest.skip("cpu-adam is not compatible")
|
|
|
|
config_dict = {
|
|
"train_batch_size": 2,
|
|
"steps_per_print": 1,
|
|
"optimizer": {
|
|
"type": 'Adam',
|
|
"params": {
|
|
"lr": 0.00015,
|
|
"betas": [0.8, 0.999],
|
|
"eps": 1e-8,
|
|
"weight_decay": 3e-7
|
|
}
|
|
},
|
|
"zero_optimization": {
|
|
"stage": zero_stage,
|
|
"cpu_offload": use_cpu_offload
|
|
}
|
|
}
|
|
if get_accelerator().is_bf16_supported():
|
|
config_dict["bf16"] = {"enabled": True}
|
|
elif get_accelerator().is_fp16_supported():
|
|
config_dict["fp16"] = {"enabled": True, "initial_scale_power": 8}
|
|
hidden_dim = 10
|
|
|
|
if zero_stage == 3:
|
|
global DeepSpeedZeroOptimizer_Stage3
|
|
from deepspeed.runtime.zero.stage3 import DeepSpeedZeroOptimizer_Stage3
|
|
with deepspeed.zero.Init(config_dict_or_path=config_dict):
|
|
models = [SimpleModel(hidden_dim, empty_grad=False) for _ in range(2)]
|
|
else:
|
|
models = [SimpleModel(hidden_dim, empty_grad=False) for _ in range(2)]
|
|
|
|
checkpoint_correctness_verification(config_dict, models, hidden_dim, tmpdir, load_optimizer_states=False)
|
|
|
|
@pytest.mark.parametrize('zero_stage', [1, 2])
|
|
def test_hybrid_optimizer_state(self, tmpdir, zero_stage):
|
|
config_dict = {
|
|
"train_micro_batch_size_per_gpu": 2,
|
|
"gradient_accumulation_steps": 2,
|
|
"steps_per_print": 1,
|
|
"zero_optimization": {
|
|
"stage": zero_stage
|
|
},
|
|
"zero_allow_untested_optimizer": True,
|
|
}
|
|
if get_accelerator().is_bf16_supported():
|
|
config_dict["bf16"] = {"enabled": True}
|
|
elif get_accelerator().is_fp16_supported():
|
|
config_dict["fp16"] = {"enabled": True, "initial_scale_power": 8}
|
|
hidden_dim = 10
|
|
models = [SimpleModel(hidden_dim=hidden_dim) for _ in range(2)]
|
|
optimizers = [HybridStateOptimizer(model.parameters()) for model in models]
|
|
|
|
checkpoint_correctness_verification(config_dict,
|
|
models=models,
|
|
base_optimizers=optimizers,
|
|
hidden_dim=hidden_dim,
|
|
tmpdir=tmpdir,
|
|
load_optimizer_states=True)
|
|
|
|
@pytest.mark.parametrize('zero_stage', [0, 1, 2, 3])
|
|
def test_load_module_only(self, tmpdir, zero_stage):
|
|
if zero_stage == 0 and get_accelerator().device_name() == "cpu":
|
|
pytest.skip("CPU Accelerator does not support this test")
|
|
config_dict = {
|
|
"train_batch_size": 2,
|
|
"optimizer": {
|
|
"type": 'Adam'
|
|
},
|
|
"zero_optimization": {
|
|
"stage": zero_stage,
|
|
}
|
|
}
|
|
if get_accelerator().is_bf16_supported():
|
|
config_dict["bf16"] = {"enabled": True}
|
|
elif get_accelerator().is_fp16_supported():
|
|
config_dict["fp16"] = {"enabled": True, "initial_scale_power": 8}
|
|
hidden_dim = 10
|
|
|
|
if zero_stage == 3:
|
|
with deepspeed.zero.Init(config_dict_or_path=config_dict):
|
|
models = [SimpleModel(hidden_dim, empty_grad=False) for _ in range(2)]
|
|
else:
|
|
models = [SimpleModel(hidden_dim, empty_grad=False) for _ in range(2)]
|
|
|
|
checkpoint_correctness_verification(config_dict, models, hidden_dim, tmpdir, load_module_only=True)
|
|
|
|
|
|
class ws4_model_checkpoint(DistributedFixture):
|
|
world_size = 4
|
|
|
|
def run(self, class_tmpdir, elastic_save, load_optim):
|
|
config_dict = {
|
|
"train_batch_size": 4,
|
|
"optimizer": {
|
|
"type": 'Adam'
|
|
},
|
|
"zero_optimization": {
|
|
"stage": 2,
|
|
"elastic_checkpoint": elastic_save
|
|
}
|
|
}
|
|
if get_accelerator().is_bf16_supported():
|
|
config_dict["bf16"] = {"enabled": True}
|
|
elif get_accelerator().is_fp16_supported():
|
|
config_dict["fp16"] = {"enabled": True, "initial_scale_power": 8}
|
|
hidden_dim = 10
|
|
model = SimpleModel(hidden_dim)
|
|
|
|
model, _, _, _ = deepspeed.initialize(config=config_dict, model=model, model_parameters=model.parameters())
|
|
data_loader = random_dataloader(model=model, total_samples=8, hidden_dim=hidden_dim, device=model.device)
|
|
for n, batch in enumerate(data_loader):
|
|
loss = model(batch[0], batch[1])
|
|
model.backward(loss)
|
|
model.step()
|
|
|
|
if load_optim:
|
|
torch.save(model.optimizer.optimizer.state_dict(), os.path.join(class_tmpdir, 'opt-state-dict'))
|
|
model.save_checkpoint(class_tmpdir)
|
|
|
|
|
|
class ws4_model_checkpoint_zeropp(DistributedFixture):
|
|
|
|
world_size = 4
|
|
|
|
def run(self, class_tmpdir):
|
|
config_dict = {
|
|
"train_batch_size": 4,
|
|
"optimizer": {
|
|
"type": 'Adam'
|
|
},
|
|
"zero_optimization": {
|
|
"stage": 3,
|
|
"zero_hpz_partition_size": 2,
|
|
}
|
|
}
|
|
|
|
hidden_dim = 10
|
|
model = SimpleModel(hidden_dim)
|
|
for param in model.parameters():
|
|
param.data = torch.ones_like(param.data, device=param.data.device, requires_grad=False)
|
|
|
|
# save model and zero checkpoint
|
|
torch.save(model.state_dict(), os.path.join(class_tmpdir, "model.pt"))
|
|
ds_model = create_deepspeed_model(config_dict=config_dict, model=model, base_optimizer=None)
|
|
ds_model.save_checkpoint(class_tmpdir)
|
|
|
|
|
|
@pytest.mark.parametrize("elastic_save", [True, False])
|
|
@pytest.mark.parametrize("elastic_load", [True, False])
|
|
@pytest.mark.parametrize("load_optim", [True, False])
|
|
class TestZeROElasticCheckpoint(DistributedTest):
|
|
world_size = 2
|
|
|
|
def test_elastic_checkpoint_fixed_dp(self, tmpdir, elastic_save, elastic_load, load_optim):
|
|
config_dict = {
|
|
"train_batch_size": 2,
|
|
"optimizer": {
|
|
"type": 'Adam'
|
|
},
|
|
"zero_optimization": {
|
|
"stage": 2,
|
|
"elastic_checkpoint": elastic_save
|
|
}
|
|
}
|
|
if get_accelerator().is_bf16_supported():
|
|
config_dict["bf16"] = {"enabled": True}
|
|
elif get_accelerator().is_fp16_supported():
|
|
config_dict["fp16"] = {"enabled": True, "initial_scale_power": 8}
|
|
hidden_dim = 10
|
|
|
|
# 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, ...
|
|
expected_mismatch_keys = [] if required_torch_version(min_version=1.4) else ['params']
|
|
models = [SimpleModel(hidden_dim) for _ in range(2)]
|
|
model, _, _, _ = deepspeed.initialize(config=config_dict,
|
|
model=models[0],
|
|
model_parameters=models[0].parameters())
|
|
run_steps = 8
|
|
data_loader = random_dataloader(model=model,
|
|
total_samples=run_steps,
|
|
hidden_dim=hidden_dim,
|
|
device=model.device)
|
|
for n, batch in enumerate(data_loader):
|
|
loss = model(batch[0], batch[1])
|
|
model.backward(loss)
|
|
model.step()
|
|
if load_optim:
|
|
opt_state_dict_file = f'opt-state-dict_rank{dist.get_rank()}'
|
|
torch.save(model.optimizer.optimizer.state_dict(), os.path.join(tmpdir, opt_state_dict_file))
|
|
model.save_checkpoint(tmpdir)
|
|
|
|
config_dict["zero_optimization"]["elastic_checkpoint"] = elastic_load
|
|
model, _, _, _ = deepspeed.initialize(config=config_dict,
|
|
model=models[1],
|
|
model_parameters=models[1].parameters())
|
|
model.load_checkpoint(tmpdir, load_optimizer_states=load_optim)
|
|
|
|
if load_optim:
|
|
saved_sd = torch.load(os.path.join(tmpdir, opt_state_dict_file), weights_only=False)
|
|
curr_sd = model.optimizer.optimizer.state_dict()
|
|
compare_opt_state_dicts(curr_sd, saved_sd, expected_mismatch_keys)
|
|
|
|
data_loader = random_dataloader(model=model, total_samples=8, hidden_dim=hidden_dim, device=model.device)
|
|
for n, batch in enumerate(data_loader):
|
|
loss = model(batch[0], batch[1])
|
|
model.backward(loss)
|
|
model.step()
|
|
|
|
def test_elastic_checkpoint_change_dp(self, ws4_model_checkpoint, class_tmpdir, elastic_save, elastic_load,
|
|
load_optim):
|
|
config_dict = {
|
|
"train_batch_size": 4,
|
|
"optimizer": {
|
|
"type": 'Adam'
|
|
},
|
|
"zero_optimization": {
|
|
"stage": 2,
|
|
"elastic_checkpoint": elastic_load
|
|
}
|
|
}
|
|
if get_accelerator().is_bf16_supported():
|
|
config_dict["bf16"] = {"enabled": True}
|
|
elif get_accelerator().is_fp16_supported():
|
|
config_dict["fp16"] = {"enabled": True, "initial_scale_power": 8}
|
|
hidden_dim = 10
|
|
model = SimpleModel(hidden_dim)
|
|
|
|
# Load checkpoint with dp world size = 2
|
|
model, _, _, _ = deepspeed.initialize(config=config_dict, model=model, model_parameters=model.parameters())
|
|
if load_optim:
|
|
with pytest.raises(deepspeed.runtime.zero.utils.ZeRORuntimeException):
|
|
model.load_checkpoint(class_tmpdir, load_optimizer_states=load_optim)
|
|
else:
|
|
model.load_checkpoint(class_tmpdir, load_optimizer_states=load_optim)
|
|
|
|
|
|
class TestZeROSaveLoadEdgeCase(DistributedTest):
|
|
world_size = 2
|
|
|
|
@pytest.mark.parametrize('zero_stage', [0, 1, 2, 3])
|
|
def test_immediate_save_load(self, tmpdir, zero_stage):
|
|
config_dict = {
|
|
"train_batch_size": 4,
|
|
"optimizer": {
|
|
"type": 'Adam'
|
|
},
|
|
"zero_optimization": {
|
|
"stage": zero_stage,
|
|
}
|
|
}
|
|
if get_accelerator().is_bf16_supported():
|
|
config_dict["bf16"] = {"enabled": True}
|
|
elif get_accelerator().is_fp16_supported():
|
|
config_dict["fp16"] = {"enabled": True, "initial_scale_power": 8}
|
|
hidden_dim = 10
|
|
model = SimpleModel(hidden_dim)
|
|
|
|
ds_model = create_deepspeed_model(config_dict=config_dict, model=model, base_optimizer=None)
|
|
ds_model.save_checkpoint(tmpdir)
|
|
ds_model.load_checkpoint(tmpdir,
|
|
load_optimizer_states=False,
|
|
load_lr_scheduler_states=False,
|
|
load_module_only=False)
|
|
|
|
@pytest.mark.parametrize('zero_stage', [0, 1, 2, 3])
|
|
def test_load_immediate_save(self, tmpdir, zero_stage):
|
|
if zero_stage == 0 and get_accelerator().device_name() == "cpu":
|
|
pytest.skip("CPU Accelerator does not support this test")
|
|
config_dict = {
|
|
"train_batch_size": 4,
|
|
"optimizer": {
|
|
"type": 'Adam'
|
|
},
|
|
"zero_optimization": {
|
|
"stage": zero_stage,
|
|
}
|
|
}
|
|
if get_accelerator().is_bf16_supported():
|
|
config_dict["bf16"] = {"enabled": True}
|
|
elif get_accelerator().is_fp16_supported():
|
|
config_dict["fp16"] = {"enabled": True, "initial_scale_power": 8}
|
|
hidden_dim = 10
|
|
model = SimpleModel(hidden_dim)
|
|
|
|
# 1. pretrain a model and save it
|
|
ds_model = create_deepspeed_model(config_dict=config_dict, model=model, base_optimizer=None)
|
|
data_loader = random_dataloader(model=ds_model, total_samples=1, hidden_dim=hidden_dim, device=ds_model.device)
|
|
for _, batch in enumerate(data_loader):
|
|
loss = ds_model(batch[0], batch[1])
|
|
ds_model.backward(loss)
|
|
ds_model.step()
|
|
|
|
ds_model.empty_partition_cache()
|
|
ds_model.save_checkpoint(tmpdir)
|
|
|
|
# 2. load and immediately save a model with a fresh ds engine
|
|
ds_model = create_deepspeed_model(config_dict=config_dict, model=model, base_optimizer=None)
|
|
ds_model.load_checkpoint(tmpdir,
|
|
load_optimizer_states=False,
|
|
load_lr_scheduler_states=False,
|
|
load_module_only=False)
|
|
ds_model.save_checkpoint(tmpdir)
|
|
|
|
@pytest.mark.parametrize('zero_stage', [0, 1, 2, 3])
|
|
def test_save_before_accum_grad_is_done(self, tmpdir, zero_stage):
|
|
config_dict = {
|
|
"optimizer": {
|
|
"type": 'Adam'
|
|
},
|
|
"zero_optimization": {
|
|
"stage": zero_stage,
|
|
"stage3_gather_fp16_weights_on_model_save": True,
|
|
},
|
|
"gradient_accumulation_steps": 2,
|
|
"train_micro_batch_size_per_gpu": 1,
|
|
"train_batch_size": 4,
|
|
}
|
|
if get_accelerator().is_bf16_supported():
|
|
config_dict["bf16"] = {"enabled": True}
|
|
elif get_accelerator().is_fp16_supported():
|
|
config_dict["fp16"] = {"enabled": True, "initial_scale_power": 8}
|
|
hidden_dim = 10
|
|
model = SimpleModel(hidden_dim)
|
|
|
|
# This test reproduces a bug where one tries to retrieve a 16bit model before grad_accum
|
|
# cycle was completed.
|
|
# So we config grad_accum=2 and step only once and save_16bit_model
|
|
ds_model = create_deepspeed_model(config_dict=config_dict, model=model, base_optimizer=None)
|
|
|
|
data_loader = random_dataloader(model=ds_model, total_samples=2, hidden_dim=hidden_dim, device=ds_model.device)
|
|
|
|
batch = next(iter(data_loader))
|
|
loss = ds_model(batch[0], batch[1])
|
|
ds_model.backward(loss)
|
|
ds_model.step()
|
|
|
|
ds_model.empty_partition_cache()
|
|
|
|
# we stepped only once, and now save 16bit model before gradient_accumulation_steps=2 is complete
|
|
ds_model.save_16bit_model(tmpdir, "model.pt")
|
|
|
|
# let's test just as well that we can save the checkpoint too
|
|
ds_model.save_checkpoint(tmpdir)
|
|
|
|
|
|
class TestZeROCheckpointFrozenWeights(DistributedTest):
|
|
world_size = 2
|
|
|
|
@pytest.mark.parametrize('zero_stage', [1, 2, 3])
|
|
def test_load_optimizer_state(self, tmpdir, zero_stage):
|
|
|
|
config_dict = {
|
|
"train_batch_size": 2,
|
|
"steps_per_print": 1,
|
|
"optimizer": {
|
|
"type": 'Adam',
|
|
"params": {
|
|
"lr": 0.00015,
|
|
"betas": [0.8, 0.999],
|
|
"eps": 1e-8,
|
|
"weight_decay": 3e-7
|
|
}
|
|
},
|
|
"wall_clock_breakdown": True,
|
|
"zero_optimization": {
|
|
"stage": zero_stage
|
|
}
|
|
}
|
|
if get_accelerator().is_bf16_supported():
|
|
config_dict["bf16"] = {"enabled": True}
|
|
elif get_accelerator().is_fp16_supported():
|
|
config_dict["fp16"] = {"enabled": True, "initial_scale_power": 8}
|
|
hidden_dim = 10
|
|
|
|
with deepspeed.zero.Init(enabled=zero_stage == 3, config_dict_or_path=config_dict):
|
|
models = [SimpleFrozenModel(hidden_dim, empty_grad=False) for _ in range(2)]
|
|
|
|
checkpoint_correctness_verification(config_dict, models, hidden_dim, tmpdir, load_optimizer_states=True)
|
|
|
|
@pytest.mark.parametrize('zero_stage', [1, 2, 3])
|
|
def test_not_load_optimizer_state(self, tmpdir, zero_stage):
|
|
|
|
config_dict = {
|
|
"train_batch_size": 2,
|
|
"steps_per_print": 1,
|
|
"optimizer": {
|
|
"type": 'Adam',
|
|
"params": {
|
|
"lr": 0.00015,
|
|
"betas": [0.8, 0.999],
|
|
"eps": 1e-8,
|
|
"weight_decay": 3e-7
|
|
}
|
|
},
|
|
"zero_optimization": {
|
|
"stage": zero_stage
|
|
}
|
|
}
|
|
if get_accelerator().is_bf16_supported():
|
|
config_dict["bf16"] = {"enabled": True}
|
|
elif get_accelerator().is_fp16_supported():
|
|
config_dict["fp16"] = {"enabled": True}
|
|
hidden_dim = 10
|
|
|
|
with deepspeed.zero.Init(enabled=zero_stage == 3, config_dict_or_path=config_dict):
|
|
models = [SimpleFrozenModel(hidden_dim, empty_grad=False) for _ in range(2)]
|
|
|
|
checkpoint_correctness_verification(config_dict, models, hidden_dim, tmpdir, load_optimizer_states=False)
|
|
|
|
@pytest.mark.parametrize('zero_stage', [1, 2, 3])
|
|
def test_load_module_only(self, tmpdir, zero_stage):
|
|
config_dict = {
|
|
"train_batch_size": 2,
|
|
"optimizer": {
|
|
"type": 'Adam'
|
|
},
|
|
"zero_optimization": {
|
|
"stage": zero_stage,
|
|
}
|
|
}
|
|
if get_accelerator().is_bf16_supported():
|
|
config_dict["bf16"] = {"enabled": True}
|
|
elif get_accelerator().is_fp16_supported():
|
|
config_dict["fp16"] = {"enabled": True, "initial_scale_power": 8}
|
|
hidden_dim = 10
|
|
|
|
with deepspeed.zero.Init(enabled=zero_stage == 3, config_dict_or_path=config_dict):
|
|
models = [SimpleFrozenModel(hidden_dim, empty_grad=False) for _ in range(2)]
|
|
|
|
checkpoint_correctness_verification(config_dict, models, hidden_dim, tmpdir, load_module_only=True)
|
|
|
|
@pytest.mark.parametrize('zero_stage', [1, 2])
|
|
def test_save_exclude_frozen_weights(self, tmpdir, zero_stage):
|
|
world_size = 1
|
|
config_dict = {
|
|
"train_micro_batch_size_per_gpu": 1,
|
|
"optimizer": {
|
|
"type": 'Adam'
|
|
},
|
|
"zero_optimization": {
|
|
"stage": zero_stage,
|
|
}
|
|
}
|
|
if get_accelerator().is_bf16_supported():
|
|
config_dict["bf16"] = {"enabled": True}
|
|
elif get_accelerator().is_fp16_supported():
|
|
config_dict["fp16"] = {"enabled": True, "initial_scale_power": 8}
|
|
hidden_dim = 10
|
|
|
|
model = SimpleFrozenModel(hidden_dim, empty_grad=False)
|
|
|
|
ds_engine, _, _, _ = deepspeed.initialize(model=model, model_parameters=model.parameters(), config=config_dict)
|
|
|
|
# Validate backwards-compatibility of including frozen parameters in checkpoint
|
|
all_ckpt_folder = os.path.join(tmpdir, 'all_params')
|
|
ds_engine.save_checkpoint(all_ckpt_folder)
|
|
all_params_ckpt_file = get_model_ckpt_name_for_rank(os.path.join(all_ckpt_folder, 'global_step0'), '00')
|
|
loaded_all_param_model = torch.load(all_params_ckpt_file, weights_only=False)['module']
|
|
all_param_names = set([n for n, p in model.named_parameters()])
|
|
assert set(loaded_all_param_model.keys()) == all_param_names
|
|
|
|
# Validate exclusion of frozen parameters
|
|
trainable_ckpt_folder = os.path.join(tmpdir, 'no_frozen_params')
|
|
ds_engine.save_checkpoint(trainable_ckpt_folder, exclude_frozen_parameters=True)
|
|
|
|
trainable_ckpt_file = get_model_ckpt_name_for_rank(os.path.join(trainable_ckpt_folder, 'global_step0'), '00')
|
|
|
|
# Excluding frozen parameters should reduce checkpoint size
|
|
assert os.path.getsize(all_params_ckpt_file) > os.path.getsize(trainable_ckpt_file)
|
|
|
|
loaded_trainable_param_model = torch.load(trainable_ckpt_file, weights_only=False)['module']
|
|
frozen_param_names = set([n for n, p in model.named_parameters() if not p.requires_grad])
|
|
loaded_trainable_param_names = set(loaded_trainable_param_model.keys())
|
|
overlap_names = set.intersection(loaded_trainable_param_names, frozen_param_names)
|
|
assert len(overlap_names) == 0
|
|
|
|
trainable_param_names = set([n for n, p in model.named_parameters() if p.requires_grad])
|
|
assert loaded_trainable_param_names == trainable_param_names
|
|
|
|
@pytest.mark.parametrize('zero_stage', [1, 2])
|
|
def test_save_exclude_custom_frozen_weights(self, tmpdir, zero_stage):
|
|
world_size = 1
|
|
config_dict = {
|
|
"train_micro_batch_size_per_gpu": 1,
|
|
"optimizer": {
|
|
"type": 'Adam'
|
|
},
|
|
"zero_optimization": {
|
|
"stage": zero_stage,
|
|
}
|
|
}
|
|
if get_accelerator().is_bf16_supported():
|
|
config_dict["bf16"] = {"enabled": True}
|
|
elif get_accelerator().is_fp16_supported():
|
|
config_dict["fp16"] = {"enabled": True, "initial_scale_power": 8}
|
|
hidden_dim = 10
|
|
|
|
model = SimpleFrozenModel(hidden_dim, empty_grad=False)
|
|
|
|
ds_engine, _, _, _ = deepspeed.initialize(model=model, model_parameters=model.parameters(), config=config_dict)
|
|
|
|
# Validate custom state_dict model
|
|
state_dict_bk = model.state_dict
|
|
model.state_dict = model.custom_state_dict
|
|
custom_state_dict_ckpt_folder = os.path.join(tmpdir, 'custom_state_dict')
|
|
ds_engine.save_checkpoint(custom_state_dict_ckpt_folder, exclude_frozen_parameters=True)
|
|
|
|
custom_state_dict_ckpt_file = get_model_ckpt_name_for_rank(
|
|
os.path.join(custom_state_dict_ckpt_folder, 'global_step0'), '00')
|
|
loaded_custom_state_dict_param_model = torch.load(custom_state_dict_ckpt_file, weights_only=False)['module']
|
|
loaded_custom_state_dict_param_names = set(loaded_custom_state_dict_param_model.keys())
|
|
|
|
custom_state_dict_param_names = set([k for k, v in model.state_dict().items()])
|
|
trainable_param_names = set([n for n, p in model.named_parameters() if p.requires_grad])
|
|
overlap_names = set.intersection(custom_state_dict_param_names, trainable_param_names)
|
|
|
|
assert loaded_custom_state_dict_param_names == overlap_names
|
|
|
|
model.state_dict = state_dict_bk
|
|
|
|
|
|
class TestSaveTensorClone(DistributedTest):
|
|
world_size = 1
|
|
|
|
@pytest.mark.parametrize('zero_stage', [1, 2])
|
|
@pytest.mark.parametrize('use_cpu_device', [True, False])
|
|
def test_save_tensor_clone(self, tmpdir, zero_stage, use_cpu_device):
|
|
|
|
config_dict = {
|
|
"optimizer": {
|
|
"type": "AdamW",
|
|
},
|
|
"zero_optimization": {
|
|
"stage": zero_stage
|
|
},
|
|
"train_batch_size": 1,
|
|
"train_micro_batch_size_per_gpu": 1
|
|
}
|
|
hidden_dim = 1024
|
|
model = SimpleModel(hidden_dim, nlayers=4).half()
|
|
ref_model_state_dict = model.state_dict()
|
|
|
|
ds_engine, _, _, _ = deepspeed.initialize(model=model, config_params=config_dict)
|
|
clone_device = torch.device('cpu') if use_cpu_device else get_accelerator().current_device()
|
|
clone_state_dict = clone_tensors_for_torch_save(ds_engine.module.state_dict())
|
|
compare_state_dicts(ref_model_state_dict, clone_state_dict)
|
|
|
|
ref_ckpt_file = os.path.join(tmpdir, 'ref_ckpt.pt')
|
|
torch.save(ref_model_state_dict, ref_ckpt_file)
|
|
clone_ckpt_file = os.path.join(tmpdir, 'clone_ckpt.pt')
|
|
torch.save(clone_state_dict, clone_ckpt_file)
|
|
|
|
compare_state_dicts(torch.load(ref_ckpt_file, weights_only=False),
|
|
torch.load(clone_ckpt_file, weights_only=False))
|
|
|
|
|
|
def test_elastic_checkpoint_is_deprecated_for_zero3(monkeypatch):
|
|
warning_messages = []
|
|
|
|
def mock_logger_warning(message, *args, **kwargs):
|
|
warning_messages.append(message)
|
|
|
|
monkeypatch.setattr("deepspeed.utils.logger.warning", mock_logger_warning)
|
|
|
|
DeepSpeedZeroConfig(stage=3, elastic_checkpoint=True)
|
|
|
|
assert any("elastic checkpointing is deprecated" in str(message).lower() for message in warning_messages)
|
|
|
|
|
|
class TestZeRONonDistributed(DistributedTest):
|
|
world_size = 1
|
|
# This test calls deepspeed.initialize(), so use the harness' file-store
|
|
# initialization instead of env:// TCP rendezvous ports under xdist.
|
|
init_distributed = True
|
|
|
|
@pytest.mark.parametrize('zero_stage', [1, 2, 3])
|
|
def test_chmod_exception_handling(self, monkeypatch, zero_stage):
|
|
|
|
config_dict = {
|
|
"optimizer": {
|
|
"type": "AdamW"
|
|
},
|
|
"train_batch_size": 1,
|
|
"zero_optimization": {
|
|
"stage": zero_stage
|
|
}
|
|
}
|
|
args = SimpleNamespace(local_rank=0)
|
|
net = SimpleModel(hidden_dim=4)
|
|
engine, _, _, _ = deepspeed.initialize(args=args,
|
|
config=config_dict,
|
|
model=net,
|
|
model_parameters=net.parameters())
|
|
|
|
log_called = False
|
|
|
|
def mock_logger_info(message, *args, **kwargs):
|
|
nonlocal log_called
|
|
log_called = True
|
|
|
|
monkeypatch.setattr("deepspeed.utils.logger.info", mock_logger_info)
|
|
"""
|
|
This is presented for use-cases like Azure Storage File Share (where permissions are not allowed)
|
|
We use a fake file for this test (file not existing would present a similar issue as not being able to chmod)
|
|
"""
|
|
fake_recovery_script_dst = os.path.join("tmp", "zero_to_fp32.py")
|
|
engine._change_recovery_script_permissions(fake_recovery_script_dst)
|
|
|
|
assert log_called, "Expected deepspeed.utils.logger.info to be called."
|
|
|
|
|
|
class TestZeROPPLoadCheckpoint(DistributedTest):
|
|
|
|
world_size = 4
|
|
|
|
def test_load_zeropp_model(self, ws4_model_checkpoint_zeropp, class_tmpdir):
|
|
config_dict = {
|
|
"train_batch_size": 4,
|
|
"optimizer": {
|
|
"type": 'Adam'
|
|
},
|
|
"zero_optimization": {
|
|
"stage": 3,
|
|
"zero_hpz_partition_size": 2,
|
|
"stage3_param_persistence_threshold": 1
|
|
}
|
|
}
|
|
|
|
# Init model and load saved model
|
|
hidden_dim = 10
|
|
with deepspeed.zero.Init(config_dict_or_path=config_dict):
|
|
model = SimpleModel(hidden_dim)
|
|
ds_model = create_deepspeed_model(config_dict=config_dict, model=model, base_optimizer=None)
|
|
|
|
with deepspeed.zero.GatheredParameters(ds_model.module.parameters(), modifier_rank=0):
|
|
if dist.get_rank() == 0:
|
|
state_dict = torch.load(os.path.join(class_tmpdir, "model.pt"))
|
|
ds_model.module.load_state_dict(state_dict)
|
|
|
|
# Check the parameters after gather
|
|
params_to_gather = [p for p in ds_model.module.parameters() if p.ds_status == ZeroParamStatus.NOT_AVAILABLE]
|
|
if len(params_to_gather) > 0:
|
|
handle = params_to_gather[0].all_gather_coalesced(params_to_gather)
|
|
handle.wait()
|
|
for ds_param in params_to_gather:
|
|
for v in ds_param.data.cpu().flatten().numpy():
|
|
assert v == 1.0
|
|
|
|
def test_load_zeropp_checkpoint(self, ws4_model_checkpoint_zeropp, class_tmpdir):
|
|
config_dict = {
|
|
"train_batch_size": 4,
|
|
"optimizer": {
|
|
"type": 'Adam'
|
|
},
|
|
"zero_optimization": {
|
|
"stage": 3,
|
|
"zero_hpz_partition_size": 2,
|
|
"stage3_param_persistence_threshold": 1
|
|
}
|
|
}
|
|
|
|
# Init model and load zero checkpoint
|
|
hidden_dim = 10
|
|
model = SimpleModel(hidden_dim)
|
|
ds_model = create_deepspeed_model(config_dict=config_dict, model=model, base_optimizer=None)
|
|
ds_model.load_checkpoint(class_tmpdir,
|
|
load_optimizer_states=True,
|
|
load_lr_scheduler_states=False,
|
|
load_module_only=False)
|
|
|
|
# Check the parameters after gather
|
|
params_to_gather = [p for p in ds_model.module.parameters() if p.ds_status == ZeroParamStatus.NOT_AVAILABLE]
|
|
if len(params_to_gather) > 0:
|
|
handle = params_to_gather[0].all_gather_coalesced(params_to_gather)
|
|
handle.wait()
|
|
for ds_param in params_to_gather:
|
|
for v in ds_param.data.cpu().flatten().numpy():
|
|
assert v == 1.0
|