admin/lerobot
1
0
Fork 0
mirror of https://github.com/huggingface/lerobot.git synced 2026-07-23 01:41:54 +00:00
Code Issues Packages Projects Releases Wiki Activity

refactor(wall-x): subclass native Transformers Qwen2.5-VL instead of vendoring it (#4035)

Browse Source
This commit is contained in:
Steven Palma
2026-07-17 19:09:12 +02:00
committed by GitHub
parent 9d82bb9871
commit a9879e69ed
9 changed files with 854 additions and 3020 deletions
Download Patch File Download Diff File Expand all files Collapse all files
src/lerobot/policies/wall_x/configuration_wall_x.py
+18 -2
View File
@@ -58,10 +58,14 @@ class WallXConfig(PreTrainedConfig):
# Action prediction mode: "diffusion" or "fast"
prediction_mode: str = "diffusion"
# Attention Implementation, options: "eager", "flash_attention_2", "sdpa"
# NOTE: flash-attn==2.7.4.post1 is required for flash_attention_2 implementation
# Wall-X's bidirectional action-token islands currently require eager attention.
attn_implementation: str = "eager"
# Vision attention is independent from the text action-token mask. ``auto`` uses
# PyTorch's packed variable-length attention when the runtime supports it and
# otherwise falls back to the native per-chunk SDPA implementation.
vision_attn_implementation: str = "auto"
# ==================== Optimizer Presets ====================
optimizer_lr: float = 2e-5
optimizer_betas: tuple[float, float] = (0.9, 0.95)
@@ -86,6 +90,18 @@ class WallXConfig(PreTrainedConfig):
if self.prediction_mode not in ["diffusion", "fast"]:
raise ValueError(f"prediction_mode must be 'diffusion' or 'fast', got {self.prediction_mode}")
if self.attn_implementation != "eager":
raise ValueError(
"Wall-X currently supports only attn_implementation='eager' because its "
"bidirectional action-token islands require an explicit attention mask."
)
if self.vision_attn_implementation not in {"auto", "sdpa", "varlen"}:
raise ValueError(
"vision_attn_implementation must be one of 'auto', 'sdpa', or 'varlen', got "
f"{self.vision_attn_implementation!r}"
)
# Assign use_fast_tokenizer based on prediction_mode
if self.prediction_mode == "fast":
self.use_fast_tokenizer = True
src/lerobot/policies/wall_x/modeling_wall_x.py
+210 -128
View File
@@ -43,11 +43,14 @@ from typing import TYPE_CHECKING, Any
import numpy as np
import torch
import torch.nn as nn
import torch.nn.functional as F
from PIL import Image
import torch.nn.functional as functional
from safetensors import SafetensorError
from safetensors.torch import load_file
from torch import Tensor
from torch.distributions import Beta
from torch.nn import CrossEntropyLoss
from torchvision.transforms import InterpolationMode
from torchvision.transforms.v2 import functional as tv_functional
from lerobot.utils.constants import ACTION, OBS_STATE
from lerobot.utils.import_utils import (
@@ -74,17 +77,17 @@ if TYPE_CHECKING or _wallx_deps_available:
from qwen_vl_utils.vision_process import smart_resize
from torchdiffeq import odeint
from transformers import AutoProcessor, BatchFeature
from transformers.cache_utils import StaticCache
from transformers.models.qwen2_5_vl.modeling_qwen2_5_vl import (
Qwen2_5_VisionTransformerPretrainedModel,
Qwen2_5_VLForConditionalGeneration,
)
from transformers.utils import is_torchdynamo_compiling
from transformers.utils import cached_file, is_torchdynamo_compiling
from .qwen_model.configuration_qwen2_5_vl import Qwen2_5_VLConfig
from .qwen_model.qwen2_5_vl_moe import (
Qwen2_5_VisionTransformerPretrainedModel,
from .qwen_model import (
Qwen2_5_VLACausalLMOutputWithPast,
Qwen2_5_VLConfig,
Qwen2_5_VLMoEModel,
configure_wall_x_vision_attention,
)
else:
LoraConfig = None
@@ -93,13 +96,14 @@ else:
odeint = None
AutoProcessor = None
BatchFeature = None
StaticCache = None
Qwen2_5_VLForConditionalGeneration = None
cached_file = None
is_torchdynamo_compiling = None
Qwen2_5_VLConfig = None
Qwen2_5_VisionTransformerPretrainedModel = None
Qwen2_5_VLACausalLMOutputWithPast = None
Qwen2_5_VLMoEModel = None
configure_wall_x_vision_attention = None
from .utils import (
get_wallx_normal_text,
@@ -111,6 +115,75 @@ from .utils import (
logger = logging.getLogger(__name__)
def _wall_x_resize_dimensions(height: int, width: int) -> tuple[int, int, int, int]:
"""Return the intermediate and final Wall-X resize dimensions as ``(H, W, H, W)``."""
if RESOLUTION == -1:
intermediate_height, intermediate_width = height, width
elif width > height:
intermediate_width = RESOLUTION
intermediate_height = int(RESOLUTION * height / width)
else:
intermediate_height = RESOLUTION
intermediate_width = int(RESOLUTION * width / height)
resized_height, resized_width = smart_resize(
intermediate_height,
intermediate_width,
factor=IMAGE_FACTOR,
min_pixels=MIN_PIXELS,
max_pixels=MAX_PIXELS,
)
return intermediate_height, intermediate_width, resized_height, resized_width
def _resize_wall_x_image_batch(images: Tensor) -> tuple[Tensor, tuple[int, int, int, int]]:
"""Quantize and resize a BCHW camera batch without leaving its current device."""
if images.ndim != 4:
raise ValueError(f"Wall-X images must be BCHW tensors, got shape {tuple(images.shape)}")
original_height, original_width = images.shape[-2:]
intermediate_height, intermediate_width, resized_height, resized_width = _wall_x_resize_dimensions(
original_height, original_width
)
if images.is_floating_point():
# Match the previous PIL path, which quantized via `(image * 255).to(torch.uint8)`.
images = (images * 255).to(torch.uint8)
elif images.dtype != torch.uint8:
raise TypeError(f"Wall-X images must be floating point or uint8, got {images.dtype}")
if images.shape[-2:] != (intermediate_height, intermediate_width):
images = tv_functional.resize(
images,
[intermediate_height, intermediate_width],
interpolation=InterpolationMode.BICUBIC,
antialias=True,
)
if images.shape[-2:] != (resized_height, resized_width):
images = tv_functional.resize(
images,
[resized_height, resized_width],
interpolation=InterpolationMode.BICUBIC,
antialias=True,
)
return images, (original_height, original_width, resized_height, resized_width)
def _prepare_wall_x_image_inputs(
batch: dict[str, Any], img_keys: list[str]
) -> tuple[list[list[Tensor]], dict[str, tuple[int, int, int, int]]]:
"""Resize each camera as a batch, then restore sample-major/camera-minor ordering."""
resized_by_key: dict[str, Tensor] = {}
dimensions_by_key: dict[str, tuple[int, int, int, int]] = {}
for key in img_keys:
resized_by_key[key], dimensions_by_key[key] = _resize_wall_x_image_batch(batch[key])
batch_size = batch[img_keys[0]].shape[0]
image_inputs = [[resized_by_key[key][i] for key in img_keys] for i in range(batch_size)]
return image_inputs, dimensions_by_key
class SinusoidalPosEmb(nn.Module):
"""Sinusoidal positional embedding for diffusion timesteps."""
@@ -246,7 +319,7 @@ class ActionHead(nn.Module):
flow = flow.to(torch.float32)
action_pred = self.action_proj_back(action_hidden_states)
loss = F.mse_loss(action_pred, flow, reduction="none")
loss = functional.mse_loss(action_pred, flow, reduction="none")
if dof_mask is not None:
dof_mask = dof_mask.reshape(-1, dof_mask.shape[-1]).to(torch.float32)
@@ -254,7 +327,7 @@ class ActionHead(nn.Module):
return loss
def proprioception_proj(self, proprioception, dof_mask=None, use_history=False):
def proprioception_proj(self, proprioception, dof_mask=None):
"""Project proprioceptive data to hidden space."""
# Ensure proper device and dtype alignment
proprioception = proprioception.to(device=self.propri_proj.weight.device).to(
@@ -264,10 +337,7 @@ class ActionHead(nn.Module):
if dof_mask is not None:
# Concatenate proprioception with DOF mask
# TODO: Use variable-based dimension checking for better flexibility
if use_history:
proprioception = torch.cat([proprioception, dof_mask], dim=-1)
else:
proprioception = torch.cat([proprioception, dof_mask], dim=-1)
proprioception = torch.cat([proprioception, dof_mask], dim=-1)
proprioception = proprioception.to(device=self.propri_proj.weight.device).to(
dtype=self.propri_proj.weight.dtype
@@ -281,7 +351,7 @@ class ActionHead(nn.Module):
_Qwen2_5_VLForAction_Base = Qwen2_5_VLForConditionalGeneration if _wallx_deps_available else nn.Module
class Qwen2_5_VLMoEForAction(_Qwen2_5_VLForAction_Base):
class Qwen2_5_VLMoEForAction(_Qwen2_5_VLForAction_Base): # noqa: N801
"""
Qwen2.5 Vision-Language Mixture of Experts model for action processing.
@@ -305,6 +375,7 @@ class Qwen2_5_VLMoEForAction(_Qwen2_5_VLForAction_Base):
config=None,
action_tokenizer_path=None,
attn_implementation: str = "eager",
vision_attn_implementation: str = "auto",
cache_dir: str | PathLike | None = None,
force_download: bool = False,
local_files_only: bool = False,
@@ -321,11 +392,14 @@ class Qwen2_5_VLMoEForAction(_Qwen2_5_VLForAction_Base):
config_path (str, optional): Configuration file path, if None will look for qwen25_config.json in pretrained_model_path
action_tokenizer_path (str, optional): Action tokenizer path, if None will load from default config
attn_implementation (str, optional): Attention implementation, if None will load from default config
vision_attn_implementation (str, optional): Vision attention backend. ``auto`` uses packed
variable-length attention when supported and otherwise falls back to SDPA.
**kwargs: Additional arguments
Returns:
Qwen2_5_VLMoEForAction: Loaded model instance
"""
Qwen2_5_VLMoEModel._require_eager_attention(attn_implementation)
if config is None:
config = cls.config_class.from_pretrained(
pretrained_name_or_path,
@@ -339,7 +413,15 @@ class Qwen2_5_VLMoEForAction(_Qwen2_5_VLForAction_Base):
)
if attn_implementation is not None:
config._attn_implementation = attn_implementation
processor = AutoProcessor.from_pretrained(pretrained_name_or_path, use_fast=True)
processor = AutoProcessor.from_pretrained(
pretrained_name_or_path,
cache_dir=cache_dir,
force_download=force_download,
local_files_only=local_files_only,
token=token,
revision=revision,
use_fast=True,
)
if action_tokenizer_path is not None:
action_tokenizer = AutoProcessor.from_pretrained(action_tokenizer_path, trust_remote_code=True)
processor.action_processor = action_tokenizer
@@ -351,41 +433,41 @@ class Qwen2_5_VLMoEForAction(_Qwen2_5_VLForAction_Base):
config.text_config.pad_token_id = processor.tokenizer.pad_token_id
# Initialize model with configuration and processor
model = cls(config, processor=processor, action_tokenizer=action_tokenizer, **kwargs)
model = cls(
config,
processor=processor,
action_tokenizer=action_tokenizer,
vision_attn_implementation=vision_attn_implementation,
**kwargs,
)
# Resize token embeddings to match processor tokenizer vocabulary size
model.resize_token_embeddings(len(processor.tokenizer))
# Try to load the model.safetensors file
print(f"Loading model from: {pretrained_name_or_path}")
logger.info("Loading Wall-X model from %s", pretrained_name_or_path)
try:
from transformers.utils import cached_file
# Try safetensors first
resolved_file = cached_file(
pretrained_name_or_path,
"model.safetensors",
cache_dir=kwargs.get("cache_dir"),
force_download=kwargs.get("force_download", False),
cache_dir=cache_dir,
force_download=force_download,
resume_download=kwargs.get("resume_download"),
proxies=kwargs.get("proxies"),
token=kwargs.get("token"),
revision=kwargs.get("revision"),
local_files_only=kwargs.get("local_files_only", False),
token=token,
revision=revision,
local_files_only=local_files_only,
)
from safetensors.torch import load_file
sd = load_file(resolved_file)
print("✓ Loaded state dict from model.safetensors")
except Exception as e:
print(f"Could not load state dict from remote files: {e}")
print("Returning model without loading pretrained weights")
return model
except (OSError, SafetensorError) as error:
raise OSError(
f"Failed to load pretrained Wall-X weights from {pretrained_name_or_path!r}"
) from error
logger.info("Loaded Wall-X state dict from model.safetensors")
state_dict = {}
# filter normalizer statistic params
del_keys = []
for key in sd.keys():
for key in sd:
if "action_preprocessor.normalizer" in key:
del_keys.append(key)
for key in del_keys:
@@ -404,6 +486,7 @@ class Qwen2_5_VLMoEForAction(_Qwen2_5_VLForAction_Base):
action_tokenizer=None,
action_mapper=None,
flow_loss_weight=1.0,
vision_attn_implementation: str = "auto",
):
"""
Initialize the Qwen2.5 VLMoE model for action processing.
@@ -416,10 +499,16 @@ class Qwen2_5_VLMoEForAction(_Qwen2_5_VLForAction_Base):
action_mapper: Action mapping utility
flow_loss_weight (float): Weight for flow loss computation
"""
Qwen2_5_VLMoEModel._require_eager_attention(config._attn_implementation)
config._attn_implementation = "eager"
# Text needs eager attention for action-token islands. Vision has no such
# constraint, so keep its portable native fallback on SDPA.
config.vision_config._attn_implementation = "sdpa"
super().__init__(config)
# Initialize vision transformer and language model components
self.visual = Qwen2_5_VisionTransformerPretrainedModel._from_config(config.vision_config)
configure_wall_x_vision_attention(self.visual, vision_attn_implementation)
self.model = Qwen2_5_VLMoEModel(config)
self.vocab_size = config.vocab_size
self.lm_head = nn.Linear(config.hidden_size, config.vocab_size, bias=False)
@@ -457,7 +546,7 @@ class Qwen2_5_VLMoEForAction(_Qwen2_5_VLForAction_Base):
params_to_keep_float32 = []
for name, param in self.named_parameters():
for name, _param in self.named_parameters():
if "input_layernorm" in name or "post_attention_layernorm" in name or "model.norm" in name:
params_to_keep_float32.append(name)
if "action_preprocessor" in name:
@@ -491,7 +580,7 @@ class Qwen2_5_VLMoEForAction(_Qwen2_5_VLForAction_Base):
"action_token_id": action_token_id,
}
def add_lora(self, r=8, lora_alpha=32, target_modules=["q_proj", "v_proj"], lora_dropout=0.1):
def add_lora(self, r=8, lora_alpha=32, target_modules=None, lora_dropout=0.1):
"""
Add LoRA (Low-Rank Adaptation) adapters to the model.
@@ -501,6 +590,9 @@ class Qwen2_5_VLMoEForAction(_Qwen2_5_VLForAction_Base):
target_modules (list): List of module names to apply LoRA to
lora_dropout (float): Dropout probability for LoRA layers
"""
if target_modules is None:
target_modules = ["q_proj", "v_proj"]
config = LoraConfig(
r=r,
lora_alpha=lora_alpha,
@@ -795,6 +887,9 @@ class Qwen2_5_VLMoEForAction(_Qwen2_5_VLForAction_Base):
)
return_dict = return_dict if return_dict is not None else self.config.use_return_dict
if rope_deltas is not None:
self.rope_deltas = rope_deltas
# Calculate RoPE position IDs if not provided
# Note: Cannot calculate rope deltas with 4D attention mask. TODO: Fix this limitation
if position_ids is None and (attention_mask is None or attention_mask.ndim == 2):
@@ -833,7 +928,7 @@ class Qwen2_5_VLMoEForAction(_Qwen2_5_VLForAction_Base):
# Process image embeddings
if pixel_values is not None:
pixel_values = pixel_values.type(self.visual.dtype)
image_embeds = self.visual(pixel_values, grid_thw=image_grid_thw)
image_embeds = self.visual(pixel_values, grid_thw=image_grid_thw).pooler_output
mask = input_ids == self.config.image_token_id
mask_unsqueezed = mask.unsqueeze(-1)
mask_expanded = mask_unsqueezed.expand_as(inputs_embeds)
@@ -845,7 +940,7 @@ class Qwen2_5_VLMoEForAction(_Qwen2_5_VLForAction_Base):
# Process video embeddings
if pixel_values_videos is not None:
pixel_values_videos = pixel_values_videos.type(self.visual.dtype)
video_embeds = self.visual(pixel_values_videos, grid_thw=video_grid_thw)
video_embeds = self.visual(pixel_values_videos, grid_thw=video_grid_thw).pooler_output
n_video_tokens = (input_ids == self.config.video_token_id).sum().item()
n_video_features = video_embeds.shape[0]
@@ -869,7 +964,6 @@ class Qwen2_5_VLMoEForAction(_Qwen2_5_VLForAction_Base):
proprioception = self.action_preprocessor.proprioception_proj(
proprioception,
agent_pos_mask,
use_history=proprioception.shape[1] > 1,
)
mask = input_ids == self.action_token_id_set["propri_token_id"]
mask_unsqueezed = mask.unsqueeze(-1)
@@ -919,6 +1013,7 @@ class Qwen2_5_VLMoEForAction(_Qwen2_5_VLForAction_Base):
output_attentions=output_attentions,
output_hidden_states=output_hidden_states,
return_dict=return_dict,
cache_position=cache_position,
)
hidden_states = outputs[0]
@@ -1107,7 +1202,7 @@ class Qwen2_5_VLMoEForAction(_Qwen2_5_VLForAction_Base):
# Process image embeddings
if pixel_values is not None:
pixel_values = pixel_values.type(self.visual.dtype)
image_embeds = self.visual(pixel_values, grid_thw=image_grid_thw)
image_embeds = self.visual(pixel_values, grid_thw=image_grid_thw).pooler_output
n_image_tokens = (input_ids == self.config.image_token_id).sum().item()
n_image_features = image_embeds.shape[0]
@@ -1128,7 +1223,7 @@ class Qwen2_5_VLMoEForAction(_Qwen2_5_VLForAction_Base):
# Process video embeddings
if pixel_values_videos is not None:
pixel_values_videos = pixel_values_videos.type(self.visual.dtype)
video_embeds = self.visual(pixel_values_videos, grid_thw=video_grid_thw)
video_embeds = self.visual(pixel_values_videos, grid_thw=video_grid_thw).pooler_output
n_video_tokens = (input_ids == self.config.video_token_id).sum().item()
n_video_features = video_embeds.shape[0]
@@ -1153,7 +1248,6 @@ class Qwen2_5_VLMoEForAction(_Qwen2_5_VLForAction_Base):
proprio_embed = self.action_preprocessor.proprioception_proj(
proprioception,
agent_pos_mask,
use_history=proprioception.shape[1] > 1,
)
proprioception_mask = input_ids == self.action_token_id_set["propri_token_id"]
proprio_embed = proprio_embed.to(torch.bfloat16)
@@ -1202,25 +1296,37 @@ class Qwen2_5_VLMoEForAction(_Qwen2_5_VLForAction_Base):
# Split input sequence for text and fast modes (not needed for diffusion)
if predict_mode == "text" or predict_mode == "fast":
# Look for generation prompt tokens: <|im_start|>assistant
generation_prompt = "<|im_start|>assistant\n"
generation_prompt_ids = torch.tensor(
[151644, 77091], device=input_ids.device, dtype=input_ids.dtype
)
matches = (input_ids[0, :-1] == generation_prompt_ids[0]) & (
input_ids[0, 1:] == generation_prompt_ids[1]
self.processor.tokenizer.encode(generation_prompt, add_special_tokens=False),
device=input_ids.device,
dtype=input_ids.dtype,
)
prompt_length = generation_prompt_ids.numel()
if prompt_length == 0:
raise ValueError(f"Tokenizer produced no tokens for generation prompt {generation_prompt!r}")
if input_ids.shape[1] < prompt_length:
matches = torch.empty(0, device=input_ids.device, dtype=torch.bool)
else:
matches = (
input_ids[0]
.unfold(dimension=0, size=prompt_length, step=1)
.eq(generation_prompt_ids)
.all(dim=-1)
)
if matches.any():
split_pos = torch.nonzero(matches, as_tuple=True)[0][0].item()
prompt_end = split_pos + prompt_length
# Extract ground truth output tokens (including newline)
gt_output_ids = input_ids[:, split_pos + 3 :]
gt_output_ids = input_ids[:, prompt_end:]
# Remove output part from input, keeping prompt
input_ids = input_ids[:, : split_pos + 3]
inputs_embeds = inputs_embeds[:, : split_pos + 3, :]
input_ids = input_ids[:, :prompt_end]
inputs_embeds = inputs_embeds[:, :prompt_end, :]
if attention_mask is not None:
attention_mask = attention_mask[:, : split_pos + 3]
attention_mask = attention_mask[:, :prompt_end]
if labels is not None:
labels = labels[:, split_pos + 3 :]
labels = labels[:, prompt_end:]
else:
raise ValueError(
"input_ids does not contain the generation prompt tokens <|im_start|>assistant"
@@ -1255,7 +1361,7 @@ class Qwen2_5_VLMoEForAction(_Qwen2_5_VLForAction_Base):
use_cache=True,
pad_token_id=self.processor.tokenizer.pad_token_id,
temperature=(1.0 if not re_generate else 0.7), # Higher temperature for regeneration
do_sample=(False if not re_generate else True), # Enable sampling for regeneration
do_sample=re_generate, # Enable sampling for regeneration
)
# Decode generated and ground truth text
@@ -1524,27 +1630,6 @@ class Qwen2_5_VLMoEForAction(_Qwen2_5_VLForAction_Base):
else:
model_inputs = {"input_ids": input_ids, "inputs_embeds": None}
# Prepare 4D causal attention mask for static cache
if isinstance(past_key_values, StaticCache) and attention_mask.ndim == 2:
if model_inputs["inputs_embeds"] is not None:
batch_size, sequence_length, _ = inputs_embeds.shape
device = inputs_embeds.device
else:
batch_size, sequence_length = input_ids.shape
device = input_ids.device
attention_mask = self.model._prepare_4d_causal_attention_mask_with_cache_position(
attention_mask,
sequence_length=sequence_length,
target_length=past_key_values.get_max_cache_shape(),
dtype=self.lm_head.weight.dtype,
device=device,
cache_position=cache_position,
batch_size=batch_size,
config=self.config,
past_key_values=past_key_values,
)
# Assemble all model inputs for generation
model_inputs.update(
{
@@ -1749,6 +1834,7 @@ class WallXPolicy(PreTrainedPolicy):
pretrained_name_or_path=config.pretrained_name_or_path,
action_tokenizer_path=config.action_tokenizer_path,
attn_implementation=config.attn_implementation,
vision_attn_implementation=config.vision_attn_implementation,
)
self.model.to(config.device)
self.model.to_bfloat16_for_selected_params()
@@ -1768,6 +1854,8 @@ class WallXPolicy(PreTrainedPolicy):
def preprocess_inputs(
self,
batch: dict[str, Any],
*,
compute_position_ids: bool = False,
) -> BatchFeature:
"""
Convert a batch of LeRobot dataset items to Wall-X model input format.
@@ -1789,50 +1877,21 @@ class WallXPolicy(PreTrainedPolicy):
# Get batch size from state tensor
batch_size = batch[OBS_STATE].shape[0]
# ==================== PROCESS ALL SAMPLES ====================
all_image_inputs = []
all_texts = []
# Find image keys in batch
img_keys = [key for key in self.config.image_features if key in batch]
if not img_keys:
raise ValueError("Wall-X requires at least one image feature in each batch")
# Resize one camera batch at a time on the tensors' current device. Reassembling
# sample-major keeps image_grid_thw aligned with each sample's image placeholders.
all_image_inputs, dimensions_by_key = _prepare_wall_x_image_inputs(batch, img_keys)
all_texts = []
# Preserve the existing grounding behavior for multi-camera inputs: the old camera
# loop left these values set to the final configured camera's dimensions.
orig_height, orig_width, resized_height, resized_width = dimensions_by_key[img_keys[-1]]
for i in range(batch_size):
# Vision preprocessing per sample
processed_frames = []
orig_height, orig_width = None, None
resized_height, resized_width = None, None
for key in img_keys:
current_obs = batch[key][i].clone() # (C, H, W)
if current_obs.dim() == 3:
current_obs = current_obs.permute(1, 2, 0) # (H, W, C)
img_pil = Image.fromarray((current_obs * 255).to(torch.uint8).cpu().numpy())
orig_width, orig_height = img_pil.size
target_size = RESOLUTION
if target_size != -1:
if orig_width > orig_height:
new_width = target_size
new_height = int(target_size * orig_height / orig_width)
else:
new_height = target_size
new_width = int(target_size * orig_width / orig_height)
img_pil = img_pil.resize((new_width, new_height))
current_width, current_height = img_pil.size
resized_height, resized_width = smart_resize(
current_height,
current_width,
factor=IMAGE_FACTOR,
min_pixels=MIN_PIXELS,
max_pixels=MAX_PIXELS,
)
resized_img = img_pil.resize((resized_width, resized_height))
processed_frames.append(resized_img)
all_image_inputs.append(processed_frames)
# Text preprocessing
task_text = batch["task"][i] if isinstance(batch["task"], list) else batch["task"]
instruction_info = {"instruction": task_text}
@@ -1859,8 +1918,8 @@ class WallXPolicy(PreTrainedPolicy):
agent_pos_mask = (~torch.isnan(agent_pos)).float()
agent_pos = agent_pos.nan_to_num(nan=0.0)
if agent_pos.shape[-1] != 20:
pad_size = 20 - agent_pos.shape[-1]
if agent_pos.shape[-1] < self.config.max_state_dim:
pad_size = self.config.max_state_dim - agent_pos.shape[-1]
agent_pos = torch.cat(
[
agent_pos,
@@ -1880,6 +1939,10 @@ class WallXPolicy(PreTrainedPolicy):
],
dim=-1,
)
elif agent_pos.shape[-1] > self.config.max_state_dim:
raise ValueError(
f"State dimension {agent_pos.shape[-1]} exceeds max_state_dim {self.config.max_state_dim}"
)
# ==================== PROCESS ACTIONS ====================
action = batch.get(ACTION) # (batch_size, chunk_size, action_dim)
@@ -1889,8 +1952,8 @@ class WallXPolicy(PreTrainedPolicy):
dof_mask = (~torch.isnan(action)).float()
action = action.nan_to_num(nan=0.0)
if action.shape[-1] != 20:
pad_size = 20 - action.shape[-1]
if action.shape[-1] < self.config.max_action_dim:
pad_size = self.config.max_action_dim - action.shape[-1]
action = torch.cat(
[action, torch.zeros(action.shape[0], action.shape[1], pad_size, device=action.device)],
dim=-1,
@@ -1902,6 +1965,10 @@ class WallXPolicy(PreTrainedPolicy):
],
dim=-1,
)
elif action.shape[-1] > self.config.max_action_dim:
raise ValueError(
f"Action dimension {action.shape[-1]} exceeds max_action_dim {self.config.max_action_dim}"
)
else:
action_dim = self.config.output_features[ACTION].shape[0]
dof_mask = torch.cat(
@@ -1910,7 +1977,10 @@ class WallXPolicy(PreTrainedPolicy):
batch_size, self.config.chunk_size, action_dim, device=batch[OBS_STATE].device
),
torch.zeros(
batch_size, self.config.chunk_size, 20 - action_dim, device=batch[OBS_STATE].device
batch_size,
self.config.chunk_size,
self.config.max_action_dim - action_dim,
device=batch[OBS_STATE].device,
),
],
dim=-1,
@@ -1930,12 +2000,26 @@ class WallXPolicy(PreTrainedPolicy):
text=all_texts,
images=all_image_inputs,
videos=None,
device=batch[OBS_STATE].device,
padding=True,
truncation=True,
return_tensors="pt",
max_length=TOKENIZER_MAX_LENGTH,
)
if compute_position_ids:
# Qwen's RoPE indexing uses Python list/scalar conversions. Run it while the
# tokenizer and grid metadata are still on CPU, then move the compact result.
position_ids, rope_deltas = self.model.get_rope_index(
inputs.input_ids,
inputs.get("image_grid_thw"),
inputs.get("video_grid_thw"),
inputs.get("second_per_grid_ts"),
inputs.attention_mask,
)
inputs["position_ids"] = position_ids
inputs["rope_deltas"] = rope_deltas
# ==================== ADDITIONAL INPUTS ====================
action_token_id = self.model.processor.tokenizer.convert_tokens_to_ids("<|action|>")
moe_token_types = inputs.input_ids == action_token_id
@@ -1952,7 +2036,7 @@ class WallXPolicy(PreTrainedPolicy):
)
# Move all tensors to the correct device
device = self.config.device
device = batch[OBS_STATE].device
for key, value in inputs.items():
if isinstance(value, torch.Tensor):
inputs[key] = value.to(device)
@@ -1972,9 +2056,7 @@ class WallXPolicy(PreTrainedPolicy):
Returns:
tuple: (loss, loss_dict)
"""
batch = self.preprocess_inputs(
batch,
)
batch = self.preprocess_inputs(batch, compute_position_ids=True)
# Call the underlying model's forward with mode="train"
outputs = self.model(**batch, mode="train")
@@ -1982,19 +2064,19 @@ class WallXPolicy(PreTrainedPolicy):
# Extract losses from output
loss = outputs.loss
loss_dict = {
"loss": loss.item() if loss is not None else 0.0,
"loss": loss.detach() if loss is not None else 0.0,
}
if outputs.flow_loss is not None:
loss_dict["flow_loss"] = outputs.flow_loss.item()
loss_dict["flow_loss"] = outputs.flow_loss.detach()
if outputs.cross_entropy_loss is not None:
loss_dict["cross_entropy_loss"] = outputs.cross_entropy_loss.item()
loss_dict["cross_entropy_loss"] = outputs.cross_entropy_loss.detach()
# Add channel losses if available
if outputs.channel_loss_dict is not None:
for key, value in outputs.channel_loss_dict.items():
if isinstance(value, torch.Tensor):
loss_dict[f"channel_{key}"] = value.item()
loss_dict[f"channel_{key}"] = value.detach()
return loss, loss_dict
src/lerobot/policies/wall_x/qwen_model/__init__.py
+45
View File
@@ -0,0 +1,45 @@
#!/usr/bin/env python
# Copyright 2025 HuggingFace Inc. team. All rights reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
from .configuration_qwen2_5_vl import (
Qwen2_5_VLConfig,
Qwen2_5_VLTextConfig,
Qwen2_5_VLVisionConfig,
)
from .qwen2_5_vl_moe import (
BlockSparseMLP,
Qwen2_5_VLACausalLMOutputWithPast,
Qwen2_5_VLDecoderLayer_with_MoE,
Qwen2_5_VLMoEModel,
SparseMoeBlock,
)
from .vision_attention import (
WallXVisionAttention,
configure_wall_x_vision_attention,
)
__all__ = [
"BlockSparseMLP",
"Qwen2_5_VLACausalLMOutputWithPast",
"Qwen2_5_VLConfig",
"Qwen2_5_VLDecoderLayer_with_MoE",
"Qwen2_5_VLMoEModel",
"Qwen2_5_VLTextConfig",
"Qwen2_5_VLVisionConfig",
"SparseMoeBlock",
"WallXVisionAttention",
"configure_wall_x_vision_attention",
]
src/lerobot/policies/wall_x/qwen_model/configuration_qwen2_5_vl.py
+103 -239
View File
@@ -1,250 +1,114 @@
from transformers.configuration_utils import PretrainedConfig
from transformers.modeling_rope_utils import rope_config_validation
#!/usr/bin/env python
# Copyright 2025 HuggingFace Inc. team. All rights reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
"""Wall-X configuration extensions for the native Transformers Qwen2.5-VL config."""
from dataclasses import dataclass
from typing import TYPE_CHECKING
from huggingface_hub.dataclasses import strict
from lerobot.utils.import_utils import _transformers_available
if TYPE_CHECKING or _transformers_available:
from transformers.models.qwen2_5_vl.configuration_qwen2_5_vl import (
Qwen2_5_VLConfig as TransformersQwen2_5_VLConfig,
Qwen2_5_VLTextConfig as TransformersQwen2_5_VLTextConfig,
Qwen2_5_VLVisionConfig,
)
else:
@dataclass
class _TransformersConfigFallback:
"""Import-safe stand-in used only when Transformers is unavailable."""
TransformersQwen2_5_VLConfig = _TransformersConfigFallback
TransformersQwen2_5_VLTextConfig = _TransformersConfigFallback
Qwen2_5_VLVisionConfig = None
# Wall-X checkpoints pre0.6.0 use the legacy, flat Qwen2.5-VL config layout. The native
# ``Qwen2_5_VLConfig`` accepts that layout and moves text-model fields into its
# ``text_config`` sub-config, so only the Wall-X-specific MoE fields need to be
# declared here.
_LEGACY_TEXT_ATTRIBUTES = {
"attention_dropout",
"attention_moe",
"dim_inputs",
"dof_config",
"experts",
"hidden_act",
"hidden_size",
"initializer_range",
"intermediate_size",
"layer_types",
"max_position_embeddings",
"max_window_layers",
"mlp_moe",
"noise_scheduler",
"num_attention_heads",
"num_experts",
"num_hidden_layers",
"num_key_value_heads",
"pad_token_id",
"rms_norm_eps",
"sliding_window",
"use_cache",
"use_sliding_window",
"vocab_size",
}
class Qwen2_5_VLVisionConfig(PretrainedConfig):
model_type = "qwen2_5_vl"
base_config_key = "vision_config"
@strict
class Qwen2_5_VLTextConfig(TransformersQwen2_5_VLTextConfig): # noqa: N801
"""Native Qwen2.5-VL text config plus Wall-X's hard-routed MoE settings."""
def __init__(
self,
depth=32,
hidden_size=3584,
hidden_act="silu",
intermediate_size=3420,
num_heads=16,
in_channels=3,
patch_size=14,
spatial_merge_size=2,
temporal_patch_size=2,
tokens_per_second=4,
window_size=112,
out_hidden_size=3584,
fullatt_block_indexes=[7, 15, 23, 31],
initializer_range=0.02,
**kwargs,
):
super().__init__(**kwargs)
num_experts: int = 4
experts: list[dict] | None = None
dof_config: dict | None = None
noise_scheduler: dict | None = None
dim_inputs: tuple[int, ...] | list[int] = (1536, 1536)
attention_moe: bool = False
mlp_moe: bool = False
self.depth = depth
self.hidden_size = hidden_size
self.hidden_act = hidden_act
self.intermediate_size = intermediate_size
self.num_heads = num_heads
self.in_channels = in_channels
self.patch_size = patch_size
self.spatial_merge_size = spatial_merge_size
self.temporal_patch_size = temporal_patch_size
self.tokens_per_second = tokens_per_second
self.window_size = window_size
self.fullatt_block_indexes = fullatt_block_indexes
self.out_hidden_size = out_hidden_size
self.initializer_range = initializer_range
def __post_init__(self, **kwargs):
self.dim_inputs = tuple(self.dim_inputs)
super().__post_init__(**kwargs)
class Qwen2_5_VLConfig(PretrainedConfig):
r"""
This is the configuration class to store the configuration of a [`Qwen2_5_VLModel`]. It is used to instantiate a
Qwen2-VL model according to the specified arguments, defining the model architecture. Instantiating a configuration
with the defaults will yield a similar configuration to that of
Qwen2-VL-7B-Instruct [Qwen/Qwen2-VL-7B-Instruct](https://huggingface.co/Qwen/Qwen2-VL-7B-Instruct).
@strict
class Qwen2_5_VLConfig(TransformersQwen2_5_VLConfig): # noqa: N801
"""Native composite Qwen2.5-VL config with a Wall-X text sub-config.
Configuration objects inherit from [`PretrainedConfig`] and can be used to control the model outputs. Read the
documentation from [`PretrainedConfig`] for more information.
The native composite loader supports both current nested configs and the
flat layout used by existing ``wall-oss-flow`` checkpoints.
"""
Args:
vocab_size (`int`, *optional*, defaults to 152064):
Vocabulary size of the Qwen2_5_VL model. Defines the number of different tokens that can be represented by the
`inputs_ids` passed when calling [`Qwen2_5_VLModel`]
hidden_size (`int`, *optional*, defaults to 8192):
Dimension of the hidden representations.
intermediate_size (`int`, *optional*, defaults to 29568):
Dimension of the MLP representations.
num_hidden_layers (`int`, *optional*, defaults to 80):
Number of hidden layers in the Transformer encoder.
num_attention_heads (`int`, *optional*, defaults to 64):
Number of attention heads for each attention layer in the Transformer encoder.
num_key_value_heads (`int`, *optional*, defaults to 8):
This is the number of key_value heads that should be used to implement Grouped Query Attention. If
`num_key_value_heads=num_attention_heads`, the model will use Multi Head Attention (MHA), if
`num_key_value_heads=1` the model will use Multi Query Attention (MQA) otherwise GQA is used. When
converting a multi-head checkpoint to a GQA checkpoint, each group key and value head should be constructed
by meanpooling all the original heads within that group. For more details checkout [this
paper](https://arxiv.org/pdf/2305.13245.pdf). If it is not specified, will default to `32`.
hidden_act (`str` or `function`, *optional*, defaults to `"silu"`):
The non-linear activation function (function or string) in the decoder.
max_position_embeddings (`int`, *optional*, defaults to 32768):
The maximum sequence length that this model might ever be used with.
initializer_range (`float`, *optional*, defaults to 0.02):
The standard deviation of the truncated_normal_initializer for initializing all weight matrices.
rms_norm_eps (`float`, *optional*, defaults to 1e-05):
The epsilon used by the rms normalization layers.
use_cache (`bool`, *optional*, defaults to `True`):
Whether or not the model should return the last key/values attentions (not used by all models). Only
relevant if `config.is_decoder=True`.
tie_word_embeddings (`bool`, *optional*, defaults to `False`):
Whether the model's input and output word embeddings should be tied.
rope_theta (`float`, *optional*, defaults to 1000000.0):
The base period of the RoPE embeddings.
use_sliding_window (`bool`, *optional*, defaults to `False`):
Whether to use sliding window attention.
sliding_window (`int`, *optional*, defaults to 4096):
Sliding window attention (SWA) window size. If not specified, will default to `4096`.
max_window_layers (`int`, *optional*, defaults to 80):
The number of layers that use SWA (Sliding Window Attention). The bottom layers use SWA while the top use full attention.
attention_dropout (`float`, *optional*, defaults to 0.0):
The dropout ratio for the attention probabilities.
vision_config (`Dict`, *optional*):
The config for the visual encoder initialization.
rope_scaling (`Dict`, *optional*):
Dictionary containing the scaling configuration for the RoPE embeddings. NOTE: if you apply new rope type
and you expect the model to work on longer `max_position_embeddings`, we recommend you to update this value
accordingly.
Expected contents:
`rope_type` (`str`):
The sub-variant of RoPE to use. Can be one of ['default', 'linear', 'dynamic', 'yarn', 'longrope',
'llama3'], with 'default' being the original RoPE implementation.
`factor` (`float`, *optional*):
Used with all rope types except 'default'. The scaling factor to apply to the RoPE embeddings. In
most scaling types, a `factor` of x will enable the model to handle sequences of length x *
original maximum pre-trained length.
`original_max_position_embeddings` (`int`, *optional*):
Used with 'dynamic', 'longrope' and 'llama3'. The original max position embeddings used during
pretraining.
`attention_factor` (`float`, *optional*):
Used with 'yarn' and 'longrope'. The scaling factor to be applied on the attention
computation. If unspecified, it defaults to value recommended by the implementation, using the
`factor` field to infer the suggested value.
`beta_fast` (`float`, *optional*):
Only used with 'yarn'. Parameter to set the boundary for extrapolation (only) in the linear
ramp function. If unspecified, it defaults to 32.
`beta_slow` (`float`, *optional*):
Only used with 'yarn'. Parameter to set the boundary for interpolation (only) in the linear
ramp function. If unspecified, it defaults to 1.
`short_factor` (`List[float]`, *optional*):
Only used with 'longrope'. The scaling factor to be applied to short contexts (<
`original_max_position_embeddings`). Must be a list of numbers with the same length as the hidden
size divided by the number of attention heads divided by 2
`long_factor` (`List[float]`, *optional*):
Only used with 'longrope'. The scaling factor to be applied to long contexts (<
`original_max_position_embeddings`). Must be a list of numbers with the same length as the hidden
size divided by the number of attention heads divided by 2
`low_freq_factor` (`float`, *optional*):
Only used with 'llama3'. Scaling factor applied to low frequency components of the RoPE
`high_freq_factor` (`float`, *optional*):
Only used with 'llama3'. Scaling factor applied to high frequency components of the RoPE
```python
>>> from transformers import Qwen2_5_VLForConditionalGeneration, Qwen2_5_VLConfig
>>> # Initializing a Qwen2_5_VL style configuration
>>> configuration = Qwen2_5_VLConfig()
>>> # Initializing a model from the Qwen2-VL-7B style configuration
>>> model = Qwen2_5_VLForConditionalGeneration(configuration)
>>> # Accessing the model configuration
>>> configuration = model.config
```"""
model_type = "qwen2_5_vl"
sub_configs = {"vision_config": Qwen2_5_VLVisionConfig}
keys_to_ignore_at_inference = ["past_key_values"]
# Default tensor parallel plan for base model `Qwen2_5_VL`
base_model_tp_plan = {
"layers.*.self_attn.q_proj": "colwise",
"layers.*.self_attn.k_proj": "colwise",
"layers.*.self_attn.v_proj": "colwise",
"layers.*.self_attn.o_proj": "rowwise",
"layers.*.mlp.gate_proj": "colwise",
"layers.*.mlp.up_proj": "colwise",
"layers.*.mlp.down_proj": "rowwise",
}
base_model_pp_plan = {
"embed_tokens": (["input_ids"], ["inputs_embeds"]),
"layers": (["hidden_states", "attention_mask"], ["hidden_states"]),
"norm": (["hidden_states"], ["hidden_states"]),
sub_configs = {
"vision_config": Qwen2_5_VLVisionConfig,
"text_config": Qwen2_5_VLTextConfig,
}
def __init__(
self,
vocab_size=152064,
hidden_size=8192,
intermediate_size=29568,
num_hidden_layers=80,
num_attention_heads=64,
num_key_value_heads=8,
hidden_act="silu",
max_position_embeddings=32768,
initializer_range=0.02,
rms_norm_eps=1e-05,
use_cache=True,
tie_word_embeddings=False,
rope_theta=1000000.0,
use_sliding_window=False,
sliding_window=4096,
max_window_layers=80,
attention_dropout=0.0,
vision_config=None,
rope_scaling=None,
num_experts=4,
experts=None,
dof_config=None,
noise_scheduler=None,
dim_inputs=(1536, 1536),
attention_moe=False,
mlp_moe=False,
**kwargs,
):
if isinstance(vision_config, dict):
self.vision_config = self.sub_configs["vision_config"](**vision_config)
elif vision_config is None:
self.vision_config = self.sub_configs["vision_config"]()
def __getattr__(self, name):
"""Keep legacy direct access to fields now owned by ``text_config``.
self.vocab_size = vocab_size
self.max_position_embeddings = max_position_embeddings
self.hidden_size = hidden_size
self.intermediate_size = intermediate_size
self.num_hidden_layers = num_hidden_layers
self.num_attention_heads = num_attention_heads
self.use_sliding_window = use_sliding_window
self.sliding_window = sliding_window
self.max_window_layers = max_window_layers
self.layer_types = ["dense"] * num_hidden_layers
# for backward compatibility
if num_key_value_heads is None:
num_key_value_heads = num_attention_heads
self.num_key_value_heads = num_key_value_heads
self.hidden_act = hidden_act
self.initializer_range = initializer_range
self.rms_norm_eps = rms_norm_eps
self.use_cache = use_cache
self.rope_theta = rope_theta
self.attention_dropout = attention_dropout
self.rope_scaling = rope_scaling
self.num_experts = num_experts
self.experts = experts
self.dof_config = dof_config
self.noise_scheduler = noise_scheduler
self.dim_inputs = tuple(dim_inputs)
self.attention_moe = attention_moe
self.mlp_moe = mlp_moe
if self.rope_scaling is not None and "type" in self.rope_scaling:
if self.rope_scaling["type"] == "mrope":
self.rope_scaling["type"] = "default"
self.rope_scaling["rope_type"] = self.rope_scaling["type"]
rope_config_validation(self, ignore_keys={"mrope_section"})
super().__init__(tie_word_embeddings=tie_word_embeddings, **kwargs)
@property
def text_config(self):
return self
__all__ = ["Qwen2_5_VLConfig"]
Wall-X historically used a flat config and accesses fields such as
``hidden_size`` and ``num_experts`` directly. Forwarding unknown
attributes preserves that API without duplicating the native config.
"""
text_config = self.__dict__.get("text_config")
if name in _LEGACY_TEXT_ATTRIBUTES and text_config is not None and hasattr(text_config, name):
return getattr(text_config, name)
raise AttributeError(f"{type(self).__name__!s} has no attribute {name!r}")
src/lerobot/policies/wall_x/qwen_model/qwen2_5_vl_moe.py
+209 -2635
View File
File diff suppressed because it is too large Load Diff
src/lerobot/policies/wall_x/qwen_model/vision_attention.py
+208
View File
@@ -0,0 +1,208 @@
#!/usr/bin/env python
# Copyright 2025 HuggingFace Inc. team. All rights reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
"""Wall-X vision attention backends.
Qwen2.5-VL's native non-Flash vision path splits a packed image sequence into
Python-level chunks before calling attention. Wall-X batches many camera frames,
so that path launches thousands of tiny attention operations per training step.
This module keeps the native SDPA path as a portable fallback and adds a packed
``torch.nn.attention.varlen`` path that consumes Qwen's existing ``cu_seqlens``
metadata directly.
"""
from __future__ import annotations
import inspect
import logging
from functools import lru_cache
from typing import TYPE_CHECKING, Literal
import torch
import torch.nn as nn
from lerobot.utils.import_utils import _transformers_available
if TYPE_CHECKING or _transformers_available:
from transformers.models.qwen2_5_vl.modeling_qwen2_5_vl import (
Qwen2_5_VLVisionAttention,
apply_rotary_pos_emb_vision,
)
else:
Qwen2_5_VLVisionAttention = nn.Module
apply_rotary_pos_emb_vision = None
try:
from torch.nn.attention.varlen import varlen_attn as _varlen_attn
except ImportError: # torch<2.10
_varlen_attn = None
_VARLEN_USES_WINDOW_SIZE = (
_varlen_attn is not None and "window_size" in inspect.signature(_varlen_attn).parameters
)
VisionAttentionBackend = Literal["auto", "sdpa", "varlen"]
logger = logging.getLogger(__name__)
@lru_cache
def _log_resolved_backend(requested: str, resolved: str) -> None:
logger.info("Wall-X vision attention backend: %s (requested: %s)", resolved, requested)
def _varlen_unavailable_reason(
hidden_states: torch.Tensor,
position_embeddings: tuple[torch.Tensor, torch.Tensor] | None,
) -> str | None:
if _varlen_attn is None:
return "torch.nn.attention.varlen is unavailable (PyTorch 2.10 or newer is required)"
if position_embeddings is None:
return "precomputed vision position embeddings were not provided"
if hidden_states.device.type != "cuda" or torch.version.cuda is None:
return "packed varlen attention requires an NVIDIA CUDA device"
if hidden_states.dtype not in {torch.float16, torch.bfloat16}:
return f"packed varlen attention requires float16 or bfloat16 inputs, got {hidden_states.dtype}"
major, _minor = torch.cuda.get_device_capability(hidden_states.device)
if major < 8:
return "packed varlen attention requires an NVIDIA Ampere GPU or newer"
return None
def _supports_varlen_attention(
hidden_states: torch.Tensor,
position_embeddings: tuple[torch.Tensor, torch.Tensor] | None,
) -> bool:
return _varlen_unavailable_reason(hidden_states, position_embeddings) is None
class WallXVisionAttention(Qwen2_5_VLVisionAttention):
"""Qwen2.5-VL vision attention with packed varlen and native SDPA fallback."""
def __init__(self, config, backend: VisionAttentionBackend):
super().__init__(config)
self.wallx_backend = backend
self._resolved_backend_key = None
self._resolved_backend = None
def _resolve_backend(
self,
hidden_states: torch.Tensor,
position_embeddings: tuple[torch.Tensor, torch.Tensor] | None,
) -> str:
key = (
hidden_states.device.type,
hidden_states.device.index,
hidden_states.dtype,
position_embeddings is not None,
)
if self._resolved_backend_key == key:
return self._resolved_backend
use_varlen = self.wallx_backend != "sdpa" and _supports_varlen_attention(
hidden_states, position_embeddings
)
if self.wallx_backend == "varlen" and not use_varlen:
reason = _varlen_unavailable_reason(hidden_states, position_embeddings)
raise RuntimeError(f"Wall-X vision_attn_implementation='varlen' cannot be used: {reason}")
resolved_backend = "varlen" if use_varlen else "sdpa"
self._resolved_backend_key = key
self._resolved_backend = resolved_backend
_log_resolved_backend(self.wallx_backend, resolved_backend)
return resolved_backend
def forward(
self,
hidden_states: torch.Tensor,
cu_seqlens: torch.Tensor,
rotary_pos_emb: torch.Tensor | None = None,
position_embeddings: tuple[torch.Tensor, torch.Tensor] | None = None,
**kwargs,
) -> torch.Tensor:
del rotary_pos_emb
if self._resolve_backend(hidden_states, position_embeddings) == "sdpa":
return super().forward(
hidden_states=hidden_states,
cu_seqlens=cu_seqlens,
position_embeddings=position_embeddings,
**kwargs,
)
seq_length = hidden_states.shape[0]
query_states, key_states, value_states = (
self.qkv(hidden_states).reshape(seq_length, 3, self.num_heads, -1).permute(1, 0, 2, 3).unbind(0)
)
cos, sin = position_embeddings
query_states, key_states = apply_rotary_pos_emb_vision(
query_states,
key_states,
cos,
sin,
)
if cu_seqlens.dtype != torch.int32:
cu_seqlens = cu_seqlens.to(dtype=torch.int32)
max_seqlen = int((cu_seqlens[1:] - cu_seqlens[:-1]).max().item())
varlen_kwargs = {"scale": self.scaling}
if _VARLEN_USES_WINDOW_SIZE:
varlen_kwargs["window_size"] = (-1, -1)
else: # Stable PyTorch 2.10 API; pre-release variants used window_size.
varlen_kwargs["is_causal"] = False
attn_output = _varlen_attn(
query_states,
key_states,
value_states,
cu_seqlens,
cu_seqlens,
max_seqlen,
max_seqlen,
**varlen_kwargs,
)
attn_output = attn_output.reshape(seq_length, -1).contiguous()
return self.proj(attn_output)
def configure_wall_x_vision_attention(
vision_model: nn.Module,
backend: VisionAttentionBackend,
) -> None:
"""Install Wall-X's scoped packed attention without changing checkpoint keys."""
if backend == "sdpa":
_log_resolved_backend(backend, "sdpa")
return
if backend == "varlen" and _varlen_attn is None:
raise RuntimeError(
"Wall-X vision_attn_implementation='varlen' requires torch.nn.attention.varlen "
"from PyTorch 2.10 or newer"
)
if backend == "auto" and _varlen_attn is None:
_log_resolved_backend(backend, "sdpa")
return
for block in vision_model.blocks:
previous_attention = block.attn
replacement = WallXVisionAttention(previous_attention.config, backend=backend)
replacement.to(
device=previous_attention.qkv.weight.device,
dtype=previous_attention.qkv.weight.dtype,
)
replacement.load_state_dict(previous_attention.state_dict(), strict=True)
replacement.train(previous_attention.training)
block.attn = replacement
src/lerobot/policies/wall_x/utils.py
+16 -13
View File
@@ -116,6 +116,7 @@ def preprocesser_call(
images: list | Any | None = None,
text: str | list[str] | None = None,
videos: list | Any | None = None,
device: torch.device | str | None = None,
padding: bool | str = False,
truncation: bool | None = None,
max_length: int | None = None,
@@ -134,6 +135,7 @@ def preprocesser_call(
images: Input images (PIL, numpy arrays, or torch tensors)
text: Text or list of texts to tokenize
videos: Input videos (numpy arrays or torch tensors)
device: Device on which image/video preprocessing should run
padding: Whether to pad sequences to same length
truncation: Whether to truncate sequences longer than max_length
max_length: Maximum length for truncation/padding
@@ -151,7 +153,11 @@ def preprocesser_call(
"""
# Process image inputs
if images is not None and len(images) > 0:
image_inputs = processor.image_processor(images=images, return_tensors=return_tensors)
image_inputs = processor.image_processor(
images=images,
return_tensors=return_tensors,
device=device,
)
image_grid_thw = image_inputs["image_grid_thw"]
else:
image_inputs = {}
@@ -159,7 +165,11 @@ def preprocesser_call(
# Process video inputs
if videos is not None:
videos_inputs = processor.image_processor(videos=videos, return_tensors=return_tensors)
videos_inputs = processor.image_processor(
videos=videos,
return_tensors=return_tensors,
device=device,
)
video_grid_thw = videos_inputs["video_grid_thw"]
else:
videos_inputs = {}
@@ -413,10 +423,7 @@ def get_task_instruction(
}
)
if priority_order is not None:
priority_order = OrderedDict(priority_order)
else:
priority_order = default_priority_order
priority_order = OrderedDict(priority_order) if priority_order is not None else default_priority_order
got_instruction = False
task_instruction = ""
@@ -424,9 +431,8 @@ def get_task_instruction(
# Sample instruction components based on priority probabilities
for key, prob in priority_order.items():
if key in frame_instruction_info and frame_instruction_info[key] != "":
if got_instruction:
if random.random() >= prob:
continue
if got_instruction and random.random() >= prob:
continue
task_instruction += f"\n{frame_instruction_info[key]}"
got_instruction = True
@@ -538,10 +544,7 @@ def img_key_mapping(img_keys: list[str]) -> list[str]:
if key in CAMERA_NAME_MAPPING:
key = CAMERA_NAME_MAPPING[key]
else:
if "view" in key:
key = key.replace("_", " ")
else:
key = key + " view"
key = key.replace("_", " ") if "view" in key else key + " view"
processed_img_keys.append(key)
return processed_img_keys