create mininal universal checkpoint info for client state #5526
Conversation
tjruwase
requested review from
samadejacobs,
tohtana and
lekurile
and removed request for
mrwyattii
|
@xylian86, thanks for this great work. Can you please add convergence curves of an HF model as demo? |
| @@ -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.
Shoudn't we show a warning when we don't have the necessary info?
It will silently produce an incorrect checkpoint if the checkpoint is loaded for TP or PP.
We can say that the converted checkpoint is only for pure DP scaling.
There was a problem hiding this comment.
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.
@xylian86, the new approach is to do the injection in the conversion script rather than during saving. The injection should be done into ds_checkpoint before this assertion. Furthermore, the injection should be enabled by command-line argument (disabled by default) so that users are fully aware of what is going on. The command-line arg could be called --inject-missing-state.
| @@ -5,7 +5,7 @@ | |||
|
|
|||
| import os | |||
| import torch | |||
| from .constants import (MODEL_FILE_PREFIX, MODEL_FILE_SUFFIX, OPTIM_FILE_SUFFIX, ZERO_FILE_PREFIX) | |||
| from .constants import (MODEL_FILE_PREFIX, MODEL_FILE_SUFFIX, OPTIM_FILE_SUFFIX, ZERO_FILE_PREFIX, UNIVERSAL_CHECKPOINT_INFO, UNIVERSAL_CHECKPOINT_VERSION_KEY, UNIVERSAL_CHECKPOINT_VERSION_VALUE) | |||
There was a problem hiding this comment.
FYI @xylian86 - can you run the precommit formatter on this branch so it will pass our Formatting check?
pre-commit run --all-files
This PR solves the Issue-5430.
The PR enables the universal checkpoint feature for other platforms like HuggingFace Trainer without requiring changes to the HuggingFace code. It does this by creating a minimal universal checkpoint info, specifically the version, as a default action for the client state.