From f463a301c923c1c9a0d82d3cbccd35eb2eaf2ded Mon Sep 17 00:00:00 2001 From: Michael Wyatt Date: Wed, 3 Aug 2022 15:06:24 -0700 Subject: [PATCH 01/26] refactor of onebit tests --- tests/unit/test_onebit.py | 1451 ++++++++++++++++--------------------- tests/unit/test_pipe.py | 11 +- 2 files changed, 638 insertions(+), 824 deletions(-) diff --git a/tests/unit/test_onebit.py b/tests/unit/test_onebit.py index b6f1f8bd4e15..41fbd1c59099 100644 --- a/tests/unit/test_onebit.py +++ b/tests/unit/test_onebit.py @@ -8,12 +8,11 @@ import numpy as np from deepspeed.runtime.pipe.topology import PipeDataParallelTopology -from deepspeed.ops.op_builder import OpBuilder PipeTopo = PipeDataParallelTopology from deepspeed.runtime.pipe.module import PipelineModule -from .common import distributed_test -from .simple_model import SimpleModel, random_dataloader, args_from_dict +from .common import DistributedTest +from .simple_model import SimpleModel, random_dataloader from .test_pipe import AlexNetPipe, train_cifar TORCH_MAJOR = int(torch.__version__.split('.')[0]) @@ -22,138 +21,92 @@ pytest.skip("NCCL-based 1-bit compression requires torch 1.8 or higher", allow_module_level=True) -rocm_version = OpBuilder.installed_rocm_version() -if rocm_version[0] > 4: - pytest.skip( - "NCCL-based 1-bit compression is not yet supported w. ROCm 5 until cupy supports ROCm 5", - allow_module_level=True) - - -def test_onebitadam_fp16_basic(tmpdir): - config_dict = { - "train_batch_size": 2, - "steps_per_print": 1, - "optimizer": { - "type": "OneBitAdam", - "params": { - "lr": 0.00015, - "weight_decay": 0.01, - "freeze_step": 2, - "cuda_aware": False, - "comm_backend_name": "nccl" - } - }, - "gradient_clipping": 1.0, - "fp16": { - "enabled": True, - "loss_scale": 0, - "initial_scale_power": 16 - } - } - args = args_from_dict(tmpdir, config_dict) - hidden_dim = 10 - - model = SimpleModel(hidden_dim) - @distributed_test(world_size=[1, 2]) - def _test_onebitadam_fp16_basic(args, model, hidden_dim): - model, _, _, _ = deepspeed.initialize(args=args, - model=model, - model_parameters=model.parameters()) - data_loader = random_dataloader(model=model, - total_samples=50, - 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() - - _test_onebitadam_fp16_basic(args=args, model=model, hidden_dim=hidden_dim) - - -def test_onebitadam_fp32_basic(tmpdir): - config_dict = { - "train_batch_size": 2, - "steps_per_print": 1, - "optimizer": { - "type": "OneBitAdam", - "params": { - "lr": 0.00015, - "weight_decay": 0.01, - "freeze_step": 2, - "cuda_aware": False, - "comm_backend_name": "nccl" +@pytest.mark.parametrize("dtype", [torch.float32, torch.float16], ids=["fp32", "fp16"]) +class TestOneBitAdamBasic(DistributedTest): + world_size = [1, 2] + + def test(self, dtype): + config_dict = { + "train_batch_size": 2, + "steps_per_print": 1, + "optimizer": { + "type": "OneBitAdam", + "params": { + "lr": 0.00015, + "weight_decay": 0.01, + "freeze_step": 2, + "cuda_aware": False, + "comm_backend_name": "nccl" + } + }, + "gradient_clipping": 1.0, + "fp16": { + "enabled": (dtype == torch.float16), + "loss_scale": 0, + "initial_scale_power": 16 } - }, - "gradient_clipping": 1.0, - } - args = args_from_dict(tmpdir, config_dict) - hidden_dim = 10 - - model = SimpleModel(hidden_dim) + } + hidden_dim = 10 - @distributed_test(world_size=[1, 2]) - def _test_onebitadam_fp32_basic(args, model, hidden_dim): - model, _, _, _ = deepspeed.initialize(args=args, + 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.float) + dtype=dtype) for n, batch in enumerate(data_loader): loss = model(batch[0], batch[1]) model.backward(loss) model.step() - _test_onebitadam_fp32_basic(args=args, model=model, hidden_dim=hidden_dim) - - -def test_onebitadam_exp_avg_mask(tmpdir): - config_dict = { - "train_batch_size": 2, - "steps_per_print": 1, - "optimizer": { - "type": "OneBitAdam", - "params": { - "lr": 0.00015, - "weight_decay": 0.01, - "freeze_step": 2, - "cuda_aware": False, - "comm_backend_name": "nccl" + +class TestOneBitAdamExpAvgMask(DistributedTest): + world_size = 2 + + def test(self): + config_dict = { + "train_batch_size": 2, + "steps_per_print": 1, + "optimizer": { + "type": "OneBitAdam", + "params": { + "lr": 0.00015, + "weight_decay": 0.01, + "freeze_step": 2, + "cuda_aware": False, + "comm_backend_name": "nccl" + } + }, + "gradient_clipping": 1.0, + "fp16": { + "enabled": True, + "loss_scale": 0, + "initial_scale_power": 16 } - }, - "gradient_clipping": 1.0, - "fp16": { - "enabled": True, - "loss_scale": 0, - "initial_scale_power": 16 } - } - args = args_from_dict(tmpdir, config_dict) - hidden_dim = 10 - - model = SimpleModel(hidden_dim) - param_optimizer = list(model.named_parameters()) - mask1 = torch.zeros_like(param_optimizer[0][1].data) - for col in range(mask1.size()[1]): - mask1[0][col] += 1 - mask1 = torch.flatten(mask1) - optimizer_grouped_parameters = [{ - 'params': [param_optimizer[0][1]], - 'weight_decay': 0.01, - 'exp_avg_mask': mask1 - }, - { - 'params': [param_optimizer[1][1]], - 'weight_decay': 0.01 - }] - - @distributed_test(world_size=[2]) - def _test_onebitadam_exp_avg_mask(args, model, hidden_dim): - model, optimizer, _, _ = deepspeed.initialize(args=args, + hidden_dim = 10 + + model = SimpleModel(hidden_dim) + param_optimizer = list(model.named_parameters()) + mask1 = torch.zeros_like(param_optimizer[0][1].data) + for col in range(mask1.size()[1]): + mask1[0][col] += 1 + mask1 = torch.flatten(mask1) + optimizer_grouped_parameters = [{ + 'params': [param_optimizer[0][1]], + 'weight_decay': 0.01, + 'exp_avg_mask': mask1 + }, + { + 'params': [param_optimizer[1][1]], + 'weight_decay': 0.01 + }] + + model, optimizer, _, _ = deepspeed.initialize(config=config_dict, model=model, model_parameters=optimizer_grouped_parameters) data_loader = random_dataloader(model=model, @@ -169,75 +122,73 @@ def _test_onebitadam_exp_avg_mask(args, model, hidden_dim): if v['exp_avg'].size() == mask1.size(): assert torch.allclose(v['exp_avg'], v['exp_avg'].mul_(mask1.to(device=v['exp_avg'].device)), atol=1e-07), f"Momentum mask is not working properly" - _test_onebitadam_exp_avg_mask(args=args, model=model, hidden_dim=hidden_dim) - - -def test_onebitadam_checkpointing(tmpdir): - config_dict = { - "train_batch_size": 2, - "steps_per_print": 1, - "optimizer": { - "type": "OneBitAdam", - "params": { - "lr": 0.00015, - "weight_decay": 0.01, - "freeze_step": 2, - "cuda_aware": False, - "comm_backend_name": "nccl" + +class TestOneBitAdamCheckpointing(DistributedTest): + world_size = 2 + + def test(self, tmpdir): + config_dict = { + "train_batch_size": 2, + "steps_per_print": 1, + "optimizer": { + "type": "OneBitAdam", + "params": { + "lr": 0.00015, + "weight_decay": 0.01, + "freeze_step": 2, + "cuda_aware": False, + "comm_backend_name": "nccl" + } + }, + "gradient_clipping": 1.0, + "fp16": { + "enabled": True, + "loss_scale": 0, + "initial_scale_power": 16 } - }, - "gradient_clipping": 1.0, - "fp16": { - "enabled": True, - "loss_scale": 0, - "initial_scale_power": 16 } - } - args = args_from_dict(tmpdir, config_dict) - hidden_dim = 10 - - model = SimpleModel(hidden_dim) - param_optimizer = list(model.named_parameters()) - mask1 = torch.zeros_like(param_optimizer[0][1].data) - mask2 = torch.zeros_like(param_optimizer[0][1].data) - for col in range(mask1.size()[1]): - mask1[0][col] += 1 - mask2[1][col] += 1 - mask1 = torch.flatten(mask1) - mask2 = torch.flatten(mask2) - - optimizer_grouped_parameters_1 = [{ - 'params': [param_optimizer[0][1]], - 'weight_decay': 0.01, - 'exp_avg_mask': mask1 - }, - { - 'params': [param_optimizer[1][1]], - 'weight_decay': 0.01 - }] - - optimizer_grouped_parameters_2 = [{ - 'params': [param_optimizer[0][1]], - 'weight_decay': 0.01, - 'exp_avg_mask': mask2 - }, - { - 'params': [param_optimizer[1][1]], - 'weight_decay': 0.01 - }] - - optimizer_grouped_parameters_3 = [{ - 'params': [param_optimizer[0][1]], - 'weight_decay': 0.01 - }, - { - 'params': [param_optimizer[1][1]], - 'weight_decay': 0.01 - }] - - @distributed_test(world_size=[2]) - def _test_onebitadam_checkpointing(mask1, mask2, args, model, hidden_dim): - model_1, optimizer_1, _, _ = deepspeed.initialize(args=args, + hidden_dim = 10 + + model = SimpleModel(hidden_dim) + param_optimizer = list(model.named_parameters()) + mask1 = torch.zeros_like(param_optimizer[0][1].data) + mask2 = torch.zeros_like(param_optimizer[0][1].data) + for col in range(mask1.size()[1]): + mask1[0][col] += 1 + mask2[1][col] += 1 + mask1 = torch.flatten(mask1) + mask2 = torch.flatten(mask2) + + optimizer_grouped_parameters_1 = [{ + 'params': [param_optimizer[0][1]], + 'weight_decay': 0.01, + 'exp_avg_mask': mask1 + }, + { + 'params': [param_optimizer[1][1]], + 'weight_decay': 0.01 + }] + + optimizer_grouped_parameters_2 = [{ + 'params': [param_optimizer[0][1]], + 'weight_decay': 0.01, + 'exp_avg_mask': mask2 + }, + { + 'params': [param_optimizer[1][1]], + 'weight_decay': 0.01 + }] + + optimizer_grouped_parameters_3 = [{ + 'params': [param_optimizer[0][1]], + 'weight_decay': 0.01 + }, + { + 'params': [param_optimizer[1][1]], + 'weight_decay': 0.01 + }] + + model_1, optimizer_1, _, _ = deepspeed.initialize(config=config_dict, model=model, model_parameters=optimizer_grouped_parameters_1) data_loader = random_dataloader(model=model_1, @@ -257,7 +208,7 @@ def _test_onebitadam_checkpointing(mask1, mask2, args, model, hidden_dim): assert torch.allclose(optimizer_1.param_groups[0]['exp_avg_mask'], mask1, atol=1e-07), f"Momentum mask should not change after saving checkpoint" - model_2, optimizer_2, _, _ = deepspeed.initialize(args=args, + model_2, optimizer_2, _, _ = deepspeed.initialize(config=config_dict, model=model, model_parameters=optimizer_grouped_parameters_2) # Test whether momentum mask stays the same after loading checkpoint @@ -274,7 +225,7 @@ def _test_onebitadam_checkpointing(mask1, mask2, args, model, hidden_dim): assert 'server_error' not in v, f"Incorrect server error" assert optimizer_2.optimizer.adam_freeze_key is True - model_3, optimizer_3, _, _ = deepspeed.initialize(args=args, + model_3, optimizer_3, _, _ = deepspeed.initialize(config=config_dict, model=model, model_parameters=optimizer_grouped_parameters_3) optimizer_3.optimizer.freeze_step = 20 @@ -300,42 +251,31 @@ def _test_onebitadam_checkpointing(mask1, mask2, args, model, hidden_dim): assert 'server_error' not in v, f"Incorrect server error" assert optimizer_3.optimizer.adam_freeze_key is False - _test_onebitadam_checkpointing(mask1, - mask2, - args=args, - model=model, - hidden_dim=hidden_dim) - - -def test_onebitadam_checkpointing_overflow(tmpdir): - config_dict = { - "train_batch_size": 2, - "steps_per_print": 1, - "optimizer": { - "type": "OneBitAdam", - "params": { - "lr": 0.00015, - "weight_decay": 0.01, - "freeze_step": 2, - "cuda_aware": False, - "comm_backend_name": "nccl" + def test_overflow(self, tmpdir): + config_dict = { + "train_batch_size": 2, + "steps_per_print": 1, + "optimizer": { + "type": "OneBitAdam", + "params": { + "lr": 0.00015, + "weight_decay": 0.01, + "freeze_step": 2, + "cuda_aware": False, + "comm_backend_name": "nccl" + } + }, + "gradient_clipping": 1.0, + "fp16": { + "enabled": True, + "loss_scale": 0, + "initial_scale_power": 16 } - }, - "gradient_clipping": 1.0, - "fp16": { - "enabled": True, - "loss_scale": 0, - "initial_scale_power": 16 } - } - args = args_from_dict(tmpdir, config_dict) - hidden_dim = 10 + hidden_dim = 10 - model = SimpleModel(hidden_dim) - - @distributed_test(world_size=[2]) - def _test_onebitadam_checkpointing_overflow(args, model, hidden_dim): - model, _, _, _ = deepspeed.initialize(args=args, + model = SimpleModel(hidden_dim) + model, _, _, _ = deepspeed.initialize(config=config_dict, model=model, model_parameters=model.parameters()) data_loader = random_dataloader(model=model, @@ -353,10 +293,6 @@ def _test_onebitadam_checkpointing_overflow(args, model, hidden_dim): dist.barrier() model.save_checkpoint(save_folder, tag=None) - _test_onebitadam_checkpointing_overflow(args=args, - model=model, - hidden_dim=hidden_dim) - @pytest.mark.parametrize('topo', [ @@ -367,45 +303,46 @@ def _test_onebitadam_checkpointing_overflow(args, model, hidden_dim): PipeTopo(num_pp=4, num_dp=1), ]) -def test_onebitadam_fp16_pipeline(topo, tmpdir): - config_dict = { - "train_batch_size": 16, - "train_micro_batch_size_per_gpu": 4, - "steps_per_print": 20, - "optimizer": { - "type": "OneBitAdam", - "params": { - "lr": 0.00001, - "betas": [0.9, - 0.999], - "eps": 1e-8, - "weight_decay": 3e-7, - "freeze_step": 200, - "cuda_aware": False, - "comm_backend_name": "nccl" +class TestOneBitAdamFP16Pipeline(DistributedTest): + world_size = 4 + + def test(self, topo): + config_dict = { + "train_batch_size": 16, + "train_micro_batch_size_per_gpu": 4, + "steps_per_print": 20, + "optimizer": { + "type": "OneBitAdam", + "params": { + "lr": 0.00001, + "betas": [0.9, + 0.999], + "eps": 1e-8, + "weight_decay": 3e-7, + "freeze_step": 200, + "cuda_aware": False, + "comm_backend_name": "nccl" + } + }, + "gradient_clipping": 1.0, + "zero_optimization": { + "stage": 0 + }, + "fp16": { + "enabled": True, + "loss_scale": 0, + "initial_scale_power": 16 + }, + "pipeline": { + "seed_layers": True, + "activation_checkpoint_interval": 1 } - }, - "gradient_clipping": 1.0, - "zero_optimization": { - "stage": 0 - }, - "fp16": { - "enabled": True, - "loss_scale": 0, - "initial_scale_power": 16 - }, - "pipeline": { - "seed_layers": True, - "activation_checkpoint_interval": 1 } - } - args = args_from_dict(tmpdir, config_dict) - # Allocate model for consistent initial weights. - init_net = AlexNetPipe() + # Allocate model for consistent initial weights. + init_net = AlexNetPipe() - @distributed_test(world_size=4) - def _helper(topo, tmpdir, steps=500): + steps = 500 assert steps >= 100 test_net = copy.deepcopy(init_net) @@ -414,147 +351,102 @@ def _helper(topo, tmpdir, steps=500): loss_fn=nn.CrossEntropyLoss()) test_losses = train_cifar(test_model, - args, + config_dict=config_dict, num_steps=steps, fp16=config_dict['fp16']['enabled']) - _helper(topo, tmpdir) - - -def test_zerooneadam_fp16_basic(tmpdir): - config_dict = { - "train_batch_size": 2, - "steps_per_print": 1, - "optimizer": { - "type": "ZeroOneAdam", - "params": { - "lr": 0.00015, - "weight_decay": 0.01, - "var_freeze_step": 4, - "var_update_scaler": 1, - "local_step_scaler": 1, - "local_step_clipper": 2, - "cuda_aware": False, - "comm_backend_name": "nccl" - } - }, - "gradient_clipping": 1.0, - "fp16": { - "enabled": True, - "loss_scale": 0, - "initial_scale_power": 16 - } - } - args = args_from_dict(tmpdir, config_dict) - hidden_dim = 10 - - model = SimpleModel(hidden_dim) - - @distributed_test(world_size=[1, 2]) - def _test_zerooneadam_fp16_basic(args, model, hidden_dim): - model, _, _, _ = deepspeed.initialize(args=args, - model=model, - model_parameters=model.parameters()) - data_loader = random_dataloader(model=model, - total_samples=50, - 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() - _test_zerooneadam_fp16_basic(args=args, model=model, hidden_dim=hidden_dim) - - -def test_zerooneadam_fp32_basic(tmpdir): - config_dict = { - "train_batch_size": 2, - "steps_per_print": 1, - "optimizer": { - "type": "ZeroOneAdam", - "params": { - "lr": 0.00015, - "weight_decay": 0.01, - "var_freeze_step": 4, - "var_update_scaler": 1, - "local_step_scaler": 1, - "local_step_clipper": 2, - "cuda_aware": False, - "comm_backend_name": "nccl" +@pytest.mark.parametrize("dtype", [torch.float32, torch.float16], ids=["fp32", "fp16"]) +class TestZeroOneAdamBasic(DistributedTest): + world_size = [1, 2] + + def test(self, dtype): + config_dict = { + "train_batch_size": 2, + "steps_per_print": 1, + "optimizer": { + "type": "ZeroOneAdam", + "params": { + "lr": 0.00015, + "weight_decay": 0.01, + "var_freeze_step": 4, + "var_update_scaler": 1, + "local_step_scaler": 1, + "local_step_clipper": 2, + "cuda_aware": False, + "comm_backend_name": "nccl" + } + }, + "gradient_clipping": 1.0, + "fp16": { + "enabled": (dtype == torch.float16), + "loss_scale": 0, + "initial_scale_power": 16 } - }, - "gradient_clipping": 1.0, - } - args = args_from_dict(tmpdir, config_dict) - hidden_dim = 10 - - model = SimpleModel(hidden_dim) + } + hidden_dim = 10 - @distributed_test(world_size=[1, 2]) - def _test_zerooneadam_fp32_basic(args, model, hidden_dim): - model, _, _, _ = deepspeed.initialize(args=args, + 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.float) + dtype=dtype) for n, batch in enumerate(data_loader): loss = model(batch[0], batch[1]) model.backward(loss) model.step() - _test_zerooneadam_fp32_basic(args=args, model=model, hidden_dim=hidden_dim) - - -def test_zerooneadam_exp_avg_mask(tmpdir): - config_dict = { - "train_batch_size": 2, - "steps_per_print": 1, - "optimizer": { - "type": "ZeroOneAdam", - "params": { - "lr": 0.00015, - "weight_decay": 0.01, - "var_freeze_step": 4, - "var_update_scaler": 1, - "local_step_scaler": 1, - "local_step_clipper": 2, - "cuda_aware": False, - "comm_backend_name": "nccl" + +class TestZeroOneAdamExpAvgMask(DistributedTest): + world_size = 2 + + def test(self): + config_dict = { + "train_batch_size": 2, + "steps_per_print": 1, + "optimizer": { + "type": "ZeroOneAdam", + "params": { + "lr": 0.00015, + "weight_decay": 0.01, + "var_freeze_step": 4, + "var_update_scaler": 1, + "local_step_scaler": 1, + "local_step_clipper": 2, + "cuda_aware": False, + "comm_backend_name": "nccl" + } + }, + "gradient_clipping": 1.0, + "fp16": { + "enabled": True, + "loss_scale": 0, + "initial_scale_power": 16 } - }, - "gradient_clipping": 1.0, - "fp16": { - "enabled": True, - "loss_scale": 0, - "initial_scale_power": 16 } - } - args = args_from_dict(tmpdir, config_dict) - hidden_dim = 10 - - model = SimpleModel(hidden_dim) - param_optimizer = list(model.named_parameters()) - mask1 = torch.zeros_like(param_optimizer[0][1].data) - for col in range(mask1.size()[1]): - mask1[0][col] += 1 - mask1 = torch.flatten(mask1) - optimizer_grouped_parameters = [{ - 'params': [param_optimizer[0][1]], - 'weight_decay': 0.01, - 'exp_avg_mask': mask1 - }, - { - 'params': [param_optimizer[1][1]], - 'weight_decay': 0.01 - }] - - @distributed_test(world_size=[2]) - def _test_zerooneadam_exp_avg_mask(args, model, hidden_dim): - model, optimizer, _, _ = deepspeed.initialize(args=args, + hidden_dim = 10 + + model = SimpleModel(hidden_dim) + param_optimizer = list(model.named_parameters()) + mask1 = torch.zeros_like(param_optimizer[0][1].data) + for col in range(mask1.size()[1]): + mask1[0][col] += 1 + mask1 = torch.flatten(mask1) + optimizer_grouped_parameters = [{ + 'params': [param_optimizer[0][1]], + 'weight_decay': 0.01, + 'exp_avg_mask': mask1 + }, + { + 'params': [param_optimizer[1][1]], + 'weight_decay': 0.01 + }] + + model, optimizer, _, _ = deepspeed.initialize(config=config_dict, model=model, model_parameters=optimizer_grouped_parameters) data_loader = random_dataloader(model=model, @@ -570,78 +462,76 @@ def _test_zerooneadam_exp_avg_mask(args, model, hidden_dim): if v['exp_avg'].size() == mask1.size(): assert torch.allclose(v['exp_avg'], v['exp_avg'].mul_(mask1.to(device=v['exp_avg'].device)), atol=1e-07), f"Momentum mask is not working properly" - _test_zerooneadam_exp_avg_mask(args=args, model=model, hidden_dim=hidden_dim) - - -def test_zerooneadam_checkpointing(tmpdir): - config_dict = { - "train_batch_size": 2, - "steps_per_print": 1, - "optimizer": { - "type": "ZeroOneAdam", - "params": { - "lr": 0.00015, - "weight_decay": 0.01, - "var_freeze_step": 4, - "var_update_scaler": 1, - "local_step_scaler": 1, - "local_step_clipper": 2, - "cuda_aware": False, - "comm_backend_name": "nccl" + +class TestZeroOneAdamCheckpointing(DistributedTest): + world_size = 2 + + def test(self, tmpdir): + config_dict = { + "train_batch_size": 2, + "steps_per_print": 1, + "optimizer": { + "type": "ZeroOneAdam", + "params": { + "lr": 0.00015, + "weight_decay": 0.01, + "var_freeze_step": 4, + "var_update_scaler": 1, + "local_step_scaler": 1, + "local_step_clipper": 2, + "cuda_aware": False, + "comm_backend_name": "nccl" + } + }, + "gradient_clipping": 1.0, + "fp16": { + "enabled": True, + "loss_scale": 0, + "initial_scale_power": 16 } - }, - "gradient_clipping": 1.0, - "fp16": { - "enabled": True, - "loss_scale": 0, - "initial_scale_power": 16 } - } - args = args_from_dict(tmpdir, config_dict) - hidden_dim = 10 - - model = SimpleModel(hidden_dim) - param_optimizer = list(model.named_parameters()) - mask1 = torch.zeros_like(param_optimizer[0][1].data) - mask2 = torch.zeros_like(param_optimizer[0][1].data) - for col in range(mask1.size()[1]): - mask1[0][col] += 1 - mask2[1][col] += 1 - mask1 = torch.flatten(mask1) - mask2 = torch.flatten(mask2) - - optimizer_grouped_parameters_1 = [{ - 'params': [param_optimizer[0][1]], - 'weight_decay': 0.01, - 'exp_avg_mask': mask1 - }, - { - 'params': [param_optimizer[1][1]], - 'weight_decay': 0.01 - }] - - optimizer_grouped_parameters_2 = [{ - 'params': [param_optimizer[0][1]], - 'weight_decay': 0.01, - 'exp_avg_mask': mask2 - }, - { - 'params': [param_optimizer[1][1]], - 'weight_decay': 0.01 - }] - - optimizer_grouped_parameters_3 = [{ - 'params': [param_optimizer[0][1]], - 'weight_decay': 0.01 - }, - { - 'params': [param_optimizer[1][1]], - 'weight_decay': 0.01 - }] - - @distributed_test(world_size=[2]) - def _test_zerooneadam_checkpointing(mask1, mask2, args, model, hidden_dim): - model_1, optimizer_1, _, _ = deepspeed.initialize(args=args, + hidden_dim = 10 + + model = SimpleModel(hidden_dim) + param_optimizer = list(model.named_parameters()) + mask1 = torch.zeros_like(param_optimizer[0][1].data) + mask2 = torch.zeros_like(param_optimizer[0][1].data) + for col in range(mask1.size()[1]): + mask1[0][col] += 1 + mask2[1][col] += 1 + mask1 = torch.flatten(mask1) + mask2 = torch.flatten(mask2) + + optimizer_grouped_parameters_1 = [{ + 'params': [param_optimizer[0][1]], + 'weight_decay': 0.01, + 'exp_avg_mask': mask1 + }, + { + 'params': [param_optimizer[1][1]], + 'weight_decay': 0.01 + }] + + optimizer_grouped_parameters_2 = [{ + 'params': [param_optimizer[0][1]], + 'weight_decay': 0.01, + 'exp_avg_mask': mask2 + }, + { + 'params': [param_optimizer[1][1]], + 'weight_decay': 0.01 + }] + + optimizer_grouped_parameters_3 = [{ + 'params': [param_optimizer[0][1]], + 'weight_decay': 0.01 + }, + { + 'params': [param_optimizer[1][1]], + 'weight_decay': 0.01 + }] + + model_1, optimizer_1, _, _ = deepspeed.initialize(config=config_dict, model=model, model_parameters=optimizer_grouped_parameters_1) data_loader = random_dataloader(model=model_1, @@ -660,7 +550,7 @@ def _test_zerooneadam_checkpointing(mask1, mask2, args, model, hidden_dim): assert torch.allclose(optimizer_1.param_groups[0]['exp_avg_mask'], mask1, atol=1e-07), f"Momentum mask should not change after saving checkpoint" - model_2, optimizer_2, _, _ = deepspeed.initialize(args=args, + model_2, optimizer_2, _, _ = deepspeed.initialize(config=config_dict, model=model, model_parameters=optimizer_grouped_parameters_2) # Test whether momentum mask stays the same after loading checkpoint @@ -676,7 +566,7 @@ def _test_zerooneadam_checkpointing(mask1, mask2, args, model, hidden_dim): assert 'worker_error' not in v, f"Incorrect worker error" assert 'server_error' not in v, f"Incorrect server error" - model_3, optimizer_3, _, _ = deepspeed.initialize(args=args, + model_3, optimizer_3, _, _ = deepspeed.initialize(config=config_dict, model=model, model_parameters=optimizer_grouped_parameters_3) optimizer_3.optimizer.freeze_step = 20 @@ -700,45 +590,34 @@ def _test_zerooneadam_checkpointing(mask1, mask2, args, model, hidden_dim): assert 'worker_error' not in v, f"Incorrect worker error" assert 'server_error' not in v, f"Incorrect server error" - _test_zerooneadam_checkpointing(mask1, - mask2, - args=args, - model=model, - hidden_dim=hidden_dim) - - -def test_zerooneadam_checkpointing_overflow(tmpdir): - config_dict = { - "train_batch_size": 2, - "steps_per_print": 1, - "optimizer": { - "type": "ZeroOneAdam", - "params": { - "lr": 0.00015, - "weight_decay": 0.01, - "var_freeze_step": 4, - "var_update_scaler": 1, - "local_step_scaler": 1, - "local_step_clipper": 2, - "cuda_aware": False, - "comm_backend_name": "nccl" + def test_overflow(self, tmpdir): + config_dict = { + "train_batch_size": 2, + "steps_per_print": 1, + "optimizer": { + "type": "ZeroOneAdam", + "params": { + "lr": 0.00015, + "weight_decay": 0.01, + "var_freeze_step": 4, + "var_update_scaler": 1, + "local_step_scaler": 1, + "local_step_clipper": 2, + "cuda_aware": False, + "comm_backend_name": "nccl" + } + }, + "gradient_clipping": 1.0, + "fp16": { + "enabled": True, + "loss_scale": 0, + "initial_scale_power": 16 } - }, - "gradient_clipping": 1.0, - "fp16": { - "enabled": True, - "loss_scale": 0, - "initial_scale_power": 16 } - } - args = args_from_dict(tmpdir, config_dict) - hidden_dim = 10 + hidden_dim = 10 - model = SimpleModel(hidden_dim) - - @distributed_test(world_size=[2]) - def _test_zerooneadam_checkpointing_overflow(args, model, hidden_dim): - model, _, _, _ = deepspeed.initialize(args=args, + model = SimpleModel(hidden_dim) + model, _, _, _ = deepspeed.initialize(config=config_dict, model=model, model_parameters=model.parameters()) data_loader = random_dataloader(model=model, @@ -756,10 +635,6 @@ def _test_zerooneadam_checkpointing_overflow(args, model, hidden_dim): dist.barrier() model.save_checkpoint(save_folder, tag=None) - _test_zerooneadam_checkpointing_overflow(args=args, - model=model, - hidden_dim=hidden_dim) - @pytest.mark.parametrize('topo', [ @@ -770,48 +645,49 @@ def _test_zerooneadam_checkpointing_overflow(args, model, hidden_dim): PipeTopo(num_pp=4, num_dp=1), ]) -def test_zerooneadam_fp16_pipeline(topo, tmpdir): - config_dict = { - "train_batch_size": 16, - "train_micro_batch_size_per_gpu": 4, - "steps_per_print": 20, - "optimizer": { - "type": "ZeroOneAdam", - "params": { - "lr": 0.00001, - "betas": [0.9, - 0.999], - "eps": 1e-8, - "weight_decay": 3e-7, - "var_freeze_step": 4, - "var_update_scaler": 1, - "local_step_scaler": 1, - "local_step_clipper": 2, - "cuda_aware": False, - "comm_backend_name": "nccl" +class TestZeroOneAdamFP16Pipeline(DistributedTest): + world_size = 4 + + def test(self, topo): + config_dict = { + "train_batch_size": 16, + "train_micro_batch_size_per_gpu": 4, + "steps_per_print": 20, + "optimizer": { + "type": "ZeroOneAdam", + "params": { + "lr": 0.00001, + "betas": [0.9, + 0.999], + "eps": 1e-8, + "weight_decay": 3e-7, + "var_freeze_step": 4, + "var_update_scaler": 1, + "local_step_scaler": 1, + "local_step_clipper": 2, + "cuda_aware": False, + "comm_backend_name": "nccl" + } + }, + "gradient_clipping": 1.0, + "zero_optimization": { + "stage": 0 + }, + "fp16": { + "enabled": True, + "loss_scale": 0, + "initial_scale_power": 16 + }, + "pipeline": { + "seed_layers": True, + "activation_checkpoint_interval": 1 } - }, - "gradient_clipping": 1.0, - "zero_optimization": { - "stage": 0 - }, - "fp16": { - "enabled": True, - "loss_scale": 0, - "initial_scale_power": 16 - }, - "pipeline": { - "seed_layers": True, - "activation_checkpoint_interval": 1 } - } - args = args_from_dict(tmpdir, config_dict) - # Allocate model for consistent initial weights. - init_net = AlexNetPipe() + # Allocate model for consistent initial weights. + init_net = AlexNetPipe() - @distributed_test(world_size=4) - def _helper(topo, tmpdir, steps=500): + steps = 500 assert steps >= 100 test_net = copy.deepcopy(init_net) @@ -820,155 +696,107 @@ def _helper(topo, tmpdir, steps=500): loss_fn=nn.CrossEntropyLoss()) test_losses = train_cifar(test_model, - args, + config_dict=config_dict, num_steps=steps, fp16=config_dict['fp16']['enabled']) - _helper(topo, tmpdir) - - -def test_onebitlamb_fp16_basic(tmpdir): - config_dict = { - "train_batch_size": 2, - "steps_per_print": 1, - "optimizer": { - "type": "OneBitLamb", - "params": { - "lr": 0.00015, - "weight_decay": 0.01, - "max_coeff": 0.3, - "min_coeff": 0.01, - "freeze_step": 2, - "cuda_aware": False, - "comm_backend_name": "nccl", - "coeff_beta": 0.9, - "factor_max": 1.0, - "factor_min": 0.5, - "factor_threshold": 0.1 - } - }, - "gradient_clipping": 1.0, - "fp16": { - "enabled": True, - "loss_scale": 0, - "initial_scale_power": 16 - } - } - args = args_from_dict(tmpdir, config_dict) - hidden_dim = 10 - - model = SimpleModel(hidden_dim) - - @distributed_test(world_size=[1, 2]) - def _test_onebitlamb_fp16_basic(args, model, hidden_dim): - model, _, _, _ = deepspeed.initialize(args=args, - model=model, - model_parameters=model.parameters()) - data_loader = random_dataloader(model=model, - total_samples=50, - 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() - _test_onebitlamb_fp16_basic(args=args, model=model, hidden_dim=hidden_dim) - - -def test_onebitlamb_fp32_basic(tmpdir): - config_dict = { - "train_batch_size": 2, - "steps_per_print": 1, - "optimizer": { - "type": "OneBitLamb", - "params": { - "lr": 0.00015, - "weight_decay": 0.01, - "max_coeff": 0.3, - "min_coeff": 0.01, - "freeze_step": 2, - "cuda_aware": False, - "comm_backend_name": "nccl", - "coeff_beta": 0.9, - "factor_max": 1.0, - "factor_min": 0.5, - "factor_threshold": 0.1 +@pytest.mark.parametrize("dtype", [torch.float32, torch.float16], ids=["fp32", "fp16"]) +class TestOneBitLambBasic(DistributedTest): + world_size = [1, 2] + + def test(self, dtype): + config_dict = { + "train_batch_size": 2, + "steps_per_print": 1, + "optimizer": { + "type": "OneBitLamb", + "params": { + "lr": 0.00015, + "weight_decay": 0.01, + "max_coeff": 0.3, + "min_coeff": 0.01, + "freeze_step": 2, + "cuda_aware": False, + "comm_backend_name": "nccl", + "coeff_beta": 0.9, + "factor_max": 1.0, + "factor_min": 0.5, + "factor_threshold": 0.1 + } + }, + "gradient_clipping": 1.0, + "fp16": { + "enabled": (dtype == torch.float16), + "loss_scale": 0, + "initial_scale_power": 16 } - }, - "gradient_clipping": 1.0, - } - args = args_from_dict(tmpdir, config_dict) - hidden_dim = 10 - - model = SimpleModel(hidden_dim) + } + hidden_dim = 10 - @distributed_test(world_size=[1, 2]) - def _test_onebitlamb_fp32_basic(args, model, hidden_dim): - model, _, _, _ = deepspeed.initialize(args=args, + 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.float) + dtype=dtype) for n, batch in enumerate(data_loader): loss = model(batch[0], batch[1]) model.backward(loss) model.step() - _test_onebitlamb_fp32_basic(args=args, model=model, hidden_dim=hidden_dim) - - -def test_onebitlamb_exp_avg_mask(tmpdir): - config_dict = { - "train_batch_size": 2, - "steps_per_print": 1, - "optimizer": { - "type": "OneBitLamb", - "params": { - "lr": 0.00015, - "weight_decay": 0.01, - "max_coeff": 0.3, - "min_coeff": 0.01, - "freeze_step": 2, - "cuda_aware": False, - "comm_backend_name": "nccl", - "coeff_beta": 0.9, - "factor_max": 1.0, - "factor_min": 0.5, - "factor_threshold": 0.1 + +class TestOneBitLampExpAvgMask(DistributedTest): + world_size = 2 + + def test(self): + config_dict = { + "train_batch_size": 2, + "steps_per_print": 1, + "optimizer": { + "type": "OneBitLamb", + "params": { + "lr": 0.00015, + "weight_decay": 0.01, + "max_coeff": 0.3, + "min_coeff": 0.01, + "freeze_step": 2, + "cuda_aware": False, + "comm_backend_name": "nccl", + "coeff_beta": 0.9, + "factor_max": 1.0, + "factor_min": 0.5, + "factor_threshold": 0.1 + } + }, + "gradient_clipping": 1.0, + "fp16": { + "enabled": True, + "loss_scale": 0, + "initial_scale_power": 16 } - }, - "gradient_clipping": 1.0, - "fp16": { - "enabled": True, - "loss_scale": 0, - "initial_scale_power": 16 } - } - args = args_from_dict(tmpdir, config_dict) - hidden_dim = 10 - - model = SimpleModel(hidden_dim) - param_optimizer = list(model.named_parameters()) - mask1 = torch.zeros_like(param_optimizer[0][1].data) - for col in range(mask1.size()[1]): - mask1[0][col] += 1 - optimizer_grouped_parameters = [{ - 'params': [param_optimizer[0][1]], - 'weight_decay': 0.01, - 'exp_avg_mask': mask1 - }, - { - 'params': [param_optimizer[1][1]], - 'weight_decay': 0.01 - }] - - @distributed_test(world_size=[2]) - def _test_onebitlamb_exp_avg_mask(args, model, hidden_dim): - model, optimizer, _, _ = deepspeed.initialize(args=args, + hidden_dim = 10 + + model = SimpleModel(hidden_dim) + param_optimizer = list(model.named_parameters()) + mask1 = torch.zeros_like(param_optimizer[0][1].data) + for col in range(mask1.size()[1]): + mask1[0][col] += 1 + optimizer_grouped_parameters = [{ + 'params': [param_optimizer[0][1]], + 'weight_decay': 0.01, + 'exp_avg_mask': mask1 + }, + { + 'params': [param_optimizer[1][1]], + 'weight_decay': 0.01 + }] + + model, optimizer, _, _ = deepspeed.initialize(config=config_dict, model=model, model_parameters=optimizer_grouped_parameters) data_loader = random_dataloader(model=model, @@ -984,79 +812,77 @@ def _test_onebitlamb_exp_avg_mask(args, model, hidden_dim): if v['exp_avg'].size() == mask1.size(): assert torch.allclose(v['exp_avg'], v['exp_avg'].mul_(mask1.to(device=v['exp_avg'].device)), atol=1e-07), f"Momentum mask is not working properly" - _test_onebitlamb_exp_avg_mask(args=args, model=model, hidden_dim=hidden_dim) - - -def test_onebitlamb_checkpointing(tmpdir): - config_dict = { - "train_batch_size": 2, - "steps_per_print": 1, - "optimizer": { - "type": "OneBitLamb", - "params": { - "lr": 0.00015, - "weight_decay": 0.01, - "max_coeff": 0.3, - "min_coeff": 0.01, - "freeze_step": 2, - "cuda_aware": False, - "comm_backend_name": "nccl", - "coeff_beta": 0.9, - "factor_max": 1.0, - "factor_min": 0.5, - "factor_threshold": 0.1 + +class TestOneBitLambCheckpointing(DistributedTest): + world_size = 2 + + def test(self, tmpdir): + config_dict = { + "train_batch_size": 2, + "steps_per_print": 1, + "optimizer": { + "type": "OneBitLamb", + "params": { + "lr": 0.00015, + "weight_decay": 0.01, + "max_coeff": 0.3, + "min_coeff": 0.01, + "freeze_step": 2, + "cuda_aware": False, + "comm_backend_name": "nccl", + "coeff_beta": 0.9, + "factor_max": 1.0, + "factor_min": 0.5, + "factor_threshold": 0.1 + } + }, + "gradient_clipping": 1.0, + "fp16": { + "enabled": True, + "loss_scale": 0, + "initial_scale_power": 16 } - }, - "gradient_clipping": 1.0, - "fp16": { - "enabled": True, - "loss_scale": 0, - "initial_scale_power": 16 } - } - args = args_from_dict(tmpdir, config_dict) - hidden_dim = 10 - - model = SimpleModel(hidden_dim) - param_optimizer = list(model.named_parameters()) - mask1 = torch.zeros_like(param_optimizer[0][1].data) - mask2 = torch.zeros_like(param_optimizer[0][1].data) - for col in range(mask1.size()[1]): - mask1[0][col] += 1 - mask2[1][col] += 1 - - optimizer_grouped_parameters_1 = [{ - 'params': [param_optimizer[0][1]], - 'weight_decay': 0.01, - 'exp_avg_mask': mask1 - }, - { - 'params': [param_optimizer[1][1]], - 'weight_decay': 0.01 - }] - - optimizer_grouped_parameters_2 = [{ - 'params': [param_optimizer[0][1]], - 'weight_decay': 0.01, - 'exp_avg_mask': mask2 - }, - { - 'params': [param_optimizer[1][1]], - 'weight_decay': 0.01 - }] - - optimizer_grouped_parameters_3 = [{ - 'params': [param_optimizer[0][1]], - 'weight_decay': 0.01 - }, - { - 'params': [param_optimizer[1][1]], - 'weight_decay': 0.01 - }] - - @distributed_test(world_size=[2]) - def _test_onebitlamb_checkpointing(mask1, mask2, args, model, hidden_dim): - model_1, optimizer_1, _, _ = deepspeed.initialize(args=args, + hidden_dim = 10 + + model = SimpleModel(hidden_dim) + param_optimizer = list(model.named_parameters()) + mask1 = torch.zeros_like(param_optimizer[0][1].data) + mask2 = torch.zeros_like(param_optimizer[0][1].data) + for col in range(mask1.size()[1]): + mask1[0][col] += 1 + mask2[1][col] += 1 + + optimizer_grouped_parameters_1 = [{ + 'params': [param_optimizer[0][1]], + 'weight_decay': 0.01, + 'exp_avg_mask': mask1 + }, + { + 'params': [param_optimizer[1][1]], + 'weight_decay': 0.01 + }] + + optimizer_grouped_parameters_2 = [{ + 'params': [param_optimizer[0][1]], + 'weight_decay': 0.01, + 'exp_avg_mask': mask2 + }, + { + 'params': [param_optimizer[1][1]], + 'weight_decay': 0.01 + }] + + optimizer_grouped_parameters_3 = [{ + 'params': [param_optimizer[0][1]], + 'weight_decay': 0.01 + }, + { + 'params': [param_optimizer[1][1]], + 'weight_decay': 0.01 + }] + + model_1, optimizer_1, _, _ = deepspeed.initialize(config=config_dict, model=model, model_parameters=optimizer_grouped_parameters_1) data_loader = random_dataloader(model=model_1, @@ -1080,7 +906,7 @@ def _test_onebitlamb_checkpointing(mask1, mask2, args, model, hidden_dim): assert torch.allclose(optimizer_1.param_groups[0]['exp_avg_mask'], mask1, atol=1e-07), f"Momentum mask should not change after saving checkpoint" - model_2, optimizer_2, _, _ = deepspeed.initialize(args=args, + model_2, optimizer_2, _, _ = deepspeed.initialize(config=config_dict, model=model, model_parameters=optimizer_grouped_parameters_2) # Test whether momentum mask stays the same after loading checkpoint @@ -1102,7 +928,7 @@ def _test_onebitlamb_checkpointing(mask1, mask2, args, model, hidden_dim): assert list(sorted(scaling_coeff_2)) == list(sorted(scaling_coeff_1)), f"Incorrect scaling_coeffs" assert optimizer_2.optimizer.lamb_freeze_key is True - model_3, optimizer_3, _, _ = deepspeed.initialize(args=args, + model_3, optimizer_3, _, _ = deepspeed.initialize(config=config_dict, model=model, model_parameters=optimizer_grouped_parameters_3) optimizer_3.optimizer.freeze_step = 20 @@ -1132,48 +958,37 @@ def _test_onebitlamb_checkpointing(mask1, mask2, args, model, hidden_dim): assert 'scaling_coeff' not in v, f"Incorrect scaling_coeff" assert optimizer_3.optimizer.lamb_freeze_key is False - _test_onebitlamb_checkpointing(mask1, - mask2, - args=args, - model=model, - hidden_dim=hidden_dim) - - -def test_onebitlamb_checkpointing_overflow(tmpdir): - config_dict = { - "train_batch_size": 2, - "steps_per_print": 1, - "optimizer": { - "type": "OneBitLamb", - "params": { - "lr": 0.00015, - "weight_decay": 0.01, - "max_coeff": 0.3, - "min_coeff": 0.01, - "freeze_step": 2, - "cuda_aware": False, - "comm_backend_name": "nccl", - "coeff_beta": 0.9, - "factor_max": 1.0, - "factor_min": 0.5, - "factor_threshold": 0.1 + def test_overflow(self, tmpdir): + config_dict = { + "train_batch_size": 2, + "steps_per_print": 1, + "optimizer": { + "type": "OneBitLamb", + "params": { + "lr": 0.00015, + "weight_decay": 0.01, + "max_coeff": 0.3, + "min_coeff": 0.01, + "freeze_step": 2, + "cuda_aware": False, + "comm_backend_name": "nccl", + "coeff_beta": 0.9, + "factor_max": 1.0, + "factor_min": 0.5, + "factor_threshold": 0.1 + } + }, + "gradient_clipping": 1.0, + "fp16": { + "enabled": True, + "loss_scale": 0, + "initial_scale_power": 16 } - }, - "gradient_clipping": 1.0, - "fp16": { - "enabled": True, - "loss_scale": 0, - "initial_scale_power": 16 } - } - args = args_from_dict(tmpdir, config_dict) - hidden_dim = 10 - - model = SimpleModel(hidden_dim) + hidden_dim = 10 - @distributed_test(world_size=[2]) - def _test_onebitlamb_checkpointing_overflow(args, model, hidden_dim): - model, _, _, _ = deepspeed.initialize(args=args, + model = SimpleModel(hidden_dim) + model, _, _, _ = deepspeed.initialize(config=config_dict, model=model, model_parameters=model.parameters()) data_loader = random_dataloader(model=model, @@ -1191,10 +1006,6 @@ def _test_onebitlamb_checkpointing_overflow(args, model, hidden_dim): dist.barrier() model.save_checkpoint(save_folder, tag=None) - _test_onebitlamb_checkpointing_overflow(args=args, - model=model, - hidden_dim=hidden_dim) - @pytest.mark.parametrize('topo', [ @@ -1205,45 +1016,46 @@ def _test_onebitlamb_checkpointing_overflow(args, model, hidden_dim): PipeTopo(num_pp=4, num_dp=1), ]) -def test_onebitlamb_fp16_pipeline(topo, tmpdir): - config_dict = { - "train_batch_size": 16, - "train_micro_batch_size_per_gpu": 4, - "steps_per_print": 20, - "optimizer": { - "type": "OneBitLamb", - "params": { - "lr": 0.00001, - "betas": [0.9, - 0.999], - "eps": 1e-8, - "weight_decay": 3e-7, - "freeze_step": 200, - "cuda_aware": False, - "comm_backend_name": "nccl" +class TestOneBitLambFP16Pipeline(DistributedTest): + world_size = 4 + + def test(self, topo): + config_dict = { + "train_batch_size": 16, + "train_micro_batch_size_per_gpu": 4, + "steps_per_print": 20, + "optimizer": { + "type": "OneBitLamb", + "params": { + "lr": 0.00001, + "betas": [0.9, + 0.999], + "eps": 1e-8, + "weight_decay": 3e-7, + "freeze_step": 200, + "cuda_aware": False, + "comm_backend_name": "nccl" + } + }, + "gradient_clipping": 1.0, + "zero_optimization": { + "stage": 0 + }, + "fp16": { + "enabled": True, + "loss_scale": 0, + "initial_scale_power": 16 + }, + "pipeline": { + "seed_layers": True, + "activation_checkpoint_interval": 1 } - }, - "gradient_clipping": 1.0, - "zero_optimization": { - "stage": 0 - }, - "fp16": { - "enabled": True, - "loss_scale": 0, - "initial_scale_power": 16 - }, - "pipeline": { - "seed_layers": True, - "activation_checkpoint_interval": 1 } - } - args = args_from_dict(tmpdir, config_dict) - # Allocate model for consistent initial weights. - init_net = AlexNetPipe() + # Allocate model for consistent initial weights. + init_net = AlexNetPipe() - @distributed_test(world_size=4) - def _helper(topo, tmpdir, steps=500): + steps = 500 assert steps >= 100 test_net = copy.deepcopy(init_net) @@ -1252,17 +1064,16 @@ def _helper(topo, tmpdir, steps=500): loss_fn=nn.CrossEntropyLoss()) test_losses = train_cifar(test_model, - args, + config_dict=config_dict, num_steps=steps, fp16=config_dict['fp16']['enabled']) - _helper(topo, tmpdir) - @pytest.mark.sequential -def test_compressed_allreduce_basic(tmpdir): - @distributed_test(world_size=[1, 2]) - def _test_compressed_allreduce_basic(): +class TestCompressedAllReduceBasic(DistributedTest): + world_size = [1, 2] + + def test(self, tmpdir): from deepspeed.runtime.comm.nccl import NcclBackend size = dist.get_world_size() rank = dist.get_rank() @@ -1327,5 +1138,3 @@ def torch_sim(a): if torch.sum(check_mag_mask) != 0: print('Fails at {} of positions'.format(torch.sum(check_mag_mask))) assert torch.sum(diff_server_mask) == 0 or torch.sum(check_mag_mask) == 0 - - _test_compressed_allreduce_basic() diff --git a/tests/unit/test_pipe.py b/tests/unit/test_pipe.py index de1bd3ff279c..74c46cd0c536 100755 --- a/tests/unit/test_pipe.py +++ b/tests/unit/test_pipe.py @@ -136,7 +136,12 @@ def cifar_trainset(fp16=False): return trainset -def train_cifar(model, args, num_steps=400, average_dp_losses=True, fp16=True, seed=123): +def train_cifar(model, + config_dict, + num_steps=400, + average_dp_losses=True, + fp16=True, + seed=123): with torch.random.fork_rng(devices=[torch.cuda.current_device()]): ds_utils.set_random_seed(seed) @@ -144,10 +149,10 @@ def train_cifar(model, args, num_steps=400, average_dp_losses=True, fp16=True, s model.eval() trainset = cifar_trainset(fp16=fp16) - args.local_rank = dist.get_rank() + config_dict['local_rank'] = dist.get_rank() engine, _, _, _ = deepspeed.initialize( - args=args, + config=config_dict, model=model, model_parameters=[p for p in model.parameters()], training_data=trainset) From 300b7a41fc170d0af64ae6d87ba1a80bc7b0e302 Mon Sep 17 00:00:00 2001 From: Michael Wyatt Date: Thu, 4 Aug 2022 10:15:17 -0700 Subject: [PATCH 02/26] fixed topology pickling issue --- tests/unit/test_onebit.py | 944 ++++++++++++++++++++++---------------- 1 file changed, 554 insertions(+), 390 deletions(-) diff --git a/tests/unit/test_onebit.py b/tests/unit/test_onebit.py index 41fbd1c59099..f7b3c69bdded 100644 --- a/tests/unit/test_onebit.py +++ b/tests/unit/test_onebit.py @@ -15,11 +15,13 @@ from .simple_model import SimpleModel, random_dataloader from .test_pipe import AlexNetPipe, train_cifar -TORCH_MAJOR = int(torch.__version__.split('.')[0]) -TORCH_MINOR = int(torch.__version__.split('.')[1]) +TORCH_MAJOR = int(torch.__version__.split(".")[0]) +TORCH_MINOR = int(torch.__version__.split(".")[1]) if TORCH_MAJOR < 1 or TORCH_MINOR < 8: - pytest.skip("NCCL-based 1-bit compression requires torch 1.8 or higher", - allow_module_level=True) + pytest.skip( + "NCCL-based 1-bit compression requires torch 1.8 or higher", + allow_module_level=True, + ) @pytest.mark.parametrize("dtype", [torch.float32, torch.float16], ids=["fp32", "fp16"]) @@ -37,27 +39,29 @@ def test(self, dtype): "weight_decay": 0.01, "freeze_step": 2, "cuda_aware": False, - "comm_backend_name": "nccl" - } + "comm_backend_name": "nccl", + }, }, "gradient_clipping": 1.0, "fp16": { "enabled": (dtype == torch.float16), "loss_scale": 0, - "initial_scale_power": 16 - } + "initial_scale_power": 16, + }, } 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=dtype) + 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=dtype, + ) for n, batch in enumerate(data_loader): loss = model(batch[0], batch[1]) model.backward(loss) @@ -78,15 +82,15 @@ def test(self): "weight_decay": 0.01, "freeze_step": 2, "cuda_aware": False, - "comm_backend_name": "nccl" - } + "comm_backend_name": "nccl", + }, }, "gradient_clipping": 1.0, "fp16": { "enabled": True, "loss_scale": 0, "initial_scale_power": 16 - } + }, } hidden_dim = 10 @@ -96,19 +100,23 @@ def test(self): for col in range(mask1.size()[1]): mask1[0][col] += 1 mask1 = torch.flatten(mask1) - optimizer_grouped_parameters = [{ - 'params': [param_optimizer[0][1]], - 'weight_decay': 0.01, - 'exp_avg_mask': mask1 - }, - { - 'params': [param_optimizer[1][1]], - 'weight_decay': 0.01 - }] - - model, optimizer, _, _ = deepspeed.initialize(config=config_dict, - model=model, - model_parameters=optimizer_grouped_parameters) + optimizer_grouped_parameters = [ + { + "params": [param_optimizer[0][1]], + "weight_decay": 0.01, + "exp_avg_mask": mask1, + }, + { + "params": [param_optimizer[1][1]], + "weight_decay": 0.01 + }, + ] + + model, optimizer, _, _ = deepspeed.initialize( + config=config_dict, + model=model, + model_parameters=optimizer_grouped_parameters, + ) data_loader = random_dataloader(model=model, total_samples=50, hidden_dim=hidden_dim, @@ -119,8 +127,12 @@ def test(self): model.step() # Test whether the momentum mask works for v in optimizer.state.values(): - if v['exp_avg'].size() == mask1.size(): - assert torch.allclose(v['exp_avg'], v['exp_avg'].mul_(mask1.to(device=v['exp_avg'].device)), atol=1e-07), f"Momentum mask is not working properly" + if v["exp_avg"].size() == mask1.size(): + assert torch.allclose( + v["exp_avg"], + v["exp_avg"].mul_(mask1.to(device=v["exp_avg"].device)), + atol=1e-07, + ), f"Momentum mask is not working properly" class TestOneBitAdamCheckpointing(DistributedTest): @@ -137,15 +149,15 @@ def test(self, tmpdir): "weight_decay": 0.01, "freeze_step": 2, "cuda_aware": False, - "comm_backend_name": "nccl" - } + "comm_backend_name": "nccl", + }, }, "gradient_clipping": 1.0, "fp16": { "enabled": True, "loss_scale": 0, "initial_scale_power": 16 - } + }, } hidden_dim = 10 @@ -159,96 +171,127 @@ def test(self, tmpdir): mask1 = torch.flatten(mask1) mask2 = torch.flatten(mask2) - optimizer_grouped_parameters_1 = [{ - 'params': [param_optimizer[0][1]], - 'weight_decay': 0.01, - 'exp_avg_mask': mask1 - }, - { - 'params': [param_optimizer[1][1]], - 'weight_decay': 0.01 - }] - - optimizer_grouped_parameters_2 = [{ - 'params': [param_optimizer[0][1]], - 'weight_decay': 0.01, - 'exp_avg_mask': mask2 - }, - { - 'params': [param_optimizer[1][1]], - 'weight_decay': 0.01 - }] - - optimizer_grouped_parameters_3 = [{ - 'params': [param_optimizer[0][1]], - 'weight_decay': 0.01 - }, - { - 'params': [param_optimizer[1][1]], - 'weight_decay': 0.01 - }] - - model_1, optimizer_1, _, _ = deepspeed.initialize(config=config_dict, - model=model, - model_parameters=optimizer_grouped_parameters_1) - data_loader = random_dataloader(model=model_1, - total_samples=10, - hidden_dim=hidden_dim, - device=model_1.device) + optimizer_grouped_parameters_1 = [ + { + "params": [param_optimizer[0][1]], + "weight_decay": 0.01, + "exp_avg_mask": mask1, + }, + { + "params": [param_optimizer[1][1]], + "weight_decay": 0.01 + }, + ] + + optimizer_grouped_parameters_2 = [ + { + "params": [param_optimizer[0][1]], + "weight_decay": 0.01, + "exp_avg_mask": mask2, + }, + { + "params": [param_optimizer[1][1]], + "weight_decay": 0.01 + }, + ] + + optimizer_grouped_parameters_3 = [ + { + "params": [param_optimizer[0][1]], + "weight_decay": 0.01 + }, + { + "params": [param_optimizer[1][1]], + "weight_decay": 0.01 + }, + ] + + model_1, optimizer_1, _, _ = deepspeed.initialize( + config=config_dict, + model=model, + model_parameters=optimizer_grouped_parameters_1, + ) + data_loader = random_dataloader( + model=model_1, + total_samples=10, + hidden_dim=hidden_dim, + device=model_1.device, + ) for n, batch in enumerate(data_loader): loss = model_1(batch[0], batch[1]) model_1.backward(loss) model_1.step() # Test whether momentum mask still exist after saving checkpoint assert optimizer_1.optimizer.adam_freeze_key is True - mask1 = mask1.to(device=optimizer_1.param_groups[0]['exp_avg_mask'].device) - assert torch.allclose(optimizer_1.param_groups[0]['exp_avg_mask'], mask1, atol=1e-07), f"Incorrect momentum mask" - save_folder = os.path.join(tmpdir, 'saved_checkpoint') + mask1 = mask1.to(device=optimizer_1.param_groups[0]["exp_avg_mask"].device) + assert torch.allclose( + optimizer_1.param_groups[0]["exp_avg_mask"], mask1, atol=1e-07 + ), f"Incorrect momentum mask" + save_folder = os.path.join(tmpdir, "saved_checkpoint") model_1.save_checkpoint(save_folder, tag=None) - assert torch.allclose(optimizer_1.param_groups[0]['exp_avg_mask'], mask1, atol=1e-07), f"Momentum mask should not change after saving checkpoint" - - - model_2, optimizer_2, _, _ = deepspeed.initialize(config=config_dict, - model=model, - model_parameters=optimizer_grouped_parameters_2) + assert torch.allclose( + optimizer_1.param_groups[0]["exp_avg_mask"], mask1, atol=1e-07 + ), f"Momentum mask should not change after saving checkpoint" + + model_2, optimizer_2, _, _ = deepspeed.initialize( + config=config_dict, + model=model, + model_parameters=optimizer_grouped_parameters_2, + ) # Test whether momentum mask stays the same after loading checkpoint - mask2 = mask2.to(device=optimizer_2.param_groups[0]['exp_avg_mask'].device) - assert torch.allclose(optimizer_2.param_groups[0]['exp_avg_mask'], mask2, atol=1e-07), f"Incorrect momentum mask" - model_2.load_checkpoint(save_folder, - tag=None, - load_optimizer_states=True, - load_lr_scheduler_states=True) - assert torch.allclose(optimizer_2.param_groups[0]['exp_avg_mask'], mask2, atol=1e-07), f"Momentum mask should not change after loading checkpoint" + mask2 = mask2.to(device=optimizer_2.param_groups[0]["exp_avg_mask"].device) + assert torch.allclose( + optimizer_2.param_groups[0]["exp_avg_mask"], mask2, atol=1e-07 + ), f"Incorrect momentum mask" + model_2.load_checkpoint( + save_folder, + tag=None, + load_optimizer_states=True, + load_lr_scheduler_states=True, + ) + assert torch.allclose( + optimizer_2.param_groups[0]["exp_avg_mask"], mask2, atol=1e-07 + ), f"Momentum mask should not change after loading checkpoint" # Test whether worker&server error is reset for v in optimizer_2.state.values(): - assert 'worker_error' not in v, f"Incorrect worker error" - assert 'server_error' not in v, f"Incorrect server error" + assert "worker_error" not in v, f"Incorrect worker error" + assert "server_error" not in v, f"Incorrect server error" assert optimizer_2.optimizer.adam_freeze_key is True - model_3, optimizer_3, _, _ = deepspeed.initialize(config=config_dict, - model=model, - model_parameters=optimizer_grouped_parameters_3) + model_3, optimizer_3, _, _ = deepspeed.initialize( + config=config_dict, + model=model, + model_parameters=optimizer_grouped_parameters_3, + ) optimizer_3.optimizer.freeze_step = 20 - data_loader = random_dataloader(model=model_3, - total_samples=50, - hidden_dim=hidden_dim, - device=model_3.device) + data_loader = random_dataloader( + model=model_3, + total_samples=50, + hidden_dim=hidden_dim, + device=model_3.device, + ) for n, batch in enumerate(data_loader): loss = model_3(batch[0], batch[1]) model_3.backward(loss) model_3.step() assert optimizer_3.optimizer.adam_freeze_key is True # Test whether momentum mask stays the same after loading checkpoint - assert 'exp_avg_mask' not in optimizer_3.param_groups[0], f"Incorrect momentum mask" - model_3.load_checkpoint(save_folder, - tag=None, - load_optimizer_states=True, - load_lr_scheduler_states=True) - assert 'exp_avg_mask' not in optimizer_3.param_groups[0], f"Momentum mask should not change after loading checkpoint" + assert ( + "exp_avg_mask" not in optimizer_3.param_groups[0] + ), f"Incorrect momentum mask" + model_3.load_checkpoint( + save_folder, + tag=None, + load_optimizer_states=True, + load_lr_scheduler_states=True, + ) + assert ( + "exp_avg_mask" not in optimizer_3.param_groups[0] + ), f"Momentum mask should not change after loading checkpoint" # Test whether worker&server error is reset for v in optimizer_3.state.values(): - assert 'worker_error' not in v, f"Incorrect worker error" - assert 'server_error' not in v, f"Incorrect server error" + assert "worker_error" not in v, f"Incorrect worker error" + assert "server_error" not in v, f"Incorrect server error" assert optimizer_3.optimizer.adam_freeze_key is False def test_overflow(self, tmpdir): @@ -262,27 +305,27 @@ def test_overflow(self, tmpdir): "weight_decay": 0.01, "freeze_step": 2, "cuda_aware": False, - "comm_backend_name": "nccl" - } + "comm_backend_name": "nccl", + }, }, "gradient_clipping": 1.0, "fp16": { "enabled": True, "loss_scale": 0, "initial_scale_power": 16 - } + }, } hidden_dim = 10 model = SimpleModel(hidden_dim) - model, _, _, _ = deepspeed.initialize(config=config_dict, - model=model, - model_parameters=model.parameters()) + model, _, _, _ = deepspeed.initialize( + config=config_dict, model=model, model_parameters=model.parameters() + ) data_loader = random_dataloader(model=model, total_samples=100, hidden_dim=hidden_dim, device=model.device) - save_folder = os.path.join(tmpdir, 'saved_checkpoint') + save_folder = os.path.join(tmpdir, "saved_checkpoint") for n, batch in enumerate(data_loader): loss = model(batch[0], batch[1]) if dist.get_rank() == 0 and n >= 10: @@ -294,19 +337,27 @@ def test_overflow(self, tmpdir): model.save_checkpoint(save_folder, tag=None) -@pytest.mark.parametrize('topo', - [ - PipeTopo(num_pp=1, - num_dp=4), - PipeTopo(num_pp=2, - num_dp=2), - PipeTopo(num_pp=4, - num_dp=1), - ]) +@pytest.mark.parametrize( + "topo_config", + [ + { + "num_pp": 1, + "num_dp": 4 + }, + { + "num_pp": 2, + "num_dp": 2 + }, + { + "num_pp": 4, + "num_dp": 1 + }, + ], +) class TestOneBitAdamFP16Pipeline(DistributedTest): world_size = 4 - def test(self, topo): + def test(self, topo_config): config_dict = { "train_batch_size": 16, "train_micro_batch_size_per_gpu": 4, @@ -321,8 +372,8 @@ def test(self, topo): "weight_decay": 3e-7, "freeze_step": 200, "cuda_aware": False, - "comm_backend_name": "nccl" - } + "comm_backend_name": "nccl", + }, }, "gradient_clipping": 1.0, "zero_optimization": { @@ -336,9 +387,11 @@ def test(self, topo): "pipeline": { "seed_layers": True, "activation_checkpoint_interval": 1 - } + }, } + topo = PipeTopo(**topo_config) + # Allocate model for consistent initial weights. init_net = AlexNetPipe() @@ -350,10 +403,12 @@ def test(self, topo): topology=topo, loss_fn=nn.CrossEntropyLoss()) - test_losses = train_cifar(test_model, - config_dict=config_dict, - num_steps=steps, - fp16=config_dict['fp16']['enabled']) + test_losses = train_cifar( + test_model, + config_dict=config_dict, + num_steps=steps, + fp16=config_dict["fp16"]["enabled"], + ) @pytest.mark.parametrize("dtype", [torch.float32, torch.float16], ids=["fp32", "fp16"]) @@ -374,27 +429,29 @@ def test(self, dtype): "local_step_scaler": 1, "local_step_clipper": 2, "cuda_aware": False, - "comm_backend_name": "nccl" - } + "comm_backend_name": "nccl", + }, }, "gradient_clipping": 1.0, "fp16": { "enabled": (dtype == torch.float16), "loss_scale": 0, - "initial_scale_power": 16 - } + "initial_scale_power": 16, + }, } 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=dtype) + 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=dtype, + ) for n, batch in enumerate(data_loader): loss = model(batch[0], batch[1]) model.backward(loss) @@ -418,15 +475,15 @@ def test(self): "local_step_scaler": 1, "local_step_clipper": 2, "cuda_aware": False, - "comm_backend_name": "nccl" - } + "comm_backend_name": "nccl", + }, }, "gradient_clipping": 1.0, "fp16": { "enabled": True, "loss_scale": 0, "initial_scale_power": 16 - } + }, } hidden_dim = 10 @@ -436,19 +493,23 @@ def test(self): for col in range(mask1.size()[1]): mask1[0][col] += 1 mask1 = torch.flatten(mask1) - optimizer_grouped_parameters = [{ - 'params': [param_optimizer[0][1]], - 'weight_decay': 0.01, - 'exp_avg_mask': mask1 - }, - { - 'params': [param_optimizer[1][1]], - 'weight_decay': 0.01 - }] - - model, optimizer, _, _ = deepspeed.initialize(config=config_dict, - model=model, - model_parameters=optimizer_grouped_parameters) + optimizer_grouped_parameters = [ + { + "params": [param_optimizer[0][1]], + "weight_decay": 0.01, + "exp_avg_mask": mask1, + }, + { + "params": [param_optimizer[1][1]], + "weight_decay": 0.01 + }, + ] + + model, optimizer, _, _ = deepspeed.initialize( + config=config_dict, + model=model, + model_parameters=optimizer_grouped_parameters, + ) data_loader = random_dataloader(model=model, total_samples=50, hidden_dim=hidden_dim, @@ -459,8 +520,12 @@ def test(self): model.step() # Test whether the momentum mask works for v in optimizer.state.values(): - if v['exp_avg'].size() == mask1.size(): - assert torch.allclose(v['exp_avg'], v['exp_avg'].mul_(mask1.to(device=v['exp_avg'].device)), atol=1e-07), f"Momentum mask is not working properly" + if v["exp_avg"].size() == mask1.size(): + assert torch.allclose( + v["exp_avg"], + v["exp_avg"].mul_(mask1.to(device=v["exp_avg"].device)), + atol=1e-07, + ), f"Momentum mask is not working properly" class TestZeroOneAdamCheckpointing(DistributedTest): @@ -480,15 +545,15 @@ def test(self, tmpdir): "local_step_scaler": 1, "local_step_clipper": 2, "cuda_aware": False, - "comm_backend_name": "nccl" - } + "comm_backend_name": "nccl", + }, }, "gradient_clipping": 1.0, "fp16": { "enabled": True, "loss_scale": 0, "initial_scale_power": 16 - } + }, } hidden_dim = 10 @@ -502,93 +567,124 @@ def test(self, tmpdir): mask1 = torch.flatten(mask1) mask2 = torch.flatten(mask2) - optimizer_grouped_parameters_1 = [{ - 'params': [param_optimizer[0][1]], - 'weight_decay': 0.01, - 'exp_avg_mask': mask1 - }, - { - 'params': [param_optimizer[1][1]], - 'weight_decay': 0.01 - }] - - optimizer_grouped_parameters_2 = [{ - 'params': [param_optimizer[0][1]], - 'weight_decay': 0.01, - 'exp_avg_mask': mask2 - }, - { - 'params': [param_optimizer[1][1]], - 'weight_decay': 0.01 - }] - - optimizer_grouped_parameters_3 = [{ - 'params': [param_optimizer[0][1]], - 'weight_decay': 0.01 - }, - { - 'params': [param_optimizer[1][1]], - 'weight_decay': 0.01 - }] - - model_1, optimizer_1, _, _ = deepspeed.initialize(config=config_dict, - model=model, - model_parameters=optimizer_grouped_parameters_1) - data_loader = random_dataloader(model=model_1, - total_samples=10, - hidden_dim=hidden_dim, - device=model_1.device) + optimizer_grouped_parameters_1 = [ + { + "params": [param_optimizer[0][1]], + "weight_decay": 0.01, + "exp_avg_mask": mask1, + }, + { + "params": [param_optimizer[1][1]], + "weight_decay": 0.01 + }, + ] + + optimizer_grouped_parameters_2 = [ + { + "params": [param_optimizer[0][1]], + "weight_decay": 0.01, + "exp_avg_mask": mask2, + }, + { + "params": [param_optimizer[1][1]], + "weight_decay": 0.01 + }, + ] + + optimizer_grouped_parameters_3 = [ + { + "params": [param_optimizer[0][1]], + "weight_decay": 0.01 + }, + { + "params": [param_optimizer[1][1]], + "weight_decay": 0.01 + }, + ] + + model_1, optimizer_1, _, _ = deepspeed.initialize( + config=config_dict, + model=model, + model_parameters=optimizer_grouped_parameters_1, + ) + data_loader = random_dataloader( + model=model_1, + total_samples=10, + hidden_dim=hidden_dim, + device=model_1.device, + ) for n, batch in enumerate(data_loader): loss = model_1(batch[0], batch[1]) model_1.backward(loss) model_1.step() # Test whether momentum mask still exist after saving checkpoint - mask1 = mask1.to(device=optimizer_1.param_groups[0]['exp_avg_mask'].device) - assert torch.allclose(optimizer_1.param_groups[0]['exp_avg_mask'], mask1, atol=1e-07), f"Incorrect momentum mask" - save_folder = os.path.join(tmpdir, 'saved_checkpoint') + mask1 = mask1.to(device=optimizer_1.param_groups[0]["exp_avg_mask"].device) + assert torch.allclose( + optimizer_1.param_groups[0]["exp_avg_mask"], mask1, atol=1e-07 + ), f"Incorrect momentum mask" + save_folder = os.path.join(tmpdir, "saved_checkpoint") model_1.save_checkpoint(save_folder, tag=None) - assert torch.allclose(optimizer_1.param_groups[0]['exp_avg_mask'], mask1, atol=1e-07), f"Momentum mask should not change after saving checkpoint" - - - model_2, optimizer_2, _, _ = deepspeed.initialize(config=config_dict, - model=model, - model_parameters=optimizer_grouped_parameters_2) + assert torch.allclose( + optimizer_1.param_groups[0]["exp_avg_mask"], mask1, atol=1e-07 + ), f"Momentum mask should not change after saving checkpoint" + + model_2, optimizer_2, _, _ = deepspeed.initialize( + config=config_dict, + model=model, + model_parameters=optimizer_grouped_parameters_2, + ) # Test whether momentum mask stays the same after loading checkpoint - mask2 = mask2.to(device=optimizer_2.param_groups[0]['exp_avg_mask'].device) - assert torch.allclose(optimizer_2.param_groups[0]['exp_avg_mask'], mask2, atol=1e-07), f"Incorrect momentum mask" - model_2.load_checkpoint(save_folder, - tag=None, - load_optimizer_states=True, - load_lr_scheduler_states=True) - assert torch.allclose(optimizer_2.param_groups[0]['exp_avg_mask'], mask2, atol=1e-07), f"Momentum mask should not change after loading checkpoint" + mask2 = mask2.to(device=optimizer_2.param_groups[0]["exp_avg_mask"].device) + assert torch.allclose( + optimizer_2.param_groups[0]["exp_avg_mask"], mask2, atol=1e-07 + ), f"Incorrect momentum mask" + model_2.load_checkpoint( + save_folder, + tag=None, + load_optimizer_states=True, + load_lr_scheduler_states=True, + ) + assert torch.allclose( + optimizer_2.param_groups[0]["exp_avg_mask"], mask2, atol=1e-07 + ), f"Momentum mask should not change after loading checkpoint" # Test whether worker&server error is reset for v in optimizer_2.state.values(): - assert 'worker_error' not in v, f"Incorrect worker error" - assert 'server_error' not in v, f"Incorrect server error" - - model_3, optimizer_3, _, _ = deepspeed.initialize(config=config_dict, - model=model, - model_parameters=optimizer_grouped_parameters_3) + assert "worker_error" not in v, f"Incorrect worker error" + assert "server_error" not in v, f"Incorrect server error" + + model_3, optimizer_3, _, _ = deepspeed.initialize( + config=config_dict, + model=model, + model_parameters=optimizer_grouped_parameters_3, + ) optimizer_3.optimizer.freeze_step = 20 - data_loader = random_dataloader(model=model_3, - total_samples=50, - hidden_dim=hidden_dim, - device=model_3.device) + data_loader = random_dataloader( + model=model_3, + total_samples=50, + hidden_dim=hidden_dim, + device=model_3.device, + ) for n, batch in enumerate(data_loader): loss = model_3(batch[0], batch[1]) model_3.backward(loss) model_3.step() # Test whether momentum mask stays the same after loading checkpoint - assert 'exp_avg_mask' not in optimizer_3.param_groups[0], f"Incorrect momentum mask" - model_3.load_checkpoint(save_folder, - tag=None, - load_optimizer_states=True, - load_lr_scheduler_states=True) - assert 'exp_avg_mask' not in optimizer_3.param_groups[0], f"Momentum mask should not change after loading checkpoint" + assert ( + "exp_avg_mask" not in optimizer_3.param_groups[0] + ), f"Incorrect momentum mask" + model_3.load_checkpoint( + save_folder, + tag=None, + load_optimizer_states=True, + load_lr_scheduler_states=True, + ) + assert ( + "exp_avg_mask" not in optimizer_3.param_groups[0] + ), f"Momentum mask should not change after loading checkpoint" # Test whether worker&server error is reset for v in optimizer_3.state.values(): - assert 'worker_error' not in v, f"Incorrect worker error" - assert 'server_error' not in v, f"Incorrect server error" + assert "worker_error" not in v, f"Incorrect worker error" + assert "server_error" not in v, f"Incorrect server error" def test_overflow(self, tmpdir): config_dict = { @@ -604,27 +700,27 @@ def test_overflow(self, tmpdir): "local_step_scaler": 1, "local_step_clipper": 2, "cuda_aware": False, - "comm_backend_name": "nccl" - } + "comm_backend_name": "nccl", + }, }, "gradient_clipping": 1.0, "fp16": { "enabled": True, "loss_scale": 0, "initial_scale_power": 16 - } + }, } hidden_dim = 10 model = SimpleModel(hidden_dim) - model, _, _, _ = deepspeed.initialize(config=config_dict, - model=model, - model_parameters=model.parameters()) + model, _, _, _ = deepspeed.initialize( + config=config_dict, model=model, model_parameters=model.parameters() + ) data_loader = random_dataloader(model=model, total_samples=100, hidden_dim=hidden_dim, device=model.device) - save_folder = os.path.join(tmpdir, 'saved_checkpoint') + save_folder = os.path.join(tmpdir, "saved_checkpoint") for n, batch in enumerate(data_loader): loss = model(batch[0], batch[1]) if dist.get_rank() == 0 and n >= 10: @@ -636,19 +732,27 @@ def test_overflow(self, tmpdir): model.save_checkpoint(save_folder, tag=None) -@pytest.mark.parametrize('topo', - [ - PipeTopo(num_pp=1, - num_dp=4), - PipeTopo(num_pp=2, - num_dp=2), - PipeTopo(num_pp=4, - num_dp=1), - ]) +@pytest.mark.parametrize( + "topo_config", + [ + { + "num_pp": 1, + "num_dp": 4 + }, + { + "num_pp": 2, + "num_dp": 2 + }, + { + "num_pp": 4, + "num_dp": 1 + }, + ], +) class TestZeroOneAdamFP16Pipeline(DistributedTest): world_size = 4 - def test(self, topo): + def test(self, topo_config): config_dict = { "train_batch_size": 16, "train_micro_batch_size_per_gpu": 4, @@ -666,8 +770,8 @@ def test(self, topo): "local_step_scaler": 1, "local_step_clipper": 2, "cuda_aware": False, - "comm_backend_name": "nccl" - } + "comm_backend_name": "nccl", + }, }, "gradient_clipping": 1.0, "zero_optimization": { @@ -681,9 +785,11 @@ def test(self, topo): "pipeline": { "seed_layers": True, "activation_checkpoint_interval": 1 - } + }, } + topo = PipeTopo(**topo_config) + # Allocate model for consistent initial weights. init_net = AlexNetPipe() @@ -695,10 +801,12 @@ def test(self, topo): topology=topo, loss_fn=nn.CrossEntropyLoss()) - test_losses = train_cifar(test_model, - config_dict=config_dict, - num_steps=steps, - fp16=config_dict['fp16']['enabled']) + test_losses = train_cifar( + test_model, + config_dict=config_dict, + num_steps=steps, + fp16=config_dict["fp16"]["enabled"], + ) @pytest.mark.parametrize("dtype", [torch.float32, torch.float16], ids=["fp32", "fp16"]) @@ -722,27 +830,29 @@ def test(self, dtype): "coeff_beta": 0.9, "factor_max": 1.0, "factor_min": 0.5, - "factor_threshold": 0.1 - } + "factor_threshold": 0.1, + }, }, "gradient_clipping": 1.0, "fp16": { "enabled": (dtype == torch.float16), "loss_scale": 0, - "initial_scale_power": 16 - } + "initial_scale_power": 16, + }, } 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=dtype) + 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=dtype, + ) for n, batch in enumerate(data_loader): loss = model(batch[0], batch[1]) model.backward(loss) @@ -769,15 +879,15 @@ def test(self): "coeff_beta": 0.9, "factor_max": 1.0, "factor_min": 0.5, - "factor_threshold": 0.1 - } + "factor_threshold": 0.1, + }, }, "gradient_clipping": 1.0, "fp16": { "enabled": True, "loss_scale": 0, "initial_scale_power": 16 - } + }, } hidden_dim = 10 @@ -786,19 +896,23 @@ def test(self): mask1 = torch.zeros_like(param_optimizer[0][1].data) for col in range(mask1.size()[1]): mask1[0][col] += 1 - optimizer_grouped_parameters = [{ - 'params': [param_optimizer[0][1]], - 'weight_decay': 0.01, - 'exp_avg_mask': mask1 - }, - { - 'params': [param_optimizer[1][1]], - 'weight_decay': 0.01 - }] - - model, optimizer, _, _ = deepspeed.initialize(config=config_dict, - model=model, - model_parameters=optimizer_grouped_parameters) + optimizer_grouped_parameters = [ + { + "params": [param_optimizer[0][1]], + "weight_decay": 0.01, + "exp_avg_mask": mask1, + }, + { + "params": [param_optimizer[1][1]], + "weight_decay": 0.01 + }, + ] + + model, optimizer, _, _ = deepspeed.initialize( + config=config_dict, + model=model, + model_parameters=optimizer_grouped_parameters, + ) data_loader = random_dataloader(model=model, total_samples=50, hidden_dim=hidden_dim, @@ -809,8 +923,12 @@ def test(self): model.step() # Test whether the momentum mask works for v in optimizer.state.values(): - if v['exp_avg'].size() == mask1.size(): - assert torch.allclose(v['exp_avg'], v['exp_avg'].mul_(mask1.to(device=v['exp_avg'].device)), atol=1e-07), f"Momentum mask is not working properly" + if v["exp_avg"].size() == mask1.size(): + assert torch.allclose( + v["exp_avg"], + v["exp_avg"].mul_(mask1.to(device=v["exp_avg"].device)), + atol=1e-07, + ), f"Momentum mask is not working properly" class TestOneBitLambCheckpointing(DistributedTest): @@ -833,15 +951,15 @@ def test(self, tmpdir): "coeff_beta": 0.9, "factor_max": 1.0, "factor_min": 0.5, - "factor_threshold": 0.1 - } + "factor_threshold": 0.1, + }, }, "gradient_clipping": 1.0, "fp16": { "enabled": True, "loss_scale": 0, "initial_scale_power": 16 - } + }, } hidden_dim = 10 @@ -853,109 +971,142 @@ def test(self, tmpdir): mask1[0][col] += 1 mask2[1][col] += 1 - optimizer_grouped_parameters_1 = [{ - 'params': [param_optimizer[0][1]], - 'weight_decay': 0.01, - 'exp_avg_mask': mask1 - }, - { - 'params': [param_optimizer[1][1]], - 'weight_decay': 0.01 - }] - - optimizer_grouped_parameters_2 = [{ - 'params': [param_optimizer[0][1]], - 'weight_decay': 0.01, - 'exp_avg_mask': mask2 - }, - { - 'params': [param_optimizer[1][1]], - 'weight_decay': 0.01 - }] - - optimizer_grouped_parameters_3 = [{ - 'params': [param_optimizer[0][1]], - 'weight_decay': 0.01 - }, - { - 'params': [param_optimizer[1][1]], - 'weight_decay': 0.01 - }] - - model_1, optimizer_1, _, _ = deepspeed.initialize(config=config_dict, - model=model, - model_parameters=optimizer_grouped_parameters_1) - data_loader = random_dataloader(model=model_1, - total_samples=10, - hidden_dim=hidden_dim, - device=model_1.device) + optimizer_grouped_parameters_1 = [ + { + "params": [param_optimizer[0][1]], + "weight_decay": 0.01, + "exp_avg_mask": mask1, + }, + { + "params": [param_optimizer[1][1]], + "weight_decay": 0.01 + }, + ] + + optimizer_grouped_parameters_2 = [ + { + "params": [param_optimizer[0][1]], + "weight_decay": 0.01, + "exp_avg_mask": mask2, + }, + { + "params": [param_optimizer[1][1]], + "weight_decay": 0.01 + }, + ] + + optimizer_grouped_parameters_3 = [ + { + "params": [param_optimizer[0][1]], + "weight_decay": 0.01 + }, + { + "params": [param_optimizer[1][1]], + "weight_decay": 0.01 + }, + ] + + model_1, optimizer_1, _, _ = deepspeed.initialize( + config=config_dict, + model=model, + model_parameters=optimizer_grouped_parameters_1, + ) + data_loader = random_dataloader( + model=model_1, + total_samples=10, + hidden_dim=hidden_dim, + device=model_1.device, + ) for n, batch in enumerate(data_loader): loss = model_1(batch[0], batch[1]) model_1.backward(loss) model_1.step() # Test whether momentum mask still exist after saving checkpoint assert optimizer_1.optimizer.lamb_freeze_key is True - mask1 = mask1.to(device=optimizer_1.param_groups[0]['exp_avg_mask'].device) - assert torch.allclose(optimizer_1.param_groups[0]['exp_avg_mask'], mask1, atol=1e-07), f"Incorrect momentum mask" + mask1 = mask1.to(device=optimizer_1.param_groups[0]["exp_avg_mask"].device) + assert torch.allclose( + optimizer_1.param_groups[0]["exp_avg_mask"], mask1, atol=1e-07 + ), f"Incorrect momentum mask" scaling_coeff_1 = [] for v in optimizer_1.state.values(): - assert 'scaling_coeff' in v, f"Incorrect scaling_coeff" - scaling_coeff_1.append(v['scaling_coeff']) - save_folder = os.path.join(tmpdir, 'saved_checkpoint') + assert "scaling_coeff" in v, f"Incorrect scaling_coeff" + scaling_coeff_1.append(v["scaling_coeff"]) + save_folder = os.path.join(tmpdir, "saved_checkpoint") model_1.save_checkpoint(save_folder, tag=None) - assert torch.allclose(optimizer_1.param_groups[0]['exp_avg_mask'], mask1, atol=1e-07), f"Momentum mask should not change after saving checkpoint" - - - model_2, optimizer_2, _, _ = deepspeed.initialize(config=config_dict, - model=model, - model_parameters=optimizer_grouped_parameters_2) + assert torch.allclose( + optimizer_1.param_groups[0]["exp_avg_mask"], mask1, atol=1e-07 + ), f"Momentum mask should not change after saving checkpoint" + + model_2, optimizer_2, _, _ = deepspeed.initialize( + config=config_dict, + model=model, + model_parameters=optimizer_grouped_parameters_2, + ) # Test whether momentum mask stays the same after loading checkpoint - mask2 = mask2.to(device=optimizer_2.param_groups[0]['exp_avg_mask'].device) - assert torch.allclose(optimizer_2.param_groups[0]['exp_avg_mask'], mask2, atol=1e-07), f"Incorrect momentum mask" - model_2.load_checkpoint(save_folder, - tag=None, - load_optimizer_states=True, - load_lr_scheduler_states=True) - assert torch.allclose(optimizer_2.param_groups[0]['exp_avg_mask'], mask2, atol=1e-07), f"Momentum mask should not change after loading checkpoint" + mask2 = mask2.to(device=optimizer_2.param_groups[0]["exp_avg_mask"].device) + assert torch.allclose( + optimizer_2.param_groups[0]["exp_avg_mask"], mask2, atol=1e-07 + ), f"Incorrect momentum mask" + model_2.load_checkpoint( + save_folder, + tag=None, + load_optimizer_states=True, + load_lr_scheduler_states=True, + ) + assert torch.allclose( + optimizer_2.param_groups[0]["exp_avg_mask"], mask2, atol=1e-07 + ), f"Momentum mask should not change after loading checkpoint" # Test whether worker&server error is reset assert len(optimizer_2.optimizer.worker_errors) == 0, f"Incorrect worker error" assert len(optimizer_2.optimizer.server_errors) == 0, f"Incorrect server error" # Test whether scaling_coeffs is loaded correctly scaling_coeff_2 = [] for v in optimizer_2.state.values(): - assert 'scaling_coeff' in v, f"Incorrect scaling_coeff" - scaling_coeff_2.append(v['scaling_coeff']) - assert list(sorted(scaling_coeff_2)) == list(sorted(scaling_coeff_1)), f"Incorrect scaling_coeffs" + assert "scaling_coeff" in v, f"Incorrect scaling_coeff" + scaling_coeff_2.append(v["scaling_coeff"]) + assert list(sorted(scaling_coeff_2)) == list( + sorted(scaling_coeff_1) + ), f"Incorrect scaling_coeffs" assert optimizer_2.optimizer.lamb_freeze_key is True - model_3, optimizer_3, _, _ = deepspeed.initialize(config=config_dict, - model=model, - model_parameters=optimizer_grouped_parameters_3) + model_3, optimizer_3, _, _ = deepspeed.initialize( + config=config_dict, + model=model, + model_parameters=optimizer_grouped_parameters_3, + ) optimizer_3.optimizer.freeze_step = 20 - data_loader = random_dataloader(model=model_3, - total_samples=50, - hidden_dim=hidden_dim, - device=model_3.device) + data_loader = random_dataloader( + model=model_3, + total_samples=50, + hidden_dim=hidden_dim, + device=model_3.device, + ) for n, batch in enumerate(data_loader): loss = model_3(batch[0], batch[1]) model_3.backward(loss) model_3.step() assert optimizer_3.optimizer.lamb_freeze_key is True # Test whether momentum mask stays the same after loading checkpoint - assert 'exp_avg_mask' not in optimizer_3.param_groups[0], f"Incorrect momentum mask" - model_3.load_checkpoint(save_folder, - tag=None, - load_optimizer_states=True, - load_lr_scheduler_states=True) - assert 'exp_avg_mask' not in optimizer_3.param_groups[0], f"Momentum mask should not change after loading checkpoint" + assert ( + "exp_avg_mask" not in optimizer_3.param_groups[0] + ), f"Incorrect momentum mask" + model_3.load_checkpoint( + save_folder, + tag=None, + load_optimizer_states=True, + load_lr_scheduler_states=True, + ) + assert ( + "exp_avg_mask" not in optimizer_3.param_groups[0] + ), f"Momentum mask should not change after loading checkpoint" # Test whether worker&server error is reset assert len(optimizer_3.optimizer.worker_errors) == 0, f"Incorrect worker error" assert len(optimizer_3.optimizer.server_errors) == 0, f"Incorrect server error" # Test whether scaling_coeffs, lamb_coeff_freeze, last_factor are reset for v in optimizer_3.state.values(): - assert v['lamb_coeff_freeze'] == 0.0, f"Incorrect lamb_coeff_freeze" - assert v['last_factor'] == 1.0, f"Incorrect last_factor" - assert 'scaling_coeff' not in v, f"Incorrect scaling_coeff" + assert v["lamb_coeff_freeze"] == 0.0, f"Incorrect lamb_coeff_freeze" + assert v["last_factor"] == 1.0, f"Incorrect last_factor" + assert "scaling_coeff" not in v, f"Incorrect scaling_coeff" assert optimizer_3.optimizer.lamb_freeze_key is False def test_overflow(self, tmpdir): @@ -975,27 +1126,27 @@ def test_overflow(self, tmpdir): "coeff_beta": 0.9, "factor_max": 1.0, "factor_min": 0.5, - "factor_threshold": 0.1 - } + "factor_threshold": 0.1, + }, }, "gradient_clipping": 1.0, "fp16": { "enabled": True, "loss_scale": 0, "initial_scale_power": 16 - } + }, } hidden_dim = 10 model = SimpleModel(hidden_dim) - model, _, _, _ = deepspeed.initialize(config=config_dict, - model=model, - model_parameters=model.parameters()) + model, _, _, _ = deepspeed.initialize( + config=config_dict, model=model, model_parameters=model.parameters() + ) data_loader = random_dataloader(model=model, total_samples=100, hidden_dim=hidden_dim, device=model.device) - save_folder = os.path.join(tmpdir, 'saved_checkpoint') + save_folder = os.path.join(tmpdir, "saved_checkpoint") for n, batch in enumerate(data_loader): loss = model(batch[0], batch[1]) if dist.get_rank() == 0 and n >= 10: @@ -1007,19 +1158,27 @@ def test_overflow(self, tmpdir): model.save_checkpoint(save_folder, tag=None) -@pytest.mark.parametrize('topo', - [ - PipeTopo(num_pp=1, - num_dp=4), - PipeTopo(num_pp=2, - num_dp=2), - PipeTopo(num_pp=4, - num_dp=1), - ]) +@pytest.mark.parametrize( + "topo_config", + [ + { + "num_pp": 1, + "num_dp": 4 + }, + { + "num_pp": 2, + "num_dp": 2 + }, + { + "num_pp": 4, + "num_dp": 1 + }, + ], +) class TestOneBitLambFP16Pipeline(DistributedTest): world_size = 4 - def test(self, topo): + def test(self, topo_config): config_dict = { "train_batch_size": 16, "train_micro_batch_size_per_gpu": 4, @@ -1034,8 +1193,8 @@ def test(self, topo): "weight_decay": 3e-7, "freeze_step": 200, "cuda_aware": False, - "comm_backend_name": "nccl" - } + "comm_backend_name": "nccl", + }, }, "gradient_clipping": 1.0, "zero_optimization": { @@ -1049,9 +1208,11 @@ def test(self, topo): "pipeline": { "seed_layers": True, "activation_checkpoint_interval": 1 - } + }, } + topo = PipeTopo(**topo_config) + # Allocate model for consistent initial weights. init_net = AlexNetPipe() @@ -1063,10 +1224,12 @@ def test(self, topo): topology=topo, loss_fn=nn.CrossEntropyLoss()) - test_losses = train_cifar(test_model, - config_dict=config_dict, - num_steps=steps, - fp16=config_dict['fp16']['enabled']) + test_losses = train_cifar( + test_model, + config_dict=config_dict, + num_steps=steps, + fp16=config_dict["fp16"]["enabled"], + ) @pytest.mark.sequential @@ -1075,6 +1238,7 @@ class TestCompressedAllReduceBasic(DistributedTest): def test(self, tmpdir): from deepspeed.runtime.comm.nccl import NcclBackend + size = dist.get_world_size() rank = dist.get_rank() backend = NcclBackend() @@ -1090,8 +1254,8 @@ def torch_sim(a): worker_error = a - a_compressed dist.all_reduce(a_compressed) a_compressed.mul_(1 / dist.get_world_size()) - a_server_sign = a_compressed.sign().add_(1).bool().float().add_(-0.5).mul_( - 2.0) + a_server_sign = ( + a_compressed.sign().add_(1).bool().float().add_(-0.5).mul_(2.0)) a_list = torch.chunk(a_compressed, chunks=dist.get_world_size()) server_scale = [ chunk_a.norm() / np.sqrt(chunk_a.numel()) for chunk_a in a_list @@ -1136,5 +1300,5 @@ def torch_sim(a): # The test would skip those numbers that are too small in compensated_server_m check_mag_mask = mpi_server[diff_server_mask] > magnitude_threshold if torch.sum(check_mag_mask) != 0: - print('Fails at {} of positions'.format(torch.sum(check_mag_mask))) + print("Fails at {} of positions".format(torch.sum(check_mag_mask))) assert torch.sum(diff_server_mask) == 0 or torch.sum(check_mag_mask) == 0 From c4cd57929174070e45085d3714ea60fcfb275eb9 Mon Sep 17 00:00:00 2001 From: Michael Wyatt Date: Thu, 4 Aug 2022 11:04:16 -0700 Subject: [PATCH 03/26] moving non-dist tests to new directory structure --- tests/unit/{test_pipe.py => alexnet_model.py} | 128 +----------------- .../{ => compression}/test_compression.py | 6 +- tests/unit/ops/adagrad/test_cpu_adagrad.py | 0 tests/unit/ops/adam/test_cpu_adam.py | 0 .../test_sparse_attention.py | 0 .../{ => runtime/fp16/onebit}/test_onebit.py | 24 ++-- tests/unit/runtime/pipe/test_pipe.py | 115 ++++++++++++++++ tests/unit/{ => runtime}/test_ds_config.py | 0 .../{ => runtime/zero}/test_zero_config.py | 0 .../{ => runtime/zero}/test_zero_tiled.py | 0 .../unit/{ => utils}/test_get_optim_files.py | 0 11 files changed, 132 insertions(+), 141 deletions(-) rename tests/unit/{test_pipe.py => alexnet_model.py} (54%) rename tests/unit/{ => compression}/test_compression.py (98%) mode change 100755 => 100644 mode change 100755 => 100644 tests/unit/ops/adagrad/test_cpu_adagrad.py mode change 100755 => 100644 tests/unit/ops/adam/test_cpu_adam.py rename tests/unit/{ => ops/sparse_attention}/test_sparse_attention.py (100%) mode change 100755 => 100644 rename tests/unit/{ => runtime/fp16/onebit}/test_onebit.py (99%) create mode 100644 tests/unit/runtime/pipe/test_pipe.py rename tests/unit/{ => runtime}/test_ds_config.py (100%) mode change 100755 => 100644 rename tests/unit/{ => runtime/zero}/test_zero_config.py (100%) mode change 100755 => 100644 rename tests/unit/{ => runtime/zero}/test_zero_tiled.py (100%) rename tests/unit/{ => utils}/test_get_optim_files.py (100%) diff --git a/tests/unit/test_pipe.py b/tests/unit/alexnet_model.py similarity index 54% rename from tests/unit/test_pipe.py rename to tests/unit/alexnet_model.py index 74c46cd0c536..7418a7f62ef4 100755 --- a/tests/unit/test_pipe.py +++ b/tests/unit/alexnet_model.py @@ -1,30 +1,11 @@ -import copy - import torch import torch.nn as nn import torch.nn.functional as F -import deepspeed.comm as dist - -import pytest - import deepspeed +import deepspeed.comm as dist import deepspeed.runtime.utils as ds_utils - -from deepspeed.runtime.pipe.topology import PipeDataParallelTopology - -PipeTopo = PipeDataParallelTopology from deepspeed.runtime.pipe.module import PipelineModule, LayerSpec -from .common import distributed_test - - -def rel_diff(A, B): - return abs(A - B) / abs(A) - - -# All models -from .simple_model import args_from_dict - class AlexNet(nn.Module): def __init__(self, num_classes=10): @@ -137,22 +118,22 @@ def cifar_trainset(fp16=False): def train_cifar(model, - config_dict, + config, num_steps=400, average_dp_losses=True, fp16=True, seed=123): - with torch.random.fork_rng(devices=[torch.cuda.current_device()]): + with torch.random.fork_rng(devices=[torch.cuda.current_device()], enabled=False): ds_utils.set_random_seed(seed) # disable dropout model.eval() trainset = cifar_trainset(fp16=fp16) - config_dict['local_rank'] = dist.get_rank() + config['local_rank'] = dist.get_rank() engine, _, _, _ = deepspeed.initialize( - config=config_dict, + config=config, model=model, model_parameters=[p for p in model.parameters()], training_data=trainset) @@ -171,102 +152,3 @@ def train_cifar(model, losses = loss_tensor.tolist() return losses - - -@pytest.mark.skip(reason="been seeing nondeterministic failures, skipping for now") -@pytest.mark.parametrize('topo', - [ - PipeTopo(num_pp=1, - num_dp=4), - PipeTopo(num_pp=2, - num_dp=2), - PipeTopo(num_pp=4, - num_dp=1), - ]) -def test_pipe_cifar10(topo, tmpdir): - config_dict = { - "train_batch_size": 16, - "train_micro_batch_size_per_gpu": 4, - "steps_per_print": 20, - "optimizer": { - "type": "Adam", - "params": { - "lr": 0.001, - "betas": [0.9, - 0.999], - "eps": 1e-8, - "weight_decay": 3e-7 - } - }, - "zero_optimization": { - "stage": 0 - }, - "fp16": { - "enabled": False - }, - "pipeline": { - "seed_layers": True, - "activation_checkpoint_interval": 1 - } - } - args = args_from_dict(tmpdir, config_dict) - - # Allocate model for consistent initial weights. - init_net = AlexNetPipe() - - @distributed_test(world_size=4) - def _helper(topo, tmpdir, steps=500): - assert steps >= 100 - - base_net = copy.deepcopy(init_net) - base_model = PipelineModule(layers=base_net.to_layers(), - num_stages=1, - loss_fn=nn.CrossEntropyLoss()) - - # Train with just data parallelism - base_losses = train_cifar(base_model, - args, - num_steps=steps, - fp16=config_dict['fp16']['enabled']) - - test_net = copy.deepcopy(init_net) - test_model = PipelineModule(layers=test_net.to_layers(), - topology=topo, - loss_fn=nn.CrossEntropyLoss()) - - #test_model = AlexNetPipe(num_classes=10, - # topology=test_topo, - # seed_layers=config_dict['pipeline']['seed_layers']) - test_losses = train_cifar(test_model, - args, - num_steps=steps, - fp16=config_dict['fp16']['enabled']) - - abs_diffs = [l0 - l1 for l0, l1 in zip(base_losses, test_losses)] - rel_diffs = [rel_diff(l0, l1) for l0, l1 in zip(base_losses, test_losses)] - if dist.get_rank() == 0: - print( - f'abs min={min(abs_diffs)} max={max(abs_diffs)} avg={sum(abs_diffs)/len(abs_diffs)}' - ) - print( - f'rel min={min(rel_diffs)} max={max(rel_diffs)} avg={sum(rel_diffs)/len(rel_diffs)}' - ) - print( - f'first: base={base_losses[0]} test={test_losses[0]} abs={abs_diffs[0]} rel={rel_diffs[0]}' - ) - - for lastX in [1, 10, 100]: - base_avg = sum(base_losses[-lastX:]) / lastX - test_avg = sum(test_losses[-lastX:]) / lastX - print( - f'last-{lastX}: base={base_avg} test={test_avg} abs={base_avg - test_avg} rel={rel_diff(base_avg, test_avg)}' - ) - - lastX = 100 - base = base_losses[-lastX:] - base_avg = sum(base) / len(base) - test = test_losses[-lastX:] - test_avg = sum(test) / len(test) - assert rel_diff(base_avg, test_avg) < 0.03 - - _helper(topo, tmpdir) diff --git a/tests/unit/test_compression.py b/tests/unit/compression/test_compression.py old mode 100755 new mode 100644 similarity index 98% rename from tests/unit/test_compression.py rename to tests/unit/compression/test_compression.py index d8d21bb630c0..3f5a72a4d30c --- a/tests/unit/test_compression.py +++ b/tests/unit/compression/test_compression.py @@ -2,10 +2,10 @@ import pytest import random import numpy as np -from .megatron_model import get_gpt2_model +from tests.unit.megatron_model import get_gpt2_model from deepspeed.compression.compress import init_compression -from .modeling import BertConfig -from .modelingpreln import BertEncoder as BertEncoderPreln +from tests.unit.modeling import BertConfig +from tests.unit.modelingpreln import BertEncoder as BertEncoderPreln from deepspeed.compression.basic_layer import LinearLayer_Compress, ColumnParallelLinear_Compress, RowParallelLinear_Compress from deepspeed.compression.helper import convert_conv1d_to_linear diff --git a/tests/unit/ops/adagrad/test_cpu_adagrad.py b/tests/unit/ops/adagrad/test_cpu_adagrad.py old mode 100755 new mode 100644 diff --git a/tests/unit/ops/adam/test_cpu_adam.py b/tests/unit/ops/adam/test_cpu_adam.py old mode 100755 new mode 100644 diff --git a/tests/unit/test_sparse_attention.py b/tests/unit/ops/sparse_attention/test_sparse_attention.py old mode 100755 new mode 100644 similarity index 100% rename from tests/unit/test_sparse_attention.py rename to tests/unit/ops/sparse_attention/test_sparse_attention.py diff --git a/tests/unit/test_onebit.py b/tests/unit/runtime/fp16/onebit/test_onebit.py similarity index 99% rename from tests/unit/test_onebit.py rename to tests/unit/runtime/fp16/onebit/test_onebit.py index f7b3c69bdded..096213478de1 100644 --- a/tests/unit/test_onebit.py +++ b/tests/unit/runtime/fp16/onebit/test_onebit.py @@ -11,9 +11,9 @@ PipeTopo = PipeDataParallelTopology from deepspeed.runtime.pipe.module import PipelineModule -from .common import DistributedTest -from .simple_model import SimpleModel, random_dataloader -from .test_pipe import AlexNetPipe, train_cifar +from tests.unit.common import DistributedTest +from tests.unit.simple_model import SimpleModel, random_dataloader +from tests.unit.alexnet_model import AlexNetPipe, train_cifar TORCH_MAJOR = int(torch.__version__.split(".")[0]) TORCH_MINOR = int(torch.__version__.split(".")[1]) @@ -391,13 +391,11 @@ def test(self, topo_config): } topo = PipeTopo(**topo_config) + steps = 500 # Must be >=100 # Allocate model for consistent initial weights. init_net = AlexNetPipe() - steps = 500 - assert steps >= 100 - test_net = copy.deepcopy(init_net) test_model = PipelineModule(layers=test_net.to_layers(), topology=topo, @@ -405,7 +403,7 @@ def test(self, topo_config): test_losses = train_cifar( test_model, - config_dict=config_dict, + config=config_dict, num_steps=steps, fp16=config_dict["fp16"]["enabled"], ) @@ -789,13 +787,11 @@ def test(self, topo_config): } topo = PipeTopo(**topo_config) + steps = 500 # Must be >=100 # Allocate model for consistent initial weights. init_net = AlexNetPipe() - steps = 500 - assert steps >= 100 - test_net = copy.deepcopy(init_net) test_model = PipelineModule(layers=test_net.to_layers(), topology=topo, @@ -803,7 +799,7 @@ def test(self, topo_config): test_losses = train_cifar( test_model, - config_dict=config_dict, + config=config_dict, num_steps=steps, fp16=config_dict["fp16"]["enabled"], ) @@ -1212,13 +1208,11 @@ def test(self, topo_config): } topo = PipeTopo(**topo_config) + steps = 500 # Must be >=100 # Allocate model for consistent initial weights. init_net = AlexNetPipe() - steps = 500 - assert steps >= 100 - test_net = copy.deepcopy(init_net) test_model = PipelineModule(layers=test_net.to_layers(), topology=topo, @@ -1226,7 +1220,7 @@ def test(self, topo_config): test_losses = train_cifar( test_model, - config_dict=config_dict, + config=config_dict, num_steps=steps, fp16=config_dict["fp16"]["enabled"], ) diff --git a/tests/unit/runtime/pipe/test_pipe.py b/tests/unit/runtime/pipe/test_pipe.py new file mode 100644 index 000000000000..ef4b79115027 --- /dev/null +++ b/tests/unit/runtime/pipe/test_pipe.py @@ -0,0 +1,115 @@ +import copy +import torch.nn as nn +import pytest + +import deepspeed.comm as dist +from deepspeed.runtime.pipe.topology import PipeDataParallelTopology +from deepspeed.runtime.pipe.module import PipelineModule +from tests.unit.alexnet_model import AlexNetPipe, train_cifar +from tests.unit.common import DistributedTest + +PipeTopo = PipeDataParallelTopology + + +def rel_diff(A, B): + return abs(A - B) / abs(A) + + +@pytest.mark.parametrize('topo_config', + [ + { + "num_pp": 1, + "num_dp": 4 + }, + { + "num_pp": 2, + "num_dp": 2 + }, + { + "num_pp": 4, + "num_dp": 1 + }, + ]) +class TestPipeCifar10(DistributedTest): + world_size = 4 + + def test(self, topo_config): + config_dict = { + "train_batch_size": 16, + "train_micro_batch_size_per_gpu": 4, + "steps_per_print": 20, + "optimizer": { + "type": "Adam", + "params": { + "lr": 0.001, + "betas": [0.9, + 0.999], + "eps": 1e-8, + "weight_decay": 3e-7 + } + }, + "zero_optimization": { + "stage": 0 + }, + "fp16": { + "enabled": False + }, + "pipeline": { + "seed_layers": True, + "activation_checkpoint_interval": 1 + } + } + + topo = PipeTopo(**topo_config) + steps = 500 # must be >=100 + + # Allocate model for consistent initial weights. + init_net = AlexNetPipe() + + base_net = copy.deepcopy(init_net) + base_model = PipelineModule(layers=base_net.to_layers(), + num_stages=1, + loss_fn=nn.CrossEntropyLoss()) + + # Train with just data parallelism + base_losses = train_cifar(base_model, + config=config_dict, + num_steps=steps, + fp16=config_dict['fp16']['enabled']) + + test_net = copy.deepcopy(init_net) + test_model = PipelineModule(layers=test_net.to_layers(), + topology=topo, + loss_fn=nn.CrossEntropyLoss()) + + test_losses = train_cifar(test_model, + config=config_dict, + num_steps=steps, + fp16=config_dict['fp16']['enabled']) + + abs_diffs = [l0 - l1 for l0, l1 in zip(base_losses, test_losses)] + rel_diffs = [rel_diff(l0, l1) for l0, l1 in zip(base_losses, test_losses)] + if dist.get_rank() == 0: + print( + f'abs min={min(abs_diffs)} max={max(abs_diffs)} avg={sum(abs_diffs)/len(abs_diffs)}' + ) + print( + f'rel min={min(rel_diffs)} max={max(rel_diffs)} avg={sum(rel_diffs)/len(rel_diffs)}' + ) + print( + f'first: base={base_losses[0]} test={test_losses[0]} abs={abs_diffs[0]} rel={rel_diffs[0]}' + ) + + for lastX in [1, 10, 100]: + base_avg = sum(base_losses[-lastX:]) / lastX + test_avg = sum(test_losses[-lastX:]) / lastX + print( + f'last-{lastX}: base={base_avg} test={test_avg} abs={base_avg - test_avg} rel={rel_diff(base_avg, test_avg)}' + ) + + lastX = 100 + base = base_losses[-lastX:] + base_avg = sum(base) / len(base) + test = test_losses[-lastX:] + test_avg = sum(test) / len(test) + assert rel_diff(base_avg, test_avg) < 0.03 diff --git a/tests/unit/test_ds_config.py b/tests/unit/runtime/test_ds_config.py old mode 100755 new mode 100644 similarity index 100% rename from tests/unit/test_ds_config.py rename to tests/unit/runtime/test_ds_config.py diff --git a/tests/unit/test_zero_config.py b/tests/unit/runtime/zero/test_zero_config.py old mode 100755 new mode 100644 similarity index 100% rename from tests/unit/test_zero_config.py rename to tests/unit/runtime/zero/test_zero_config.py diff --git a/tests/unit/test_zero_tiled.py b/tests/unit/runtime/zero/test_zero_tiled.py similarity index 100% rename from tests/unit/test_zero_tiled.py rename to tests/unit/runtime/zero/test_zero_tiled.py diff --git a/tests/unit/test_get_optim_files.py b/tests/unit/utils/test_get_optim_files.py similarity index 100% rename from tests/unit/test_get_optim_files.py rename to tests/unit/utils/test_get_optim_files.py From 17029cddf4cf755a9d976f2ee22707650650b487 Mon Sep 17 00:00:00 2001 From: Michael Wyatt Date: Thu, 4 Aug 2022 11:13:50 -0700 Subject: [PATCH 04/26] enable torch-latest CI workflow --- .github/workflows/amd.yml | 4 ++-- .github/workflows/nv-torch-latest-v100.yml | 4 ++-- 2 files changed, 4 insertions(+), 4 deletions(-) diff --git a/.github/workflows/amd.yml b/.github/workflows/amd.yml index bb2a72f8f09e..0c303824d458 100644 --- a/.github/workflows/amd.yml +++ b/.github/workflows/amd.yml @@ -66,5 +66,5 @@ jobs: run: | if [[ -d ./torch-extensions ]]; then rm -rf ./torch-extensions; fi cd tests - TORCH_EXTENSIONS_DIR=./torch-extensions pytest --color=yes --durations=0 --verbose -n 4 unit/{comm,inference,monitor,ops,profiling,runtime} - #TORCH_EXTENSIONS_DIR=./torch-extensions pytest --color=yes --durations=0 --verbose -m 'sequential' unit/{comm,inference,monitor,ops,profiling,runtime} + TORCH_EXTENSIONS_DIR=./torch-extensions pytest --color=yes --durations=0 --verbose -n 4 unit/{comm,compression,inference,monitor,ops,profiling,runtime,utils} + TORCH_EXTENSIONS_DIR=./torch-extensions pytest --color=yes --durations=0 --verbose -m 'sequential' unit/{comm,compression,inference,monitor,ops,profiling,runtime,utils} diff --git a/.github/workflows/nv-torch-latest-v100.yml b/.github/workflows/nv-torch-latest-v100.yml index 028b05b43748..74ee237688c4 100644 --- a/.github/workflows/nv-torch-latest-v100.yml +++ b/.github/workflows/nv-torch-latest-v100.yml @@ -61,5 +61,5 @@ jobs: if [[ -d ./torch-extensions ]]; then rm -rf ./torch-extensions; fi cd tests EXPECTED_TORCH=$(pip index versions torch | grep -oP -m1 "^\s*LATEST.*\s\K\d+\.\d+") - TORCH_EXTENSIONS_DIR=./torch-extensions pytest --color=yes --durations=0 --forked --verbose -n 4 unit/ --torch_ver=$EXPECTED_TROCH --cuda_ver="11.3" - TORCH_EXTENSIONS_DIR=./torch-extensions pytest --color=yes --durations=0 --forked --verbose -m 'sequential' unit/ --torch_ver=$EXPECTED_TORCH --cuda_ver="11.3" + TORCH_EXTENSIONS_DIR=./torch-extensions pytest --color=yes --durations=0 --forked --verbose -n 4 unit/{comm,compression,inference,monitor,ops,profiling,runtime,utils} --torch_ver=$EXPECTED_TROCH --cuda_ver="11.3" + TORCH_EXTENSIONS_DIR=./torch-extensions pytest --color=yes --durations=0 --forked --verbose -m 'sequential' unit/{comm,compression,inference,monitor,ops,profiling,runtime,utils} --torch_ver=$EXPECTED_TORCH --cuda_ver="11.3" From 2d46be0debffc6fc69c293aea54bba91e0734adb Mon Sep 17 00:00:00 2001 From: Michael Wyatt Date: Thu, 4 Aug 2022 11:26:23 -0700 Subject: [PATCH 05/26] hard-code expected torch version --- .github/workflows/nv-torch-latest-v100.yml | 5 ++--- 1 file changed, 2 insertions(+), 3 deletions(-) diff --git a/.github/workflows/nv-torch-latest-v100.yml b/.github/workflows/nv-torch-latest-v100.yml index 74ee237688c4..b41da667f3ee 100644 --- a/.github/workflows/nv-torch-latest-v100.yml +++ b/.github/workflows/nv-torch-latest-v100.yml @@ -60,6 +60,5 @@ jobs: unset TORCH_CUDA_ARCH_LIST # only jit compile for current arch if [[ -d ./torch-extensions ]]; then rm -rf ./torch-extensions; fi cd tests - EXPECTED_TORCH=$(pip index versions torch | grep -oP -m1 "^\s*LATEST.*\s\K\d+\.\d+") - TORCH_EXTENSIONS_DIR=./torch-extensions pytest --color=yes --durations=0 --forked --verbose -n 4 unit/{comm,compression,inference,monitor,ops,profiling,runtime,utils} --torch_ver=$EXPECTED_TROCH --cuda_ver="11.3" - TORCH_EXTENSIONS_DIR=./torch-extensions pytest --color=yes --durations=0 --forked --verbose -m 'sequential' unit/{comm,compression,inference,monitor,ops,profiling,runtime,utils} --torch_ver=$EXPECTED_TORCH --cuda_ver="11.3" + TORCH_EXTENSIONS_DIR=./torch-extensions pytest --color=yes --durations=0 --forked --verbose -n 4 unit/{comm,compression,inference,monitor,ops,profiling,runtime,utils} --torch_ver="1.12" --cuda_ver="11.3" + TORCH_EXTENSIONS_DIR=./torch-extensions pytest --color=yes --durations=0 --forked --verbose -m 'sequential' unit/{comm,compression,inference,monitor,ops,profiling,runtime,utils} --torch_ver="1.12" --cuda_ver="11.3" From 86e32d13f7a5b0101d3f7e00db6c74a1af58a352 Mon Sep 17 00:00:00 2001 From: Michael Wyatt Date: Thu, 4 Aug 2022 12:51:53 -0700 Subject: [PATCH 06/26] fix for unable to pickle lambda function --- tests/unit/alexnet_model.py | 7 ++++++- 1 file changed, 6 insertions(+), 1 deletion(-) diff --git a/tests/unit/alexnet_model.py b/tests/unit/alexnet_model.py index 7418a7f62ef4..c99c3cacf1bd 100755 --- a/tests/unit/alexnet_model.py +++ b/tests/unit/alexnet_model.py @@ -84,6 +84,11 @@ def __init__(self, num_classes=10, **kwargs): super().__init__(layers=specs, loss_fn=nn.CrossEntropyLoss(), **kwargs) +# Define this here because we cannot pickle local lambda functions +def cast_to_half(x): + return x.half() + + def cifar_trainset(fp16=False): import torchvision import torchvision.transforms as transforms @@ -98,7 +103,7 @@ def cifar_trainset(fp16=False): 0.5)), ] if fp16: - transform_list.append(torchvision.transforms.Lambda(lambda x: x.half())) + transform_list.append(torchvision.transforms.Lambda(cast_to_half)) transform = transforms.Compose(transform_list) From 871600b43294fa635ad9eb8d915e1c7aae037f45 Mon Sep 17 00:00:00 2001 From: Michael Wyatt Date: Thu, 4 Aug 2022 13:08:27 -0700 Subject: [PATCH 07/26] fix for PipeLine tests on AMD --- .github/workflows/amd.yml | 1 + 1 file changed, 1 insertion(+) diff --git a/.github/workflows/amd.yml b/.github/workflows/amd.yml index 0c303824d458..532be7889a41 100644 --- a/.github/workflows/amd.yml +++ b/.github/workflows/amd.yml @@ -38,6 +38,7 @@ jobs: pip install torch torchvision --extra-index-url https://download.pytorch.org/whl/rocm5.1.1 python -c "import torch; print('torch:', torch.__version__, torch)" python -c "import torch; print('CUDA available:', torch.cuda.is_available())" + HCC_AMDGPU_TARGET=gfx908 __HIP_PLATFORM_HCC__="" CUPY_INSTALL_USE_HIP=1 pip install -U cupy sudo apt-get update sudo apt-get install -y libaio-dev From bebbf950bafcb7eabe5f8ca782b004c2ff830a04 Mon Sep 17 00:00:00 2001 From: Michael Wyatt Date: Thu, 4 Aug 2022 13:11:31 -0700 Subject: [PATCH 08/26] fix for p40 runner --- .github/workflows/nv-torch12-p40.yml | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/.github/workflows/nv-torch12-p40.yml b/.github/workflows/nv-torch12-p40.yml index 944ba3beb19d..b726046e6f21 100644 --- a/.github/workflows/nv-torch12-p40.yml +++ b/.github/workflows/nv-torch12-p40.yml @@ -31,7 +31,7 @@ jobs: nvcc --version pip install --upgrade pip pip uninstall --yes torch torchvision - pip install torch==1.2.0 torchvision==0.4.0 + pip install torch==1.2.0 torchvision==0.5.0 python -c "import torch; print('torch:', torch.__version__, torch)" python -c "import torch; print('CUDA available:', torch.cuda.is_available())" From 0f44d31b07f86110408f39a37f33043ed7984f56 Mon Sep 17 00:00:00 2001 From: Michael Wyatt Date: Thu, 4 Aug 2022 13:12:09 -0700 Subject: [PATCH 09/26] revert previous change --- tests/unit/alexnet_model.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tests/unit/alexnet_model.py b/tests/unit/alexnet_model.py index c99c3cacf1bd..71273041040f 100755 --- a/tests/unit/alexnet_model.py +++ b/tests/unit/alexnet_model.py @@ -128,7 +128,7 @@ def train_cifar(model, average_dp_losses=True, fp16=True, seed=123): - with torch.random.fork_rng(devices=[torch.cuda.current_device()], enabled=False): + with torch.random.fork_rng(devices=[torch.cuda.current_device()]): ds_utils.set_random_seed(seed) # disable dropout From 84161ffccf98fc745f58ae42bbc7a6cee1dc1c42 Mon Sep 17 00:00:00 2001 From: Michael Wyatt Date: Thu, 4 Aug 2022 13:27:20 -0700 Subject: [PATCH 10/26] fix for cupy on amd --- .github/workflows/amd.yml | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/.github/workflows/amd.yml b/.github/workflows/amd.yml index 532be7889a41..5fe27ea15cc5 100644 --- a/.github/workflows/amd.yml +++ b/.github/workflows/amd.yml @@ -38,7 +38,7 @@ jobs: pip install torch torchvision --extra-index-url https://download.pytorch.org/whl/rocm5.1.1 python -c "import torch; print('torch:', torch.__version__, torch)" python -c "import torch; print('CUDA available:', torch.cuda.is_available())" - HCC_AMDGPU_TARGET=gfx908 __HIP_PLATFORM_HCC__="" CUPY_INSTALL_USE_HIP=1 pip install -U cupy + pip install https://github.com/cupy/cupy/releases/download/v11.0.0/cupy_rocm_5_0-11.0.0-cp37-cp37m-manylinux1_x86_64.whl sudo apt-get update sudo apt-get install -y libaio-dev From c14bcef25f833a5b4f2fc69712a0b820b081500c Mon Sep 17 00:00:00 2001 From: Michael Wyatt Date: Thu, 4 Aug 2022 13:39:34 -0700 Subject: [PATCH 11/26] skip added for p40 tests --- .github/workflows/nv-torch12-p40.yml | 2 +- tests/unit/alexnet_model.py | 5 +++++ 2 files changed, 6 insertions(+), 1 deletion(-) diff --git a/.github/workflows/nv-torch12-p40.yml b/.github/workflows/nv-torch12-p40.yml index b726046e6f21..944ba3beb19d 100644 --- a/.github/workflows/nv-torch12-p40.yml +++ b/.github/workflows/nv-torch12-p40.yml @@ -31,7 +31,7 @@ jobs: nvcc --version pip install --upgrade pip pip uninstall --yes torch torchvision - pip install torch==1.2.0 torchvision==0.5.0 + pip install torch==1.2.0 torchvision==0.4.0 python -c "import torch; print('torch:', torch.__version__, torch)" python -c "import torch; print('CUDA available:', torch.cuda.is_available())" diff --git a/tests/unit/alexnet_model.py b/tests/unit/alexnet_model.py index 71273041040f..7bac6fee03ae 100755 --- a/tests/unit/alexnet_model.py +++ b/tests/unit/alexnet_model.py @@ -1,3 +1,4 @@ +import pytest import torch import torch.nn as nn import torch.nn.functional as F @@ -93,6 +94,10 @@ def cifar_trainset(fp16=False): import torchvision import torchvision.transforms as transforms + tv_maj_ver, tv_min_ver = torchvision.__version__.split(".")[:2] + if not (tv_maj_ver > 0) and (tv_min_ver < 5): + pytest.skip("torchvision>=0.5.0 required") + transform_list = [ transforms.ToTensor(), transforms.Normalize((0.5, From 0b82169a4e44f5e18a0c535afcdad69c5cb0b6af Mon Sep 17 00:00:00 2001 From: Michael Wyatt Date: Thu, 4 Aug 2022 13:45:43 -0700 Subject: [PATCH 12/26] another fix... --- tests/unit/alexnet_model.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tests/unit/alexnet_model.py b/tests/unit/alexnet_model.py index 7bac6fee03ae..380821817663 100755 --- a/tests/unit/alexnet_model.py +++ b/tests/unit/alexnet_model.py @@ -94,7 +94,7 @@ def cifar_trainset(fp16=False): import torchvision import torchvision.transforms as transforms - tv_maj_ver, tv_min_ver = torchvision.__version__.split(".")[:2] + tv_maj_ver, tv_min_ver = map(int, torchvision.__version__.split(".")[:2]) if not (tv_maj_ver > 0) and (tv_min_ver < 5): pytest.skip("torchvision>=0.5.0 required") From a01d7ab8ca1041b23105daa9b326fce815e259c7 Mon Sep 17 00:00:00 2001 From: Michael Wyatt Date: Thu, 4 Aug 2022 14:35:08 -0700 Subject: [PATCH 13/26] moved more tests, some fixes --- .github/workflows/amd.yml | 5 +- .github/workflows/nv-torch-latest-v100.yml | 4 +- tests/unit/alexnet_model.py | 6 +- .../unit/{ => autotuning}/test_autotuning.py | 0 .../test_reshape_checkpoint.py | 0 tests/unit/{ => launcher}/test_run.py | 0 tests/unit/runtime/fp16/onebit/test_onebit.py | 11 +++- .../{ => runtime/pipe}/test_pipe_schedule.py | 39 +++++-------- .../{ => runtime/sparse_tensor}/test_csr.py | 0 tests/unit/{ => runtime}/test_autocast.py | 19 ++----- tests/unit/test_groups.py | 57 ------------------- tests/unit/utils/test_groups.py | 53 +++++++++++++++++ tests/unit/{ => utils}/test_init_on_device.py | 2 +- 13 files changed, 85 insertions(+), 111 deletions(-) rename tests/unit/{ => autotuning}/test_autotuning.py (100%) rename tests/unit/{ => checkpoint}/test_reshape_checkpoint.py (100%) rename tests/unit/{ => launcher}/test_run.py (100%) rename tests/unit/{ => runtime/pipe}/test_pipe_schedule.py (88%) rename tests/unit/{ => runtime/sparse_tensor}/test_csr.py (100%) rename tests/unit/{ => runtime}/test_autocast.py (80%) delete mode 100644 tests/unit/test_groups.py create mode 100644 tests/unit/utils/test_groups.py rename tests/unit/{ => utils}/test_init_on_device.py (92%) diff --git a/.github/workflows/amd.yml b/.github/workflows/amd.yml index 5fe27ea15cc5..67f2802fee36 100644 --- a/.github/workflows/amd.yml +++ b/.github/workflows/amd.yml @@ -38,7 +38,6 @@ jobs: pip install torch torchvision --extra-index-url https://download.pytorch.org/whl/rocm5.1.1 python -c "import torch; print('torch:', torch.__version__, torch)" python -c "import torch; print('CUDA available:', torch.cuda.is_available())" - pip install https://github.com/cupy/cupy/releases/download/v11.0.0/cupy_rocm_5_0-11.0.0-cp37-cp37m-manylinux1_x86_64.whl sudo apt-get update sudo apt-get install -y libaio-dev @@ -67,5 +66,5 @@ jobs: run: | if [[ -d ./torch-extensions ]]; then rm -rf ./torch-extensions; fi cd tests - TORCH_EXTENSIONS_DIR=./torch-extensions pytest --color=yes --durations=0 --verbose -n 4 unit/{comm,compression,inference,monitor,ops,profiling,runtime,utils} - TORCH_EXTENSIONS_DIR=./torch-extensions pytest --color=yes --durations=0 --verbose -m 'sequential' unit/{comm,compression,inference,monitor,ops,profiling,runtime,utils} + TORCH_EXTENSIONS_DIR=./torch-extensions pytest --color=yes --durations=0 --verbose -n 4 unit/{autotuning,checkpointing,comm,compression,inference,launcher,monitor,ops,profiling,runtime,utils} + TORCH_EXTENSIONS_DIR=./torch-extensions pytest --color=yes --durations=0 --verbose -m 'sequential' unit/{autotuning,checkpointing,comm,compression,inference,launcher,monitor,ops,profiling,runtime,utils} diff --git a/.github/workflows/nv-torch-latest-v100.yml b/.github/workflows/nv-torch-latest-v100.yml index b41da667f3ee..cced62edefc5 100644 --- a/.github/workflows/nv-torch-latest-v100.yml +++ b/.github/workflows/nv-torch-latest-v100.yml @@ -60,5 +60,5 @@ jobs: unset TORCH_CUDA_ARCH_LIST # only jit compile for current arch if [[ -d ./torch-extensions ]]; then rm -rf ./torch-extensions; fi cd tests - TORCH_EXTENSIONS_DIR=./torch-extensions pytest --color=yes --durations=0 --forked --verbose -n 4 unit/{comm,compression,inference,monitor,ops,profiling,runtime,utils} --torch_ver="1.12" --cuda_ver="11.3" - TORCH_EXTENSIONS_DIR=./torch-extensions pytest --color=yes --durations=0 --forked --verbose -m 'sequential' unit/{comm,compression,inference,monitor,ops,profiling,runtime,utils} --torch_ver="1.12" --cuda_ver="11.3" + TORCH_EXTENSIONS_DIR=./torch-extensions pytest --color=yes --durations=0 --forked --verbose -n 4 unit/{autotuning,checkpointing,comm,compression,inference,launcher,monitor,ops,profiling,runtime,utils} --torch_ver="1.12" --cuda_ver="11.3" + TORCH_EXTENSIONS_DIR=./torch-extensions pytest --color=yes --durations=0 --forked --verbose -m 'sequential' unit/{autotuning,checkpointing,comm,compression,inference,launcher,monitor,ops,profiling,runtime,utils} --torch_ver="1.12" --cuda_ver="11.3" diff --git a/tests/unit/alexnet_model.py b/tests/unit/alexnet_model.py index 380821817663..ab51a4630b7f 100755 --- a/tests/unit/alexnet_model.py +++ b/tests/unit/alexnet_model.py @@ -91,13 +91,9 @@ def cast_to_half(x): def cifar_trainset(fp16=False): - import torchvision + torchvision = pytest.importorskip("torchvision", minversion="0.5.0") import torchvision.transforms as transforms - tv_maj_ver, tv_min_ver = map(int, torchvision.__version__.split(".")[:2]) - if not (tv_maj_ver > 0) and (tv_min_ver < 5): - pytest.skip("torchvision>=0.5.0 required") - transform_list = [ transforms.ToTensor(), transforms.Normalize((0.5, diff --git a/tests/unit/test_autotuning.py b/tests/unit/autotuning/test_autotuning.py similarity index 100% rename from tests/unit/test_autotuning.py rename to tests/unit/autotuning/test_autotuning.py diff --git a/tests/unit/test_reshape_checkpoint.py b/tests/unit/checkpoint/test_reshape_checkpoint.py similarity index 100% rename from tests/unit/test_reshape_checkpoint.py rename to tests/unit/checkpoint/test_reshape_checkpoint.py diff --git a/tests/unit/test_run.py b/tests/unit/launcher/test_run.py similarity index 100% rename from tests/unit/test_run.py rename to tests/unit/launcher/test_run.py diff --git a/tests/unit/runtime/fp16/onebit/test_onebit.py b/tests/unit/runtime/fp16/onebit/test_onebit.py index 096213478de1..6c150117ac71 100644 --- a/tests/unit/runtime/fp16/onebit/test_onebit.py +++ b/tests/unit/runtime/fp16/onebit/test_onebit.py @@ -8,13 +8,14 @@ import numpy as np from deepspeed.runtime.pipe.topology import PipeDataParallelTopology - -PipeTopo = PipeDataParallelTopology +from deepspeed.ops.op_builder import OpBuilder from deepspeed.runtime.pipe.module import PipelineModule from tests.unit.common import DistributedTest from tests.unit.simple_model import SimpleModel, random_dataloader from tests.unit.alexnet_model import AlexNetPipe, train_cifar +PipeTopo = PipeDataParallelTopology + TORCH_MAJOR = int(torch.__version__.split(".")[0]) TORCH_MINOR = int(torch.__version__.split(".")[1]) if TORCH_MAJOR < 1 or TORCH_MINOR < 8: @@ -23,6 +24,12 @@ allow_module_level=True, ) +rocm_version = OpBuilder.installed_rocm_version() +if rocm_version[0] > 4: + pytest.skip( + "NCCL-based 1-bit compression is not yet supported w. ROCm 5 until cupy supports ROCm 5", + allow_module_level=True) + @pytest.mark.parametrize("dtype", [torch.float32, torch.float16], ids=["fp32", "fp16"]) class TestOneBitAdamBasic(DistributedTest): diff --git a/tests/unit/test_pipe_schedule.py b/tests/unit/runtime/pipe/test_pipe_schedule.py similarity index 88% rename from tests/unit/test_pipe_schedule.py rename to tests/unit/runtime/pipe/test_pipe_schedule.py index 8c65f3d7614b..1a02e003120f 100644 --- a/tests/unit/test_pipe_schedule.py +++ b/tests/unit/runtime/pipe/test_pipe_schedule.py @@ -1,5 +1,4 @@ import pytest - import deepspeed.runtime.pipe.schedule as schedule @@ -23,28 +22,26 @@ def test_pipe_train_schedule_singlestage(): sched = schedule.TrainSchedule(micro_batches=4, stages=1, stage_id=0) assert sched.num_micro_batches == 4 full = list(iter(sched)) - print() for idx, cmds in enumerate(full): - print(idx, cmds) - #assert len(cmds) == 2 - #assert type(cmds[0]) == schedule.LoadMicroBatch - #assert type(cmds[1]) == schedule.ForwardPass - #assert cmds[0].buffer_id == cmds[1].buffer_id - #assert len(full) == sched.num_micro_batches + if (idx % 2) != 0: + assert (len(cmds) == 1) or (len(cmds) == 4) + assert type(cmds[0]) == schedule.BackwardPass + else: + assert len(cmds) == 2 + assert type(cmds[0]) == schedule.LoadMicroBatch + assert type(cmds[1]) == schedule.ForwardPass + assert cmds[0].buffer_id == cmds[1].buffer_id + assert len(full) == sched.num_micro_batches * 2 @pytest.mark.parametrize('micro_batches', [1, 3, 8, 10]) -def test_pipe_inference_schedule_firststage(micro_batches, stages=3, verbose=False): +def test_pipe_inference_schedule_firststage(micro_batches, stages=3): sched = schedule.InferenceSchedule(micro_batches=micro_batches, stages=stages, stage_id=0) assert sched.num_micro_batches == micro_batches full = list(iter(sched)) - if verbose: - print() for idx, cmds in enumerate(full): - if verbose: - print(idx, cmds) # Ensure we don't send an activation the first step if idx == 0: assert len(cmds) == 2 @@ -73,17 +70,13 @@ def test_pipe_inference_schedule_firststage(micro_batches, stages=3, verbose=Fal @pytest.mark.parametrize('micro_batches', [1, 3, 8, 10]) -def test_pipe_inference_schedule_midstage(micro_batches, stages=3, verbose=False): +def test_pipe_inference_schedule_midstage(micro_batches, stages=3): sched = schedule.InferenceSchedule(micro_batches=micro_batches, stages=stages, stage_id=1) full = list(iter(sched)) - if verbose: - print() for idx, cmds in enumerate(full): - if verbose: - print(idx, cmds) if idx < sched.stage: assert len(cmds) == 0 continue @@ -103,16 +96,12 @@ def test_pipe_inference_schedule_midstage(micro_batches, stages=3, verbose=False @pytest.mark.parametrize('micro_batches', [1, 3, 8, 10]) -def test_pipe_inference_schedule_laststage(micro_batches, stages=3, verbose=False): +def test_pipe_inference_schedule_laststage(micro_batches, stages=3): sched = schedule.InferenceSchedule(micro_batches=micro_batches, stages=stages, stage_id=2) full = list(iter(sched)) - if verbose: - print() for idx, cmds in enumerate(full): - if verbose: - print(idx, cmds) if idx < sched.stage or idx > sched.stage + sched.num_micro_batches: assert len(cmds) == 0 continue @@ -135,10 +124,8 @@ def test_pipe_schedule_firststage(): def test_pipe_schedule_laststage(): sched = schedule.TrainSchedule(stages=3, micro_batches=4, stage_id=2) - #assert len(sched) == 2 * (sched.micro_batches + sched.stages - 1) - print() + assert len(list(iter(sched))) == 2 * (sched.micro_batches + sched.stages) for cmds in sched: - print(cmds) assert all(instr.__class__ != schedule.SendActivation for instr in cmds) assert all(instr.__class__ != schedule.RecvGrad for instr in cmds) diff --git a/tests/unit/test_csr.py b/tests/unit/runtime/sparse_tensor/test_csr.py similarity index 100% rename from tests/unit/test_csr.py rename to tests/unit/runtime/sparse_tensor/test_csr.py diff --git a/tests/unit/test_autocast.py b/tests/unit/runtime/test_autocast.py similarity index 80% rename from tests/unit/test_autocast.py rename to tests/unit/runtime/test_autocast.py index 7bffad14530d..f402486455ca 100644 --- a/tests/unit/test_autocast.py +++ b/tests/unit/runtime/test_autocast.py @@ -3,15 +3,6 @@ from deepspeed.runtime.zero.linear import LinearModuleForZeroStage3 -def _skip_autocast_test(): - try: - from torch.cuda.amp import custom_fwd, custom_bwd # noqa: F401 - except (ImportError, AttributeError) as exp: - return True - - return False - - @pytest.mark.parametrize('half_op', [False, True]) def test_missing_amp_autocast(tmpdir, half_op): hidden_dim = 4 @@ -28,8 +19,7 @@ def test_missing_amp_autocast(tmpdir, half_op): @pytest.mark.parametrize('half_op', [False, True]) def test_disable_autocast_linear(tmpdir, half_op): - if _skip_autocast_test(): - pytest.skip("amp autocast is not available") + amp = pytest.importorskip("torch.cuda.amp") hidden_dim = 4 if half_op: @@ -39,7 +29,7 @@ def test_disable_autocast_linear(tmpdir, half_op): input = torch.randn(hidden_dim).cuda() ds_linear = LinearModuleForZeroStage3(hidden_dim, hidden_dim).cuda() - with torch.cuda.amp.autocast(False): + with amp.autocast(False): output = ds_linear(input) assert output.dtype == ds_linear.weight.dtype @@ -54,8 +44,7 @@ def test_disable_autocast_linear(tmpdir, half_op): (True, True)]) def test_autocast_linear(tmpdir, half_input, half_weight): - if _skip_autocast_test(): - pytest.skip("amp autocast is not available") + amp = pytest.importorskip("torch.cuda.amp") hidden_dim = 4 input = torch.randn(hidden_dim).cuda() @@ -67,6 +56,6 @@ def test_autocast_linear(tmpdir, half_input, half_weight): if half_weight: ds_linear = ds_linear.half() - with torch.cuda.amp.autocast(): + with amp.autocast(): output = ds_linear(input) assert output.dtype == torch.half diff --git a/tests/unit/test_groups.py b/tests/unit/test_groups.py deleted file mode 100644 index 2769da74436d..000000000000 --- a/tests/unit/test_groups.py +++ /dev/null @@ -1,57 +0,0 @@ -import unittest - -from deepspeed.utils.groups import _get_expert_parallel_ranks - - -class TestGroups(unittest.TestCase): - def test_get_expert_parallel_ranks(self): - """ - Example - E + M + D parallel - world_size = 16 - model_degree = 2 - expert_degree = 4 # number of experts in same group - mp_group = [0, 1], [2,3], [4,5] ... - data_parallel_group =[0,2,4,6,8,10, 12,14], [1,3,5,7,9,11,13,15] - expert_parallel_group = [0,2,4,6], [8,10,12,14] [1,3,5,7], [9,11,13,15] - expert_data_parallel_group = [0,8],[2,10],[4,12],[6,14], [1,9],[3,11],[5,13],[7,15] - """ - expert_parallel_groups, expert_data_parallel_groups = _get_expert_parallel_ranks( - world_size=16, model_parallel_size_=2, expert_parallel_size_=4) - self.assertEqual(expert_parallel_groups, - [[0, - 2, - 4, - 6], - [8, - 10, - 12, - 14], - [1, - 3, - 5, - 7], - [9, - 11, - 13, - 15]]) - self.assertEqual(expert_data_parallel_groups, - [[0, - 8], - [2, - 10], - [4, - 12], - [6, - 14], - [1, - 9], - [3, - 11], - [5, - 13], - [7, - 15]]) - - -if __name__ == '__main__': - unittest.main() diff --git a/tests/unit/utils/test_groups.py b/tests/unit/utils/test_groups.py new file mode 100644 index 000000000000..b2f33cf436d3 --- /dev/null +++ b/tests/unit/utils/test_groups.py @@ -0,0 +1,53 @@ +from deepspeed.utils.groups import _get_expert_parallel_ranks + + +def test_get_expert_parallel_ranks(): + """ + Example - E + M + D parallel + world_size = 16 + model_degree = 2 + expert_degree = 4 # number of experts in same group + mp_group = [0, 1], [2,3], [4,5] ... + data_parallel_group =[0,2,4,6,8,10, 12,14], [1,3,5,7,9,11,13,15] + expert_parallel_group = [0,2,4,6], [8,10,12,14] [1,3,5,7], [9,11,13,15] + expert_data_parallel_group = [0,8],[2,10],[4,12],[6,14], [1,9],[3,11],[5,13],[7,15] + """ + expert_parallel_groups, expert_data_parallel_groups = _get_expert_parallel_ranks( + world_size=16, model_parallel_size_=2, expert_parallel_size_=4 + ) + assert expert_parallel_groups == [ + [0, + 2, + 4, + 6], + [8, + 10, + 12, + 14], + [1, + 3, + 5, + 7], + [9, + 11, + 13, + 15], + ] + assert expert_data_parallel_groups == [ + [0, + 8], + [2, + 10], + [4, + 12], + [6, + 14], + [1, + 9], + [3, + 11], + [5, + 13], + [7, + 15], + ] diff --git a/tests/unit/test_init_on_device.py b/tests/unit/utils/test_init_on_device.py similarity index 92% rename from tests/unit/test_init_on_device.py rename to tests/unit/utils/test_init_on_device.py index 2306d315f9e4..8bbd70c7f7fc 100644 --- a/tests/unit/test_init_on_device.py +++ b/tests/unit/utils/test_init_on_device.py @@ -1,6 +1,6 @@ import torch import pytest -from .simple_model import SimpleModel +from tests.unit.simple_model import SimpleModel from deepspeed import OnDevice from packaging import version as pkg_version From af27b01f951bf210d5bb6a28173611c37048652e Mon Sep 17 00:00:00 2001 From: Michael Wyatt Date: Thu, 4 Aug 2022 14:42:03 -0700 Subject: [PATCH 14/26] fix broken import --- tests/unit/autotuning/test_autotuning.py | 2 +- .../sparse_attention/test_sparse_attention.py | 140 ------------------ 2 files changed, 1 insertion(+), 141 deletions(-) diff --git a/tests/unit/autotuning/test_autotuning.py b/tests/unit/autotuning/test_autotuning.py index 681c3108b15b..5ece541e1bb7 100644 --- a/tests/unit/autotuning/test_autotuning.py +++ b/tests/unit/autotuning/test_autotuning.py @@ -1,6 +1,6 @@ import os import pytest -from .simple_model import create_config_from_dict +from tests.unit.simple_model import create_config_from_dict from deepspeed.launcher import runner as dsrun from deepspeed.autotuning.autotuner import Autotuner from deepspeed.autotuning.scheduler import ResourceManager diff --git a/tests/unit/ops/sparse_attention/test_sparse_attention.py b/tests/unit/ops/sparse_attention/test_sparse_attention.py index 8ff843c93169..740dfacdd0de 100644 --- a/tests/unit/ops/sparse_attention/test_sparse_attention.py +++ b/tests/unit/ops/sparse_attention/test_sparse_attention.py @@ -13,146 +13,6 @@ allow_module_level=True) -def test_sparse_attention_module_availability(): - return True - try: - from deepspeed.ops import sparse_attention # noqa: F401 - except ImportError: - print("Sparse Attention Module is not installed!") - return False - return True - - -def test_matmul_module_availability(): - return True - try: - from deepspeed.ops.sparse_attention.matmul import MatMul # noqa: F401 - except ImportError: - print("Sparse MatMul Module is not installed!") - return False - return True - - -def test_softmax_module_availability(): - return True - try: - from deepspeed.ops.sparse_attention.softmax import Softmax # noqa: F401 - except ImportError: - print("Sparse Softmax Module is not installed!") - return False - return True - - -def test_sparsityconfig_module_availability(): - return True - try: - from deepspeed.ops.sparse_attention import SparsityConfig # noqa: F401 - except ImportError: - print("SparsityConfig Module is not installed!") - return False - return True - - -def test_densesparsityconfig_module_availability(): - return True - try: - from deepspeed.ops.sparse_attention import DenseSparsityConfig # noqa: F401 - except ImportError: - print("DenseSparsityConfig Module is not installed!") - return False - return True - - -def test_fixedsparsityconfig_module_availability(): - return True - try: - from deepspeed.ops.sparse_attention import FixedSparsityConfig # noqa: F401 - except ImportError: - print("FixedSparsityConfig Module is not installed!") - return False - return True - - -def test_variablesparsityconfig_module_availability(): - return True - try: - from deepspeed.ops.sparse_attention import VariableSparsityConfig # noqa: F401 - except ImportError: - print("VariableSparsityConfig Module is not installed!") - return False - return True - - -def test_bigbirdsparsityconfig_module_availability(): - return True - try: - from deepspeed.ops.sparse_attention import BigBirdSparsityConfig # noqa: F401 - except ImportError: - print("BigBirdSparsityConfig Module is not installed!") - return False - return True - - -def test_bslongformersparsityconfig_module_availability(): - return True - try: - from deepspeed.ops.sparse_attention import BSLongformerSparsityConfig # noqa: F401 - except ImportError: - print("BSLongformerSparsityConfig Module is not installed!") - return False - return True - - -def test_localwindowsparsityconfig_module_availability(): - return True - try: - from deepspeed.ops.sparse_attention import LocalSlidingWindowSparsityConfig # noqa: F401 - except ImportError: - print("LocalSlidingWindowSparsityConfig Module is not installed!") - return False - return True - - -def test_sparseselfattention_module_availability(): - return True - try: - from deepspeed.ops.sparse_attention import SparseSelfAttention # noqa: F401 - except ImportError: - print("SparseSelfAttention Module is not installed!") - return False - return True - - -def test_bertsparseselfattention_module_availability(): - return True - try: - from deepspeed.ops.sparse_attention import BertSparseSelfAttention # noqa: F401 - except ImportError: - print("BertSparseSelfAttention Module is not installed!") - return False - return True - - -def test_sparseattentionutils_availability(): - return True - try: - from deepspeed.ops.sparse_attention import SparseAttentionUtils # noqa: F401 - except ImportError: - print("SparseAttentionUtils Module is not installed!") - return False - return True - - -def test_cpp_utils_availability(): - return True - try: - from deepspeed.ops.sparse_attention import cpp_utils # noqa: F401 - except ImportError: - print("Sparse Attention cpp_utils Module is not installed!") - return False - return True - - def dense_to_sparse(w, mask, block): """Converts dense matrix with explicit zeros to sparse matrix """ From 43ead12a536a4209b1b20ddae376db779165a0f6 Mon Sep 17 00:00:00 2001 From: Michael Wyatt Date: Thu, 4 Aug 2022 15:05:19 -0700 Subject: [PATCH 15/26] broken path --- .github/workflows/amd.yml | 4 ++-- .github/workflows/nv-torch-latest-v100.yml | 4 ++-- tests/unit/runtime/pipe/test_pipe_schedule.py | 2 +- 3 files changed, 5 insertions(+), 5 deletions(-) diff --git a/.github/workflows/amd.yml b/.github/workflows/amd.yml index 67f2802fee36..3b6ba57fc455 100644 --- a/.github/workflows/amd.yml +++ b/.github/workflows/amd.yml @@ -66,5 +66,5 @@ jobs: run: | if [[ -d ./torch-extensions ]]; then rm -rf ./torch-extensions; fi cd tests - TORCH_EXTENSIONS_DIR=./torch-extensions pytest --color=yes --durations=0 --verbose -n 4 unit/{autotuning,checkpointing,comm,compression,inference,launcher,monitor,ops,profiling,runtime,utils} - TORCH_EXTENSIONS_DIR=./torch-extensions pytest --color=yes --durations=0 --verbose -m 'sequential' unit/{autotuning,checkpointing,comm,compression,inference,launcher,monitor,ops,profiling,runtime,utils} + TORCH_EXTENSIONS_DIR=./torch-extensions pytest --color=yes --durations=0 --verbose -n 4 unit/{autotuning,checkpoint,comm,compression,inference,launcher,monitor,ops,profiling,runtime,utils} + TORCH_EXTENSIONS_DIR=./torch-extensions pytest --color=yes --durations=0 --verbose -m 'sequential' unit/{autotuning,checkpoint,comm,compression,inference,launcher,monitor,ops,profiling,runtime,utils} diff --git a/.github/workflows/nv-torch-latest-v100.yml b/.github/workflows/nv-torch-latest-v100.yml index cced62edefc5..9cabd2e66560 100644 --- a/.github/workflows/nv-torch-latest-v100.yml +++ b/.github/workflows/nv-torch-latest-v100.yml @@ -60,5 +60,5 @@ jobs: unset TORCH_CUDA_ARCH_LIST # only jit compile for current arch if [[ -d ./torch-extensions ]]; then rm -rf ./torch-extensions; fi cd tests - TORCH_EXTENSIONS_DIR=./torch-extensions pytest --color=yes --durations=0 --forked --verbose -n 4 unit/{autotuning,checkpointing,comm,compression,inference,launcher,monitor,ops,profiling,runtime,utils} --torch_ver="1.12" --cuda_ver="11.3" - TORCH_EXTENSIONS_DIR=./torch-extensions pytest --color=yes --durations=0 --forked --verbose -m 'sequential' unit/{autotuning,checkpointing,comm,compression,inference,launcher,monitor,ops,profiling,runtime,utils} --torch_ver="1.12" --cuda_ver="11.3" + TORCH_EXTENSIONS_DIR=./torch-extensions pytest --color=yes --durations=0 --forked --verbose -n 4 unit/{autotuning,checkpoint,comm,compression,inference,launcher,monitor,ops,profiling,runtime,utils} --torch_ver="1.12" --cuda_ver="11.3" + TORCH_EXTENSIONS_DIR=./torch-extensions pytest --color=yes --durations=0 --forked --verbose -m 'sequential' unit/{autotuning,checkpoint,comm,compression,inference,launcher,monitor,ops,profiling,runtime,utils} --torch_ver="1.12" --cuda_ver="11.3" diff --git a/tests/unit/runtime/pipe/test_pipe_schedule.py b/tests/unit/runtime/pipe/test_pipe_schedule.py index 1a02e003120f..68d13c2d9ba6 100644 --- a/tests/unit/runtime/pipe/test_pipe_schedule.py +++ b/tests/unit/runtime/pipe/test_pipe_schedule.py @@ -124,7 +124,7 @@ def test_pipe_schedule_firststage(): def test_pipe_schedule_laststage(): sched = schedule.TrainSchedule(stages=3, micro_batches=4, stage_id=2) - assert len(list(iter(sched))) == 2 * (sched.micro_batches + sched.stages) + assert len(list(iter(sched))) == 2 * (sched.micro_batches + sched.stages - 1) for cmds in sched: assert all(instr.__class__ != schedule.SendActivation for instr in cmds) assert all(instr.__class__ != schedule.RecvGrad for instr in cmds) From 9bb657c71c2636363d3bd14051a47b90a38ba752 Mon Sep 17 00:00:00 2001 From: Michael Wyatt Date: Thu, 4 Aug 2022 16:17:12 -0700 Subject: [PATCH 16/26] refactoring more tests --- .github/workflows/amd.yml | 4 +- .github/workflows/nv-torch-latest-v100.yml | 4 +- tests/unit/{ => elasticity}/test_elastic.py | 191 +++++------ tests/unit/runtime/pipe/test_pipe.py | 4 +- tests/unit/runtime/test_bf16.py | 301 ++++++++++++++++ tests/unit/{ => runtime}/test_data.py | 36 +- tests/unit/{ => runtime}/test_pld.py | 114 +++---- .../unit/{ => runtime}/test_runtime_utils.py | 20 +- tests/unit/test_bf16.py | 320 ------------------ 9 files changed, 477 insertions(+), 517 deletions(-) rename tests/unit/{ => elasticity}/test_elastic.py (69%) create mode 100644 tests/unit/runtime/test_bf16.py rename tests/unit/{ => runtime}/test_data.py (69%) rename tests/unit/{ => runtime}/test_pld.py (50%) rename tests/unit/{ => runtime}/test_runtime_utils.py (86%) delete mode 100644 tests/unit/test_bf16.py diff --git a/.github/workflows/amd.yml b/.github/workflows/amd.yml index 3b6ba57fc455..379bb20af2be 100644 --- a/.github/workflows/amd.yml +++ b/.github/workflows/amd.yml @@ -66,5 +66,5 @@ jobs: run: | if [[ -d ./torch-extensions ]]; then rm -rf ./torch-extensions; fi cd tests - TORCH_EXTENSIONS_DIR=./torch-extensions pytest --color=yes --durations=0 --verbose -n 4 unit/{autotuning,checkpoint,comm,compression,inference,launcher,monitor,ops,profiling,runtime,utils} - TORCH_EXTENSIONS_DIR=./torch-extensions pytest --color=yes --durations=0 --verbose -m 'sequential' unit/{autotuning,checkpoint,comm,compression,inference,launcher,monitor,ops,profiling,runtime,utils} + TORCH_EXTENSIONS_DIR=./torch-extensions pytest --color=yes --durations=0 --verbose -n 4 unit/{autotuning,checkpoint,comm,compression,elasticity,inference,launcher,monitor,ops,profiling,runtime,utils} + TORCH_EXTENSIONS_DIR=./torch-extensions pytest --color=yes --durations=0 --verbose -m 'sequential' unit/{autotuning,checkpoint,comm,compression,elasticity,inference,launcher,monitor,ops,profiling,runtime,utils} diff --git a/.github/workflows/nv-torch-latest-v100.yml b/.github/workflows/nv-torch-latest-v100.yml index 9cabd2e66560..1dc535f9b327 100644 --- a/.github/workflows/nv-torch-latest-v100.yml +++ b/.github/workflows/nv-torch-latest-v100.yml @@ -60,5 +60,5 @@ jobs: unset TORCH_CUDA_ARCH_LIST # only jit compile for current arch if [[ -d ./torch-extensions ]]; then rm -rf ./torch-extensions; fi cd tests - TORCH_EXTENSIONS_DIR=./torch-extensions pytest --color=yes --durations=0 --forked --verbose -n 4 unit/{autotuning,checkpoint,comm,compression,inference,launcher,monitor,ops,profiling,runtime,utils} --torch_ver="1.12" --cuda_ver="11.3" - TORCH_EXTENSIONS_DIR=./torch-extensions pytest --color=yes --durations=0 --forked --verbose -m 'sequential' unit/{autotuning,checkpoint,comm,compression,inference,launcher,monitor,ops,profiling,runtime,utils} --torch_ver="1.12" --cuda_ver="11.3" + TORCH_EXTENSIONS_DIR=./torch-extensions pytest --color=yes --durations=0 --forked --verbose -n 4 unit/{autotuning,checkpoint,comm,compression,elasticity,inference,launcher,monitor,ops,profiling,runtime,utils} --torch_ver="1.12" --cuda_ver="11.3" + TORCH_EXTENSIONS_DIR=./torch-extensions pytest --color=yes --durations=0 --forked --verbose -m 'sequential' unit/{autotuning,checkpoint,comm,compression,elasticity,inference,launcher,monitor,ops,profiling,runtime,utils} --torch_ver="1.12" --cuda_ver="11.3" diff --git a/tests/unit/test_elastic.py b/tests/unit/elasticity/test_elastic.py similarity index 69% rename from tests/unit/test_elastic.py rename to tests/unit/elasticity/test_elastic.py index 4ed2c0dd0c95..13fcd288ede0 100644 --- a/tests/unit/test_elastic.py +++ b/tests/unit/elasticity/test_elastic.py @@ -1,9 +1,9 @@ import pytest import deepspeed -from .common import distributed_test +from tests.unit.common import DistributedTest from deepspeed.git_version_info import version as ds_version import os -from .simple_model import SimpleModel, args_from_dict +from tests.unit.simple_model import SimpleModel base_ds_config = { "elasticity": { @@ -185,123 +185,116 @@ def test_proper_mbsz(): assert mbsize == 3 -def test_non_elastic_batch_params(tmpdir): - config_dict = { - "train_batch_size": 2, - "steps_per_print": 1, - "optimizer": { - "type": "Lamb", - "params": { - "lr": 0.00015 +class TestNonElasticBatchParams(DistributedTest): + world_size = [1, 2] + + def test(self): + config_dict = { + "train_batch_size": 2, + "steps_per_print": 1, + "optimizer": { + "type": "Lamb", + "params": { + "lr": 0.00015 + } + }, + "gradient_clipping": 1.0, + "elasticity": { + "enabled": True, + "max_train_batch_size": 4, + "micro_batch_sizes": [1, + 2, + 3, + 4], + "min_gpus": 1, + "max_gpus": 4, + "min_time": 20, + "version": 0.1 } - }, - "gradient_clipping": 1.0, - "elasticity": { - "enabled": True, - "max_train_batch_size": 4, - "micro_batch_sizes": [1, - 2, - 3, - 4], - "min_gpus": 1, - "max_gpus": 4, - "min_time": 20, - "version": 0.1 } - } - args = args_from_dict(tmpdir, config_dict) - hidden_dim = 10 + hidden_dim = 10 - model = SimpleModel(hidden_dim, empty_grad=False) + model = SimpleModel(hidden_dim, empty_grad=False) - @distributed_test(world_size=[1, 2]) - def _test_elastic(args, model, hidden_dim): with pytest.raises(deepspeed.elasticity.config.ElasticityError): - model, _, _,_ = deepspeed.initialize(args=args, + model, _, _,_ = deepspeed.initialize(config=config_dict, model=model, model_parameters=model.parameters()) - _test_elastic(args=args, model=model, hidden_dim=hidden_dim) - -def test_non_elastic_batch_params_w_override(tmpdir): - config_dict = { - "train_batch_size": 2, - "steps_per_print": 1, - "optimizer": { - "type": "Lamb", - "params": { - "lr": 0.00015 +class TestNonElasticBatchParamsWithOverride(DistributedTest): + world_size = [1, 2] + + def test(self): + config_dict = { + "train_batch_size": 2, + "steps_per_print": 1, + "optimizer": { + "type": "Lamb", + "params": { + "lr": 0.00015 + } + }, + "gradient_clipping": 1.0, + "elasticity": { + "enabled": True, + "max_train_batch_size": 4, + "micro_batch_sizes": [1, + 2, + 3, + 4], + "min_gpus": 1, + "max_gpus": 4, + "min_time": 20, + "version": 0.1, + "ignore_non_elastic_batch_info": True } - }, - "gradient_clipping": 1.0, - "elasticity": { - "enabled": True, - "max_train_batch_size": 4, - "micro_batch_sizes": [1, - 2, - 3, - 4], - "min_gpus": 1, - "max_gpus": 4, - "min_time": 20, - "version": 0.1, - "ignore_non_elastic_batch_info": True } - } - args = args_from_dict(tmpdir, config_dict) - hidden_dim = 10 - - model = SimpleModel(hidden_dim, empty_grad=False) + hidden_dim = 10 - @distributed_test(world_size=[1, 2]) - def _test_elastic(args, model, hidden_dim): - model, _, _,_ = deepspeed.initialize(args=args, + model = SimpleModel(hidden_dim, empty_grad=False) + model, _, _,_ = deepspeed.initialize(config=config_dict, model=model, model_parameters=model.parameters()) - _test_elastic(args=args, model=model, hidden_dim=hidden_dim) - -def test_elastic_config_changed(tmpdir): - config_dict = { - "train_batch_size": 2, - "steps_per_print": 1, - "optimizer": { - "type": "Lamb", - "params": { - "lr": 0.00015 +class TestElasticConfigChanged(DistributedTest): + world_size = [1, 2] + + def test(self): + config_dict = { + "train_batch_size": 2, + "steps_per_print": 1, + "optimizer": { + "type": "Lamb", + "params": { + "lr": 0.00015 + } + }, + "gradient_clipping": 1.0, + "elasticity": { + "enabled": True, + "max_train_batch_size": 4, + "micro_batch_sizes": [1, + 2, + 3, + 4], + "min_gpus": 1, + "max_gpus": 4, + "min_time": 20, + "version": 0.1, + "ignore_non_elastic_batch_info": True } - }, - "gradient_clipping": 1.0, - "elasticity": { - "enabled": True, - "max_train_batch_size": 4, - "micro_batch_sizes": [1, - 2, - 3, - 4], - "min_gpus": 1, - "max_gpus": 4, - "min_time": 20, - "version": 0.1, - "ignore_non_elastic_batch_info": True } - } - import json, os - scheduler_elastic_config = config_dict.copy() - scheduler_elastic_config["elasticity"]["max_train_batch_size"] = 27 - os.environ['DEEPSPEED_ELASTICITY_CONFIG'] = json.dumps(scheduler_elastic_config) - args = args_from_dict(tmpdir, config_dict) - hidden_dim = 10 + import json, os + scheduler_elastic_config = config_dict.copy() + scheduler_elastic_config["elasticity"]["max_train_batch_size"] = 27 + os.environ['DEEPSPEED_ELASTICITY_CONFIG'] = json.dumps(scheduler_elastic_config) + hidden_dim = 10 - model = SimpleModel(hidden_dim, empty_grad=False) + model = SimpleModel(hidden_dim, empty_grad=False) - @distributed_test(world_size=[1, 2]) - def _test_elastic(args, model, hidden_dim): with pytest.raises(deepspeed.elasticity.config.ElasticityError): - model, _, _,_ = deepspeed.initialize(args=args, + model, _, _,_ = deepspeed.initialize(config=config_dict, model=model, model_parameters=model.parameters()) - - _test_elastic(args=args, model=model, hidden_dim=hidden_dim) diff --git a/tests/unit/runtime/pipe/test_pipe.py b/tests/unit/runtime/pipe/test_pipe.py index ef4b79115027..7c097b1522a5 100644 --- a/tests/unit/runtime/pipe/test_pipe.py +++ b/tests/unit/runtime/pipe/test_pipe.py @@ -112,4 +112,6 @@ def test(self, topo_config): base_avg = sum(base) / len(base) test = test_losses[-lastX:] test_avg = sum(test) / len(test) - assert rel_diff(base_avg, test_avg) < 0.03 + assert rel_diff( + base_avg, + test_avg) < 0.05 # Originally 0.03, but seeing instability with AMD results diff --git a/tests/unit/runtime/test_bf16.py b/tests/unit/runtime/test_bf16.py new file mode 100644 index 000000000000..db3c0c4ab1c8 --- /dev/null +++ b/tests/unit/runtime/test_bf16.py @@ -0,0 +1,301 @@ +import torch +import deepspeed +import pytest +from deepspeed.ops.adam import FusedAdam +from tests.unit.common import DistributedTest +from deepspeed.ops.op_builder import CPUAdamBuilder +from testing.unit.simple_model import SimpleModel, SimpleOptimizer, random_dataloader +from tests.unit.util import bf16_required_version_check + + +@pytest.mark.parametrize('zero_stage, use_cpu_offload', [(2, False)]) +class TestAdamBF16ZeroOneCycleCompatibility(DistributedTest): + world_size = 1 + + def test(self, zero_stage, use_cpu_offload): + 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_batch_size": 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() + + +@pytest.mark.parametrize('zero_stage, use_cpu_offload', [(2, False)]) +class TestZeroAllowUntestedOptimizer(DistributedTest): + world_size = 1 + + def test(self, zero_stage, use_cpu_offload): + 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_batch_size": 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()) + + +@pytest.mark.parametrize('zero_stage, use_cpu_offload', [(2, False)]) +class TestZeroEmptyPartition(DistributedTest): + world_size = 3 + + def test(self, zero_stage, use_cpu_offload): + 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('zero_stage, optimizer_constructor', + [(2, + torch.optim.Adam), + (2, + FusedAdam)]) +class TestZeroSupportedClientOptimizer(DistributedTest): + world_size = 1 + + def test(self, zero_stage, optimizer_constructor): + 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_batch_size": 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_batch_size": 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() + + +@pytest.mark.parametrize('stage', [2]) +class TestZeroEmptyGrad(DistributedTest): + world_size = 1 + + def test(self, 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" + ) + + config_dict = { + "train_batch_size": 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() diff --git a/tests/unit/test_data.py b/tests/unit/runtime/test_data.py similarity index 69% rename from tests/unit/test_data.py rename to tests/unit/runtime/test_data.py index 93510e557450..c15d2f48eeef 100644 --- a/tests/unit/test_data.py +++ b/tests/unit/runtime/test_data.py @@ -2,8 +2,8 @@ import torch import pytest import deepspeed -from .common import distributed_test -from .simple_model import SimpleModel, args_from_dict, random_dataset +from tests.unit.common import DistributedTest +from tests.unit.simple_model import SimpleModel, random_dataset def test_repeating_loader(): @@ -25,26 +25,26 @@ def test_repeating_loader(): False), (4, False)]) -def test_dataloader_drop_last(tmpdir, train_batch_size, drop_last): - config_dict = { - "train_batch_size": train_batch_size, - "dataloader_drop_last": drop_last, - "steps_per_print": 1 - } - args = args_from_dict(tmpdir, config_dict) - hidden_dim = 10 - - model = SimpleModel(hidden_dim) - - @distributed_test(world_size=[1]) - def _test_dataloader_drop_last(args, model, hidden_dim): +class TestDataLoaderDropLast(DistributedTest): + world_size = 1 + + def test(self, train_batch_size, drop_last): + config_dict = { + "train_batch_size": train_batch_size, + "dataloader_drop_last": drop_last, + "steps_per_print": 1 + } + hidden_dim = 10 + + model = SimpleModel(hidden_dim) optimizer = torch.optim.AdamW(params=model.parameters()) - #TODO: Figure out why this breaks with cuda device + # TODO: no way to set DeepSpeedEngine.deepspeed_io params, need to use + # pin_memory=False for cuda device train_dataset = random_dataset(total_samples=50, hidden_dim=hidden_dim, device=torch.device('cpu'), dtype=torch.float32) - model, _, training_dataloader, _ = deepspeed.initialize(args=args, + model, _, training_dataloader, _ = deepspeed.initialize(config=config_dict, model=model, training_data=train_dataset, optimizer=optimizer) @@ -54,5 +54,3 @@ def _test_dataloader_drop_last(args, model, hidden_dim): loss = model(x, y) model.backward(loss) model.step() - - _test_dataloader_drop_last(args=args, model=model, hidden_dim=hidden_dim) diff --git a/tests/unit/test_pld.py b/tests/unit/runtime/test_pld.py similarity index 50% rename from tests/unit/test_pld.py rename to tests/unit/runtime/test_pld.py index 0953b648dce4..736fb3459438 100755 --- a/tests/unit/test_pld.py +++ b/tests/unit/runtime/test_pld.py @@ -3,8 +3,8 @@ import pytest from deepspeed.runtime.progressive_layer_drop import ProgressiveLayerDrop -from .common import distributed_test -from .simple_model import SimpleModel, PLD_SimpleModel, random_dataloader, args_from_dict +from tests.unit.common import DistributedTest +from tests.unit.simple_model import SimpleModel, PLD_SimpleModel, random_dataloader @pytest.mark.parametrize('theta', [0, 0.1, 0.9, 1.0]) @@ -20,35 +20,33 @@ def test_pld_schedule(tmpdir, theta): @pytest.mark.parametrize('theta', [0, 0.1, 0.9, 1.0]) -def test_pld_model(tmpdir, theta): - gamma = 0.001 - config_dict = { - "train_batch_size": 1, - "steps_per_print": 1, - "optimizer": { - "type": 'Adam', - "params": { - "lr": 0.0001 +class TestPLDModel(DistributedTest): + world_size = 1 + + def test_pld_model(self, theta): + gamma = 0.001 + config_dict = { + "train_batch_size": 1, + "steps_per_print": 1, + "optimizer": { + "type": 'Adam', + "params": { + "lr": 0.0001 + } + }, + "fp16": { + "enabled": True + }, + "progressive_layer_drop": { + "enabled": True, + "theta": theta, + "gamma": gamma } - }, - "fp16": { - "enabled": True - }, - "progressive_layer_drop": { - "enabled": True, - "theta": theta, - "gamma": gamma } - } - - args = args_from_dict(tmpdir, config_dict) - hidden_dim = 10 + hidden_dim = 10 - model = PLD_SimpleModel(hidden_dim, empty_grad=False) - - @distributed_test(world_size=[1]) - def _test_pld_model(args, model, hidden_dim, theta, gamma): - model, _, _, _ = deepspeed.initialize(args=args, + model = PLD_SimpleModel(hidden_dim, empty_grad=False) + model, _, _, _ = deepspeed.initialize(config=config_dict, model=model, model_parameters=model.parameters()) @@ -66,43 +64,35 @@ def _test_pld_model(args, model, hidden_dim, theta, gamma): actual_theta = model.get_pld_theta() assert expected_theta == actual_theta - _test_pld_model(args=args, - model=model, - hidden_dim=hidden_dim, - theta=theta, - gamma=gamma) - -def test_non_pld_model(tmpdir): - gamma = 0.001 - theta = 0.5 - config_dict = { - "train_batch_size": 1, - "steps_per_print": 1, - "optimizer": { - "type": 'Adam', - "params": { - "lr": 0.0001 +class TestNonPLDModel(DistributedTest): + world_size = 1 + + def test_non_pld_model(self): + gamma = 0.001 + theta = 0.5 + config_dict = { + "train_batch_size": 1, + "steps_per_print": 1, + "optimizer": { + "type": 'Adam', + "params": { + "lr": 0.0001 + } + }, + "fp16": { + "enabled": True + }, + "progressive_layer_drop": { + "enabled": True, + "theta": theta, + "gamma": gamma } - }, - "fp16": { - "enabled": True - }, - "progressive_layer_drop": { - "enabled": True, - "theta": theta, - "gamma": gamma } - } - - args = args_from_dict(tmpdir, config_dict) - hidden_dim = 10 + hidden_dim = 10 - model = SimpleModel(hidden_dim, empty_grad=False) - - @distributed_test(world_size=[1]) - def _test_non_pld_model(args, model, hidden_dim): - model, _, _, _ = deepspeed.initialize(args=args, + model = SimpleModel(hidden_dim, empty_grad=False) + model, _, _, _ = deepspeed.initialize(config=config_dict, model=model, model_parameters=model.parameters()) @@ -114,5 +104,3 @@ def _test_non_pld_model(args, model, hidden_dim): for i, batch in enumerate(data_loader): with pytest.raises(TypeError): loss = model(batch[0], batch[1]) - - _test_non_pld_model(args=args, model=model, hidden_dim=hidden_dim) diff --git a/tests/unit/test_runtime_utils.py b/tests/unit/runtime/test_runtime_utils.py similarity index 86% rename from tests/unit/test_runtime_utils.py rename to tests/unit/runtime/test_runtime_utils.py index 8e1697105000..e59ede26c8df 100644 --- a/tests/unit/test_runtime_utils.py +++ b/tests/unit/runtime/test_runtime_utils.py @@ -6,7 +6,7 @@ import deepspeed.runtime.utils as ds_utils import deepspeed.utils.groups as groups -from .common import distributed_test +from tests.unit.common import DistributedTest def test_call_to_str(): @@ -20,9 +20,10 @@ def test_call_to_str(): assert c2s('hello', 1138, val=3) == 'hello(1138, val=3)' -def test_clip_grad_norm_(): - @distributed_test(world_size=[2]) - def _test_clip_grad_norm_() -> None: +class TestClibGradNorm(DistributedTest): + world_size = 2 + + def test(self): param1 = torch.nn.Parameter(torch.Tensor([0])) param1.grad = torch.Tensor([1]) param2 = torch.nn.Parameter(torch.Tensor([0])) @@ -44,13 +45,12 @@ def _test_clip_grad_norm_() -> None: assert gathered_norm[0] == gathered_norm[1], "norm at rank 0 does not match the norm at rank 1" - return _test_clip_grad_norm_() - @pytest.mark.parametrize("check_using_norm", [(False), (True)]) -def test_CheckOverflow(check_using_norm): - @distributed_test(world_size=[2]) - def _test_CheckOverflow(check_using_norm: bool): +class TestCheckOverflow(DistributedTest): + world_size = 2 + + def test(self, check_using_norm): groups._create_expert_and_data_parallel(2) param1 = torch.nn.Parameter(torch.Tensor([0])) @@ -72,5 +72,3 @@ def _test_CheckOverflow(check_using_norm: bool): overflow_checker = ds_utils.CheckOverflow([parameters]) overflow = overflow_checker.check() assert overflow - - return _test_CheckOverflow(check_using_norm) diff --git a/tests/unit/test_bf16.py b/tests/unit/test_bf16.py deleted file mode 100644 index 4930a74640de..000000000000 --- a/tests/unit/test_bf16.py +++ /dev/null @@ -1,320 +0,0 @@ -import torch -import deepspeed -import pytest -from deepspeed.ops.adam import FusedAdam -from .common import distributed_test -from deepspeed.ops.op_builder import CPUAdamBuilder -from .simple_model import SimpleModel, SimpleOptimizer, random_dataloader, args_from_dict -from .util import bf16_required_version_check - - -@pytest.mark.parametrize('zero_stage, use_cpu_offload', [(2, False)]) -def test_adam_bf16_zero_onecycle_compatibility(tmpdir, zero_stage, use_cpu_offload): - 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_batch_size": 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 - } - } - - args = args_from_dict(tmpdir, config_dict) - hidden_dim = 10 - - @distributed_test(world_size=[1]) - def _test_adam_bf16_zero_onecycle_compatibility(args, zero_stage, hidden_dim): - model = SimpleModel(hidden_dim) - - model, _, _, _ = deepspeed.initialize(args=args, - 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() - - _test_adam_bf16_zero_onecycle_compatibility(args=args, - zero_stage=zero_stage, - hidden_dim=hidden_dim) - - -@pytest.mark.parametrize('zero_stage, use_cpu_offload', [(2, False)]) -def test_zero_allow_untested_optimizer(tmpdir, zero_stage, use_cpu_offload): - 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_batch_size": 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 - } - args = args_from_dict(tmpdir, config_dict) - - @distributed_test(world_size=[1]) - def _test_zero_allow_untested_optimizer(args, zero_stage): - hidden_dim = 10 - model = SimpleModel(hidden_dim) - optimizer = SimpleOptimizer(model.parameters()) - with pytest.raises(AssertionError): - model, optim, _, _ = deepspeed.initialize(args=args, - model=model, - optimizer=optimizer, - model_parameters=model.parameters()) - - _test_zero_allow_untested_optimizer(args, zero_stage) - - -@pytest.mark.parametrize('zero_stage, use_cpu_offload', [(2, False)]) -def test_zero_empty_partition(tmpdir, zero_stage, use_cpu_offload): - 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 - } - } - args = args_from_dict(tmpdir, config_dict) - - @distributed_test(world_size=[3]) - def _test_zero_empty_partition(args, zero_stage): - 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(args=args, - 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() - - _test_zero_empty_partition(args=args, zero_stage=zero_stage) - - -@pytest.mark.parametrize('zero_stage, optimizer_constructor', - [(2, - torch.optim.Adam), - (2, - FusedAdam)]) -def test_zero_supported_client_optimizer(tmpdir, zero_stage, optimizer_constructor): - 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_batch_size": 2, - "steps_per_print": 1, - "fp16": { - "enabled": False - }, - "bf16": { - "enabled": True - }, - "zero_optimization": { - "stage": zero_stage - } - } - args = args_from_dict(tmpdir, config_dict) - hidden_dim = 10 - - @distributed_test(world_size=[1]) - def _test_zero_supported_client_optimizer(args, zero_stage, optimizer_constructor): - model = SimpleModel(hidden_dim) - - client_optimizer = optimizer_constructor(params=model.parameters()) - model, _, _, _ = deepspeed.initialize(args=args, - model=model, - optimizer=client_optimizer) - - _test_zero_supported_client_optimizer(args=args, - zero_stage=zero_stage, - optimizer_constructor=optimizer_constructor) - - -def test_zero2_reduce_scatter_off(tmpdir): - 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_batch_size": 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 - } - } - args = args_from_dict(tmpdir, config_dict) - hidden_dim = 10 - - model = SimpleModel(hidden_dim) - - @distributed_test(world_size=[2]) - def _helper(args, model, hidden_dim): - model, _, _, _ = deepspeed.initialize(args=args, - 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() - - _helper(args=args, model=model, hidden_dim=hidden_dim) - - -@pytest.mark.parametrize('stage', [2]) -def test_zero_empty_grad(tmpdir, 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" - ) - - config_dict = { - "train_batch_size": 1, - "steps_per_print": 1, - "fp16": { - "enabled": False - }, - "bf16": { - "enabled": True - }, - "zero_optimization": { - "stage": stage - } - } - args = args_from_dict(tmpdir, config_dict) - hidden_dim = 10 - - model = SimpleModel(hidden_dim) - - @distributed_test(world_size=[1]) - def _go(args, model, hidden_dim): - optimizer = torch.optim.Adam(model.parameters()) - model, _, _, _ = deepspeed.initialize(args=args, - 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() - - _go(args=args, model=model, hidden_dim=hidden_dim) From 4a5910f664a2e49e0fc6a2161c08acf2be51ccf6 Mon Sep 17 00:00:00 2001 From: Michael Wyatt Date: Fri, 5 Aug 2022 09:17:27 -0700 Subject: [PATCH 17/26] fix for merge conflict --- tests/unit/runtime/test_bf16.py | 301 ------------------------------ tests/unit/test_bf16.py | 320 ++++++++++++++++++++++++++++++++ 2 files changed, 320 insertions(+), 301 deletions(-) delete mode 100644 tests/unit/runtime/test_bf16.py create mode 100644 tests/unit/test_bf16.py diff --git a/tests/unit/runtime/test_bf16.py b/tests/unit/runtime/test_bf16.py deleted file mode 100644 index db3c0c4ab1c8..000000000000 --- a/tests/unit/runtime/test_bf16.py +++ /dev/null @@ -1,301 +0,0 @@ -import torch -import deepspeed -import pytest -from deepspeed.ops.adam import FusedAdam -from tests.unit.common import DistributedTest -from deepspeed.ops.op_builder import CPUAdamBuilder -from testing.unit.simple_model import SimpleModel, SimpleOptimizer, random_dataloader -from tests.unit.util import bf16_required_version_check - - -@pytest.mark.parametrize('zero_stage, use_cpu_offload', [(2, False)]) -class TestAdamBF16ZeroOneCycleCompatibility(DistributedTest): - world_size = 1 - - def test(self, zero_stage, use_cpu_offload): - 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_batch_size": 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() - - -@pytest.mark.parametrize('zero_stage, use_cpu_offload', [(2, False)]) -class TestZeroAllowUntestedOptimizer(DistributedTest): - world_size = 1 - - def test(self, zero_stage, use_cpu_offload): - 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_batch_size": 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()) - - -@pytest.mark.parametrize('zero_stage, use_cpu_offload', [(2, False)]) -class TestZeroEmptyPartition(DistributedTest): - world_size = 3 - - def test(self, zero_stage, use_cpu_offload): - 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('zero_stage, optimizer_constructor', - [(2, - torch.optim.Adam), - (2, - FusedAdam)]) -class TestZeroSupportedClientOptimizer(DistributedTest): - world_size = 1 - - def test(self, zero_stage, optimizer_constructor): - 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_batch_size": 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_batch_size": 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() - - -@pytest.mark.parametrize('stage', [2]) -class TestZeroEmptyGrad(DistributedTest): - world_size = 1 - - def test(self, 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" - ) - - config_dict = { - "train_batch_size": 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() diff --git a/tests/unit/test_bf16.py b/tests/unit/test_bf16.py new file mode 100644 index 000000000000..4930a74640de --- /dev/null +++ b/tests/unit/test_bf16.py @@ -0,0 +1,320 @@ +import torch +import deepspeed +import pytest +from deepspeed.ops.adam import FusedAdam +from .common import distributed_test +from deepspeed.ops.op_builder import CPUAdamBuilder +from .simple_model import SimpleModel, SimpleOptimizer, random_dataloader, args_from_dict +from .util import bf16_required_version_check + + +@pytest.mark.parametrize('zero_stage, use_cpu_offload', [(2, False)]) +def test_adam_bf16_zero_onecycle_compatibility(tmpdir, zero_stage, use_cpu_offload): + 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_batch_size": 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 + } + } + + args = args_from_dict(tmpdir, config_dict) + hidden_dim = 10 + + @distributed_test(world_size=[1]) + def _test_adam_bf16_zero_onecycle_compatibility(args, zero_stage, hidden_dim): + model = SimpleModel(hidden_dim) + + model, _, _, _ = deepspeed.initialize(args=args, + 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() + + _test_adam_bf16_zero_onecycle_compatibility(args=args, + zero_stage=zero_stage, + hidden_dim=hidden_dim) + + +@pytest.mark.parametrize('zero_stage, use_cpu_offload', [(2, False)]) +def test_zero_allow_untested_optimizer(tmpdir, zero_stage, use_cpu_offload): + 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_batch_size": 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 + } + args = args_from_dict(tmpdir, config_dict) + + @distributed_test(world_size=[1]) + def _test_zero_allow_untested_optimizer(args, zero_stage): + hidden_dim = 10 + model = SimpleModel(hidden_dim) + optimizer = SimpleOptimizer(model.parameters()) + with pytest.raises(AssertionError): + model, optim, _, _ = deepspeed.initialize(args=args, + model=model, + optimizer=optimizer, + model_parameters=model.parameters()) + + _test_zero_allow_untested_optimizer(args, zero_stage) + + +@pytest.mark.parametrize('zero_stage, use_cpu_offload', [(2, False)]) +def test_zero_empty_partition(tmpdir, zero_stage, use_cpu_offload): + 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 + } + } + args = args_from_dict(tmpdir, config_dict) + + @distributed_test(world_size=[3]) + def _test_zero_empty_partition(args, zero_stage): + 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(args=args, + 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() + + _test_zero_empty_partition(args=args, zero_stage=zero_stage) + + +@pytest.mark.parametrize('zero_stage, optimizer_constructor', + [(2, + torch.optim.Adam), + (2, + FusedAdam)]) +def test_zero_supported_client_optimizer(tmpdir, zero_stage, optimizer_constructor): + 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_batch_size": 2, + "steps_per_print": 1, + "fp16": { + "enabled": False + }, + "bf16": { + "enabled": True + }, + "zero_optimization": { + "stage": zero_stage + } + } + args = args_from_dict(tmpdir, config_dict) + hidden_dim = 10 + + @distributed_test(world_size=[1]) + def _test_zero_supported_client_optimizer(args, zero_stage, optimizer_constructor): + model = SimpleModel(hidden_dim) + + client_optimizer = optimizer_constructor(params=model.parameters()) + model, _, _, _ = deepspeed.initialize(args=args, + model=model, + optimizer=client_optimizer) + + _test_zero_supported_client_optimizer(args=args, + zero_stage=zero_stage, + optimizer_constructor=optimizer_constructor) + + +def test_zero2_reduce_scatter_off(tmpdir): + 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_batch_size": 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 + } + } + args = args_from_dict(tmpdir, config_dict) + hidden_dim = 10 + + model = SimpleModel(hidden_dim) + + @distributed_test(world_size=[2]) + def _helper(args, model, hidden_dim): + model, _, _, _ = deepspeed.initialize(args=args, + 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() + + _helper(args=args, model=model, hidden_dim=hidden_dim) + + +@pytest.mark.parametrize('stage', [2]) +def test_zero_empty_grad(tmpdir, 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" + ) + + config_dict = { + "train_batch_size": 1, + "steps_per_print": 1, + "fp16": { + "enabled": False + }, + "bf16": { + "enabled": True + }, + "zero_optimization": { + "stage": stage + } + } + args = args_from_dict(tmpdir, config_dict) + hidden_dim = 10 + + model = SimpleModel(hidden_dim) + + @distributed_test(world_size=[1]) + def _go(args, model, hidden_dim): + optimizer = torch.optim.Adam(model.parameters()) + model, _, _, _ = deepspeed.initialize(args=args, + 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() + + _go(args=args, model=model, hidden_dim=hidden_dim) From 640e2bd30279f9eacb8750d87dcab50ed07cfb9c Mon Sep 17 00:00:00 2001 From: Michael Wyatt Date: Fri, 5 Aug 2022 09:42:18 -0700 Subject: [PATCH 18/26] more test refactors --- .../{ => runtime/pipe}/test_pipe_module.py | 21 +- tests/unit/runtime/test_bf16.py | 370 ++++++++++++++++ tests/unit/test_bf16.py | 396 ------------------ tests/unit/test_sparse_grads.py | 54 ++- 4 files changed, 405 insertions(+), 436 deletions(-) rename tests/unit/{ => runtime/pipe}/test_pipe_module.py (86%) create mode 100644 tests/unit/runtime/test_bf16.py delete mode 100644 tests/unit/test_bf16.py diff --git a/tests/unit/test_pipe_module.py b/tests/unit/runtime/pipe/test_pipe_module.py similarity index 86% rename from tests/unit/test_pipe_module.py rename to tests/unit/runtime/pipe/test_pipe_module.py index 1cba989b54e8..786a2234b158 100644 --- a/tests/unit/test_pipe_module.py +++ b/tests/unit/runtime/pipe/test_pipe_module.py @@ -15,8 +15,8 @@ from deepspeed.pipe import PipelineModule from deepspeed.utils import RepeatingLoader -from .common import distributed_test -from .simple_model import args_from_dict +from tests.unit.common import DistributedTest +from tests.unit.simple_model import args_from_dict HIDDEN_DIM = 32 LAYERS = 8 @@ -34,7 +34,7 @@ def sequential_model(): @pytest.fixture -def simple_args(tmpdir): +def simple_config(): config_dict = { "train_batch_size": 1, "train_micro_batch_size_per_gpu": 1, @@ -53,15 +53,14 @@ def simple_args(tmpdir): "activation_checkpoint_interval": 1 } } - args = args_from_dict(tmpdir, config_dict) - return args + return config_dict -def test_pipe_module_sequential(sequential_model, simple_args): - batch_input = torch.randn(1, HIDDEN_DIM) +class TestPipeModuleSequential(DistributedTest): + world_size = 4 - @distributed_test(world_size=4) - def _helper(): + def test(self, sequential_model, simple_config): + batch_input = torch.randn(1, HIDDEN_DIM) base_model = copy.deepcopy(sequential_model) base_input = batch_input.clone().detach() base_output = base_model(base_input) @@ -79,7 +78,7 @@ def _helper(): assert total_pipe_params == base_params pipe_model, _, _, _ = deepspeed.initialize( - args=simple_args, + config=simple_config, model=pipe_model, model_parameters=[p for p in pipe_model.parameters()]) @@ -98,5 +97,3 @@ def _helper(): pipe_output = pipe_output.to('cpu') assert torch.allclose(base_output, pipe_output, atol=1e-4) - - _helper() diff --git a/tests/unit/runtime/test_bf16.py b/tests/unit/runtime/test_bf16.py new file mode 100644 index 000000000000..6a19e8d3124a --- /dev/null +++ b/tests/unit/runtime/test_bf16.py @@ -0,0 +1,370 @@ +import torch +import deepspeed +import pytest +from deepspeed.ops.adam import FusedAdam +from tests.unit.common import DistributedTest +from deepspeed.ops.op_builder import CPUAdamBuilder +from tests.unit.simple_model import SimpleModel, SimpleOptimizer, random_dataloader +from tests.unit.util import bf16_required_version_check +from deepspeed import comm as dist + + +@pytest.mark.parametrize('zero_stage, use_cpu_offload', [(2, False)]) +class TestAdamBF16ZeroOneCycleCompatibility(DistributedTest): + world_size = 1 + + def test(self, zero_stage, use_cpu_offload): + 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_batch_size": 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() + + +@pytest.mark.parametrize('zero_stage, use_cpu_offload', [(2, False)]) +class TestZeroAllowUntestedOptimizer(DistributedTest): + world_size = 1 + + def test(self, zero_stage, use_cpu_offload): + 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_batch_size": 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()) + + +@pytest.mark.parametrize('zero_stage, use_cpu_offload', [(2, False)]) +class TestZeroEmptyPartition(DistributedTest): + world_size = 3 + + def test(self, zero_stage, use_cpu_offload): + 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('zero_stage, optimizer_constructor', + [(2, + torch.optim.Adam), + (2, + FusedAdam)]) +class TestZeroSupportedClientOptimizer(DistributedTest): + world_size = 1 + + def test(self, zero_stage, optimizer_constructor): + 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_batch_size": 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_batch_size": 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() + + +@pytest.mark.parametrize('stage', [2]) +class TestZeroEmptyGrad(DistributedTest): + world_size = 1 + + def test(self, 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" + ) + + config_dict = { + "train_batch_size": 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() + + +def custom_reduce(self, tensor, dst, op=dist.ReduceOp.SUM, group=None, async_op=False): + assert tensor.dtype == comm_torch_dtype + return orig_torch_reduce(tensor, dst, op, group, async_op) + +@pytest.mark.parametrize('comp_type_str, comm_type_str', + [("fp16", + "fp16"), + ("bfp16", + "bfp16"), + ("fp16", + "bfp16"), + ("bfp16", + "fp16"), + ("fp32", + "fp16"), + ("fp32", + "bfp16")]) +class TestZeroDtypeCocktail(DistributedTest): + world_size = 2 + + def test(self, comp_type_str, comm_type_str): + torch_dtype_dict = { + "fp16": torch.float16, + "bfp16": torch.bfloat16, + "fp32": torch.float + } + comp_torch_dtype = torch_dtype_dict[comp_type_str] + comm_torch_dtype = torch_dtype_dict[comm_type_str] + + if comp_torch_dtype == torch.bfloat16 or comm_torch_dtype == 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" + ) + + config_dict = { + "train_batch_size": 2, + "steps_per_print": 1, + "fp16": { + "enabled": torch_dtype_dict[comp_type_str] == torch.float16 + }, + "bf16": { + "enabled": torch_dtype_dict[comp_type_str] == torch.bfloat16 + }, + "zero_optimization": { + "stage": 2 + }, + "communication_data_type": comm_type_str + } + hidden_dim = 10 + + model = SimpleModel(hidden_dim) + orig_torch_reduce = dist.reduce + 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_dtype) + 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 diff --git a/tests/unit/test_bf16.py b/tests/unit/test_bf16.py deleted file mode 100644 index 2ff084933820..000000000000 --- a/tests/unit/test_bf16.py +++ /dev/null @@ -1,396 +0,0 @@ -import torch -import deepspeed -import pytest -from deepspeed.ops.adam import FusedAdam -from .common import distributed_test -from deepspeed.ops.op_builder import CPUAdamBuilder -from .simple_model import SimpleModel, SimpleOptimizer, random_dataloader, args_from_dict -from .util import bf16_required_version_check -from deepspeed import comm as dist - - -@pytest.mark.parametrize('zero_stage, use_cpu_offload', [(2, False)]) -def test_adam_bf16_zero_onecycle_compatibility(tmpdir, zero_stage, use_cpu_offload): - 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_batch_size": 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 - } - } - - args = args_from_dict(tmpdir, config_dict) - hidden_dim = 10 - - @distributed_test(world_size=[1]) - def _test_adam_bf16_zero_onecycle_compatibility(args, zero_stage, hidden_dim): - model = SimpleModel(hidden_dim) - - model, _, _, _ = deepspeed.initialize(args=args, - 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() - - _test_adam_bf16_zero_onecycle_compatibility(args=args, - zero_stage=zero_stage, - hidden_dim=hidden_dim) - - -@pytest.mark.parametrize('zero_stage, use_cpu_offload', [(2, False)]) -def test_zero_allow_untested_optimizer(tmpdir, zero_stage, use_cpu_offload): - 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_batch_size": 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 - } - args = args_from_dict(tmpdir, config_dict) - - @distributed_test(world_size=[1]) - def _test_zero_allow_untested_optimizer(args, zero_stage): - hidden_dim = 10 - model = SimpleModel(hidden_dim) - optimizer = SimpleOptimizer(model.parameters()) - with pytest.raises(AssertionError): - model, optim, _, _ = deepspeed.initialize(args=args, - model=model, - optimizer=optimizer, - model_parameters=model.parameters()) - - _test_zero_allow_untested_optimizer(args, zero_stage) - - -@pytest.mark.parametrize('zero_stage, use_cpu_offload', [(2, False)]) -def test_zero_empty_partition(tmpdir, zero_stage, use_cpu_offload): - 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 - } - } - args = args_from_dict(tmpdir, config_dict) - - @distributed_test(world_size=[3]) - def _test_zero_empty_partition(args, zero_stage): - 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(args=args, - 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() - - _test_zero_empty_partition(args=args, zero_stage=zero_stage) - - -@pytest.mark.parametrize('zero_stage, optimizer_constructor', - [(2, - torch.optim.Adam), - (2, - FusedAdam)]) -def test_zero_supported_client_optimizer(tmpdir, zero_stage, optimizer_constructor): - 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_batch_size": 2, - "steps_per_print": 1, - "fp16": { - "enabled": False - }, - "bf16": { - "enabled": True - }, - "zero_optimization": { - "stage": zero_stage - } - } - args = args_from_dict(tmpdir, config_dict) - hidden_dim = 10 - - @distributed_test(world_size=[1]) - def _test_zero_supported_client_optimizer(args, zero_stage, optimizer_constructor): - model = SimpleModel(hidden_dim) - - client_optimizer = optimizer_constructor(params=model.parameters()) - model, _, _, _ = deepspeed.initialize(args=args, - model=model, - optimizer=client_optimizer) - - _test_zero_supported_client_optimizer(args=args, - zero_stage=zero_stage, - optimizer_constructor=optimizer_constructor) - - -def test_zero2_reduce_scatter_off(tmpdir): - 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_batch_size": 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 - } - } - args = args_from_dict(tmpdir, config_dict) - hidden_dim = 10 - - model = SimpleModel(hidden_dim) - - @distributed_test(world_size=[2]) - def _helper(args, model, hidden_dim): - model, _, _, _ = deepspeed.initialize(args=args, - 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() - - _helper(args=args, model=model, hidden_dim=hidden_dim) - - -@pytest.mark.parametrize('stage', [2]) -def test_zero_empty_grad(tmpdir, 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" - ) - - config_dict = { - "train_batch_size": 1, - "steps_per_print": 1, - "fp16": { - "enabled": False - }, - "bf16": { - "enabled": True - }, - "zero_optimization": { - "stage": stage - } - } - args = args_from_dict(tmpdir, config_dict) - hidden_dim = 10 - - model = SimpleModel(hidden_dim) - - @distributed_test(world_size=[1]) - def _go(args, model, hidden_dim): - optimizer = torch.optim.Adam(model.parameters()) - model, _, _, _ = deepspeed.initialize(args=args, - 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() - - _go(args=args, model=model, hidden_dim=hidden_dim) - - -@pytest.mark.parametrize('comp_type_str, comm_type_str', - [("fp16", - "fp16"), - ("bfp16", - "bfp16"), - ("fp16", - "bfp16"), - ("bfp16", - "fp16"), - ("fp32", - "fp16"), - ("fp32", - "bfp16")]) -def test_zero_dtype_cocktail(tmpdir, comp_type_str, comm_type_str): - torch_dtype_dict = { - "fp16": torch.float16, - "bfp16": torch.bfloat16, - "fp32": torch.float - } - comp_torch_dtype = torch_dtype_dict[comp_type_str] - comm_torch_dtype = torch_dtype_dict[comm_type_str] - - if comp_torch_dtype == torch.bfloat16 or comm_torch_dtype == 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" - ) - - config_dict = { - "train_batch_size": 2, - "steps_per_print": 1, - "fp16": { - "enabled": torch_dtype_dict[comp_type_str] == torch.float16 - }, - "bf16": { - "enabled": torch_dtype_dict[comp_type_str] == torch.bfloat16 - }, - "zero_optimization": { - "stage": 2 - }, - "communication_data_type": comm_type_str - } - - args = args_from_dict(tmpdir, config_dict) - hidden_dim = 10 - - model = SimpleModel(hidden_dim) - - orig_torch_reduce = dist.reduce - - def custom_reduce(tensor, dst, op=dist.ReduceOp.SUM, group=None, async_op=False): - assert tensor.dtype == comm_torch_dtype - return orig_torch_reduce(tensor, dst, op, group, async_op) - - @distributed_test(world_size=[2]) - def _go(args, model, hidden_dim, comp_dtype): - optimizer = torch.optim.Adam(model.parameters()) - model, _, _, _ = deepspeed.initialize(args=args, - model=model, - optimizer=optimizer) - data_loader = random_dataloader(model=model, - total_samples=2, - hidden_dim=hidden_dim, - device=model.device, - dtype=comp_dtype) - 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 - - _go(args=args, model=model, hidden_dim=hidden_dim, comp_dtype=comp_torch_dtype) diff --git a/tests/unit/test_sparse_grads.py b/tests/unit/test_sparse_grads.py index 5be8ec3968fb..546f0f133b19 100644 --- a/tests/unit/test_sparse_grads.py +++ b/tests/unit/test_sparse_grads.py @@ -1,42 +1,42 @@ import torch import deepspeed -from .common import distributed_test +from tests.unit.common import DistributedTest import deepspeed.utils.groups as groups -def test_sparse_adam(tmpdir): - config_dict = {"train_batch_size": 2, "steps_per_print": 1, "sparse_gradients": True} +class Model(torch.nn.Module): + def __init__(self): + super().__init__() + self.emb = torch.nn.EmbeddingBag(10, 3, mode="sum", sparse=True) + self.linear = torch.nn.Linear(3, 1) - class Model(torch.nn.Module): - def __init__(self): - super().__init__() - self.emb = torch.nn.EmbeddingBag(10, 3, mode="sum", sparse=True) - self.linear = torch.nn.Linear(3, 1) + def forward(self, x, offsets): + return self.linear(self.emb(x, offsets)) - def forward(self, x, offsets): - return self.linear(self.emb(x, offsets)) +class Adam(torch.optim.Optimizer): + def __init__(self, dense_params, sparse_params): + super().__init__(dense_params + sparse_params, defaults={}) + self.adam = torch.optim.Adam(dense_params) + self.adam_sparse = torch.optim.SparseAdam(sparse_params) - class Adam(torch.optim.Optimizer): - def __init__(self, dense_params, sparse_params): - super().__init__(dense_params + sparse_params, defaults={}) - self.adam = torch.optim.Adam(dense_params) - self.adam_sparse = torch.optim.SparseAdam(sparse_params) + @torch.no_grad() + def step(self, closure=None): + loss_1 = self.adam.step(closure) + loss_2 = self.adam_sparse.step(closure) - @torch.no_grad() - def step(self, closure=None): - loss_1 = self.adam.step(closure) - loss_2 = self.adam_sparse.step(closure) + if loss_1 is not None and loss_2 is not None: + return loss_1 + loss_2 + return loss_1 or loss_2 - if loss_1 is not None and loss_2 is not None: - return loss_1 + loss_2 - return loss_1 or loss_2 +class TestSparseAdam(DistributedTest): + world_size = 2 - model = Model() - optimizer = Adam(list(model.linear.parameters()), list(model.emb.parameters())) + def test(self): + config_dict = {"train_batch_size": 2, "steps_per_print": 1, "sparse_gradients": True} - @distributed_test(world_size=[2]) - def _test(model, optimizer): + model = Model() + optimizer = Adam(list(model.linear.parameters()), list(model.emb.parameters())) engine, _, _, _ = deepspeed.initialize(model=model, optimizer=optimizer, config=config_dict) @@ -64,5 +64,3 @@ def _test(model, optimizer): ] for res in results: assert torch.allclose(res[0], res[1]) - - _test(model, optimizer) From 216d4a6ff9d91ad65327667818644f9c6d15e434 Mon Sep 17 00:00:00 2001 From: Michael Wyatt Date: Fri, 5 Aug 2022 09:50:53 -0700 Subject: [PATCH 19/26] formatting --- tests/unit/{runtime => }/pipe/test_pipe_module.py | 1 - tests/unit/runtime/test_bf16.py | 11 ++++++----- tests/unit/test_sparse_grads.py | 8 +++++++- 3 files changed, 13 insertions(+), 7 deletions(-) rename tests/unit/{runtime => }/pipe/test_pipe_module.py (98%) diff --git a/tests/unit/runtime/pipe/test_pipe_module.py b/tests/unit/pipe/test_pipe_module.py similarity index 98% rename from tests/unit/runtime/pipe/test_pipe_module.py rename to tests/unit/pipe/test_pipe_module.py index 786a2234b158..7a0b37acace6 100644 --- a/tests/unit/runtime/pipe/test_pipe_module.py +++ b/tests/unit/pipe/test_pipe_module.py @@ -16,7 +16,6 @@ from deepspeed.utils import RepeatingLoader from tests.unit.common import DistributedTest -from tests.unit.simple_model import args_from_dict HIDDEN_DIM = 32 LAYERS = 8 diff --git a/tests/unit/runtime/test_bf16.py b/tests/unit/runtime/test_bf16.py index 6a19e8d3124a..4d9f7aa7746c 100644 --- a/tests/unit/runtime/test_bf16.py +++ b/tests/unit/runtime/test_bf16.py @@ -300,10 +300,6 @@ def test(self, stage): model.step() -def custom_reduce(self, tensor, dst, op=dist.ReduceOp.SUM, group=None, async_op=False): - assert tensor.dtype == comm_torch_dtype - return orig_torch_reduce(tensor, dst, op, group, async_op) - @pytest.mark.parametrize('comp_type_str, comm_type_str', [("fp16", "fp16"), @@ -361,7 +357,12 @@ def test(self, comp_type_str, comm_type_str): total_samples=2, hidden_dim=hidden_dim, device=model.device, - dtype=comp_dtype) + dtype=comp_torch_dtype) + + def custom_reduce(tensor, dst, op=dist.ReduceOp.SUM, group=None, async_op=False): + assert tensor.dtype == comm_torch_dtype + return orig_torch_reduce(tensor, dst, op, group, async_op) + dist.reduce = custom_reduce for n, batch in enumerate(data_loader): loss = model(batch[0], batch[1]) diff --git a/tests/unit/test_sparse_grads.py b/tests/unit/test_sparse_grads.py index 546f0f133b19..94cc389db7b5 100644 --- a/tests/unit/test_sparse_grads.py +++ b/tests/unit/test_sparse_grads.py @@ -14,6 +14,7 @@ def __init__(self): def forward(self, x, offsets): return self.linear(self.emb(x, offsets)) + class Adam(torch.optim.Optimizer): def __init__(self, dense_params, sparse_params): super().__init__(dense_params + sparse_params, defaults={}) @@ -29,11 +30,16 @@ def step(self, closure=None): return loss_1 + loss_2 return loss_1 or loss_2 + class TestSparseAdam(DistributedTest): world_size = 2 def test(self): - config_dict = {"train_batch_size": 2, "steps_per_print": 1, "sparse_gradients": True} + config_dict = { + "train_batch_size": 2, + "steps_per_print": 1, + "sparse_gradients": True + } model = Model() optimizer = Adam(list(model.linear.parameters()), list(model.emb.parameters())) From 1ccae81947f646c69f535566868973e026ad5ab3 Mon Sep 17 00:00:00 2001 From: Michael Wyatt Date: Fri, 5 Aug 2022 13:12:00 -0700 Subject: [PATCH 20/26] test for AMD fix --- .github/workflows/amd.yml | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/.github/workflows/amd.yml b/.github/workflows/amd.yml index 379bb20af2be..7f7d3e2c0f8d 100644 --- a/.github/workflows/amd.yml +++ b/.github/workflows/amd.yml @@ -66,5 +66,5 @@ jobs: run: | if [[ -d ./torch-extensions ]]; then rm -rf ./torch-extensions; fi cd tests - TORCH_EXTENSIONS_DIR=./torch-extensions pytest --color=yes --durations=0 --verbose -n 4 unit/{autotuning,checkpoint,comm,compression,elasticity,inference,launcher,monitor,ops,profiling,runtime,utils} + TORCH_EXTENSIONS_DIR=./torch-extensions pytest --color=yes --durations=0 --verbose unit/{autotuning,checkpoint,comm,compression,elasticity,inference,launcher,monitor,ops,profiling,runtime,utils} TORCH_EXTENSIONS_DIR=./torch-extensions pytest --color=yes --durations=0 --verbose -m 'sequential' unit/{autotuning,checkpoint,comm,compression,elasticity,inference,launcher,monitor,ops,profiling,runtime,utils} From 5aea9672dacdca6509cdbed9f25c99d3db7e3f2c Mon Sep 17 00:00:00 2001 From: Michael Wyatt Date: Tue, 9 Aug 2022 14:00:56 -0700 Subject: [PATCH 21/26] working out CI failures --- tests/unit/elasticity/test_elastic.py | 70 ++++++++++++--------------- tests/unit/pipe/test_pipe_module.py | 5 +- 2 files changed, 34 insertions(+), 41 deletions(-) diff --git a/tests/unit/elasticity/test_elastic.py b/tests/unit/elasticity/test_elastic.py index 13fcd288ede0..e7c42de7d317 100644 --- a/tests/unit/elasticity/test_elastic.py +++ b/tests/unit/elasticity/test_elastic.py @@ -5,24 +5,27 @@ import os from tests.unit.simple_model import SimpleModel -base_ds_config = { - "elasticity": { - "enabled": True, - "max_train_batch_size": 10000, - "micro_batch_sizes": [8, - 12, - 16, - 17], - "min_gpus": 32, - "max_gpus": 1500, - "min_time": 20, - "version": 0.1 + +@pytest.fixture +def ds_config(): + config_dict = { + "elasticity": { + "enabled": True, + "max_train_batch_size": 10000, + "micro_batch_sizes": [8, + 12, + 16, + 17], + "min_gpus": 32, + "max_gpus": 1500, + "min_time": 20, + "version": 0.1 + } } -} + return config_dict -def test_basic_10k(): - ds_config = base_ds_config.copy() +def test_basic_10k(ds_config): final_batch_size, valid_gpus = deepspeed.elasticity.compute_elastic_config( ds_config=ds_config, target_deepspeed_version=ds_version) @@ -42,16 +45,14 @@ def test_basic_10k(): assert final_batch_size == 9792 -def test_old_version(): - ds_config = base_ds_config.copy() +def test_old_version(ds_config): with pytest.raises(deepspeed.elasticity.config.ElasticityError): final_batch_size, valid_gpus = deepspeed.elasticity.compute_elastic_config( ds_config=ds_config, target_deepspeed_version="0.2") -def test_disabled(): - ds_config = base_ds_config.copy() +def test_disabled(ds_config): ds_config['elasticity']['enabled'] = False with pytest.raises(deepspeed.elasticity.config.ElasticityError): final_batch_size, valid_gpus = deepspeed.elasticity.compute_elastic_config( @@ -59,8 +60,7 @@ def test_disabled(): target_deepspeed_version=ds_version) -def test_valid_world_size(): - ds_config = base_ds_config.copy() +def test_valid_world_size(ds_config): final_batch_size, valid_gpus, mbsize = deepspeed.elasticity.compute_elastic_config( ds_config=ds_config, target_deepspeed_version=ds_version, @@ -68,8 +68,7 @@ def test_valid_world_size(): assert mbsize == 17 -def test_invalid_world_size(): - ds_config = base_ds_config.copy() +def test_invalid_world_size(ds_config): with pytest.raises(deepspeed.elasticity.config.ElasticityIncompatibleWorldSize): final_batch_size, valid_gpus, mbsize = deepspeed.elasticity.compute_elastic_config( ds_config=ds_config, @@ -77,24 +76,21 @@ def test_invalid_world_size(): world_size=128) -def test_future_elastic_version(): - ds_config = base_ds_config.copy() +def test_future_elastic_version(ds_config): ds_config['elasticity']['version'] = '0.3' with pytest.raises(deepspeed.elasticity.config.ElasticityError): deepspeed.elasticity.compute_elastic_config(ds_config=ds_config, target_deepspeed_version=ds_version) -def test_missing_max_batch(): - ds_config = base_ds_config.copy() +def test_missing_max_batch(ds_config): del ds_config['elasticity']['max_train_batch_size'] with pytest.raises(deepspeed.elasticity.config.ElasticityError): deepspeed.elasticity.compute_elastic_config(ds_config=ds_config, target_deepspeed_version=ds_version) -def test_missing_micro_batch(): - ds_config = base_ds_config.copy() +def test_missing_micro_batch(ds_config): del ds_config['elasticity']['micro_batch_sizes'] with pytest.raises(deepspeed.elasticity.config.ElasticityError): deepspeed.elasticity.compute_elastic_config(ds_config=ds_config, @@ -108,8 +104,7 @@ def test_empty_config(): target_deepspeed_version=ds_version) -def test_model_parallel_v1_invalid(): - ds_config = base_ds_config.copy() +def test_model_parallel_v1_invalid(ds_config): ds_config["elasticity"]["model_parallel_size"] = 4 ds_config["elasticity"]["num_gpus_per_node"] = 8 ds_config["elasticity"]["version"] = 0.1 @@ -119,8 +114,7 @@ def test_model_parallel_v1_invalid(): target_deepspeed_version=ds_version) -def test_model_parallel_v2_invalid(): - ds_config = base_ds_config.copy() +def test_model_parallel_v2_invalid(ds_config): ds_config["elasticity"]["model_parallel_size"] = 16 ds_config["elasticity"]["num_gpus_per_node"] = 8 ds_config["elasticity"]["version"] = 0.2 @@ -131,9 +125,7 @@ def test_model_parallel_v2_invalid(): world_size=16) -def test_model_parallel_v2_valid(): - - ds_config = base_ds_config.copy() +def test_model_parallel_v2_valid(ds_config): ds_config["elasticity"]["model_parallel_size"] = 4 ds_config["elasticity"]["num_gpus_per_node"] = 8 ds_config["elasticity"]["version"] = 0.2 @@ -165,16 +157,14 @@ def test_model_parallel_v2_valid(): [2, 0.5, 4])]) -def test_invalid_config_values(key, value): - ds_config = base_ds_config.copy() +def test_invalid_config_values(key, value, ds_config): ds_config['elasticity'][key] = value with pytest.raises(deepspeed.elasticity.config.ElasticityError): deepspeed.elasticity.compute_elastic_config(ds_config=ds_config, target_deepspeed_version=ds_version) -def test_proper_mbsz(): - ds_config = base_ds_config.copy() +def test_proper_mbsz(ds_config): ds_config["elasticity"]["max_train_batch_size"] = 32 ds_config["elasticity"]["micro_batch_sizes"] = [1, 2, 3, 7] ds_config["elasticity"]["min_gpus"] = 1 diff --git a/tests/unit/pipe/test_pipe_module.py b/tests/unit/pipe/test_pipe_module.py index 7a0b37acace6..4d7341245a4d 100644 --- a/tests/unit/pipe/test_pipe_module.py +++ b/tests/unit/pipe/test_pipe_module.py @@ -95,4 +95,7 @@ def test(self, sequential_model, simple_config): base_output = base_output.to('cpu') pipe_output = pipe_output.to('cpu') - assert torch.allclose(base_output, pipe_output, atol=1e-4) + assert torch.allclose( + base_output, + pipe_output, + atol=1e-3) # TODO: original atol=1e-4, figure out why this started failing From 50d97ee1584f96954e0dd583c4392ea173564668 Mon Sep 17 00:00:00 2001 From: Michael Wyatt Date: Tue, 9 Aug 2022 14:32:03 -0700 Subject: [PATCH 22/26] Update amd.yml --- .github/workflows/amd.yml | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/.github/workflows/amd.yml b/.github/workflows/amd.yml index 7f7d3e2c0f8d..97b96f2d5875 100644 --- a/.github/workflows/amd.yml +++ b/.github/workflows/amd.yml @@ -67,4 +67,4 @@ jobs: if [[ -d ./torch-extensions ]]; then rm -rf ./torch-extensions; fi cd tests TORCH_EXTENSIONS_DIR=./torch-extensions pytest --color=yes --durations=0 --verbose unit/{autotuning,checkpoint,comm,compression,elasticity,inference,launcher,monitor,ops,profiling,runtime,utils} - TORCH_EXTENSIONS_DIR=./torch-extensions pytest --color=yes --durations=0 --verbose -m 'sequential' unit/{autotuning,checkpoint,comm,compression,elasticity,inference,launcher,monitor,ops,profiling,runtime,utils} + #TORCH_EXTENSIONS_DIR=./torch-extensions pytest --color=yes --durations=0 --verbose -m 'sequential' unit/{autotuning,checkpoint,comm,compression,elasticity,inference,launcher,monitor,ops,profiling,runtime,utils} From 093ed5d9a67db423285b3d88315ca5ccd4a6020d Mon Sep 17 00:00:00 2001 From: Michael Wyatt Date: Tue, 9 Aug 2022 16:02:10 -0700 Subject: [PATCH 23/26] move randn call outside distributed test --- tests/unit/pipe/test_pipe_module.py | 22 +++++++++------------- 1 file changed, 9 insertions(+), 13 deletions(-) diff --git a/tests/unit/pipe/test_pipe_module.py b/tests/unit/pipe/test_pipe_module.py index 4d7341245a4d..99862a96a2d3 100644 --- a/tests/unit/pipe/test_pipe_module.py +++ b/tests/unit/pipe/test_pipe_module.py @@ -7,11 +7,6 @@ import pytest import deepspeed - -from deepspeed.runtime.pipe.topology import PipeDataParallelTopology - -PipeTopo = PipeDataParallelTopology - from deepspeed.pipe import PipelineModule from deepspeed.utils import RepeatingLoader @@ -55,11 +50,15 @@ def simple_config(): return config_dict +@pytest.fixture +def batch_input(): + return torch.randn(1, HIDDEN_DIM) + + class TestPipeModuleSequential(DistributedTest): - world_size = 4 + world_size = 2 - def test(self, sequential_model, simple_config): - batch_input = torch.randn(1, HIDDEN_DIM) + def test(self, sequential_model, simple_config, batch_input): base_model = copy.deepcopy(sequential_model) base_input = batch_input.clone().detach() base_output = base_model(base_input) @@ -67,7 +66,7 @@ def test(self, sequential_model, simple_config): base_params = sum(p.numel() for p in base_model.parameters()) pipe_model = copy.deepcopy(sequential_model) - pipe_model = PipelineModule(layers=pipe_model, num_stages=4) + pipe_model = PipelineModule(layers=pipe_model, num_stages=2) # Ensure all parameters are accounted for. my_params = sum(p.numel() for p in pipe_model.parameters()) @@ -95,7 +94,4 @@ def test(self, sequential_model, simple_config): base_output = base_output.to('cpu') pipe_output = pipe_output.to('cpu') - assert torch.allclose( - base_output, - pipe_output, - atol=1e-3) # TODO: original atol=1e-4, figure out why this started failing + assert torch.allclose(base_output, pipe_output, atol=1e-4) From efbb9feff43b0f49fad33366e4e0e03ca927cadf Mon Sep 17 00:00:00 2001 From: Michael Wyatt Date: Tue, 9 Aug 2022 16:47:47 -0700 Subject: [PATCH 24/26] minor adjustments --- tests/unit/elasticity/test_elastic.py | 6 +- tests/unit/runtime/fp16/onebit/test_onebit.py | 8 +-- tests/unit/runtime/test_bf16.py | 72 ++++++++----------- tests/unit/runtime/test_pld.py | 0 4 files changed, 35 insertions(+), 51 deletions(-) mode change 100755 => 100644 tests/unit/runtime/test_pld.py diff --git a/tests/unit/elasticity/test_elastic.py b/tests/unit/elasticity/test_elastic.py index e7c42de7d317..5702e6e4d7d8 100644 --- a/tests/unit/elasticity/test_elastic.py +++ b/tests/unit/elasticity/test_elastic.py @@ -176,7 +176,7 @@ def test_proper_mbsz(ds_config): class TestNonElasticBatchParams(DistributedTest): - world_size = [1, 2] + world_size = 2 def test(self): config_dict = { @@ -213,7 +213,7 @@ def test(self): class TestNonElasticBatchParamsWithOverride(DistributedTest): - world_size = [1, 2] + world_size = 2 def test(self): config_dict = { @@ -249,7 +249,7 @@ def test(self): class TestElasticConfigChanged(DistributedTest): - world_size = [1, 2] + world_size = 2 def test(self): config_dict = { diff --git a/tests/unit/runtime/fp16/onebit/test_onebit.py b/tests/unit/runtime/fp16/onebit/test_onebit.py index 6c150117ac71..139ac79c468e 100644 --- a/tests/unit/runtime/fp16/onebit/test_onebit.py +++ b/tests/unit/runtime/fp16/onebit/test_onebit.py @@ -33,7 +33,7 @@ @pytest.mark.parametrize("dtype", [torch.float32, torch.float16], ids=["fp32", "fp16"]) class TestOneBitAdamBasic(DistributedTest): - world_size = [1, 2] + world_size = 2 def test(self, dtype): config_dict = { @@ -418,7 +418,7 @@ def test(self, topo_config): @pytest.mark.parametrize("dtype", [torch.float32, torch.float16], ids=["fp32", "fp16"]) class TestZeroOneAdamBasic(DistributedTest): - world_size = [1, 2] + world_size = 2 def test(self, dtype): config_dict = { @@ -814,7 +814,7 @@ def test(self, topo_config): @pytest.mark.parametrize("dtype", [torch.float32, torch.float16], ids=["fp32", "fp16"]) class TestOneBitLambBasic(DistributedTest): - world_size = [1, 2] + world_size = 2 def test(self, dtype): config_dict = { @@ -1235,7 +1235,7 @@ def test(self, topo_config): @pytest.mark.sequential class TestCompressedAllReduceBasic(DistributedTest): - world_size = [1, 2] + world_size = 2 def test(self, tmpdir): from deepspeed.runtime.comm.nccl import NcclBackend diff --git a/tests/unit/runtime/test_bf16.py b/tests/unit/runtime/test_bf16.py index 4d9f7aa7746c..d75b92441117 100644 --- a/tests/unit/runtime/test_bf16.py +++ b/tests/unit/runtime/test_bf16.py @@ -9,11 +9,10 @@ from deepspeed import comm as dist -@pytest.mark.parametrize('zero_stage, use_cpu_offload', [(2, False)]) class TestAdamBF16ZeroOneCycleCompatibility(DistributedTest): world_size = 1 - def test(self, zero_stage, use_cpu_offload): + 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" @@ -23,7 +22,6 @@ def test(self, zero_stage, use_cpu_offload): pytest.skip("cpu-adam is not compatible") config_dict = { - "train_batch_size": 1, "steps_per_print": 1, "optimizer": { "type": "Adam", @@ -73,11 +71,10 @@ def test(self, zero_stage, use_cpu_offload): model.step() -@pytest.mark.parametrize('zero_stage, use_cpu_offload', [(2, False)]) class TestZeroAllowUntestedOptimizer(DistributedTest): world_size = 1 - def test(self, zero_stage, use_cpu_offload): + 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" @@ -112,11 +109,10 @@ def test(self, zero_stage, use_cpu_offload): model_parameters=model.parameters()) -@pytest.mark.parametrize('zero_stage, use_cpu_offload', [(2, False)]) class TestZeroEmptyPartition(DistributedTest): world_size = 3 - def test(self, zero_stage, use_cpu_offload): + 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" @@ -172,15 +168,11 @@ def test(self, zero_stage, use_cpu_offload): model.step() -@pytest.mark.parametrize('zero_stage, optimizer_constructor', - [(2, - torch.optim.Adam), - (2, - FusedAdam)]) +@pytest.mark.parametrize("optimizer_constructor", [torch.optim.Adam, FusedAdam]) class TestZeroSupportedClientOptimizer(DistributedTest): world_size = 1 - def test(self, zero_stage, optimizer_constructor): + 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" @@ -259,11 +251,10 @@ def test(self): model.step() -@pytest.mark.parametrize('stage', [2]) class TestZeroEmptyGrad(DistributedTest): world_size = 1 - def test(self, stage): + 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" @@ -300,55 +291,47 @@ def test(self, stage): model.step() -@pytest.mark.parametrize('comp_type_str, comm_type_str', - [("fp16", - "fp16"), - ("bfp16", - "bfp16"), - ("fp16", - "bfp16"), - ("bfp16", - "fp16"), - ("fp32", - "fp16"), - ("fp32", - "bfp16")]) +@pytest.mark.parametrize("comp_type", + [torch.float16, + torch.bfloat16, + torch.float], + ids=["fp16", + "bfp16", + "fp32"]) +@pytest.mark.parametrize("comm_type", + [torch.float16, + torch.bfloat16], + ids=["fp16", + "bfp16"]) class TestZeroDtypeCocktail(DistributedTest): world_size = 2 - def test(self, comp_type_str, comm_type_str): - torch_dtype_dict = { - "fp16": torch.float16, - "bfp16": torch.bfloat16, - "fp32": torch.float - } - comp_torch_dtype = torch_dtype_dict[comp_type_str] - comm_torch_dtype = torch_dtype_dict[comm_type_str] - - if comp_torch_dtype == torch.bfloat16 or comm_torch_dtype == torch.bfloat16: + 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" ) + type_str = {torch.float16: "fp16", torch.bfloat16: "bfp16"} + config_dict = { "train_batch_size": 2, "steps_per_print": 1, "fp16": { - "enabled": torch_dtype_dict[comp_type_str] == torch.float16 + "enabled": comp_type == torch.float16 }, "bf16": { - "enabled": torch_dtype_dict[comp_type_str] == torch.bfloat16 + "enabled": comp_type == torch.bfloat16 }, "zero_optimization": { "stage": 2 }, - "communication_data_type": comm_type_str + "communication_data_type": type_str[comm_type] } hidden_dim = 10 model = SimpleModel(hidden_dim) - orig_torch_reduce = dist.reduce optimizer = torch.optim.Adam(model.parameters()) model, _, _, _ = deepspeed.initialize(config=config_dict, model=model, @@ -357,12 +340,13 @@ def test(self, comp_type_str, comm_type_str): total_samples=2, hidden_dim=hidden_dim, device=model.device, - dtype=comp_torch_dtype) + dtype=comp_type) def custom_reduce(tensor, dst, op=dist.ReduceOp.SUM, group=None, async_op=False): - assert tensor.dtype == comm_torch_dtype + 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]) diff --git a/tests/unit/runtime/test_pld.py b/tests/unit/runtime/test_pld.py old mode 100755 new mode 100644 From 4d605a48350200c6ad0cac1cd10de14a8ba7b496 Mon Sep 17 00:00:00 2001 From: Michael Wyatt Date: Wed, 10 Aug 2022 13:43:55 -0700 Subject: [PATCH 25/26] improvements to DistributedTest to allow kwargs in test method defs --- tests/unit/comm/test_dist.py | 15 +++++++++++++-- tests/unit/common.py | 27 +++++++++++++++++++++++---- 2 files changed, 36 insertions(+), 6 deletions(-) diff --git a/tests/unit/comm/test_dist.py b/tests/unit/comm/test_dist.py index 34a86a9e0c0c..32ea9126d75b 100644 --- a/tests/unit/comm/test_dist.py +++ b/tests/unit/comm/test_dist.py @@ -44,13 +44,24 @@ def test_one(self, number): assert dist.get_world_size() == 2 assert number == 1138 - @pytest.mark.parametrize("color", ["purple"]) - def test_two(self, number, color): + def test_two(self, number, color="purple"): assert dist.get_world_size() == 2 assert number == 1138 assert color == "purple" +# Demonstration of world_size override +class TestWorldSizeOverrideDistTest(DistributedTest): + world_size = 2 + + def test_world_size_2(self): + assert dist.get_world_size() == 2 + + @pytest.mark.world_size(1) + def test_world_size_1(self): + assert dist.get_world_size() == 1 + + class TestDistAllReduce(DistributedTest): world_size = [1, 2, 4] diff --git a/tests/unit/common.py b/tests/unit/common.py index 7b418ebbf645..c86122ae2a21 100644 --- a/tests/unit/common.py +++ b/tests/unit/common.py @@ -12,6 +12,7 @@ import pytest from _pytest.outcomes import Skipped +from _pytest.fixtures import FixtureLookupError # Worker timeout *after* the first worker has completed. DEEPSPEED_UNIT_WORKER_TIMEOUT = 120 @@ -67,12 +68,27 @@ class DistributedTest(ABC): world_size = 2 backend = "nccl" + # Temporary directory that can be shared among test methods in a class + @pytest.fixture(autouse=True, scope="class") + def class_tmpdir(self, tmpdir_factory): + fn = tmpdir_factory.mktemp(self.__class__.__name__) + return fn + def _run_test(self, request): self.current_test = self._get_current_test_func(request) self.test_kwargs = self._get_test_kwargs(request) - if isinstance(self.world_size, int): - self.world_size = [self.world_size] - for procs in self.world_size: + + # Catch world_size override pytest mark + for mark in getattr(request.function, "pytestmark", []): + if mark.name == "world_size": + world_size = mark.args[0] + break + else: + world_size = self.world_size + + if isinstance(world_size, int): + world_size = [world_size] + for procs in world_size: self._launch_procs(procs) time.sleep(0.5) @@ -87,7 +103,10 @@ def _get_test_kwargs(self, request): params = inspect.getfullargspec(self.current_test).args params.remove("self") for p in params: - test_kwargs[p] = request.getfixturevalue(p) + try: + test_kwargs[p] = request.getfixturevalue(p) + except FixtureLookupError: + pass # test methods can have kwargs that are not fixtures return test_kwargs def _launch_procs(self, num_procs): From 3c92a5a2e1dec362d166c7853e3edbe70222d3fb Mon Sep 17 00:00:00 2001 From: Michael Wyatt Date: Wed, 10 Aug 2022 13:51:52 -0700 Subject: [PATCH 26/26] Update common.py --- tests/unit/common.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tests/unit/common.py b/tests/unit/common.py index c86122ae2a21..36a1a7743156 100644 --- a/tests/unit/common.py +++ b/tests/unit/common.py @@ -68,7 +68,7 @@ class DistributedTest(ABC): world_size = 2 backend = "nccl" - # Temporary directory that can be shared among test methods in a class + # Temporary directory that is shared among test methods in a class @pytest.fixture(autouse=True, scope="class") def class_tmpdir(self, tmpdir_factory): fn = tmpdir_factory.mktemp(self.__class__.__name__)