fix: stop DeepSpeedConfig writing max_grad_norm back into the caller's config dict

[ebarkhordar]

What happens

deepspeed.initialize() writes into the dict the caller passed as config. When
optimizer.params.max_grad_norm is set to a positive value, DeepSpeedConfig._do_warning_check
assigns 0.0 into self.optimizer_params, and that is the caller's own
config["optimizer"]["params"] object rather than a copy.

Measured in a clean python:3.11-slim container at HEAD 11b518a00, torch 2.13.0+cpu,
deepspeed installed with pip install -e . from the checkout
(deepspeed.__file__ = /src/deepspeed/__init__.py, deepspeed.__version__ = 0.19.6+unknown):

import os, json, copy
os.environ.update(MASTER_ADDR="127.0.0.1", MASTER_PORT="29517",
                  RANK="0", LOCAL_RANK="0", WORLD_SIZE="1")
import torch, deepspeed

cfg = {"train_micro_batch_size_per_gpu": 1,
       "optimizer": {"type": "AdamW", "params": {"lr": 1e-3, "max_grad_norm": 1.0}}}
model = torch.nn.Linear(4, 4)
client_opt = torch.optim.AdamW(model.parameters(), lr=1e-3)

print("BEFORE:", json.dumps(cfg["optimizer"]["params"]))
engine, *_ = deepspeed.initialize(model=model, optimizer=client_opt, config=cfg)
print("AFTER :", json.dumps(cfg["optimizer"]["params"]))
print("gradient_clipping in force:", engine.gradient_clipping())

Observed:

BEFORE: {"lr": 0.001, "max_grad_norm": 1.0}
[WARNING] [config.py:1068:_do_warning_check] DeepSpeedConfig: In FP32 mode, DeepSpeed does not permit MAX_GRAD_NORM (1.0) > 0, setting to zero
AFTER : {"lr": 0.001, "max_grad_norm": 0.0}
gradient_clipping in force: 1.0

Expected: initialize leaves the caller's dict as it found it.

Passing a client optimizer is what makes this visible, because _configure_basic_optimizer is
then never called and the usual ValueError never fires, so initialization succeeds with the
caller's config quietly rewritten. The same run without a client optimizer still raises the
ValueError, and still leaves 0.0 behind in the caller's dict, because the zeroing happens
during config construction and the engine's check tests for the key's presence rather than its
value.

The warning is also no longer accurate. The value it claims to zero is not read by anything:
get_optimizer_gradient_clipping (config.py:458) is its only reader and has no callers
anywhere in deepspeed/ or tests/ (checked with an AST scan for Call nodes, not a text
search). The clipping actually applied comes from gradient_clipping, which is why the run
above reports 1.0. The engine side of this behaviour was removed in abe2204d (#232, 2020)
and replaced by the hard ValueError; the config side predates that change and was not
revisited with it.

Why the fix looks like this

_configure_basic_optimizer already declares the invariant this line breaks, at
engine.py:2088, added three months ago in 3c337b542 (#8010):

# Copy so the pop() calls below (torch_adam, adam_w_mode, fp32_optimizer_states) do not
# mutate the shared config dict returned by optimizer_params().
optimizer_parameters = dict(self.optimizer_params() or {})

Enumerating every writer of that dict across deepspeed/ by AST (subscript assignment plus
update/pop/setdefault/clear/popitem calls):

	writers
before	4: config.py:1071 on the caller's dict, plus engine.py:2096, 2097, 2115 on the copy
after	3: engine.py:2096, 2097, 2115, all on the copy

The FP16 and FP32 branches were a pair with the zeroing, so removing it leaves them without a
distinction to draw. Nothing passes max_grad_norm to an FP16 wrapper today: the only
consumers of a max_grad_norm param group are the Lamb and OneBit optimizers, which take it
as a constructor argument. The two branches therefore collapse into one warning carrying the
same remedy that _configure_basic_optimizer already raises.

If you would rather keep both original messages and drop only the assignment, or split the
message change into its own PR, say so and I will rework it.

Tests

tests/unit/runtime/test_ds_config_dict.py::test_max_grad_norm_leaves_caller_config_untouched
pins the caller's dict directly. In the same container, on master it fails with
assert 0.0 == 1.0; with this change it passes.

The rest of that file is unaffected: 27 passed, 5 skipped. TestArgs needs --shm-size above
the Docker default and fails with OSError: [Errno 28] No space left on device without it, on
master and on this branch alike.

yapf --style .style.yapf --diff and flake8 --config .flake8 are both clean on the two
changed files.

One limit worth stating: the unit test covers config construction, which is where the write
happens. The full deepspeed.initialize path is covered by the container run above rather
than by a unit test, since it needs a built comm extension.

[chatgpt-codex-connector]

💡 Codex Review

Here are some automated review suggestions for this pull request.

Reviewed commit: 675a8726b7

ℹ️ 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 sign-off trailer

This is a single-parent, non-merge commit, but its commit message has no Signed-off-by trailer. Add the author identity from git config user.name and git config user.email using --signoff so the commit satisfies the repository's mandatory commit policy.

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

Useful? React with 👍 / 👎.

[tohtana]

Good catch, thank you @ebarkhordar!