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

Here's a quick benchmark:

Synthetic checkpoint (FLUX.1-shaped, 1.33B params, 2.65GB bf16):

mode tp_degree load time peak GPU/rank peak CPU/rank

streamed (from_pretrained(parallel_config=...)) 2 1.92s ± 0.11 1.71GB 2.70GB

materialize+reshard (from_pretrained + enable_parallelism) 2 3.06s ± 0.02 1.69GB 4.08GB

streamed (from_pretrained(parallel_config=...)) 4 1.45s ± 0.03 1.23GB 2.20GB

materialize+reshard (from_pretrained + enable_parallelism) 4 3.29s ± 0.02 1.23GB 4.08GB

Speed: streaming is ~37% faster at tp_degree=2, ~56% faster at tp_degree=4 on the synthetic checkpoint (1.92s vs 3.06s, 1.45s vs 3.29s).

GPU memory: basically the same between the two.

CPU/host memory: 34% less at tp_degree=2, 46% less at tp_degree=4 (2.70GB vs 4.08GB, 2.20GB vs 4.08GB).

Okay nice. So, it's the speed that matters for now. Let's include these numbers in the docs, then?

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?

It's mostly about scoping... I didn't want to handle "quantization + TP" sharding in the scale of this PR. IMO, it could be possible if we shard and then quantize, and then according to the quantization configs and the backend, some might work, some doesn't... I would rather raise for now, we could consider supporting the combo properly if the community shows interest.

Yeah raising is totally fine. Maybe we should put this blob of text into a "> [!CAUTION]" block?

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

Thanks for keeping the PR very scoped and easy to review. Left some feedback which should be straightforward to resolve.

Ran the tests on HF Jobs, too, and they pass: https://huggingface.co/jobs/sayakpaul/6aba42356b030d633f69c66d

sayakpaul · 2026-09-28T09:12:45Z

-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.


Okay nice. So, it's the speed that matters for now. Let's include these numbers in the docs, then?

sayakpaul · 2026-09-28T09:17:51Z


 `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.


Yeah raising is totally fine. Maybe we should put this blob of text into a "> [!CAUTION]" block?

sayakpaul · 2026-09-28T09:18:12Z


+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 isn't supported yet, and [`~ModelMixin.save_pretrained`] raises on one. Save the model before sharding it.


sayakpaul · 2026-09-28T09:36:21Z

+    for block_size in block_sizes:
+        # An uneven split is rejected rather than silently handed to `Shard`, which pads the tail
+        # and would break the paired colwise/rowwise matmul.
+        if block_size % tp_size != 0:
+            raise ValueError(
+                f"Cannot shard a block of size {block_size} across {tp_size} tensor-parallel ranks: "
+                f"{block_size} is not divisible by {tp_size}."
+            )


Let's raise this error earlier.

sayakpaul · 2026-09-28T09:38:31Z

-                blocks = _blocks if _blocks is not None else getattr(module, "_tp_packed_col_blocks")
-                rank = device_mesh.get_local_rank()
-                tp_size = device_mesh.size()
+                blocks = _blocks if _blocks is not None else module._tp_packed_col_blocks


Should we not raise if module._tp_packed_col_blocks doesn't have anything?

sayakpaul · 2026-09-28T09:56:32Z

+            non_safetensors = [f for f in resolved_model_file if not str(f).endswith(".safetensors")]
+            if non_safetensors:
+                raise ValueError(
+                    f"A tensor-parallel `parallel_config` requires safetensors weights, so that each rank can "
+                    f"read only its own slice of each tensor. Got {non_safetensors}."
+                )


Is this raise sufficiently early in the stack?

sayakpaul · 2026-09-28T09:58:15Z

        model.eval()

-        if parallel_config is not None:
+        if tp_shard_specs is not None:


I think it's better to keep the same conditionals depend on the same variable, readability wise. So, we could keep parallel_config and then condition on tp_shards_specs?

sayakpaul · 2026-09-28T10:00:21Z

+    return resolved
+
+
+def _check_tp_supported(model_name: str, tp_plan: "dict | None", num_heads: "int | None", tp_config) -> None:


I feel like the validation methods are a bit scattered along different files. It makes sense to centralize them.

sayakpaul · 2026-09-28T10:02:08Z

            )

+    @classmethod
+    def _check_tp_streaming_supported(


This could probably go to check_tp_supported and check_tp_supported_model_state could be consolidated into it?

sayakpaul · 2026-09-28T10:03:07Z

+        return self.linear_2(self.linear_1(hidden_states))
+
+
+@pytest.fixture(scope="module")


We don't test these methods. Let's remove this test.

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 Sep 2, 2026

sayakpaul moved this to In Progress in Diffusers Roadmap 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

Comment thread src/diffusers/hooks/tensor_parallel.py Outdated

Comment thread src/diffusers/models/model_loading_utils.py

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

9511143

JingyaHuang mentioned this pull request Sep 10, 2026

[TPU] TorchTPU backend integration - eager / torch.compile / tp #14039

Open

DN6 added this to the Release 0.41.0 milestone Sep 17, 2026

JingyaHuang and others added 3 commits September 19, 2026 11:06

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

53fd823

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

4112d18

review: remove tp save/ dcp related

f741782

github-actions Bot removed the utils label Sep 23, 2026

JingyaHuang and others added 3 commits September 23, 2026 15:47

review: add a new helper .

1b4d1d7

review: moved the TP checks into enable_parallelism

4e9a927

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

0cca9a3

JingyaHuang requested a review from sayakpaul September 24, 2026 13:29

JingyaHuang mentioned this pull request Sep 24, 2026

[core] Tensor parallelism for Qwen-Image-2.1 #14865

Open

13 tasks

sayakpaul reviewed Sep 28, 2026

View reviewed changes

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

7ea8b97

mode	tp_degree	load time	peak GPU/rank	peak CPU/rank
streamed (`from_pretrained(parallel_config=...)`)	2	1.92s ± 0.11	1.71GB	2.70GB
materialize+reshard (`from_pretrained` + `enable_parallelism`)	2	3.06s ± 0.02	1.69GB	4.08GB
streamed (`from_pretrained(parallel_config=...)`)	4	1.45s ± 0.03	1.23GB	2.20GB
materialize+reshard (`from_pretrained` + `enable_parallelism`)	4	3.29s ± 0.02	1.23GB	4.08GB


		`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.


		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 isn't supported yet, and [`~ModelMixin.save_pretrained`] raises on one. Save the model before sharding it.

		return resolved


		def _check_tp_supported(model_name: str, tp_plan: "dict \| None", num_heads: "int \| None", tp_config) -> None:

		return self.linear_2(self.linear_1(hidden_states))


		@pytest.fixture(scope="module")

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!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

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!

Choose a reason for hiding this comment

Uh oh!

This branch has not been deployed

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

4 participants

JingyaHuang commented Aug 20, 2026 •

edited

Loading