Merge pull request huggingface#140 from zanussbaum/zero_init

fix: count params when zero init'd
dengdifan · Feb 27, 2023 · aa18556 · aa18556
2 parents 681ce93 + e6bf09d
commit aa18556
Showing 1 changed file with 6 additions and 1 deletion.
diff --git a/src/peft/peft_model.py b/src/peft/peft_model.py
@@ -267,7 +267,12 @@ def print_trainable_parameters(self):
         trainable_params = 0
         all_param = 0
         for _, param in self.named_parameters():
-            all_param += param.numel()
+            num_params = param.numel()
+            # if using DS Zero 3 and the weights are initialized empty
+            if num_params == 0 and hasattr(param, "ds_numel"):
+                num_params = param.ds_numel
+
+            all_param += num_params
             if param.requires_grad:
                 trainable_params += param.numel()
         print(