[NLP]LLaMA与LLamMA2解读

摘要

Meta最近提出了LLaMA(开放和高效的基础语言模型)模型参数包括从7B到65B等多个版本。最值得注意的是,LLaMA-13B的性能优于GPT-3,而体积却小了10倍以上,LLaMA-65B与Chinchilla-70B和PaLM-540B具有竞争性。

一、引言

一般而言,模型越大,效果越好。然而有文献指出[1],当给定计算量的预算之后,最好的performance,并不是最大的模型,而是在一个小模型上用更多的数据进行训练。针对给定的计算量预算,scaling laws可以计算如何选择数据量的大小和模型的大小。然而这忽略了inference的预算,而这一点在模型推理时非常关键。当给定一个模型performance目标之后,最好的模型不是训练最快的模型,而是推理最快的模型。尽管在这种情况下,训练一个更大的模型成本会更低。

文献[2]中推荐,训练一个 10B 的模型,需要 200B 的 tokens,而本文的实验发现,一个7B的模型,经过 1T tokens 训练之后,performance 仍然在增加。本文的目标在于,通过在超大规模的数据上训练,给出一系列可能最好 performance 的 LLM。

二、预训练数据

2.1 数据集

训练语料是混合的开源语料,中文占比很低,几乎不支持中文。详细占比为:CommonCrawl 67%, C4 15%, GitHub 4.5%, Wikipedia 4.5%, Books 4.5%, ArXiv 2.5%, Stack Exchange 2%.

一共有1.4T的tokens,大部分的训练数据都只用了一次,除了Wikipedia 和 Books 使用了大概2个epochs。

2.2 Tokenizer

使用byte pair encoding (BPE) 算法,使用的是Sentence-Piece的实现。所有数字被拆分为单独的digit,所有未知的UTF-8 字符,回退到字节来进行分解。因此,LLaMA 可以通过byte 的方式,构造出很多不在 vocab 中的字符,从而也具有较好的多语言能力。

三、网络结构改进

优化器

论文的模型使用AdamW优化器(Loshchilov和Hutter,2017)进行训练,具有以下超参数:

使用余弦学习率计划,使得最终学习率等于最大学习率的10%。论文使用0.1的权重衰减和1.0的梯度剪裁。使用2000个预热步骤,并随着模型的大小而改变学习率和批次大小。

使用了基于transformer的架构,并做了如下3点改进:

3.1 Pre-normalization

为了提高训练的稳定性,对每个transformer层的输入进行归一化,而不是输出进行归一化。

同时,使用 RMS Norm 归一化函数。RMS Norm 的全称为 Root Mean Square layer normalization。与 layer Norm 相比,RMS Norm的主要区别在于去掉了减去均值的部分,计算公式为:

 RMS Norm 的作者认为这种模式在简化了Layer Norm 的计算,可以在减少约 7%∼64% 的计算时间[3]

class LlamaRMSNorm(nn.Module):
    def __init__(self, hidden_size, eps=1e-6):
        """
        LlamaRMSNorm is equivalent to T5LayerNorm
        """
        super().__init__()
        self.weight = nn.Parameter(torch.ones(hidden_size))
        self.variance_epsilon = eps

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

        return (self.weight * hidden_states).to(input_dtype)

3.2 SwiGLU

使用SwiGLU替代了ReLU作为激活函数。和PaLM中不同,维度采用 234� 而不是 4� 。

SwiGLU 在论文[4] 中提出,相比于其他的激活函数变体,可以取得 log-perplexity 的最优值(和 GEGLU 并列)。

class LlamaMLP(nn.Module):
    def __init__(
        self,
        hidden_size: int,
        intermediate_size: int,
        hidden_act: str,
    ):
        super().__init__()
        self.gate_proj = nn.Linear(hidden_size, intermediate_size, bias=False)
        self.down_proj = nn.Linear(intermediate_size, hidden_size, bias=False)
        self.up_proj = nn.Linear(hidden_size, intermediate_size, bias=False)
        # config 中 hidden_act = 'silu'
        # 'silu' 和 'swish' 对应的激活函数均为:SiLUActivation 
        # https://github.com/huggingface/transformers/blob/717dadc6f36be9f50abc66adfd918f9b0e6e3502/src/transformers/activations.py#L229
        self.act_fn = ACT2FN[hidden_act]

    def forward(self, x):
        # 对应上述公式的 SwiGLU
        return self.down_proj(self.act_fn(self.gate_proj(x)) * self.up_proj(x))

从代码可以看到 LlamaMLP 中一共有 3 个 Linear 层,原因就在于 SwiGLU 激活函数比类似 ReLU 的激活函数,需要多一个 Linear 层进行门控。

