Drop the documented grad_hooks ZeRO option, which does not exist

[vineethsaivs]

Problem

config-json.md documents grad_hooks as a ZeRO option with a default of True:

grad_hooks: [boolean]
For use with ZeRO stage 1, enable backward hooks to reduce gradients during the backward pass or wait until the end of the backward pass. Default True

DeepSpeedZeroConfig has no such field, and DeepSpeedConfigModel sets extra="forbid", so anyone who follows the docs gets a hard failure out of deepspeed.initialize:

DeepSpeedZeroConfig(**{"stage": 1, "grad_hooks": False})

pydantic_core._pydantic_core.ValidationError: 1 validation error for DeepSpeedZeroConfig
grad_hooks
  Extra inputs are not permitted [type=extra_forbidden, input_value=False, input_type=bool]

The option was never wired up rather than removed later, so there is nothing to restore. cfa63f5da ("ZeRO stage 1 refresh", #1042, thanks @ebarkhordar for pinning it down) added the doc entry and DeepSpeedEngine.zero_grad_hooks() in the same commit and never added the field to deepspeed/runtime/zero/config.py or to the ZeRO constants module. It is absent from both in every release back to v0.3.0. zero_grad_hooks() reads self._config.zero_config.grad_hooks, which can only raise AttributeError, and nothing in the package or the tests calls it.

Fix

Remove the doc entry and the dead accessor. If the intent is that the option should exist, that is a feature rather than a fix and I would rather leave it to you than invent a semantic for it.

Test

No test. An earlier revision of this PR added a guard to tests/unit/runtime/zero/test_zero_config.py that checked every ZeRO key config-json.md documents is one DeepSpeedZeroConfig accepts. @tohtana asked for it to be dropped as too fragile for CI, which is fair, so it is gone and this is now a two file deletion.

What was checked instead, on CPU:

• The ValidationError above reproduces on master and the doc entry is what invites it.
• An ast sweep over every .py in the tree finds zero_grad_hooks exactly once, its own definition at deepspeed/runtime/engine.py:1290, and .grad_hooks exactly once, on the line inside it. Removing it cannot break a caller, since any call raises AttributeError today.
• python -m pytest tests/unit/runtime/zero/test_zero_config.py gives 6 passed before and after, unchanged, and that file is now untouched by this PR.
• yapf --style .style.yapf and flake8 --config .flake8 are clean on deepspeed/runtime/engine.py, and clean on the unmodified tree as a control.

How it was found

Cross-checking every <i>**key**</i> in config-json.md against a grep for that literal in deepspeed/. Of 126 documented keys only two came back unreferenced: Compression, which is a section name, and this one.

For scope rather than as a request: the drift runs the other way too, with 28 fields DeepSpeedZeroConfig declares having no entry in that doc section. That direction misleads rather than crashes, so it is left alone here.

[ebarkhordar]

The central claim holds, and the history makes it stronger than the body puts it. grad_hooks was not removed at some point, it was never wired up. cfa63f5da ("ZeRO stage 1 refresh", #1042) added the doc entry and zero_grad_hooks() in the same commit, and added the field to neither the ZeRO config nor the constants module:

$ git show cfa63f5da:deepspeed/runtime/zero/config.py    | grep -i grad_hook   # no output
$ git show cfa63f5da:deepspeed/runtime/zero/constants.py | grep -i grad_hook   # no output
$ git show cfa63f5da -- deepspeed/runtime/engine.py docs/_pages/config-json.md | grep grad_hooks
+    def zero_grad_hooks(self):
+        return self._config.zero_config.grad_hooks
+<i>**grad_hooks**</i>: [boolean]

Naming that commit in the body would answer the "did we break this, and should we restore it instead" question up front.

The dead-accessor claim holds too. An ast sweep over every .py in the master tree finds zero_grad_hooks exactly once, its own definition at deepspeed/runtime/engine.py:1290, and .grad_hooks exactly once, on the line inside it. Removing it cannot break a caller, since any call raises AttributeError: 'DeepSpeedZeroConfig' object has no attribute 'grad_hooks' today.

I ran your new guard, because CI has not. On this head sha cpu-torch-latest, Formatting, python, nv-pre-compile-ops and DCO / required are all sitting at action_required, so modal-torch-latest is the only leg that has executed. In a clean python:3.11-slim container with torch 2.13.0+cpu and pydantic 2.13.4:

$ python -m pytest tests/unit/runtime/zero/test_zero_config.py -q
7 passed

It also bites, which is the part worth knowing: re-adding the grad_hooks doc block and rerunning gives AssertionError: documented but not accepted by DeepSpeedZeroConfig: ['grad_hooks']. Your parser finds 22 keys against its floor of 15, and the six stage3_* entries among them are accepted through their pydantic aliases rather than as declared fields, so filtering on extra_forbidden is the right call there.

For scope, not a request: the drift runs both ways. 28 fields DeepSpeedZeroConfig declares have no entry in that doc section, zenflow, leaf_module, sub_group_size and mics_shard_size among them. This PR guards the direction that actually breaks users, which seems like the right half to take on here.

[tohtana]

Hi @vineethsaivs,
Thank you for this PR! It is good to remove grad_hooks.

The test could be useful to prevent regressions, but it's a bit fragile and can cause false failures. I think this type of validation shouldn't run in CI. Would it be okay to remove the test?

[vineethsaivs]

Thanks @tohtana, that is a fair call and the test is gone. Pushed as 0449b34.

You are right about the fragility. The guard parsed config-json.md with a regex, so a change to the doc format would have broken it in CI on a PR that had nothing to do with ZeRO config, and the person paying for that would not have been the person who caused it. The floor assertion I added to catch that case turns a silent no-op into a red build, which is not obviously the better failure of the two.

The PR is now a two file deletion: the doc entry and DeepSpeedEngine.zero_grad_hooks().

@ebarkhordar, thank you for finding cfa63f5da, that answers the question I could not. I had only established that grad_hooks was absent from config.py back to v0.3.0, which left open whether it had been dropped by accident and should be restored instead. Since #1042 added the doc entry and the accessor in the same commit and never added the field, it was never wired up, so there is nothing to restore and deleting is the right direction. I have put that in the PR body and the commit message.

Both of your verification notes match what I get, including the documented but not accepted: ['grad_hooks'] failure and the six stage3_* aliases, so at least the removed test agreed with itself while it existed.

[tohtana]

Thank you for the update!