Fix: properly error when DDP + Dtensor model (#3629)

* Feat: add check

* Refactor: nits

Author: S1ro1

src/accelerate/accelerator.py

@@ -108,6 +108,7 @@
     is_xpu_available,
     load_fsdp_model,
     load_fsdp_optimizer,
+    model_has_dtensor,
     pad_across_processes,
     parse_choice_from_env,
     recursively_apply,
@@ -1631,6 +1632,10 @@ def prepare_model(self, model: torch.nn.Module, device_placement: bool = None, e
                 DistributedType.MULTI_XPU,
                 DistributedType.MULTI_HPU,
             ):
+                if model_has_dtensor(model):
+                    raise 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 'cpu' instead."
+                    )
                 if any(p.requires_grad for p in model.parameters()):
                     kwargs = self.ddp_handler.to_kwargs() if self.ddp_handler is not None else {}
                     # TODO: Look at enabling native TP training directly with a proper config

src/accelerate/utils/__init__.py

@@ -282,6 +282,7 @@
     is_port_in_use,
     load,
     merge_dicts,
+    model_has_dtensor,
     recursive_getattr,
     save,
     wait_for_everyone,

src/accelerate/utils/other.py

@@ -194,6 +194,26 @@ def compile_regions_deepspeed(module: torch.nn.Module, **compile_kwargs):
         module.compile(**compile_kwargs)
 
 
+def model_has_dtensor(model: torch.nn.Module) -> bool:
+    """
+    Check if the model has DTensor parameters.
+
+    Args:
+        model (`torch.nn.Module`):
+            The model to check.
+
+    Returns:
+        `bool`: Whether the model has DTensor parameters.
+    """
+    if is_torch_version(">=", "2.5.0"):
+        from torch.distributed.tensor import DTensor
+    else:
+        # from torch 2.0.0 (oldest supported accelerate torch version), DTensor is in torch.distributed._tensor
+        from torch.distributed._tensor import DTensor
+
+    return any(isinstance(p, DTensor) for p in model.parameters())
+
+
 def extract_model_from_parallel(
     model, keep_fp32_wrapper: bool = True, keep_torch_compile: bool = True, recursive: bool = False
 ):