3.3 RoPE

RoPE 的核心思想是“通过绝对位置编码的方式实现相对位置编码”,可以说是具备了绝对位置编码的方便性,同时可以表示不同 token 之间的相对位置关系。[5] 不同于原始 Transformers 论文中,将 pos embedding 和 token embedding 进行相加,RoPE 是将位置编码和 query (或者 key) 进行相乘。具体如下:

 

# 代码增加了注释,可以看到和原始公式的对应关系。
class LlamaRotaryEmbedding(torch.nn.Module):
    def __init__(self, dim, max_position_embeddings=2048, base=10000, device=None):
        super().__init__()
        # 此处 inv_freq 对应公式中的 theta
        inv_freq = 1.0 / (base ** (torch.arange(0, dim, 2).float().to(device) / dim))
        self.register_buffer("inv_freq", inv_freq)

        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 对应公式中的 m * theta, t 对应公式中的 m,表示位置
        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
        # 此处和原始公式不同,theta_0 和 theta_0 不再相邻
        # 而是分在向量的前半部分和后半部分
        emb = torch.cat((freqs, freqs), dim=-1)
        dtype = torch.get_default_dtype()
        self.register_buffer("cos_cached", emb.cos()[None, None, :, :].to(dtype), persistent=False)
        self.register_buffer("sin_cached", emb.sin()[None, None, :, :].to(dtype), persistent=False)

    def forward(self, x, seq_len=None):
        # x: [bs, num_attention_heads, seq_len, head_size]
        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, :, :].to(x.dtype), persistent=False)
            self.register_buffer("sin_cached", emb.sin()[None, None, :, :].to(x.dtype), persistent=False)
        # 大部分情况下,直接从这里返回
        return (
            self.cos_cached[:, :, :seq_len, ...].to(dtype=x.dtype),
            self.sin_cached[:, :, :seq_len, ...].to(dtype=x.dtype),
        )


