Use persistent workers if async saving is enabled #14465
Conversation
Signed-off-by: Ananth Subramaniam <[email protected]>
Signed-off-by: Ananth Subramaniam <[email protected]>
ananthsub
force-pushed
the
persistent-worker-ckpt
branch
from
a3431d5 to
989337c
Compare
| use_sharp: bool = False, | ||
| save_ckpt_format: str = "torch_dist", | ||
| ckpt_async_save: bool = True, | ||
| ckpt_async_save: bool = False, |
There was a problem hiding this comment.
why False by default?
There was a problem hiding this comment.
one concern now is persistent workers are consuming more memory, risking OOMs if async save is the default. that being said, all the recipes explicitly set this to True so until the memory regression is resolved, making sync save the default is the safer option
Signed-off-by: Ananth Subramaniam <[email protected]>
|
This PR is stale because it has been open for 14 days with no activity. Remove stale label or comment or update or this will be closed in 7 days. |
|
This PR was closed because it has been inactive for 7 days since being marked as stale. |
| """ | ||
|
|
||
| def __init__(self, checkpoint_io: AsyncCompatibleCheckpointIO) -> None: | ||
| def __init__(self, checkpoint_io: AsyncCompatibleCheckpointIO, persistent_workers: bool = False) -> None: |
There was a problem hiding this comment.
I need to double check in a clean installation but I was trying to quickly apply this patch/changeset to my branch. When using persistent_workers=True with save_last=True in ModelCheckpoint, my test was stuck with the below stack trace:
Traceback (most recent call last):
File "/usr/lib/python3.12/multiprocessing/process.py", line 314, in _bootstrap
self.run()
File "/usr/lib/python3.12/multiprocessing/process.py", line 108, in run
self._target(*self._args, **self._kwargs)
File "/usr/lib/python3.12/contextlib.py", line 81, in inner
return func(*args, **kwds)
^^^^^^^^^^^^^^^^^^^
File "/opt/megatron-lm/megatron/core/dist_checkpointing/strategies/async_utils.py", line 448, in async_loop
item = queue.get()
^^^^^^^^^^^
File "/usr/lib/python3.12/multiprocessing/queues.py", line 103, in get
res = self._recv_bytes()
^^^^^^^^^^^^^^^^^^
File "/usr/lib/python3.12/multiprocessing/connection.py", line 216, in recv_bytes
buf = self._recv_bytes(maxlength)
^^^^^^^^^^^^^^^^^^^^^^^^^^^
File "/usr/lib/python3.12/multiprocessing/connection.py", line 430, in _recv_bytes
buf = self._recv(4)
^^^^^^^^^^^^^
File "/usr/lib/python3.12/multiprocessing/connection.py", line 395, in _recv
chunk = read(handle, remaining)
^^^^^^^^^^^^^^^^^^^^^^^Possible that my envirorment has some issues but just mentioning here so that you can double check.
There was a problem hiding this comment.
Just to confirm here: I tested this branch/PR with the nemo 25.09 and it hangs at the end of training when using async checkpointing + persistent workers. cc: @maanug-nv
6 participants
Important
The
Update branchbutton must only be pressed in very rare occassions.An outdated branch is never blocking the merge of a PR.
Please reach out to the automation team before pressing that button.
What does this PR do ?
Async saving should use persistent workers for writing out the checkpoint
Collection: [Note which collection this PR will affect]
Changelog
Usage
# Add a code snippet demonstrating how to use thisGitHub 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:
PR Type:
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