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

JingyaHuang · 2026-09-29T15:38:08Z

Just validated the changes on cuda / neuron / tpu.

sayakpaul

Just a few last comments and I think we're good after that!

sayakpaul · 2026-09-30T05:04:07Z

+    from torch.distributed.tensor import Replicate, distribute_tensor
    from torch.distributed.tensor.parallel import ColwiseParallel, RowwiseParallel

    def _make_packed_col(marker: PackedColwiseParallel) -> ColwiseParallel:


I like the separation here.

sayakpaul

Thanks, last set of comment, one being major.

sayakpaul · 2026-10-01T02:36:31Z

+| tp_degree | method | load time | peak CPU/rank |
+|---|---|---|---|
+| 4 | `from_pretrained(parallel_config=...)` | 12.5s | 6.8GB |
+| 4 | `from_pretrained` + `enable_parallelism` | 30.4s | 64.1GB |


Am I reading it right that it leads to a decrease of ~10x in the peak CPU memory per rank? 😳

Just reproduced it myself. Wow, this is very cool!

sayakpaul · 2026-10-01T02:44:30Z

    def test_tensor_parallel_batch_inputs(self):
        self.test_tensor_parallel_inference(batch_size=2)

+    def _tp_checkpoint_and_reference(self, tmp_path, world_size):


Hmm this is good but this doesn't ensure that the model checkpoint is saved with sharded state dict files, though. This is something that reflects real-world use cases.

We usually test tiny models so the sharding is not obvious. To do it explicitly, we need to specify the shard_size argument:

diffusers/tests/models/testing_utils/common.py

Line 657 in c60830e

def test_sharded_checkpoints(self, base_model_output, tmp_path, atol=1e-5, rtol=0):

I pushed some updates to in 95c44cc. Hope, that is okay with you :) I ran the tests and they passed.

sayakpaul · 2026-10-01T02:46:22Z

+            sharded = {k: v for k, v in model.state_dict().items() if isinstance(v, DTensor)}
+            assert sharded, "No parameter was sharded into a DTensor by the streaming load."


Can be done in a follow-up but can we go even specific to assert DTensor on the keys we expect to be sharded based on the _tp_plan?

Sure, will improve the test in #14039

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

JingyaHuang and others added 12 commits September 28, 2026 16:54

review: add loading time numbers

e3fdafd

review: move tp size dividende check

8a6ff80

review: delete test

64c9e86

review: consolidate tp checks

b52a272

review: consolidate the tp checks, all in tensor_parallel.py

068d7ab

review: improve conditionals

6a56929

review: debug info

85dc21b

review: check weights format before instantiating the model

aad9e9e

review: restore duplicated _find_mismatched_keys since irrelevant

fc76c2c

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

cb01688

revert: delete duplicated helper

30d4184

review: remove the dead duplicate _find_mismatched_keys again

8f577c6

JingyaHuang requested a review from sayakpaul September 29, 2026 15:37

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

8b79c5b

sayakpaul reviewed Sep 30, 2026

View reviewed changes

JingyaHuang and others added 4 commits September 30, 2026 14:13

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

2123a63

review: update doc for loading benchmark of flux2-dev

55e6318

review: improve the raising check

5eb09a4

review: add an args section

b66b9c5

JingyaHuang requested a review from sayakpaul September 30, 2026 14:55

JingyaHuang mentioned this pull request Sep 30, 2026

Allow tensor parallelism and context parallelism in one ParallelConfig #14725

Open

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

68560c4

sayakpaul reviewed Oct 1, 2026

View reviewed changes

Test tensor-parallel loading from sharded checkpoints

95c44cc

sayakpaul approved these changes Oct 1, 2026

View reviewed changes

sayakpaul merged commit 863092f into huggingface:main Oct 1, 2026
40 of 41 checks passed

sayakpaul added the performance Anything related to performance improvements, profiling and benchmarking label Oct 1, 2026

JingyaHuang deleted the add-shard-ckpt-loading branch October 1, 2026 13:08

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

		sharded = {k: v for k, v in model.state_dict().items() if isinstance(v, DTensor)}
		assert sharded, "No parameter was sharded into a DTensor by the streaming load."

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!

Uh oh!

Uh oh!

sayakpaul Sep 2, 2026

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

JingyaHuang commented Sep 29, 2026

Uh oh!

sayakpaul left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

sayakpaul Sep 30, 2026

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Uh oh!

sayakpaul left a comment

Choose a reason for hiding this comment

Uh oh!

sayakpaul Oct 1, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

sayakpaul Oct 1, 2026

Choose a reason for hiding this comment

Uh oh!

sayakpaul Oct 1, 2026

Choose a reason for hiding this comment

Uh oh!

sayakpaul Oct 1, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

sayakpaul Oct 1, 2026

Choose a reason for hiding this comment

Uh oh!

JingyaHuang Oct 1, 2026

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

4 participants

JingyaHuang commented Aug 20, 2026 •

edited

Loading

sayakpaul Oct 1, 2026 •

edited

Loading

sayakpaul Oct 1, 2026 •

edited

Loading