wehub-resource-sync
137 lines
5.1 KiB
Python
137 lines
5.1 KiB
Python
# Copyright (c) DeepSpeed Team.
|
|
# SPDX-License-Identifier: Apache-2.0
|
|
|
|
# DeepSpeed Team
|
|
|
|
import pytest
|
|
import deepspeed.comm as dist
|
|
from deepspeed.accelerator import get_accelerator
|
|
|
|
from unit.common import DistributedTest
|
|
from unit.simple_model import SimpleModel, random_dataloader
|
|
import deepspeed
|
|
|
|
|
|
class BaseZenFlowTest:
|
|
hidden_dim = 10
|
|
batch_size = 4
|
|
grad_acc_steps = 1
|
|
|
|
def get_config_dict(self, stage, offload_selective_optimizer, select_strategy, select_interval, update_interval,
|
|
full_warm_up_rounds):
|
|
config = {
|
|
"train_batch_size": self.batch_size,
|
|
"gradient_accumulation_steps": self.grad_acc_steps,
|
|
"steps_per_print": 1,
|
|
"optimizer": {
|
|
"type": "Adam",
|
|
"params": {
|
|
"lr": 1e-4
|
|
}
|
|
},
|
|
"zero_optimization": {
|
|
"stage": stage,
|
|
"offload_optimizer": {
|
|
"device": "cpu"
|
|
},
|
|
"overlap_comm": True,
|
|
"zenflow": {
|
|
"topk_ratio": 0.2,
|
|
"select_strategy": select_strategy,
|
|
"select_interval": select_interval,
|
|
"update_interval": update_interval,
|
|
"overlap_step": False,
|
|
"offload": offload_selective_optimizer,
|
|
"auto_ratio": 0.99,
|
|
"full_warm_up_rounds": full_warm_up_rounds,
|
|
}
|
|
},
|
|
"zero_allow_untested_optimizer": True,
|
|
}
|
|
|
|
if get_accelerator().is_bf16_supported():
|
|
config["bf16"] = {"enabled": True}
|
|
return config
|
|
|
|
def run_training_distributed(self, config_dict):
|
|
|
|
if get_accelerator().device_name() == "cpu":
|
|
return
|
|
|
|
model = SimpleModel(self.hidden_dim)
|
|
model, _, _, _ = deepspeed.initialize(model=model, model_parameters=model.parameters(), config=config_dict)
|
|
train_dataloader = random_dataloader(model=model,
|
|
total_samples=20,
|
|
hidden_dim=self.hidden_dim,
|
|
device=model.device)
|
|
|
|
dist.barrier()
|
|
|
|
for step, batch in enumerate(train_dataloader):
|
|
loss = model(batch[0], batch[1])
|
|
model.backward(loss)
|
|
model.step()
|
|
model.destroy()
|
|
|
|
|
|
@pytest.mark.parametrize("stage", [1, 2, 3])
|
|
@pytest.mark.parametrize("full_warm_up_rounds", [0, 3])
|
|
@pytest.mark.parametrize("offload_selective_optimizer", [True, False])
|
|
@pytest.mark.parametrize("select_strategy,select_interval,update_interval", [
|
|
("auto", "auto", "auto"),
|
|
("step", 10, 3),
|
|
("epoch", 1, 4),
|
|
])
|
|
class TestZenFlowSingleGPU(DistributedTest, BaseZenFlowTest):
|
|
world_size = 1
|
|
|
|
def test_zenflow_single_gpu(self, stage, offload_selective_optimizer, select_strategy, select_interval,
|
|
update_interval, full_warm_up_rounds):
|
|
tester = BaseZenFlowTest()
|
|
config_dict = tester.get_config_dict(stage, offload_selective_optimizer, select_strategy, select_interval,
|
|
update_interval, full_warm_up_rounds)
|
|
tester.run_training_distributed(config_dict)
|
|
|
|
|
|
@pytest.mark.parametrize("stage", [1, 2, 3])
|
|
@pytest.mark.parametrize("full_warm_up_rounds", [0, 3])
|
|
@pytest.mark.parametrize("offload_selective_optimizer", [True, False])
|
|
@pytest.mark.parametrize("select_strategy,select_interval,update_interval", [
|
|
("auto", "auto", "auto"),
|
|
("step", 10, 3),
|
|
("epoch", 1, 4),
|
|
])
|
|
class TestZenFlowDistributed(DistributedTest, BaseZenFlowTest):
|
|
world_size = 2
|
|
|
|
def test_zenflow_distributed(self, stage, offload_selective_optimizer, select_strategy, select_interval,
|
|
update_interval, full_warm_up_rounds):
|
|
config_dict = self.get_config_dict(stage, offload_selective_optimizer, select_strategy, select_interval,
|
|
update_interval, full_warm_up_rounds)
|
|
self.run_training_distributed(config_dict)
|
|
|
|
|
|
@pytest.mark.parametrize(
|
|
"cores,perc,expected_zf,expected_pt",
|
|
[
|
|
# Normal split: ceil(0.25 * 8) = 2 cores reserved for training.
|
|
([0, 1, 2, 3, 4, 5, 6, 7], 0.25, [2, 3, 4, 5, 6, 7], [0, 1]),
|
|
# Rounds up: ceil(0.1 * 8) = 1.
|
|
([0, 1, 2, 3, 4, 5, 6, 7], 0.1, [1, 2, 3, 4, 5, 6, 7], [0]),
|
|
# Two cores, half each.
|
|
([10, 11], 0.5, [11], [10]),
|
|
# Reserve rounds to 0 -> both sides share the full set.
|
|
([0, 1, 2, 3], 0.0, [0, 1, 2, 3], [0, 1, 2, 3]),
|
|
# Reserve rounds to every core -> both sides share the full set.
|
|
([0, 1, 2, 3], 1.0, [0, 1, 2, 3], [0, 1, 2, 3]),
|
|
])
|
|
def test_split_affinity(cores, perc, expected_zf, expected_pt):
|
|
from deepspeed.runtime.zenflow.zenflow_utils import _split_affinity
|
|
zf, pt = _split_affinity(cores, perc)
|
|
assert zf == expected_zf
|
|
assert pt == expected_pt
|
|
# When the sides are actually isolated they must partition the cores exactly.
|
|
if zf != pt:
|
|
assert sorted(zf + pt) == sorted(cores)
|
|
assert not (set(zf) & set(pt))
|