[NeMo-UX] Support `save_last="link"` #10548

ashors1 · 2024-09-20T01:02:03Z

What does this PR do ?

Adds support for creating a symlink for -last checkpoints. Implementation is compatible with synchronous and asynchronous checkpointing.

Collection: llm

Changelog

Add specific line by line info of high level changes in this PR.

Usage

You can potentially add a usage example below

# Add a code snippet demonstrating how to use this

GitHub Actions CI

The Jenkins CI system has been replaced by GitHub Actions self-hosted runners.

The GitHub Actions CI will run automatically when the "Run CICD" label is added to the PR.
To re-run CI remove and add the label again.
To run CI on an untrusted fork, a NeMo user with write access must first click "Approve and run".

Before your PR is "Ready for review"

Pre checks:

Make sure you read and followed Contributor guidelines
Did you write any new necessary tests?
Did you add or update any necessary documentation?
Does the PR affect components that are optional to install? (Ex: Numba, Pynini, Apex etc)
- Reviewer: Does the PR have correct import guards for all optional libraries?

PR Type:

New Feature
Bugfix
Documentation

If you haven't finished some of the above items you can still open "Draft" PR.

Who can review?

Anyone in the community is free to review the PR once the checks have passed.
Contributor guidelines contains specific people who can review PRs to various areas.

Additional Information

Related to # (issue)

Signed-off-by: ashors1 <[email protected]>

Signed-off-by: Anna Shors <[email protected]>

Signed-off-by: ashors1 <[email protected]>

…t-ckpt

Signed-off-by: ashors1 <[email protected]>

nemo/lightning/pytorch/callbacks/model_checkpoint.py

mikolajblaz · 2024-09-23T13:37:36Z

nemo/lightning/pytorch/callbacks/model_checkpoint.py

                logging.info(f'Scheduled async checkpoint save for {filepath}')
            else:
                finalize_fn()

+    def _save_last_checkpoint(self, trainer: "pl.Trainer", monitor_candidates: Dict[str, torch.Tensor]) -> None:


I'm wondering if there is some way to avoid overriding the whole method. This is always risky, since we lose touch with the upstream.

How is our flow different from the one in PTL which makes us add saved_current_step logic and also not rely on self.last_model_path?
Is it because PTL links to any available last checkpoint (not necessarily from the last iteration)?

Yeah, I made two changes:

Made sure to add a symlink only when the current step was actually saved. As you suggested, PTL always links to the last checkpoint saved, which might not correspond to the latest step

Added these lines which fix last_model_path saved to the *-last checkpoint state dict when using symlinks

I'll think about whether we can make these fixes without overwriting the entire _save_last_checkpoint method

Thanks. Maybe overwriting save_last_checkpoint is inevitable in which case current version is ok

Running some final tests now, but I think I was able to avoid overwriting _save_last_checkpoint. Please let me know if you have any concerns with the current approach

Great, thanks, this is great.

Do you know how last_model_path is used during restart? I'm wondering if the loaded state dict will be valid if e.g. failure happens between the regular and "last" ckpt save

I think last_model_path is only used when removing the previous -last model to ensure we only retain a single -last checkpoint. If failure happens between the regular and last checkpoint save, I don't think the state dict will be valid, but I also don't think this is a concern, because we'd end up restoring from the previously saved -last checkpoint which does have the correct state dict.

nemo/lightning/pytorch/callbacks/model_checkpoint.py

Signed-off-by: ashors1 <[email protected]>

…o ashors/symlink-last-ckpt

Signed-off-by: ashors1 <[email protected]>

nemo/lightning/pytorch/callbacks/model_checkpoint.py

Signed-off-by: ashors1 <[email protected]>

…o ashors/symlink-last-ckpt

Signed-off-by: ashors1 <[email protected]>

tests/lightning/pytorch/callbacks/test_model_checkpoint.py

Signed-off-by: ashors1 <[email protected]>

tests/lightning/pytorch/callbacks/test_model_checkpoint.py

Signed-off-by: ashors1 <[email protected]>

…o ashors/symlink-last-ckpt

athitten

LGTM, thank you @ashors1

github-actions · 2024-09-28T16:07:19Z

[🤖]: Hi @ashors1 👋,

I just wanted to let you know that, you know, a CICD pipeline for this PR just finished successfully ✨

So it might be time to merge this PR or like to get some approvals 🚀

But I'm just a 🤖 so I'll leave it you what to do next.

Have a great day!

//cc @ko3n1g

