fix saving to mulitple shards

ymcui · May 8, 2023 · 11e5d1f · 11e5d1f
1 parent 08d2ee2
commit 11e5d1f
Showing 1 changed file with 6 additions and 2 deletions.
diff --git a/scripts/merge_llama_with_chinese_lora.py b/scripts/merge_llama_with_chinese_lora.py
@@ -134,8 +134,12 @@ def save_shards(model_sd, num_shards: int):
                         splits = v.split(v.size(1)//num_shards,dim=1)
                     elif new_k=='output.weight':
                         print(f"Processing {new_k}")
-                        splits = v.split(v.size(0)//num_shards,dim=0)
-
+                        if v.size(0)%num_shards==0:
+                            splits = v.split(v.size(0)//num_shards,dim=0)
+                        else:
+                            size_list = [v.size(0)//num_shards] * num_shards
+                            size_list[-1] += v.size(0)%num_shards
+                            splits = v.split(size_list, dim=0) # 13B: size_list == [24976,24977]
                     elif new_k=='norm.weight':
                         print(f"Processing {new_k}")
                         splits = [v] * num_shards