mirror of
https://github.com/NVIDIA/Model-Optimizer.git
synced 2026-10-02 03:14:52 +08:00
Default limit max hf quantzed safetensors file to 10GB (#1087)
### What does this PR do? Type of change: New feature Adds a `max_shard_size` parameter to `export_hf_checkpoint()` (and the internal `_export_diffusers_checkpoint()`) that controls the maximum size of each exported safetensors shard file. Defaults to `"10GB"`. Previously, the shard size was not explicitly controlled, which could result in very large single-file checkpoints [E.g. Qwen3.5] that are difficult to handle (e.g., slow uploads, memory issues, or exceeding file size limits on model hubs). This change ensures exported checkpoints are automatically sharded into ≤10GB files by default, while allowing users to customize the threshold. ### Usage ```python import modelopt.torch.quantization as mtq from modelopt.torch.export import export_hf_checkpoint # Default: shards capped at 10GB export_hf_checkpoint(model, dtype=torch.float16, export_dir="./output") # Custom shard size export_hf_checkpoint(model, dtype=torch.float16, export_dir="./output", max_shard_size="5GB") ``` ### Testing Verified that the `max_shard_size` parameter is correctly propagated to both the HF `save_pretrained()` (for transformers models) and diffusers component `save_pretrained()` calls. ### Before your PR is "*Ready for review*" Make sure you read and follow [Contributor guidelines](https://github.com/NVIDIA/Model-Optimizer/blob/main/CONTRIBUTING.md) and your commits are signed (`git commit -s -S`). Make sure you read and follow the [Security Best Practices](https://github.com/NVIDIA/Model-Optimizer/blob/main/SECURITY.md#security-coding-practices-for-contributors) (e.g. avoiding hardcoded `trust_remote_code=True`, `torch.load(..., weights_only=False)`, `pickle`, etc.). - Is this change backward compatible?: ✅ (new optional parameter with default matching previous behavior of HF's `save_pretrained`) - If you copied code from any other sources or added a new PIP dependency, did you follow guidance in `CONTRIBUTING.md`: N/A - Did you write any new necessary tests?: ❌ - Did you update [Changelog](https://github.com/NVIDIA/Model-Optimizer/blob/main/CHANGELOG.rst)?: ❌ ### Additional Information The 10GB default was chosen to keep exported safetensors files within common file size limits while minimizing unnecessary sharding for most models. <!-- This is an auto-generated comment: release notes by coderabbit.ai --> ## Summary by CodeRabbit * **New Features** * Added configurable maximum shard size parameter for Hugging Face checkpoint exports. When exporting both transformer and diffusion models, users can now specify the maximum size of each safetensors shard, controlling how checkpoint files are divided. The default shard size is set to 10GB. This feature works with both quantized and non-quantized models. <!-- end of auto-generated comment: release notes by coderabbit.ai --> Signed-off-by: Chenjie Luo <chenjiel@nvidia.com>
This commit is contained in:
@@ -1118,6 +1118,7 @@ def export_hf_checkpoint(
|
||||
save_modelopt_state: bool = False,
|
||||
components: list[str] | None = None,
|
||||
extra_state_dict: dict[str, torch.Tensor] | None = None,
|
||||
max_shard_size: int | str = "10GB",
|
||||
**kwargs,
|
||||
):
|
||||
"""Export quantized HuggingFace model checkpoint (transformers or diffusers).
|
||||
@@ -1136,6 +1137,7 @@ def export_hf_checkpoint(
|
||||
components: Only used for diffusers pipelines. Optional list of component names
|
||||
to export. If None, all quantized components are exported.
|
||||
extra_state_dict: Extra state dictionary to add to the exported model.
|
||||
max_shard_size: Maximum size of each safetensors shard file. Defaults to "10GB".
|
||||
**kwargs: Internal-only keyword arguments. Supported key: merged_base_safetensor_path
|
||||
(str, optional). When provided, merges the exported diffusion transformer
|
||||
weights with non-transformer components (VAE, vocoder, text encoders, etc.)
|
||||
@@ -1153,7 +1155,7 @@ def export_hf_checkpoint(
|
||||
is_diffusers_obj = is_diffusers_object(model)
|
||||
if is_diffusers_obj:
|
||||
_export_diffusers_checkpoint(
|
||||
model, dtype, export_dir, components, merged_base_safetensor_path
|
||||
model, dtype, export_dir, components, merged_base_safetensor_path, max_shard_size
|
||||
)
|
||||
return
|
||||
|
||||
@@ -1183,6 +1185,7 @@ def export_hf_checkpoint(
|
||||
export_dir,
|
||||
state_dict={**post_state_dict, **(extra_state_dict or {})},
|
||||
save_modelopt_state=save_modelopt_state,
|
||||
max_shard_size=max_shard_size,
|
||||
)
|
||||
finally:
|
||||
_unpatch_revert_weight_conversion(_patches)
|
||||
|
||||
Reference in New Issue
Block a user