FSDP orchestration: apply + loading/saving by 3outeille · Pull Request #46990 · huggingface/transformers

3outeille · 2026-07-01T05:09:05Z

Summary

FSDP:
- Only 1 model for now has base_fsdp_plan. Will do another PR to edit every other models later
- now wired through from_pretrained
- For FSDP: loading through shard-on-Read + saving like TP (DCP optional)
- Add FSDP Ci
- DistributedMixin
TP:
- Wired DistributedConfig everywhere (no more tp_plan=auto)
- TP left untouched (no Dtensor yet)

Stack

Based on Shard on read Dtensor aware #46717 (A-PR-3 dual-path loading)

Wire distributed_config from_pretrained/save_pretrained alongside the legacy tp_plan path, add distributed/utils.py for mesh orchestration and checkpoint I/O, and extend sharding_utils with DTensor gather/optimizer fusion helpers needed by save/load.

…orchestration

HuggingFaceDocBuilderDev · 2026-07-01T05:52:17Z

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.

Wire FSDP tests into the dynamic PR CI caller, mirroring tests_tensor_parallel_ci: detect tests_fsdp_ci_test_list.txt, run with is_fsdp_test marker and RUN_FSDP_TESTS, and exclude FSDP tests from the tests_torch job. Companion to huggingface/transformers#46990 (tests_fetcher changes stay in transformers). Co-authored-by: Cursor <cursoragent@cursor.com>

3outeille · 2026-07-03T03:52:49Z

run-slow: cohere2_moe, deepseek_v4, glm_moe_dsa, gpt_oss

github-actions · 2026-07-03T03:54:11Z

Workflow Run ⚙️

This comment contains run-slow, running the specified jobs:

models: ["models/cohere2_moe", "models/deepseek_v4", "models/glm_moe_dsa", "models/gpt_oss"]
quantizations: []

…orchestration

github-actions · 2026-07-03T05:02:23Z

CI Results

Workflow Run ⚙️

Commit Info

Context	Commit	Description
RUN	1bcc0fb0	workflow commit (merge commit)
PR	37df13ee	branch commit (from PR)
main	70544cd9	base commit (on `main`)

⚠️ Model CI failed to report results

The test failure analysis could not be completed. Please check the workflow run for details.

…orchestration

Wire FSDP tests into the dynamic PR CI caller, mirroring tests_tensor_parallel_ci: detect tests_fsdp_ci_test_list.txt, run with is_fsdp_test marker and RUN_FSDP_TESTS, and exclude FSDP tests from the tests_torch job. Companion to huggingface/transformers#46990 (tests_fetcher changes stay in transformers). Co-authored-by: Cursor <cursoragent@cursor.com>

github-actions · 2026-07-06T06:42:09Z

[For maintainers] Suggested jobs to run (before merge)

run-slow: cohere2_moe, deepseek_v4, glm_moe_dsa, gpt_oss

github-actions · 2026-07-06T07:04:25Z

CI recap

Dashboard: View test results in Grafana
Latest run: 28772915580:1
Result: failure | Jobs: 2 | Tests: 21 | Failures: 1 | Duration: 2m 16s

Code quality check failed: test jobs were skipped. Fix the code quality issues and push again to run tests.

ArthurZucker

ty, isolate changes as much as possible

ArthurZucker · 2026-07-06T06:53:53Z

-        # Get the model state_dict
+        # --- FSDP path: DCP save (large models, torch >= 2.7) ---
+        if distributed_checkpoint:
+            self.save_distributed_checkpoint(model_to_save, save_directory)


Suggested change

self.save_distributed_checkpoint(model_to_save, save_directory)

self.save_distributed_checkpoint(model_to_save, save_directory, push_to_hub, save_on_this_rank)

ArthurZucker · 2026-07-06T06:54:01Z

