NeMo already provides S3CheckpointIO for synchronous and asynchronous checkpoint storage. I would like to propose adding support for saving and restoring NeMo checkpoints with Hugging Face Storage Buckets.
Storage Buckets are a strong fit for NeMo workloads because they provide read/write object storage for large distributed checkpoints, Xet-based deduplication across similar checkpoint versions, access from different training environments, and integration with Hugging Face authentication and tooling.
huggingface_hub is already a NeMo dependency, and Storage Buckets are accessible through both native hf://buckets/ paths and an S3-compatible API.
One possible native user experience would be:
exp_manager:
checkpoint_callback_params:
dirpath: hf://buckets/my-org/nemo-checkpoints/experiment-1
hf_bucket_checkpointing:
enable_async_checkpointing: true
max_read_concurrency: 15
max_write_concurrency: 10
chunk_size_MB: 2048
Authentication could use the normal Hugging Face mechanisms:
hf auth login
export HF_TOKEN=...
One implementation option would be an HFBucketCheckpointIO backend using huggingface_hub.HfFileSystem or the native bucket APIs. This would support hf://buckets/ URLs and allow users to authenticate through HF_TOKEN without generating separate S3 credentials.
Another option would be to extend S3CheckpointIO with a Hugging Face provider or endpoint configuration:
s3_checkpointing:
provider: huggingface
namespace: my-org
bucket: nemo-checkpoints
prefix: experiment-1
enable_async_checkpointing: true
The S3-compatible approach could reuse most of the existing checkpoint implementation. The gateway requires the endpoint https://s3.hf.co/<namespace>, region us-east-1, path-style addressing, and checksum calculation and validation set to when_required. Large multipart chunks are recommended.
It would be useful for the integration to cover rank-zero checkpoint discovery and broadcast, synchronous and asynchronous writes, multipart uploads for large checkpoint shards, restoration after interrupted training, checkpoint rotation, read-only tokens, and concurrent reads from many ranks.
The S3-compatible gateway does not support every AWS S3 feature, including object versioning, ACLs, lifecycle rules, server-side encryption, object tagging, and UploadPartCopy. Existing AWS S3 behavior should remain unchanged when Hugging Face is not selected.
I'm happy to submit a PR if a maintainer indicates interest
NeMo already provides
S3CheckpointIOfor synchronous and asynchronous checkpoint storage. I would like to propose adding support for saving and restoring NeMo checkpoints with Hugging Face Storage Buckets.Storage Buckets are a strong fit for NeMo workloads because they provide read/write object storage for large distributed checkpoints, Xet-based deduplication across similar checkpoint versions, access from different training environments, and integration with Hugging Face authentication and tooling.
huggingface_hubis already a NeMo dependency, and Storage Buckets are accessible through both nativehf://buckets/paths and an S3-compatible API.One possible native user experience would be:
Authentication could use the normal Hugging Face mechanisms:
hf auth login export HF_TOKEN=...One implementation option would be an
HFBucketCheckpointIObackend usinghuggingface_hub.HfFileSystemor the native bucket APIs. This would supporthf://buckets/URLs and allow users to authenticate throughHF_TOKENwithout generating separate S3 credentials.Another option would be to extend
S3CheckpointIOwith a Hugging Face provider or endpoint configuration:The S3-compatible approach could reuse most of the existing checkpoint implementation. The gateway requires the endpoint
https://s3.hf.co/<namespace>, regionus-east-1, path-style addressing, and checksum calculation and validation set towhen_required. Large multipart chunks are recommended.It would be useful for the integration to cover rank-zero checkpoint discovery and broadcast, synchronous and asynchronous writes, multipart uploads for large checkpoint shards, restoration after interrupted training, checkpoint rotation, read-only tokens, and concurrent reads from many ranks.
The S3-compatible gateway does not support every AWS S3 feature, including object versioning, ACLs, lifecycle rules, server-side encryption, object tagging, and
UploadPartCopy. Existing AWS S3 behavior should remain unchanged when Hugging Face is not selected.I'm happy to submit a PR if a maintainer indicates interest