huangyanpei/VLA
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

add mot (#83)

Browse Source 
* add mot

* update libero example

* translate zh to en

* fix load model from hf

* lint

* lint

---------

Co-authored-by: yangping <yangping@x2robot.com>
This commit is contained in:
suolyer
2026-02-03 11:35:25 +08:00
committed by GitHub
co-authored by yangping
parent 05b6d8dcf7
commit d18fa65fa1
26 changed files with 8509 additions and 1179 deletions
Download Patch File Download Diff File Expand all files Collapse all files
wall_x/model/qwen2_5_based/configuration_qwen2_5_vl.py
+42 -5
View File
@@ -21,6 +21,9 @@ class Qwen2_5_VLVisionConfig(PretrainedConfig):
window_size=112,
out_hidden_size=3584,
fullatt_block_indexes=[7, 15, 23, 31],
initializer_range=0.02,
_attn_implementation="flash_attention_2",
attn_deterministic=False,
**kwargs,
):
super().__init__(**kwargs)
@@ -38,6 +41,9 @@ class Qwen2_5_VLVisionConfig(PretrainedConfig):
self.window_size = window_size
self.fullatt_block_indexes = fullatt_block_indexes
self.out_hidden_size = out_hidden_size
self.initializer_range = initializer_range
self._attn_implementation = _attn_implementation
self.attn_deterministic = attn_deterministic
class Qwen2_5_VLConfig(PretrainedConfig):
@@ -169,6 +175,8 @@ class Qwen2_5_VLConfig(PretrainedConfig):
self,
vocab_size=152064,
hidden_size=8192,
action_hidden_size=2048,
state_hidden_size=2048,
intermediate_size=29568,
num_hidden_layers=80,
num_attention_heads=64,
@@ -193,16 +201,25 @@ class Qwen2_5_VLConfig(PretrainedConfig):
dim_inputs=(1536, 1536),
attention_moe=False,
mlp_moe=False,
norm_moe=False,
mot_opt=False,
flow_loss_weight=10,
use_state_string_representation=False,
use_adarms=False,
proj_with_mask=True,
use_flow_action_expert=True,
adarms_cond_dim=None,
action_horizon_flow=32,
causal_action_attention_mask=False,
use_x_pred=False,
attn_deterministic=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"]()
self.vocab_size = vocab_size
self.max_position_embeddings = max_position_embeddings
self.hidden_size = hidden_size
self.action_hidden_size = action_hidden_size
self.state_hidden_size = state_hidden_size
self.intermediate_size = intermediate_size
self.num_hidden_layers = num_hidden_layers
self.num_attention_heads = num_attention_heads
@@ -230,6 +247,19 @@ class Qwen2_5_VLConfig(PretrainedConfig):
self.dim_inputs = tuple(dim_inputs)
self.attention_moe = attention_moe
self.mlp_moe = mlp_moe
self.norm_moe = norm_moe
self.mot_opt = mot_opt
self.flow_loss_weight = flow_loss_weight
self.use_state_string_representation = use_state_string_representation
self.use_adarms = use_adarms
self.adarms_cond_dim = adarms_cond_dim
self.proj_with_mask = proj_with_mask
self.use_flow_action_expert = use_flow_action_expert
self.action_horizon_flow = action_horizon_flow
self.causal_action_attention_mask = causal_action_attention_mask
self.use_x_pred = use_x_pred
self.attn_deterministic = attn_deterministic
# Validate the correctness of rotary position embeddings parameters
# BC: if there is a 'type' field, move it to 'rope_type'.
@@ -244,5 +274,12 @@ class Qwen2_5_VLConfig(PretrainedConfig):
super().__init__(tie_word_embeddings=tie_word_embeddings, **kwargs)
# move vision config initialization after super init to avoid recursively set in latest transformers version
# TODO: make it better
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"]()
__all__ = ["Qwen2_5_VLConfig"]
wall_x/model/qwen2_5_based/modeling_qwen2_5_vl.py
+97 -15
View File
@@ -120,23 +120,68 @@ class Qwen2_5_VisionRotaryEmbedding(nn.Module):
class Qwen2RMSNorm(nn.Module):
def __init__(self, hidden_size, eps=1e-6):
def __init__(
self, hidden_size: int, eps: float = 1e-6, cond_dim: Optional[int] = None
):
"""
Qwen2RMSNorm is equivalent to T5LayerNorm
Qwen2RMSNorm with optional conditional input support, equivalent to T5LayerNorm
"""
super().__init__()
self.weight = nn.Parameter(torch.ones(hidden_size))
self.variance_epsilon = eps
self.hidden_size = hidden_size
self.cond_dim = cond_dim
def forward(self, hidden_states):
# Dense layer for adaptive normalization (if cond_dim is provided)
if cond_dim is not None:
self.dense = nn.Linear(cond_dim, hidden_size * 3, bias=True)
nn.init.zeros_(self.dense.weight)
else:
self.dense = None
self.weight = nn.Parameter(torch.ones(hidden_size))
def _norm(self, x):
# Compute variance in float32 for numerical stability
variance = x.pow(2).mean(-1, keepdim=True)
# Compute normalization
normed_inputs = x * torch.rsqrt(variance + self.variance_epsilon)
return normed_inputs
def forward(self, hidden_states, cond=None):
input_dtype = hidden_states.dtype
hidden_states = hidden_states.to(torch.float32)
variance = hidden_states.pow(2).mean(-1, keepdim=True)
hidden_states = hidden_states * torch.rsqrt(variance + self.variance_epsilon)
return self.weight * hidden_states.to(input_dtype)
normed_inputs = self._norm(hidden_states)
if cond is None or self.dense is None:
# Regular RMSNorm
normed_inputs = self.weight * normed_inputs
return normed_inputs.to(input_dtype), None
# Adaptive RMSNorm
if cond.shape[-1] != self.cond_dim:
raise ValueError(
f"Expected cond dimension {self.cond_dim}, got {cond.shape[-1]}"
)
# Compute modulation parameters
cond = cond.to(dtype=self.dense.weight.dtype)
modulation = self.dense(cond)
if len(hidden_states.shape) == 3: # [batch, seq, features]
modulation = modulation.unsqueeze(1)
scale, shift, gate = torch.chunk(modulation, 3, dim=-1)
# Apply adaptive normalization
normed_inputs = normed_inputs * (1 + scale.to(torch.float32)) + shift.to(
torch.float32
)
return normed_inputs.to(input_dtype), gate.to(input_dtype)
def extra_repr(self):
return f"{tuple(self.weight.shape)}, eps={self.variance_epsilon}"
repr_str = f"{tuple(self.weight.shape)}, eps={self.variance_epsilon}"
if self.dense is not None:
repr_str += f", adaptive=True, cond_dim={self.cond_dim}"
return repr_str
class Qwen2_5_VLPatchMerger(nn.Module):
@@ -151,7 +196,7 @@ class Qwen2_5_VLPatchMerger(nn.Module):
)
def forward(self, x: torch.Tensor) -> torch.Tensor:
x = self.mlp(self.ln_q(x).view(-1, self.hidden_size))
x = self.mlp(self.ln_q(x)[0].view(-1, self.hidden_size))
return x
@@ -381,13 +426,13 @@ class Qwen2_5_VLVisionBlock(nn.Module):
position_embeddings: Optional[Tuple[torch.Tensor, torch.Tensor]] = None,
) -> torch.Tensor:
hidden_states = hidden_states + self.attn(
self.norm1(hidden_states),
self.norm1(hidden_states)[0],
cu_seqlens=cu_seqlens,
max_seqlen=max_seqlen,
rotary_pos_emb=rotary_pos_emb,
position_embeddings=position_embeddings,
)
hidden_states = hidden_states + self.mlp(self.norm2(hidden_states))
hidden_states = hidden_states + self.mlp(self.norm2(hidden_states)[0])
return hidden_states
@@ -1082,16 +1127,41 @@ class Qwen2_5_VLSdpaAttention(Qwen2_5_VLAttention):
"cos": cos,
"cache_position": cache_position,
} # Specific to RoPE models
key_states, value_states = past_key_value.update(
key_states, value_states, self.layer_idx, cache_kwargs
)
if use_cache:
key_states, value_states = past_key_value.update(
key_states, value_states, self.layer_idx, cache_kwargs
)
else:
past_key_states, past_value_states = past_key_value[self.layer_idx]
key_states = torch.cat([past_key_states, key_states], dim=-2)
value_states = torch.cat([past_value_states, value_states], dim=-2)
key_states = repeat_kv(key_states, self.num_key_value_groups)
value_states = repeat_kv(value_states, self.num_key_value_groups)
causal_mask = attention_mask
if attention_mask is not None: # no matter the length, we just slice it
causal_mask = attention_mask[:, :, :, : key_states.shape[-2]]
# Ensure the attention_mask correctly matches the head dimension
if len(attention_mask.shape) == 2: # [batch_size, seq_len]
# Expand to [batch_size, 1, seq_len, seq_len] causal mask format
bsz, seq_len = attention_mask.shape
causal_mask = attention_mask.view(bsz, 1, 1, seq_len).expand(
bsz, 1, seq_len, seq_len
)
elif len(attention_mask.shape) == 3: # [batch_size, seq_len, seq_len]
# Add head dimension: [batch_size, 1, seq_len, seq_len]
causal_mask = attention_mask.unsqueeze(1)
elif (
len(attention_mask.shape) == 4
): # [batch_size, num_heads, seq_len, seq_len]
causal_mask = attention_mask
else:
raise ValueError(
f"Unsupported attention_mask dim: {attention_mask.shape}"
)
# Convert the attention mask to boolean type
causal_mask = causal_mask.to(torch.bool)
# SDPA with memory-efficient backend is currently (torch==2.1.2) bugged with non-contiguous inputs with custom attn_mask,
# Reference: https://github.com/pytorch/pytorch/issues/112577.
@@ -1105,6 +1175,18 @@ class Qwen2_5_VLSdpaAttention(Qwen2_5_VLAttention):
# The q_len > 1 is necessary to match with AttentionMaskConverter.to_causal_4d that does not create a causal mask in case q_len == 1.
is_causal = True if causal_mask is None and q_len > 1 else False
if q_len == 1:
is_causal = False
causal_mask = torch.ones(
bsz,
1,
1,
key_states.shape[2],
device=hidden_states.device,
dtype=hidden_states.dtype,
).contiguous()
causal_mask = causal_mask.to(torch.bool)
attn_output = torch.nn.functional.scaled_dot_product_attention(
query_states,
key_states,
wall_x/model/qwen2_5_based/modeling_qwen2_5_vl_act.py
+898 -585
View File
File diff suppressed because it is too large Load Diff