Raise a clear error when FSDP is enabled on a mesh with no shard dimension

[qgallouedec]

When FSDP2 is enabled together with a ParallelismConfig whose ranks are fully consumed by other dimensions (for example tp_size == world_size, which is also what transformers' expert parallelism produces), fsdp2_prepare_model slices the mesh with fsdp_dim_names and fails with a confusing error:

KeyError: "Invalid mesh_dim_names ('dp_shard_cp',) specified. Valid mesh_dim_names are ['tp']."

This PR raises an explicit error instead:

ValueError: FSDP is enabled but the device mesh is ('tp',), which has no dimension for FSDP to shard across (both `dp_shard_size` and `cp_size` are 1). This usually means a model that is already parallelized another way -- e.g. loaded with `DistributedConfig(tp_size=N)` or `enable_expert_parallel=True`, which makes the whole world size tensor/expert parallel -- was launched under an FSDP config. Either launch it without the FSDP config, or leave ranks for FSDP to use by lowering `tp_size`.

Reproduction (2 GPUs):

# torchrun --nproc_per_node 2 repro.py
import torch
import torch.nn as nn
from accelerate import Accelerator
from accelerate.parallelism_config import ParallelismConfig
from accelerate.utils import FullyShardedDataParallelPlugin

accelerator = Accelerator(
    parallelism_config=ParallelismConfig(tp_size=2),
    fsdp_plugin=FullyShardedDataParallelPlugin(fsdp_version=2),
)
model = nn.Linear(8, 8)
optimizer = torch.optim.AdamW(model.parameters())
model, optimizer = accelerator.prepare(model, optimizer)

[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.

[SunMarc]

Thanks ! Left a comment

[SunMarc]

Can't we put that in _validate_accelerator ? i feel like should detect this way earlier

[qgallouedec]

Good call, moved to _validate_accelerator.
The condition only needs the config (is_fsdp2 + neither dp_shard nor cp enabled), so it now raises at Accelerator() init instead of inside prepare_model.

[SunMarc]

Thanks