diff --git a/.github/workflows/amd.yml b/.github/workflows/amd.yml index bb2a72f8f09e..97b96f2d5875 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 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 028b05b43748..1dc535f9b327 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/ --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/{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_pipe.py b/tests/unit/alexnet_model.py similarity index 52% rename from tests/unit/test_pipe.py rename to tests/unit/alexnet_model.py index de1bd3ff279c..ab51a4630b7f 100755 --- a/tests/unit/test_pipe.py +++ b/tests/unit/alexnet_model.py @@ -1,30 +1,12 @@ -import copy - +import pytest 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): @@ -103,8 +85,13 @@ 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 + torchvision = pytest.importorskip("torchvision", minversion="0.5.0") import torchvision.transforms as transforms transform_list = [ @@ -117,7 +104,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) @@ -136,7 +123,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, + 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 +136,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['local_rank'] = dist.get_rank() engine, _, _, _ = deepspeed.initialize( - args=args, + config=config, model=model, model_parameters=[p for p in model.parameters()], training_data=trainset) @@ -166,102 +158,3 @@ def train_cifar(model, args, num_steps=400, average_dp_losses=True, fp16=True, s 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_autotuning.py b/tests/unit/autotuning/test_autotuning.py similarity index 98% rename from tests/unit/test_autotuning.py rename to tests/unit/autotuning/test_autotuning.py index 681c3108b15b..5ece541e1bb7 100644 --- a/tests/unit/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/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/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..36a1a7743156 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 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__) + 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): 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/test_elastic.py b/tests/unit/elasticity/test_elastic.py similarity index 59% rename from tests/unit/test_elastic.py rename to tests/unit/elasticity/test_elastic.py index 4ed2c0dd0c95..5702e6e4d7d8 100644 --- a/tests/unit/test_elastic.py +++ b/tests/unit/elasticity/test_elastic.py @@ -1,28 +1,31 @@ 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 - -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 +from tests.unit.simple_model import SimpleModel + + +@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 @@ -185,123 +175,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 = 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 = 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 = 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/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/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 70% rename from tests/unit/test_sparse_attention.py rename to tests/unit/ops/sparse_attention/test_sparse_attention.py index 8ff843c93169..740dfacdd0de --- a/tests/unit/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 """ diff --git a/tests/unit/test_pipe_module.py b/tests/unit/pipe/test_pipe_module.py similarity index 82% rename from tests/unit/test_pipe_module.py rename to tests/unit/pipe/test_pipe_module.py index 1cba989b54e8..99862a96a2d3 100644 --- a/tests/unit/test_pipe_module.py +++ b/tests/unit/pipe/test_pipe_module.py @@ -7,16 +7,10 @@ import pytest import deepspeed - -from deepspeed.runtime.pipe.topology import PipeDataParallelTopology - -PipeTopo = PipeDataParallelTopology - 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 HIDDEN_DIM = 32 LAYERS = 8 @@ -34,7 +28,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 +47,18 @@ 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) +@pytest.fixture +def batch_input(): + return torch.randn(1, HIDDEN_DIM) - @distributed_test(world_size=4) - def _helper(): + +class TestPipeModuleSequential(DistributedTest): + world_size = 2 + + 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) @@ -69,7 +66,7 @@ def _helper(): 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()) @@ -79,7 +76,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 +95,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/fp16/onebit/test_onebit.py b/tests/unit/runtime/fp16/onebit/test_onebit.py new file mode 100644 index 000000000000..139ac79c468e --- /dev/null +++ b/tests/unit/runtime/fp16/onebit/test_onebit.py @@ -0,0 +1,1305 @@ +import torch +import torch.nn as nn +import deepspeed.comm as dist +import deepspeed +import pytest +import copy +import os +import numpy as np + +from deepspeed.runtime.pipe.topology import 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: + 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) + + +@pytest.mark.parametrize("dtype", [torch.float32, torch.float16], ids=["fp32", "fp16"]) +class TestOneBitAdamBasic(DistributedTest): + world_size = 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, + }, + } + 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, + ) + for n, batch in enumerate(data_loader): + loss = model(batch[0], batch[1]) + model.backward(loss) + model.step() + + +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 + }, + } + 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, + 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 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" + + +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 + }, + } + 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, + 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") + 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, + ) + # 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" + # 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 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, + ) + optimizer_3.optimizer.freeze_step = 20 + 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" + # 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 optimizer_3.optimizer.adam_freeze_key is False + + 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 + }, + } + 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=100, + hidden_dim=hidden_dim, + device=model.device) + 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: + loss = loss * 1000000.0 + model.backward(loss) + dist.barrier() + model.step() + dist.barrier() + model.save_checkpoint(save_folder, tag=None) + + +@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_config): + 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 + }, + } + + topo = PipeTopo(**topo_config) + steps = 500 # Must be >=100 + + # Allocate model for consistent initial weights. + init_net = AlexNetPipe() + + 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"], + ) + + +@pytest.mark.parametrize("dtype", [torch.float32, torch.float16], ids=["fp32", "fp16"]) +class TestZeroOneAdamBasic(DistributedTest): + world_size = 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, + }, + } + 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, + ) + for n, batch in enumerate(data_loader): + loss = model(batch[0], batch[1]) + model.backward(loss) + model.step() + + +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 + }, + } + 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, + 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 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" + + +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 + }, + } + 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, + 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") + 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, + ) + # 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" + # 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, + ) + optimizer_3.optimizer.freeze_step = 20 + 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" + # 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" + + 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 + }, + } + 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=100, + hidden_dim=hidden_dim, + device=model.device) + 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: + loss = loss * 1000000.0 + model.backward(loss) + dist.barrier() + model.step() + dist.barrier() + model.save_checkpoint(save_folder, tag=None) + + +@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_config): + 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 + }, + } + + topo = PipeTopo(**topo_config) + steps = 500 # Must be >=100 + + # Allocate model for consistent initial weights. + init_net = AlexNetPipe() + + 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"], + ) + + +@pytest.mark.parametrize("dtype", [torch.float32, torch.float16], ids=["fp32", "fp16"]) +class TestOneBitLambBasic(DistributedTest): + world_size = 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, + }, + } + 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, + ) + for n, batch in enumerate(data_loader): + loss = model(batch[0], batch[1]) + model.backward(loss) + model.step() + + +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 + }, + } + 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, + 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 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" + + +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 + }, + } + 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, + 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" + 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") + 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, + ) + # 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" + # 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 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, + ) + optimizer_3.optimizer.freeze_step = 20 + 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" + # 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 optimizer_3.optimizer.lamb_freeze_key is False + + 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 + }, + } + 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=100, + hidden_dim=hidden_dim, + device=model.device) + 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: + loss = loss * 1000000.0 + model.backward(loss) + dist.barrier() + model.step() + dist.barrier() + model.save_checkpoint(save_folder, tag=None) + + +@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_config): + 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 + }, + } + + topo = PipeTopo(**topo_config) + steps = 500 # Must be >=100 + + # Allocate model for consistent initial weights. + init_net = AlexNetPipe() + + 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"], + ) + + +@pytest.mark.sequential +class TestCompressedAllReduceBasic(DistributedTest): + world_size = 2 + + def test(self, tmpdir): + from deepspeed.runtime.comm.nccl import NcclBackend + + size = dist.get_world_size() + rank = dist.get_rank() + backend = NcclBackend() + local_rank = dist.get_rank() + device = torch.device("cuda", dist.get_rank()) + + # A simulated compression function using deepspeed.comm + def torch_sim(a): + a_sign = a.sign().add_(1).bool().float().add_(-0.5).mul_(2.0) + scale = a.norm() / np.sqrt(a.numel()) + a_compressed = scale * a_sign + a_sign = None + 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_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 + ] + a_sign_list = torch.chunk(a_server_sign, dist.get_world_size()) + a_server_compressed = torch.cat( + [server_scale[i] * a_sign_list[i] for i in range(dist.get_world_size())]) + rank = dist.get_rank() + server_error = a_list[rank] - server_scale[rank] * a_sign_list[rank] + torch.cuda.synchronize() + dist.barrier() + return a_server_compressed, worker_error, server_error + + tensor_size = 300 * 2**20 + server_size = int(tensor_size / size) + if tensor_size % (8 * size) != 0: + right_tensor_size = tensor_size + (8 * size - (tensor_size % (8 * size))) + else: + right_tensor_size = tensor_size + right_server_size = right_tensor_size // size + + # Adding bias to the initialization of the gradient we are communicating + # In order to get rid of the case where some elements in the gradient are too small + a = (torch.rand(tensor_size, device=device) - 0.5) + 0.01 * rank + + worker_error = torch.zeros(right_tensor_size, device=device) + server_error = torch.zeros(right_server_size, device=device) + + a_torch, worker_error_torch, server_error_torch = torch_sim(a) + torch.cuda.empty_cache() + + a_after = backend.compressed_allreduce(a, worker_error, server_error, local_rank) + + threshold = 1e-6 + magnitude_threshold = 1e-6 + diff_mask = (a_after - a_torch) > threshold + diff_server_mask = torch.chunk(diff_mask, size)[rank] + mpi_server = torch.chunk(a_after, size)[rank] + server_error + torch_server = torch.chunk(a_torch, size)[rank] + server_error_torch + + # If the number in the compensated_server_m is too small (e.g 1e-8), then calling sign() might be problematic + # 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))) + assert torch.sum(diff_server_mask) == 0 or torch.sum(check_mag_mask) == 0 diff --git a/tests/unit/runtime/pipe/test_pipe.py b/tests/unit/runtime/pipe/test_pipe.py new file mode 100644 index 000000000000..7c097b1522a5 --- /dev/null +++ b/tests/unit/runtime/pipe/test_pipe.py @@ -0,0 +1,117 @@ +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.05 # Originally 0.03, but seeing instability with AMD results 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..68d13c2d9ba6 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 - 1) 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/runtime/test_bf16.py b/tests/unit/runtime/test_bf16.py new file mode 100644 index 000000000000..d75b92441117 --- /dev/null +++ b/tests/unit/runtime/test_bf16.py @@ -0,0 +1,355 @@ +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 + + +class TestAdamBF16ZeroOneCycleCompatibility(DistributedTest): + world_size = 1 + + def test(self, zero_stage=2, use_cpu_offload=False): + if not bf16_required_version_check(): + pytest.skip( + " DeepSpeed BFloat16 tests need torch >= 1.10, NCCL >= 2.10.3, CUDA > =11.0 and HW support for BFloat16 to run correctly" + ) + + if use_cpu_offload and not deepspeed.ops.__compatible_ops__[CPUAdamBuilder.NAME]: + pytest.skip("cpu-adam is not compatible") + + config_dict = { + "steps_per_print": 1, + "optimizer": { + "type": "Adam", + "params": { + "lr": 0.00015 + } + }, + "scheduler": { + "type": "OneCycle", + "params": { + "cycle_first_step_size": 16000, + "cycle_first_stair_count": 8000, + "decay_step_size": 16000, + "cycle_min_lr": 1e-06, + "cycle_max_lr": 3e-05, + "decay_lr_rate": 1e-07, + "cycle_min_mom": 0.85, + "cycle_max_mom": 0.99, + "decay_mom_rate": 0.0 + } + }, + "fp16": { + "enabled": False + }, + "bf16": { + "enabled": True + }, + "zero_optimization": { + "stage": zero_stage, + "cpu_offload": use_cpu_offload + } + } + + hidden_dim = 10 + model = SimpleModel(hidden_dim) + model, _, _, _ = deepspeed.initialize(config=config_dict, + model=model, + model_parameters=model.parameters()) + data_loader = random_dataloader(model=model, + total_samples=50, + hidden_dim=hidden_dim, + device=model.device, + dtype=torch.bfloat16) + for n, batch in enumerate(data_loader): + loss = model(batch[0], batch[1]) + model.backward(loss) + model.step() + + +class TestZeroAllowUntestedOptimizer(DistributedTest): + world_size = 1 + + def test(self, zero_stage=2, use_cpu_offload=False): + if not bf16_required_version_check(): + pytest.skip( + " DeepSpeed BFloat16 tests need torch >= 1.10, NCCL >= 2.10.3, CUDA > =11.0 and HW support for BFloat16 to run correctly" + ) + + if use_cpu_offload and not deepspeed.ops.__compatible_ops__[CPUAdamBuilder.NAME]: + pytest.skip("cpu-adam is not compatible") + + config_dict = { + "train_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()) + + +class TestZeroEmptyPartition(DistributedTest): + world_size = 3 + + def test(self, zero_stage=2, use_cpu_offload=False): + if not bf16_required_version_check(): + pytest.skip( + " DeepSpeed BFloat16 tests need torch >= 1.10, NCCL >= 2.10.3, CUDA > =11.0 and HW support for BFloat16 to run correctly" + ) + + if use_cpu_offload and not deepspeed.ops.__compatible_ops__[CPUAdamBuilder.NAME]: + pytest.skip("cpu-adam is not compatible") + + if zero_stage == 3: + pytest.skip("skip for now") + + config_dict = { + "train_micro_batch_size_per_gpu": 1, + "gradient_accumulation_steps": 1, + "fp16": { + "enabled": False + }, + "bf16": { + "enabled": True + }, + "optimizer": { + "type": "Adam", + "params": { + "lr": 0.00015 + } + }, + "zero_optimization": { + "stage": zero_stage, + "cpu_offload": use_cpu_offload, + "reduce_bucket_size": 100, + "allgather_bucket_size": 100 + } + } + + hidden_dim = 1 + model = SimpleModel(hidden_dim) + + # Ensure model has 2 parameters, to cause empty partition with DP=3 + assert len(list(model.parameters())) == 2 + model, _, _, _ = deepspeed.initialize(config=config_dict, + model=model, + model_parameters=model.parameters()) + + # Now make sure things work.. + data_loader = random_dataloader(model=model, + total_samples=1, + hidden_dim=hidden_dim, + device=model.device, + dtype=torch.bfloat16) + for n, batch in enumerate(data_loader): + loss = model(batch[0], batch[1]) + model.backward(loss) + model.step() + + +@pytest.mark.parametrize("optimizer_constructor", [torch.optim.Adam, FusedAdam]) +class TestZeroSupportedClientOptimizer(DistributedTest): + world_size = 1 + + def test(self, optimizer_constructor, zero_stage=2): + if not bf16_required_version_check(): + pytest.skip( + " DeepSpeed BFloat16 tests need torch >= 1.10, NCCL >= 2.10.3, CUDA > =11.0 and HW support for BFloat16 to run correctly" + ) + + config_dict = { + "train_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() + + +class TestZeroEmptyGrad(DistributedTest): + world_size = 1 + + def test(self, stage=2): + if not bf16_required_version_check(): + pytest.skip( + " DeepSpeed BFloat16 tests need torch >= 1.10, NCCL >= 2.10.3, CUDA > =11.0 and HW support for BFloat16 to run correctly" + ) + + config_dict = { + "train_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() + + +@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, 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": comp_type == torch.float16 + }, + "bf16": { + "enabled": comp_type == torch.bfloat16 + }, + "zero_optimization": { + "stage": 2 + }, + "communication_data_type": type_str[comm_type] + } + hidden_dim = 10 + + model = SimpleModel(hidden_dim) + optimizer = torch.optim.Adam(model.parameters()) + model, _, _, _ = deepspeed.initialize(config=config_dict, + model=model, + optimizer=optimizer) + data_loader = random_dataloader(model=model, + total_samples=2, + hidden_dim=hidden_dim, + device=model.device, + dtype=comp_type) + + def custom_reduce(tensor, dst, op=dist.ReduceOp.SUM, group=None, async_op=False): + assert tensor.dtype == comm_type + return orig_torch_reduce(tensor, dst, op, group, async_op) + + orig_torch_reduce = dist.reduce + dist.reduce = custom_reduce + for n, batch in enumerate(data_loader): + loss = model(batch[0], batch[1]) + model.backward(loss) + model.step() + dist.reduce = orig_torch_reduce 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_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_pld.py b/tests/unit/runtime/test_pld.py old mode 100755 new mode 100644 similarity index 50% rename from tests/unit/test_pld.py rename to tests/unit/runtime/test_pld.py index 0953b648dce4..736fb3459438 --- 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_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_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_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/test_onebit.py b/tests/unit/test_onebit.py deleted file mode 100644 index b6f1f8bd4e15..000000000000 --- a/tests/unit/test_onebit.py +++ /dev/null @@ -1,1331 +0,0 @@ -import torch -import torch.nn as nn -import deepspeed.comm as dist -import deepspeed -import pytest -import copy -import os -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 .test_pipe import AlexNetPipe, train_cifar - -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) - -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" - } - }, - "gradient_clipping": 1.0, - } - args = args_from_dict(tmpdir, config_dict) - hidden_dim = 10 - - model = SimpleModel(hidden_dim) - - @distributed_test(world_size=[1, 2]) - def _test_onebitadam_fp32_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, - dtype=torch.float) - 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" - } - }, - "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, - model=model, - model_parameters=optimizer_grouped_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 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" - - _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" - } - }, - "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, - 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') - 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(args=args, - 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" - # 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 optimizer_2.optimizer.adam_freeze_key is True - - model_3, optimizer_3, _, _ = deepspeed.initialize(args=args, - 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) - 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" - # 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 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" - } - }, - "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=[2]) - def _test_onebitadam_checkpointing_overflow(args, model, hidden_dim): - model, _, _, _ = deepspeed.initialize(args=args, - 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') - for n, batch in enumerate(data_loader): - loss = model(batch[0], batch[1]) - if dist.get_rank() == 0 and n >= 10: - loss = loss * 1000000.0 - model.backward(loss) - dist.barrier() - model.step() - 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', - [ - PipeTopo(num_pp=1, - num_dp=4), - PipeTopo(num_pp=2, - num_dp=2), - 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" - } - }, - "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() - - @distributed_test(world_size=4) - def _helper(topo, tmpdir, steps=500): - assert steps >= 100 - - 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, - args, - 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" - } - }, - "gradient_clipping": 1.0, - } - args = args_from_dict(tmpdir, config_dict) - hidden_dim = 10 - - model = SimpleModel(hidden_dim) - - @distributed_test(world_size=[1, 2]) - def _test_zerooneadam_fp32_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, - dtype=torch.float) - 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" - } - }, - "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, - model=model, - model_parameters=optimizer_grouped_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 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" - - _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" - } - }, - "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, - 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') - 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(args=args, - 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" - # 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(args=args, - 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) - 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" - # 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" - - _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" - } - }, - "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=[2]) - def _test_zerooneadam_checkpointing_overflow(args, model, hidden_dim): - model, _, _, _ = deepspeed.initialize(args=args, - 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') - for n, batch in enumerate(data_loader): - loss = model(batch[0], batch[1]) - if dist.get_rank() == 0 and n >= 10: - loss = loss * 1000000.0 - model.backward(loss) - dist.barrier() - model.step() - 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', - [ - PipeTopo(num_pp=1, - num_dp=4), - PipeTopo(num_pp=2, - num_dp=2), - 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" - } - }, - "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() - - @distributed_test(world_size=4) - def _helper(topo, tmpdir, steps=500): - assert steps >= 100 - - 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, - args, - 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 - } - }, - "gradient_clipping": 1.0, - } - args = args_from_dict(tmpdir, config_dict) - hidden_dim = 10 - - model = SimpleModel(hidden_dim) - - @distributed_test(world_size=[1, 2]) - def _test_onebitlamb_fp32_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, - dtype=torch.float) - 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 - } - }, - "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, - model=model, - model_parameters=optimizer_grouped_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 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" - - _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 - } - }, - "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, - 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" - 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') - 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(args=args, - 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" - # 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 optimizer_2.optimizer.lamb_freeze_key is True - - model_3, optimizer_3, _, _ = deepspeed.initialize(args=args, - 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) - 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" - # 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 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 - } - }, - "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=[2]) - def _test_onebitlamb_checkpointing_overflow(args, model, hidden_dim): - model, _, _, _ = deepspeed.initialize(args=args, - 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') - for n, batch in enumerate(data_loader): - loss = model(batch[0], batch[1]) - if dist.get_rank() == 0 and n >= 10: - loss = loss * 1000000.0 - model.backward(loss) - dist.barrier() - model.step() - 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', - [ - PipeTopo(num_pp=1, - num_dp=4), - PipeTopo(num_pp=2, - num_dp=2), - 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" - } - }, - "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() - - @distributed_test(world_size=4) - def _helper(topo, tmpdir, steps=500): - assert steps >= 100 - - 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, - args, - 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(): - from deepspeed.runtime.comm.nccl import NcclBackend - size = dist.get_world_size() - rank = dist.get_rank() - backend = NcclBackend() - local_rank = dist.get_rank() - device = torch.device("cuda", dist.get_rank()) - - # A simulated compression function using deepspeed.comm - def torch_sim(a): - a_sign = a.sign().add_(1).bool().float().add_(-0.5).mul_(2.0) - scale = a.norm() / np.sqrt(a.numel()) - a_compressed = scale * a_sign - a_sign = None - 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_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 - ] - a_sign_list = torch.chunk(a_server_sign, dist.get_world_size()) - a_server_compressed = torch.cat( - [server_scale[i] * a_sign_list[i] for i in range(dist.get_world_size())]) - rank = dist.get_rank() - server_error = a_list[rank] - server_scale[rank] * a_sign_list[rank] - torch.cuda.synchronize() - dist.barrier() - return a_server_compressed, worker_error, server_error - - tensor_size = 300 * 2**20 - server_size = int(tensor_size / size) - if tensor_size % (8 * size) != 0: - right_tensor_size = tensor_size + (8 * size - (tensor_size % (8 * size))) - else: - right_tensor_size = tensor_size - right_server_size = right_tensor_size // size - - # Adding bias to the initialization of the gradient we are communicating - # In order to get rid of the case where some elements in the gradient are too small - a = (torch.rand(tensor_size, device=device) - 0.5) + 0.01 * rank - - worker_error = torch.zeros(right_tensor_size, device=device) - server_error = torch.zeros(right_server_size, device=device) - - a_torch, worker_error_torch, server_error_torch = torch_sim(a) - torch.cuda.empty_cache() - - a_after = backend.compressed_allreduce(a, worker_error, server_error, local_rank) - - threshold = 1e-6 - magnitude_threshold = 1e-6 - diff_mask = (a_after - a_torch) > threshold - diff_server_mask = torch.chunk(diff_mask, size)[rank] - mpi_server = torch.chunk(a_after, size)[rank] + server_error - torch_server = torch.chunk(a_torch, size)[rank] + server_error_torch - - # If the number in the compensated_server_m is too small (e.g 1e-8), then calling sign() might be problematic - # 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))) - assert torch.sum(diff_server_mask) == 0 or torch.sum(check_mag_mask) == 0 - - _test_compressed_allreduce_basic() diff --git a/tests/unit/test_sparse_grads.py b/tests/unit/test_sparse_grads.py index 5be8ec3968fb..94cc389db7b5 100644 --- a/tests/unit/test_sparse_grads.py +++ b/tests/unit/test_sparse_grads.py @@ -1,42 +1,48 @@ 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 - model = Model() - optimizer = Adam(list(model.linear.parameters()), list(model.emb.parameters())) - @distributed_test(world_size=[2]) - def _test(model, optimizer): +class TestSparseAdam(DistributedTest): + world_size = 2 + + def test(self): + 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())) engine, _, _, _ = deepspeed.initialize(model=model, optimizer=optimizer, config=config_dict) @@ -64,5 +70,3 @@ def _test(model, optimizer): ] for res in results: assert torch.allclose(res[0], res[1]) - - _test(model, optimizer) 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 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