* provide support for save_last='link' Signed-off-by: ashors1 <[email protected]> * fix symlinks when top_k checkpoint not saved Signed-off-by: ashors1 <[email protected]> * support symlinks with async checkpointing Signed-off-by: ashors1 <[email protected]> * only unlink on rank 0 Signed-off-by: Anna Shors <[email protected]> * fix race condition Signed-off-by: ashors1 <[email protected]> * force linked checkpoint to correspond to last finalized checkpoint Signed-off-by: ashors1 <[email protected]> * fix last_model_path after restore Signed-off-by: ashors1 <[email protected]> * move symlink removal to strategy Signed-off-by: ashors1 <[email protected]> * remove unneeded lines Signed-off-by: ashors1 <[email protected]> * add some more documentation Signed-off-by: ashors1 <[email protected]> * Apply isort and black reformatting Signed-off-by: ashors1 <[email protected]> * address some comments Signed-off-by: ashors1 <[email protected]> * fix syntax Signed-off-by: ashors1 <[email protected]> * avoid overwriting _save_last_checkpoint Signed-off-by: ashors1 <[email protected]> * fix base call Signed-off-by: ashors1 <[email protected]> * small fix Signed-off-by: ashors1 <[email protected]> * Apply isort and black reformatting Signed-off-by: ashors1 <[email protected]> * add test for save_last=link Signed-off-by: ashors1 <[email protected]> * Apply isort and black reformatting Signed-off-by: ashors1 <[email protected]> * clean up test Signed-off-by: ashors1 <[email protected]> * use megatroncheckpointio in test Signed-off-by: ashors1 <[email protected]> * add async test and clean up Signed-off-by: ashors1 <[email protected]> * fix remaining merge conflicts Signed-off-by: ashors1 <[email protected]> * Apply isort and black reformatting Signed-off-by: ashors1 <[email protected]> * check number of saved checkpoints Signed-off-by: ashors1 <[email protected]> * remove unused import Signed-off-by: ashors1 <[email protected]> * run test on gpu only Signed-off-by: ashors1 <[email protected]> * fix a small bug and add a resume test Signed-off-by: ashors1 <[email protected]> * Apply isort and black reformatting Signed-off-by: ashors1 <[email protected]> * remove old comment Signed-off-by: ashors1 <[email protected]> --------- Signed-off-by: ashors1 <[email protected]> Signed-off-by: Anna Shors <[email protected]> Signed-off-by: ashors1 <[email protected]> Co-authored-by: ashors1 <[email protected]>

* provide support for save_last='link' Signed-off-by: ashors1 <[email protected]> * fix symlinks when top_k checkpoint not saved Signed-off-by: ashors1 <[email protected]> * support symlinks with async checkpointing Signed-off-by: ashors1 <[email protected]> * only unlink on rank 0 Signed-off-by: Anna Shors <[email protected]> * fix race condition Signed-off-by: ashors1 <[email protected]> * force linked checkpoint to correspond to last finalized checkpoint Signed-off-by: ashors1 <[email protected]> * fix last_model_path after restore Signed-off-by: ashors1 <[email protected]> * move symlink removal to strategy Signed-off-by: ashors1 <[email protected]> * remove unneeded lines Signed-off-by: ashors1 <[email protected]> * add some more documentation Signed-off-by: ashors1 <[email protected]> * Apply isort and black reformatting Signed-off-by: ashors1 <[email protected]> * address some comments Signed-off-by: ashors1 <[email protected]> * fix syntax Signed-off-by: ashors1 <[email protected]> * avoid overwriting _save_last_checkpoint Signed-off-by: ashors1 <[email protected]> * fix base call Signed-off-by: ashors1 <[email protected]> * small fix Signed-off-by: ashors1 <[email protected]> * Apply isort and black reformatting Signed-off-by: ashors1 <[email protected]> * add test for save_last=link Signed-off-by: ashors1 <[email protected]> * Apply isort and black reformatting Signed-off-by: ashors1 <[email protected]> * clean up test Signed-off-by: ashors1 <[email protected]> * use megatroncheckpointio in test Signed-off-by: ashors1 <[email protected]> * add async test and clean up Signed-off-by: ashors1 <[email protected]> * fix remaining merge conflicts Signed-off-by: ashors1 <[email protected]> * Apply isort and black reformatting Signed-off-by: ashors1 <[email protected]> * check number of saved checkpoints Signed-off-by: ashors1 <[email protected]> * remove unused import Signed-off-by: ashors1 <[email protected]> * run test on gpu only Signed-off-by: ashors1 <[email protected]> * fix a small bug and add a resume test Signed-off-by: ashors1 <[email protected]> * Apply isort and black reformatting Signed-off-by: ashors1 <[email protected]> * remove old comment Signed-off-by: ashors1 <[email protected]> --------- Signed-off-by: ashors1 <[email protected]> Signed-off-by: Anna Shors <[email protected]> Signed-off-by: ashors1 <[email protected]> Co-authored-by: ashors1 <[email protected]> Signed-off-by: Lifu Zhang <[email protected]>

