Repository navigation
Should Trainer support models FSDP2-sharded at load time (DistributedConfig(fsdp_size=N))? #48210
Description
System info
- transformers: main (
d56c55bf56) - torch: 2.13.0+cu130
- accelerate: 1.15.0.dev0
- python: 3.13.13, Linux, 2x H100
Who can help?
Reproduction
Load a model sharded with FSDP2 at load time and hand it to Trainer:
# torchrun --nproc_per_node 2 repro.py
import torch
from torch.utils.data import Dataset
from transformers import AutoModelForCausalLM, Trainer, TrainingArguments
from transformers.distributed import DistributedConfig
class RandomTokens(Dataset):
def __len__(self):
return 64
def __getitem__(self, i):
g = torch.Generator().manual_seed(i)
ids = torch.randint(0, 50000, (128,), generator=g)
return {"input_ids": ids, "labels": ids.clone()}
model = AutoModelForCausalLM.from_pretrained(
"Qwen/Qwen3-0.6B",
dtype=torch.bfloat16,
distributed_config=DistributedConfig(fsdp_size=2),
)
trainer = Trainer(
model=model,
args=TrainingArguments(
output_dir="/tmp/fsdp-trainer",
per_device_train_batch_size=1,
max_steps=3,
logging_steps=1,
report_to=[],
save_strategy="no",
),
train_dataset=RandomTokens(),
)
trainer.train()Result:
Traceback (most recent call last):
File "repro.py", line 37, in <module>
trainer.train()
File ".../transformers/trainer.py", line 1452, in train
return inner_training_loop(...)
File ".../transformers/trainer.py", line 1491, in _inner_training_loop
self._prepare_for_training(...)
File ".../transformers/trainer.py", line 1627, in _prepare_for_training
... self.accelerator.prepare(...)
File ".../accelerate/accelerator.py", line 1557, in prepare
File ".../accelerate/accelerator.py", line 1879, in prepare_model
raise ValueError(
ValueError: Your model contains `DTensor` parameters, which is incompatible with DDP. Maybe you loaded your model with `device_map='auto'`? Specify `device_map='cuda'` or 'xpu' or 'cpu' instead.
Expected behavior
Either the Trainer should tell accelerate about the existing sharding instead of letting it wrap the DTensor model in DDP, or this combination is not meant to be supported and a clear error saying so would help. The Trainer does have the analogous wiring for tensor parallel models in create_accelerator_and_postprocess, which suggests load-time sharding plus Trainer is intended, but there is nothing equivalent for fsdp_size and the model does not even expose the sharding degree (there is a tp_size property but no fsdp_size).
So the question is what the intended design is: should accelerate adopt an already-sharded model, or should FSDP always go through the accelerate plugin? Context: found while working on FSDP2 x expert parallelism in #48204, where the 2-D case is wired up. Happy to help once the direction is decided.
Activity
greekera1000 commentedon Aug 22, 2026
hi @qgallouedec , Following the thread from #48204 — since that PR is transformers-side only, I went looking at what happens on the accelerate side once
_fsdp_sizeis stamped on the model. Sharing in case it's useful. Line refs are againstmainas of today, so they'll drift.Where it stops being a transformers problem
#48204 already sets
model._tp_size/model._fsdp_sizeinmaybe_distribute_modelfor the Trainer to mirror intoParallelismConfig. Assuming that lands andTrainer._build_accelerator_argsgains an FSDP block alongside the existing TP one (trainer.py, ~L748), the model still hits this inaccelerate/accelerator.pyprepare_model(~L1877):The guard short-circuits only on
tp_enabled.ParallelismConfig.dp_shard_enabledexists but is consulted exactly once in that file (L827, a mesh accessor) — never here. So a model sharded byfully_shardat load time still falls into the DDP branch and trips the DTensor check, regardless of whatParallelismConfigis carrying.The adoption logic already exists
fsdp2_prepare_model(utils/fsdp_utils.py, ~L733) opens with:Since
apply_fully_sharded_data_parallelismuses composablefully_shard, the model is anFSDPModule— so "don't re-shard, just adopt it" is already implemented and would do the right thing. It's simply unreachable: the guard above diverts to DDP first, andis_fsdp2isdistributed_type == FSDP and fsdp_plugin.fsdp_version == 2(state.py, L1101), so aParallelismConfigcarryingdp_shard_sizewon't route there without an FSDP plugin also being present.On the design question
Answering it directly as best I can from outside: accelerate looks close to already supporting this — widening the
prepare_modelguard pasttp_enabled, and deciding whetherdp_shard_enabledor the FSDP plugin is the thing that signals "already sharded."Which of those two owns that signal seems like the real fork. The manual gradient averaging and mesh-aware clipping in #48204's
trainer.pychanges read a bit like symptoms of the same ambiguity — work accelerate would normally own, done in Trainer because the model can't be handed over cleanly.One knock-on worth flagging if the answer ends up being "
ParallelismConfig, no plugin":Trainersetsis_fsdp_enabled = accelerator.state.fsdp_plugin is not None(trainer.py, ~L840), which gates the FSDP-aware checkpoint paths. A plugin-free route would leave thatFalseand silently bypass them — a quieter failure than the current one.Caveat: this is a source read, not a repro — no 2×H100 here to verify any of it.
github-actions commentedon Sep 21, 2026
This issue has been automatically marked as stale because it has not had recent activity. If you think this still needs to be addressed please comment on this thread.
Please note that issues that do not follow the contributing guidelines are likely to be ignored.