Deprecate sparse_gradients (#8504)
Remove the `sparse_gradients` flag
Tracking: #8489
## Summary
- Remove the `sparse_gradients` config flag, which compressed **dense**
`torch.nn.Embedding` / `EmbeddingBag` gradients into DeepSpeed's sparse
reduction path.
- Reject the removed key in `_REMOVED_TOP_LEVEL_CONFIG_KEYS` so a
leftover
`sparse_gradients` entry raises `DeepSpeedConfigError` instead of being
silently ignored.
- Drop the `sparse_tensor_module_names` bookkeeping and its checkpoint
field, which existed
only to serve that flag.
- Sweep the leftover docs, tutorial output, and tests.
- Tracking: #8489
### What the flag did, and why it goes
The flag's entire contribution was the dense-input branch of
`SparseTensor.__init__`:
```python
result = torch.sum(dense_tensor, dim=1)
self.indices = result.nonzero().flatten()
self.values = dense_tensor[self.indices]
```
It selected rows whose values **sum** to something non-zero, so a row
like `[1.0, -1.0]`
was discarded as empty even though it carries a real gradient. That
silently dropped updates.
`config-json.md` has described the feature as "essentially deprecated as
we don't see use
cases for it as much anymore" since #1418 (Sept 2021), and
`deepspeed/runtime/zero/` contains no
sparse handling at all, so the flag never worked together with ZeRO.
### What is deliberately NOT removed
**Sparse gradient reduction still works at ZeRO stage 0.** A model that
builds
`torch.nn.Embedding(..., sparse=True)` itself produces a natively sparse
gradient, and that
still reaches the full DeepSpeed sparse path.
Byte-identical in this PR:
- `split_half_float_double_sparse`
- `sparse_allreduce`, `sparse_allreduce_bucket`, `sparse_all_gather`
- the `if grad_data.is_sparse:` guard and the `SparseTensor(param.grad)`
call itself
Simplified, but behavior-preserving for a sparse gradient:
- `sparse_allreduce_no_retain` drops an `else` branch that densified and
copied back. A
`SparseTensor` is now only ever built from an already-sparse gradient,
so `is_sparse` is
always true and that branch was unreachable.
- `SparseTensor.__init__` drops its dense-input branch in favor of an
assertion, and
`to_dense()` goes with it, since that removed `else` branch was its only
caller.
Both remaining sparse tests are kept and still cover that path. They
already built
`EmbeddingBag(..., sparse=True)`, so the only change to them is dropping
the now-rejected
`"sparse_gradients": true` line from their config.
The single construction site is `engine.py`:
```python
grad_data = param.grad.data
if grad_data.is_sparse:
grad_data = SparseTensor(param.grad)
```
Reached through `backward()` or `step()` → `allreduce_gradients()` →
`buffered_allreduce_fallback()` → `_get_gradients_for_reduction()`. All
of the gates in
`allreduce_gradients` must pass, which in practice means: DeepCompile
inactive, ZeRO stage 0,
and a gradient accumulation boundary. ZeRO stages 1, 2, and 3 all divert
to the ZeRO optimizer
before that line.
**Open question for review.** Because it only works at ZeRO stage 0,
that whole path is a
candidate for deprecation too. Doing so would remove `sparse_tensor.py`,
`split_half_float_double_sparse`, the four `sparse_allreduce*` methods,
`sparse_all_gather`,
and the two remaining sparse tests. This PR deliberately leaves that
decision open rather
than bundling it.
Three `SparseTensor` members now have no caller anywhere: `add()`,
`sparse_size()`, and
`type()`. They are independent of the removed flag and are left for a
separate dead-code pass.
### Behavior changes
| Before | After |
|---|---|
| `"sparse_gradients": true` compressed dense embedding gradients |
raises `DeepSpeedConfigError` |
| `"sparse_gradients": false` parsed and did nothing | raises
`DeepSpeedConfigError` |
| `Embedding(sparse=True)` reduced sparsely at ZeRO stage 0 | unchanged
|
| checkpoints carried a `sparse_tensor_module_names` field | no longer
written |
Presence of the key raises regardless of its value, matching how #8490
treats
`mics_hierarchical_params_gather`, also a boolean whose default is
`False`.
A checkpoint written by an older release still loads. Its leftover
`sparse_tensor_module_names` key is no longer recognized, so it now
appears in the
`client_state` returned by `load_checkpoint` rather than being filtered
out.
## Test plan
Run with `slurm_scripts/test_deprecate_sparse_gradients.slurm` (1 node,
2 GPUs). The job
asserts that the branch under test is the tree actually imported before
running anything.
- [ ] `tests/unit/runtime/test_ds_config_dict.py` — the new
`test_sparse_gradients_config_is_rejected`, plus the four existing
removed-key tests
as a regression check on `config.py`
- [ ] `tests/unit/runtime/sparse_tensor/` — sparse reduction at ZeRO
stage 0 still correct
with the config key gone
- [ ] `tests/unit/checkpoint/test_latest_checkpoint.py`,
`tests/unit/checkpoint/test_lr_scheduler.py` — cover the edited
`tests/unit/checkpoint/common.py` helper
- [ ] CI unit / accelerator workflows on this PR
## Diffstat
```
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 ----------------
tests/unit/runtime/sparse_tensor/test_csr.py | 55 -------------
.../sparse_tensor/test_averaging_sparse_gradients.py | 4 +-
tests/unit/runtime/sparse_tensor/test_sparse_grads.py | 4 +-
tests/unit/runtime/test_ds_config_dict.py | 10 +++
13 files changed, 34 insertions(+), 229 deletions(-)
```
Two deleted test files asserted only on removed members:
`test_sparse.py` on
`sparse_tensor_module_names`, and `test_csr.py` on the dense constructor
and `to_dense()`.
Signed-off-by: pengdurice <pengduhit@gmail.com>
