fix(modeling_base): partial loading of a sharded checkpoint (#578)

CarperAI · Oct 26, 2023 · 91a0f43 · 91a0f43
1 parent 0575ce5
commit 91a0f43
Showing 1 changed file with 8 additions and 6 deletions.
diff --git a/trlx/models/modeling_base.py b/trlx/models/modeling_base.py
@@ -240,11 +240,15 @@ def from_pretrained(  # noqa: max-complexity
                     )
                     logger.info("Trained peft adapter loaded")
 
+            # No peft
             if base_model is None:
-                # No peft
+                # Disable warnings about missing weights when loading the base model
+                verbosity = transformers.logging.get_verbosity()
+                transformers.logging.set_verbosity_error()
                 base_model = cls._auto_model_parent_class.from_pretrained(
                     pretrained_model_name_or_path, *model_args, **from_pretrained_kwargs
                 )
+                transformers.logging.set_verbosity(verbosity)
 
         elif isinstance(pretrained_model_name_or_path, transformers.PreTrainedModel):
             base_model = pretrained_model_name_or_path
@@ -288,11 +292,9 @@ def from_pretrained(  # noqa: max-complexity
                         )
                     with open(index_file_name, "r") as f:
                         index = json.load(f)
-                    # Collect files containing weights from supported modules
-                    files_to_download = set()
-                    for k, v in index["weight_map"].items():
-                        if any([module in k for module in cls._supported_modules]):
-                            files_to_download.add(v)
+
+                    # Load all weights from the shards
+                    files_to_download = set(index["weight_map"].values())
                     is_sharded = True
 
             if is_sharded: