llama-2大模型结构

llama-2大模型结构

llama大模型是Meta公司开源的大语言模型，目前最新的为llama3.1参数量达到了405B，非常恐怖

以7B模型为例，梳理llama-2的模型结构，重点关注其对传统Transformer模型的结构修改

LlamaConfig {
"_name_or_path": "llama2/llama-2-7b",
"architectures": [
"LlamaForCausalLM"
],
"attention_bias": false,
"bos_token_id": 1,
"eos_token_id": 2,
"hidden_act": "silu",
"hidden_size": 4096,
"initializer_range": 0.02,
"intermediate_size": 11008,
"max_position_embeddings": 2048,
"model_type": "llama",
"num_attention_heads": 32,
"num_hidden_layers": 32,
"num_key_value_heads": 32,
"pad_token_id": 0,
"pretraining_tp": 1,
"rms_norm_eps": 1e-05,
"rope_scaling": null,
"rope_theta": 10000.0,
"tie_word_embeddings": false,
"torch_dtype": "bfloat16",
"transformers_version": "4.31.0",
"use_cache": true,
"vocab_size": 32000
}

从上述config文件可以看到，

• llama2 采用的激活函数是 silu
• 词表大小为 32000
• 预训练参数的数据类型是 bfloat16

Pytorch定义好模型后，print(model) 的输出：

LlamaForCausalLM(
(model): LlamaModel(
(embed_tokens): Embedding(32000, 4096, padding_idx=0)
(layers): ModuleList(
(0–31): 32 x LlamaDecoderLayer(
(self_attn): LlamaAttention(
(q_proj): Linear(in_features=4096, out_features=4096, bias=False)
(k_proj): Linear(in_features=4096, out_features=4096, bias=False)
(v_proj): Linear(in_features=4096, out_features=4096, bias=False)
(o_proj): Linear(in_features=4096, out_features=4096, bias=False)
(rotary_emb): LlamaRotaryEmbedding()
)
(mlp): LlamaMLP(
(gate_proj): Linear(in_features=4096, out_features=11008, bias=False)
(up_proj): Linear(in_features=4096, out_features=11008, bias=False)
(down_proj): Linear(in_features=11008, out_features=4096, bias=False)
(act_fn): SiLUActivation()
)
(input_layernorm): LlamaRMSNorm()
(post_attention_layernorm): LlamaRMSNorm()
)
)
(norm): LlamaRMSNorm()
)
(lm_head): Linear(in_features=4096, out_features=32000, bias=False)
)

可以看到 LLama-2 7b 是由 32 个 LLamaDecoderLayer 堆叠而成；

LLamaDecoderLayer 结构图： LLama中采用了 RMSNorm 而非传统的 LayerNorm；

LLamaAttention结构图： LLama将传统Transformer中Transformer Block前的 position embedding 放入到了 Transformer Block中；

并且采用了 LlamaRotaryEmbedding ；

LLamaMLP结构： 对MLP结果也做了修改；

传统的MLP只包含两个FC层，第一个FC层后有激活函数，第一个FC层起到了隐式地 gating 作用；

而，LLamaMLP中引入了显式地 gate_proj 层，其后接激活函数，其结果与 up_proj 进行 element-wise 乘法，起到 gating 的作用；