Add Wall-X serving and Turtle2 TCP WebSocket bridge

wall_x/model/core/ops/moe.py

@@ -5,6 +5,7 @@ Mixture-of-Experts processing.
 """
 
 import logging
+import os
 
 import torch
 
@@ -13,6 +14,14 @@ from wall_x.model.core.ops.base import OpsProxy
 logger = logging.getLogger(__name__)
 
 
+def _force_pytorch_backend() -> bool:
+    return os.environ.get("WALL_X_FORCE_TORCH_MOE_OPS", "").lower() in {
+        "1",
+        "true",
+        "yes",
+    }
+
+
 class PermuteOp(OpsProxy):
     """Reorder tokens by expert assignment for MoE processing.
 
@@ -21,9 +30,11 @@ class PermuteOp(OpsProxy):
 
     @property
     def _external_accel_name(self):
-        return "permute"
+        return None if _force_pytorch_backend() else "permute"
 
     def _get_cuda_kernel(self):
+        if _force_pytorch_backend():
+            return None
         try:
             from wall_x.model.core.ops._cuda_wrappers import permute_kernel
 
@@ -64,9 +75,11 @@ class UnpermuteOp(OpsProxy):
 
     @property
     def _external_accel_name(self):
-        return "unpermute"
+        return None if _force_pytorch_backend() else "unpermute"
 
     def _get_cuda_kernel(self):
+        if _force_pytorch_backend():
+            return None
         try:
             from wall_x.model.core.ops._cuda_wrappers import unpermute_kernel
 

wall_x/model/core/vla_mixin.py

@@ -321,6 +321,20 @@ class ActionModelMixMin:
             noisy_action_emb = noisy_action_emb.to(
                 inputs_embeds.device, inputs_embeds.dtype
             )
+            action_token_count = int(mask.sum().item())
+            if noisy_action_emb.ndim != 3 or (
+                action_token_count * inputs_embeds.shape[-1]
+                != noisy_action_emb.numel()
+            ):
+                raise ValueError(
+                    "Flow action placeholder/embedding mismatch: "
+                    f"input_ids={tuple(input_ids.shape)}, "
+                    f"action_tokens={action_token_count}, "
+                    f"action_chunk={tuple(action_chunk.shape)}, "
+                    f"inputs_embeds={tuple(inputs_embeds.shape)}, "
+                    f"noisy_action_emb={tuple(noisy_action_emb.shape)}, "
+                    f"configured_horizon={self.config.action_horizon_flow}"
+                )
             inputs_embeds = inputs_embeds.masked_scatter(action_mask, noisy_action_emb)
 
         return inputs_embeds, flow, adarms_cond
@@ -530,7 +544,10 @@ class ActionGenerationMixin(GenerationMixin):
             target_modules=target_modules,
             lora_dropout=lora_dropout,
             bias="none",
-            task_type="CAUSAL_LM",
+            # ``self.model`` is Wall-X's decoder, not the outer generation model.
+            # Leave task_type unset so PEFT uses its generic wrapper and forwards
+            # Wall-X's custom MoE arguments unchanged.
+            task_type=None,
         )
         self.model = get_peft_model(self.model, config)
         # Log trainable parameter information

wall_x/model/qact/qwen2_5/modeling_qwen2_5_vl.py

