L1–146VLA 模型的 MoE(Mixture of Experts)架构核心模块：包含非学习路由器 TokenTypeRouter、单层 MLP 专家 BlockSparseMLP、以及动态专家编排 SparseMoeBlock。支持梯度检查点和 FSDP 分布式训练，实现图文与动作 token 的异构 MoE 处理。
导入 PyTorch 核心库：torch 主包、nn 模块、checkpoint 梯度重算工具(激活值重算降低显存峰值)。
5
from torch.distributed.fsdp import MixedPrecision as MP
6
from torch.distributed.fsdp import FullyShardedDataParallel as FSDP
导入分布式训练工具：FSDP(FullyShardedDataParallel)与 MP(MixedPrecision)，支持多 GPU/TPU 上的大模型训练。
10
from peft import LoraConfig, get_peft_model
11
from typing import Optional, Union, Dict
12
from packaging import version
导入 LoRA 微调配置与工具、类型提示、版本比较库；后续可能用于参数高效微调。
14
from transformers import GenerationMixin
15
from transformers.activations import ACT2FN
16
from transformers.modeling_utils import AttentionInterface
18
from transformers.utils import logging, is_torch_xla_available
导入 Transformers 库的生成 mixin、激活函数映射表(ACT2FN)、注意力接口、日志工具和 XLA 检测函数。
20
from wall_x.model.action_head import ActionProcessor
21
from wall_x.model.model_utils import find_first_last_ones
导入模型模块：ActionProcessor(处理流匹配动作生成与归一化)与辅助函数 find_first_last_ones。
23
ALL_ATTENTION_FUNCTIONS: AttentionInterface = AttentionInterface()
24
logger = logging.get_logger(__name__)
初始化全局注意力函数表与模块日志器；注意力表用于模型中注意力实现的多态查表。
27
X2ROBOT_ATTENTION_FUNCTIONS = []
28
ATTENTION_TYPES_WITH_2D_MASK = [
29
"sdpa",
30
]
31
ATTENTION_TYPES_WITH_FLASH_MASK = []
定义注意力实现列表与支持的 mask 类型常量：X2ROBOT_ATTENTION_FUNCTIONS 为空列表(可扩展)，ATTENTION_TYPES_WITH_2D_MASK=['sdpa']支持 2D mask(用于图文与动作 token 混合注意力)，ATTENTION_TYPES_WITH_FLASH_MASK 为空。
34
class TokenTypeRouter(nn.Module):
35
def __init__(self, num_experts: int):
36
super().__init__()
37
self.num_experts = num_expertsTokenTypeRouter 类定义开始：非学习的固定路由器，按 token 类型分配至不同专家。__init__ 接收 num_experts(=2，对应图文与动作)，存储为实例变量。
39
def forward(self, token_types: torch.Tensor) -> torch.Tensor:
40
"""41
Assigns tokens to different experts based on `token_type`.42
Args:43
token_types (torch.Tensor): A tensor of shape (batch_size, seq_length) representing the type of each token.45
Returns:46
experts_indices (torch.Tensor): A tensor of shape (batch_size, seq_length) representing the expert index assigned to each token.47
"""48
experts_indices = token_types % self.num_experts49
return experts_indicesforward 方法与类结尾：输入 token_types(B,S)，返回 experts_indices(B,S)。docstring 说明语义(基于 token_type 分配专家)，实现用简单取模 token_types % num_experts(0→expert0图文，1→expert1动作)，无学习参数。LIBERO 场景中这是固定的类型路由。行 51 空行分隔类。
52
class BlockSparseMLP(nn.Module):
53
def __init__(self, config, use_selective_recompute: bool = False):
54
super().__init__()
55
self.hidden_size = config["hidden_size"]
56
self.intermediate_size = config["intermediate_size"]
57
self.hidden_act = config["hidden_act"]
59
self.use_selective_recompute = use_selective_recomputeBlockSparseMLP 类初始化：单个 MLP 专家模块。从 config dict 读取 hidden_size(标准 2048)、intermediate_size(图文 11008/动作 2048)、hidden_act(激活函数名)；use_selective_recompute 控制是否启用梯度检查点。
61
self.gate_proj = nn.Linear(self.hidden_size, self.intermediate_size, bias=False)
62
self.up_proj = nn.Linear(self.hidden_size, self.intermediate_size, bias=False)
63
self.down_proj = nn.Linear(self.intermediate_size, self.hidden_size, bias=False)
65
self.act_fn = ACT2FN[self.hidden_act]
定义三个线性层与激活函数：gate_proj 和 up_proj 用于门控 MLP(类似 GLU 结构)，down_proj 投回隐维；act_fn 从字符串映射至实际激活函数对象(如'gelu'→F.gelu)。shape 流向：(B*N,H)→(B*N,I)两条路→逐元相乘→(B*N,H)。行 66 空行分隔方法。
67
def _full_mlp(self, hidden_state):
68
gate_out = self.gate_proj(hidden_state)69
up_out = self.up_proj(hidden_state)70
act_out = self.act_fn(gate_out) * up_out71
return self.down_proj(act_out)
_full_mlp 内部方法：完整的门控 MLP 计算。gate_out 和 up_out 分别通过两条线性路径 (B*N,H)→(B*N,I)，激活后逐元相乘(Gated Linear Unit)，再通过 down_proj 回到 (B*N,H)。
73
def forward(self, hidden_state):
74
if self.use_selective_recompute:
75
# Perform checkpoint recalculation for the entire expert MLP.76
return cp.checkpoint(77
self._full_mlp,78
hidden_state,
79
use_reentrant=False,80
)
81
else:82
return self._full_mlp(hidden_state)
forward 方法与类结尾：根据 use_selective_recompute 开关决定是否用 gradient checkpoint 包装 _full_mlp。checkpoint(..., use_reentrant=False) 在前向时保存输入，反向时重算激活值但不重算权重 GEMM，适合 FSDP。LIBERO 微调时该开关控制显存优化策略。行 84 空行分隔类。
85
class SparseMoeBlock(nn.Module):
86
def __init__(self, config, num_experts: int, use_selective_recompute: bool = False):
87
super().__init__()
88
self.num_experts = num_experts89
self.use_selective_recompute = use_selective_recompute91
# Pass the `use_selective_recompute` parameter to each expert.SparseMoeBlock 类定义：多专家动态编排模块。__init__ 接收 config、num_experts(=2)、use_selective_recompute；存储这些参数后声明要创建专家列表和路由逻辑。
92
self.experts = nn.ModuleList(93
[
94
BlockSparseMLP(
95
config.experts[i], use_selective_recompute=use_selective_recompute
96
)
97
for i in range(num_experts)
98
]
99
)
创建专家列表 nn.ModuleList：包含 num_experts 个 BlockSparseMLP 实例，每个从 config.experts[i] 独立配置(包含 hidden_size、intermediate_size、hidden_act)，并透传 use_selective_recompute 开关。图文专家中间维 11008(继承 Qwen)，动作专家 2048(共 0.45B 参数)。行 100 空行分隔逻辑。
101
if not hasattr(config, "dim_inputs") or not config.dim_inputs:
102
raise ValueError("Configuration must contain a valid dim_inputs")
104
self.dim_inputs = config.dim_inputs105
self.permuted = config.mot_opt校验并存储 dim_inputs：每个专家的有效输入维度列表(可能 <hidden_size，用于不对称 MoE)。若 config 缺 dim_inputs 则抛 ValueError；else 存入 self.dim_inputs。同时从 config.mot_opt 读取 permuted 标志，控制是否使用预排列模式。
107
def forward(
108
self,109
hidden_states: torch.Tensor,
110
experts_indices: torch.Tensor,
111
start_indices: torch.Tensor,
112
end_indices: torch.Tensor,
113
) -> torch.Tensor:
115
if self.permuted:
forward 方法签名：输入 hidden_states(B,S,H)为隐状态、experts_indices(B,S)为路由结果(来自 TokenTypeRouter)、start_indices/end_indices(num_experts,)为各专家在排列后张量中的起止位置。返回处理后的 hidden_states(B,S,H)。
116
permuted_inputs = hidden_states
117
else:118
batch_size, seq_length, hidden_dim = hidden_states.shape
120
flat_hidden = hidden_states.reshape(-1, hidden_dim)121
experts_indices = experts_indices.reshape(-1)122
probs = torch.ones_like(experts_indices, dtype=torch.float32).reshape(-1, 1)
123
permuted_inputs, row_id_map = ops.permute(flat_hidden, experts_indices)
125
# buffer条件分支处理排列模式：若 self.permuted=true(预排列)则跳过 permute 步，直接用原始 hidden_states；否则将 (B,S,H)→(B*S,H) flatten，专家索引也 reshape(-1)，调用 ops.permute(flat_hidden, experts_indices) 将相同专家的 token 分组连续存储，返回排列后张量(B*S,H)与行映射 row_id_map。同时构造权重张量 probs 用于后续 unpermute。
126
final_output = torch.zeros_like(permuted_inputs)
128
# Expert forward contain selective recompute初始化输出缓冲区 final_output 为零张量，shape 与 permuted_inputs 相同，用于逐专家填充处理结果。
129
for expert_idx, expert in enumerate(self.experts):
130
start, end = start_indices[expert_idx], end_indices[expert_idx]
131
if start == end:132
continue134
dim_input = self.dim_inputs[expert_idx]135
expert_input = permuted_inputs[start:end, :dim_input]
137
partial_output = expert(expert_input)
138
final_output[start:end, :dim_input] = partial_output[:, :dim_input]
140
if self.permuted:
循环遍历每个专家进行前向传播：获取当前专家分配的 token 范围 (start,end)，若为空(start==end)则跳过；否则从 permuted_inputs[start:end,:dim_input[expert_idx]] 取出该专家的 input(可能 <H 维)，调用专家(含梯度检查点)得 partial_output，写回 final_output 同位置。LIBERO 中各专家的 dim_input 可能不同(图文/动作)，在此处实现不对称路由。
141
return final_output142
else:143
final_output = ops.unpermute(final_output, row_id_map, probs)
144
return final_output.reshape(batch_size, seq_length, hidden_dim)条件分支处理输出：若 permuted=true 直接返回 final_output；否则通过 ops.unpermute(final_output, row_id_map, probs) 逆向排列恢复原始 token 顺序，再 reshape(B,S,H) 回到原始形状。LIBERO 实战中这步确保多 token 的动作输出与输入序列的对齐关系不被破坏。行 145-146 为类结尾空行。
L147–178ActionModelMixMin Mixin 类的定义和四个核心方法：(1) 类定义和属性注解；(2) 初始化并注册 ActionProcessor/TokenTypeRouter/SparseMoeBlock 三组件；(3) set_normalizer 方法代理动作和本体状态归一化器；(4) _apply_mlp_moe 根据 mlp_moe 标志选择 MoE 或普通 MLP 路由。
147
class ActionModelMixMin:
148
# config: Qwen2_5_VLConfig149
action_preprocessor: ActionProcessor
150
router: TokenTypeRouter
151
moe: SparseMoeBlock
ActionModelMixMin 是 VLA 模型中 MoE 路由和动作生成的核心 Mixin 类。class 定义无继承（因为会被动态混入到具体的 VLM 前向流程中）。三个类属性用 Python 3.6+ 类型注解声明：action_preprocessor（ActionProcessor，处理动作/状态），router（TokenTypeRouter，根据 token type 分配专家），moe（SparseMoeBlock，双专家 MoE）。注释说明 config 参数类型为 Qwen2_5_VLConfig。这些属性在 __init__ 中真正赋值。
153
def __init__(self, config, action_preprocessor, router, moe):
154
self.config = config155
self.action_preprocessor = action_preprocessor156
self.router = router157
self.moe = moe158
self._mot_opt_warned = False
152 行空行（并入相邻组）。153-158 是 __init__ 初始化方法，接收 config（Qwen 配置）和三个核心组件，分别赋值到实例属性（self.config/action_preprocessor/router/moe）。_mot_opt_warned 初始化为 False，后续用于防止 mot_opt 配置变更的重复警告。在 LIBERO 实战中，router 按 token_type 分配图文（type=0，走专家0）和动作/状态（type=1，走专家1）；MoE 内专家0 中间层 11008（继承 Qwen），专家1 中间层 2048；ActionProcessor 处理 7 维 dof（前 7 个 dof_mask=1，其余 pad）。
160
def set_normalizer(self, normalizer_action, normalizer_propri):
161
if hasattr(self, "action_preprocessor"):
162
self.action_preprocessor.set_normalizer(163
normalizer_action, normalizer_propri
164
)
165
else:166
logger.warning(
167
"ActionModelMixMin.set_normalizer is called but action_preprocessor is not set"
168
)
159 行空行（并入相邻组）。160-168 是 set_normalizer(normalizer_action, normalizer_propri) 方法，用 hasattr 防御式检查 action_preprocessor 是否已装配（Mixin 可能在某些分支被动态组装或跳过），若存在则代理调用 action_preprocessor.set_normalizer 注册两个归一化器。normalizer_action 是动作空间统计值（min/delta），normalizer_propri 是本体状态统计值。若不存在则记录 warning。在 LIBERO 中通过此方法传入 libero_8/libero_goal 等数据集特定的统计值，作用于 ActionProcessor.proprioception_proj 和 flow matching 的 action 投影反归一化。
170
def _apply_mlp_moe(self, hidden_states, token_types, start_indices, end_indices):
171
if self.config.mlp_moe:
172
hidden_states = self.moe(173
hidden_states, token_types, start_indices, end_indices
174
)
175
else:176
hidden_states = self.mlp(hidden_states)177
return hidden_states169 行空行（并入相邻组）。170-177 是 _apply_mlp_moe(hidden_states, token_types, start_indices, end_indices) 方法，根据 config.mlp_moe 标志选择路由策略：若 mlp_moe=True 则调用 self.moe（SparseMoeBlock.forward），内部按 token_types 动态排列→专家处理→恢复，tensor shape (B,S,H) 保持不变；否则调用 self.mlp（全共享 FFN，对比配置）。返回处理后的 hidden_states。start_indices/end_indices 标记各专家负责的 token 范围。178 行空行标记方法结束，后续接 _apply_norm_moe。在 LIBERO 部署中 mlp_moe=True，专家1 处理动作 token 需配合 dof_mask 隐藏无关 dof。
L179–355_apply_norm_moe 方法实现 MoE 感知的 LayerNorm，支持两种路由模式(expert-wise 段式 vs token-level 掩码) 和可选的选择性激活重计算。通过 norm_moe 配置选择是否 expert-wise norm 或共享 norm；通过 mot_opt 配置选择排列优化或掩码方案。特殊处理 adarms 条件下的动作专家，支持流匹配的时间步展开。最终返回归一化后的 hidden_states 及其关联的 gate/gate_mask，用于下游 gated residual 调制。LIBERO 微调典型配置为 norm_moe=true + mot_opt=true + use_adarms=false，此时 gate/gate_mask 多为 None。
179
def _apply_norm_moe(
180
self,181
hidden_states,
182
token_types,
183
adarms_conds,
184
norms, # list of norm layers (expert-wise)185
norm, # shared norm if not norm_moe186
start_indices=None,187
end_indices=None,188
use_selective_recompute=False,189
):
函数签名：_apply_norm_moe 是 ActionModelMixMin 类的 MoE 感知归一化方法。核心参数：hidden_states(B*S,D 或 B,S,D) 隐状态、token_types(B*S)token 类型掩码(0=图文，1=动作)、adarms_conds list 条件、norms list expert-wise norm 层、norm shared norm 层、start_indices/end_indices 路由段界、use_selective_recompute 控制是否激活重计算。被 modeling_qwen2_5_vl_act.py 多处调用。
190
"""191
MoE-aware LayerNorm with optional selective activation recomputation.193
Only activation math is recomputed. No GEMM is recomputed.194
Safe for FSDP (use_reentrant=False).195
"""Docstring：MoE 感知的 LayerNorm 支持可选的选择性激活重计算。说明重计算仅涉及激活函数，不重计算 GEMM。强调与 FSDP(use_reentrant=False) 兼容，这对分布式训练显存优化至关重要。
初始化 gate 和 gate_mask 为 None。gate 用于存储来自 expert norm 的条件门控张量(通常 adarms 专家返回)，gate_mask 记录该 gate 作用的 token 位置。两者最后随 hidden_states 返回给调用者。
200
# -------------------------201
# Case 1: norm_moe=True (expert-wise norm)202
# -------------------------203
if self.config.norm_moe:
大分支注释和 if self.config.norm_moe 条件。norm_moe=true 时每个专家有独立 LayerNorm；否则全局共享一个 norm。这决定了下游是否分别初始化各专家的归一化或应用统一 norm。
205
# ---------------------------------------------------------206
# Case 1A: mot_opt=True (segments assigned by start/end)207
# ---------------------------------------------------------208
if self.config.mot_opt:
209
new_hidden_states = torch.zeros_like(hidden_states)
Case 1A：mot_opt(motion optimization)=true 分支。按 start_indices/end_indices 段式分配(已排列优化)。初始化 new_hidden_states=zeros_like(hidden_states)，后续 for 循环逐专家填充该张量，提高缓存局部性。此模式要求输入 hidden_states 已被 permute 过。
211
for expert_idx, expert_norm in enumerate(norms):
212
start = start_indices[expert_idx]
213
end = end_indices[expert_idx]
214
if start == end:215
continuefor 循环枚举专家和其 norm 层。取 start/end 索引；若 start==end 该专家无 token 分配，continue 跳过。dim_input 从 config 取出该专家的输入维(e.g. 图文专家 2048，动作专家 2048)。selected = hidden_states[start:end] 切出该专家的 token 段，shape: [K, D] 其中 K 是该段 token 数。
217
dim_input = self.config.dim_inputs[expert_idx]218
selected = hidden_states[start:end] # [K, D]220
# ====== reshape if adarms on flow expert ======221
if self.config.use_adarms and expert_idx == 1:
222
selected = selected.view(
223
-1,224
self.config.action_horizon_flow,225
selected.shape[-1],226
)
227
input_slice = selected[:, :, :dim_input]
228
cond = adarms_conds[expert_idx]
229
else:230
input_slice = selected[:, :dim_input]
231
cond = adarms_conds[expert_idx]
处理 use_adarms and expert_idx==1 的动作专家。若条件成立，将 selected reshape 为 3D (batch_token, action_horizon_flow, D) 以保留流匹配的时间结构。input_slice 提取前 dim_input 维，cond 取对应专家的条件张量。否则直接用 2D 切片。这是流匹配框架在离散化后的时间步展开。
233
if use_selective_recompute:235
def norm_chunk(t_x, t_cond, expert_norm=expert_norm):
236
if t_cond is None or (
237
isinstance(t_cond, torch.Tensor) and t_cond.numel() == 0
238
):
239
out, _ = expert_norm(t_x)
240
else:241
out, _ = expert_norm(t_x, t_cond)
242
return out244
cond_for_cp = (
245
cond
246
if cond is not None
247
else torch.empty(0, device=input_slice.device)
248
)
249
processed = cp.checkpoint(
250
norm_chunk,
251
input_slice,
252
cond_for_cp,
253
use_reentrant=False,254
)
use_selective_recompute=true 时启用激活重计算。定义 norm_chunk 闭包将 input_slice 和条件 t_cond 传入 expert_norm；若条件为 None 或空张量则无条件调用，否则有条件调用(adarms case)。cond_for_cp 兼容处理空条件(创建同设备空张量)。cp.checkpoint 带 use_reentrant=false 启用激活重计算，显存友好且 FSDP 安全。
else 分支：不启用重计算。直接调用 expert_norm(input_slice, cond)，同时获得 processed 张量和可能的 gate 返回值(某些 norm 层支持 gate 输出)。
258
# reshape back if needed259
if self.config.use_adarms and expert_idx == 1:
260
processed = processed.view(-1, dim_input)262
new_hidden_states[start:end, :dim_input] = processed.to(
263
hidden_states.dtype
264
)
若 use_adarms and expert_idx==1，processed 从 3D (batch_token, action_horizon_flow, dim_input) reshape 回 2D (-1, dim_input)。赋值 new_hidden_states[start:end, :dim_input] = processed.to(原 dtype)，确保混精度兼容。处理后的张量形状：(K, dim_input)，其中仅前 dim_input 列被更新，其余 D-dim_input 列保持零。
Case 1A 末尾：hidden_states = new_hidden_states，用逐专家归一化后的张量替换输入，完成 mot_opt=true 路径。
268
# ---------------------------------------------------------269
# Case 1B: mot_opt=False (token-level mask)270
# ---------------------------------------------------------271
else:273
new_hidden_states = torch.zeros_like(hidden_states)
274
B, S, D = hidden_states.shape
Case 1B (else 块)：mot_opt=false 时的 token-level 掩码方案。初始化 new_hidden_states，解包 shape 为 (B, S, D)。此方案适合 token 跨专家分布不均的场景，每个位置 (b, s) 独立判断其 token 类型。
276
for expert_idx, expert_norm in enumerate(norms):
277
mask = token_types == expert_idx
278
if mask.sum() == 0:
279
continuefor 循环枚举专家。mask = token_types == expert_idx 用布尔掩码标记属于该专家的 token(在二维 (B,S) 上broadcast)。若该专家无分配 token(mask.sum()==0)，continue 跳过以避免空操作。
281
dim_input = self.config.dim_inputs[expert_idx]282
selected = hidden_states[mask] # [K, D]284
if self.config.use_adarms and expert_idx == 1:
285
gate_mask = mask
286
selected = selected.view(
287
-1,288
self.config.action_horizon_flow,289
selected.shape[-1],290
)
291
input_slice = selected[:, :, :dim_input]
292
cond = adarms_conds[expert_idx]
293
else:294
input_slice = selected[:, :dim_input]
295
cond = adarms_conds[expert_idx]
selected = hidden_states[mask] 用布尔掩码提取 token，shape: [K, D]。若 use_adarms and expert_idx==1：记录 gate_mask=mask(后续可用于 gated residual)，reshape 为 3D (batch_token, action_horizon_flow, D)，提取 input_slice 前 dim_input 维。adarms_conds[expert_idx] 取条件张量。其他专家直接 2D 切片，cond 同样取。
297
if use_selective_recompute:299
def norm_chunk(t_x, t_cond, expert_norm=expert_norm):
300
if t_cond is None or (
301
isinstance(t_cond, torch.Tensor) and t_cond.numel() == 0
302
):
303
out, _ = expert_norm(t_x)
304
else:305
out, _ = expert_norm(t_x, t_cond)
306
return out308
cond_for_cp = (
309
cond
310
if cond is not None
311
else torch.empty(0, device=input_slice.device)
312
)
314
processed = cp.checkpoint(
315
norm_chunk,
316
input_slice,
317
cond_for_cp,
318
use_reentrant=False,319
)
use_selective_recompute=true 时的重计算分支(逻辑同 Case 1A)。定义 norm_chunk 处理条件判断。cond_for_cp 兼容空条件。cp.checkpoint 启用激活重计算。此处与 Case 1A 的区别仅在于输入提取方式(掩码 vs 切片)，重计算策略相同。
320
else:321
processed, gate = expert_norm(input_slice, cond)
323
if self.config.use_adarms and expert_idx == 1:
324
processed = processed.view(-1, dim_input)326
# scatter back327
b_id, s_id = torch.where(mask)
328
new_hidden_states[b_id, s_id, :dim_input] = processed.to(
329
hidden_states.dtype
330
)
非重计算分支：直接 processed, gate = expert_norm(input_slice, cond)。若 use_adarms and expert_idx==1，reshape processed 回 2D。scatter 回原位：b_id, s_id = torch.where(mask) 得到原始 (batch_idx, seq_idx) 坐标，new_hidden_states[b_id, s_id, :dim_input] = processed.to(dtype)。形状变化：(B,S,D)-[mask]->(K,D)->norm->(K,dim_input)->scatter->(B,S,dim_input)。
Case 1B 末尾：hidden_states = new_hidden_states，用掩码分配后的张量替换，完成 mot_opt=false 路径。
334
# -------------------------335
# Case 2: norm_moe=False (single LN)336
# -------------------------337
else:Case 2 (else)：norm_moe=false 时的单 norm 分支。使用全局共享的 norm 层替代 expert-wise 分支。定义 norm_chunk_shared(t_x, dummy, norm_module) 闭包，dummy 参数保持与 Case 1 checkpoint 签名兼容(两位置参数)。
339
def norm_chunk_shared(t_x, dummy, norm_module=norm):
340
out, _ = norm_module(t_x)
341
return out343
if use_selective_recompute:344
dummy = torch.empty(0, device=hidden_states.device)345
hidden_states = cp.checkpoint(
346
norm_chunk_shared,
347
hidden_states,
348
dummy,
349
use_reentrant=False,350
)
351
else:352
hidden_states, gate = norm(hidden_states)
354
return hidden_states, gate, gate_maskuse_selective_recompute=true：创建空 dummy 张量，cp.checkpoint 调用 norm_chunk_shared 启用激活重计算。若 false，直接调用 norm(hidden_states) 返回 hidden_states 和 gate。gate 通常为 None(标准 LN 无 gate 输出)，但代码框架预留此返回以支持条件 norm。整个 _apply_norm_moe 统一的返回逻辑体现在此。
return hidden_states, gate, gate_mask：三元组返回。hidden_states 是经过相应路由和归一化的张量(shape 与输入相同)；gate 来自 adarms 专家的条件信息(用于 _gated_residual 调制)；gate_mask 记录 gate 作用范围。LIBERO 微调中若 use_adarms=false，后两者通常为 None。调用方在 modeling_qwen2_5_vl_act.py 中接收这三个值。
L356–473五个工具方法实现：(1)_gated_residual 对残差连接施加门控调制，用于 flow 专家的输出；(2)scatter_proprioception_embeddings 将本体状态嵌入代替掩码位置；(3)scatter_flow_action_embeddings 将生成的动作嵌入代替掩码位置，同时返回 flow 监督目标和 adarms 条件；(4)_update_position_ids 调整 flow token 的 position_id 与 AR token 部分对齐；(5)_update_joint_attention_mask_2d 将 1D attention_mask 扩展为 2D 因果掩码。
356
def _gated_residual(self, x, y, gate, start_indices=None, end_indices=None):
357
"""358
Applies gated residual connection with optional gate parameter.360
Args:361
x: Input tensor (residual)362
y: Output tensor to be added363
gate: Optional gate tensor to modulate the addition365
Returns:366
x + y if gate is None, otherwise x + y * gate367
"""def _gated_residual(self, x, y, gate, start_indices=None, end_indices=None): 及 docstring 定义门控残差连接方法。被 forward():249,271 调用。当 gate 非空时，对动作专家的输出 y 施加门控调制后加到 x（图文特征）上。
368
if x is None and y is None:
369
return None
370
if x is None or y is None:
371
return x if x is not None else y
372
if gate is None:
373
return x + y四个条件分支处理残差连接简化情形：(1)x,y 都是 None→返回 None；(2)x 或 y 之一是 None→返回非空者；(3)gate 是 None→简单相加 x+y 无门控。避免后续复杂的 reshape 逻辑。第 374 行空行为段落分隔。
375
new_y = y.clone()
376
selected_y = y[start_indices[1] : end_indices[1]]
377
selected_y = selected_y.view(
378
-1, self.config.action_horizon_flow, selected_y.shape[-1]
379
)[:, :, : self.config.dim_inputs[1]]
new_y = y.clone() 克隆 y 以安全修改。selected_y = y[start_indices[1]:end_indices[1]] 取第 1 个专家（动作专家）的张量切片。view(-1, action_horizon_flow, shape[-1])[:, :, :dim_inputs[1]] 重新形状为 (batch*seq, action_horizon_flow, 20) 再取前 dim_inputs[1]=20 维。tensor shape: (K,H)→(K/A, A, H)→(K/A, A, D_action).
380
selected_y = selected_y.to(torch.float32) * gate
381
new_y[start_indices[1] : end_indices[1], : self.config.dim_inputs[1]] = (
382
selected_y.view(-1, self.config.dim_inputs[1]).to(new_y.dtype)
383
)
selected_y 转为 float32 后逐元素乘以 gate（来自 norm 的门控值）。reshape(-1, dim_inputs[1]) 并转回原 dtype，scatter 回 new_y 中 new_y[start_indices[1]:end_indices[1], :dim_inputs[1]] = ...。实战：gate 对动作向量逐 dof 维度加权，不确定维度权重低。第 384 行空行。
387
def scatter_proprioception_embeddings(
388
self, input_ids, inputs_embeds, proprioception, dataset_names, agent_pos_mask389
):
def scatter_proprioception_embeddings(self, input_ids, inputs_embeds, proprioception, dataset_names, agent_pos_mask): 将本体状态嵌入到输入序列中。被 forward():1416 调用，将本体状态特征散布到 propri_token_id 的位置。
390
if (391
proprioception is not None
392
and not self.config.use_state_string_representation
393
):
394
proprioception = proprioception.to(inputs_embeds.device).to(
395
inputs_embeds.dtype
396
)
397
agent_pos_mask = agent_pos_mask.to(inputs_embeds.device).to(
398
inputs_embeds.dtype
399
)
400
proprioception = self.action_preprocessor.proprioception_proj(401
proprioception,
402
dataset_names,
403
agent_pos_mask,
404
use_history=proprioception.shape[1] > 1,
405
)
406
mask = input_ids == self.action_token_id_set["propri_token_id"]
407
mask_unsqueezed = mask.unsqueeze(-1)408
mask_expanded = mask_unsqueezed.expand_as(inputs_embeds)
409
proprioception_mask = mask_expanded.to(inputs_embeds.device)
条件检查：proprioception 非空且未使用状态字符串表示。proprioception 和 agent_pos_mask 转移到 device 和 dtype。调用 action_preprocessor.proprioception_proj(proprioception, dataset_names, agent_pos_mask, use_history=proprioception.shape[1]>1) 进行投影和归一化。实战：LIBERO 中 agent_pos_mask 标记可见关节，proprioception_proj 按 dataset_name 查表 normalize。第 410 行空行。
411
proprioception = proprioception.to(
412
inputs_embeds.device, inputs_embeds.dtype
413
)
414
inputs_embeds = inputs_embeds.masked_scatter(
415
proprioception_mask, proprioception
416
)
mask = input_ids == propri_token_id 找本体状态 token 位置，shape (B,S)。unsqueeze(-1) 变为 (B,S,1)，expand_as 广播到 (B,S,hidden_dim)。proprioception 再转移到设备和 dtype（保险）。inputs_embeds.masked_scatter(proprioception_mask, proprioception) 将值填入被掩码标记位置。第 417 行空行。
420
def scatter_flow_action_embeddings(
421
self, input_ids, inputs_embeds, action_chunk, dataset_names, dof_mask422
):
def scatter_flow_action_embeddings(self, input_ids, inputs_embeds, action_chunk, dataset_names, dof_mask): 将生成的动作嵌入散布到动作 token 位置。被 forward():1420 调用。返回 (inputs_embeds, flow, adarms_cond)，其中 flow 是监督目标，adarms_cond 是条件向量。
423
if not self.config.use_flow_action_expert:
424
return inputs_embeds, None, None
425
adarms_cond, flow = None, None
早期返回：如果 use_flow_action_expert=False（动作使用通常 llm 生成非 flow matching），直接返回原始嵌入和两个 None。adarms_cond, flow = None, None 预初始化。
426
if action_chunk is not None:
427
action_chunk = action_chunk.to(inputs_embeds.device)
428
dof_mask = dof_mask.to(inputs_embeds.device)
429
noisy_action_emb, flow, adarms_cond = self.action_preprocessor(430
action_chunk, dataset_names, dof_mask
431
)
if action_chunk is not None: 有真实动作数据。action_chunk 和 dof_mask 转移到 device。调用 action_preprocessor(action_chunk, dataset_names, dof_mask) 返回 (noisy_action_emb, flow, adarms_cond)。noisy_action_emb 是加噪动作嵌入，flow 是动作-噪声目标，adarms_cond 是时间步条件。实战：LIBERO 中 dof_mask 标记前 7 维，后 13 维零化。
432
mask = input_ids == self.action_token_id_set["action_token_id"]
433
mask_unsqueezed = mask.unsqueeze(-1)434
mask_expanded = mask_unsqueezed.expand_as(inputs_embeds)
435
action_mask = mask_expanded.to(inputs_embeds.device)
mask = input_ids == action_token_id 找动作 token 位置，shape (B,S)。unsqueeze(-1) 变为 (B,S,1)，expand_as 广播到 (B,S,hidden_dim)。action_mask 转移到设备。第 436 行空行。
437
noisy_action_emb = noisy_action_emb.to(
438
inputs_embeds.device, inputs_embeds.dtype
439
)
440
inputs_embeds = inputs_embeds.masked_scatter(action_mask, noisy_action_emb)
noisy_action_emb 转移到 device 和 dtype。inputs_embeds.masked_scatter(action_mask, noisy_action_emb) 将噪声化动作嵌入填入被掩码标记位置。此后 inputs_embeds 中动作 token 位置为 (batch_size, action_horizon_flow, hidden_dim) 的张量。第 441 行空行。
return (inputs_embeds, flow, adarms_cond) 返回修改后的嵌入、flow 监督目标、adarms 时间步条件。flow 用于训练 MSE 损失；adarms_cond 是 action_head 的时间编码。第 443 行空行。
444
@staticmethod445
def _update_position_ids(
446
position_ids,
447
moe_token_types,
448
positional_masks,
449
):
@staticmethod def _update_position_ids(position_ids, moe_token_types, positional_masks): 静态方法，调整 flow token 的 position_id 与 AR token 部分同步。被 forward():452 调用。无 self，参数为 (position_ids, moe_token_types, positional_masks)。
450
if (451
positional_masks is None
452
or "ar_predict_token_positions" not in positional_masks
453
):
454
return position_ids早期返回条件：如果没有 positional_masks 或其中没有 'ar_predict_token_positions' 键，直接返回原始 position_ids 不做调整。第 455 行空行。
456
new_position_ids = position_ids.clone()
457
ar_predict_token_positions = positional_masks["ar_predict_token_positions"]
458
flow_mask = moe_token_types == 1460
start_ar_pos, end_ar_pos = find_first_last_ones(ar_predict_token_positions)
461
start_flow_pos, end_flow_pos = find_first_last_ones(flow_mask)
new_position_ids = position_ids.clone() 安全克隆。ar_predict_token_positions = positional_masks['ar_predict_token_positions'] 获取 AR 预测 token 位置，shape (B,S) 布尔张量。flow_mask = moe_token_types == 1 创建 flow token 掩码（token_type=1）。find_first_last_ones() 返回每个 batch item 的第一个和最后一个 1 的位置，不存在则 -1。第 462 行空行。
463
for bs_i in range(position_ids.shape[1]):
464
if start_ar_pos[bs_i] != -1 and end_ar_pos[bs_i] != -1:
465
start_ar_ids = new_position_ids[:, bs_i, start_ar_pos[bs_i]]
466
start_flow_ids = new_position_ids[:, bs_i, start_flow_pos[bs_i]]
467
diff = start_flow_ids - start_ar_ids
468
new_position_ids[:, bs_i, start_flow_pos[bs_i] :] = position_ids[
469
:, bs_i, start_flow_pos[bs_i] :
470
] - diff.unsqueeze(-1)for bs_i in range(position_ids.shape[1]): 遍历 batch。若该 batch item 中存在 AR token，获取 AR 开始位置和 flow 开始位置的 position_id，计算差值。调整 flow 部分 position_ids：new_position_ids[:, bs_i, start_flow_pos[bs_i]:] = position_ids[:, bs_i, start_flow_pos[bs_i]:] - diff.unsqueeze(-1)。作用：使 flow token 相对位置编码与 AR 部分对齐，避免位置编码冲突。第 471 行空行。
return new_position_ids 返回调整后的 position_ids，flow token 部分的位置 id 已与 AR 部分同步。第 473 行空行为方法间隔。
L474–620两个关键的 attention mask 更新方法：_update_joint_attention_mask_2d 用于标准 2D mask（支持 padding/ar_predict/valid_flow 三层过滤），_update_joint_attention_flash_mask 用于 Flash Attention 优化格式（编码为 [L,R] 行索引对，避免矩阵级 O(S^2) 内存）。两者核心逻辑一致：动作专家(moe1)可双向注意，文本专家(moe0)严格因果；图文前缀与动作序列之间的跨专家注意由 causal_action_attention_mask 控制；LIBERO 实战中 valid_flow_action_positions 根据 action_horizon=10 后的有效长度动态掩码。
474
def _update_joint_attention_mask_2d(
475
self,476
attention_mask,
477
moe_token_types,
478
positional_masks,
479
):
480
if attention_mask.dim() == 3: # bs, seq_len, seq_len
481
return attention_mask483
bs, seq_len = moe_token_types.shape[0], moe_token_types.shape[1]
484
# Create a lower triangular matrix as a causal mask.485
causal_mask = torch.tril(
486
torch.ones(
487
(seq_len, seq_len), dtype=torch.bfloat16, device=moe_token_types.device
488
)
489
)
490
# Extended to the batch dimension.491
attention_mask = causal_mask.unsqueeze(0).expand(bs, -1, -1)
方法 _update_joint_attention_mask_2d 的签名与因果 mask 初始化（474-479 签名；480-481 若已是 3D 则直接返回；482 空行；483-491 从 moe_token_types 提取 B、S，生成下三角因果 mask，dtype=bfloat16，expand 到 (B,S,S)）。
493
if positional_masks is not None and "padding_positions" in positional_masks:
494
padding_positions = positional_masks["padding_positions"]
495
# The padding is set to zero.496
attention_mask = torch.where(
497
padding_positions[:, None, :],498
torch.zeros_like(attention_mask),
499
attention_mask,
500
)
501
# The padding is set to zero.502
attention_mask = torch.where(
503
padding_positions[:, :, None],504
torch.zeros_like(attention_mask),
505
attention_mask,
506
)
处理 padding_positions 掩码（492 空行；493-506 两个 torch.where 分别作用 [:, None, :] 行和 [:, :, None] 列，padding token 位置全设为 0，完全隐形，实战中 LIBERO 序列长度不一对齐时用）。
508
# Set all values in the moe1 section to 1, and disable the fast section.509
moe1_mask = (moe_token_types[:, :, None]) & (moe_token_types[:, None, :])
511
if (512
not self.config.causal_action_attention_mask
513
): # If a causal action attention mask is not used, then all elements in the moe1 section are set to 1.514
attention_mask = torch.where(
515
moe1_mask, torch.ones_like(attention_mask), attention_mask
516
)
构造动作-动作注意掩码与非因果处理（507 空行；508-509 moe1_mask=(moe[:,:,None])&(moe[:, None, :]) 提取双端都是动作 token 的位置，广播至 (B,S,S)；510-516 若 causal_action_attention_mask=False，把 moe1_mask 设为 1 覆盖因果约束，LIBERO 推荐 False 因为动作流是平行的）。
518
if (519
positional_masks is not None
520
and "ar_predict_token_positions" in positional_masks
521
):
522
ar_predict_token_positions = positional_masks["ar_predict_token_positions"]
523
moe1_mask = (moe_token_types[:, :, None]) & (524
ar_predict_token_positions[:, None, :]525
)
526
attention_mask = torch.where(
527
moe1_mask, torch.zeros_like(attention_mask), attention_mask
528
)
处理自回归预测 token 掩码（517 空行；518-528 ar_predict_token_positions 标记预测 token，新 moe1_mask=(moe_token_types[:,:,None])&(ar_predict[:, None,:]) 捕捉 query 为动作、key 为预测的位置，这些位置设为 0 防止信息泄露）。
530
if (531
positional_masks is not None
532
and "valid_flow_action_positions" in positional_masks
533
):
534
# true in moe_token_types but false in valid_flow_action_positions535
nonvalid_flow_action_positions = (
536
moe_token_types & ~positional_masks["valid_flow_action_positions"]
537
)
538
attention_mask = torch.where(
539
nonvalid_flow_action_positions[:, None, :],540
torch.zeros_like(attention_mask),
541
attention_mask,
542
)
543
attention_mask = torch.where(
544
nonvalid_flow_action_positions[:, :, None],545
torch.zeros_like(attention_mask),
546
attention_mask,
547
)
549
return attention_mask处理非有效的动作 token 掩码与返回（529 空行；530-547 valid_flow_action_positions 根据 action_horizon 截断标记有效区域，nonvalid 部分两个 torch.where 全设为 0，LIBERO 中 action_horizon=10 对应 10 步有效窗口；548 空行；549 返回更新的 (B,S,S) mask）。
551
def _update_joint_attention_flash_mask(
552
self,553
attention_mask,
554
moe_token_types,
555
positional_masks,
556
debug=False,557
):
558
device = moe_token_types.device
559
B, S = moe_token_types.shape
560
i32 = torch.int32
562
# ---- Return vector initialization ----563
LTS = torch.ones((B, S), device=device, dtype=i32) * S
564
UTE = (
565
torch.arange(S, device=device, dtype=i32).unsqueeze(0).expand(B, S).clone()566
)
方法 _update_joint_attention_flash_mask 的签名与初始化（550 空行；551-557 方法定义，Flash Attention 优化格式返回 [L,R] 边界对而非 O(S^2) 矩阵；558-560 提取 device、B、S、i32 类型；561 空行；562-566 初始化 LTS(Left-To-Start, 默认 S) 和 UTE(Up-To-End, 默认 arange)，编码可 attend 的行范围 [L,R]）。
568
# Handling padding positions569
if positional_masks is not None and "padding_positions" in positional_masks:
570
padding_positions = positional_masks["padding_positions"]
571
LTS[padding_positions] = 0572
UTE[padding_positions] = S
处理 padding_positions 的行索引映射（567-568 空行；569-572 padding 位置 LTS=0、UTE=S，实现 L>R 的交叉屏蔽使其完全不可见）。
574
# Handling ar predict tokens575
if (576
positional_masks is not None
577
and "ar_predict_token_positions" in positional_masks
578
):
579
start_ar_pos, end_ar_pos = find_first_last_ones(
580
positional_masks["ar_predict_token_positions"]
581
)
582
for bs_i in range(B):
583
if end_ar_pos[bs_i] != -1:
584
LTS[bs_i, positional_masks["ar_predict_token_positions"][bs_i]] = (
585
end_ar_pos[bs_i].to(i32) + 1586
)
处理自回归预测 token 的行索引映射（573-574 空行；575-586 find_first_last_ones 找 ar_predict 首末位置，对应位置的 query 设置 LTS=end_ar_pos+1，实现 L>R 使其不可见，防止 attend 预测 token）。
588
# Handling flow action bidirectional mask589
flow_mask = moe_token_types == 1590
if not self.config.causal_action_attention_mask:
591
start_flow_pos, end_flow_pos = find_first_last_ones(flow_mask)
592
for bs_i in range(B):
593
if start_flow_pos[bs_i] != -1:
594
UTE[bs_i, flow_mask[bs_i]] = start_flow_pos[bs_i].to(i32)
处理动作 token 双向注意的行索引映射（587-588 空行；589-594 flow_mask=(moe==1) 提取动作 token，若 causal_action_attention_mask=False，对这些位置设 UTE=start_flow_pos 使得 L>R 自动屏蔽，LIBERO 推荐此设置使动作段内互不 attend）。
596
# Handling validate flow597
if (598
positional_masks is not None
599
and "valid_flow_action_positions" in positional_masks
600
):
601
flow_mask = moe_token_types == 1602
nonvalid_flow_action_positions = (
603
flow_mask & ~positional_masks["valid_flow_action_positions"]
604
)
605
if nonvalid_flow_action_positions.any():606
LTS[nonvalid_flow_action_positions] = 0607
UTE[nonvalid_flow_action_positions] = S
处理有效动作区域的行索引映射（595-596 空行；597-607 valid_flow_action_positions 标记有效窗口（LIBERO action_horizon=10），nonvalid 部分 LTS=0、UTE=S 实现完全屏蔽，确保预测超出 horizon 的动作被掩码）。
609
LTS = LTS.unsqueeze(-1)610
UTE = UTE.unsqueeze(-1)612
startend_row_indices = torch.cat([LTS, UTE], dim=-1)维度调整与拼接（608 空行；609-612 LTS、UTE unsqueeze(-1) 至 (B,S,1)，torch.cat 沿最后维拼接成 (B,S,2) 的 [L,R] 对；注释掉的 startend_row_indices=LTS 是调试代码）。
613
# startend_row_indices = LTS615
# add num_heads dimension616
startend_row_indices = startend_row_indices.unsqueeze(1)618
return startend_row_indiceshead 维度插入与返回（613-615 空行；616 unsqueeze(1) 插入 num_heads 维度 (B,S,2)→(B,1,S,2)，Flash Attention kernel 广播到所有注意力头；617 空行；618 返回 startend_row_indices，相比 2D mask 的 O(S^2) 内存只需 O(2S)，sequence_length=4096 时节省 100 倍内存，部署推理时 Flash kernel 根据此范围动态裁剪 attention 计算）。
L621–746ActionGenerationMixin 类的初始化部分，包含混精度策略配置(to_bfloat16_for_selected_params)和动作 token 映射(define_action_token_id)，用于在 FSDP 或非分布式场景下管理模型精度和动作 token 的转换。
ActionGenerationMixin 类继承自 GenerationMixin；action_preprocessor 是 ActionProcessor 类型的类变量，在实例化时注入动作处理器。
624
def to_bfloat16_for_selected_params(self, fsdp_plugin=None, accelerator=None):
625
"""626
Keep some model parameters as float32, and convert others to bfloat16.627
- If `fsdp_plugin` exists, use FSDP v1's `mixed_precision` wrapper.628
- Otherwise, directly modify the parameter dtype.629
"""to_bfloat16_for_selected_params 函数签名和 docstring。该函数有两个执行路径：若有 fsdp_plugin 则用 FSDP v1 的 mixed_precision 包装器，否则手动转换参数精度。目的是让 action_preprocessor(流 matching 的 Action Head)、layer norm 等保持 fp32 精度以提高稳定性，其他层转为 bfloat16 节省显存。
631
def _assign_child(root_module, dotted_name: str, new_child):
632
parts = dotted_name.split(".")
633
parent = root_module
634
for p in parts[:-1]:
635
parent = getattr(parent, p)636
setattr(parent, parts[-1], new_child)
内嵌函数 _assign_child 用于按 dotted name 路径(如'model.encoder.layer0')递归定位父模块并替换子模块，是 FSDP 包装后重新挂载 wrapped module 的工具函数。
638
if fsdp_plugin:639
fsdp_version = getattr(fsdp_plugin, "fsdp_version", None)
640
if fsdp_version != 1:
641
raise RuntimeError("Only FSDP v1 is supported (fsdp_version=1).")
FSDP 分支入口。验证 fsdp_version==1，否则抛异常(仅支持 FSDP v1，v2/default 因 API 差异不兼容)。
643
device = getattr(644
accelerator, "device", torch.device("cuda", torch.cuda.current_device())
645
)
646
if isinstance(device, torch.device) and device.type == "cuda":
647
if device.index is not None:
648
torch.cuda.set_device(device.index)
649
device_id = device.index
从 accelerator 获取计算设备(默认 cuda:current)，若为 CUDA 则调用 torch.cuda.set_device() 激活目标卡。device_id 存储卡号用于后续 FSDP 初始化。shape 无关，纯设备管理逻辑。
全模型移到指定 device，为后续 FSDP 包装做准备。FSDP 通常要求模型已在目标设备上。
654
# Define the mixed-precision strategy.655
bf16_policy = MP(
656
param_dtype=torch.bfloat16,
657
reduce_dtype=torch.float32,
658
buffer_dtype=torch.bfloat16,
659
cast_forward_inputs=False,660
cast_root_forward_inputs=False,661
)
定义 bf16_policy：参数、buffer 用 bfloat16，梯度累加用 fp32(reduce_dtype)确保精度；cast_forward_inputs=False 避免在转发入口处自动转，由编码侧手控。图文分支(专家 0)用 bf16，省显存且 Qwen2.5-VL 本身对量化友好。
663
fp32_policy = MP(
664
param_dtype=torch.float32,
665
reduce_dtype=torch.float32,
666
buffer_dtype=torch.float32,
667
cast_forward_inputs=False,668
cast_root_forward_inputs=False,669
)
定义 fp32_policy：参数、buffer、reduce 全 fp32，保护 action_preprocessor 的数值稳定性(flow matching 的去噪步数敏感于精度)。
671
# Step 1️⃣: Identify the top-level ActionProcessor module and wrap it separately with FSDP (FP32).672
for name, module in list(self.named_modules()):
673
if isinstance(module, nn.Module) and any(
674
k in name.lower() for k in ["action_preprocessor"]
675
):
676
if any(True for _ in module.children()):
677
continue678
if getattr(module, "_fsdp_wrapped", False):
679
continue681
print(f"[FSDP v1] wrapping module in FP32: {name}")
682
wrapped = FSDP(
683
module,
684
mixed_precision=fp32_policy,
685
sharding_strategy=torch.distributed.fsdp.ShardingStrategy.SHARD_GRAD_OP,
686
backward_prefetch="BACKWARD_PRE",
687
device_id=device_id,
688
use_orig_params=True,689
)
690
_assign_child(self, name, wrapped)691
setattr(wrapped, "_fsdp_wrapped", True)
Step 1 枚举模块树，找 action_preprocessor，单独用 fp32_policy 包装成 FSDP。条件：名字含 'action_preprocessor'、有子模块、未被包装过；跳过已包装的避免重复。wrapped 后记标志位 _fsdp_wrapped，用 _assign_child 重新挂到原位置。LIBERO 实战里 action_preprocessor 约 0.45B 参数，fp32 包装后显存约占总模型(8B 左右)的 5%，权衡清晰。
693
# Step 2️⃣: The outermost layer uses unified FSDP (BF16 strategy).694
print("[FSDP v1] wrapping root model with bf16 mixed precision...")
695
self = FSDP(696
self,697
mixed_precision=bf16_policy,
698
sharding_strategy=torch.distributed.fsdp.ShardingStrategy.SHARD_GRAD_OP,
699
backward_prefetch="BACKWARD_PRE",
700
device_id=device_id,
701
use_orig_params=True,702
)
704
return self
Step 2 对根模型(已嵌入 fp32 的 action_preprocessor 子模块)用 bf16_policy 包装，FSDP 会递归扫描子树但遇 already-FSDP 子模块会跳过，避免嵌套。SHARD_GRAD_OP 按梯度分片、BACKWARD_PRE 预取(加速反向)、use_orig_params=True 使梯度与原参数对齐便于监听。return self 返回 FSDP 包装后的模型。
707
else:708
print("[INFO] Running manual dtype conversion (no FSDP).")
709
self.to(dtype=torch.float32)711
params_to_keep_float32 = []
712
for name, _ in self.named_parameters():
713
if any(
714
k in name715
for k in [
716
"input_layernorm",
717
"post_attention_layernorm",
718
"model.norm",
719
"action_preprocessor",
720
]
721
):
722
params_to_keep_float32.append(name)
724
for name, param in self.named_parameters():
725
if name not in params_to_keep_float32:
726
param.data = param.data.to(torch.bfloat16)
728
return self
else 分支(无 fsdp_plugin)：全模型先转 fp32，再扫参数列表筛选需保持 fp32 的项(action_preprocessor、input/post_attention_layernorm、model.norm)，其余参数转 bfloat16。无分布式框架时最简单的混精度方案，但参数转换涉及全模型遍历，对超大模型影响显著。LIBERO 场景下 3090 显存约 24G，分布式 fp32 完整参数会超限，此路径主要用单卡测试。
730
def define_action_token_id(self):
731
action_token_list = []
732
if self.action_tokenizer_type:
733
for i in range(self.action_tokenizer.vocab_size):
734
action_token_id = self.processor.tokenizer.convert_tokens_to_ids(735
f"<|action_token_{i}|>"
736
)
737
action_token_list.append(action_token_id)
739
action_token_id = self.processor.tokenizer.convert_tokens_to_ids("<|action|>")
740
propri_token_id = self.processor.tokenizer.convert_tokens_to_ids("<|propri|>")
741
self.action_token_id_set = {742
"action_token_list": action_token_list,
743
"propri_token_id": propri_token_id,
744
"action_token_id": action_token_id,
745
}
define_action_token_id 函数初始化动作 token 映射表。逻辑：若有 action_tokenizer，遍历其 vocab 为每个 token 翻译 <|action_token_i|> 的真实 id；无论如何查询 <|action|> 和 <|propri|> 的 token id，存入字典。这些 id 后续在 compute_loss(853行) 中用来掩码动作 token 位置、提取隐状态送 flow_loss。字典持久化到 self.action_token_id_set 供整个前向/后向流程访问。
L747–875ActionGenerationMixin 类的 add_lora 和 compute_loss 两个方法。add_lora 为 VLM 模型应用低秩适配器（PEFT-LoRA）；compute_loss 计算混合损失（文本 cross_entropy + 动作 flow 匹配），处理多数据集追踪、action token 准确率统计、dof_mask 加权的 flow MSE 损失。
747
def add_lora(
748
self, r=8, lora_alpha=32, target_modules=["q_proj", "v_proj"], lora_dropout=0.1
749
):
add_lora() 函数签名。被 __init__ 或外部脚本调用以启用 LoRA 微调。参数：r=秩（低秩维度8）、lora_alpha=缩放因子32、target_modules=作用目标(默认 Q/V 投影)、lora_dropout=LoRA 适配器内置 dropout(0.1)。
750
"""Add LoRA adapter"""简洁 docstring，说明方法功能为添加 LoRA 适配器到模型。实现将在后续行完成。
751
config = LoraConfig(
752
r=r,
753
lora_alpha=lora_alpha,
754
target_modules=target_modules,
755
lora_dropout=lora_dropout,
756
bias="none",
757
task_type="CAUSAL_LM",
758
)
创建 LoraConfig 对象，配置 PEFT 库的 LoRA 超参。参数包括：秩 r、缩放 lora_alpha、target_modules(Q/V 投影)、dropout、bias='none'(无偏置)、task_type='CAUSAL_LM'(因果语言建模)。这 8 行参数配置对应 peft 的标准接口。
759
self.model = get_peft_model(self.model, config)
调用 get_peft_model() 用 LoRA config 包装 self.model，返回可训练的 PEFT 模型。self.model 被原地替换为 LoRA 模型，后续只有适配器权重和 target_modules 可梯度更新。
打印 LoRA 模型的可训练参数统计。peft 库的 print_trainable_parameters() 输出可训练参数数、总参数数及可训练比例，用于验证 LoRA 是否正确应用。第 762 行为空行，并入此组。
763
def compute_loss(
764
self,765
hidden_states,
766
logits,
767
input_ids=None,768
dataset_names=None,769
labels=None,770
action_chunk=None,771
dof_mask=None,772
flow=None,773
flow_loss_mask=None,774
**kwargs,
775
):
compute_loss() 函数签名及参数列表。核心参数：hidden_states(形状 [B,S,H])、logits(形状 [B,S,vocab_size])、labels(形状 [B,S])、action_chunk(动作目标)、dof_mask(形状 [B,A,20])、flow(光流目标)、flow_loss_mask(损失掩码)、dataset_names(数据集标签)。被 forward() 或训练循环调用，返回汇总损失及各子损失/指标。
776
if input_ids is not None:
777
batch_size, seq_length = input_ids.shape
779
loss = 0780
cross_entropy_loss, flow_loss = None, None
从 input_ids 张量(形状 [B,S])提取 batch_size B 和 seq_length S，并初始化总损失与子损失。loss=0；cross_entropy_loss/flow_loss=None。这两个值后续用于 reshape 中间张量及 channel_loss 操作。第 778、781 行为空行，并入此组。
782
# if dataset_names is not None:783
# unique_datasets_name = list(set(dataset_names))784
# channel_loss_dict = {785
# dataset_name: torch.tensor(0.0, device=logits.device)786
# for dataset_name in _ACTION_DATASET_NAMES + _MULTIMODAL_DATASET_NAMES787
# }788
# channel_loss_count_dict = {789
# dataset_name: torch.tensor(0, device=logits.device)790
# for dataset_name in _ACTION_DATASET_NAMES + _MULTIMODAL_DATASET_NAMES791
# }792
# else:793
unique_datasets_name, channel_loss_dict, channel_loss_count_dict = (
794
None,795
None,796
None,797
)
注释掉的多数据集 channel_loss 初始化代码块（遗留代码，第 782-792 行），最后统一初始化为 None（第 793-797 行）。第 798 行为空行。因为 dataset_names 通常不传，所以这些值保持 None，线 819-830 的 for 循环实际不执行。LIBERO 实战中单数据集场景，这些都为 None。
if labels is not None 条件判断：如果有 token 预测标签，则计算文本损失。初始化 action_accuracy=0，待后续填充。第 801 行为空行，并入此组。
802
shift_logits = logits[..., :-1, :].contiguous()803
shift_labels = labels[..., 1:].contiguous()804
shift_logits = shift_logits.view(-1, self.config.vocab_size)
805
shift_labels = shift_labels.view(-1)806
# Enable model parallelism807
shift_labels = shift_labels.to(shift_logits.device)
808
non_ignored_mask = shift_labels != -100生成 shift logits/labels 并 reshape 为 1D、转设备、创建非 ignore 掩码。logits 去掉最后一个 token→[B,S-1,vocab_size]→[B*(S-1),vocab_size]；labels 去掉首个 token→[B,S-1]→[B*(S-1)]。contiguous() 确保内存连续。non_ignored_mask 标记 != -100 的有效 token。
809
_cross_entropy_loss = self.loss_fct(shift_logits, shift_labels)810
cross_entropy_loss = (
811
_cross_entropy_loss[non_ignored_mask].mean()
812
if non_ignored_mask.any()813
else torch.tensor(0.0, device=shift_logits.device)
814
)
计算 cross_entropy_loss，reshape 为 2D。loss_fct(reduction='none') 返回 [B*(S-1)] 的逐 token 损失，通过 non_ignored_mask 过滤有效 token 后 mean()；无有效 token 则返回 0。随后 reshape 到 [B, S-1] 以供 channel_loss 计算。第 815 行为空行，并入此组。
816
# compute channel loss817
_cross_entropy_loss = _cross_entropy_loss.view(batch_size, seq_length - 1)818
non_ignored_mask = non_ignored_mask.view(batch_size, seq_length - 1)819
for dataset_name_i in unique_datasets_name:
820
dataset_mask = torch.tensor(
821
[name == dataset_name_i for name in dataset_names],
822
device=logits.device,
823
)
824
combined_mask = dataset_mask.unsqueeze(1) & non_ignored_mask825
channel_loss_dict[dataset_name_i] = (
826
_cross_entropy_loss[combined_mask].sum()
827
if combined_mask.any()828
else torch.tensor(0.0, device=shift_logits.device)
829
)
830
channel_loss_count_dict[dataset_name_i] += combined_mask.sum()
reshape 掩码为 2D（[B, S-1]），按 dataset_name 分组计算 channel_loss。逐数据集创建二阶掩码 combined_mask，累加总损失和样本计数。**但因为 unique_datasets_name=None（第 793-797 行），这个 for 循环永远不会执行**。代码保留为未来多数据集支持预留。LIBERO 单数据集实战被跳过。第 831 行为空行，并入此组。
832
if not torch.isnan(cross_entropy_loss):
833
loss += cross_entropy_loss
834
else:835
with torch.no_grad():836
cross_entropy_loss.detach()
NaN 检查：if not torch.isnan(cross_entropy_loss) 则累加到 loss；否则 detach()（逻辑有缺陷，detach 后未赋回）。该段代码可能是容错代码，实际在无 NaN 时直接走 loss += cross_entropy_loss。第 837 行为空行，并入此组。
if 条件：检查是否定义了 action token（即 len(self.action_token_id_set['action_token_list']) > 0）。如有，则计算该批 action token 的预测准确率。
840
shift_logits = logits[..., :-1, :].contiguous()841
action_preds = shift_logits.argmax(dim=-1)842
shift_labels = labels[..., 1:].contiguous()843
action_mask = (
844
shift_labels > self.action_token_id_set["action_token_list"][0]
845
)
846
correct_preds = (action_preds == shift_labels) & action_mask
847
action_accuracy = (
848
correct_preds.sum().float() / action_mask.sum().float()
849
)
850
channel_loss_dict["action_accuracy"] = action_accuracy
重新提取 shift logits/labels、创建 action_mask，计算 action token 准确率。argmax(dim=-1) 得预测 ID→[B,S-1]；action_mask 标记 > action_token_list[0] 的 token。correct_preds 是预测与标签匹配且处于 action token 范围的位置；action_accuracy = 正确数 / 总数。赋到 channel_loss_dict['action_accuracy']（虽然此时为 None，实际会崩）。LIBERO 因 action_tokenizer_type 为空通常跳过。
852
if action_chunk is not None:
853
action_mask = input_ids == self.action_token_id_set["action_token_id"]
854
if action_mask.any():855
action_hidden_states = hidden_states[action_mask].to(torch.float32)
856
flow = flow.reshape(-1, flow.shape[-1])
if action_chunk is not None 条件及动作隐状态提取。action_mask 找出所有 action_token_id 的位置 [B,S]；if action_mask.any() 确保至少有一个。action_hidden_states = hidden_states[action_mask]→[K,H]（K=action_token_总数，H=2048），转 float32 以减少数值误差（flow matching 对精度敏感）。flow reshape 为 [K,D]。
857
_flow_loss = self.action_preprocessor.flow_loss(858
action_hidden_states, flow, action_chunk, dof_mask, flow_loss_mask
859
)
860
if isinstance(_flow_loss, torch.Tensor):
861
flow_loss = _flow_loss.mean()
862
loss += flow_loss * self.config.flow_loss_weight调用 self.action_preprocessor.flow_loss()，传入：action_hidden_states([K,H])、flow([K,D])、action_chunk(动作目标)、dof_mask(形状[B,A,D])、flow_loss_mask(可选)。action_preprocessor 是 ActionProcessor 实例，flow_loss 内部投影隐状态→速度场、计算 MSE、乘以 dof_mask 和 flow_loss_mask，返回加权损失[K,D]。若返回 Tensor，mean() 后乘 flow_loss_weight 累加到 loss。
863
_flow_loss = _flow_loss.view(
864
dof_mask.shape[0], dof_mask.shape[1], dof_mask.shape[2]
865
)
867
return (868
loss,
869
cross_entropy_loss,
870
flow_loss,
871
channel_loss_dict,
872
channel_loss_count_dict,
873
)
reshape _flow_loss 回 [B,A,D] 形状（用 dof_mask.shape 的各维）供后续日志/可视化；随后返回 5 元组。返回值：(1) loss=总加权损失；(2) cross_entropy_loss=文本 token 损失；(3) flow_loss=动作流匹配损失均值；(4) channel_loss_dict=按数据集分组损失（通常 None）；(5) channel_loss_count_dict=各数据集样本计数（通常 None）。调用方在训练循环中使用 loss 反向传播、在日志中记录各子损失和指标。注意 _flow_loss 可能已 mean()，reshape 前应检查形状避免 mismatch。第 874-875 行为空行，并入此组。
L876–987AttentionsSelectorMixin 类实现注意力机制的自动选择和验证。核心方法 _autoset_attn_implementation 按优先级尝试 flash_attention_2、flex_attention、SDPA、eager 等实现；处理用户显式指定、硬件兼容性检查（HIP/ROCM）和向后兼容。第二个方法 _check_and_adjust_attn_implementation 进行简单的类型验证。
878
@classmethod879
def _autoset_attn_implementation(
880
cls,881
config,
882
use_flash_attention_2: bool = False,
883
torch_dtype: Optional[torch.dtype] = None,884
device_map: Optional[Union[str, Dict[str, int]]] = None,
885
check_device_map: bool = True,
886
):
@classmethod 装饰符 + _autoset_attn_implementation 方法签名。入参：config(模型配置)、use_flash_attention_2(废弃参数)、torch_dtype(精度)、device_map(设备映射)、check_device_map(是否检查设备映射)。
887
"""888
Automatically checks and dispatches to a default attention implementation. In order of priority:889
1. An implementation specified in `config._attn_implementation` (due for example to the argument attn_implementation="sdpa" in from_pretrained).890
2. DEPRECATED: if use_flash_attention_2 is set to `True` and `flash_attn` is available, flash attention. (`LlamaFlashAttention` for example)891
3. SDPA implementation, if available and supported by the model type. (`LlamaSdpaAttention` for example)892
4. The default model's implementation otherwise (`LlamaAttention` for example) .893
"""Docstring 说明该方法的作用：自动检查并分发最合适的注意力实现。按优先级列出四种选项：(1)配置显式指定、(2)废弃的 flash_attention_2 参数、(3)SDPA(如果支持)、(4)默认实现。
894
# Here we use config._attn_implementation_internal to check whether the attention implementation was explicitly set by the user.895
# The property `PretrainedConfig._attn_implementation` is never `None`, for backward compatibility (always fall back on "eager").896
# The `hasattr` here is used as some Transformers tests for some reason do not call PretrainedConfig __init__ (e.g. test_no_super_init_config_and_model)897
requested_attn_implementation = None898
if (899
hasattr(config, "_attn_implementation_internal")
900
and config._attn_implementation_internal is not None
901
):
三行注释 + 初始化 requested_attn_implementation=None。注释说明：检查 _attn_implementation_internal 属性来判断用户是否显式指定实现；PretrainedConfig._attn_implementation 永不为 None(向后兼容总是回落到 eager)；hasattr 用来处理某些 Transformers 测试不调用 __init__ 的奇怪情况。
902
if (903
config._attn_implementation != "flash_attention_2"
904
and use_flash_attention_2905
):
906
raise ValueError(907
f'Both attn_implementation="{config._attn_implementation}" and `use_flash_attention_2=True` were used when loading the model, which are not compatible.'
908
' We recommend to just use `attn_implementation="flash_attention_2"` when loading the model.'
909
)
嵌套 if 块：冲突检查。如果 config 指定了 attn_implementation 且不是 flash_attention_2，但同时传了 use_flash_attention_2=True，则抛 ValueError 提示两者不兼容，建议用 attn_implementation 参数。
911
if (912
not isinstance(config._attn_implementation, dict)
913
and config._attn_implementation914
not in ["eager"]
915
+ ALL_ATTENTION_FUNCTIONS.valid_keys()
916
+ X2ROBOT_ATTENTION_FUNCTIONS
917
):
第二个 if 块：验证用户指定的实现是否被支持。条件：(1)config._attn_implementation 不是字典、(2)不在 ['eager'] + ALL_ATTENTION_FUNCTIONS + X2ROBOT_ATTENTION_FUNCTIONS 的列表中。
918
message = f'Specified `attn_implementation="{config._attn_implementation}"` is not supported. The only possible arguments are `attn_implementation="eager"` (manual attention implementation)'
919
if cls._supports_flash_attn_2:
920
message += ', `"attn_implementation=flash_attention_2"` (implementation using flash attention 2)'
921
if cls._supports_sdpa:
922
message += ', `"attn_implementation=sdpa"` (implementation using torch.nn.functional.scaled_dot_product_attention)'
923
if cls._supports_flex_attn:
924
message += ', `"attn_implementation=flex_attention"` (implementation using torch\'s flex_attention)'
925
raise ValueError(message + ".")
构建错误消息并抛异常。消息模板列出所有支持的实现：eager(默认)、flash_attention_2(若模型支持)、sdpa(若支持)、flex_attention(若支持)。根据模型类的属性 _supports_flash_attn_2、_supports_sdpa、_supports_flex_attn 动态添加。
927
# If a config is passed with a preset attn_implementation, we skip the automatic dispatch and use the user-provided config, with hard checks that the requested attention implementation is available.928
requested_attn_implementation = config._attn_implementation_internal
空行 + 保存用户显式请求的实现到 requested_attn_implementation = config._attn_implementation_internal，以供后续优先级决策使用。
930
if use_flash_attention_2:931
logger.warning_once(
932
'The model was loaded with use_flash_attention_2=True, which is deprecated and may be removed in a future release. Please use `attn_implementation="flash_attention_2"` instead.'
933
)
934
config._attn_implementation = "flash_attention_2"
936
if config._attn_implementation == "flash_attention_2":
第一个 if-elif 链的顶级条件：config._attn_implementation 已被设为 flash_attention_2（由参数或前序逻辑）。调用 cls._check_and_enable_flash_attn_2 检查并启用，传入 hard_check_only=False 允许自动启用。
elif 处理 flex_attention：用户显式请求了 flex_attention。调用 _check_and_enable_flex_attn，hard_check_only=True 表示仅验证不自动启用。
939
torch_dtype=torch_dtype,
940
device_map=device_map,
941
hard_check_only=False,942
check_device_map=check_device_map,
943
)
944
elif requested_attn_implementation == "flex_attention":
945
config = cls._check_and_enable_flex_attn(config, hard_check_only=True)
946
elif (947
requested_attn_implementation in [None, "sdpa"]
948
and not is_torch_xla_available()
949
):
950
# use_flash_attention_2 takes priority over SDPA, hence SDPA treated in this elif.elif 处理 SDPA：用户请求 None(自动选择)或显式 sdpa，且不在 XLA(Google TPU 编译)环境。调用 _check_and_enable_sdpa，hard_check_only 根据 requested_attn_implementation 决定：None 时为 False(自动选)、显式 sdpa 时为 True(仅验证)。
951
config = cls._check_and_enable_sdpa(952
config,
953
hard_check_only=(
954
False if requested_attn_implementation is None else True
955
),
956
)
958
if (959
torch.version.hip is not None
960
and config._attn_implementation == "sdpa"
961
and torch.cuda.device_count() > 1
HIP/ROCM(AMD GPU)多卡特殊处理：如果是 HIP 环境、config 已设为 SDPA、多 GPU 且 PyTorch < 2.4.1，则警告 SDPA 在 ROCM 多卡下性能差(FA backend 问题)，禁用 flash_sdp 回落到其他后端。
962
and version.parse(torch.__version__) < version.parse("2.4.1")
963
):
964
logger.warning_once(
965
"Using the `SDPA` attention implementation on multi-gpu setup with ROCM may lead to performance issues due to the FA backend. Disabling it to use alternative backends."
处理自定义实现：三个 elif 分别处理(1)requested_attn_implementation 在 ALL_ATTENTION_FUNCTIONS 中，则采用它；(2)是字典类型(自定义路由)，设为 None；(3)在 X2ROBOT_ATTENTION_FUNCTIONS 中，pass 保持原状；(4)其他情况默认回落到 eager。
966
)
967
torch.backends.cuda.enable_flash_sdp(False)968
elif requested_attn_implementation in ALL_ATTENTION_FUNCTIONS.valid_keys():
969
config._attn_implementation = requested_attn_implementation
970
elif isinstance(requested_attn_implementation, dict):
标记 config._attn_implementation_autoset=True 表示实现已自动设置，返回更新后的 config。这个标志用来区分用户显式指定 vs 系统自动选择。
971
config._attn_implementation = None972
elif config._attn_implementation in X2ROBOT_ATTENTION_FUNCTIONS:
973
pass974
else:975
config._attn_implementation = "eager"
977
config._attn_implementation_autoset = True978
return config980
def _check_and_adjust_attn_implementation(
981
self, attn_implementation: Optional[str], is_init_check: bool = False
982
) -> str:983
assert (984
attn_implementation
985
in ["eager", "flash_attention_2", "sdpa"] + X2ROBOT_ATTENTION_FUNCTIONS
986
)
987
return attn_implementation第二个方法 _check_and_adjust_attn_implementation：入参 attn_implementation(可选字符串)、is_init_check(初始化检查标志)。断言 attn_implementation 在允许列表中(eager、flash_attention_2、sdpa + X2ROBOT_ATTENTION_FUNCTIONS)，返回原值。方法名暗示可调整但实现仅验证，为扩展预留接口。
源码零改写,与仓库 wall-x/wall_x/model/vla_mixin.py 逐字节一致(commit 97406f2)。
生成于 wall-x LIBERO 微调项目 · ← 返回导读