[feature request] Native checkpointing to/from `s3://`

[vadimkantorov]

🚀 The feature, motivation and pitch

Sometimes it's beneficial to directly stream checkpoint data to a cloud storage (rather than dump it localy and have some background process handle sync/upload/cleanup) or load weights from s3:// or gs:// path checkpoint weights. I wonder if this also can be related to the recent native HFStorageReader support #154518

I also found that torchsnapshot library (seems abandonned now - last commit 6 months ago) supports this: https://docs.pytorch.org/torchsnapshot/main/getting_started.html

---

Also there exist a fairly popular packge fsspec which itself wraps some cloud storage libraries and provides caching functionalities, there was some discussion in torchsnapshot on supporting in natively:

• FSSpec support for TorchSnapshot? torchsnapshot#102
• Add fsspec storage plugin torchsnapshot#114

And some older (before HF) discussions:

• Proxy/cache server option/hooks for downloading model checkpoints and dataset archive files in cloud environment #91965
• Potential race conditions between multiple workers trying to download and cache the same file in torch.hub.load_state_dict_from_url and torch.hub.download_url_to_file <- duplicate dataset/model downloads across DDP workers #68320

I wonder if some HF utils on checkpointing / HF Hub blob caching structure could be upstreamed to PyTorch. E.g. for loading pretrained weights, this should be good. And maybe some hf:// management/caching could be made to plug into fsspec interface? Then hf:// could be used in all fsspec-using places.

Alternatives

No response

Additional context

No response

cc @mruberry @mikaylagawarecki @LucasLLC @pradeepfn @MeetVadakkanchery @mhorowitz @ekr0

[Skylion007]

Fsspec works natively if you pass an object opened with fsspec.open to torch.load or torch.save. It will write directly to the object storage. We don't have an URI handler for fsspec at the moment, but I think distributed checkpoint already supports this natively.

[vadimkantorov]

It would be great to showcase this somewhere for DCP, yes

And maybe in torchtune/torchtitan...

Btw, do you know what is the status of torchsnapshot and if it's relevant for torchtune/torchtitan?

[vadimkantorov]

Am I right that a separate package https://github.com/awslabs/s3-connector-for-pytorch needs to be installed for s3 checkpoints with DCP?

Given that S3-API interfaces is so wide-spread (and not only for Amazon), it would be great to at least have some official tutorials on proper integration with loading/saving HF-format checkpoints from S3, DCP, torch.load, torch.save. And then maybe some nicer integration (currently one would probably need both fsspec and s3-connector-for-pytorch)

[malfet]

Added triage review as I'm not sure about the label

[ankitageorge]

@vadimkantorov a separate package doesn't need to be installed. If you use the FsspecWriter and FsspecReader (

pytorch/torch/distributed/checkpoint/_fsspec_filesystem.py

Line 1 in da910e6

# Mypy will not try inferring the types of any 3rd party libraries installed.

), it should work natively.
You're right that this is not well documented though, have assigned to myself to add documentation here.
Let me know if you have questions or issues with the fsspecwriter and reader