[Umbrella] FSDP2 × expert-parallelism fixes (do not merge) - #4178
Draft
qgallouedec wants to merge 7 commits into
Draft
[Umbrella] FSDP2 × expert-parallelism fixes (do not merge)#4178qgallouedec wants to merge 7 commits into
qgallouedec wants to merge 7 commits into
Conversation
…nsion `ParallelismConfig.fsdp_dim_names` always asks for `dp_shard_cp`, but that joint dimension is only flattened into the device mesh when `dp_shard` or `cp` is enabled. Launching an already-parallelized model -- e.g. one loaded with `DistributedConfig(tp_size=world_size)` or `enable_expert_parallel=True` -- under an FSDP config therefore slices a tp-only mesh by a name it does not contain, and the user sees a bare `KeyError` from `torch.distributed.device_mesh`, several frames below any code they wrote. Say what is actually wrong instead.
is_peft_available() checks importlib metadata, so a peft provided only through PYTHONPATH (source checkout, shipped environment) makes is_peft_model() return False even while a PeftModel instance is right there in the model tree. Callers then treat the wrapper as a plain model; in TRL this patched a fused-loss forward onto the PeftModel, where transformers' base_model property resolves to peft's LoraModel and training crashes with an unrelated-looking AttributeError. If peft is already in sys.modules, isinstance against it directly; fall back to the metadata check only to trigger the first import.
A model sharded at load time (e.g. transformers DistributedConfig with FSDP2, optionally CPU-offloaded) manages its own parameter placement, and .to() on FSDP2-managed parameters raises RuntimeError: _apply(): Couldn't swap FSDPLinear.weight. model_has_dtensor already exists and is checked a few lines below for the DDP case; use it to skip the device move too.
…ensor-to' into fsdp2-ep-integration
qgallouedec
marked this pull request as draft
August 22, 2026 19:57
|
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. |
… metadata" This reverts commit a9e2e6b.
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
2 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.
Install-only umbrella for FSDP2 × expert parallelism (huggingface/transformers#48204): merges
model.to(device)inprepare_modelfor DTensor-sharded models)so colleagues can install one branch. Not meant to merge.
Used by the runnable 100B–753B MoE training examples in huggingface/trl#6869.