2023年的深度学习入门攻略(20) – LLaMA 2模型解析

上一节咱们把LLaMA 2的生成进程以及封装的进程的代码简略介绍了下。还差LLaMA 2的模型部分没有介绍。这一节咱们就来介绍下LLaMA 2的模型部分。 这一部分需要一些深度神经网络的基础知识，不明白的话不用着急，后面的文章咱们都会介绍到。

均平方根标准化

RMSNorm是一种改进的LayerNorm技能，LayerNorm是Layer normalization，意思是层归一化。。层归一化用于帮助安稳练习并促进模型收敛，由于它具有处理输入和权重矩阵的从头居中和从头缩放的才能。

RMSNorm是2019年的论文《Root Mean Square Layer Normalization》中提出的。它假设LayerNorm中的从头居中性质并不是必需的，所以RMSNorm根据均方根（RMS）对某一层中的神经元输入进行规范化，赋予模型从头缩放的不变性属性和隐式学习率自适应才能。相比LayerNorm，RMSNorm在核算上更简略，因而愈加高效。

了解了之后，咱们来看下RMSNorm的代码完成。我把注释直接写在代码里。

class RMSNorm(torch.nn.Module):
    # 类的构造函数，它接受两个参数：dim和eps。dim是期望标准化的特征维度，eps是一个十分小的数，用于避免除以零的过错。
    def __init__(self, dim: int, eps: float = 1e-6):
        super().__init__()
        # 设置类的属性。eps是构造函数传入的参数。
        self.eps = eps
        # weight是一个可学习的参数，它是一个由1填充的张量，尺度为dim。
        self.weight = nn.Parameter(torch.ones(dim))
    def _norm(self, x):
        # 首要，对输入x求平方并核算终究一个维度的平均值，然后加上一个十分小的数self.eps避免出现零，接着对成果开平方根并求倒数，终究将成果与原始输入x相乘。
        return x * torch.rsqrt(x.pow(2).mean(-1, keepdim=True) + self.eps)
    def forward(self, x):
        # 首要，它将输入x转化为浮点数并进行标准化，然后将标准化的成果转化回x的类型。终究，将成果与权重self.weight相乘，得到终究的输出。
        output = self._norm(x.float()).type_as(x)
        return output * self.weight

方位编码

咱们温习下第3讲从前介绍过的Transformer结构。

方位编码是Transformer中的一个重要组成部分，它的作用是为输入序列中的每个方位提供一个方位向量，以便Transformer能够区分不同方位的单词。

Transformer中的方位编码是经过将正弦和余弦函数的值作为方位向量的元从来完成的。这些函数的周期是不同的，因而它们的值在不同的方位是不同的。这样，Transformer就能够经过方位编码来区分不同方位的单词。

LLaMA并没有运用正弦函数。

