From e6bf09db808de87371d73d5297f4c8aa98f56abe Mon Sep 17 00:00:00 2001
From: zanussbaum <zanussbaum@gmail.com>
Date: Sun, 26 Feb 2023 22:31:20 -0500
Subject: [PATCH] fix: count params when zero init'd

---
 src/peft/peft_model.py | 7 ++++++-
 1 file changed, 6 insertions(+), 1 deletion(-)

diff --git a/src/peft/peft_model.py b/src/peft/peft_model.py
index be08ed2..4703059 100644
--- a/src/peft/peft_model.py
+++ b/src/peft/peft_model.py
@@ -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(