🚨🚨 [saving] Default to 50GB shards, and remove non-safe serialization
Author: Cyrilvallez
Pull request: yes
Labels: name: for_v5?

# What does this PR do?

As per the title.
This PR fixes both feature requests https://github.com/huggingface/transformers/issues/42555 and https://github.com/huggingface/transformers/issues/42556.

Note that increasing the shard size was a decision that was only taken after *very careful* considerations, and the following benchmarks:

## Benchmarking scripts

First, the following script was used to create model checkpoints of several sizes and several shard sizes:

<details>
<summary>create_models.py</summary>

```python
import os
from tqdm import tqdm

import torch
from transformers import LlamaConfig, LlamaModel

# This given combination gives us a layer size of 1.008 ~ 1 GiB in bf16, which is very easy to work with
hidden_size = 2048 * 4
intermediate_size = 2048 * 8
num_attention_heads = 64
num_key_value_heads = 2
head_dim = hidden_size // num_attention_heads

mlp_size = 3 * hidden_size * intermediate_size
attn_size = 2 * hidden_size ** 2 + 2 * hidden_size * head_dim * num_key_value_heads
norm_size = hidden_size
layer_size = mlp_size + attn_size + 2 * norm_size

# Safeguard, make sure the layer size is almost exactly 1GiB
assert round(layer_size * 2 / 1024**3, 3) == 1.008

# All the model sizes we want to test out
MODEL_PARAMS = [8, 30, 70, 120]
# All the shard sizes we want to test out, for each model
SHARD_SIZES = [5, 10, 20, 35, 50]

for param_number in tqdm(MODEL_PARAMS, desc="Models"):
    model_size = param_number * 2  # 2x here as we computed layer size in GB (so 8B model -> 16 GiB in bf16 -> 16 layers)
    config = LlamaConfig(
        num_hidden_layers=model_size, 
        hidden_size=hidden_size,
        intermediate_size=intermediate_size,
        num_attention_heads=num_attention_heads,
        num_key_value_heads=num_key_value_heads,
    )
    # The fastest to skip all torch primitives initialization and ours, is to init on meta then fill with empty
    with torch.device("meta"):
        model = LlamaModel(config)
    model = model.to(dtype=torch.bfloat16)
    model = model.to_empty(device="cpu")

    # Safeguard
    assert all(v.dtype == torch.bfloat16 for v in model.parameters())
    
    last_shard_number = None
    for shard_size in tqdm(SHARD_SIZES, desc="Shard size", leave=False):
        name = f"{param_number}B_shard_{shard_size}"
        folder = os.path.join("random_models", name)
        # Always +1 here as it will never be exactly divisible anyway as layers are not precisely 1 GiB
        N_shards = (model_size // shard_size) + 1

        # In this case break the loop, model size is too small so increasing shard file give the same number of shards
        if last_shard_number is not None and N_shards == last_shard_number:
            break
        last_shard_number = N_shards

        # Unfortunately, hf_hub uses GB instead if GiB for now, so the model will be slightly more in GB compared to GiB
        # (but not important at all for the current study)
        model.save_pretrained(folder, max_shard_size=f"{shard_size}GB")
```

</details>

Then the following benchmarking script was run to obtain the data:

<details>
<summary>benchmark.py</summary>

