Minimind的推理过程学习记录

[RyanSunn]

Minimind的推理过程学习记录

大佬的minimind真是让我受益匪浅，由于之前没接触过LLM，和DDP等，这次学到的知识真不少，尤其是在推理时感受到minimind极大的学习价值，所以在此对推理过程进行一个详细的记录。因为我是刚接触，所以以下记录的应该会有错误，若有看到，请各位大佬指正。再次感谢jingyaogong大佬。在记录推理过程前，我先对DDP的学习进行一些记录。

一、DDP的使用记录

由于我有两张NVIDIA Tesla P100所以就想用分布式训练，但是按照jingyaogong大佬的分布式训练命令就报错，于是对DDP的使用进行了学习。
当然为了学习网络结构，对于1-pretrain.py我是先在单卡上进行了debug，然后在进行的DDP分布式训练。debug的时候我将df = df.sample(frac=1)改为df = df.sample(frac=0.005)进行起来更方便。
最后我也是没有完成整个训练过程，下载了大佬训练好的模型。
以下是我遇到的一些问题。

因为我用的NVIDIA显卡，所以我把dist.init_process_group(backend="gloo")改为dist.init_process_group(backend="nccl")。
然后发现①Windows系统不能下载nccl。然后我通过查询我尝试安装Docker，然后安装时显示因为wsl和系统的问题，又失败了。
最后发现②NVIDIA显卡也可以使用dist.init_process_group(backend="gloo")。
之后运行torchrun --nproc_per_node 2 1-pretrain.py依旧一直报错，经过查询后改为torchrun --standalon --nproc_per_node 2 1-pretrain.py就好了，也就是指定单机训练就好了。
以上问题记录可能显得我比较笨，还请担待。

二、推理过程

这部分来记录一下对于0-eval_pretrain.py或3-full_sft.py的推理过程，两者特别相似，以下以0-eval_pretrain.py为例，并下载大佬的minimind-v1模型权重进行推理。因为我也是刚接触LLM相关知识，所以说的可能比较琐碎，还请见谅。

1、修改模型配置

因为我选用minimind-v1而不是minimind-v1-small，所以首先在LMConfig中最如下修改：

2、加载参数

temperature和top_k在使用的时候自然就会一目了然。

3、加载模型

模型加载自然要将加载路径换成自己下载好的权重参数的路径：

4、回答方式

answer_way = 0与下方图片对应，自然也可改为1,手动输入问题。

5、准备输入数据

添加起始标记，转为token，转为指定设备上的tensor。

6、生成函数解读

接下来看到如下函数，以为是将x变量直接传进来进行推理

但不是的，跳转到到这个函数位置看看

首先对这个函数的参数做个解读：

①无梯度计算：使用 @torch.inference_mode() 装饰器，表示在推理模式下运行，不会计算梯度，从而节省内存和计算资源。

②重复惩罚：通过 rp 参数对已经生成的 token 进行惩罚，防止生成重复的内容。

③温度采样：通过 temperature 参数控制生成的随机性。温度越高，生成的文本越随机；温度为 0 时，选择概率最高的 token。

④Top-k 采样：通过 top_k 参数限制每次采样时考虑的 token 数量，只从概率最高的 top_k 个 token 中采样。

⑤流式输出：通过 stream 参数控制是否以流式方式输出生成的文本片段。

⑥缓存机制：通过 kv_cache 参数控制是否使用缓存机制，以提高生成效率。

⑦生成终止条件：生成过程会在达到最大 token 数量 max_new_tokens 或遇到结束标记 eos 时终止。

其次这个函数并没有return，而是使用yield说明这个函数是一个生成器函数：

生成器在调用时不会立即执行，而是返回一个生成器对象。每次调用生成器对象的 next() 方法时，生成器函数会执行到下一个 yield 语句并返回其值，对应上方第6节第一个图片的函数初始化和y = next(res_y)。

更须注意的是生成器函数在每次暂停时会保存其执行状态，包括局部变量、指令指针等。下次继续执行时会从上次暂停的地方继续，这使得生成器函数可以在多次调用之间保持状态。简单来说就是，上次在yield idx[:, index:]返回了值，下次进入这个函数还用yield idx[:, index:]继续向下执行。

7、第一次推理过程

开始第一次的y = next(res_y)进入生成函数，并执行前两行，有如下效果：

继续执行，跳入下图所示位置：

这里的self也就是这个Transformer模型，传入self也就是下次会跳转到这个模型的forward函数，得到的inference_res就是模型的预测结果，也就是下一个token的分类情况。

