qwen2.py — Qwen2 注意力层 Ulysses SP 适配¶
文件路径¶
verl/models/transformers/qwen2.py
文件概述¶
为 Qwen2 模型的注意力层提供支持 Ulysses 序列并行的替换函数。结构与 llama.py 几乎完全相同,主要区别在于 Qwen2 支持 滑动窗口注意力 (Sliding Window Attention)。
关键代码讲解¶
与 LLaMA 的主要差异:滑动窗口¶
# Qwen2 特有的滑动窗口逻辑
if (
self.config.use_sliding_window
and getattr(self.config, "sliding_window", None) is not None
and self.layer_idx >= self.config.max_window_layers
):
sliding_window = self.config.sliding_window
else:
sliding_window = None
Qwen2 的一部分层使用滑动窗口注意力(只关注窗口内的 token),而上层使用全局注意力。max_window_layers 控制从第几层开始使用滑动窗口。
GQA (Grouped Query Attention) 处理¶
# 在 4.45-4.47 版本中,显式重复 KV 头
key_states = repeat_kv(key_states, self.num_key_value_groups)
value_states = repeat_kv(value_states, self.num_key_value_groups)
Qwen2 使用 GQA,即多个 Query 头共享一组 KV 头。在 4.45-4.47 版本中需要手动 repeat KV 头来匹配 Query 的头数。
两个版本¶
# transformers 4.45-4.47
def qwen2_flash_attn_forward(self, hidden_states, attention_mask, position_ids, ...)
# transformers 4.48+
def qwen2_attn_forward(self, hidden_states, position_embeddings, attention_mask, ...)
4.48+ 版本使用新的 attention_interface 机制,通过 ALL_ATTENTION_FUNCTIONS 注册表选择注意力实现,并且 sliding_window 参数直接传给 attention interface。
# 4.48+ 版本,sliding_window 传给 attention_interface
attn_output, attn_weights = attention_interface(
self, query_states, key_states, value_states, attention_mask,
dropout=0.0 if not self.training else self.attention_dropout,
scaling=self.scaling,
sliding_window=sliding_window, # main diff with Llama
**kwargs,
)
核心函数列表¶
| 函数名 | 作用 |
|---|---|
qwen2_flash_attn_forward() |
transformers 4.45-4.47 的 Qwen2 注意力替换函数 |
qwen2_attn_forward() |
transformers 4.48+ 的 Qwen2 注意力替换函数 |
与其他模块的关系¶
- 被
monkey_patch.py中的apply_monkey_patch()调用 - 与
llama.py共享相同的 Ulysses SP 通信逻辑 - 依赖
verl.utils.ulysses的 AlltoAll 原语
小结¶
Qwen2 的注意力适配与 LLaMA 几乎相同,核心区别是增加了滑动窗口注意力的支持。Ulysses SP 的 AlltoAll 通信模式完全一致。