本文介绍来自 Meta AI 的 LLaMa 模型，类似于 OPT，也是一种完全开源的大语言模型。LLaMa 的参数量级从 7B 到 65B 大小不等，是在数万亿个 token 上面训练得到。值得一提的是，LLaMa 虽然只使用公共的数据集，依然取得了强悍的性能。LLaMA-13B 在大多数基准测试中都优于 GPT-3 (175B)，LLaMA65B 与最佳模型 Chinchilla-70B 和 PaLM-540B 相比具有竞争力。
1 LLaMa：开源高效的大语言模型
论文名称：LLaMA: Open and Efficient Foundation Language Models
论文地址：
https://arxiv.org/pdf/2302.13971.pdf
代码链接：
https://github.com/facebookresearch/llama
1.1 背景：模型参数量级的积累，或者训练数据的增加，哪个对性能提升帮助更大？
以 GPT-3 为代表的大语言模型 (Large language models, LLMs) 在海量文本集合上训练，展示出了惊人的涌现能力以及零样本迁移和少样本学习能力。GPT-3 把模型的量级缩放到了 175B，也使得后面的研究工作继续去放大语言模型的量级。大家好像有一个共识，就是：模型参数量级的增加就会带来同样的性能提升。
但是事实确实如此吗？
最近的 "Training Compute-Optimal Large Language Models[1]" 这篇论文提出一种缩放定律 (Scaling Law)：
训练大语言模型时，在计算成本达到最优情况下，模型大小和训练数据 (token) 的数量应该比例相等地缩放，即：如果模型的大小加倍，那么训练数据的数量也应该加倍。
翻译过来就是：当我们给定特定的计算成本预算的前提下，语言模型的最佳性能不仅仅可以通过设计较大的模型搭配小一点的数据集得到，也可以通过设计较小的模型配合大量的数据集得到。
那么，相似成本训练 LLM，是大 LLM 配小数据训练，还是小 LLM 配大数据训练更好？
缩放定律 (Scaling Law) 告诉我们对于给定的特定的计算成本预算，如何去匹配最优的模型和数据的大小。但是本文作者团队认为，这个功能只考虑了总体的计算成本，忽略了推理时候的成本。因为大部分社区用户其实没有训练 LLM 的资源，他们更多的是拿着训好的 LLM 来推理。在这种情况下，我们首选的模型应该不是训练最快的，而应该是推理最快的 LLM。呼应上题，本文认为答案就是：小 LLM 配大数据训练更好，因为小 LLM 推理更友好。
1.2 LLaMa 做到了什么
LLaMa 沿着小 LLM 配大数据训练的指导思想，训练了一系列性能强悍的语言模型，参数量从 7B 到 65B。例如，LLaMA-13B 比 GPT-3 小10倍，但是在大多数基准测试中都优于 GPT-3。大一点的 65B 的 LLaMa 模型也和 Chinchilla 或者 PaLM-540B 的性能相当。
同时，LLaMa 模型只使用了公开数据集，开源之后可以复现。但是大多数现有的模型都依赖于不公开或未记录的数据完成训练。
1.3 LLaMa 预训练数据
LLaMa 预训练数据大约包含 1.4T tokens，对于绝大部分的训练数据，在训练期间模型只见到过1次，Wikipedia 和 Books 这两个数据集见过2次。
如下图1所示是 LLaMa 预训练数据的含量和分布，其中包含了 CommonCrawl 和 Books 等不同域的数据。
图1：LLaMa 预训练数据的含量和分布
CommonCrawl (占 67%)： 包含 2017 到 2020 的5个版本，预处理部分包含：删除重复数据，去除掉非英文的数据，并通过一个 n-gram 语言模型过滤掉低质量内容。
C4 (占 15%)： 在探索性实验中，作者观察到使用不同的预处理 CommonCrawl 数据集可以提高性能，因此在预训练数据集中加了 C4。预处理部分包含：删除重复数据，过滤的方法有一些不同，主要依赖于启发式方法，例如标点符号的存在或网页中的单词和句子的数量。
Github (占 4.5%)： 在 Github 中，作者只保留在 Apache、BSD 和 MIT 许可下的项目。此外，作者使用基于行长或字母数字字符比例的启发式方法过滤低质量文件，并使用正则表达式删除标题。最后使用重复数据删除。
Wikipedia (占 4.5%)： 作者添加了 2022 年 6-8 月的 Wikipedia 数据集，包括 20 种语言，作者处理数据以删除超链接、评论和其他格式样板。
Gutenberg and Books3 (占 4.5%)： 作者添加了两个书的数据集，分别是 Gutenberg 以及 ThePile (训练 LLM 的常用公开数据集) 中的 Book3 部分。处理数据时作者执行重复数据删除，删除内容重叠超过 90% 的书籍。
ArXiv (占 2.5%)： 为了添加一些科学数据集，作者处理了 arXiv Latex 文件。作者删除了第一部分之前的所有内容，以及参考文献。还删除了 .tex 文件的评论，以及用户编写的内联扩展定义和宏，以增加论文之间的一致性。
Stack Exchange (占 2%)： 作者添加了 Stack Exchange，这是一个涵盖各种领域的高质量问题和答案网站，范围从计算机科学到化学。作者从 28 个最大的网站保留数据，从文本中删除 HTML 标签并按分数对答案进行排序。
Tokenizer 的做法基于 SentencePieceProcessor[2]，使用 bytepair encoding (BPE) 算法。
LLaMa 的 PyTorch 代码如下，用到了 sentencepiece 这个库。
class Tokenizer:
def __init__(self, model_path: str):
# reload tokenizer
assert os.path.isfile(model_path), model_path
self.sp_model = SentencePieceProcessor(model_file=model_path)
logger.info(f"Reloaded SentencePiece model from {model_path}")
# BOS / EOS token IDs
self.n_words: int = self.sp_model.vocab_size()
self.bos_id: int = self.sp_model.bos_id()
self.eos_id: int = self.sp_model.eos_id()
self.pad_id: int = self.sp_model.pad_id()
logger.info(
f"#words: {self.n_words} - BOS ID: {self.bos_id} - EOS ID: {self.eos_id}"
)
assert self.sp_model.vocab_size() == self.sp_model.get_piece_size()
def encode(self, s: str, bos: bool, eos: bool) -> List[int]:
assert type(s) is str
t = self.sp_model.encode(s)
if bos:
t = [self.bos_id] + t
if eos:
t = t + [self.eos_id]
return t
def decode(self, t: List[int]) -> str:
return self.sp_model.decode(t)1.4 LLaMa 模型架构
Pre-normalization [受 GPT3 的启发]：
为了提高训练稳定性，LLaMa 对每个 Transformer 的子层的输入进行归一化，而不是对输出进行归一化。使用 RMSNorm[3] 归一化函数。
class RMSNorm(torch.nn.Module):
def __init__(self, dim: int, eps: float = 1e-6):
super().__init__()
self.eps = eps
self.weight = nn.Parameter(torch.ones(dim))
def _norm(self, x):
return x * torch.rsqrt(x.pow(2).mean(-1, keepdim=True) + self.eps)
def forward(self, x):
output = self._norm(x.float()).type_as(x)
return output * self.weight
常规的 Layer Normalization：
式中, gi和 是bi LN 的 scale 和 shift 参数, u和 
的计算如下式所示:
RMSNorm：
相当于是去掉了 u这一项。
看上去就这一点小小的改动，有什么作用呢？RMSNorm 的原始论文进行了一些不变性的分析和梯度上的分析。
SwiGLU 激活函数 [受 PaLM 的启发]：
LLaMa 使用 SwiGLU 激活函数[4]替换 ReLU 非线性以提高性能，维度从4d 变为 
。
Rotary Embeddings [受 GPTNeo 的启发]：
LLaMa 去掉了绝对位置编码，使用旋转位置编码 (Rotary Positional Embeddings, RoPE)[5]，这里的 RoPE 来自苏剑林老师，其原理略微复杂，感兴趣的读者可以参考苏神的原始论文和官方博客介绍：
https://spaces.ac.cn/archives/8265
Self-Attention 的 PyTorch 代码：
class Attention(nn.Module):
def __init__(self, args: ModelArgs):
super().__init__()
self.n_local_heads = args.n_heads // fs_init.get_model_parallel_world_size()
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,
args.n_heads * self.head_dim,
bias=False,
gather_output=False,
init_method=lambda x: x,
)
self.wv = ColumnParallelLinear(
args.dim,
args.n_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_heads, self.head_dim)
).cuda()
self.cache_v = torch.zeros(
(args.max_batch_size, args.max_seq_len, self.n_local_heads, self.head_dim)
).cuda()
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_heads, self.head_dim)
xv = xv.view(bsz, seqlen, self.n_local_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]
xq = xq.transpose(1, 2)
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, slen, cache_len + slen)
scores = F.softmax(scores.float(), dim=-1).type_as(xq)
output = torch.matmul(scores, values) # (bs, n_local_heads, slen, head_dim)
output = output.transpose(
1, 2
).contiguous().view(bsz, seqlen, -1)
return self.wo(output)这里有几个地方值得注意一下：
首先是 model.py 文件里面从 fairscale 中 import 了3个类，分别是：ParallelEmbedding，RowParallelLinear，和 ColumnParallelLinear。
Fairscale 链接如下，是一个用于高性能大规模预训练的库，LLaMa 使用了其 ParallelEmbedding 去替换 Embedding， 使用了其 RowParallelLinear 和 ColumnParallelLinear 去替换 nn.Linear，猜测可能是为了加速吧。
