FSDP x TP Saving / Loading by michaelbenayoun · Pull Request #48802 · huggingface/transformers

michaelbenayoun · 2026-09-14T16:55:04Z

What does this PR do?

This PR handles a few thing related to distributed training.

1. `_StridedShard` placement issue

The _StridedShard placement own multiple disjoint regions of a tensor.
It produces overlapping checkpoint metadata, causing failure because the default PyTorch checkpoint adapter was designed with the idea that each rank's shard is a single contiguous chunk.

The propose solution consists in redistributing _StridedShard -> Shard. While it forces additional communication and the temporary allocation of the full tensor, it has the advantage to be a super lightweight solution.

2. Distributed checkpointing / Consolidation

Load and save distributed checkpoints for both the model.

3. Optimizer checkpointing

Handled in #49256

Script

"""
torchrun --nproc_per_node=4 tp_fsdp_overfit.py

The script overfit one sentence following the steps:
    - Train first half using FSDP=2+TP=2
    - Save the model and optimizer in distributed checkpoint
    - Reload the model and optimizer from the distributed checkpoint
    - Train the rest in TP=4 (change distributed config)
    - Save the model and optimizer in distributed checkpoint
    - Reload the model in a single safetensors file.
    - Do inference in TP=4  and assert greedy generation reproduces the sentence verbatim
"""

import os

import torch
from torch.distributed.tensor import DTensor

from transformers import AutoModelForCausalLM, AutoTokenizer
from transformers.distributed import DistributedConfig
from transformers.distributed.utils import (
    clip_grad_norm_,
    load_optimizer_distributed,
    save_optimizer_distributed,
)


def create_optimizer(model):
    dtensor_params, tensor_params = [], []
    for parameter in model.parameters():
        (dtensor_params if isinstance(parameter, DTensor) else tensor_params).append(parameter)
    groups = []
    if dtensor_params:
        groups.append({"params": dtensor_params})
    if tensor_params:
        groups.append({"params": tensor_params})
    return torch.optim.AdamW(groups, lr=1e-3, foreach=True)


NAME = "Isotonic/TinyMixtral-4x248M-MoE"
TEXT = "In a quiet village nestled between rolling hills and a slow river, the autumn mornings arrived with mist that hung low over the fields and a sky that turned from grey to pale gold as the sun climbed."
CKPT = "./checkpoints"
PADDED_MODEL = "./padded_model"
OPT = os.path.join(CKPT, "optimizer")
STEPS = 10
HALF = STEPS // 2

rank, local_rank = int(os.environ["RANK"]), int(os.environ["LOCAL_RANK"])
torch.cuda.set_device(local_rank)
torch.distributed.init_process_group("nccl", device_id=torch.device(f"cuda:{local_rank}"))

tokenizer = AutoTokenizer.from_pretrained(NAME)
ids = tokenizer(TEXT, return_tensors="pt").input_ids.to(f"cuda:{local_rank}")

# Pad before applying TP: both TP=2 and TP=4 need a divisible lm_head size.
if rank == 0:
    model = AutoModelForCausalLM.from_pretrained(NAME, dtype=torch.bfloat16, device_map="cpu")
    model.resize_token_embeddings(pad_to_multiple_of=4)
    model.save_pretrained(PADDED_MODEL)
    del model
torch.distributed.barrier()

# Train first half, distributed-save model + optimizer.
model = AutoModelForCausalLM.from_pretrained(
    PADDED_MODEL,
    distributed_config=DistributedConfig(tp_size=2, fsdp_size=2, enable_sequence_parallel=True),
    dtype=torch.bfloat16,
)
optimizer = create_optimizer(model)
model.train()
for step in range(0, HALF):
    loss = model(ids, labels=ids).loss
    loss.backward()
    total_norm = clip_grad_norm_(model.parameters(), max_norm=1.0)
    optimizer.step()
    optimizer.zero_grad()
    if rank == 0:
        print(f"step {step:>2} | loss {loss.item():.5f} grad norm {total_norm.item():.5f}")

model.save_pretrained(CKPT, distributed_checkpoint=True)
save_optimizer_distributed(model, optimizer, OPT)
del model, optimizer
torch.cuda.empty_cache()

