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:
Chenjie Luo
2026-03-20 16:11:32 -07:00
committed by GitHub
parent e0f8b53905
commit 07bc4852a8
+4 -1
View File
@@ -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)