+            if push_to_hub and save_on_this_rank:
+                model_card = create_and_tag_model_card(repo_id, self.model_tags, token=token)
+                model_card.save(os.path.join(save_directory, "README.md"))
+                self._upload_modified_files(
+                    save_directory,
+                    repo_id,
+                    files_timestamps,
+                    commit_message=commit_message,
+                    token=token,
+                    create_pr=create_pr,
+                )
+            return


Suggested change

if push_to_hub and save_on_this_rank:

model_card = create_and_tag_model_card(repo_id, self.model_tags, token=token)

model_card.save(os.path.join(save_directory, "README.md"))

self._upload_modified_files(

save_directory,

repo_id,

files_timestamps,

commit_message=commit_message,

token=token,

create_pr=create_pr,

)

return

ArthurZucker · 2026-07-06T06:54:17Z


-        # Get the model state_dict
+        # --- FSDP path: DCP save (large models, torch >= 2.7) ---
+        if distributed_checkpoint:


anything happening in this if -> in a separated function

ArthurZucker · 2026-07-06T06:54:42Z

+                filename = os.path.join(save_directory, shard_file)
+                shard_state_dict = {}
+                for tensor_name in tensor_names:
+                    # Get the tensor, and remove it from state_dict to avoid keeping the ref
+                    tensor = state_dict.pop(tensor_name)
+
+                    # If the param was offloaded, we need to load it back from disk to resave it. It's a strange pattern,
+                    # but it would otherwise not be contained in the saved shard if we were to simply move the file
+                    # or something
+                    if is_offloaded and tensor.device.type == "meta":
+                        tensor = load_offloaded_parameter(model_to_save, tensor_name)
+
+                    # only do contiguous after it's permuted correctly in case of TP
+                    shard_state_dict[tensor_name] = tensor.contiguous()
+
+                # TODO: it would be very nice to do the writing concurrently, but safetensors never releases the GIL,
+                # so it's not possible for now....
+                # Write the shard to disk
+                safe_save_file(shard_state_dict, filename, metadata=metadata)
+                # Cleanup the data before next loop (important with offloading, so we don't blowup cpu RAM)
+                del shard_state_dict
+
+            if index is None:
+                path_to_weights = os.path.join(save_directory, weights_name)
+                logger.info(f"Model weights saved in {path_to_weights}")
+            else:
+                save_index_file = SAFE_WEIGHTS_INDEX_NAME
+                save_index_file = os.path.join(save_directory, _add_variant(save_index_file, variant))
+                # Save the index as well
+                with open(save_index_file, "w", encoding="utf-8") as f:
+                    content = json.dumps(index, indent=2, sort_keys=True) + "\n"
+                    f.write(content)
+                logger.info(
+                    f"The model is bigger than the maximum size per checkpoint ({max_shard_size}) and is going to be "
+                    f"split in {len(state_dict_split.filename_to_tensors)} checkpoint shards. You can find where each parameters has been saved in the "
+                    f"index located at {save_index_file}."


should be in distributed

3outeille and others added 4 commits July 1, 2026 05:08

Merge branch 'split/a-pr-3-dual-path-loading' into split/a-pr-4-fsdp-…

17c6d40

…orchestration

add fsdp plan to 2 models for now

00eb116

add tests fsdp mixin

be296dd

3outeille mentioned this pull request Jul 1, 2026

Add native FSDP2 module + migration #46707

Open

linting

05900d6

3outeille added 8 commits July 1, 2026 06:15

refactor test fsdp mixin

fc2423b

test fsdp mixin cleaning

5bbd820

remove fsdp policy in tests + trim down further

b6d0b67

test fsdp clean

ea36123

restore test_modeling_utils

bec4d23

linting

8d3d329

start trim down stuff

6316ee1

fix

6e9004e

3outeille marked this pull request as draft July 1, 2026 08:49

3outeille added 13 commits July 2, 2026 04:38

breaking: cleaning modeling_utils.py

68df491

load path with fsdp (dtensor) and tp (old tp) is linked

16b0b29

linting

e976a44

add saving

a2fb155

styling

5f52f19

fix tp ci

7f54301

add fsdp to ci

99f79ac

linting

b451490

pick one model only for this PR

54ff4d1

restore

3399539

trigger fsdp ci

11cf79c

doc cleaning + tp_size remove

5b7ac3e

fix tp ci for ep

06b0c39

3outeille mentioned this pull request Jul 3, 2026

Add FSDP CI job to PR workflow huggingface/transformers-test-ci#80

Merged

edit doc

7a94d77

3outeille changed the title ~~FSDP orchestration: mesh init, distribute-before-load, DCP save~~ FSDP orchestration: apply + loading/saving Jul 3, 2026

3outeille added 2 commits July 3, 2026 03:17

move distributed function to utils + guarding

f3c742b

linting

37df13e

3outeille marked this pull request as ready for review July 3, 2026 03:49

Merge branch 'split/a-pr-3-dual-path-loading' into split/a-pr-4-fsdp-…

86875d2

…orchestration

Merge branch 'split/a-pr-3-dual-path-loading' into split/a-pr-4-fsdp-…

450579b

…orchestration

3outeille requested a review from ArthurZucker July 3, 2026 05:10

add DistributedMixin

91260ba

3outeille commented Jul 6, 2026

View reviewed changes

Comment thread src/transformers/modeling_utils.py

solve merge conflict

84dd107

3outeille force-pushed the split/a-pr-4-fsdp-orchestration branch from d70939a to 84dd107 Compare July 6, 2026 06:58

ArthurZucker reviewed Jul 6, 2026

View reviewed changes

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Uh oh!

FSDP orchestration: apply + loading/saving#46990

FSDP orchestration: apply + loading/saving#46990
3outeille wants to merge 33 commits into
split/a-pr-3-dual-path-loadingfrom
split/a-pr-4-fsdp-orchestration

3outeille commented Jul 1, 2026 •

edited by github-actions Bot

Loading

Uh oh!

HuggingFaceDocBuilderDev commented Jul 1, 2026

Uh oh!

3outeille commented Jul 3, 2026

Uh oh!

github-actions Bot commented Jul 3, 2026

Uh oh!

github-actions Bot commented Jul 3, 2026

Uh oh!

github-actions Bot commented Jul 6, 2026

Uh oh!

Uh oh!

github-actions Bot commented Jul 6, 2026

Uh oh!

ArthurZucker left a comment

Uh oh!

ArthurZucker Jul 6, 2026

Uh oh!

ArthurZucker Jul 6, 2026

Uh oh!

ArthurZucker Jul 6, 2026

Uh oh!

ArthurZucker Jul 6, 2026

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

4 participants

	self.save_distributed_checkpoint(model_to_save, save_directory)
	self.save_distributed_checkpoint(model_to_save, save_directory, push_to_hub, save_on_this_rank)

Uh oh!

Conversation

3outeille commented Jul 1, 2026 • edited by github-actions Bot Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Summary

Stack

Uh oh!

HuggingFaceDocBuilderDev commented Jul 1, 2026

Uh oh!

3outeille commented Jul 3, 2026

Uh oh!

github-actions Bot commented Jul 3, 2026

Uh oh!

github-actions Bot commented Jul 3, 2026

CI Results

Commit Info

Uh oh!

github-actions Bot commented Jul 6, 2026

Uh oh!

Uh oh!

github-actions Bot commented Jul 6, 2026

CI recap

Uh oh!

ArthurZucker left a comment

Choose a reason for hiding this comment

Uh oh!

ArthurZucker Jul 6, 2026

Choose a reason for hiding this comment

Uh oh!

ArthurZucker Jul 6, 2026

Choose a reason for hiding this comment

Uh oh!

ArthurZucker Jul 6, 2026

Choose a reason for hiding this comment

Uh oh!

ArthurZucker Jul 6, 2026

Choose a reason for hiding this comment

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

4 participants

3outeille commented Jul 1, 2026 •

edited by github-actions Bot

Loading