```python
import os
import gc
import time
from tqdm import tqdm
from collections import defaultdict
import json

import torch
import matplotlib.pyplot as plt
import numpy as np
from transformers import LlamaModel

# Model sizes and shard sizes we test
MODEL_PARAMS = [8, 30, 70, 120]
SHARD_SIZES = [5, 10, 20, 35, 50]
# How many time to repeat loading for both warmup and main experiment
EXPERIMENT_REPEAT = 10
WARMUP_REPEAT = 4
# Plotting specifics
LINE_FMT = {
    "8B": "ro:",
    "30B": "bv:",
    "70B": "g^:",
    "120B": "m*:",
}
# Dtype name
DTYPE_NAMES = {
    torch.float16: "fp16",
    torch.bfloat16: "bf16",
    torch.float32: "fp32",
}

def default_dict_structure():
    return {"x": [], "y": [], "yerr": []}

def benchmark(device_map, dtype = torch.bfloat16):
    """Perform a full benchmark and save results and figures, given a `device_map` and `dtype` for laoding."""
    results = defaultdict(default_dict_structure)
    warmups = defaultdict(default_dict_structure)
    for param_number in tqdm(MODEL_PARAMS, desc="Models"):
        label = f"{param_number}B"
        for shard_size in tqdm(SHARD_SIZES, desc="Shard size", leave=False):
            name = f"{param_number}B_shard_{shard_size}"
            folder = os.path.join("random_models", name)
            # Some combinations of model size/shard size are not valid because redundant
            if not os.path.isdir(folder):
                break
            
            # Warmup the checkpoints into ssd for fair comparison
            measurements = []
            for _ in range(WARMUP_REPEAT):
                torch.cuda.synchronize()
                t0 = time.time()
                model = LlamaModel.from_pretrained(folder, dtype=dtype, device_map=device_map)
                dt = time.time() - t0
                measurements.append(dt)
                # Force memory release
                del model
                gc.collect()
                # Also release all cuda memory from torch so we take cudaMalloc into account in loading time
                torch.cuda.empty_cache()

            average_time = np.mean(measurements).item()
            std_time = np.std(measurements).item()
            # Add the results for plotting
            warmups[label]["x"].append(shard_size)
            warmups[label]["y"].append(average_time)
            warmups[label]["yerr"].append(std_time)

            # Perform actual benchmark
            measurements = []
            for _ in range(EXPERIMENT_REPEAT):
                torch.cuda.synchronize()
                t0 = time.time()
                model = LlamaModel.from_pretrained(folder, dtype=dtype, device_map=device_map)
                torch.cuda.synchronize()
                dt = time.time() - t0
                measurements.append(dt)
                # Force memory release
                del model
                gc.collect()
                # Also release all cuda memory from torch so we take cudaMalloc into account in loading time
                torch.cuda.empty_cache()
            
            average_time = np.mean(measurements).item()
            std_time = np.std(measurements).item()
            # Add the results for plotting
            results[label]["x"].append(shard_size)
            results[label]["y"].append(average_time)
            results[label]["yerr"].append(std_time)

    result_name = f"device_map_{device_map}_dtype_{DTYPE_NAMES[dtype]}"
    # Dump raw results to disk if needed to adjust plots later
    with open(f"raw_results_{result_name}.json", "w") as f:
        f.write(json.dumps({"results": results, "warmups": warmups}))

    # Plot everything
    plt.figure()
    for label, data in results.items():
        plt.errorbar(data["x"], data["y"], yerr=data["yerr"], label=label, fmt=LINE_FMT[label], capsize=3)
    plt.xlabel("Shard size [GB]")
    plt.ylabel("Average loading time [s]")
    plt.legend()
    plt.grid()
    plt.title(f"device_map={device_map}, dtype={dtype}")
    # Force integer ticks for x
    locs, labels = plt.xticks()
    plt.xticks(locs, [x.get_text().split(".")[0] for x in labels])
    plt.savefig(f"main_{result_name}.png", bbox_inches="tight")


    plt.figure()
    for label, data in warmups.items():
        plt.errorbar(data["x"], data["y"], yerr=data["yerr"], label=label, fmt=LINE_FMT[label], capsize=3)
    plt.xlabel("Shard size [GB]")
    plt.ylabel("Average warmup loading time [s]")
    plt.legend()
    plt.grid()
    plt.title(f"device_map={device_map}, dtype={dtype}")
    # Force integer ticks for x
    locs, labels = plt.xticks()
    plt.xticks(locs, [x.get_text().split(".")[0] for x in labels])
    plt.savefig(f"warmup_{result_name}.png", bbox_inches="tight")


# Actually perform the benchmarks
benchmark(device_map="auto")
benchmark(device_map="cpu")
benchmark(device_map="auto", dtype=torch.float16)  # dtype in checkpoint is bf16, so this forces copy to switch dtype
benchmark(device_map="cpu", dtype=torch.float16)  # dtype in checkpoint is bf16, so this forces copy to switch dtype
```

</details>

## Results

TLDR: we don't find any significative impact on loading time coming from the shard size (in different scenarios, i.e. all 4 combinations of cpu/gpu loading and native/new dtype), so we are free to increase to the now much better 50GB shard size recommendation coming from the hub.

Here, we present the results obtained. Note that all the models were saved in `bf16` (`torch.bfloat16`), so loading them back in any other dtype requires an additional copy of the data, which is expected to take longer!
We always present main results when the checkpoints are hot in ssd, as well as timings during warmup loading runs. Note that this is not always perfectly reliable as we absolutely cannot control how the OS will cache the files, but 4 warmup loading runs should be enough. We still notice bigger errors bars on main runs for a few datapoints, meaning that the checkpoints were not perfectly cached yet.

Note: these results were obtained on a node of our cluster as of December 2025, i.e. a node with 8xH100 GPUs, and 2TiB cpu RAM.

### `device_map="auto"`, `dtype=torch.bfloat16` (native dtype)

 
<img width=49% height=auto alt="main_device_map_auto_dtype_bf16" src="https://github.com/user-attachments/assets/f7e1ab4c-2ae1-46ab-b691-286fa3e1a347" />
<img width=49% height=auto alt="warmup_device_map_auto_dtype_bf16" src="https://github.com/user-attachments/assets/5b990e9b-39cc-4162-b1c3-b3fb382512f7" />

