Skip to content

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

Open
qgallouedec wants to merge 2 commits into
mainfrom
fix-prepare-model-dtensor-to
Open

prepare_model: don't move DTensor-sharded models to the device#4181
qgallouedec wants to merge 2 commits into
mainfrom
fix-prepare-model-dtensor-to

Conversation

@qgallouedec

Copy link
Copy Markdown
Member

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.

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.
@HuggingFaceDocBuilderDev

Copy link
Copy Markdown

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.

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants