[core] Shard tensor-parallel checkpoints on load and save by JingyaHuang · Pull Request #14544 · huggingface/diffusers

JingyaHuang · 2026-08-20T14:19:39Z

What does this PR do?

This is a follow-up of the Tensor Parallelism support in #13781, TP previously required loading the whole checkpoint on every rank and resharding it afterwards, so per-rank memory was the full model size. In this PR, we adapt the shard loading (.from_pretrained()) and saving (.save_pretrained()) to be tp-aware:

Load: from_pretrained(..., parallel_config=...) shards while reading, each rank slices only its own part of every _tp_plan weight off disk, straight into a DTensor on its device. Unsupported combinations: device_map / quantization / use_flashpack / DDUF / non-safetensors -> raise.
Save: save_pretrained() gathers the shards back to a normal checkpoint
- collective: all on all ranks, rank 0 writes
- save_pretrained(..., dcp=True) writes per-rank

import torch
import torch.distributed as dist

from diffusers import Flux2Transformer2DModel, TensorParallelConfig

dist.init_process_group()                                  # "nccl" on CUDA, "neuron" on Trainium
tp = TensorParallelConfig(tp_degree=dist.get_world_size())

# LOAD: shard while reading, never materialize the full checkpoint ----
model = Flux2Transformer2DModel.from_pretrained(
    "black-forest-labs/FLUX.2-klein-9B",
    subfolder="transformer",
    torch_dtype=torch.bfloat16,
    parallel_config=tp,
)

# SAVE: gather back
model.save_pretrained("out/full")

# SAVE (b): sharded distributed checkpoint, nothing gathered
# model.save_pretrained("out/sharded", dcp=True)

# LOAD BACK a distributed ckpt
# model = Flux2Transformer2DModel.from_pretrained("out/sharded", parallel_config=tp)

Besides above:

_check_tp_model_state, called from apply_tensor_parallel, rejects a model that is quantized, group-offloaded, placed by accelerate (device_map or CPU offload), or has PEFT layers injected.
load_lora_adapter refuses a tensor-parallel model.
save_pretrained refuses a quantized TP model.

Before submitting

Who can review?

Anyone in the community is free to review the PR once the tests have passed. Feel free to tag
members/contributors who may be interested in your PR.

HuggingFaceDocBuilderDev · 2026-08-20T14:25:44Z

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.

Stream each rank's slice of a tensor-parallel checkpoint straight off disk instead of materializing the full checkpoint on every rank and resharding it afterwards, and gather the shards back on save. - `from_pretrained(..., parallel_config=TensorParallelConfig(...))` resolves the shard specs on the still-meta model, then slices each safetensors tensor before the dtype cast, so host memory peaks at ~1/tp_degree of the checkpoint. - `save_pretrained` all-gathers the DTensors into an ordinary checkpoint, or writes a distributed checkpoint with `dcp=True` so no full tensor is ever formed. The writing `tp_degree` is recorded, since a packed weight's stored layout is interleaved by it. - Factor the plan interpretation out of the Neuron pre-shard path into shared `TPShardSpec` / `resolve_tp_shard_specs` / `_local_shard` / `_hooks_only_styles` helpers, so both backends and both the load and save paths shard identically.

…ng or LoRA Addresses the remaining two items of the review on huggingface#13718: tensor parallelism was rejected alongside quantization and `device_map` only on the `from_pretrained` streaming path, while `enable_parallelism` — which the quantization error message itself recommended — accepted a quantized, offloaded or adapter-injected model and sharded it anyway. - Add `_check_tp_model_state`, called from `apply_tensor_parallel`, the one chokepoint every TP entry point funnels through. It rejects a model that is quantized, group-offloaded, placed by accelerate (`device_map` or CPU offload), or has PEFT layers injected. Placed before the device-type check so the reported reason is the useful one. - Guard the reverse order too: `enable_group_offload`, the two pipeline CPU-offload methods, and `load_lora_adapter` now refuse a tensor-parallel model. - `save_pretrained` refuses a quantized tensor-parallel model. Previously the `dcp=True` branch returned before the quantizer's serialization step, writing shards with no quantization metadata and no error. - The DCP load guard checked the `quantization_config` kwarg only, so a pre-quantized checkpoint directory loaded silently; check the config's own entry too, and add the missing `_tp_plan` check that otherwise surfaced as a raw `AttributeError`. - Correct the `from_pretrained` message and the doc sentence that pointed at `enable_parallelism` as a way to shard a quantized model. The new tests are the first tensor-parallel tests that need neither an accelerator nor more than one rank: every case asserts a raise before any collective, so they run single-process on gloo.

