Support S3 checkpointing for the torch strategy in distributed checkpointing #748
Conversation
This commit adds support for saving checkpoints to cloud storage (e.g., S3) and loading checkpoints from cloud storage for the torch strategy in distributed checkpointing. It does so by replacing pathlib.Path with cloudpathlib.AnyPath, FileSystemReader with FsspecSystemReader, and FileSytemWriter with FsspecSystemWriter. The commit enables cloud checkpointing, but makes little attempt to optimize it.
|
@mikolajblaz Would you be the right one to review this? Thanks in advance! |
There was a problem hiding this comment.
Looks good in general, thanks.
Non-optimized cloud upload is understandable, but I want to make sure that the baseline scenario does not regress in terms of performance (only worried about 2-stage torch.save) and usability (added cloudpathlib requirement for MCore).
| @@ -345,7 +343,7 @@ class TorchDistSaveShardedStrategy(SaveShardedStrategy): | |||
| """ | |||
|
|
|||
| def __init__( | |||
| self, backend: str, version: int, keep_only_main_replica: bool = True, thread_count: int = 2 | |||
| self, backend: str, version: int, keep_only_main_replica: bool = True, thread_count: int = 1 | |||
| @@ -382,27 +380,27 @@ def save(self, sharded_state_dict: ShardedStateDict, checkpoint_dir: Path): | |||
| # Use PyT saving mechanism | |||
| save_state_dict( | |||
| pyt_state_dict, | |||
| FileSystemWriter(checkpoint_dir, thread_count=self.thread_count), | |||
| FsspecWriter(path=checkpoint_dir, thread_count=self.thread_count), | |||
There was a problem hiding this comment.
This need to be configurable, because I think Fsspec is slower than FilesystemWriter. Also, please note that there is a change on main which uses a custom save_state_dict and a custom FileSystemAsyncWriter for performance.
| @@ -3,8 +3,9 @@ | |||
| """ Module for managing distributed checkpoints metadata. """ | |||
|
|
|||
| import json | |||
| import os | |||
| from cloudpathlib import AnyPath | |||
There was a problem hiding this comment.
I assume it requires extra dependencies.
Can we make it optional?
| @@ -57,6 +59,22 @@ | |||
| logger = logging.getLogger(__name__) | |||
|
|
|||
|
|
|||
| def _save(obj: object, path: AnyPath): | |||
| fileobj = io.BytesIO() | |||
There was a problem hiding this comment.
I'm worried 2-stage writing might affect performance in the baseline scenario (no cloud).
Can we make it optional as well?
Thanks for the review! Quick question: do you have a suite of benchmarks that you run to detect performance regressions? If not, I can run some of my own benchmarks and report back. |
Please run your own benchmarks |
This PR adds support for saving checkpoints to cloud storage (e.g., S3) and loading checkpoints from cloud storage for the torch strategy in distributed checkpointing. It does so by replacing pathlib.Path with cloudpathlib.AnyPath, FileSystemReader with FsspecSystemReader, and FileSytemWriter with FsspecSystemWriter. The PR enables cloud checkpointing, but makes little attempt to optimize it.