[Trainer] Force is_model_parallel when model is loaded in multiple GPUs using accelerate

src/transformers/models/t5/modeling_t5.py

@@ -1746,7 +1746,7 @@ def forward(
         loss = None
         if labels is not None:
             loss_fct = CrossEntropyLoss(ignore_index=-100)
-            loss = loss_fct(lm_logits.view(-1, lm_logits.size(-1)), labels.view(-1))
+            loss = loss_fct(lm_logits.view(-1, lm_logits.size(-1)).to(labels.device), labels.view(-1))
             # TODO(thom): Add z_loss https://github.com/tensorflow/mesh/blob/fa19d69eafc9a482aff0b59ddd96b025c0cb207d/mesh_tensorflow/layers.py#L666
 
         if not return_dict:

src/transformers/trainer.py

@@ -370,6 +370,19 @@ def __init__(
         else:
             self.is_model_parallel = False
 
+        if (
+            getattr(model, "hf_device_map", None) is not None
+            and len(set(model.hf_device_map.values())) > 1
+            and not self.is_model_parallel
+        ):
+            self.is_model_parallel = True
+
+            # warn users
+            logger.warning(
+                "You have loaded a model on multiple GPUs. `is_model_parallel` attribute will be force-set",
+                " to `True` to avoid any unexpected behavior such as device placement mismatching.",
+            )
+
         # At this stage the model is already loaded
         if getattr(model, "is_loaded_in_8bit", False):
             if getattr(model, "_is_int8_training_enabled", False):