diff --git a/tests/unit/test_adam_acuracy.py b/tests/unit/test_adam_acuracy.py index f61b6ecba58b..d1d03f500af9 100755 --- a/tests/unit/test_adam_acuracy.py +++ b/tests/unit/test_adam_acuracy.py @@ -6,8 +6,12 @@ import pytest import copy +import deepspeed from deepspeed.ops.adam import DeepSpeedCPUAdam +if not deepspeed.ops.__installed_ops__['cpu-adam']: + pytest.skip("cpu-adam is not installed", allow_module_level=True) + def check_equal(first, second, atol=1e-2, verbose=False): x = first.detach().numpy() diff --git a/tests/unit/test_checkpointing.py b/tests/unit/test_checkpointing.py index d08addb936d8..c7dee1fa2c53 100755 --- a/tests/unit/test_checkpointing.py +++ b/tests/unit/test_checkpointing.py @@ -1,4 +1,5 @@ import torch + import torch.distributed as dist import deepspeed @@ -151,6 +152,8 @@ def checkpoint_correctness_verification(args, compare_lr_scheduler_states(trained_model, loaded_model) +@pytest.mark.skipif(not deepspeed.ops.__installed_ops__['lamb'], + reason="lamb is not installed") def test_checkpoint_unfused_optimizer(tmpdir): config_dict = { "train_batch_size": 2, @@ -264,6 +267,9 @@ def _test_checkpoint_fused_optimizer(args, model, hidden_dim, load_optimizer_sta 'deepspeed_adam'), ]) def test_checkpoint_zero_optimizer(tmpdir, zero_stage, use_cpu_offload, adam_optimizer): + if use_cpu_offload and not deepspeed.ops.__installed_ops__['cpu-adam']: + pytest.skip("cpu-adam is not installed") + config_dict = { "train_batch_size": 2, "steps_per_print": 1, @@ -320,6 +326,9 @@ def test_checkpoint_zero_no_optimizer(tmpdir, zero_stage, use_cpu_offload, adam_optimizer): + if use_cpu_offload and not deepspeed.ops.__installed_ops__['cpu-adam']: + pytest.skip("cpu-adam is not installed") + config_dict = { "train_batch_size": 2, "steps_per_print": 1, @@ -379,6 +388,9 @@ def _test_checkpoint_zero_no_optimizer(args, 'deepspeed_adam'), ]) def test_checkpoint_lr_scheduler(tmpdir, zero_stage, use_cpu_offload, adam_optimizer): + if use_cpu_offload and not deepspeed.ops.__installed_ops__['cpu-adam']: + pytest.skip("cpu-adam is not installed") + config_dict = { "train_batch_size": 2, "steps_per_print": 1, @@ -450,6 +462,9 @@ def _test_checkpoint_lr_scheduler(args, 'deepspeed_adam'), ]) def test_checkpoint_no_lr_scheduler(tmpdir, zero_stage, use_cpu_offload, adam_optimizer): + if use_cpu_offload and not deepspeed.ops.__installed_ops__['cpu-adam']: + pytest.skip("cpu-adam is not installed") + config_dict = { "train_batch_size": 2, "steps_per_print": 1, diff --git a/tests/unit/test_cuda_backward.py b/tests/unit/test_cuda_backward.py index bf0e5955d62c..498466cf63f8 100755 --- a/tests/unit/test_cuda_backward.py +++ b/tests/unit/test_cuda_backward.py @@ -12,9 +12,13 @@ from modeling import BertEncoder as BertEncoderPostln from modeling import BertConfig, BertLayerNorm from deepspeed import DeepSpeedTransformerLayer, DeepSpeedTransformerConfig +import deepspeed import sys +if not deepspeed.ops.__installed_ops__['transformer']: + pytest.skip("transformer kernels are not installed", allow_module_level=True) + def check_equal(first, second, atol=1e-2, verbose=False): diction_x = {} diff --git a/tests/unit/test_cuda_forward.py b/tests/unit/test_cuda_forward.py index 4e995a34448f..e340a6165c26 100755 --- a/tests/unit/test_cuda_forward.py +++ b/tests/unit/test_cuda_forward.py @@ -12,9 +12,13 @@ from modeling import BertEncoder as BertEncoderPostln from modeling import BertLayerNorm, BertConfig from deepspeed import DeepSpeedTransformerLayer, DeepSpeedTransformerConfig +import deepspeed import sys +if not deepspeed.ops.__installed_ops__['transformer']: + pytest.skip("transformer kernels are not installed", allow_module_level=True) + def check_equal(first, second, atol=1e-2, verbose=False): if verbose: diff --git a/tests/unit/test_dynamic_loss_scale.py b/tests/unit/test_dynamic_loss_scale.py index 7575d6b49454..5bfe3353dcbb 100755 --- a/tests/unit/test_dynamic_loss_scale.py +++ b/tests/unit/test_dynamic_loss_scale.py @@ -8,6 +8,9 @@ from common import distributed_test from simple_model import SimpleModel, args_from_dict +lamb_available = pytest.mark.skipif(not deepspeed.ops.__installed_ops__['lamb'], + reason="lamb is not installed") + def run_model_step(model, gradient_list): for value in gradient_list: @@ -165,6 +168,7 @@ def _test_fused_some_overflow(args): _test_fused_some_overflow(args) +@lamb_available def test_unfused_no_overflow(tmpdir): config_dict = { "train_batch_size": 1, @@ -208,6 +212,7 @@ def _test_unfused_no_overflow(args): _test_unfused_no_overflow(args) +@lamb_available def test_unfused_all_overflow(tmpdir): config_dict = { "train_batch_size": 1, @@ -253,6 +258,7 @@ def _test_unfused_all_overflow(args): _test_unfused_all_overflow(args) +@lamb_available def test_unfused_some_overflow(tmpdir): config_dict = { "train_batch_size": 1, diff --git a/tests/unit/test_fp16.py b/tests/unit/test_fp16.py index f5176294a549..47344bfbc9a1 100755 --- a/tests/unit/test_fp16.py +++ b/tests/unit/test_fp16.py @@ -8,7 +8,11 @@ from common import distributed_test from simple_model import SimpleModel, SimpleOptimizer, random_dataloader, args_from_dict +lamb_available = pytest.mark.skipif(not deepspeed.ops.__installed_ops__['lamb'], + reason="lamb is not installed") + +@lamb_available def test_lamb_fp32_grad_clip(tmpdir): config_dict = { "train_batch_size": 2, @@ -44,6 +48,7 @@ def _test_lamb_fp32_grad_clip(args, model, hidden_dim): _test_lamb_fp32_grad_clip(args=args, model=model, hidden_dim=hidden_dim) +@lamb_available def test_lamb_fp16_basic(tmpdir): config_dict = { "train_batch_size": 2, @@ -81,6 +86,7 @@ def _test_lamb_fp16_basic(args, model, hidden_dim): _test_lamb_fp16_basic(args=args, model=model, hidden_dim=hidden_dim) +@lamb_available def test_lamb_fp16_empty_grad(tmpdir): config_dict = { "train_batch_size": 2, @@ -228,6 +234,8 @@ def _test_adamw_fp16_empty_grad(args, model, hidden_dim): True), ]) def test_adam_fp16_zero_onecycle_compatibility(tmpdir, zero_stage, use_cpu_offload): + if use_cpu_offload and not deepspeed.ops.__installed_ops__['cpu-adam']: + pytest.skip("cpu-adam is not installed") config_dict = { "train_batch_size": 1, "steps_per_print": 1, @@ -294,6 +302,8 @@ def _test_adam_fp16_zero_onecycle_compatibility(args, model, hidden_dim): True), ]) def test_zero_static_scale(tmpdir, zero_stage, use_cpu_offload): + if use_cpu_offload and not deepspeed.ops.__installed_ops__['cpu-adam']: + pytest.skip("cpu-adam is not installed") config_dict = { "train_batch_size": 4, "steps_per_print": 1, @@ -392,6 +402,8 @@ def _test_zero_static_scale(args): True), ]) def test_zero_allow_untested_optimizer(tmpdir, zero_stage, use_cpu_offload): + if use_cpu_offload and not deepspeed.ops.__installed_ops__['cpu-adam']: + pytest.skip("cpu-adam is not installed") config_dict = { "train_batch_size": 4, "steps_per_print": 1, @@ -430,6 +442,8 @@ def _test_zero_allow_untested_optimizer(args): True), ]) def test_zero_empty_partition(tmpdir, zero_stage, use_cpu_offload): + if use_cpu_offload and not deepspeed.ops.__installed_ops__['cpu-adam']: + pytest.skip("cpu-adam is not installed") config_dict = { "train_micro_batch_size_per_gpu": 1, "gradient_accumulation_steps": 1, @@ -500,6 +514,7 @@ def _test_adam_amp_basic(args, model, hidden_dim): _test_adam_amp_basic(args=args, model=model, hidden_dim=hidden_dim) +@lamb_available def test_lamb_amp_basic(tmpdir): config_dict = { "train_batch_size": 2, diff --git a/tests/unit/test_sparse_attention.py b/tests/unit/test_sparse_attention.py index f113e3b72ecd..cb68c20cf731 100644 --- a/tests/unit/test_sparse_attention.py +++ b/tests/unit/test_sparse_attention.py @@ -5,6 +5,10 @@ import pytest import torch +import deepspeed + +if not deepspeed.ops.__installed_ops__['sparse-attn']: + pytest.skip("cpu-adam is not installed", allow_module_level=True) def test_sparse_attention_module_availability():