[Cherry-pick] Add is_distributed field in sharding reshard param_meta (…

…PaddlePaddle#9028)
Mangodadada · Sep 10, 2024 · ee0446e · ee0446e
1 parent 542b0b2
commit ee0446e
Showing 1 changed file with 2 additions and 1 deletion.
diff --git a/paddlenlp/trainer/utils/sharding_io.py b/paddlenlp/trainer/utils/sharding_io.py
@@ -584,7 +584,8 @@ def _gather_sharding_metas(self):
         param_meta = {}
         for k, v in model.state_dict().items():
             structure_name_mapping[k] = v.name
-            param_meta[k] = (v.shape, int(v.dtype))
+            is_distributed = getattr(v, "is_distributed", False)
+            param_meta[k] = (v.shape, int(v.dtype), is_distributed)
 
         sharding_metas = {}
         sharding_meta = {}