Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Fix loading with only state dict and low_cpu_mem_usage = True #35217

Merged
merged 7 commits into from
Dec 18, 2024
Merged

Fix loading with only state dict and low_cpu_mem_usage = True #35217

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
7 commits
Select commit Hold shift + click to select a range
ce86779
fix loading with only state dict and config
SunMarc Dec 11, 2024
b352074
style
SunMarc Dec 11, 2024
739c478
Merge branch 'main' into load-model-state-dict
sayakpaul Dec 12, 2024
84eb5d7
add tests
SunMarc Dec 12, 2024
55e1865
Merge remote-tracking branch 'origin/load-model-state-dict' into load…
SunMarc Dec 12, 2024
e47af0b
Merge branch 'main' into load-model-state-dict
sayakpaul Dec 13, 2024
f1f8a0f
Merge branch 'main' into load-model-state-dict
sayakpaul Dec 18, 2024
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
9 changes: 6 additions & 3 deletions src/transformers/modeling_utils.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -4022,8 +4022,11 @@ def from_pretrained(
loaded_state_dict_keys = sharded_metadata["all_checkpoint_keys"]
else:
loaded_state_dict_keys = list(state_dict.keys())

if gguf_path is None and (low_cpu_mem_usage or (use_keep_in_fp32_modules and is_accelerate_available())):
if (
gguf_path is None
and (low_cpu_mem_usage or (use_keep_in_fp32_modules and is_accelerate_available()))
and pretrained_model_name_or_path is not None
):
# In case some weights need to be kept in float32 and accelerate is not installed,
# we later on want to take the path where state_dict is not None, that is the one
# that do not require accelerate.
Expand Down Expand Up @@ -4679,7 +4682,7 @@ def _find_mismatched_keys(
)

# For GGUF models `state_dict` is never set to None as the state dict is always small
if gguf_path:
if gguf_path or low_cpu_mem_usage:
fixed_state_dict = cls._fix_state_dict_keys_on_load(state_dict)
error_msgs, offload_index, state_dict_index = _load_state_dict_into_meta_model(
model_to_load,
Expand Down
20 changes: 20 additions & 0 deletions tests/utils/test_modeling_utils.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -1750,6 +1750,26 @@ def test_save_and_load_config_with_custom_generation(self):
new_model.generate(random_ids, max_new_tokens=3)
self.assertTrue(len(w) == 0)

def test_load_model_with_state_dict_only(self):
model = BertModel.from_pretrained("hf-internal-testing/tiny-random-bert")
state_dict = model.state_dict()
config = model.config

model_loaded = BertModel.from_pretrained(
pretrained_model_name_or_path=None, config=config, state_dict=state_dict
)
self.assertTrue(check_models_equal(model, model_loaded))

def test_load_model_with_state_dict_only_low_cpu_mem_usage(self):
model = BertModel.from_pretrained("hf-internal-testing/tiny-random-bert")
state_dict = model.state_dict()
config = model.config

model_loaded = BertModel.from_pretrained(
pretrained_model_name_or_path=None, config=config, state_dict=state_dict, low_cpu_mem_usage=True
)
self.assertTrue(check_models_equal(model, model_loaded))


@slow
@require_torch
Expand Down
Loading