…sers into add-shard-ckpt-loading

Resolve conflict in `_load_pretrained_model`: keep the tensor-parallel `load_fn` branch from this PR, and drop `dduf_entries` from the ordinary branch — DDUF loading was removed upstream. The dangling `dduf_entries` references this PR added (the DCP unsupported-options list and the `_check_tp_streaming_supported` guard) go away with it, since the kwarg no longer exists. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>

…sers into add-shard-ckpt-loading

sayakpaul

Left some high-level comments.

My main comment is if we want to ship the advanced features of rank aware save and load yet. I am leaning towards raising when we encouter those situations and simplify the code a bit. This way, we can see if the community wants this feature and ship it when we have enough interest. But I would like to double-check with @DN6 on this too.

sayakpaul · 2026-09-02T11:20:41Z

-Pass a [`TensorParallelConfig`] to [`~ModelMixin.enable_parallelism`]. `tp_degree` is the number of devices to shard across and must divide the model's number of attention heads. The model must define a `_tp_plan` (a flat mapping of module-name globs to a `"colwise"`/`"rowwise"` style).
+Pass a [`TensorParallelConfig`] to the `parallel_config` argument of the model's [`~ModelMixin.from_pretrained`]. `tp_degree` is the number of devices to shard across and must divide the model's number of attention heads. The model must define a `_tp_plan` (a flat mapping of module-name globs to a `"colwise"`/`"rowwise"` style).
+
+Loading this way shards the checkpoint *while reading it*: each rank reads only its own slice of each sharded weight and places it straight onto its own device. Nothing full-size is ever materialized, so per-rank memory falls as `tp_degree` rises.


Oh really! This is very cool. Could we also present a small comparison between the loading time with and without this way of loading?

sayakpaul · 2026-09-02T11:21:41Z


 `tp_degree` is taken from `world_size` above, so `--nproc-per-node 4` shards the transformer across 4 devices.

+A tensor-parallel `parallel_config` cannot be combined with `device_map`, `quantization_config`, `low_cpu_mem_usage=False`, `use_flashpack=True`, or non-safetensors weights; each raises rather than quietly falling back to loading the full checkpoint. Tensor parallelism also cannot be combined with quantization, offloading, or LoRA adapters at all — the parameters it shards have to be plain parameters owned by the model — so those raise however the model is sharded. To shard a model that is already in memory, call [`~ModelMixin.enable_parallelism`] with the same config instead — that loads everything first and reshards it, so it costs full checkpoint memory on every rank.


Interesting that we cannot load TP with quantization. Do we know why?

sayakpaul · 2026-09-02T11:25:46Z

+### Saving a tensor-parallel model
+
+[`~ModelMixin.save_pretrained`] gathers the shards back into ordinary full tensors, so the result is a normal checkpoint that loads with or without tensor parallelism. Gathering is a collective, so call it on **every** rank; only rank 0 writes.


That is cool! However, do we have to ship this yet? I don't have any strong opinions. @DN6 WDYT?

sayakpaul · 2026-09-02T11:27:52Z