# Reload model + optimizer from the distributed checkpoint, train the rest.

model = AutoModelForCausalLM.from_pretrained(
    CKPT,
    distributed_config=DistributedConfig(tp_size=4, enable_sequence_parallel=True),
    dtype=torch.bfloat16,
)
optimizer = create_optimizer(model)
load_optimizer_distributed(model, optimizer, OPT)

model.train()
for step in range(HALF, STEPS):
    loss = model(ids, labels=ids).loss
    loss.backward()
    total_norm = clip_grad_norm_(model.parameters(), max_norm=1.0)
    optimizer.step()
    optimizer.zero_grad()
    if rank == 0:
        print(f"step {step:>2} | loss {loss.item():.5f} grad norm {total_norm.item():.5f}")

# INFERENCE in TP=4
model.save_pretrained(CKPT)
save_optimizer_distributed(model, optimizer, OPT + "_tp4")
del model, optimizer
torch.cuda.empty_cache()

model = AutoModelForCausalLM.from_pretrained(
    CKPT,
    distributed_config=DistributedConfig(tp_size=4),
    dtype=torch.bfloat16,
)
model.eval()
prompt = tokenizer("In a quiet village", return_tensors="pt").to(f"cuda:{local_rank}")
out = model.generate(**prompt, max_new_tokens=ids.shape[-1] - prompt.input_ids.shape[-1], do_sample=False)

got, want = out[0].tolist(), ids[0].tolist()
if rank == 0:
    print(f"generated: {tokenizer.decode(got, skip_special_tokens=True)!r}")
    print(f"expected: {tokenizer.decode(want, skip_special_tokens=True)!r}")
assert got == want, (
    f"generation mismatch at index {next((i for i, (g, e) in enumerate(zip(got, want)) if g != e), -1)}"
)

torch.distributed.destroy_process_group()

And the output is:

warning: No `requires-python` value found in the workspace. Defaulting to `>=3.13`.
W0916 18:52:21.262000 454691 torch/distributed/run.py:874]
W0916 18:52:21.262000 454691 torch/distributed/run.py:874] *****************************************
W0916 18:52:21.262000 454691 torch/distributed/run.py:874] Setting OMP_NUM_THREADS environment variable for each process to be 1 in default, to avoid your system being overloaded, please further tune the variable for optimal performance in your application as needed.
W0916 18:52:21.262000 454691 torch/distributed/run.py:874] *****************************************
Loading weights: 100%|█████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 111/111 [00:00<00:00, 135.09it/s]
[transformers] The new embeddings will be initialized from a multivariate normal distribution that has old embeddings' mean and covariance. As described in this article: https://nlp.stanford.edu/~johnhew/vocab-expansion.html. To disable this, use `mean_resizing=False`
[transformers] The new lm_head weights will be initialized from a multivariate normal distribution that has old embeddings' mean and covariance. As described in this article: https://nlp.stanford.edu/~johnhew/vocab-expansion.html. To disable this, use `mean_resizing=False`
Writing model shards: 100%|█████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 1/1 [00:02<00:00,  2.63s/it]
Loading weights: 100%|████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 111/111 [00:00<00:00, 5646.87it/s]
Loading weights: 100%|████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 111/111 [00:00<00:00, 5430.94it/s]
Loading weights: 100%|████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 111/111 [00:00<00:00, 5556.63it/s]

/home/michael/transformers/.venv/lib/python3.13/site-packages/torch/distributed/fsdp/_fully_shard/_fsdp_state.py:331: UserWarning: FSDP2-wrapped module (FSDPMixtralRMSNorm, FSDPLinear) returned a view tensor. An in-place op on this view (e.g., `x += y`) will silently drop the pre-backward hook and skip the all-gather, which can cause backward to fail or produce wrong gradients. Use out-of-place ops (`out = out + y`, not `out += y`) or `.clone()` the output before any in-place op.
  output = self._register_pre_backward_hook(output)
