用code去探索理解Llama架构的简单又实用的方法

news2024/11/25 0:45:01

除了白月光我们也需要朱砂痣

我最近也在反思，可能有时候算法和论文也不是每个读者都爱看，我也会在今后的文章中加点code或者debug模型的内容，也许还有一些好玩的应用demo，会提升这部分在文章类型中的比例

今天带着大家通过代码角度看一下Llama,或者说看一下Casual-LLM的Transfomer到底长啥样

对Transfomer架构需要更了解的读者，可以先看这个系列

小周带你读论文-2之"草履虫都能看懂的Transformer老活儿新整"Attention is all you need(1) (qq.com)

小周带你读论文-2之"草履虫都能看懂的Transformer老活儿新整"Attention is all you need(2) (qq.com)

小周带你读论文-2之"草履虫都能看懂的Transformer老活儿新整"Attention is all you need(3) (qq.com)

小周带你读论文-2之"草履虫都能看懂的Transformer老活儿新整"Attention is all you need(4) (qq.com)

友情提示，看代码和debug都不需要GPU，在你的PC上就可以做

首先先安装transfomer库

pip install transfomers

然后进入到库下面，一般在这

进去就能找到Transfomers的库，往下拉到models，可以发现各种模型都在里面

找到Llama，包含以下文件

点进modeling_llama.py发现1200多行，根本没法看（一般人没那么多耐心，但是其实仔细看两遍还是很有收获的），然后点击左边outline 大纲，（或者ctrl+shift+o），就可以有选择的看你想要研究的具体网络层,这样感觉压力瞬间小了百分之80以上

我想查某个函数的代码块，想对它加深了解，举个例子，看起来比较怪异命名的，这个线性扩展RoPE embedding的函数

然后ctl按住函数名字就能查到它作用于attetion的机制里面，在这种情况下即使我不知道它到底干啥，也能猜个89不离10，至少和什么主模块相关我清楚了

本章的话，我们先从模型主体部分看起，点LlamaModel，就能看到非常清晰的逻辑

主体部分包含3个子模块:

先要embedding token
再要包含一个通过for循环，不断的持续经过的decoder层
还要包含一个Normal（RMSNorm）

从大面上看也就这么3个操作

初始化部分，初始了哪些模块我们看完了，再看看细节，从forward看起（大家看任何网络都要重点看forward）

forward输入部分，要求输入的参数:

input_ids：输入序列标记的索引，形状为(batch_size, sequence_length)的torch.LongTensor。
attention_mask：避免对填充标记进行注意力计算的掩码，形状为(batch_size, sequence_length)的torch.Tensor
position_ids：输入序列标记在位置嵌入中的索引，形状为(batch_size, sequence_length)的torch.LongTensor
past_key_values：包含预先计算的隐藏状态（自注意力块和交叉注意力块中的键和值），用于加速顺序解码的tuple，推理用的，当use_cache=True时，会返回这个参数，训练不用管
inputs_embeds：直接传入embedding表示而不是input_ids，形状为(batch_size, sequence_length, hidden_size)的torch.FloatTensor
use_cache：是否使用缓存加速解码的布尔值，当设置为True时，past_key_values的键值状态将被返回，用于加速解码
output_attentions：是否返回所有注意力层的注意力张量，布尔值，当设置为True时，会在返回的结果中包含注意力张量
output_hidden_states：是否返回所有层的隐藏状态，布尔值，当设置为True时，会在返回的结果中包含隐藏状态
return_dict：是否返回utils.ModelOutput对象而不是普通的元组，布尔值，当设置为True时，会返回一个ModelOutput对象

我们继续看，下面就是一些操作步骤，输入的input_id，会被向量化，生成hidden_states

hidden_states然后就被扔进了若干个hidden_layer被for循环来回的操作，比如Llama7B的32层

我们简单写一段逻辑描述上述的代码

比如在把"我爱你"已经分词的情况下我=100，爱=200，你=300

input_ids = [100,200,300]

input_ids -> nn.Emebdding(dims=3) -> hidden_states

hidden_states = [[0.1,0.2,0.3],[0.4,0.5,0.6],[0.7.0.8.1.1]]

hidden_states ->layer1 ->layer2 ------>layer32

最终还是hidden_states

也就是最终的shape和初始的hidden_states的shape的相同的

hidden_states -> Norm -> hidden_states

其实没必要把hidden_states理解的那么悬，就当它是个中间变量就可以了

Layer里面都有什么呢？我们点进去Layer里面就能看到

包含了我们在Transfomer里学到的attention层，MLP层,LayerNorm这些层

继续往下

我们可以看到首先定义了残差，然后hidden_states在没过layer之前就先被Normal了一下（这个知识点以前也讲过，Llama RMS LN是前LN）,然后过attetion，过MLP，最后hidden_states=residual+hidden_states, 这样就把位置编码啥的也都带过来了

然后我们点击进入attetion这块，看看它是咋做的