# dim是特征的维度，end应该是预核算的时序方位的数量，theta是一个常数，用于调整频率的尺度
def precompute_freqs_cis(dim: int, end: int, theta: float = 10000.0):
    # 首要生成一个从0到dim的步长为2的整数序列，然后取前dim // 2个元素，将这些元素转换为浮点类型，然后除以dim
    # 得到的成果再次被用作theta的指数，终究取其倒数，得到一组频率
    freqs = 1.0 / (theta ** (torch.arange(0, dim, 2)[: (dim // 2)].float() / dim))
    # 生成一个从0到end的整数序列，这个序列在同一个设备上创立，这个设备是freqs的设备
    t = torch.arange(end, device=freqs.device)  # type: ignore
    # 核算t和freqs的外积，然后将成果转换为浮点类型
    freqs = torch.outer(t, freqs).float()  # type: ignore
    # 将freqs从直角坐标系转换为极坐标系
    # torch.polar(r, theta)的功用是根据极径r和极角theta生成复数，这里的r是freqs的形状的全1张量，theta则是freqs。
    freqs_cis = torch.polar(torch.ones_like(freqs), freqs)  # complex64
    return freqs_cis

将方位坐标核算出来之后，咱们还需要将其变型成与输入的形状共同。 咱们来看下是如何完成的：

def reshape_for_broadcast(freqs_cis: torch.Tensor, x: torch.Tensor):
    # 获取输入张量x的维度数
    ndim = x.ndim
    assert 0 <= 1 < ndim
    assert freqs_cis.shape == (x.shape[1], x.shape[-1])
    # 关于每个维度，假如它是第二个维度或终究一个维度，则保留原来的巨细；不然，将其设置为 1
    shape = [d if i == 1 or i == ndim - 1 else 1 for i, d in enumerate(x.shape)]
    # 将freqs_cis调整为shape指定的形状，并回来成果
    return freqs_cis.view(*shape)

然后，将矩阵和方位编码相乘起来：

def apply_rotary_emb(
    xq: torch.Tensor,
    xk: torch.Tensor,
    freqs_cis: torch.Tensor,
) -> Tuple[torch.Tensor, torch.Tensor]:
    # 将 xq 和 xk 转换为复数张量，并将它们的形状调整为终究一个维度为 2 的形状
    xq_ = torch.view_as_complex(xq.float().reshape(*xq.shape[:-1], -1, 2))
    xk_ = torch.view_as_complex(xk.float().reshape(*xk.shape[:-1], -1, 2))
    # 调用刚刚讲过的 reshape_for_broadcast 函数来将 freqs_cis 调整为与 xq_ 兼容的形状
    freqs_cis = reshape_for_broadcast(freqs_cis, xq_)
    # 将xq_和xk_与freqs_cis进行逐元素的复数乘法，然后将得到的成果视为实数，终究将终究两个维度合并为一个维度
    xq_out = torch.view_as_real(xq_ * freqs_cis).flatten(3)
    xk_out = torch.view_as_real(xk_ * freqs_cis).flatten(3)
    # 运用 type_as 办法将其转换回与输入相同的数据类型
    return xq_out.type_as(xq), xk_out.type_as(xk)

LLaMA的注意力机制

class Attention(nn.Module):
    def __init__(self, args: ModelArgs):
        super().__init__()
        self.n_kv_heads = args.n_heads if args.n_kv_heads is None else args.n_kv_heads
        model_parallel_size = fs_init.get_model_parallel_world_size()
        self.n_local_heads = args.n_heads // model_parallel_size
        self.n_local_kv_heads = self.n_kv_heads // model_parallel_size
        self.n_rep = self.n_local_heads // self.n_local_kv_heads
        self.head_dim = args.dim // args.n_heads
        self.wq = ColumnParallelLinear(
            args.dim,
            args.n_heads * self.head_dim,
            bias=False,
            gather_output=False,
            init_method=lambda x: x,
        )
        self.wk = ColumnParallelLinear(
            args.dim,
            self.n_kv_heads * self.head_dim,
            bias=False,
            gather_output=False,
            init_method=lambda x: x,
        )
        self.wv = ColumnParallelLinear(
            args.dim,
            self.n_kv_heads * self.head_dim,
            bias=False,
            gather_output=False,
            init_method=lambda x: x,
        )
        self.wo = RowParallelLinear(
            args.n_heads * self.head_dim,
            args.dim,
            bias=False,
            input_is_parallel=True,
            init_method=lambda x: x,
        )
        self.cache_k = torch.zeros(
            (
                args.max_batch_size,
                args.max_seq_len,
                self.n_local_kv_heads,
                self.head_dim,
            )
        ).cuda()
        self.cache_v = torch.zeros(
            (
                args.max_batch_size,
                args.max_seq_len,
                self.n_local_kv_heads,
                self.head_dim,
            )
        ).cuda()

代码尽管多，但是都是完成head和q,k,v,很好了解：

• 初始化时定义了n_kv_heads, n_local_heads等表明head数量的变量。
• wq,wk,wv三个线性层别离用于生成query,key和value。采用ColumnParallelLinear完成散布并行。
• wo线性层对多头attention的输出做交融,采用RowParallelLinear完成散布并行。
• cache_k和cache_v用于缓存key和value,加快自注意力的核算。
• 并行线性层的运用以及caching机制,能够加快自注意力在大batch巨细场景下的练习和推理。
• 全体规划完成了高效的散布式并行自注意力核算,能够扩展到大规模多GPU/机器环境,处理长序列任务。

然后，咱们将各注意力的子模块集成起来：

    def forward(
        self,
        x: torch.Tensor,
        start_pos: int,
        freqs_cis: torch.Tensor,
        mask: Optional[torch.Tensor],
    ):
        bsz, seqlen, _ = x.shape
        xq, xk, xv = self.wq(x), self.wk(x), self.wv(x)
        xq = xq.view(bsz, seqlen, self.n_local_heads, self.head_dim)
        xk = xk.view(bsz, seqlen, self.n_local_kv_heads, self.head_dim)
        xv = xv.view(bsz, seqlen, self.n_local_kv_heads, self.head_dim)
        xq, xk = apply_rotary_emb(xq, xk, freqs_cis=freqs_cis)
        self.cache_k = self.cache_k.to(xq)
        self.cache_v = self.cache_v.to(xq)
        self.cache_k[:bsz, start_pos : start_pos + seqlen] = xk
        self.cache_v[:bsz, start_pos : start_pos + seqlen] = xv
        keys = self.cache_k[:bsz, : start_pos + seqlen]
        values = self.cache_v[:bsz, : start_pos + seqlen]
        # repeat k/v heads if n_kv_heads < n_heads
        keys = repeat_kv(keys, self.n_rep)  # (bs, seqlen, n_local_heads, head_dim)
        values = repeat_kv(values, self.n_rep)  # (bs, seqlen, n_local_heads, head_dim)
        xq = xq.transpose(1, 2)  # (bs, n_local_heads, seqlen, head_dim)
        keys = keys.transpose(1, 2)
        values = values.transpose(1, 2)
        scores = torch.matmul(xq, keys.transpose(2, 3)) / math.sqrt(self.head_dim)
        if mask is not None:
            scores = scores + mask  # (bs, n_local_heads, seqlen, cache_len + seqlen)
        scores = F.softmax(scores.float(), dim=-1).type_as(xq)
        output = torch.matmul(scores, values)  # (bs, n_local_heads, seqlen, head_dim)
        output = output.transpose(1, 2).contiguous().view(bsz, seqlen, -1)
        return self.wo(output)

首要流程如下:

• 核算query、key、value的线性映射表明xq、xk、xv
• 对xq和xk使用方位参数
• 将xk、xv写入cache
• 从cache读取key和value,重复其head维度以匹配query的head数
• 核算query和key的点积获得相关度得分
• 对scores加mask并softmax归一化
• 将scores与value做权重和,得到多头自注意力输出
• 将多头输出拼接并线性映射,便是Self-Attention的成果

其间用到了一个函数repeat_kv，它的作用是将key和value的head维度重复n_rep次，以匹配query的head数。

def repeat_kv(x: torch.Tensor, n_rep: int) -> torch.Tensor:
    """torch.repeat_interleave(x, dim=2, repeats=n_rep)"""
    bs, slen, n_kv_heads, head_dim = x.shape
    if n_rep == 1:
        return x
    return (
        x[:, :, :, None, :]
        .expand(bs, slen, n_kv_heads, n_rep, head_dim)
        .reshape(bs, slen, n_kv_heads * n_rep, head_dim)
    )

repeat_kv函数运用 expand 办法将输入张量在第四个维度上扩展 n_rep 次，并运用 reshape 办法将其调整为恰当的形状。

LLaMA的Transformer结构

中心的自注意力模块完成了之后，咱们就能够像搭积木一样，将其组装成Transformer结构了。

首要咱们看看全衔接网络：

class FeedForward(nn.Module):
    def __init__(
        self,
        dim: int,
        hidden_dim: int,
        multiple_of: int,
        ffn_dim_multiplier: Optional[float],
    ):
        super().__init__()
        hidden_dim = int(2 * hidden_dim / 3)
        # custom dim factor multiplier
        if ffn_dim_multiplier is not None:
            hidden_dim = int(ffn_dim_multiplier * hidden_dim)
        hidden_dim = multiple_of * ((hidden_dim + multiple_of - 1) // multiple_of)
        self.w1 = ColumnParallelLinear(
            dim, hidden_dim, bias=False, gather_output=False, init_method=lambda x: x
        )
        self.w2 = RowParallelLinear(
            hidden_dim, dim, bias=False, input_is_parallel=True, init_method=lambda x: x
        )
        self.w3 = ColumnParallelLinear(
            dim, hidden_dim, bias=False, gather_output=False, init_method=lambda x: x
        )
    def forward(self, x):
        return self.w2(F.silu(self.w1(x)) * self.w3(x))

LLaMA的前馈神经网络首要是立足于并行化。

首要次序为:

• 初始化时构建了3个线性层w1,w2,w3。其间w1和w3运用ColumnParallelLinear完成散布式并行,w2运用RowParallelLinear。
• forward时,先过w1做第一次线性投影,然后运用SiLU激活函数。
• 跟一个w3对原输入做的线性投影加起来,完成残差衔接。
• 终究过w2线性层输出。

这样的结构形成了一个带残差衔接的两层前馈网络。它结合并行核算和残差衔接,使模型对长序列任务拟合作用更佳。

然后，咱们将前馈全衔接网络和之前讲的自注意力机制结合起来，构建Transformer块：

class TransformerBlock(nn.Module):
    def __init__(self, layer_id: int, args: ModelArgs):
        super().__init__()
        self.n_heads = args.n_heads
        self.dim = args.dim
        self.head_dim = args.dim // args.n_heads
        self.attention = Attention(args)
        self.feed_forward = FeedForward(
            dim=args.dim,
            hidden_dim=4 * args.dim,
            multiple_of=args.multiple_of,
            ffn_dim_multiplier=args.ffn_dim_multiplier,
        )
        self.layer_id = layer_id
        self.attention_norm = RMSNorm(args.dim, eps=args.norm_eps)
        self.ffn_norm = RMSNorm(args.dim, eps=args.norm_eps)
    def forward(
        self,
        x: torch.Tensor,
        start_pos: int,
        freqs_cis: torch.Tensor,
        mask: Optional[torch.Tensor],
    ):
        h = x + self.attention.forward(
            self.attention_norm(x), start_pos, freqs_cis, mask
        )
        out = h + self.feed_forward.forward(self.ffn_norm(h))
        return out

咱们来分块解说一下。

self.n_heads = args.n_heads
self.dim = args.dim
self.head_dim = args.dim // args.n_heads
self.attention = Attention(args)

首要，从参数目标中获取必要的参数，然后创立一个Attention目标。

self.feed_forward = FeedForward(
    dim=args.dim,
    hidden_dim=4 * args.dim,
    multiple_of=args.multiple_of,
    ffn_dim_multiplier=args.ffn_dim_multiplier,
)

然后，创立一个FeedForward目标，这个目标完成了前馈神经网络。

self.layer_id = layer_id
self.attention_norm = RMSNorm(args.dim, eps=args.norm_eps)
self.ffn_norm = RMSNorm(args.dim, eps=args.norm_eps)

接着，保存层的ID，并创立两个用于归一化的RMSNorm目标。

h = x + self.attention.forward(
    self.attention_norm(x), start_pos, freqs_cis, mask
)
out = h + self.feed_forward.forward(self.ffn_norm(h))
return out

终究，经过注意力机制和前馈神经网络，核算出输出数据。在注意力机制和前馈神经网络的前后，都运用了归一化操作，这有助于改善模型的练习安稳性。

终究，咱们将上面所有的集成在一起，构建出LLaMA的Transformer结构：

class Transformer(nn.Module):
    def __init__(self, params: ModelArgs):
        super().__init__()
        self.params = params
        self.vocab_size = params.vocab_size
        self.n_layers = params.n_layers
        self.tok_embeddings = ParallelEmbedding(
            params.vocab_size, params.dim, init_method=lambda x: x
        )
        self.layers = torch.nn.ModuleList()
        for layer_id in range(params.n_layers):
            self.layers.append(TransformerBlock(layer_id, params))
        self.norm = RMSNorm(params.dim, eps=params.norm_eps)
        self.output = ColumnParallelLinear(
            params.dim, params.vocab_size, bias=False, init_method=lambda x: x
        )
        self.freqs_cis = precompute_freqs_cis(
            self.params.dim // self.params.n_heads, self.params.max_seq_len * 2
        )
    @torch.inference_mode()
    def forward(self, tokens: torch.Tensor, start_pos: int):
        _bsz, seqlen = tokens.shape
        h = self.tok_embeddings(tokens)
        self.freqs_cis = self.freqs_cis.to(h.device)
        freqs_cis = self.freqs_cis[start_pos : start_pos + seqlen]
        mask = None
        if seqlen > 1:
            mask = torch.full(
                (1, 1, seqlen, seqlen), float("-inf"), device=tokens.device
            )
            mask = torch.triu(mask, diagonal=start_pos + 1).type_as(h)
        for layer in self.layers:
            h = layer(h, start_pos, freqs_cis, mask)
        h = self.norm(h)
        output = self.output(h).float()
        return output

到了大结局阶段，可解释的就不用了。 终究的模块唯一增加的组件便是词嵌入部分:

        self.tok_embeddings = ParallelEmbedding(
            params.vocab_size, params.dim, init_method=lambda x: x
        )

然后把Transformer块打包在一起：

        self.layers = torch.nn.ModuleList()
        for layer_id in range(params.n_layers):
            self.layers.append(TransformerBlock(layer_id, params))

加上归一化：

        self.norm = RMSNorm(params.dim, eps=params.norm_eps)
        self.output = ColumnParallelLinear(
            params.dim, params.vocab_size, bias=False, init_method=lambda x: x
        )

终究，所有层都走一遍，再来一遍归一化，功德圆满：

for layer in self.layers:
    h = layer(h, start_pos, freqs_cis, mask)
h = self.norm(h)
output = self.output(h).float()
return output

小结

至此，LLaMA2的首要代码咱们就走马观花地学习了一遍。哪怕有些细节还不能了解，最少咱们掌握了一个真实的大模型代码的地图。

大家有不了解的当地也不要紧。一方面，后面咱们会针对框架的通用技能再进行一些介绍。另一方面，咱们还要解析多个其它的开源大模型的源代码。突变引起突变，大家多思考，多实验，就一定能了解大模型的代码。