/home/michael/transformers/.venv/lib/python3.13/site-packages/torch/distributed/fsdp/_fully_shard/_fsdp_state.py:331: UserWarning: FSDP2-wrapped module (FSDPMixtralRMSNorm, FSDPLinear) returned a view tensor. An in-place op on this view (e.g., `x += y`) will silently drop the pre-backward hook and skip the all-gather, which can cause backward to fail or produce wrong gradients. Use out-of-place ops (`out = out + y`, not `out += y`) or `.clone()` the output before any in-place op.
  output = self._register_pre_backward_hook(output)
/home/michael/transformers/.venv/lib/python3.13/site-packages/torch/distributed/fsdp/_fully_shard/_fsdp_state.py:331: UserWarning: FSDP2-wrapped module (FSDPMixtralRMSNorm, FSDPLinear) returned a view tensor. An in-place op on this view (e.g., `x += y`) will silently drop the pre-backward hook and skip the all-gather, which can cause backward to fail or produce wrong gradients. Use out-of-place ops (`out = out + y`, not `out += y`) or `.clone()` the output before any in-place op.
  output = self._register_pre_backward_hook(output)
/home/michael/transformers/.venv/lib/python3.13/site-packages/torch/distributed/fsdp/_fully_shard/_fsdp_state.py:331: UserWarning: FSDP2-wrapped module (FSDPMixtralRMSNorm, FSDPLinear) returned a view tensor. An in-place op on this view (e.g., `x += y`) will silently drop the pre-backward hook and skip the all-gather, which can cause backward to fail or produce wrong gradients. Use out-of-place ops (`out = out + y`, not `out += y`) or `.clone()` the output before any in-place op.
  output = self._register_pre_backward_hook(output)
/home/michael/transformers/.venv/lib/python3.13/site-packages/torch/distributed/fsdp/_fully_shard/_fsdp_state.py:331: UserWarning: FSDP2-wrapped module (FSDPMixtralForCausalLM) returned a view tensor. An in-place op on this view (e.g., `x += y`) will silently drop the pre-backward hook and skip the all-gather, which can cause backward to fail or produce wrong gradients. Use out-of-place ops (`out = out + y`, not `out += y`) or `.clone()` the output before any in-place op.
  output = self._register_pre_backward_hook(output)
/home/michael/transformers/.venv/lib/python3.13/site-packages/torch/distributed/fsdp/_fully_shard/_fsdp_state.py:331: UserWarning: FSDP2-wrapped module (FSDPMixtralForCausalLM) returned a view tensor. An in-place op on this view (e.g., `x += y`) will silently drop the pre-backward hook and skip the all-gather, which can cause backward to fail or produce wrong gradients. Use out-of-place ops (`out = out + y`, not `out += y`) or `.clone()` the output before any in-place op.
  output = self._register_pre_backward_hook(output)
/home/michael/transformers/.venv/lib/python3.13/site-packages/torch/distributed/fsdp/_fully_shard/_fsdp_state.py:331: UserWarning: FSDP2-wrapped module (FSDPMixtralForCausalLM) returned a view tensor. An in-place op on this view (e.g., `x += y`) will silently drop the pre-backward hook and skip the all-gather, which can cause backward to fail or produce wrong gradients. Use out-of-place ops (`out = out + y`, not `out += y`) or `.clone()` the output before any in-place op.
  output = self._register_pre_backward_hook(output)
/home/michael/transformers/.venv/lib/python3.13/site-packages/torch/distributed/fsdp/_fully_shard/_fsdp_state.py:331: UserWarning: FSDP2-wrapped module (FSDPMixtralForCausalLM) returned a view tensor. An in-place op on this view (e.g., `x += y`) will silently drop the pre-backward hook and skip the all-gather, which can cause backward to fail or produce wrong gradients. Use out-of-place ops (`out = out + y`, not `out += y`) or `.clone()` the output before any in-place op.
  output = self._register_pre_backward_hook(output)
