【Llama源码】旋转位置编码ROPE--源码阅读

原创

已于 2024-05-29 22:21:44 修改 · 1.4k 阅读

13 ·

CC 4.0 BY-SA版权

文章标签：

#llama #python #pytorch

于 2024-05-29 22:16:32 首次发布

旋转矩阵计算

rotary_emb 对应 $L l ama R o t a ry E mb e dd in g$ 层，其中内置 $ini t$ 初始化方法和 $f or w a r d$ 前向调用，负责生成旋转矩阵中的 $cos$ 和 $s in$ 。

代码

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