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)

Open
#279 0 comments 1 reaction 0 assignees View on GitHub
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.

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.