Shard on read Dtensor aware - #46717
Merged
Merged
Conversation
Move FSDP2 wrapping and plan verification to distributed/fsdp.py, keep integrations/fsdp.py as a backward-compatible re-export, and update core call sites to import from transformers.distributed.fsdp.
Contributor
|
View the CircleCI Test Summary for this PR: https://huggingface.co/spaces/transformers-community/circle-ci-viz?pr=46717&sha=672bbb |
|
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. |
…uggingface/transformers into split/a-pr-3-dual-path-loading
3outeille
marked this pull request as ready for review
June 23, 2026 05:13
3outeille
requested review from
ArthurZucker
and removed request for
Rocketknight1 and
ydshieh
June 23, 2026 07:43
…uggingface/transformers into split/a-pr-3-dual-path-loading
Co-authored-by: Arthur <48595927+ArthurZucker@users.noreply.github.com>
…uggingface/transformers into split/a-pr-3-dual-path-loading
3outeille
changed the base branch from
main
to
fix-distributed-config-sp-ep-docstring
July 10, 2026 05:16
ArthurZucker
approved these changes
Jul 10, 2026
| sharding_op = mapping.distributed_operation | ||
| materialize_device = device_map[""] | ||
|
|
||
| future_or_tensor = spawn_materialize( |
Contributor
CI recapDashboard: View test results in Grafana |
stevhliu
pushed a commit
to stevhliu/transformers
that referenced
this pull request
Jul 30, 2026
* add distributed config * Add native FSDP2 module and migrate FSDP imports (Phase A PR-2). Move FSDP2 wrapping and plan verification to distributed/fsdp.py, keep integrations/fsdp.py as a backward-compatible re-export, and update core call sites to import from transformers.distributed.fsdp. * linting * unecessary * copyright edit * revert * add shard on read * jsut shard on read * cleaning * linting * fix * fix * remove redundant test file * Update src/transformers/distributed/fsdp.py naming Co-authored-by: Arthur <48595927+ArthurZucker@users.noreply.github.com> * avoid looping, just look at dict * expand_fsdp returns reshard_targets, no_reshard_targets right away * better _resolve_tied_embed_lm_head_plan * cleaning * ruff * more robust detection of embed and lm_head * cleaning * ruff * typo * cleaner * cleaner * typo * refactor dense path + apply_contiguous_shard * linting * cleaning * refactor _apply_strided_shard * better * refactor _slice_and_cat * better comment * refactor moe dtensor shard ops * better comment * comment * cleaning * linting * expand_fsdp_plan iterate over modules * comment about tie embedding * add comment tied embedding * some cleaning * cleaning + comment * rename function for clarity * Apply suggestion from @ArthurZucker Co-authored-by: Arthur <48595927+ArthurZucker@users.noreply.github.com> * doc * comment --------- Co-authored-by: Arthur <48595927+ArthurZucker@users.noreply.github.com>
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
3 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
What does this PR do ?
This PR adds
shard-on-readsupport for loading DTensor-sharded model parameters duringfrom_pretrained. Instead of materializing the full checkpoint tensor on every rank and sharding afterward, each rank reads only the slice it needs from the checkpoint, based on the target parameter’s DTensor mesh and placements.For now I don't want to break the integration with the existing TP on main (since we haven't merged the inference-compatible DTensor TP yet). So here's what I'm doing: