From 5804413c4c06eb598d0fb6c4071f5c7e96214892 Mon Sep 17 00:00:00 2001 From: Valter Silva Date: Sat, 27 Jun 2026 14:26:11 +0800 Subject: [PATCH] [Bug] Clip gradients before optimizer.step() so --max-grad-norm takes effect MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit 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 #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 #1366 --- python/graphstorm/trainer/ep_trainer.py | 4 +- python/graphstorm/trainer/glem_np_trainer.py | 4 +- python/graphstorm/trainer/lp_trainer.py | 4 +- python/graphstorm/trainer/mt_trainer.py | 4 +- python/graphstorm/trainer/np_trainer.py | 4 +- tests/unit-tests/test_trainer.py | 84 ++++++++++++++++++++ 6 files changed, 94 insertions(+), 10 deletions(-) diff --git a/python/graphstorm/trainer/ep_trainer.py b/python/graphstorm/trainer/ep_trainer.py index 8e1697a812..76e20564ef 100644 --- a/python/graphstorm/trainer/ep_trainer.py +++ b/python/graphstorm/trainer/ep_trainer.py @@ -227,11 +227,11 @@ def fit(self, train_loader, num_epochs, self.optimizer.zero_grad() loss.backward() rt_profiler.record('train_backward') + if max_grad_norm is not None: + th.nn.utils.clip_grad_norm_(model.parameters(), max_grad_norm, grad_norm_type) self.optimizer.step() 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) self.log_metric("Train loss", loss.item(), total_steps) if i % 20 == 0 and get_rank() == 0: diff --git a/python/graphstorm/trainer/glem_np_trainer.py b/python/graphstorm/trainer/glem_np_trainer.py index 939bdd21e3..ccbc3939c2 100644 --- a/python/graphstorm/trainer/glem_np_trainer.py +++ b/python/graphstorm/trainer/glem_np_trainer.py @@ -238,11 +238,11 @@ def _prepare_batch(input_nodes, seeds, blocks, is_labeled=True): self.optimizer.zero_grad(optimize_sparse_params=module.training_sparse_embed) loss.backward() profiler.record('train_backward') + if max_grad_norm is not None: + th.nn.utils.clip_grad_norm_(model.parameters(), max_grad_norm, grad_norm_type) self.optimizer.step(optimize_sparse_params=module.training_sparse_embed) 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) self.log_metric("Train loss", loss.item(), total_steps) if i % 20 == 0 and get_rank() == 0: diff --git a/python/graphstorm/trainer/lp_trainer.py b/python/graphstorm/trainer/lp_trainer.py index 97778716f7..0ccfe49737 100644 --- a/python/graphstorm/trainer/lp_trainer.py +++ b/python/graphstorm/trainer/lp_trainer.py @@ -231,11 +231,11 @@ def fit(self, train_loader, num_epochs, self.optimizer.zero_grad() loss.backward() rt_profiler.record('train_backward') + if max_grad_norm is not None: + th.nn.utils.clip_grad_norm_(model.parameters(), max_grad_norm, grad_norm_type) self.optimizer.step() 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) self.log_metric("Train loss", loss.item(), total_steps) if i % 20 == 0 and get_rank() == 0: rt_profiler.print_stats() diff --git a/python/graphstorm/trainer/mt_trainer.py b/python/graphstorm/trainer/mt_trainer.py index 2befa17cd3..307d85298b 100644 --- a/python/graphstorm/trainer/mt_trainer.py +++ b/python/graphstorm/trainer/mt_trainer.py @@ -443,11 +443,11 @@ def fit(self, train_loader, self.optimizer.zero_grad() loss.backward() rt_profiler.record('train_backward') + if max_grad_norm is not None: + th.nn.utils.clip_grad_norm_(model.parameters(), max_grad_norm, grad_norm_type) self.optimizer.step() 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) self.log_metric("Train loss", loss.item(), total_steps) if i % 20 == 0 and get_rank() == 0: diff --git a/python/graphstorm/trainer/np_trainer.py b/python/graphstorm/trainer/np_trainer.py index 8e5c4f8772..42ce8ea2e1 100644 --- a/python/graphstorm/trainer/np_trainer.py +++ b/python/graphstorm/trainer/np_trainer.py @@ -214,11 +214,11 @@ def fit(self, self.optimizer.zero_grad() loss.backward() rt_profiler.record('train_backward') + if max_grad_norm is not None: + th.nn.utils.clip_grad_norm_(model.parameters(), max_grad_norm, grad_norm_type) self.optimizer.step() 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) self.log_metric("Train loss", loss.item(), total_steps) if i % 20 == 0 and get_rank() == 0: diff --git a/tests/unit-tests/test_trainer.py b/tests/unit-tests/test_trainer.py index 71c10d9529..22ca47773a 100644 --- a/tests/unit-tests/test_trainer.py +++ b/tests/unit-tests/test_trainer.py @@ -1042,6 +1042,90 @@ def test_rgcn_nc4ef(): th.distributed.destroy_process_group() dgl.distributed.kvstore.close_kvstore() +def test_node_trainer_grad_clip_before_step(): + """ Test that ``--max-grad-norm`` clips gradients *before* the optimizer step. + + Gradient clipping (``th.nn.utils.clip_grad_norm_``) must run between + ``loss.backward()`` and ``optimizer.step()``; otherwise the optimizer applies + the un-clipped gradients and the clip becomes a no-op for that update (issue + #1366). This test spies on ``clip_grad_norm_`` and ``GSOptimizer.step`` and + asserts the clip is observed before the step on every training iteration. + """ + from graphstorm.model import GSOptimizer + + # initialize the torch distributed environment + th.distributed.init_process_group(backend='gloo', + init_method='tcp://127.0.0.1:23456', + rank=0, + world_size=1) + + setup_device(0) + device = get_device() + + with tempfile.TemporaryDirectory() as tmpdirname: + _, part_config = generate_dummy_dist_graph(tmpdirname) + gdata = GSgnnData(part_config=part_config) + + create_config4ef(Path(tmpdirname), 'gnn_nc.yaml', use_ef=False) + args = Namespace(yaml_config_file=os.path.join(Path(tmpdirname), 'gnn_nc.yaml'), + local_rank=0) + config = GSConfig(args) + + model = create_builtin_node_gnn_model(gdata.g, config, True) + trainer = GSgnnNodePredictionTrainer(model) + trainer.setup_device(device) + + train_dataloader = GSgnnNodeDataLoader( + gdata, + target_idx=gdata.get_node_train_set(config.target_ntype), + fanout=config.fanout, + batch_size=config.batch_size, + label_field=config.label_field, + node_feats=config.node_feat_name, + edge_feats=None, + train_task=True) + + # Record the relative order in which clipping and the optimizer step are + # called; both spies call through to the originals so fit() runs for real. + call_order = [] + orig_clip = th.nn.utils.clip_grad_norm_ + orig_step = GSOptimizer.step + + def spy_clip(*args, **kwargs): + call_order.append("clip") + return orig_clip(*args, **kwargs) + + def spy_step(self, *args, **kwargs): + call_order.append("step") + return orig_step(self, *args, **kwargs) + + with patch.object(th.nn.utils, "clip_grad_norm_", side_effect=spy_clip), \ + patch.object(GSOptimizer, "step", autospec=True, side_effect=spy_step): + trainer.fit( + train_loader=train_dataloader, + num_epochs=1, + max_grad_norm=1.0) + + # At every point in the call sequence the number of observed clips must be + # >= the number of observed steps; this holds iff each clip precedes its + # paired step. On the buggy ordering the first event is a step, so the + # invariant fails immediately. + clips = steps = 0 + for event in call_order: + if event == "clip": + clips += 1 + else: + steps += 1 + assert clips >= steps, ( + "clip_grad_norm_ must be called before optimizer.step() on every " + f"iteration, but observed a step before its clip. Order: {call_order}") + assert steps > 0, "optimizer.step() was never called; the training loop did not run" + assert clips == steps, ( + f"expected one clip per step, got {clips} clips and {steps} steps") + + th.distributed.destroy_process_group() + dgl.distributed.kvstore.close_kvstore() + def test_rgat_nc4ef(): """ Test RGAT model Node Classification traning pipeline with/without edge features. """