Fix missing params in unconsolidated models

facebookresearch · Sep 19, 2023 · b5455b9 · b5455b9
1 parent 164cc0f
commit b5455b9
Showing 1 changed file with 4 additions and 3 deletions.
diff --git a/fairscale/nn/data_parallel/fully_sharded_data_parallel.py b/fairscale/nn/data_parallel/fully_sharded_data_parallel.py
@@ -2235,9 +2235,10 @@ def consolidate_shard_weights(
         # of synchronization between shards or that all shards buffers are equivalent).
         if with_module_buffers:
             for buffer_name in shard_metadata[0]["buffer_names"]:
-                if buffer_name not in shard_weights[0] and (not strict):
-                    continue
-                consolidated_weights[buffer_name] = shard_weights[0][buffer_name]
+                if buffer_name in shard_weights[0]:
+                    consolidated_weights[buffer_name] = shard_weights[0][buffer_name]
+                else:
+                    logging.info(f"Missing buffer weights: {buffer_name}")
 
         return consolidated_weights