Enable XPU inference for InternVLA-M1 (backend-safe)

InternVLA/model/framework/M1.py

@@ -26,7 +26,7 @@
 IGNORE_INDEX = -100
 
 from InternVLA.model.framework.base_framework import baseframework
-from InternVLA.model.modules.vlm.QWen2_5 import get_qwen2_5_interface
+from InternVLA.model.modules.vlm.QWen2_5 import get_qwen2_5_interface, _get_autocast_device
 from InternVLA.model.modules.projector.QFormer import get_layerwise_qformer
 from InternVLA.model.modules.action_model.DiTActionHeader import get_action_model
 from InternVLA.model.modules.dino_model.dino import get_dino_model
@@ -106,7 +106,7 @@ def forward(
 
         # Step 1: QWenVL input format
         qwen_inputs = self.qwen_vl_interface.build_qwenvl_inputs(images=batch_images, instructions=instructions)
-        with torch.autocast("cuda", dtype=torch.bfloat16):
+        with torch.autocast(_get_autocast_device(self.qwen_vl_interface.model), dtype=torch.bfloat16):
             qwenvl_outputs = self.qwen_vl_interface(
                 **qwen_inputs,
                 output_attentions=False,
@@ -138,7 +138,7 @@ def forward(
         action_condition = self.layer_qformer(cat_conditions)  # [B, 64, D_action]
 
         # Step 4: Action Expert Forward and Loss
-        with torch.autocast("cuda", dtype=torch.float32):
+        with torch.autocast(_get_autocast_device(self.qwen_vl_interface.model), dtype=torch.float32):
 
             # here is a tips to accelerate training speed, by repeating each sample for several times @ref to CogACT
             actions = torch.tensor(np.array(actions), device=action_condition.device)  # [B, chunk, 7]
@@ -204,7 +204,7 @@ def predict_action(
         inferface_inputs = self.qwen_vl_interface.build_qwenvl_inputs(images=batch_images, instructions=instructions)
         qwen_inputs = inferface_inputs
 
-        with torch.autocast("cuda", dtype=torch.bfloat16):
+        with torch.autocast(_get_autocast_device(self.qwen_vl_interface.model), dtype=torch.bfloat16):
             qwenvl_outputs = self.qwen_vl_interface(
                 **qwen_inputs,
                 output_hidden_states=True,
@@ -219,7 +219,7 @@ def predict_action(
             dino_encoded_features = dino_features.reshape(B, -1, dino_features.shape[-1])  # [B, num_view * token, dim]
             dino_encoded_features = self.dino_pro(dino_encoded_features)  # [B, 256, D]
 
-        with torch.autocast("cuda", dtype=torch.bfloat16):
+        with torch.autocast(_get_autocast_device(self.qwen_vl_interface.model), dtype=torch.bfloat16):
 
             start_layer = self.config.framework.layer_qformer.qformer_start_layer
             end_layer = self.config.framework.layer_qformer.qformer_end_layer
@@ -296,10 +296,15 @@ def chat_with_M1(
         image: Image.Image,
         text: str,
         max_new_tokens: int = 128,
-        device: Optional[str] = "cuda",
+        device: Optional[str] = None,
     ) -> List[str]:
         processor = getattr(self.qwen_vl_interface, "processor", None)
         model = getattr(self.qwen_vl_interface, "model", None)
+        if device is None:
+            try:
+                device = next(model.parameters()).device
+            except (StopIteration, AttributeError):
+                device = "cpu"
         # if processor is None or model is None:
         #     raise RuntimeError("qwen_vl_interface 缺少 processor 或 model。")
 
@@ -394,6 +399,10 @@ def build_model_framework(config: dict = {}) -> InternVLA_M1:
         break
 
     # try get model
-    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+    device = torch.device(
+        "xpu" if hasattr(torch, 'xpu') and torch.xpu.is_available()
+        else "cuda" if torch.cuda.is_available()
+        else "cpu"
+    )
     model = model.to(device)
     model(batch)

InternVLA/model/modules/vlm/QWen2_5.py

@@ -24,6 +24,19 @@
 
 import torch.nn as nn
 
+def _get_autocast_device(model_or_module):
+    """Detect autocast device type from model parameters."""
+    try:
+        p = next(model_or_module.parameters())
+        if p.device.type == "xpu":
+            return "xpu"
+        elif p.device.type == "cuda":
+            return "cuda"
+    except StopIteration:
+        pass
+    return "cpu"
+
+
 
 class _QWen_VL_Interface(nn.Module):
     """
@@ -77,11 +90,23 @@ def __init__(self, config: Optional[dict] = None, **kwargs):
         qwenvl_config = config.framework.get("qwenvl", {})
         model_id = qwenvl_config.get("base_vlm", "Qwen/Qwen2.5-VL-3B-Instruct")
 
+        # Backend-safe: flash_attention_2 + CUDA placement on CUDA; sdpa + XPU/CPU otherwise.
+        try:
+            from transformers.utils import is_flash_attn_2_available
+            _attn_impl = "flash_attention_2" if is_flash_attn_2_available() else "sdpa"
+        except Exception:
+            _attn_impl = "sdpa"
+        if torch.cuda.is_available():
+            _device_map = "cuda"
+        elif hasattr(torch, "xpu") and torch.xpu.is_available():
+            _device_map = "xpu"
+        else:
+            _device_map = "cpu"
         model = Qwen2_5_VLForConditionalGeneration.from_pretrained(
             model_id,
-            attn_implementation="flash_attention_2",
+            attn_implementation=_attn_impl,
             torch_dtype="auto",
-            device_map="cuda",
+            device_map=_device_map,
         )
         processor = AutoProcessor.from_pretrained(model_id)
         processor.tokenizer.padding_side = "left"
@@ -131,7 +156,7 @@ def forward(
             - Hidden states required for auxiliary alignment or feature extraction modules.
         """
 
-        with torch.autocast("cuda", dtype=torch.bfloat16):
+        with torch.autocast(_get_autocast_device(self.model), dtype=torch.bfloat16):
             outputs = self.model(
                 input_ids=input_ids,
                 attention_mask=attention_mask,
@@ -178,7 +203,7 @@ def generate(
             - Uses autocast(float16); relies on attribute enable_mixed_precision_training.
             - For iterative dialogue, caller manages past_key_values externally.
         """
-        with torch.autocast("cuda", enabled=self.enable_mixed_precision_training, dtype=torch.float16):
+        with torch.autocast(_get_autocast_device(self.model), enabled=self.enable_mixed_precision_training, dtype=torch.float16):
             generation_output = self.model.generate(
                 input_ids=input_ids,
                 attention_mask=attention_mask,