@@ -65,7 +65,7 @@ from transformers.utils import (
 from .configuration_qwen2_5_vl import Qwen2_5_VLConfig, Qwen2_5_VLVisionConfig
 from wall_x.model.core.attention.selector import AttentionsSelectorMixin
 from wall_x.model.core.ops import rot_pos_emb, get_window_index, m_rope
-
+##这一部分代码是对flash_attn的import的处理，在确认可使用的情况下再将flash_attn相关库做引入
 if is_flash_attn_2_available():
     from flash_attn import flash_attn_func
     from flash_attn import flash_attn_varlen_func
@@ -91,8 +91,8 @@ except ImportError:
 logger = logging.get_logger(__name__)
 
 _CONFIG_FOR_DOC = "Qwen2_5_VLConfig"
-
-
+#GEMM对齐实际上是将模型参数按照128一组对齐，实际上是tile的处理匹配，那么训练时没有对应的处理，所以下载之后需要做pad
+#针对VL模型中的MLP部分做实现与优化，做了GEMM对齐，gate_up合并Linear，做了checkpoint的pad / strip，以及activate的选择性重算
 class Qwen2_5_VLMLP(nn.Module):
     # Align intermediate_size to this value for cuBLAS GEMM tile efficiency.
     # E.g. 3420 -> 3456 (= 128 * 27), improving MFU from 34.8% to ~60%+.
@@ -1000,9 +1000,13 @@ class Qwen2_5_VLRotaryEmbedding(nn.Module):
             else "cpu"
         )
         with torch.autocast(device_type=device_type, enabled=False):
-            freqs = (
-                inv_freq_expanded.float() @ position_ids_expanded.float()
-            ).transpose(2, 3)
+            # This is an outer product: (..., rotary_dim, 1) x
+            # (..., 1, sequence_length). Broadcasting is mathematically
+            # identical to batched matmul and avoids a cuBLAS failure seen
+            # intermittently under multi-GPU FSDP.
+            freqs = (inv_freq_expanded.float() * position_ids_expanded.float()).transpose(
+                2, 3
+            )
             emb = torch.cat((freqs, freqs), dim=-1)
             cos = emb.cos()
             sin = emb.sin()

wall_x/model/qact/qwen2_5/modeling_qwen2_5_vl_act.py

@@ -1,3 +1,4 @@
+import os
 import re
 import torch
 import torch.nn as nn
@@ -809,6 +810,11 @@ class Qwen2_5_VLMoEForAction(
         use_selective_recompute=False,
     ):
         super().__init__(config)
+        if not getattr(config, "use_cuda_moe_ops", True):
+            os.environ["WALL_X_FORCE_TORCH_MOE_OPS"] = "1"
+            logger.warning_once(
+                "Using PyTorch MoE routing ops because use_cuda_moe_ops=false"
+            )
         self.visual = self._build_visual(config, use_selective_recompute)
         self.model = Qwen2_5_VLMoEModel(
             config, use_selective_recompute=use_selective_recompute
@@ -850,6 +856,8 @@ class Qwen2_5_VLMoEForAction(
                 target_modules=config.lora_target_modules,
                 lora_dropout=config.lora_dropout,
             )
+            if getattr(config, "lora_train_action_expert", False):
+                self._set_lora_action_expert_trainable()
         # Initialize weights and apply final processing
         self.post_init()
         self._post_init_engine(config)
@@ -861,6 +869,34 @@ class Qwen2_5_VLMoEForAction(
             use_selective_recompute=use_selective_recompute,
         )
 
+    def _set_lora_action_expert_trainable(self) -> None:
+        """Train LoRA adapters and the action expert while freezing the VLM base."""
+        action_expert_keywords = (
+            "action_preprocessor",
+            "moe.experts.1",
+            "qkv_proj_experts.1",
+            "o_proj_experts.1",
+            "input_layernorms.1",
+            "post_attention_layernorms.1",
+            "model.norms.1",
+        )
+        trainable_count = 0
+        trainable_parameters = 0
+        for name, param in self.named_parameters():
+            is_lora = "lora_" in name
+            is_action_expert = any(keyword in name for keyword in action_expert_keywords)
+            # Dataset normalizer statistics are model Parameters but must remain fixed.
+            is_normalizer_stat = "normalizer" in name and not is_lora
+            param.requires_grad = (is_lora or is_action_expert) and not is_normalizer_stat
+            if param.requires_grad:
+                trainable_count += 1
+                trainable_parameters += param.numel()
+        logger.info(
+            "LoRA action-expert mode: %d trainable tensors (%d parameters)",
+            trainable_count,
+            trainable_parameters,
+        )
+
     def _post_init_engine(self, config):
         """Hook called at the end of __init__. Subclasses may override to add engine-specific state."""
         pass