Deprecate sparse_gradients

[pengdurice]

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: Deprecate unused DeepSpeed features #8489

What the flag did, and why it goes

The flag's entire contribution was the dense-input branch of SparseTensor.__init__:

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:

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().

[chatgpt-codex-connector]

💡 Codex Review

Here are some automated review suggestions for this pull request.

Reviewed commit: 0bb2a5916f

ℹ️ About Codex in GitHub

Codex has been enabled to automatically review pull requests in this repo. Reviews are triggered when you

• Open a pull request for review
• Mark a draft as ready
• Comment "@codex review".

If Codex has suggestions, it will comment; otherwise it will react with 👍.

When you sign up for Codex through ChatGPT, Codex can also answer questions or update the PR, like "@codex address that feedback".

[chatgpt-codex-connector]

Add the required commit sign-off

This is a non-merge commit, but its message contains no Signed-off-by trailer, so it does not satisfy the repository's mandatory DCO/CI requirement. Add the author identity using --signoff before merging.

AGENTS.md reference: AGENTS.md:L8-L8

Useful? React with 👍 / 👎.