* provide support for save_last='link' Signed-off-by: ashors1 <[email protected]> * fix symlinks when top_k checkpoint not saved Signed-off-by: ashors1 <[email protected]> * support symlinks with async checkpointing Signed-off-by: ashors1 <[email protected]> * only unlink on rank 0 Signed-off-by: Anna Shors <[email protected]> * fix race condition Signed-off-by: ashors1 <[email protected]> * force linked checkpoint to correspond to last finalized checkpoint Signed-off-by: ashors1 <[email protected]> * fix last_model_path after restore Signed-off-by: ashors1 <[email protected]> * move symlink removal to strategy Signed-off-by: ashors1 <[email protected]> * remove unneeded lines Signed-off-by: ashors1 <[email protected]> * add some more documentation Signed-off-by: ashors1 <[email protected]> * Apply isort and black reformatting Signed-off-by: ashors1 <[email protected]> * address some comments Signed-off-by: ashors1 <[email protected]> * fix syntax Signed-off-by: ashors1 <[email protected]> * avoid overwriting _save_last_checkpoint Signed-off-by: ashors1 <[email protected]> * fix base call Signed-off-by: ashors1 <[email protected]> * small fix Signed-off-by: ashors1 <[email protected]> * Apply isort and black reformatting Signed-off-by: ashors1 <[email protected]> * add test for save_last=link Signed-off-by: ashors1 <[email protected]> * Apply isort and black reformatting Signed-off-by: ashors1 <[email protected]> * clean up test Signed-off-by: ashors1 <[email protected]> * use megatroncheckpointio in test Signed-off-by: ashors1 <[email protected]> * add async test and clean up Signed-off-by: ashors1 <[email protected]> * fix remaining merge conflicts Signed-off-by: ashors1 <[email protected]> * Apply isort and black reformatting Signed-off-by: ashors1 <[email protected]> * check number of saved checkpoints Signed-off-by: ashors1 <[email protected]> * remove unused import Signed-off-by: ashors1 <[email protected]> * run test on gpu only Signed-off-by: ashors1 <[email protected]> * fix a small bug and add a resume test Signed-off-by: ashors1 <[email protected]> * Apply isort and black reformatting Signed-off-by: ashors1 <[email protected]> * remove old comment Signed-off-by: ashors1 <[email protected]> --------- Signed-off-by: ashors1 <[email protected]> Signed-off-by: Anna Shors <[email protected]> Signed-off-by: ashors1 <[email protected]> Co-authored-by: ashors1 <[email protected]> Signed-off-by: Hainan Xu <[email protected]>

ashors1 added 11 commits September 10, 2024 12:52

provide support for save_last='link'

86ae11e

Signed-off-by: ashors1 <[email protected]>

fix symlinks when top_k checkpoint not saved

0f01d49

Signed-off-by: ashors1 <[email protected]>

support symlinks with async checkpointing

a953982

Signed-off-by: ashors1 <[email protected]>

only unlink on rank 0

b752074

Signed-off-by: Anna Shors <[email protected]>

fix race condition

3d62934

Signed-off-by: ashors1 <[email protected]>

Merge branch 'main' of github.com:NVIDIA/NeMo into ashors/symlink-las…

8ef9f8c

…t-ckpt

force linked checkpoint to correspond to last finalized checkpoint

ec9a92c

Signed-off-by: ashors1 <[email protected]>

fix last_model_path after restore

d72a838

Signed-off-by: ashors1 <[email protected]>

move symlink removal to strategy

4e7b3a0

Signed-off-by: ashors1 <[email protected]>

remove unneeded lines

cdcaeb7

Signed-off-by: ashors1 <[email protected]>

add some more documentation

ca1d6ee

Signed-off-by: ashors1 <[email protected]>

ashors1 requested review from hemildesai and mikolajblaz September 20, 2024 01:02

Apply isort and black reformatting

200cae3

Signed-off-by: ashors1 <[email protected]>

ashors1 added the Run CICD label Sep 20, 2024

mikolajblaz reviewed Sep 23, 2024

View reviewed changes

ashors1 and others added 7 commits September 23, 2024 09:19

address some comments

dec10ea

Signed-off-by: ashors1 <[email protected]>

Merge branch 'ashors/symlink-last-ckpt' of github.com:NVIDIA/NeMo int…

0cc6846

…o ashors/symlink-last-ckpt

fix syntax

9b213b7

Signed-off-by: ashors1 <[email protected]>

avoid overwriting _save_last_checkpoint

da26e98

Signed-off-by: ashors1 <[email protected]>

fix base call

f80f492

Signed-off-by: ashors1 <[email protected]>

small fix

bc322f6

Signed-off-by: ashors1 <[email protected]>

Apply isort and black reformatting

1aa6bf1

