Skip to content

[Feature] Support Hugging Face Storage Buckets for checkpointing #16016

Description

@abidlabs

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

Metadata

Metadata

Assignees

No one assigned

    Type

    No type

    Projects

    No projects

    Milestone

    No milestone

    Relationships

    None yet

    Development

    No branches or pull requests

    Issue actions