wehub-resource-sync
509 lines
19 KiB
Python
509 lines
19 KiB
Python
# Copyright (c) Microsoft Corporation.
|
|
# SPDX-License-Identifier: Apache-2.0
|
|
|
|
# DeepSpeed Team
|
|
|
|
import torch
|
|
import deepspeed
|
|
import pytest
|
|
from deepspeed.ops.adam import FusedAdam
|
|
from unit.common import DistributedTest
|
|
from deepspeed.ops.op_builder import CPUAdamBuilder
|
|
from unit.simple_model import SimpleModel, SimpleOptimizer, random_dataloader
|
|
from unit.util import bf16_required_version_check
|
|
from deepspeed import comm as dist
|
|
from deepspeed.accelerator import get_accelerator
|
|
from unit.v1.zero.test_zero_user_backward import (initialize_distributed, create_ddp_model, collect_ddp_gradients,
|
|
collect_gradients_safe, compare_gradients)
|
|
|
|
|
|
class TestAdamBF16ZeroOneCycleCompatibility(DistributedTest):
|
|
world_size = 1
|
|
|
|
def test(self, zero_stage=2, use_cpu_offload=False):
|
|
if 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"
|
|
)
|
|
|
|
if use_cpu_offload and not deepspeed.ops.__compatible_ops__[CPUAdamBuilder.NAME]:
|
|
pytest.skip("cpu-adam is not compatible")
|
|
|
|
config_dict = {
|
|
"train_micro_batch_size_per_gpu": 1,
|
|
"steps_per_print": 1,
|
|
"optimizer": {
|
|
"type": "Adam",
|
|
"params": {
|
|
"lr": 0.00015
|
|
}
|
|
},
|
|
"scheduler": {
|
|
"type": "OneCycle",
|
|
"params": {
|
|
"cycle_first_step_size": 16000,
|
|
"cycle_first_stair_count": 8000,
|
|
"decay_step_size": 16000,
|
|
"cycle_min_lr": 1e-06,
|
|
"cycle_max_lr": 3e-05,
|
|
"decay_lr_rate": 1e-07,
|
|
"cycle_min_mom": 0.85,
|
|
"cycle_max_mom": 0.99,
|
|
"decay_mom_rate": 0.0
|
|
}
|
|
},
|
|
"fp16": {
|
|
"enabled": False
|
|
},
|
|
"bf16": {
|
|
"enabled": True
|
|
},
|
|
"zero_optimization": {
|
|
"stage": zero_stage,
|
|
"cpu_offload": use_cpu_offload
|
|
}
|
|
}
|
|
|
|
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=50,
|
|
hidden_dim=hidden_dim,
|
|
device=model.device,
|
|
dtype=torch.bfloat16)
|
|
for n, batch in enumerate(data_loader):
|
|
loss = model(batch[0], batch[1])
|
|
model.backward(loss)
|
|
model.step()
|
|
|
|
|
|
class TestZeroAllowUntestedOptimizer(DistributedTest):
|
|
world_size = 1
|
|
|
|
def test(self, zero_stage=2, use_cpu_offload=False):
|
|
if 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"
|
|
)
|
|
|
|
if use_cpu_offload and not deepspeed.ops.__compatible_ops__[CPUAdamBuilder.NAME]:
|
|
pytest.skip("cpu-adam is not compatible")
|
|
|
|
config_dict = {
|
|
"train_micro_batch_size_per_gpu": 4,
|
|
"steps_per_print": 1,
|
|
"fp16": {
|
|
"enabled": False,
|
|
},
|
|
"bf16": {
|
|
"enabled": True
|
|
},
|
|
"zero_optimization": {
|
|
"stage": zero_stage,
|
|
"cpu_offload": use_cpu_offload
|
|
},
|
|
"zero_allow_untested_optimizer": False
|
|
}
|
|
|
|
hidden_dim = 10
|
|
model = SimpleModel(hidden_dim)
|
|
optimizer = SimpleOptimizer(model.parameters())
|
|
with pytest.raises(AssertionError):
|
|
model, optim, _, _ = deepspeed.initialize(config=config_dict,
|
|
model=model,
|
|
optimizer=optimizer,
|
|
model_parameters=model.parameters())
|
|
|
|
|
|
class TestZeroEmptyPartition(DistributedTest):
|
|
world_size = 3
|
|
|
|
def test(self, zero_stage=2, use_cpu_offload=False):
|
|
if 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"
|
|
)
|
|
|
|
if use_cpu_offload and not deepspeed.ops.__compatible_ops__[CPUAdamBuilder.NAME]:
|
|
pytest.skip("cpu-adam is not compatible")
|
|
|
|
if zero_stage == 3:
|
|
pytest.skip("skip for now")
|
|
|
|
config_dict = {
|
|
"train_micro_batch_size_per_gpu": 1,
|
|
"gradient_accumulation_steps": 1,
|
|
"fp16": {
|
|
"enabled": False
|
|
},
|
|
"bf16": {
|
|
"enabled": True
|
|
},
|
|
"optimizer": {
|
|
"type": "Adam",
|
|
"params": {
|
|
"lr": 0.00015
|
|
}
|
|
},
|
|
"zero_optimization": {
|
|
"stage": zero_stage,
|
|
"cpu_offload": use_cpu_offload,
|
|
"reduce_bucket_size": 100,
|
|
"allgather_bucket_size": 100
|
|
}
|
|
}
|
|
|
|
hidden_dim = 1
|
|
model = SimpleModel(hidden_dim)
|
|
|
|
# Ensure model has 2 parameters, to cause empty partition with DP=3
|
|
assert len(list(model.parameters())) == 2
|
|
model, _, _, _ = deepspeed.initialize(config=config_dict, model=model, model_parameters=model.parameters())
|
|
|
|
# Now make sure things work..
|
|
data_loader = random_dataloader(model=model,
|
|
total_samples=1,
|
|
hidden_dim=hidden_dim,
|
|
device=model.device,
|
|
dtype=torch.bfloat16)
|
|
for n, batch in enumerate(data_loader):
|
|
loss = model(batch[0], batch[1])
|
|
model.backward(loss)
|
|
model.step()
|
|
|
|
|
|
@pytest.mark.parametrize("optimizer_constructor", [torch.optim.Adam, FusedAdam])
|
|
class TestZeroSupportedClientOptimizer(DistributedTest):
|
|
world_size = 1
|
|
|
|
def test(self, optimizer_constructor, zero_stage=2):
|
|
if 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"
|
|
)
|
|
|
|
config_dict = {
|
|
"train_micro_batch_size_per_gpu": 2,
|
|
"steps_per_print": 1,
|
|
"fp16": {
|
|
"enabled": False
|
|
},
|
|
"bf16": {
|
|
"enabled": True
|
|
},
|
|
"zero_optimization": {
|
|
"stage": zero_stage
|
|
}
|
|
}
|
|
hidden_dim = 10
|
|
|
|
model = SimpleModel(hidden_dim)
|
|
client_optimizer = optimizer_constructor(params=model.parameters())
|
|
model, _, _, _ = deepspeed.initialize(config=config_dict, model=model, optimizer=client_optimizer)
|
|
|
|
|
|
class TestZero2ReduceScatterOff(DistributedTest):
|
|
world_size = 2
|
|
|
|
def test(self):
|
|
if 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"
|
|
)
|
|
|
|
config_dict = {
|
|
"train_micro_batch_size_per_gpu": 2,
|
|
"steps_per_print": 1,
|
|
"optimizer": {
|
|
"type": "Adam",
|
|
"params": {
|
|
"lr": 0.00015
|
|
}
|
|
},
|
|
"gradient_clipping": 1.0,
|
|
"zero_optimization": {
|
|
"stage": 2,
|
|
"contiguous_gradients": True,
|
|
"allgather_bucket_size": 2000000000,
|
|
"reduce_bucket_size": 200000000,
|
|
"overlap_comm": False,
|
|
"reduce_scatter": False
|
|
},
|
|
"fp16": {
|
|
"enabled": False
|
|
},
|
|
"bf16": {
|
|
"enabled": True
|
|
}
|
|
}
|
|
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=50,
|
|
hidden_dim=hidden_dim,
|
|
device=model.device,
|
|
dtype=torch.bfloat16)
|
|
for n, batch in enumerate(data_loader):
|
|
loss = model(batch[0], batch[1])
|
|
model.backward(loss)
|
|
model.step()
|
|
|
|
|
|
class TestZeroEmptyGrad(DistributedTest):
|
|
world_size = 1
|
|
|
|
def test(self, stage=2):
|
|
if 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"
|
|
)
|
|
|
|
config_dict = {
|
|
"train_micro_batch_size_per_gpu": 1,
|
|
"steps_per_print": 1,
|
|
"fp16": {
|
|
"enabled": False
|
|
},
|
|
"bf16": {
|
|
"enabled": True
|
|
},
|
|
"zero_optimization": {
|
|
"stage": stage
|
|
}
|
|
}
|
|
hidden_dim = 10
|
|
|
|
model = SimpleModel(hidden_dim)
|
|
optimizer = torch.optim.Adam(model.parameters())
|
|
model, _, _, _ = deepspeed.initialize(config=config_dict, model=model, optimizer=optimizer)
|
|
data_loader = random_dataloader(model=model,
|
|
total_samples=50,
|
|
hidden_dim=hidden_dim,
|
|
device=model.device,
|
|
dtype=torch.bfloat16)
|
|
for n, batch in enumerate(data_loader):
|
|
loss = model(batch[0], batch[1])
|
|
model.backward(loss)
|
|
model.step()
|
|
|
|
|
|
@pytest.mark.parametrize("comp_type", [torch.float16, torch.bfloat16, torch.float], ids=["fp16", "bf16", "fp32"])
|
|
@pytest.mark.parametrize("comm_type", [torch.float16, torch.bfloat16, None], ids=["fp16", "bf16", "default"])
|
|
class TestZeroDtypeCocktail(DistributedTest):
|
|
world_size = 2
|
|
|
|
def test(self, comp_type, comm_type):
|
|
if comp_type == torch.bfloat16 or comm_type == torch.bfloat16:
|
|
if 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"
|
|
)
|
|
|
|
if comp_type == torch.float16 or comm_type == torch.float16:
|
|
if not get_accelerator().is_fp16_supported():
|
|
pytest.skip("fp16 is not supported")
|
|
|
|
type_str = {torch.float16: "fp16", torch.bfloat16: "bf16"}
|
|
|
|
config_dict = {
|
|
"train_micro_batch_size_per_gpu": 2,
|
|
"steps_per_print": 1,
|
|
"fp16": {
|
|
"enabled": comp_type == torch.float16
|
|
},
|
|
"bf16": {
|
|
"enabled": comp_type == torch.bfloat16
|
|
},
|
|
"zero_optimization": {
|
|
"stage": 2
|
|
},
|
|
}
|
|
if comm_type is not None:
|
|
config_dict["communication_data_type"] = type_str[comm_type]
|
|
else:
|
|
comm_type = comp_type
|
|
hidden_dim = 10
|
|
|
|
model = SimpleModel(hidden_dim)
|
|
optimizer = torch.optim.Adam(model.parameters())
|
|
model, _, _, _ = deepspeed.initialize(config=config_dict, model=model, optimizer=optimizer)
|
|
data_loader = random_dataloader(model=model,
|
|
total_samples=2,
|
|
hidden_dim=hidden_dim,
|
|
device=model.device,
|
|
dtype=comp_type)
|
|
|
|
def custom_reduce(tensor, dst, op=dist.ReduceOp.SUM, group=None, async_op=False):
|
|
assert tensor.dtype == comm_type
|
|
return orig_torch_reduce(tensor, dst, op, group, async_op)
|
|
|
|
orig_torch_reduce = dist.reduce
|
|
dist.reduce = custom_reduce
|
|
for n, batch in enumerate(data_loader):
|
|
loss = model(batch[0], batch[1])
|
|
model.backward(loss)
|
|
model.step()
|
|
dist.reduce = orig_torch_reduce
|
|
|
|
|
|
@pytest.mark.parametrize("bf16_optimizer_states,use_cpu_offload,zero_stage", [
|
|
pytest.param(False, True, 1, id="zero_stage_1_cpu_offload"),
|
|
pytest.param(True, False, 1, id="zero_stage_1_bf16_opt_states_True"),
|
|
pytest.param(True, True, 1, id="zero_stage_1_bf16_opt_states_cpu_offload"),
|
|
pytest.param(False, True, 2, id="zero_stage_2_cpu_offload"),
|
|
pytest.param(True, False, 2, id="zero_stage_2_bf16_opt_states_True"),
|
|
pytest.param(True, True, 2, id="zero_stage_2_bf16_opt_states_cpu_offload"),
|
|
pytest.param(False, True, 3, id="zero_stage_3_cpu_offload"),
|
|
pytest.param(True, False, 3, id="zero_stage_3_bf16_opt_states_True"),
|
|
pytest.param(True, True, 3, id="zero_stage_3_bf16_opt_states_cpu_offload"),
|
|
])
|
|
class TestBF16MasterWeightsGradients(DistributedTest):
|
|
world_size = 2
|
|
|
|
def test_gradients_match_ddp(self, bf16_optimizer_states, use_cpu_offload, zero_stage):
|
|
if 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"
|
|
)
|
|
|
|
if use_cpu_offload and not deepspeed.ops.__compatible_ops__[CPUAdamBuilder.NAME]:
|
|
pytest.skip("cpu-adam is not compatible")
|
|
|
|
hidden_dim = 6
|
|
lr = 1e-3
|
|
seed = 123
|
|
|
|
device, rank, dtype = initialize_distributed()
|
|
|
|
model_ddp, optimizer_ddp = create_ddp_model(SimpleModel,
|
|
device,
|
|
rank,
|
|
dtype,
|
|
seed=seed,
|
|
lr=lr,
|
|
hidden_dim=hidden_dim,
|
|
nlayers=2)
|
|
|
|
torch.manual_seed(seed)
|
|
ds_model = SimpleModel(hidden_dim, nlayers=2)
|
|
|
|
bf16_config = {
|
|
"enabled": True,
|
|
"bf16_master_weights_and_grads": True,
|
|
}
|
|
if bf16_optimizer_states:
|
|
bf16_config["bf16_optimizer_states"] = True
|
|
|
|
zero_config = {"stage": zero_stage}
|
|
if use_cpu_offload:
|
|
zero_config["cpu_offload"] = True
|
|
|
|
config_dict = {
|
|
"train_micro_batch_size_per_gpu": 2,
|
|
"gradient_accumulation_steps": 1,
|
|
"steps_per_print": 1,
|
|
"optimizer": {
|
|
"type": "Adam",
|
|
"params": {
|
|
"lr": lr
|
|
}
|
|
},
|
|
"bf16": bf16_config,
|
|
"zero_optimization": zero_config
|
|
}
|
|
|
|
engine, _, _, _ = deepspeed.initialize(config=config_dict,
|
|
model=ds_model,
|
|
model_parameters=ds_model.parameters())
|
|
|
|
data_loader = random_dataloader(model=engine,
|
|
total_samples=8,
|
|
hidden_dim=hidden_dim,
|
|
device=device,
|
|
dtype=torch.bfloat16)
|
|
batch = next(iter(data_loader))
|
|
|
|
optimizer_ddp.zero_grad()
|
|
loss_ddp = model_ddp(batch[0], batch[1])
|
|
loss_ddp.backward()
|
|
grads_ddp = collect_ddp_gradients(model_ddp)
|
|
|
|
loss_ds = engine(batch[0], batch[1])
|
|
loss_ds.backward()
|
|
grads_ds = collect_gradients_safe(engine)
|
|
|
|
compare_gradients(
|
|
grads_ddp,
|
|
grads_ds,
|
|
step_info=
|
|
f"bf16_optimizer_states={bf16_optimizer_states}, cpu_offload={use_cpu_offload}, zero_stage={zero_stage}")
|
|
|
|
optimizer_ddp.step()
|
|
optimizer_ddp.zero_grad()
|
|
engine.step()
|
|
engine.zero_grad()
|
|
|
|
if bf16_optimizer_states and use_cpu_offload:
|
|
# With CPU offload the Adam moments must be allocated in bf16 on the host so the
|
|
# offloaded optimizer-state footprint is smaller than with fp32 moments.
|
|
cpu_adam_state = engine.optimizer.optimizer.state
|
|
moment_tensors = []
|
|
for param_state in cpu_adam_state.values():
|
|
for moment_key in ("exp_avg", "exp_avg_sq"):
|
|
if moment_key in param_state:
|
|
moment_tensors.append(param_state[moment_key])
|
|
assert moment_tensors, "expected Adam moment tensors to be allocated after a step"
|
|
for moment in moment_tensors:
|
|
assert moment.dtype == torch.bfloat16, f"expected bf16 moment, got {moment.dtype}"
|
|
assert moment.device.type == "cpu", f"expected moment on cpu, got {moment.device}"
|
|
|
|
engine.destroy()
|
|
|
|
|
|
@pytest.mark.parametrize("zero_stage", [1, 2, 3])
|
|
class TestBF16OptimizerStatesOffloadValidation(DistributedTest):
|
|
world_size = 1
|
|
|
|
def test_user_cpu_adam_must_enable_bf16_states(self, zero_stage):
|
|
"""A user-provided DeepSpeedCPUAdam must be built with fp32_optimizer_states=False
|
|
to combine bf16_optimizer_states with ZeRO-Offload, otherwise the moments would
|
|
silently stay fp32 and the memory benefit would be lost."""
|
|
if 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"
|
|
)
|
|
if not deepspeed.ops.__compatible_ops__[CPUAdamBuilder.NAME]:
|
|
pytest.skip("cpu-adam is not compatible")
|
|
|
|
from deepspeed.ops.adam import DeepSpeedCPUAdam
|
|
|
|
hidden_dim = 6
|
|
model = SimpleModel(hidden_dim, nlayers=2)
|
|
# fp32_optimizer_states defaults to True, which keeps fp32 moments and is
|
|
# incompatible with bf16_optimizer_states under ZeRO-Offload.
|
|
optimizer = DeepSpeedCPUAdam(model.parameters())
|
|
|
|
config_dict = {
|
|
"train_micro_batch_size_per_gpu": 2,
|
|
"steps_per_print": 1,
|
|
"bf16": {
|
|
"enabled": True,
|
|
"bf16_master_weights_and_grads": True,
|
|
"bf16_optimizer_states": True,
|
|
},
|
|
# offload_optimizer is the current config key for ZeRO optimizer offload
|
|
# (TestBF16MasterWeightsGradients above still uses the legacy cpu_offload alias).
|
|
"zero_optimization": {
|
|
"stage": zero_stage,
|
|
"offload_optimizer": {
|
|
"device": "cpu"
|
|
},
|
|
},
|
|
}
|
|
|
|
with pytest.raises(AssertionError, match="fp32_optimizer_states=False"):
|
|
deepspeed.initialize(config=config_dict, model=model, optimizer=optimizer)
|