大模型常见算子定义

news2024/10/20 15:46:01

本文将汇总大模型常用的算子定义，方便快速根据定义公式评估其计算量。

LayerNorm

这是在BERT、GPT等模型中广泛使用的LayerNorm：

RMSNorm

RMSNorm(root mean square)发现LayerNorm的中心偏移没什么用(减去均值等操作)。将其去掉之后，效果几乎不变，但是速度提升了40%。最终公式为：

注意除了没有减均值，加偏置以外，分母上求的RMS而不是方差

SwiGLU/SiLU

LLaMA没有使用ReLU，而是使用了SwiGLU，有时也被称为SiLU，效果类似平滑版的ReLU。公式如下：

y = sigmoid(x) * x

RoPE

LLaMA使用了Rotary Position Embedding。对于Q的第m个位置向量q，通过以下方法注入相对位置编码

公式中第二、四项的计算代码：

class LlamaRotaryEmbedding(torch.nn.Module):
    def __init__(self, dim, max_position_embeddings=2048, base=10000, device=None):
        super().__init__()
        inv_freq = 1.0 / (base ** (torch.arange(0, dim, 2).float().to(device) / dim))
        self.register_buffer("inv_freq", inv_freq)

        # Build here to make `torch.jit.trace` work.
        self.max_seq_len_cached = max_position_embeddings
        t = torch.arange(self.max_seq_len_cached, device=self.inv_freq.device, dtype=self.inv_freq.dtype)
        freqs = torch.einsum("i,j->ij", t, self.inv_freq)
        # Different from paper, but it uses a different permutation in order to obtain the same calculation
        emb = torch.cat((freqs, freqs), dim=-1)
        self.register_buffer("cos_cached", emb.cos()[None, None, :, :], persistent=False)
        self.register_buffer("sin_cached", emb.sin()[None, None, :, :], persistent=False)

    def forward(self, x, seq_len=None):
        # x: [bs, num_attention_heads, seq_len, head_size]
        # This `if` block is unlikely to be run after we build sin/cos in `__init__`. Keep the logic here just in case.
        if seq_len > self.max_seq_len_cached:
            self.max_seq_len_cached = seq_len
            t = torch.arange(self.max_seq_len_cached, device=x.device, dtype=self.inv_freq.dtype)
            freqs = torch.einsum("i,j->ij", t, self.inv_freq)
            # Different from paper, but it uses a different permutation in order to obtain the same calculation
            emb = torch.cat((freqs, freqs), dim=-1).to(x.device)
            self.register_buffer("cos_cached", emb.cos()[None, None, :, :], persistent=False)
            self.register_buffer("sin_cached", emb.sin()[None, None, :, :], persistent=False)
        return (
            self.cos_cached[:, :, :seq_len, ...].to(dtype=x.dtype),
            self.sin_cached[:, :, :seq_len, ...].to(dtype=x.dtype),
        )


# 在LlamaAttention通过以下命令调用：
cos, sin = self.rotary_emb(seq_len=kv_seq_len)

公式中第三项的计算代码

# 在接下来的apply_rotary_pos_emb函数里调用

def rotate_half(x):
    x1 = x[..., : x.shape[-1] // 2]
    x2 = x[..., x.shape[-1] // 2 :]
    return torch.cat((-x2, x1), dim=-1)

最后通过以下代码得到结合了位置编码的Q,K(K和Q使用同样的方式进行位置编码)。

def apply_rotary_pos_emb(q, k, cos, sin, position_ids):
    q_embed = (q * cos[position_ids]) + (rotate_half(q) * sin[position_ids])
    k_embed = (k * cos[position_ids]) + (rotate_half(k) * sin[position_ids])
    return q_embed, k_embed

# 在LLamaAttention中通过以下命令调用：
query_states, key_states = apply_rotary_pos_emb(query_states, key_states, cos, sin, position_ids)

下图蓝色字体是原始论文RoPE的推导公式，两种都是常用的RoPE实现。