Fix LoRA on FSDP2-sharded models: shape inference and module hooks

[qgallouedec]

Two fixes for applying LoRA to a model whose parameters are already FSDP2-sharded (fully_shard) when get_peft_model runs, e.g. a model loaded with transformers' DistributedConfig(fsdp_size=..., tp_size=..., enable_expert_parallel=True), which shards inside from_pretrained (huggingface/transformers#48204).

1. _get_in_out_features reads the local shard shape on FSDP2 weights

peft/src/peft/tuners/tuners_utils.py

Lines 176 to 178 in 5d602fd

	if _torch_supports_distributed and isinstance(module.weight, torch.distributed.tensor.DTensor):
	# If Tensor Parallel is used, the weight is sharded, so we need to get the local shape
	out_features, in_features = module.weight.to_local().shape

This is correct for TP, where the module computes on its local shard. Under FSDP2 the storage is sharded (Shard(0)) but the module still computes the full projection, so LoRA B gets built with out_features / fsdp_size and the forward crashes:

RuntimeError: The size of tensor a (4096) must match the size of tensor b (2048) at non-singleton dimension 2

(Qwen3-30B-A3B q_proj: 2048 → 4096, FSDP2-sharded 2-way → local shape (2048, 2048), so lora_B was created 16 → 2048 against a 4096-dim base output.)

Minimal reproduction — stock transformers main, 2 GPUs (any model whose targeted projection is non-square):

# torchrun --nproc_per_node 2 repro.py
import torch
from peft import LoraConfig, get_peft_model

from transformers import AutoModelForCausalLM
from transformers.distributed import DistributedConfig

model = AutoModelForCausalLM.from_pretrained(
    "Qwen/Qwen3-0.6B",
    dtype=torch.bfloat16,
    distributed_config=DistributedConfig(fsdp_size=2),
)
model = get_peft_model(model, LoraConfig(r=8, target_modules=["q_proj"]))
ids = torch.randint(0, 1000, (1, 16), device="cuda")
model(input_ids=ids)

RuntimeError: The size of tensor a (2048) must match the size of tensor b (1024) at non-singleton dimension 2

(Qwen3-0.6B q_proj is 1024 → 2048; with this PR the same script runs.)

Fix: when the weight's mesh has only a fsdp dimension (or no named dimensions), use the module's in_features/out_features; keep the local-shape path for TP meshes.

2. BaseTuner.forward skips module hooks

Calling .forward() directly bypasses nn.Module.__call__, so no module hooks run on the wrapped model, including the root pre-forward hook fully_shard registers. FSDP2's lazy init then never sees its root and a later non-root state fails with:

RuntimeError: FSDP requires a single root module but got (FSDPQwen3MoeRMSNorm((2048,), eps=1e-06), FSDPLinear(in_features=2048, out_features=151936, bias=False))

Fix: simply call self.model(*args, **kwargs).

Validation

With both fixes (and the transformers 2-D draft), LoRA fine-tuning of Qwen3-30B-A3B under FSDP2 × expert parallelism on 4×H100 trains with healthy losses and finite gradient norms for 8/8 steps (losses ~12.85, grad norms ~1.0). Without fix 1 the first forward crashes with the shape mismatch above; with fix 1 but not fix 2, FSDP2 lazy init raises the single-root error.

[HuggingFaceDocBuilderDev]

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.

[BenjaminBossan]

Thanks for fixing this issue, I tested it and could reproduce the error and the fix. I have some questions and small comments before merging, please check.

[BenjaminBossan]

So here we match set(mesh_dim_names) == {"fsdp"} or set(mesh_dim_names) == set(). In what circumstances can we have the latter?

[qgallouedec]

Plain fully_shard(model) without an explicit mesh: torch builds the default mesh with mesh_dim_names=None, see https://github.com/pytorch/pytorch/blob/v2.13.0/torch/distributed/fsdp/_fully_shard/_fsdp_init.py#L198-L211.

import os
import torch
import torch.distributed as dist
from torch.distributed.fsdp import fully_shard

os.environ.setdefault("MASTER_ADDR", "localhost")
os.environ.setdefault("MASTER_PORT", "29501")
dist.init_process_group(rank=0, world_size=1)

model = torch.nn.Linear(8, 8)
fully_shard(model)
print(type(model.weight))  # <class 'torch.distributed.tensor.DTensor'>
print(model.weight.device_mesh.mesh_dim_names)  # None

Named dims only exist when the caller passes a mesh, e.g. transformers DistributedConfig names its dim "fsdp".
So the unnamed case is the vanilla torch API path, and it's FSDP-sharded storage like the named one, hence the same full-shape branch.
In theory a TP weight could also sit on an unnamed mesh and be misclassified here, but torch's TP APIs are documented and used with named meshes, and before this PR that case was broken the other way around for everyone.

[BenjaminBossan]

I see, thanks for explaining. Could you please add a comment to explain the named vs unnamed case?

[BenjaminBossan]

Thanks for providing further context. I just have one small comment left, otherwise the PR LGTM:

[BenjaminBossan]

I see, thanks for explaining. Could you please add a comment to explain the named vs unnamed case?

[qgallouedec]

Added in 8a764bc.

[BenjaminBossan]

Thanks for these FSDP fixes and for explaining what's going on, LGTM.