接上文..

https://github.com/facebookresearch/fairscale

另一个需要注意的点是：cache 的缓存机制，可以看到在构造函数里面定义了下面两个东西：
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()

关键其实就是这几行代码：
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]

在训练的时候，因为每次都是输入完整的一句话，所以 cache 机制其实是不发挥作用的。
在推理的时候，比如要生成 "I have a cat"，过程是：
1 输入 <s>，生成 <s> I。
2 输入 <s> I，生成 <s> I have。
3 输入 <s> I have，生成 <s> I have a。
4 输入 <s> I have a，生成 <s> I have a cat。

在执行3这一步时，计算 "a" 的信息时，还要计算 <s> I have 的 Attention 信息，比较复杂。因此，cache 的作用就是在执行2这一步时，提前把 <s> I have 的 keys 和 values 算好，并保存在 self.cache_k 和 self.cache_v 中。在执行3这一步时，计算 Attention 所需的 keys 和 values 是直接从这里面取出来的：
keys = self.cache_k[:bsz, : start_pos + seqlen]
values = self.cache_v[:bsz, : start_pos + seqlen]
只需要额外地计算 "a" 的 keys 和 values 即可，这对模型的快速推理是至关重要的。

还有一个值得注意的点：self.cache_k = self.cache_k.to(xq)
这里使用的是 to() 函数的一种不太常见的用法：torch.to(other, non_blocking=False, copy=False)→Tensor
Returns a Tensor with same torch.dtype and torch.device as the Tensor other.

FFN 的 PyTorch 代码：

class FeedForward(nn.Module):
    def __init__(
        self,
        dim: int,
        hidden_dim: int,
        multiple_of: int,
    ):
        super().__init__()
        hidden_dim = int(2 * hidden_dim / 3)
        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))

这里需要注意的点是：
激活函数用的是 F.silu()，也就是 Swish 激活函数。
self.w2(F.silu(self.w1(x)) * self.w3(x)) 的实现也就是 SwiGLU 激活函数。

图2：silu 激活函数

Transformer Block 的 PyTorch 代码：

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
        )
        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

Transformer 的 PyTorch 代码：

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[:, -1, :])  # only compute last logits
        return output.float()

self.tok_embeddings 用的是 ParallelEmbedding 这个函数，把 ids 变为词向量。
mask 部分通过 torch.full() 函数和 torch.triu() 函数得到一个上三角矩阵，用于注意力的计算。
通过 torch.nn.ModuleList() 函数定义所有的 Transformer Block。
所有的 norm 函数都使用 RMSNorm 去定义。

生成过程的 PyTorch 代码：

class LLaMA:
    def __init__(self, model: Transformer, tokenizer: Tokenizer):
        self.model = model
        self.tokenizer = tokenizer

    def generate(
        self,
        prompts: List[str],
        max_gen_len: int,
        temperature: float = 0.8,
        top_p: float = 0.95,
    ) -> List[str]:
        bsz = len(prompts)
        params = self.model.params
        assert bsz <= params.max_batch_size, (bsz, params.max_batch_size)

        prompt_tokens = [self.tokenizer.encode(x, bos=True, eos=False) for x in prompts]

        min_prompt_size = min([len(t) for t in prompt_tokens])
        max_prompt_size = max([len(t) for t in prompt_tokens])

        total_len = min(params.max_seq_len, max_gen_len + max_prompt_size)

        tokens = torch.full((bsz, total_len), self.tokenizer.pad_id).cuda().long()
        for k, t in enumerate(prompt_tokens):
            tokens[k, : len(t)] = torch.tensor(t).long()
        input_text_mask = tokens != self.tokenizer.pad_id
        start_pos = min_prompt_size
        prev_pos = 0
        for cur_pos in range(start_pos, total_len):
            logits = self.model.forward(tokens[:, prev_pos:cur_pos], prev_pos)
            if temperature > 0:
                probs = torch.softmax(logits / temperature, dim=-1)
                next_token = sample_top_p(probs, top_p)
            else:
                next_token = torch.argmax(logits, dim=-1)
            next_token = next_token.reshape(-1)
            # only replace token if prompt has already been generated
            next_token = torch.where(
                input_text_mask[:, cur_pos], tokens[:, cur_pos], next_token
            )
            tokens[:, cur_pos] = next_token
            prev_pos = cur_pos

        decoded = []
        for i, t in enumerate(tokens.tolist()):
            # cut to max gen len
            t = t[: len(prompt_tokens[i]) + max_gen_len]
            # cut to eos tok if any
            try:
                t = t[: t.index(self.tokenizer.eos_id)]
            except ValueError:
                pass
            decoded.append(self.tokenizer.decode(t))
        return decoded


def sample_top_p(probs, p):
    probs_sort, probs_idx = torch.sort(probs, dim=-1, descending=True)
    probs_sum = torch.cumsum(probs_sort, dim=-1)
    mask = probs_sum - probs_sort > p
    probs_sort[mask] = 0.0
    probs_sort.div_(probs_sort.sum(dim=-1, keepdim=True))
    next_token = torch.multinomial(probs_sort, num_samples=1)
    next_token = torch.gather(probs_idx, -1, next_token)
    return next_token

这里需要注意的是：
torch.multinomial() 函数用于按照一定的概率 (probs_sort) 采样一定数量 (num_samples) 的 Tensor。
torch.gather() 函数是一个抽数据的函数，按照 probs_idx 的索引和 dim=-1 的维度。

1.5 LLaMa 的优化

AdamW, , 使用 cosine 学习率衰减策略, 2000 步的 warm-up, 最终学习率等于最大学习率的10% , 使用 0.1 的权重衰减和 1.0 的梯度裁剪。

