Add support for saving HF format tensors with DCP #1351
Conversation
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
facebook-github-bot
added
the
CLA Signed
ankitageorge
changed the title
|
@Saiteja64 This will conflict with your PR. |
fegin
reviewed
Jun 27, 2025
There was a problem hiding this comment.
Overall the logic LGTM, please address comments and ensure that this PR doesn't conflict with the PR from @Saiteja64. Please also add a test result -- save a hf checkpoint and load one back and check the accuracy.
| @@ -12,14 +12,19 @@ | |||
| import shutil | |||
| import threading | |||
| import time | |||
| from typing import Any | |||
| from concurrent.futures import Future | |||
| from typing import Any, Optional | |||
There was a problem hiding this comment.
We use 3.10 type checking, so you don't need Optional.
| checkpoint_id: str, | ||
| is_async: bool, | ||
| hf_safetensors_format: bool, | ||
| pg: Optional[dist.ProcessGroup] = None, |
There was a problem hiding this comment.
Suggested change
| pg: Optional[dist.ProcessGroup] = None, | |
| pg: dist.ProcessGroup | None = None, |
| def dcp_save( | ||
| state_dict: dict[str, Any], | ||
| checkpoint_id: str, | ||
| is_async: bool, |
There was a problem hiding this comment.
@Saiteja64 do we also need anther argument for ZOC?
| is_async: bool, | ||
| hf_safetensors_format: bool, | ||
| pg: Optional[dist.ProcessGroup] = None, | ||
| ) -> Optional[Future]: |
There was a problem hiding this comment.
Suggested change
| ) -> Optional[Future]: | |
| ) -> Future | None: |
| hf_safetensors_format: bool, | ||
| pg: Optional[dist.ProcessGroup] = None, | ||
| ) -> Optional[Future]: | ||
| """Save the checkpoint with dcp. |
| checkpoint_id (str): The checkpoint id to save. | ||
| is_async (bool): Whether the checkpoint is async. | ||
| hf_safetensors_format (bool): Whether to use the HuggingFace safetensors format. | ||
| pg (Optional[dist.ProcessGroup]): The process group to use. |
There was a problem hiding this comment.
Add the return value as well.
Comment on lines
+116
to
+130
| if hf_safetensors_format: | ||
| storage_writer = HuggingFaceStorageWriter(path=checkpoint_id, save_sharded=True) | ||
| if is_async: | ||
| return dcp.async_save( | ||
| state_dict, storage_writer=storage_writer, process_group=pg | ||
| ) | ||
| else: | ||
| return dcp.save(state_dict, storage_writer=storage_writer) | ||
| else: | ||
| if is_async: | ||
| return dcp.async_save( | ||
| state_dict, checkpoint_id=checkpoint_id, process_group=pg | ||
| ) | ||
| else: | ||
| return dcp.save(state_dict, checkpoint_id=checkpoint_id) |
There was a problem hiding this comment.
We should simplify the function as follow
Suggested change
| if hf_safetensors_format: | |
| storage_writer = HuggingFaceStorageWriter(path=checkpoint_id, save_sharded=True) | |
| if is_async: | |
| return dcp.async_save( | |
| state_dict, storage_writer=storage_writer, process_group=pg | |
| ) | |
| else: | |
| return dcp.save(state_dict, storage_writer=storage_writer) | |
| else: | |
| if is_async: | |
| return dcp.async_save( | |
| state_dict, checkpoint_id=checkpoint_id, process_group=pg | |
| ) | |
| else: | |
| return dcp.save(state_dict, checkpoint_id=checkpoint_id) | |
| storage_writer = HuggingFaceStorageWriter(path=checkpoint_id, save_sharded=True) if hf_safetensors_format else None | |
| checkpoint_id = checkpoint_id if not hf_safetensors_format else None | |
| if is_async: | |
| return dcp.async_save( | |
| state_dict, storage_writer=storage_writer, checkpoint_id=checkpoint_id, process_group=pg | |
| ) | |
| else: | |
| return dcp.save(state_dict, storage_writer=storage_writer, checkpoint_id=checkpoint_id) |
| def dcp_load( | ||
| state_dict: dict[str, Any], checkpoint_id: str, hf_safetensors_format: bool | ||
| ) -> None: | ||
| """Load the checkpoint with dcp. |
| enable_hf_safetensors_format: bool = False | ||
| """ | ||
| Enable the use of safetensors format for checkpointing. This will save checkpoints | ||
| in safetensors format instead of the default DCP format. The default value is False. |
There was a problem hiding this comment.
Can we also mention the possible performance penalty? It's not cost free, right?
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.
This suggestion is invalid because no changes were made to the code.
Suggestions cannot be applied while the pull request is closed.
Suggestions cannot be applied while viewing a subset of changes.
Only one suggestion per line can be applied in a batch.
Add this suggestion to a batch that can be applied as a single commit.
Applying suggestions on deleted lines is not supported.
You must change the existing code in this line in order to create a valid suggestion.
Outdated suggestions cannot be applied.
This suggestion has been applied or marked resolved.
Suggestions cannot be applied from pending reviews.
Suggestions cannot be applied on multi-line comments.
Suggestions cannot be applied while the pull request is queued to merge.
Suggestion cannot be applied right now. Please check back later.
If checkpoint.enable_hf_safetensors_format is set, then save the checkpoint with DCP HF components that will save the checkpoint in .safetensors files instead of regular DCP format.