Summary
PyTorch Distributed Checkpointing (DCP) is taking significant strides to enhance interoperability within its ecosystem, particularly by integrating popular formats like HuggingFace safetensors. As HuggingFace has emerged as a leading format for inference and fine-tuning in machine learning, DCP’s support for safetensors is a vital addition. The first beneficiary of these improvements is torchtune, which has reported an enhanced user experience thanks to the newfound ability to seamlessly read and write HuggingFace models directly through DCP APIs.
The Problem
With over 5 million users utilizing HuggingFace, many machine learning engineers find themselves wanting to save and load their checkpoints in the safetensors format, which aligns better with the HuggingFace ecosystem. By implementing native support for the safetensors format in DCP, the checkpointing process becomes significantly streamlined, making it more accessible for users. Here are the primary challenges that DCP addresses:
- Previously, DCP maintained a custom format. Users aiming to leverage DCP’s performance enhancements while working with HuggingFace models had to create their own converters and components, leading to unnecessary complexity.
- Instead of needing to rely on local storage for every interaction, users can now save and load their HuggingFace models directly into any fsspec-compatible storage, enhancing workflow efficiency.
How to Use
From the user standpoint, the transition to using safetensors is straightforward. The primary adjustment involves using the new load planner and storage reader when loading models, as well as the save planner and storage writer for saving. The load and save APIs are structured as follows:
load(
state_dict=state_dict,
storage_reader=HuggingFaceStorageReader(path=path),
)
save(
state_dict=state_dict,
storage_writer=HuggingFaceStorageWriter(
path=path,
fqn_to_index_mapping=mapping
),
)
The HuggingFaceStorageReader and HuggingFaceStorageWriter are capable of handling any fsspec-based path, allowing read and write operations in the HF safetensors format for various backend storage options, including both local and HF storage. While the metadata capabilities of HuggingFace safetensors currently lack some depth compared to DCP’s metadata, DCP is actively working on plans for enhancing native support for distributed checkpoints.
Torchtune’s Experience
The first adopter of the new HuggingFace DCP support is torchtune, a post-training library built using native PyTorch. For torchtune users, the typical process involves fetching model weights from the Hugging Face Hub. Previously, this required manual downloads and uploads of trained checkpoints through additional CLI commands. With the introduction of the new DCP APIs, users now enjoy the ability to directly read from and write to HuggingFace, resulting in a significantly more user-friendly experience.
Additionally, integrating safetensor serialization into DCP has simplified the checkpointing process within torchtune, eliminating the need for format-specific checkpointing solutions. This shift not only enhances user experience but also boosts developer productivity within the project.
Future Work
Looking ahead, DCP aims to refine how distributed loading and saving of HuggingFace safetensors checkpoints are handled, particularly through resharding techniques. Furthermore, plans are in place to enable the generation of a consolidated final checkpoint into a single file, streamlining the publishing process for users. As developments move forward, improvements in DCP’s integration with the HuggingFace ecosystem will continue to enhance the overall experience for machine learning engineers.
Inspired by: Source
