bigscience-workshop / bigscience-workshop/Megatron-DeepSpeed
[new feature] add a `tag` to the arguments to load the checkpoint from a specific step (not necessarly the latest)
- Dominant language
- Python
- Stars
- 1.4k
- Forks
- 226
- PR merge metrics
- No merged PRs in 30d
Description
It might be useful to be able to add a new argument to load the checkpoint from a specific step. Currently, the only way to do it - to the best of my knowledge - is to change manually the `latest` file inside the folder containing all the checkpoints.
Some thoughts shared by @thomasw21 to do it:
- add a `tag` argument in the `arguments.py` file and use it at [this line](https://github.com/bigscience-workshop/Megatron-DeepSpeed/blob/efefd8a41a4a0fde621e3cf7fa64548a106c68d9/megatron/checkpointing.py#L276)
- see the corresponding docstring of the `load_checkpoint` DeepSpeed's method [here](https://github.com/microsoft/DeepSpeed/blob/89b8c8872bad3f3e11a846e97d4cb1fb796e554d/deepspeed/runtime/engine.py#L2382-L2388)
Contributor guide
No contributing guide indexed for this repository
Assessment
This issue has not been assessed yet.