From 16e06026ce1e89bdad8a93bb67d358e13dd4ee84 Mon Sep 17 00:00:00 2001 From: pengdurice Date: Sat, 12 Sep 2026 18:24:08 +0000 Subject: [PATCH] Deprecate sparse_gradients config flag Remove the `sparse_gradients` flag Tracking: #8489 Signed-off-by: pengdurice --- deepspeed/runtime/config.py | 10 +-- deepspeed/runtime/constants.py | 4 - deepspeed/runtime/engine.py | 55 ++---------- deepspeed/runtime/sparse_tensor.py | 20 ++--- deepspeed/runtime/utils.py | 5 +- docs/_pages/config-json.md | 6 -- docs/_tutorials/cifar-10.md | 1 - tests/unit/checkpoint/common.py | 1 - tests/unit/checkpoint/test_sparse.py | 88 ------------------- .../test_averaging_sparse_gradients.py | 4 +- tests/unit/runtime/sparse_tensor/test_csr.py | 55 ------------ .../sparse_tensor/test_sparse_grads.py | 4 +- tests/unit/runtime/test_ds_config_dict.py | 10 +++ 13 files changed, 34 insertions(+), 229 deletions(-) delete mode 100644 tests/unit/checkpoint/test_sparse.py delete mode 100644 tests/unit/runtime/sparse_tensor/test_csr.py diff --git a/deepspeed/runtime/config.py b/deepspeed/runtime/config.py index 2dcb549d7d62..c6f2a17e2485 100755 --- a/deepspeed/runtime/config.py +++ b/deepspeed/runtime/config.py @@ -104,6 +104,11 @@ class DeepSpeedConfigError(Exception): "quantize_training": "Mixture-of-Quantization (MoQ) / 'quantize_training' has been removed. See " f"{_REMOVED_FEATURES_ISSUE}.", + "sparse_gradients": + "Sparse compression of dense torch.nn.Embedding gradients has been removed. A leftover " + "'sparse_gradients' flag would be ignored and those gradients would be reduced dense. " + "Gradients from an embedding constructed with sparse=True are still reduced sparsely and " + f"need no config flag. See {_REMOVED_FEATURES_ISSUE}.", } _REMOVED_ZERO_CONFIG_KEYS = { "mics_shard_size": @@ -230,10 +235,6 @@ def get_managed_gradient_accumulation(param_dict): return get_scalar_param(param_dict, MANAGED_GRADIENT_ACCUMULATION, MANAGED_GRADIENT_ACCUMULATION_DEFAULT) -def get_sparse_gradients_enabled(param_dict): - return get_scalar_param(param_dict, SPARSE_GRADIENTS, SPARSE_GRADIENTS_DEFAULT) - - def get_communication_data_type(param_dict, comm_type=COMMUNICATION_DATA_TYPE, comm_data_type_default=COMMUNICATION_DATA_TYPE_DEFAULT): @@ -819,7 +820,6 @@ def _initialize_params(self, param_dict): self.prescale_gradients = get_prescale_gradients(param_dict) self.gradient_predivide_factor = get_gradient_predivide_factor(param_dict) self.gradient_allreduce_op = get_gradient_allreduce_op(param_dict) - self.sparse_gradients_enabled = get_sparse_gradients_enabled(param_dict) self.zero_config = get_zero_config(param_dict) self.zero_optimization_stage = self.zero_config.stage diff --git a/deepspeed/runtime/constants.py b/deepspeed/runtime/constants.py index 08a355c4bcde..84e3ed3535bf 100755 --- a/deepspeed/runtime/constants.py +++ b/deepspeed/runtime/constants.py @@ -114,10 +114,6 @@ MANAGED_GRADIENT_ACCUMULATION = "managed_gradient_accumulation" MANAGED_GRADIENT_ACCUMULATION_DEFAULT = True -# DeepSpeed CSR gradient sparsity -SPARSE_GRADIENTS = "sparse_gradients" -SPARSE_GRADIENTS_DEFAULT = False - ######################################### # BFLOAT16 support ######################################### diff --git a/deepspeed/runtime/engine.py b/deepspeed/runtime/engine.py index 020f180920a4..4d2d67c24b62 100644 --- a/deepspeed/runtime/engine.py +++ b/deepspeed/runtime/engine.py @@ -693,14 +693,6 @@ def __init__(self, elif self.bfloat16_enabled(): self.optimizer = self._configure_bf16_optimizer(optimizer=None) - # Bookkeeping for sparse support - self.sparse_tensor_module_names = set() - # if self.sparse_gradients_enabled(): - for name, module in self.module.named_modules(): - if isinstance(module, (torch.nn.Embedding, torch.nn.EmbeddingBag)) and self.sparse_gradients_enabled(): - self.sparse_tensor_module_names.add(name + ".weight") - logger.info("Will convert {} to sparse tensor during training".format(name)) - self._optimized_linear_offload_setup() self.save_non_zero_checkpoint = False @@ -1440,9 +1432,6 @@ def autotuning_profile_model_info(self): ) and self._config.autotuning_config.model_info and self._config.autotuning_config.model_info.get( "profile", False) - def sparse_gradients_enabled(self): - return self._config.sparse_gradients_enabled - def train_batch_size(self): return self._config.train_batch_size @@ -3875,7 +3864,7 @@ def _get_gradients_for_reduction(self): for key in self.expert_data_parallel_group.keys(): expert_grads[key] = [] - for param_name, param in self.module.named_parameters(): + for param in self.module.parameters(): if not param.requires_grad: continue @@ -3893,7 +3882,7 @@ def _get_gradients_for_reduction(self): param.grad = torch.zeros(param.size(), dtype=param.dtype, device=param.device) grad_data = param.grad.data - if param_name in self.sparse_tensor_module_names or grad_data.is_sparse: + if grad_data.is_sparse: # Call param.grad without data to avoid problem with setting of updated grads grad_data = SparseTensor(param.grad) @@ -3966,12 +3955,9 @@ def buffered_allreduce_fallback(self, grads=None, elements_per_buffer=500000000) def sparse_allreduce_no_retain(self, bucket, dp_group, dp_world_size=None): allreduced_sparses = self.sparse_allreduce_bucket(bucket, dp_group, dp_world_size) - # Densify sparse tensor and copy back to original location + # Copy the reduced sparse tensor back to the original location for tensor in allreduced_sparses: - if tensor.is_sparse: - tensor.orig_dense_tensor.data = tensor.to_coo_tensor() - else: - tensor.orig_dense_tensor.copy_(tensor.to_dense()) + tensor.orig_dense_tensor.data = tensor.to_coo_tensor() def sparse_allreduce_bucket(self, bucket, dp_group, dp_world_size=None): sparse_list = [] @@ -4734,41 +4720,13 @@ def _load_checkpoint(self, ) and 'data_sampler' in checkpoint: self.training_dataloader.data_sampler.load_state_dict(checkpoint['data_sampler']) - def get_sparse_tensor_module_names(original_set, loaded_set, original_parameters, loaded_parameters): - result = set() - - for name in original_set: - if name in loaded_parameters and name not in loaded_set: - continue # parameter existed in previous model and was not sparse - result.add(name) - - for name in loaded_set: - if name in original_parameters: - result.add(name) # parameter exists in both configs and it was sparse - - return result - - if 'sparse_tensor_module_names' in checkpoint: - sparse_tensor_module_names = checkpoint['sparse_tensor_module_names'] - elif 'csr_tensor_module_names' in checkpoint: - sparse_tensor_module_names = checkpoint['csr_tensor_module_names'] - else: - sparse_tensor_module_names = None - if sparse_tensor_module_names is not None: - if load_module_strict: - self.sparse_tensor_module_names = sparse_tensor_module_names - else: - self.sparse_tensor_module_names = get_sparse_tensor_module_names( - self.sparse_tensor_module_names, sparse_tensor_module_names, - dict(self.module.named_parameters()), checkpoint["module"]) - self.global_steps = checkpoint['global_steps'] self.global_samples = checkpoint.get('global_samples', self.global_steps * self.train_batch_size()) self.skipped_steps = checkpoint['skipped_steps'] self.loaded_checkpoint_mp_world_size = checkpoint['mp_world_size'] deepspeed_states = [ - 'module', 'sparse_tensor_module_names', 'skipped_steps', 'global_steps', 'dp_world_size', - 'mp_world_size', 'data_sampler', 'random_ltd' + 'module', 'skipped_steps', 'global_steps', 'dp_world_size', 'mp_world_size', 'data_sampler', + 'random_ltd' ] client_state = {} @@ -5082,7 +5040,6 @@ def _common_checkpoint_state(self, module_state_dict, zero_optimizer_state, save data_sampler=self.training_dataloader.data_sampler.state_dict() if (self.training_dataloader is not None and self.curriculum_learning_enabled()) else None, random_ltd=self.random_ltd_scheduler.state_dict() if self.random_ltd_enabled() else None, - sparse_tensor_module_names=self.sparse_tensor_module_names, skipped_steps=self.skipped_steps, global_steps=self.global_steps, global_samples=self.global_samples, diff --git a/deepspeed/runtime/sparse_tensor.py b/deepspeed/runtime/sparse_tensor.py index f87c4e418e8b..14abdd1d9d63 100644 --- a/deepspeed/runtime/sparse_tensor.py +++ b/deepspeed/runtime/sparse_tensor.py @@ -16,16 +16,13 @@ class SparseTensor(object): def __init__(self, dense_tensor=None): self.orig_dense_tensor = dense_tensor if dense_tensor is not None: - self.is_sparse = dense_tensor.is_sparse + assert dense_tensor.is_sparse, \ + "SparseTensor is only built from a natively sparse gradient" + self.is_sparse = True self.dtype = self.orig_dense_tensor.dtype - if self.is_sparse: - dense_tensor = dense_tensor.coalesce() - self.indices = dense_tensor.indices().flatten() - self.values = dense_tensor.values() - else: - result = torch.sum(dense_tensor, dim=1) - self.indices = result.nonzero().flatten() - self.values = dense_tensor[self.indices] + dense_tensor = dense_tensor.coalesce() + self.indices = dense_tensor.indices().flatten() + self.values = dense_tensor.values() self.dense_size = list(dense_tensor.size()) else: self.indices = None @@ -39,11 +36,6 @@ def to_coo_tensor(self): def type(): return "deepspeed.SparseTensor" - def to_dense(self): - it = self.indices.unsqueeze(1) - full_indices = torch.cat([it for _ in range(self.dense_size[1])], dim=1) - return self.values.new_zeros(self.dense_size).scatter_add_(0, full_indices, self.values) - def sparse_size(self): index_size = list(self.indices.size()) index_size = index_size[0] diff --git a/deepspeed/runtime/utils.py b/deepspeed/runtime/utils.py index 54a8ddc60d26..4d23aac3e21e 100644 --- a/deepspeed/runtime/utils.py +++ b/deepspeed/runtime/utils.py @@ -235,8 +235,9 @@ def has_inf_or_nan(x): ``bool`` get one from ordinary truthiness. An empty tensor holds no non-finite value, and the guard is required rather than defensive: ``amax`` - raises on an empty input. A sparse tensor reaches this on the ``sparse_gradients`` path, where only the - stored entries can be non-finite -- the implicit zeros cannot -- so the reduction runs over those. + raises on an empty input. A sparse tensor reaches this when the model itself produces natively sparse + gradients, where only the stored entries can be non-finite -- the implicit zeros cannot -- so the + reduction runs over those. """ if x.is_sparse: # ``amax``/``amin`` have no sparse implementation, and the public ``values()`` rejects uncoalesced diff --git a/docs/_pages/config-json.md b/docs/_pages/config-json.md index 40c3c8bad285..d2ccefb48805 100644 --- a/docs/_pages/config-json.md +++ b/docs/_pages/config-json.md @@ -156,12 +156,6 @@ Example of **scheduler** | ------------------------------------------------------------------------------------------------------------------------------------------------- | ------- | | Before mean gradient allreduce, predivide gradients by a specified factor; this can sometimes help with fp16 stability when scaling to large numbers of GPUs | `1.0` | -**sparse_gradients**: [boolean] - -| Description | Default | -| ----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------- | ------- | -| Enable sparse compression of [torch.nn.Embedding](https://pytorch.org/docs/stable/nn.html#torch.nn.Embedding) gradients. This feature is essentially deprecated as we don't see use cases for it as much anymore. It should be noted that this feature is not compatible with [torch.sparse](https://pytorch.org/docs/stable/sparse.html) related features. | `false` | - ### FP16 training options **Note:** this mode cannot be combined with the `amp` mode described below. diff --git a/docs/_tutorials/cifar-10.md b/docs/_tutorials/cifar-10.md index 2bd06abf0e89..144483655922 100644 --- a/docs/_tutorials/cifar-10.md +++ b/docs/_tutorials/cifar-10.md @@ -260,7 +260,6 @@ worker-0: optimizer_params ............. {'lr': 0.001, 'betas': [0.8, 0.999], worker-0: prescale_gradients ........... False worker-0: scheduler_name ............... WarmupLR worker-0: scheduler_params ............. {'warmup_min_lr': 0, 'warmup_max_lr': 0.001, 'warmup_num_steps': 1000} -worker-0: sparse_gradients_enabled ..... False worker-0: steps_per_print .............. 2000 worker-0: tensorboard_enabled .......... False worker-0: tensorboard_job_name ......... DeepSpeedJobName diff --git a/tests/unit/checkpoint/common.py b/tests/unit/checkpoint/common.py index 0daa1b070850..c654ccd4a30d 100644 --- a/tests/unit/checkpoint/common.py +++ b/tests/unit/checkpoint/common.py @@ -23,7 +23,6 @@ def compare_deepspeed_states(saved_model, loaded_model): # These are compared in more depth in other places assert hasattr(loaded_model, 'module') - assert saved_model.sparse_tensor_module_names == loaded_model.sparse_tensor_module_names assert saved_model.skipped_steps == loaded_model.skipped_steps assert saved_model.global_steps == loaded_model.global_steps diff --git a/tests/unit/checkpoint/test_sparse.py b/tests/unit/checkpoint/test_sparse.py deleted file mode 100644 index 19fbcd81e473..000000000000 --- a/tests/unit/checkpoint/test_sparse.py +++ /dev/null @@ -1,88 +0,0 @@ -# Copyright (c) Microsoft Corporation. -# SPDX-License-Identifier: Apache-2.0 - -# DeepSpeed Team - -import deepspeed - -from unit.common import DistributedTest -from unit.simple_model import * - -import pytest - - -class TestSparseCheckpoint(DistributedTest): - world_size = 2 - - @pytest.mark.parametrize(["to_save_model_has_embedding", "to_save_model_sparse"], [ - [False, False], - [True, False], - [True, True], - ]) - @pytest.mark.parametrize(["destination_has_embedding", "destination_sparse"], [ - [False, False], - [True, False], - [True, True], - ]) - def test_non_strict_load_sparse(self, tmpdir, to_save_model_has_embedding, to_save_model_sparse, - destination_has_embedding, destination_sparse): - - class ModelNoEmbedding(torch.nn.Module): - - def __init__(self): - super().__init__() - self.linear = torch.nn.Linear(3, 1) - - def forward(self, x): - return self.linear(x) - - class ModelEmbedding(torch.nn.Module): - - def __init__(self): - super().__init__() - self.emb = torch.nn.Embedding(10, 3) - self.linear = torch.nn.Linear(3, 1) - - def forward(self, x, offsets): - return self.linear(self.emb(x, offsets)) - - if to_save_model_has_embedding: - model_to_save = ModelEmbedding() - else: - model_to_save = ModelNoEmbedding() - if destination_has_embedding: - model_destination = ModelEmbedding() - else: - model_destination = ModelNoEmbedding() - - engine_to_save, _, _, _ = deepspeed.initialize(model=model_to_save, - config={ - "train_batch_size": 2, - "sparse_gradients": to_save_model_sparse - }) - engine_destination, _, _, _ = deepspeed.initialize(model=model_destination, - config={ - "train_batch_size": 2, - "sparse_gradients": destination_sparse - }) - - save_folder = os.path.join(tmpdir, 'saved_checkpoint') - save_tag = '1' - - engine_to_save.save_checkpoint(save_folder, tag=save_tag) - - is_sparse_destination = isinstance(model_destination, ModelEmbedding) and destination_sparse - if isinstance(model_destination, ModelEmbedding) and model_destination.emb.sparse: - assert "emb.weight" in engine_destination.sparse_tensor_module_names - engine_destination.load_checkpoint(save_folder, - tag=save_tag, - load_module_strict=False, - load_optimizer_states=False, - load_lr_scheduler_states=False, - load_module_only=False) - if isinstance(model_destination, ModelEmbedding) and isinstance(model_to_save, ModelEmbedding): - assert engine_destination.sparse_tensor_module_names == engine_to_save.sparse_tensor_module_names - elif isinstance(model_destination, ModelEmbedding): - assert not is_sparse_destination or "emb.weight" in engine_destination.sparse_tensor_module_names - else: - assert len(engine_destination.sparse_tensor_module_names) == 0 diff --git a/tests/unit/runtime/sparse_tensor/test_averaging_sparse_gradients.py b/tests/unit/runtime/sparse_tensor/test_averaging_sparse_gradients.py index badd0bcee549..8f851f8ed7bd 100644 --- a/tests/unit/runtime/sparse_tensor/test_averaging_sparse_gradients.py +++ b/tests/unit/runtime/sparse_tensor/test_averaging_sparse_gradients.py @@ -11,7 +11,7 @@ from deepspeed.accelerator import get_accelerator if get_accelerator().device_name() == 'hpu': - pytest.skip("sparse_gradients not supported by HPU.", allow_module_level=True) + pytest.skip("sparse gradients not supported by HPU.", allow_module_level=True) class Model(torch.nn.Module): @@ -62,7 +62,7 @@ class TestSparseAdam(DistributedTest): def test(self): skip_on_arch(min_arch=7) - config_dict = {"train_batch_size": 2, "steps_per_print": 1, "sparse_gradients": True} + config_dict = {"train_batch_size": 2, "steps_per_print": 1} model, optimizer = get_model_optimizer() loss = torch.nn.BCEWithLogitsLoss() engine, _, _, _ = deepspeed.initialize(model=model, optimizer=optimizer, config=config_dict) diff --git a/tests/unit/runtime/sparse_tensor/test_csr.py b/tests/unit/runtime/sparse_tensor/test_csr.py deleted file mode 100644 index 937b981735a5..000000000000 --- a/tests/unit/runtime/sparse_tensor/test_csr.py +++ /dev/null @@ -1,55 +0,0 @@ -# Copyright (c) Microsoft Corporation. -# SPDX-License-Identifier: Apache-2.0 - -# DeepSpeed Team - -import torch -import random -from deepspeed.runtime.sparse_tensor import SparseTensor - - -def test_csr_addition_self(): - row_count = 10 - random.seed(1234) - - x = torch.ones(1, 5) - for i in range(row_count - 1): - if random.random() > 0.75: - x = torch.cat([x, torch.ones(1, 5)]) - else: - x = torch.cat([x, torch.zeros(1, 5)]) - dense_x = x.clone() - cx = SparseTensor(x) - - assert torch.all(dense_x == cx.to_dense()) - - cx.add(cx) - assert torch.all(dense_x + dense_x == cx.to_dense()) - - -def test_csr_addition_different(): - row_count = 10 - random.seed(1234) - - x = torch.ones(1, 5) - for i in range(row_count - 1): - if random.random() > 0.75: - x = torch.cat([x, torch.ones(1, 5)]) - else: - x = torch.cat([x, torch.zeros(1, 5)]) - dense_x = x.clone() - cx = SparseTensor(x) - - y = torch.ones(1, 5) - for i in range(row_count - 1): - if random.random() > 0.75: - y = torch.cat([y, torch.ones(1, 5)]) - else: - y = torch.cat([y, torch.zeros(1, 5)]) - dense_y = y.clone() - cy = SparseTensor(y) - - dense_sum = dense_x + dense_y - cx.add(cy) - - assert torch.all(dense_sum == cx.to_dense()) diff --git a/tests/unit/runtime/sparse_tensor/test_sparse_grads.py b/tests/unit/runtime/sparse_tensor/test_sparse_grads.py index 6338a16b8dbb..466ce049d34d 100644 --- a/tests/unit/runtime/sparse_tensor/test_sparse_grads.py +++ b/tests/unit/runtime/sparse_tensor/test_sparse_grads.py @@ -11,7 +11,7 @@ import deepspeed.utils.groups as groups if get_accelerator().device_name() == 'hpu': - pytest.skip("sparse_gradients not supported by HPU.", allow_module_level=True) + pytest.skip("sparse gradients not supported by HPU.", allow_module_level=True) class Model(torch.nn.Module): @@ -46,7 +46,7 @@ class TestSparseAdam(DistributedTest): world_size = 2 def test(self): - config_dict = {"train_batch_size": 2, "steps_per_print": 1, "sparse_gradients": True} + config_dict = {"train_batch_size": 2, "steps_per_print": 1} model = Model() optimizer = Adam(list(model.linear.parameters()), list(model.emb.parameters())) diff --git a/tests/unit/runtime/test_ds_config_dict.py b/tests/unit/runtime/test_ds_config_dict.py index ad9f1f1f8c8e..faa602f630d6 100644 --- a/tests/unit/runtime/test_ds_config_dict.py +++ b/tests/unit/runtime/test_ds_config_dict.py @@ -271,6 +271,16 @@ def test_mics_zero_config_is_rejected(): DeepSpeedConfig(config_dict) +def test_sparse_gradients_config_is_rejected(): + config_dict = { + "train_micro_batch_size_per_gpu": 1, + "sparse_gradients": True, + } + + with pytest.raises(DeepSpeedConfigError, match="sparse_gradients"): + DeepSpeedConfig(config_dict) + + def test_compression_helper_shim_reexports_module_utils(): with pytest.warns(FutureWarning, match="deepspeed.compression.helper"): from deepspeed.compression.helper import recursive_getattr, recursive_setattr