step  0 | loss 4.97490 grad norm 10.75000
step  1 | loss 1.45883 grad norm 5.71875
step  2 | loss 0.23285 grad norm 2.67188
step  3 | loss 0.21926 grad norm 3.20312
step  4 | loss 0.38996 grad norm 10.81250
[rank3]:W0916 18:52:40.667000 454786 torch/distributed/tensor/_redistribute.py:390] While redistributing from (_StridedShard(dim=0, sf=2), Shard(dim=0)) to (Shard(dim=0), Shard(dim=0)), 2 sequential all_gather operations will be performed. This is suboptimal: multiple collective operations have higher latency (separate kernel launches and synchronization points) and may give inconsistent results between ranks due to different reduction orders. it is not possible to merge non-ascending order all_gather operations.
[rank2]:W0916 18:52:40.668000 454785 torch/distributed/tensor/_redistribute.py:390] While redistributing from (_StridedShard(dim=0, sf=2), Shard(dim=0)) to (Shard(dim=0), Shard(dim=0)), 2 sequential all_gather operations will be performed. This is suboptimal: multiple collective operations have higher latency (separate kernel launches and synchronization points) and may give inconsistent results between ranks due to different reduction orders. it is not possible to merge non-ascending order all_gather operations.
[rank1]:W0916 18:52:40.669000 454784 torch/distributed/tensor/_redistribute.py:390] While redistributing from (_StridedShard(dim=0, sf=2), Shard(dim=0)) to (Shard(dim=0), Shard(dim=0)), 2 sequential all_gather operations will be performed. This is suboptimal: multiple collective operations have higher latency (separate kernel launches and synchronization points) and may give inconsistent results between ranks due to different reduction orders. it is not possible to merge non-ascending order all_gather operations.
[rank0]:W0916 18:52:40.674000 454783 torch/distributed/tensor/_redistribute.py:390] While redistributing from (_StridedShard(dim=0, sf=2), Shard(dim=0)) to (Shard(dim=0), Shard(dim=0)), 2 sequential all_gather operations will be performed. This is suboptimal: multiple collective operations have higher latency (separate kernel launches and synchronization points) and may give inconsistent results between ranks due to different reduction orders. it is not possible to merge non-ascending order all_gather operations.
Loading weights: 100%|████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 111/111 [00:00<00:00, 1784.55it/s]
Loading weights: 100%|████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 111/111 [00:00<00:00, 1709.28it/s]
Loading weights: 100%|████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 111/111 [00:00<00:00, 1762.23it/s]
Loading weights: 100%|████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 111/111 [00:00<00:00, 1561.45it/s]
step  5 | loss 0.00217 grad norm 0.06689
step  6 | loss 0.00039 grad norm 0.01117
step  7 | loss 0.00025 grad norm 0.00674
step  8 | loss 0.00023 grad norm 0.00922
step  9 | loss 0.00008 grad norm 0.00201
Writing model shards: 100%|█████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 1/1 [00:02<00:00,  2.96s/it]
Loading weights: 100%|████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 111/111 [00:00<00:00, 1942.67it/s]
Loading weights: 100%|████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 111/111 [00:00<00:00, 1496.11it/s]
Loading weights: 100%|████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 111/111 [00:00<00:00, 1503.06it/s]
Loading weights: 100%|████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 111/111 [00:00<00:00, 1467.24it/s]
generated: 'In a quiet village nestled between rolling hills and a slow river, the autumn mornings arrived with mist that hung low over the fields and a sky that turned from grey to pale gold as the sun climbed.'
expected: 'In a quiet village nestled between rolling hills and a slow river, the autumn mornings arrived with mist that hung low over the fields and a sky that turned from grey to pale gold as the sun climbed.'

HuggingFaceDocBuilderDev · 2026-09-14T17:09:30Z

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.

3outeille · 2026-09-18T10:31:29Z

    # being emitted if the function is not used
    import torch.distributed.checkpoint as dcp
-    from torch.distributed.checkpoint.state_dict import get_optimizer_state_dict
+    from torch.distributed.checkpoint.state_dict import StateDictOptions, get_optimizer_state_dict


same import top + guarding distributed as well. Make sure as well to check torch version availabilities for the func

But the comment at the top mentions that it is imported here on purpose to avoid warnings.

okay then add guarding. We have to make sure the user downloaded a torch version that has DCP

The guard is already there, line 514.

michaelbenayoun

I've addressed your comments. WDYT?

3outeille · 2026-09-29T18:22:28Z

-            transformers_explicit_filename=getattr(config, "transformers_weights", None),
-            tqdm_class=tqdm_class,
-        )
+        # Local checkpoints saved with `save_pretrained(..., distributed_checkpoint=True)` hold rank-local shards


