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 the checkpoint saving issues when zero3 is enabled #99

Closed
wants to merge 2 commits into from
Closed

Fix the checkpoint saving issues when zero3 is enabled #99

Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
88f0bb1
update to save zero3 ckpt
lyxok1 Dec 5, 2023
bccb6ad
fix zero3 saving issue; in compatible with transformers
lyxok1 Dec 9, 2023
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
24 changes: 16 additions & 8 deletions chatllms/utils/model_utils.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -11,6 +11,7 @@
from transformers.trainer_utils import get_last_checkpoint
from transformers.generation.logits_process import LogitsProcessor
from transformers.generation.utils import LogitsProcessorList
from transformers.deepspeed import is_deepspeed_zero3_enabled
from chatllms.data.data_utils import (DEFAULT_BOS_TOKEN, DEFAULT_EOS_TOKEN,
DEFAULT_PAD_TOKEN, DEFAULT_UNK_TOKEN)

Expand Down Expand Up @@ -308,13 +309,20 @@ def get_logits_processor() -> LogitsProcessorList:
return logits_processor



def safe_save_model_for_hf_trainer(trainer: Trainer, output_dir: str):
"""Collects the state dict and dump to disk."""
state_dict = trainer.model.state_dict()
if trainer.args.should_save:
cpu_state_dict = {
key: value.cpu()
for key, value in state_dict.items()
}
del state_dict
trainer._save(output_dir, state_dict=cpu_state_dict) # noqa
trainer.save_model(output_dir)

# state_dict = trainer.model.state_dict()
# if not is_deepspeed_zero3_enabled() and trainer.args.should_save:
# cpu_state_dict = {
# key: value.cpu()
# for key, value in state_dict.items()
# }
# del state_dict
# trainer._save(output_dir, state_dict=cpu_state_dict) # noqa
# elif is_deepspeed_zero3_enabled():
# # save for deepspeed ZeRO3 checkpoint
# if not trainer.wrapped_model.save_16bit_model(output_dir):
# trainer.wrapped_model.save_checkpoint(output_dir, save_latest=True)