-
Notifications
You must be signed in to change notification settings - Fork 4.2k
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.
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
create mininal universal checkpoint info for client state #5526
Changes from 1 commit
File filter
Filter by extension
Conversations
Jump to
Diff view
Diff view
There are no files selected for viewing
Original file line number | Diff line number | Diff line change |
---|---|---|
|
@@ -59,6 +59,7 @@ | |
WEIGHT_QUANTIZE_VERBOSE, \ | ||
WEIGHT_QUANTIZE_KERNEL | ||
from deepspeed.checkpoint.constants import OPTIMIZER_STATE_DICT, FROZEN_PARAM_FRAGMENTS | ||
from deepspeed.checkpoint.utils import inject_universal_info | ||
from deepspeed.runtime.sparse_tensor import SparseTensor | ||
|
||
from deepspeed.runtime import lr_schedules | ||
|
@@ -3319,6 +3320,7 @@ def _save_checkpoint(self, save_dir, tag, client_state={}, exclude_frozen_parame | |
ds_config=self.config, | ||
ds_version=version) | ||
state.update(client_state) | ||
inject_universal_info(state) | ||
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. Shoudn't we show a warning when we don't have the necessary info? There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. After I discussed with Tunji, we are considering another approach for this. He will share a new approach. I keep this comment but just disregard it. There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. @xylian86, the new approach is to do the injection in the conversion script rather than during saving. The injection should be done into |
||
|
||
if self.save_non_zero_checkpoint: | ||
log_dist(message=f'Saving model checkpoint: {save_path}', ranks=[0, 1]) | ||
|
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
FYI @xylian86 - can you run the precommit formatter on this branch so it will pass our Formatting check?