prepare_model: don't move DTensor-sharded models to the device

[qgallouedec]

A model sharded at load time (e.g. transformers' DistributedConfig with FSDP2, optionally CPU-offloaded) manages its own parameter placement, but prepare_model unconditionally calls model = model.to(self.device) when device_placement is set. .to() recurses _apply through FSDP2-managed parameters and raises:

RuntimeError: Attempted to set the storage of a tensor on device "cpu" to a storage on different device "cuda:0".
RuntimeError: _apply(): Couldn't swap FSDPLinear.weight

model_has_dtensor already exists and is consulted a few lines below for the DDP case; use it to skip the device move too. Hit while fine-tuning GLM-5.2 (753B) with fsdp_cpu_offload=True through Trainer (huggingface/transformers#48204); on-device sharded models are unaffected because the move was a same-device no-op for them.

Reproduction

/!\ A standalone toy does not trigger this: plain FSDP2 (fresh or lazy-initialized, with or without CPU offload) tolerates the .to().

The failing _apply swap needs the parameter arrangement that transformers' sharded loading produces (DTensor expert weights on a tp mesh combined with FSDP2-managed dense parameters, CPU-offloaded). With the transformers side installed (huggingface/transformers#48204):

# torchrun --nproc_per_node 4 repro.py
import torch
from transformers import AutoModelForCausalLM, Trainer, TrainingArguments
from transformers.distributed import DistributedConfig

model = AutoModelForCausalLM.from_pretrained(
    "Qwen/Qwen3-30B-A3B", dtype=torch.bfloat16,
    distributed_config=DistributedConfig(tp_size=2, fsdp_size=2, enable_expert_parallel=True, fsdp_cpu_offload=True),
)
Trainer(model=model, args=TrainingArguments(output_dir="/tmp/x")).accelerator.prepare(model)
# RuntimeError: Attempted to set the storage of a tensor on device "cpu" to a storage on
# different device "cuda:0"  /  _apply(): Couldn't swap FSDPLinear.weight

The guard is correct independently of the trigger's rarity: a model holding DTensor parameters manages its own placement, and model_has_dtensor is already the codebase's test for exactly that condition a few lines below.

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

Indeed, thanks !

[SunMarc]

maybe put the check in device_placement ? we already check for fsdp there

[qgallouedec]

Done in f74c62f, moved into the device_placement resolution next to the FSDP check.

Note an explicit device_placement=True still attempts the move (same as FSDP today): the resolution is a default, not a guard, which I think is the right semantics.

[SunMarc]

here

[SunMarc]

@bot /style

[github-actions]

Style bot fixed some files and pushed the changes.