取出logits这个结果后进行重复惩罚，rp=1，所以这里并无影响。

如下图所示，temperature为0时，选出概率最大的类，temperature不为0时，进行logits的放大，并选出前top_k个概率最大的类。

将这些类除外的logits所对应的值都设备负无穷，这样下面进行softmax的时候除了最大的top_k个类处的值都是0，可注意下方的softmax的公式和简易图像。
$$
Softmax(zi)=exp(zi)/Σexp(zj)
$$

最后是在这个top_k中随机算一个，作为结果。

稍后在判断是否是结束符后与原输入tokens进行了一个拼接，并通过yield idx[:, index:]返回。

8、返回后输出

上述返回后进入以下程序：

在这里进行了预测结果——“下一个字”的打印

在这段代码中，主要进行了以下判断和操作：

①初始化 history_idx:history_idx 被初始化为 0，用于跟踪已经打印的答案的索引。
②循环处理生成的答案:使用 while y != None: 循环来处理生成的答案，直到 y 为 None。
③解码生成的答案:answer = tokenizer.decode(y[0].tolist()) 将生成的张量 y 解码为字符串 answer。
④处理不完整的字符:如果 answer 的最后一个字符是 �（表示不完整的字符），则尝试获取下一个生成的 y，并继续循环。
⑤检查答案长度:如果 answer 的长度为零，则尝试获取下一个生成的 y，并继续循环。
⑥打印答案:打印从 history_idx 开始的 answer，并刷新输出缓冲区。
⑦获取下一个生成的 y:在每次打印后，尝试获取下一个生成的 y。
⑧更新 history_idx:更新 history_idx 为当前 answer 的长度，以便在下次打印时从正确的位置开始。
⑨检查 stream 标志:如果 stream 为 False，则在打印一次后退出循环。
⑩打印换行:在循环结束后，打印一个换行符以分隔不同的回答。

可见在下一次的推理生成也就是⑦处，出来之后依旧按照这个循环进行打印，如此反复就可以打印整个词语接龙的句子。

9、第n次推理（n>=2）

截止到上一节，整个推理过程已经比较完整了，按理没有其他的了，但是这里还是要补充一下，因为如果我们第二次进入y = next(res_y)就会发现情况有所不同。第二次进入生成函数，根据生成器函数的特性，从yield idx[:, index:]继续执行，直接进入到while idx.shape[1] < max_new_tokens - 1:循环这一行。

但我们继续执行就会发现，由于第一次if init_inference or not kv_cache:这一行下面init_inference = False的赋值，会进入以下位置，并注意到这次传入的并不是整个tokens，而是最后一个，也就是上次生成的token，并且注意到kv_cache为True，也可以会想到当时训练的时候这个值一直为False。同时也注意到了current_idx的传入。

继续进入这个forward过程来看看，进入之后对current_idx从参数进行了赋值，代表当前token的位置，然后在下面的位置编码处只取了当前token位置的位置编码。
然后进行embeding等直到进入TransformerBlock部分。
到这里我们还有个疑惑就是：无论训练还是第一次的推理都是传入的整个句子的tokens，这次只传入最后一个token，也只有一个位置编码，可行吗？

继续进入到TransformerBlock，没发现特殊情况，然后在进入attention，在进行正常的xq、xk、xv的生成和加入位置编码后进入了以下图片中蓝色位置，这个位置之前并没有进入过。

到这里看到self.k_cache和self.v_cache的shape都是torch.Size([1, 11, 8, 48])，再看一下上个图片的记录的current_idx=11，也就是第十二个位置，所以这里shape中的11代表的好像是之前的tokens的数量。
没错，因为上次虽然没有进入到蓝色行位置，但是因为kv_cache为True也进入到了黄色框位置，这里记录了上次的xk和xv
到这里一切都说得通了：拿这次的一个token产生的xk、xv和之前的拼接就得到了整个句子的xk、xv，然后就可以和以前一样可以继续进行下去了。
因此kv_cache通过缓存键和值，可以避免在每次推理时重新计算整个序列的键和值，尤其是在序列长度较长时，这种优化可以节省大量计算资源。
此外，这里的xq还只是代表一个token的向量，亦可以和xk和xv进行计算，在下图位置进行计算后正常输出即可。

到此，整个过程结束。

[jingyaogong]

很不错的记录，谢谢！

[RyanSunn]

竟然放在这里，无比感激，受益匪浅，感谢^ω^

[FangKQ]

感谢大佬细致的讲解，尤其是KVcache，学到了！