### `device_map="cpu"`, `dtype=torch.bfloat16` (native dtype)

<img width=49% height=auto alt="main_device_map_cpu_dtype_bf16" src="https://github.com/user-attachments/assets/2ea9dfd4-a926-4e42-b32d-9be0b2fe94aa" />
<img width=49% height=auto alt="warmup_device_map_cpu_dtype_bf16" src="https://github.com/user-attachments/assets/fd5e8e54-f5cf-478e-a858-73aef07ac048" />

### `device_map="auto"`, `dtype=torch.float16` (NON-native dtype, needs copy)

<img width=49% height=auto alt="main_device_map_auto_dtype_fp16" src="https://github.com/user-attachments/assets/3272d165-a2b5-4e7b-a2cc-96e6e4501753" />
<img width=49% height=auto alt="warmup_device_map_auto_dtype_fp16" src="https://github.com/user-attachments/assets/14c7dd0a-0454-43a2-9588-eac00220b001" />

### `device_map="cpu"`, `dtype=torch.float16` (NON-native dtype, needs copy)

(not sure what happened for 120B on this one, looks like the cpu got saturated a bit - but the trend is the same anyway)

<img width=49% height=auto alt="main_device_map_cpu_dtype_fp16" src="https://github.com/user-attachments/assets/cc7311a9-5c95-4f1e-880d-5bb6033e9df0" />
<img width=49% height=auto alt="warmup_device_map_cpu_dtype_fp16" src="https://github.com/user-attachments/assets/938aeb6b-fd29-4d25-acca-fbe283e462b9" />


Comments:
The docs for this PR live [here](https://moon-ci-docs.huggingface.co/docs/transformers/pr_42734). All of your documentation changes will be reflected on that endpoint. The docs are available until 30 days after the last update.

Spotted a few places where `safe_serialization` still exists:
- in base `Pipeline` https://github.com/huggingface/transformers/blob/05c0e1d39082ee8b69064ed4ea9c239cd17405e9/src/transformers/pipelines/base.py#L956
- in `PushToHubMixin` https://github.com/huggingface/transformers/blob/ba1ad5359394c6dab864d5fb56658f35516c5fa5/src/transformers/utils/hub.py#L725
- in 5 test files including many occurrences in `test_modeling_utils.py` https://github.com/huggingface/transformers/blob/ba1ad5359394c6dab864d5fb56658f35516c5fa5/tests/utils/test_modeling_utils.py#L762
- in the `Trainer` (as `self.args.save_safetensors`) https://github.com/huggingface/transformers/blob/6564633bc87b1f808de5773a74db2ec01f681520/src/transformers/trainer.py#L4087
- in many convert scripts (dunno if those are maintained?) e.g. https://github.com/huggingface/transformers/blob/69cff312f5c8026fea13029bb45b139385a88b4c/src/transformers/models/deepseek_vl/convert_deepseek_vl_weights_to_hf.py#L310
- in many examples/ scripts as well

Would be great to entirely get rid of them

Yep, and much more in the tests 🥲 Should have removed them all now

Thanks @Cyrilvallez for the update. I've pushed https://github.com/huggingface/transformers/pull/42734/commits/bf8680d02372a88e5008af25a1ccc52b8951f2d9 as well to update the musicgen docs.

Otherwise the [torchao docs](https://github.com/huggingface/transformers/blob/80134e6e663db962c30d0e45c54265b3141bc7c7/docs/source/en/quantization/torchao.md?plain=1#L640) are mentioning:

> [Safetensors](https://huggingface.co/docs/safetensors/en/index) serialization and deserialization does not work with torchao.
> 
> (...)
> 
> ```py
> # don't serialize model with Safetensors
> output_dir = "llama3-8b-int4wo-128"
> quantized_model.save_pretrained("llama3-8b-int4wo-128", safe_serialization=False)
> ```
> 
> (...) + in 3 other places

What should be done with this since `safe_serialization=False` cannot be passed anymore? 


@Wauplin the doc is wrong here, it does support it since version > 0.15.0

> @Wauplin the doc is wrong here, it does support it since version > 0.15.0

(I'll let you update it, I'm not exactly sure what should be written)

**[For maintainers]** Suggested jobs to run (before merge)

run-slow: aimv2, aria, chameleon, csm, deepseek_vl, deepseek_vl_hybrid, depth_pro, efficientloftr, emu3, eomt, fuyu, gemma, gemma2, gemma3, gemma3n, gpt_oss

View the CircleCI Test Summary for this PR:

https://huggingface.co/spaces/transformers-community/circle-ci-viz?pr=42734&sha=f3ea9e