moving that in mixin ?

At this point, the model is not even created so it is not even distributed, hence not a DistributedMixin.
I could, of course, create functions or static methods, but is it worth it?

We have 2 blocks for this logic in from_pretrained.

I was thinking to have something like:

distributed_checkpoint_dir = get_distributed_checkpoint_dir(pretrained_model_name_or_path, subfolder, gguf_file, state_dict) if distributed_checkpoint_dir is not None: checkpoint_files, sharded_metadata = _get_resolved_distributed_checkpoint_files(distributed_checkpoint_dir, hf_quantizer) else: checkpoint_files, sharded_metadata = _get_resolved_checkpoint_files( pretrained_model_name_or_path=pretrained_model_name_or_path, ... )

3outeille · 2026-09-29T18:22:35Z

        )
-        loading_info, disk_offload_index = cls._load_pretrained_model(model, state_dict, checkpoint_files, load_config)
-        loading_info = cls._finalize_model_loading(model, load_config, loading_info)
+        if distributed_checkpoint_dir is not None:


moving that in mixin ?

same here we can do like above

if distributed_checkpoint_dir is not None: loading_info, disk_offload_index = cls._load_and_finalize_distributed_pretrained_model( model, distributed_checkpoint_dir, load_config ) else: loading_info, disk_offload_index = cls._load_pretrained_model( model, state_dict, checkpoint_files, load_config ) loading_info = cls._finalize_model_loading(model, load_config, loading_info)

3outeille · 2026-09-29T18:33:14Z

+        _load_sharded_checkpoint_in_distributed_model(model, checkpoint_dir, strict=strict)
+    elif is_sharded_checkpoint(os.path.join(checkpoint_dir, "sharded")):
+        _load_sharded_checkpoint_in_distributed_model(model, os.path.join(checkpoint_dir, "sharded"), strict=strict)
+    elif os.path.isfile(safe_index_file):


why is there 2 paths when checkpoint is consolidate ?

It depends if the file is one safetensors file or a bunch of shards.
Refactored to separate the logic of getting the filenames and calling _load_sharded_checkpoint_in_distributed_model.

ArthurZucker · 2026-10-05T05:33:50Z

Sorry a bit late, reviewing now!

3outeille · 2026-10-06T01:14:05Z

    import torch


+def _check_distributed_checkpointing_available(raise_if_not: bool = True) -> bool:


i prefer not having the raise_if_not boolean and have the error raise happening where this is called. If we keep it, we find ourselves with pattern in the code like

if not _check_distributed_checkpointing_available(raise_if_not=False): raise OSError("Loading a distributed checkpoint requires torch>=2.7.")

which is not consistent (as there is a double raise but one was turn off)

3outeille · 2026-10-06T01:19:06Z

+    from transformers.distributed import DistributedConfig
+    from transformers.distributed.checkpoint import load_model_checkpoint_distributed
+
+    if dist.is_available():


use is_torch_distributed_available() as it uses across the codebase

github-actions · 2026-10-06T01:28:10Z

CI recap

Dashboard: View test results in Grafana
Latest run: 37397912686:1
Result: success | Jobs: 16 | Tests: 193,976 | Failures: 0 | Duration: 13h 48m

3outeille · 2026-10-06T06:04:45Z

-                token=token,
-                **hub_kwargs,
+                # Native DCP checkpoints remain sharded, use `distributed_checkpoint=False` for an interoperable checkpoint.
+                consolidate=False,


we should have an utils function called consolidate_distributed_checkpoint() in transformers.distributed.checkpoint

3outeille · 2026-10-06T06:36:35Z

