fix: count params when zero init'd

2026-09-10 12:20:21 +08:00 · 2023-02-26 22:31:20 -05:00
parent 681ce93cc1
commit e6bf09db80
1 changed files with 6 additions and 1 deletions
@@ -267,7 +267,12 @@ class PeftModel(PushToHubMixin, torch.nn.Module):
        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(