Implement async_checkpoint #313
Conversation
Summary: This PR implements 2 different async checkpoint. The first one is to use DCP.async_save another one is to use pinned memory + a seperate process to avoid GILs issue. ghstack-source-id: 87fb6c28d7bc3e514c0bee7646be5188f1f66bbd Pull Request resolved: #313
facebook-github-bot
added
the
CLA Signed
Summary: This PR implements 2 different async checkpoint. The first one is to use DCP.async_save another one is to use pinned memory + a seperate process to avoid GILs issue. ghstack-source-id: 87fb6c28d7bc3e514c0bee7646be5188f1f66bbd Pull Request resolved: #313
|
It would be good to add an integration test for async checkpoint cc: @fegin |
Summary: This PR implements 2 different async checkpoint. The first one is to use DCP.async_save another one is to use pinned memory + a seperate process to avoid GILs issue. ghstack-source-id: 87fb6c28d7bc3e514c0bee7646be5188f1f66bbd Pull Request resolved: pytorch#313
Summary: This PR implements 2 different async checkpoint. The first one is to use DCP.async_save another one is to use pinned memory + a seperate process to avoid GILs issue. ghstack-source-id: 87fb6c28d7bc3e514c0bee7646be5188f1f66bbd Pull Request resolved: pytorch#313
| self._async_with_pinned_memory(checkpoint_id) | ||
| elif self.async_mode == AsyncMode.ASYNC: | ||
| self.async_future = dcp.async_save( | ||
| self.states, checkpoint_id=checkpoint_id, process_group=self.pg |
There was a problem hiding this comment.
@fegin Why did you choose to use the GLOO process group for the async save? Is it expected to make this more efficient?
Neither the DCP docs nor https://discuss.pytorch.org/t/distributed-w-torchtitan-optimizing-checkpointing-efficiency-with-pytorch-dcp/211250 mention or recommend this.
I'm curious to know if this was on purpose and if you have any numbers to show
Thanks!
There was a problem hiding this comment.
We don't want the checkpointing to affect the training, which involve NCCL. So we choose gloo for any checkpointing communication. Also the main bottleneck of checkpointing is unlikely to be the communication. The storage read/write (or upload/download) will be the major overhead.
Stack from ghstack (oldest at bottom):
Summary:
This PR implements 2 different async checkpoint. The first one is to use
DCP.async_save another one is to use pinned memory + a seperate process
to avoid GILs issue.