Skip to content

fix(trainer): clip gradients before optimizer.step() so --max-grad-norm takes effect#1384

Open
valter-silva-au wants to merge 1 commit into
awslabs:mainfrom
valter-silva-au:fix/graphstorm-max-grad-norm-clip-after-step
Open

fix(trainer): clip gradients before optimizer.step() so --max-grad-norm takes effect#1384
valter-silva-au wants to merge 1 commit into
awslabs:mainfrom
valter-silva-au:fix/graphstorm-max-grad-norm-clip-after-step

Conversation

@valter-silva-au

Copy link
Copy Markdown

Summary

Fixes #1366.

--max-grad-norm is meant to clip gradients for training stability, but every
trainer calls th.nn.utils.clip_grad_norm_ after optimizer.step():

self.optimizer.zero_grad()
loss.backward()
rt_profiler.record('train_backward')
self.optimizer.step()          # <-- optimizer already used the un-clipped grads
rt_profiler.record('train_step')

if max_grad_norm is not None:
    th.nn.utils.clip_grad_norm_(model.parameters(), max_grad_norm, grad_norm_type)

Since optimizer.step() has already consumed the gradients to update the
parameters, 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-norm silently does nothing.

The correct contract is backward → clip → step, so the optimizer applies the
clipped gradients.

Changes

Moved the clip block to run immediately after loss.backward() and before
optimizer.step(). The issue names np_trainer, but the same wrong ordering is
present in all five trainers, so I've fixed them together:

  • python/graphstorm/trainer/np_trainer.py
  • python/graphstorm/trainer/ep_trainer.py
  • python/graphstorm/trainer/lp_trainer.py
  • python/graphstorm/trainer/mt_trainer.py
  • python/graphstorm/trainer/glem_np_trainer.py

This 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 user
explicitly sets the flag, where it was already broken.

Scope note for maintainers: the issue only mentions np_trainer. I scoped
the fix to all five trainers since the bug is identical in each, but I'm happy
to narrow this PR to just np_trainer (or np/ep/lp) if you'd prefer a smaller
change — just let me know.

Testing

Added tests/unit-tests/test_trainer.py::test_node_trainer_grad_clip_before_step:
it spies on clip_grad_norm_ and GSOptimizer.step (both call through to the
originals so fit() runs for real), runs trainer.fit(num_epochs=1, max_grad_norm=1.0) on the existing tiny CPU/gloo dummy graph, and asserts the
clip is observed before the step on every iteration.

  • Fails on the current ordering (records ['step', 'clip']).
  • Passes after the reorder.
  • Full test_trainer.py suite: 17 passed (CPU/gloo, world_size=1).
  • pylint --rcfile=tests/lint/pylintrc python/graphstorm/trainer/: 10.00/10.

…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
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

possible max_grad_norm implementation error

1 participant