1.6 LLaMa 的高效实现

快速的注意力机制： LLaMa 采用了高效的 causal multi-head attention (基于 xformers[6])，不存储注意力权重，且不计算 mask 掉的 query 和 key 的值。

手动实现反向传播过程，不使用 PyTorch autograd： 使用 checkpointing 技术减少反向传播中的激活值的计算，更准确地说，LLaMa 保存计算代价较高的激活值，例如线性层的输出。

通过使用模型和序列并行减少模型的内存使用。此外，LLaMa 还尽可能多地重叠激活的计算和网络上的 GPU 之间的通信。

LLaMa-65B 的模型使用 2048 块 80G 的 A100 GPU，在 1.4T token 的数据集上训练 21 天。

1.7 LLaMa 实验结果

LLaMa 在 20 个标准的 Zero-Shot 和 Few-Shot 任务上面做了评测。在评测时的任务包括自由形式的生成任务和多项选择任务。多项选择任务的目标是根据提供的上下文在一组给定选项中选择最合适的答案。

Zero-Shot 在评测时，作者提供了任务和测试示例的文本描述。LLaMa 要么使用开放式生成提供答案，要么对给定的答案进行排名。Few-Shot 在评测时，作者提供了任务的几个示例 (在 1 到 64 之间) 和一个测试示例。LLaMa 将此文本作为输入并生成答案或者排名不同的选项。

1.7.1 常识推理实验结果

作者考虑了8个标准的常识推理基准：BoolQ, PIQA, SIQA, WinoGrande 等，采用标准的 Zero-Shot 的设定进行评估。结果如图3所示，LLaMA-65B 在除了 BoolQ 的所有基准测试中都优于 Chinchilla-70B，在除了 BoolQ 和 WinoGrande 的任何地方都超过了 PaLM540B。LLAMA-13B 模型在大多数基准测试中也优于 GPT-3。

图3：常识推理实验结果

1.7.2 封闭式问答实验结果

如下图3和4所示是封闭式问答实验结果，图4是 Natural Questions 数据集，图5是 TriviaQA 数据集，报告的是报告精确匹配性能，即：模型无法访问包含回答问题证据的文档。在这两个基准测试中，LLaMA-65B 在零样本和少样本设置中实现了最先进的性能，而且 LLaMa-13B 的性能也同样具备竞争力。

图4：Natural Questions 封闭式问答实验结果

图5：TriviaQA 封闭式问答实验结果

1.7.3 阅读理解实验结果

阅读理解任务在 RACE 数据集上做评测，结果如图6所示。LLaMA-65B 与 PaLM-540B 具有竞争力，LLaMA-13B 的性能比 GPT-3 好几个百分点。

图6：阅读理解实验结果

1.7.4 数学推理实验结果

作者在 MATH 和 GSM8k 两个任务上面做数学推理任务，MATH 是一个 12K 中学和高中数学问题的数据集，用 LaTeX 编写。GSM8k 是一组中学数学问题。在 GSM8k 上，尽管 LLaMA-65B 从没在数学数据上进行微调，但可以观察到 LLaMA-65B 优于 Minerva-62B。

图7：数学推理实验结果

1.7.5 代码生成实验结果

作者在 HumanEval 和 MBPP 两个任务上面做代码生成任务，对于这两个任务，模型接收几个句子中的程序描述，以及一些输入输出示例。模型需要生成一个符合描述并满足测试用例的 Python 程序。图7将 LLaMa 与尚未在代码上微调的现有语言模型 (PaLM 和 LaMDA) 进行比较，PaLM 和 LLAMA 在包含相似数量代码标记的数据集上进行训练。对于相似数量的参数，LLaMa 优于其他通用模型，例如 LaMDA 和 PaLM，这些模型没有专门针对代码进行训练或微调。具有 13B 参数的 LLAMA，在 HumanEval 和 MBPP 上都优于 LaMDA 137B。LLaMA 65B 也超过了训练时间更长的 PaLM 62B。

1.7.6 大规模多任务语言理解实验结果

MMLU 大规模多任务语言理解基准由涵盖各种知识领域的多项选择题组成，包括人文、STEM 和社会科学。作者使用基准提供的示例在 5-shot 设置中评估我们的模型，结果如图7所示。可以观察到 LLaMa-65B 在大多数领域平均落后于 Chinchilla70B 和 PaLM-540B 几个百分点。一个潜在的解释是，LLaMa 在预训练数据中只使用了有限数量的书籍和学术论文，即 ArXiv、Gutenberg 和 Books3，总计只有 177GB，而其他的模型训练了多达 2TB 的书籍。

作者还发现加入一些微调指令也能够提升 大规模多任务语言理解的性能。尽管 LLaMA-65B 的非微调版本已经能够遵循基本指令，但可以观察到非常少量的微调提高了 MMLU 的性能，并进一步提高了模型遵循指令的能力。

如下图8所示，尽管这里使用的指令微调方法很简单，但在 MMLU 上达到了 68.9%。LLAMA-I (65B) 优于 MMLU 现有中等大小的指令微调模型，但仍远未达到最先进的水平。

图8：大规模多任务语言理解实验结果

1.8 训练期间的性能变化

如下图9所示是 7B、13B、33B 和 65B 这几个模型在一些问答和常识基准的表现随着 training token 的变化，图10是 7B、13B、33B 和 65B 这几个模型的 training loss 随着 training token 的变化。在大多数基准测试中，性能稳步提高，并且与模型的训练困惑度相关。

图9：7B、13B、33B 和 65B 这几个模型在一些问答和常识基准的表现随着 training token 的变化

图10：7B、13B、33B 和 65B 这几个模型的 training loss 随着 training token 的变化