fix galore layerwise with frozen params (#29743)

a1a74541 · peterjc123 · GitHub · 8692aa88 · a1a74541 · a1a74541
Unverified Commit a1a74541 authored Mar 20, 2024 by peterjc123 Committed by GitHub Mar 20, 2024
Hide whitespace changes
Inline Side-by-side

Showing with 4 additions and 2 deletions

src/transformers/optimization.py src/transformers/optimization.py +2 -1

src/transformers/trainer.py src/transformers/trainer.py +2 -1

No files found.
--- a/src/transformers/optimization.py
+++ b/src/transformers/optimization.py
@@ -385,7 +385,8 @@ def get_scheduler(
                scheduler_dict[param].step()
        for param in optimizer_dict.keys():
-            param.register_post_accumulate_grad_hook(scheduler_hook)
+            if param.requires_grad:
+                param.register_post_accumulate_grad_hook(scheduler_hook)
        return LayerWiseDummyScheduler()

--- a/src/transformers/trainer.py
+++ b/src/transformers/trainer.py
@@ -1303,7 +1303,8 @@ class Trainer:
                        optimizer_dict[param].zero_grad()
                for param in model.parameters():
-                    param.register_post_accumulate_grad_hook(optimizer_hook)
+                    if param.requires_grad:
+                        param.register_post_accumulate_grad_hook(optimizer_hook)
                optimizer_cls = LayerWiseDummyOptimizer
                optimizer_kwargs.update({"optimizer_dict": optimizer_dict})