Unverified Commit 645f1742 authored by Sylvain Gugger's avatar Sylvain Gugger Committed by GitHub
Browse files

Exit early in load if no weights are in the sharded state dict (#18937)

parent 660e0b97
......@@ -418,6 +418,9 @@ def _load_state_dict_into_model(model_to_load, state_dict, start_prefix):
def load(module: nn.Module, state_dict, prefix=""):
local_metadata = {} if metadata is None else metadata.get(prefix[:-1], {})
args = (state_dict, prefix, local_metadata, True, [], [], error_msgs)
# Parameters of module and children will start with prefix. We can exit early if there are none in this
# state_dict
if len([key for key in state_dict if key.startswith(prefix)]) > 0:
if is_deepspeed_zero3_enabled():
import deepspeed
......
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment