Deprecate sparse_gradients

deepspeed/runtime/config.py

@@ -105,6 +105,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}.",
     "eigenvalue":
     "Eigenvalue-based Mixture-of-Quantization (MoQ) has been removed; the standalone "
     f"'eigenvalue' configuration block is no longer supported. See {_REMOVED_FEATURES_ISSUE}.",
@@ -241,10 +246,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):
@@ -594,7 +595,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

deepspeed/runtime/constants.py

@@ -78,10 +78,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
 #########################################

deepspeed/runtime/engine.py

@@ -684,14 +684,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
@@ -1432,9 +1424,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
 
@@ -3862,7 +3851,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
 
@@ -3880,7 +3869,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)
 
@@ -3953,12 +3942,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 = []
@@ -4721,41 +4707,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 = {}
 
@@ -5069,7 +5027,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,

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]

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

docs/_pages/config-json.md

@@ -156,12 +156,6 @@ Example of <i>**scheduler**</i>
 | ------------------------------------------------------------------------------------------------------------------------------------------------- | ------- |
 | 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`   |
 
-<i>**sparse_gradients**</i>: [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.

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

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
 

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)