Support fp32 grad clipping and fix max_grad_norm confusion - #232
Merged
Merged
Conversation
tjruwase
approved these changes
May 26, 2020
samyam
reviewed
May 27, 2020
samyam
approved these changes
May 27, 2020
This was referenced Aug 20, 2026
banxingmjj
pushed a commit
to openanolis/DeepSpeed
that referenced
this pull request
Aug 23, 2026
…s config dict (deepspeedai#8289) ## 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`): ```python 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` (deepspeedai#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` (deepspeedai#8010): ```python # 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. Signed-off-by: Ehsan Barkhordar <realbarkhordar@gmail.com> Co-authored-by: Guokai Ma <guokai.ma@intel.com>
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
3 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
No description provided.