fix(trainer): clip gradients before optimizer.step() so --max-grad-norm takes effect#1384
Open
valter-silva-au wants to merge 1 commit into
Open
Conversation
…rm takes effect The `--max-grad-norm` option clips gradients via th.nn.utils.clip_grad_norm_, but every trainer called it *after* optimizer.step(). Since the optimizer has already consumed the (un-clipped) gradients to update the parameters, clipping afterward is a no-op for that update and the feature silently does nothing (issue awslabs#1366). Move the clip block to run between loss.backward() and optimizer.step() in all five trainers (np / ep / lp / mt / glem). This is a pure statement reorder: the `if max_grad_norm is not None` guard is preserved, so the default (max_grad_norm=None) path is unchanged and behavior changes only when the user explicitly sets the flag — where it was already broken. Add a CPU unit test that spies on clip_grad_norm_ and GSOptimizer.step and asserts clipping is observed before the step on every training iteration. The test fails on the old ordering and passes after the fix. Fixes awslabs#1366
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
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.
Summary
Fixes #1366.
--max-grad-normis meant to clip gradients for training stability, but everytrainer calls
th.nn.utils.clip_grad_norm_afteroptimizer.step():Since
optimizer.step()has already consumed the gradients to update theparameters, clipping afterward is a no-op for that update — the gradients are
zeroed at the start of the next iteration before they're ever used. The result
is that
--max-grad-normsilently does nothing.The correct contract is
backward → clip → step, so the optimizer applies theclipped gradients.
Changes
Moved the clip block to run immediately after
loss.backward()and beforeoptimizer.step(). The issue namesnp_trainer, but the same wrong ordering ispresent in all five trainers, so I've fixed them together:
python/graphstorm/trainer/np_trainer.pypython/graphstorm/trainer/ep_trainer.pypython/graphstorm/trainer/lp_trainer.pypython/graphstorm/trainer/mt_trainer.pypython/graphstorm/trainer/glem_np_trainer.pyThis is a pure statement reorder — no signature or API change. The
if max_grad_norm is not None:guard is preserved, so the default(
max_grad_norm=None) path is unchanged; behavior only changes when a userexplicitly sets the flag, where it was already broken.
Testing
Added
tests/unit-tests/test_trainer.py::test_node_trainer_grad_clip_before_step:it spies on
clip_grad_norm_andGSOptimizer.step(both call through to theoriginals so
fit()runs for real), runstrainer.fit(num_epochs=1, max_grad_norm=1.0)on the existing tiny CPU/gloo dummy graph, and asserts theclip is observed before the step on every iteration.
['step', 'clip']).test_trainer.pysuite: 17 passed (CPU/gloo,world_size=1).pylint --rcfile=tests/lint/pylintrc python/graphstorm/trainer/: 10.00/10.