跳转至

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 通信模式完全一致。