From 07bc4852a85c72f2712d33afa7987840d7816cd1 Mon Sep 17 00:00:00 2001 From: Chenjie Luo <108829653+cjluo-nv@users.noreply.github.com> Date: Fri, 20 Mar 2026 16:11:32 -0700 Subject: [PATCH] Default limit max hf quantzed safetensors file to 10GB (#1087) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit ### 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. ## 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. Signed-off-by: Chenjie Luo --- modelopt/torch/export/unified_export_hf.py | 5 ++++- 1 file changed, 4 insertions(+), 1 deletion(-) diff --git a/modelopt/torch/export/unified_export_hf.py b/modelopt/torch/export/unified_export_hf.py index 78c8874a0..9d7b75eb1 100644 --- a/modelopt/torch/export/unified_export_hf.py +++ b/modelopt/torch/export/unified_export_hf.py @@ -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)