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.

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

[Flash Checkpoint] Integrated flash checkpoint #155

Open
wants to merge 1 commit into
base: main
Choose a base branch
from

Conversation

Caozhou1995
Copy link
Collaborator

@Caozhou1995 Caozhou1995 commented Jul 14, 2024

This PR integrated flash checkpoint, referring to dlrover.
Dlrover's core capabilities are invoked, but FlagScale customizes file paths, data organization, and so on.
To use the save and load functions of flash checkpoint, see the following example:
from flagscale.train.checkpointing import save_checkpoint
from flagscale.train.checkpointing import load_checkpoint
Interfaces are compatible with megatron and dlrover, the results are shown below:

  • save
    img_v3_02cp_a5618315-15a9-4d48-9f7d-3d81a7c2b84g
    img_v3_02cp_9c9974ae-aaf4-4f0b-84de-d1695300d1dg

  • load
    img_v3_02cp_c2bb102a-ca68-4951-8bfc-02c521d0f60g
    img_v3_02cp_e0db12da-867d-44fd-ae00-a865780697bg

If you don't want to use flash checkpoint, just set flash=False, the save and load functions provided by megatron are used.

@aoyulong aoyulong changed the title 【Flash Checkpoint】Integrated flash checkpoint [Flash Checkpoint] Integrated flash checkpoint Jul 14, 2024
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Labels
None yet
Projects
None yet
Development

Successfully merging this pull request may close these issues.

1 participant