Skip to content

Commit

Permalink
Add max shard size to transformers save_pretrained
Browse files Browse the repository at this point in the history
  • Loading branch information
b-chu committed Nov 7, 2024
1 parent a16b93d commit 3dd899a
Showing 1 changed file with 2 additions and 2 deletions.
4 changes: 2 additions & 2 deletions llmfoundry/callbacks/hf_checkpointer.py
Original file line number Diff line number Diff line change
Expand Up @@ -743,7 +743,7 @@ def tensor_hook(
) if is_te_imported and state.precision == Precision.AMP_FP8 else contextlib.nullcontext(
)
with context_manager:
new_model_instance.save_pretrained(temp_save_dir)
new_model_instance.save_pretrained(temp_save_dir, max_shard_size='1GB')
if original_tokenizer is not None:
assert isinstance(
original_tokenizer,
Expand Down Expand Up @@ -799,7 +799,7 @@ def tensor_hook(
new_model_instance = self.transform_model_pre_registration(
new_model_instance,
)
new_model_instance.save_pretrained(register_save_dir)
new_model_instance.save_pretrained(register_save_dir, max_shard_size='1GB')
if original_tokenizer:
original_tokenizer.save_pretrained(register_save_dir)

Expand Down

0 comments on commit 3dd899a

Please sign in to comment.