Fix

huggingface · Feb 21, 2025 · a59744d · a59744d
1 parent 2c2a101
commit a59744d
Showing 1 changed file with 12 additions and 2 deletions.
diff --git a/tests/distributed/test_common.py b/tests/distributed/test_common.py
@@ -227,7 +227,12 @@ def move_grads_to_cpu(parameters):
                     xm.mark_step()
 
                     if max_grad_norm is not None:
-                        accelerator.clip_grad_norm_(model.local_parameters(), max_norm=max_grad_norm, norm_type=2)
+                        accelerator.clip_grad_norm_(
+                            model.local_parameters(),
+                            max_norm=max_grad_norm,
+                            norm_type=2,
+                            postpone_clipping_to_optimizer_step=True,
+                        )
 
                     # Checking that at least some of the parameters have a gradient.
                     grads_on_cpu = move_grads_to_cpu(model.local_parameters())
@@ -259,7 +264,12 @@ def move_grads_to_cpu(parameters):
                     loss.backward()
 
                     if max_grad_norm is not None:
-                        accelerator.clip_grad_norm_(model.parameters(), max_norm=max_grad_norm, norm_type=2)
+                        accelerator.clip_grad_norm_(
+                            model.parameters(),
+                            max_norm=max_grad_norm,
+                            norm_type=2,
+                            postpone_clipping_to_optimizer_step=True,
+                        )
 
                     # Checking that at least some of the parameters have a gradient.
                     grads_on_cpu = move_grads_to_cpu(model.parameters())