+pipeline.transformer.save_pretrained("flux2-transformer-dcp", dcp=True)
+```
+
+`from_pretrained` detects such a directory automatically, and reads it back with the same `parallel_config` you saved it under. Because a packed projection's shards are stored interleaved by the writing degree, the checkpoint only loads at that same `tp_degree`, and only with tensor parallelism — anything else raises rather than silently returning wrong weights. It is also local-only: a distributed checkpoint is recognized by the `.metadata` file in its directory, so it cannot be pushed to or loaded from the Hub. To lift any of these restrictions, re-save with the default (gathered) path, which produces an ordinary checkpoint.


Not sure if we want to ship this right away because it is introducing quite a bit of code changes for something we don't know to be impactful yet.

sayakpaul · 2026-09-02T11:31:11Z

-                    dist_param = nn.Parameter(
-                        DTensor.from_local(local, device_mesh, [Shard(0)], run_check=False),
-                        requires_grad=param.requires_grad,
+                    local = _local_shard(full, 0, _blocks_to_block_sizes(full.shape[0], blocks), device_mesh)


This reads a bit difficult. Can we break them further a bit?

sayakpaul · 2026-09-02T11:34:14Z

+    # Before the device-type check below, so that a quantized or offloaded model reports what is actually wrong with
+    # it rather than being turned away for its device type.
+    _check_tp_model_state(model)


But should we rather have them in enable_parallelism() like the validations performed by CP?

sayakpaul · 2026-09-02T11:41:40Z

+    # Before the device-type check below, so that a quantized or offloaded model reports what is actually wrong with
+    # it rather than being turned away for its device type.
+    _check_tp_model_state(model)


But should we rather have them in enable_parallelism() like the validations performed by CP?

sayakpaul · 2026-09-02T11:42:14Z

    return offload_index, state_dict_index, mismatched_keys, error_msgs


+def _load_shard_file_tp(


Pausing the review on these since they are related to save/load which we're still deciding on.

github-actions Bot added documentation Improvements or additions to documentation fixes-issue size/M PR with diff < 200 LOC labels Aug 20, 2026

github-actions Bot added models tests utils hooks size/L PR with diff > 200 LOC and removed size/M PR with diff < 200 LOC labels Aug 20, 2026

JingyaHuang force-pushed the add-shard-ckpt-loading branch from a5e135c to acdf4bf Compare August 20, 2026 16:09

JingyaHuang and others added 3 commits August 20, 2026 18:11

Merge branch 'main' into add-shard-ckpt-loading

0b9686b

Merge branch 'add-shard-ckpt-loading' of github.com:JingyaHuang/diffu…

0760934

…sers into add-shard-ckpt-loading

github-actions Bot added lora pipelines labels Aug 21, 2026

Merge branch 'main' into add-shard-ckpt-loading

842c643

JingyaHuang marked this pull request as ready for review August 26, 2026 16:19

JingyaHuang and others added 4 commits August 26, 2026 18:19

Merge branch 'main' into add-shard-ckpt-loading

949cbdc

Merge branch 'add-shard-ckpt-loading' of github.com:JingyaHuang/diffu…

6055292

…sers into add-shard-ckpt-loading

JingyaHuang requested a review from sayakpaul September 1, 2026 08:58

sayakpaul added this to Diffusers Roadmap 0.41 Sep 2, 2026

sayakpaul moved this to In Progress in Diffusers Roadmap 0.41 Sep 2, 2026

sayakpaul reviewed Sep 2, 2026

View reviewed changes

Merge branch 'main' into add-shard-ckpt-loading

fd31c34

sayakpaul requested a review from DN6 September 2, 2026 11:36

sayakpaul reviewed Sep 2, 2026

View reviewed changes

Navigation Menu

Uh oh!

[core] Shard tensor-parallel checkpoints on load and save - #14544

[core] Shard tensor-parallel checkpoints on load and save#14544
JingyaHuang wants to merge 10 commits into
huggingface:mainfrom
JingyaHuang:add-shard-ckpt-loading

JingyaHuang commented Aug 20, 2026 •

edited

Loading

Uh oh!

HuggingFaceDocBuilderDev commented Aug 20, 2026

Uh oh!

sayakpaul left a comment

Uh oh!

sayakpaul Sep 2, 2026

Uh oh!

sayakpaul Sep 2, 2026

Uh oh!

sayakpaul Sep 2, 2026

Uh oh!

sayakpaul Sep 2, 2026

Uh oh!

sayakpaul Sep 2, 2026

Uh oh!

sayakpaul Sep 2, 2026

Uh oh!

sayakpaul Sep 2, 2026

Uh oh!

sayakpaul Sep 2, 2026

Uh oh!

sayakpaul Sep 2, 2026

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants


		`tp_degree` is taken from `world_size` above, so `--nproc-per-node 4` shards the transformer across 4 devices.

		A tensor-parallel `parallel_config` cannot be combined with `device_map`, `quantization_config`, `low_cpu_mem_usage=False`, `use_flashpack=True`, or non-safetensors weights; each raises rather than quietly falling back to loading the full checkpoint. Tensor parallelism also cannot be combined with quantization, offloading, or LoRA adapters at all — the parameters it shards have to be plain parameters owned by the model — so those raise however the model is sharded. To shard a model that is already in memory, call [`~ModelMixin.enable_parallelism`] with the same config instead — that loads everything first and reshards it, so it costs full checkpoint memory on every rank.

		### Saving a tensor-parallel model

		[`~ModelMixin.save_pretrained`] gathers the shards back into ordinary full tensors, so the result is a normal checkpoint that loads with or without tensor parallelism. Gathering is a collective, so call it on every rank; only rank 0 writes.

		return offload_index, state_dict_index, mismatched_keys, error_msgs


		def _load_shard_file_tp(

Uh oh!

Conversation

JingyaHuang commented Aug 20, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

What does this PR do?

Before submitting

Who can review?

Uh oh!

HuggingFaceDocBuilderDev commented Aug 20, 2026

Uh oh!

sayakpaul left a comment

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants

JingyaHuang commented Aug 20, 2026 •

edited

Loading