+
+
+def _prepare_state_dict_for_dcp(state_dict):
+    """


the docstring example doesnt make it obvious to understand what's going on (especially with __create_write_items__, __create_chunk_list__ and __get_tensor_shard__)

3outeille · 2026-10-06T06:53:06Z

-            transformers_explicit_filename=getattr(config, "transformers_weights", None),
-            tqdm_class=tqdm_class,
-        )
+        # Local checkpoints saved with `save_pretrained(..., distributed_checkpoint=True)` hold rank-local shards


I was thinking to have something like:

distributed_checkpoint_dir = get_distributed_checkpoint_dir(pretrained_model_name_or_path, subfolder, gguf_file, state_dict) if distributed_checkpoint_dir is not None: checkpoint_files, sharded_metadata = _get_resolved_distributed_checkpoint_files(distributed_checkpoint_dir, hf_quantizer) else: checkpoint_files, sharded_metadata = _get_resolved_checkpoint_files( pretrained_model_name_or_path=pretrained_model_name_or_path, ... )

3outeille · 2026-10-06T06:57:49Z

        )
-        loading_info, disk_offload_index = cls._load_pretrained_model(model, state_dict, checkpoint_files, load_config)
-        loading_info = cls._finalize_model_loading(model, load_config, loading_info)
+        if distributed_checkpoint_dir is not None:


same here we can do like above

if distributed_checkpoint_dir is not None: loading_info, disk_offload_index = cls._load_and_finalize_distributed_pretrained_model( model, distributed_checkpoint_dir, load_config ) else: loading_info, disk_offload_index = cls._load_pretrained_model( model, state_dict, checkpoint_files, load_config ) loading_info = cls._finalize_model_loading(model, load_config, loading_info)

3outeille · 2026-10-06T07:21:35Z

+    safe_index_file = os.path.join(checkpoint_dir, SAFE_WEIGHTS_INDEX_NAME)
+    safe_weights_file = os.path.join(checkpoint_dir, SAFE_WEIGHTS_NAME)
+
+    for dcp_dir in (checkpoint_dir, os.path.join(checkpoint_dir, "sharded")):


since we hardcode the save_distributed_checkpoint path to consolidate=False, we only have 1 case which is the non-consolidated DCP checkpoint:

self.save_distributed_checkpoint( model_to_save, save_directory, # Native DCP checkpoints remain sharded, use `distributed_checkpoint=False` for an interoperable checkpoint. consolidate=False, )

no need the for loop then

3outeille · 2026-10-06T07:32:34Z

+            return
+
+    checkpoint_files = []
+    if os.path.isfile(safe_index_file):


do we really need this part ? i dont see how this part can be triggered with the current save_pretrained(distributed_checkpoint=True|False)/from_pretrained() API .

3outeille · 2026-10-06T07:33:35Z

+    set_model_state_dict(model, state)
+
+
+def load_model_checkpoint_distributed(model, checkpoint_dir: str | os.PathLike, strict: bool = True) -> None:


I think we want load_model_checkpoint_distributed and save_model_checkpoint_distributed to be private function so that users dont call those functions and rely on save_pretrained(distributed_checkpoint=True|False)/from_pretrained() API only

3outeille · 2026-10-06T07:35:49Z

We are getting there !

Main points:

consolidate_distributed_checkpoint()in transformers.distributed.checkpoint
up-to-date script for PR description
- 1. train + save dcp
- 1. resume from the DCP checkpoint (same process group), train, save dcp.
- 1. consolidate with consolidate_distributed_checkpoint()
- 1. generate in different distributed layout than training
better explanation of doctstring
move things to mixin
remove potential code paths that are not triggered anymore

michaelbenayoun added 6 commits September 14, 2026 14:00

wip: add script to track potential changes

4454e4b

feat: workaround for _StridedShard

ec979c4

feat: handle optimizer loading / saving

6d9a239

wip: add clip_grad_norm_

6868815

doc: add docstring to clip_grad_norm_

f20be8d

Merge branch 'main' into fsdp_tp_saving

fec373e

michaelbenayoun commented Sep 14, 2026

View reviewed changes

Comment thread tp_fsdp_overfit.py Outdated

michaelbenayoun changed the title ~~Fsdp tp saving~~ FSDP x TP Saving / Loading Sep 15, 2026

michaelbenayoun added 2 commits September 15, 2026 12:53

feat: uniformize the clip_grad_norm_ implem

f9341a8

refactor: use the proper function for clip_grad_norm_

b09c7ee

michaelbenayoun mentioned this pull request Sep 15, 2026

Align Accelerator.clip_grad_norm_ with Transformers huggingface/accelerate#4266

Open

michaelbenayoun added 6 commits September 15, 2026 14:29

feat: save shards / consolidate utility functions

9f677ef

feat: save / load sharded / consolidated optimizer states

9c742be

feat: add checkpoint state check

dc192d9

Merge branch 'main' into fsdp_tp_saving

24558d8

feat: add distributed checkpoint loading

b271952

test: add saving, loading and clipping tests

6a5075c

michaelbenayoun mentioned this pull request Sep 15, 2026

FSDP x TP Saving / Loading, Planner / Writer based solution #48830

Closed

michaelbenayoun added 3 commits September 16, 2026 10:02

wip: remove .bin support

78d55a4

doc: update docstring

607d7f9

chore: remove testing script

55f0814

3outeille requested review from 3outeille and ArthurZucker September 16, 2026 09:45

3outeille reviewed Sep 18, 2026

View reviewed changes

michaelbenayoun added 5 commits September 18, 2026 14:54

Merge branch 'main' into fsdp_tp_saving

b5890fe

align: match planner branch loader surface

6724f5a

fix: remove method that was directly calling a util

2818860

doc: rename word in docstring

61896f0

wip: add guards in utility functions

b58cdbb

michaelbenayoun added 3 commits September 29, 2026 15:04

fix: tiny fixes

c975115

feat: add loading path in from_pretrained

3d5932d

tests: tests for checkpointing

151017d

michaelbenayoun commented Sep 29, 2026

View reviewed changes

michaelbenayoun added 2 commits September 29, 2026 16:24

chore: restore fsdp mixin

676e1b2

fix: broken imports in tests

7b88913

3outeille force-pushed the fsdp_tp_saving branch from 0e1cc55 to 7b88913 Compare September 29, 2026 16:10

3outeille reviewed Sep 29, 2026

View reviewed changes

Comment thread src/transformers/distributed/checkpoint.py Outdated

3outeille reviewed Sep 29, 2026

View reviewed changes

Comment thread src/transformers/distributed/checkpoint.py Outdated

3outeille reviewed Sep 29, 2026

View reviewed changes

michaelbenayoun added 6 commits October 1, 2026 11:45

fix: remove the support for full optimizer loading

15913b5

fix: remove consolidation support for optimizer saving

8682742

refactor: rename to load_model_checkpoint_distributed

14cfa9a

refactor: group two if statements into a for loop

81f34cd

refactor: separate the logic of finding the files and loading

332b3b3

refactor: restore optimizer changes for stacking

e4b7ba0

This was referenced Oct 1, 2026

FSDP x TP Saving / Loading, Optimizer michaelbenayoun/transformers#1

Closed

FSDP x TP Saving / Loading, Optimizer #49256

Open

style: ruff

e8bdcde

Merge branch 'main' into fsdp_tp_saving

0a46816

3outeille reviewed Oct 6, 2026

View reviewed changes

		import torch


		def _check_distributed_checkpointing_available(raise_if_not: bool = True) -> bool:

		set_model_state_dict(model, state)


		def load_model_checkpoint_distributed(model, checkpoint_dir: str \| os.PathLike, strict: bool = True) -> None:

Conversation

michaelbenayoun commented Sep 14, 2026 • edited by github-actions Bot Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

What does this PR do?

1. _StridedShard placement issue

2. Distributed checkpointing / Consolidation

3. Optimizer checkpointing

Script

Uh oh!

Uh oh!

HuggingFaceDocBuilderDev commented Sep 14, 2026

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

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!

michaelbenayoun left a comment

Choose a reason for hiding this comment

Uh oh!

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!

3outeille Oct 6, 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!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

ArthurZucker commented Oct 5, 2026

Uh oh!

Choose a reason for hiding this comment

Uh oh!

3outeille Oct 6, 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!

github-actions Bot commented Oct 6, 2026

CI recap

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!

3outeille Oct 6, 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!

3outeille Oct 6, 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!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

3outeille commented Oct 6, 2026

michaelbenayoun commented Sep 14, 2026 •

edited by github-actions Bot

Loading

1. `_StridedShard` placement issue

3outeille Oct 6, 2026 •

edited

Loading

3outeille Oct 6, 2026 •

edited

Loading

3outeille Oct 6, 2026 •

edited

Loading

3outeille Oct 6, 2026 •

edited

Loading