From d9b80599c9b0e741889d65d0245b2873c1128c79 Mon Sep 17 00:00:00 2001 From: padatta Date: Mon, 6 Jul 2026 08:03:04 +0000 Subject: [PATCH] Enable XPU inference for InternVLA-M1 --- InternVLA/model/framework/M1.py | 23 ++++++++++++------ InternVLA/model/modules/vlm/QWen2_5.py | 33 ++++++++++++++++++++++---- 2 files changed, 45 insertions(+), 11 deletions(-) diff --git a/InternVLA/model/framework/M1.py b/InternVLA/model/framework/M1.py index 34ab8ce..a92ad53 100644 --- a/InternVLA/model/framework/M1.py +++ b/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) diff --git a/InternVLA/model/modules/vlm/QWen2_5.py b/InternVLA/model/modules/vlm/QWen2_5.py index 8d34e37..69f7dd1 100644 --- a/InternVLA/model/modules/vlm/QWen2_5.py +++ b/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,