Llama开源代码详细解读（3）

expand_mask模块

def _expand_mask(mask: torch.Tensor, dtype: torch.dtype, tgt_len: Optional[int] = None):
    """
    Expands attention_mask from `[bsz, seq_len]` to `[bsz, 1, tgt_seq_len, src_seq_len]`.
    """
    bsz, src_len = mask.size()
    tgt_len = tgt_len if tgt_len is not None else src_len

    expanded_mask = mask[:, None, None, :].expand(bsz, 1, tgt_len, src_len).to(dtype)

    inverted_mask = 1.0 - expanded_mask

    return inverted_mask.masked_fill(inverted_mask.to(torch.bool), torch.finfo(dtype).min)

这个函数接收一个张量mask，并将其变换为特定的形状。输入三个参数分别为：mask：大小为[bsz, seq_len]。dtype：数据类型。tgt_len：目标序列长度。以下是函数的运行方式。

获取mask参数

  bsz, src_len = mask.size()

• .size()函数获取了mask张量的行数、列数，即bsz，src_len。

确定目标序列长度

tgt_len = tgt_len if tgt_len is not None else src_len

• 如果tgt_len没有被指定，则赋值为src_len。

扩展掩码

expanded_mask = mask[:, None, None, :].expand(bsz, 1, tgt_len, src_len).to(dtype)

• mask[:, None, None, :]将mask的维度从[bsz, seq_len]拓展为[bsz,1,1,seq_len]，expand(bsz, 1, tgt_len, src_len)将拓展后的矩阵继续拓展为[bsz,1,tgt_len,src_len),to(dtype)转换为指定的数据类型。

生成反转掩码

inverted_mask = 1.0 - expanded_mask

将掩码中0和1的位置互换。

填充反转掩码

return inverted_mask.masked_fill(inverted_mask.to(torch.bool), torch.finfo(dtype).min)

inverted_mask.to(torch.bool)将反转掩码转换为布尔类型

masked_fill(inverted_mask.to(torch.bool), torch.finfo(dtype).min)用dtype这个类型的最小值填充其中为true的位置。

返回一个经过填充处理的反转掩码张量，形状为 [bsz, 1, tgt_len, src_len]，数据类型为 dtype。

RMSNorm归一化模块

class LlamaRMSNorm(nn.Module):
    def __init__(self, hidden_size, eps=1e-6):
        # eps一个很小的数，用于避免除零错误
        """
        LlamaRMSNorm is equivalent to T5LayerNorm
        """
        super().__init__()
        self.weight = nn.Parameter(torch.ones(hidden_size))
        # weight是一个可训练的参数，初始化为一个大小为 hidden_size 的全 1 张量
        self.variance_epsilon = eps

    def forward(self, hidden_states):
        input_dtype = hidden_states.dtype
        # 存储hidden_states的数据类型
        hidden_states = hidden_states.to(torch.float32)
        variance = hidden_states.pow(2).mean(-1, keepdim=True)
        # hidden_states一般是batch size * sequence length * hidden size，这里的mean是按最后一维取平均
        hidden_states = hidden_states * torch.rsqrt(variance + self.variance_epsilon)
        return (self.weight * hidden_states).to(input_dtype)

这段代码定义了一个RMSNorm类，用于实现归一化。

构造函数__init__()

def __init__(self, hidden_size, eps=1e-6):
    super().__init__()
    self.weight = nn.Parameter(torch.ones(hidden_size))
    self.variance_epsilon = eps

• hidden_size表示隐藏层的维度大小，eps表示一个很小的数，用于防止除零错误，self.weight表示为一个可训练的权重参数，初始化为一个和hidden_size大小一致的全1张量。

前向传播forward()

def forward(self, hidden_states):
    input_dtype = hidden_states.dtype
    hidden_states = hidden_states.to(torch.float32)
    variance = hidden_states.pow(2).mean(-1, keepdim=True)
    hidden_states = hidden_states * torch.rsqrt(variance + self.variance_epsilon)
    return (self.weight * hidden_states).to(input_dtype)

• input_dtype储存隐藏层的数据类型
• hidden_states = hidden_states.to(torch.float32)将数据类型转换为bf32，确保计算的稳定性
• variance = hidden_states.pow(2).mean(-1, keepdim=True)计算隐藏层在最后一层的方差，具体步骤是对每个元素平方再取平均。
• hidden_states = hidden_states * torch.rsqrt(variance + self.variance_epsilon)计算反向标准差（1/标准差），并与hidden_state元素逐个相乘，从而实现归一化。
• return (self.weight * hidden_states).to(input_dtype)将归一化后的隐藏层乘以权重，并转换回输入的数据类型。