def rotate_half(x):
    """Rotates half the hidden dims of the input."""
    # 此次和原始推导中不同,正负号不是间隔的,而是分前半部分和后半部分。但对于结果没有影响
    x1 = x[..., : x.shape[-1] // 2]
    x2 = x[..., x.shape[-1] // 2 :]
    return torch.cat((-x2, x1), dim=-1)


def apply_rotary_pos_emb(q, k, cos, sin, position_ids):
    # The first two dimensions of cos and sin are always 1, so we can `squeeze` them.
    cos = cos.squeeze(1).squeeze(0)  # [seq_len, dim]
    sin = sin.squeeze(1).squeeze(0)  # [seq_len, dim]
    cos = cos[position_ids].unsqueeze(1)  # [bs, 1, seq_len, dim]
    sin = sin[position_ids].unsqueeze(1)  # [bs, 1, seq_len, dim]
    # 对应上图中 RoPE 的简化计算
    q_embed = (q * cos) + (rotate_half(q) * sin)
    k_embed = (k * cos) + (rotate_half(k) * sin)
    return q_embed, k_embed

四、高效实现

加速训练:

  • 使用 随机多头注意力机制(causal multi-head attention) 提高模型的训练速度。该机制的实现借用了xformers库,它的思路是不存储注意力权重,不计算其中注意力得分
  • 手动实现了Transformer的激活函数,而没有用pytorch库的autograd,以得到更优的训练速度。同时使用了并行化技术提高训练速度。
  • 减少了activation checkpointing 中,重新计算 activation 的计算量。手动实现 transformer 层的反向传递函数,保存了计算成本高的 activations,例如线性层的输出。
  • 通过使用 model parallelism 和 sequence parallelism 来减少显存的使用量。
  • 尽可能地将 activations 的计算和GPU之间的通讯进行并行。

加速效果:

  • 65B的模型,在2048个80G的A100 GPU上,可以达到380 tokens/sec/GPU的速度。训练1.4T tokens需要21天。

五、主要结果与结论

LLaMA-13B 优于 GPT-3,尽管只有1/10大小。 LLaMA-65B 是可以与 Chinchilla-70B 和 PaLM-540B 这种最佳的LLM相竞争的模型。经过微调之后,LLaMA的效果有显著的提升。

未来打算发布在更大的语料上预训练上的更大的模型,因为随着数据和模型的增大,可以看到 performance 的稳定提升。

五、高效实现

LLaMA2的开源地址:https://github.com/facebookresearch/llama/blob/main/MODEL_CARD.md

LLaMA2的下载地址:https://ai.meta.com/resources/models-and-libraries/llama-downloads/

LLaMA2的官方博客地址:https://ai.meta.com/resources/models-and-libraries/llama/

5.1 技术细节记录

  1. 模型尺寸:LLaMA 2提供了三种不同的模型尺寸:7B、13B和70B。其中,7B和13B的架构与LLaMA 1相同,可直接用于商业应用。
  2. 训练:LLaMA 2模型经过了2万亿个标记的训练,其上下文长度是LLaMA 1的两倍。此外,LLaMA-2-chat模型还接受了超过100万个新的人类注释的训练。LLaMA 2的训练语料比LLaMA 1多出40%,上下文长度从2048增加到4096,使其能够理解和生成更长的文本。
  3. 预训练:LLaMA 2使用公开的在线数据进行预训练,然后通过监督微调创建LLaMA-2-chat的初始版本。接下来,LLaMA-2-chat使用人类反馈强化学习(RLHF)进行迭代细化,其中包括拒绝采样和近端策略优化(PPO)。
  4. 模型架构:LLaMA 2采用了LLaMA 1的大部分预训练设置和模型架构,使用标准Transformer架构,使用RMSNorm应用预归一化、使用SwiGLU激活函数和旋转位置嵌入RoPE。与LLaMA 1的主要架构差异包括增加了上下文长度和分组查询注意力(GQA)。
  5. 分组查询注意力(GQA):这是一个新的注意力机制,可以提高大模型的推理可扩展性。它的工作原理是将键和值投影在多个头之间共享,而不会大幅降低性能。可以使用具有单个KV投影的原始多查询格式(MQA)或具有8KV投影的分组查询注意力变体(GQA)。
  6. 超参数:使用AdamW优化器进行训练,其中β1=0.9,β2=0.95,eps=10−5。使用余弦学习率计划,预热2000步,衰减最终学习率降至峰值学习率的10%。使用0.1的权重衰减和1.0的梯度裁剪。
  7. 分词器:LLaMA 2使用与LLaMA 1相同的分词器;它采用字节对编码(BPE)算法,使用SentencePiece实现。与LLaMA 1一样,将所有数字拆分为单独的数字,并使用字节来分解未知的UTF-8字符。总数词汇量为32k个token。
  8. 微调:LLaMA 2-Chat是数月实验研究和对齐技术迭代应用的结果,包括指令微调和RLHF,需要大量的计算和数据标注资源。有监督微调指令数据质量非常重要,包括多样性,注重隐私安全不包含任何元用户数据。
  9. 安全性:该研究使用三个常用基准评估了Llama 2的安全性,针对三个关键维度:真实性,指语言模型是否会产生错误信息,采用TruthfulQA基准;毒性,指语言模型是否会产生「有毒」、粗鲁、有害的内容,采用ToxiGen基准;偏见,指语言模型是否会产生存在偏见的内容,采用BOLD基准。

5.2 llama1&2 论文对比

  • llama1论文详细介绍了训练数据的工作方式,读者可以复制其训练过程。
  • llama2论文主要集中在介绍模型训练的方法论,对于数据部分的透明度较低。
  • llama1由67%的公共爬虫数据和15%的精选数据组成。
  • 关于llama2的数据,论文中提供的信息较少,但meta提到llama2的预训练语料库比llama1大40%,可能与llama1的数据集类似。

5.3 数据质量的重要性

  • llama2加大了数据集中高可信度的权重,导致了模型的可信度更高;
  • 清理和整理出更好的数据是确保模型生成准确文本的关键步骤;
  • 随着对数据质量评估和去除重复数据等方法的进步,是否会继续增加训练数据量引发讨论;
  • 预训练数据量和Chinchilla Optimal的关系:在训练过程中,llama2预训练数据的量超过了Chinchilla Optimal所设定的限制,并且在训练结束时损失曲线仍在下降。这引发了对Chinchilla Optimal是否仍然适用的疑问——

参考

Training Compute-Optimal Large Language Models https://arxiv.org/abs/2203.15556

Root Mean Square Layer Normalization https://arxiv.org/pdf/1910.07467.pdf

GLU Variants Improve Transformer https://arxiv.org/pdf/2002.05202.pdf

Transformer升级之路:2、博采众长的旋转式位置编码 - 科学空间|Scientific Spaces

Llama2|得数据者得天下——Latent Space上LLaMA2讨论的精彩整理 - 知乎 (zhihu.com)

LLM/Meta的LLAMA-2__简介及论文翻译 - 知乎 (zhihu.com)

LLaMA2 RLHF 技术细节 - 知乎 (zhihu.com)

本图文内容来源于网友网络收集整理提供,作为学习参考使用,版权属于原作者。
THE END
分享
二维码
< <上一篇
下一篇>>