apertus.py — Apertus 注意力层 Ulysses SP 适配

文件路径

verl/models/transformers/apertus.py

文件概述

为 Apertus 模型（SwissAI 开发）的注意力层提供 Ulysses 序列并行支持。结构与 llama.py 类似，关键区别在于 Apertus 在 RoPE 之前会进行 QK 归一化 (QK Normalization)。

关键代码讲解

QK 归一化 -- 与 LLaMA 的核心区别

def apertus_attn_forward(self, hidden_states, position_embeddings, attention_mask, ...):
    query_states = self.q_proj(hidden_states).view(bsz, q_len, -1, self.head_dim).transpose(1, 2)
    key_states = self.k_proj(hidden_states).view(bsz, q_len, -1, self.head_dim).transpose(1, 2)
    value_states = self.v_proj(hidden_states).view(bsz, q_len, -1, self.head_dim).transpose(1, 2)

    # Apertus 特有：QK 归一化，在 AlltoAll 之前执行
    query_states = self.q_norm(query_states)
    key_states = self.k_norm(key_states)

    ########## AlltoAll for Ulysses ##########
    # ... 后续与 llama.py 完全相同

注意顺序：QK 归一化必须在 AlltoAll 通信 之前 执行。因为归一化是逐头操作的，在 AlltoAll 之前每个 GPU 拥有所有头，可以正确进行归一化。

其余流程

AlltoAll 通信、RoPE、attention 计算、反向 AlltoAll 的流程与 llama_attn_forward() 完全相同，此处不再重复。

核心函数列表

函数名	作用
apertus_attn_forward()	Apertus 模型的注意力替换函数（含 QK 归一化 + Ulysses SP）

与其他模块的关系

• 被 monkey_patch.py 中的 apply_monkey_patch() 调用
• 仅支持 transformers 4.48+ 版本
• 与 llama.py 共享 Ulysses SP 逻辑

小结

Apertus 的注意力适配是 LLaMA 适配的变体，唯一区别是在 AlltoAll 通信之前插入了 QK 归一化步骤。QK 归一化有助于训练稳定性，是部分大模型（如 PaLM、ViT-22B）采用的技术。