🚨🚨 [saving] Default to 50GB shards, and remove non-safe serialization

[Cyrilvallez]

What does this PR do?

As per the title.
This PR fixes both feature requests #42555 and #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:

create_models.py

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")

Then the following benchmarking script was run to obtain the data:

benchmark.py

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

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)

device_map="cpu", dtype=torch.bfloat16 (native dtype)

device_map="auto", dtype=torch.float16 (NON-native dtype, needs copy)

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)

[HuggingFaceDocBuilderDev]

The docs for this PR live here. All of your documentation changes will be reflected on that endpoint. The docs are available until 30 days after the last update.

[ArthurZucker]

LGTM put a lot of 🔴 on top please

[Wauplin]

Spotted a few places where safe_serialization still exists:

• in base Pipeline

  transformers/src/transformers/pipelines/base.py

  Line 956 in 05c0e1d

  safe_serialization: bool = True,

• in PushToHubMixin

  transformers/src/transformers/utils/hub.py

  Line 725 in ba1ad53

  safe_serialization: bool = True,

• in 5 test files including many occurrences in test_modeling_utils.py

  transformers/tests/utils/test_modeling_utils.py

  Line 762 in ba1ad53

  model.save_pretrained(tmp_dir, max_shard_size=max_size, safe_serialization=False)

• in the Trainer (as self.args.save_safetensors)

  transformers/src/transformers/trainer.py

  Line 4087 in 6564633

  safe_serialization=self.args.save_safetensors,

• in many convert scripts (dunno if those are maintained?) e.g.

  transformers/src/transformers/models/deepseek_vl/convert_deepseek_vl_weights_to_hf.py

  Line 310 in 69cff31

  model.save_pretrained(output_dir, safe_serialization=safe_serialization)

• in many examples/ scripts as well

Would be great to entirely get rid of them

[Cyrilvallez]

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

[SunMarc]

Thanks ! Left some minor comments !

[SunMarc]

cc @jerryzh168 we will remove bin serialization from now on.

[Wauplin]

Thanks @Cyrilvallez for the update. I've pushed bf8680d as well to update the musicgen docs.

Otherwise the torchao docs are mentioning:

Safetensors serialization and deserialization does not work with torchao.

(...)

# 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?

[Cyrilvallez]

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

[Wauplin]

@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)

[McPatate]

chatted with @ArthurZucker, we might update safetensors to handle aliasing better than the current state of affairs.
will ping when I have more on that!

[github-actions]

[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

[Wauplin]

Thanks for the updates!

[github-actions]

View the CircleCI Test Summary for this PR:

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