Signed-off-by: ashors1 <[email protected]>

athitten reviewed Sep 23, 2024

View reviewed changes

nemo/lightning/pytorch/callbacks/model_checkpoint.py Show resolved Hide resolved

ashors1 and others added 3 commits September 23, 2024 22:36

add test for save_last=link

0c7f066

Signed-off-by: ashors1 <[email protected]>

Merge branch 'ashors/symlink-last-ckpt' of github.com:NVIDIA/NeMo int…

b5ce506

…o ashors/symlink-last-ckpt

Apply isort and black reformatting

dae8c1f

Signed-off-by: ashors1 <[email protected]>

github-advanced-security bot found potential problems Sep 24, 2024

View reviewed changes

ashors1 added 2 commits September 24, 2024 09:34

clean up test

a5cb615

Signed-off-by: ashors1 <[email protected]>

use megatroncheckpointio in test

96f235e

Signed-off-by: ashors1 <[email protected]>

ashors1 and others added 2 commits September 24, 2024 10:43

fix remaining merge conflicts

e7d22d1

Signed-off-by: ashors1 <[email protected]>

Apply isort and black reformatting

afc2000

Signed-off-by: ashors1 <[email protected]>

ashors1 added Run CICD and removed Run CICD labels Sep 24, 2024

github-advanced-security bot found potential problems Sep 24, 2024

View reviewed changes

tests/lightning/pytorch/callbacks/test_model_checkpoint.py Fixed Show fixed Hide fixed

athitten reviewed Sep 24, 2024

View reviewed changes

tests/lightning/pytorch/callbacks/test_model_checkpoint.py Outdated Show resolved Hide resolved

ashors1 added 2 commits September 24, 2024 15:08

check number of saved checkpoints

93e7735

Signed-off-by: ashors1 <[email protected]>

remove unused import

e6e8fd2

Signed-off-by: ashors1 <[email protected]>

ashors1 added Run CICD and removed Run CICD labels Sep 24, 2024

run test on gpu only

fd727f4

Signed-off-by: ashors1 <[email protected]>

ashors1 added Run CICD and removed Run CICD labels Sep 25, 2024

mikolajblaz previously approved these changes Sep 26, 2024

View reviewed changes

fix a small bug and add a resume test

1888f54

Signed-off-by: ashors1 <[email protected]>

ashors1 dismissed mikolajblaz’s stale review via 1888f54 September 27, 2024 23:27

ashors1 and others added 3 commits September 27, 2024 23:28

Apply isort and black reformatting

0ea4411

Signed-off-by: ashors1 <[email protected]>

remove old comment

a2638f3

Signed-off-by: ashors1 <[email protected]>

Merge branch 'ashors/symlink-last-ckpt' of github.com:NVIDIA/NeMo int…

ff14905

…o ashors/symlink-last-ckpt

ashors1 added Run CICD and removed Run CICD labels Sep 28, 2024

athitten approved these changes Sep 28, 2024

View reviewed changes

athitten merged commit d664b74 into main Sep 30, 2024
153 of 160 checks passed

athitten deleted the ashors/symlink-last-ckpt branch September 30, 2024 17:33

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

[NeMo-UX] Support `save_last="link"` #10548

[NeMo-UX] Support `save_last="link"` #10548

ashors1 commented Sep 20, 2024 •

edited

Loading

mikolajblaz Sep 23, 2024

ashors1 Sep 23, 2024

mikolajblaz Sep 23, 2024

ashors1 Sep 23, 2024

mikolajblaz Sep 24, 2024 •

edited

Loading

ashors1 Sep 24, 2024

athitten left a comment

github-actions bot commented Sep 28, 2024

[NeMo-UX] Support save_last="link" #10548

[NeMo-UX] Support save_last="link" #10548

Conversation

ashors1 commented Sep 20, 2024 • edited Loading

What does this PR do ?

Changelog

Usage

GitHub Actions CI

Before your PR is "Ready for review"

Who can review?

Additional Information

mikolajblaz Sep 23, 2024

Choose a reason for hiding this comment

ashors1 Sep 23, 2024

Choose a reason for hiding this comment

mikolajblaz Sep 23, 2024

Choose a reason for hiding this comment

ashors1 Sep 23, 2024

Choose a reason for hiding this comment

mikolajblaz Sep 24, 2024 • edited Loading

Choose a reason for hiding this comment

ashors1 Sep 24, 2024

Choose a reason for hiding this comment

athitten left a comment

Choose a reason for hiding this comment

github-actions bot commented Sep 28, 2024

[NeMo-UX] Support `save_last="link"` #10548

[NeMo-UX] Support `save_last="link"` #10548

ashors1 commented Sep 20, 2024 •

edited

Loading

mikolajblaz Sep 24, 2024 •

edited

Loading