LLMs源码阅读(二)Baichuan

LLMs源码阅读(二)Baichuan

sliderSun:灵魂拷问之word2vec

sliderSun:关于CNN、RNN、LSTM、Transformer、BERT参数计算的那些疑问

sliderSun:关于Transformer的那些个为什么

sliderSun:关于BERT中的那些为什么

sliderSun:Transformer、Like-Bert、对比学习、ChatGPT相关面试集锦

sliderSun:知识盛宴:探秘LLMs、Sora和LWM的神奇世界

sliderSun:LLMs源码阅读(一)ChatGLM

sliderSun:LLMs源码阅读(二)Baichuan

sliderSun:LLMs源码阅读之(三)LLaMA

sliderSun:LLMs源码阅读之(四)Mistral系列

sliderSun:LangChain:代码世界的魔法师,源码解读带你笑看技术黑洞

一、Baichuan

纯llama结构,有自己搞的NormHead,设置了padding,llama是没有padding的

  • 位置编码:7B的位置编码采用RoPE,13B位置编码采用ALiBi。主要是因为两种位置编码对模型效果基本没影响,所以继承了Baichuan1的7B和13B的位置编码。
  • 激活函数:采用SwiGLU激活函数,不同于传统FFN的2个矩阵,SwiGLU有三个矩阵,因此缩小了隐藏层维度,由原来的4倍变成8/3倍,再调整为128的整数。
  • 归一化:对Transformer的输入进行采用层归一化,提高warm-up的鲁棒性,并用RMSNorm实现。
  • NormHead:为了提高模型训练的稳定性,对输出的embedding进行归一化,主要解决:(1)稀有标记的embedding在训练过程中变小,干扰了训练的动态;(2)分析发现输出表示语义关系受余弦相似度计算明显,对L2距离不明显,归一化可以减少线性分类器通过点积计算logits时,L2距离的影响。
  • 去掉dropout和bias
  • xformer,加速attention
  • Max-z loss:预训练时,logits可能会变的非常大,会导致推理过程中对惩罚因子不敏感,受NormSoftmax启发,对logits进行归约。主要有助于稳定训练并使推理对超参数更具鲁棒性。

L_{max-z}=2e^{-4}*z^{2} \\

BaiChuanConfig


添加图片注释,不超过 140 字(可选)

{
  "_from_model_config": true,
  "architectures": [
    "BaichuanForCausalLM"
  ],
  "auto_map": {
    "AutoConfig": "configuration_baichuan.BaichuanConfig",
    "AutoModelForCausalLM": "modeling_baichuan.BaichuanForCausalLM"
  },
  "tokenizer_class": "BaichuanTokenizer",
  "bos_token_id": 1,
  "eos_token_id": 2,
  "gradient_checkpointing": false,
  "hidden_act": "silu",
  "hidden_size": 5120,
  "initializer_range": 0.02,
  "intermediate_size": 13696,
  "model_max_length": 4096,
  "model_type": "baichuan",
  "num_attention_heads": 40,
  "num_hidden_layers": 40,
  "pad_token_id": 0,
  "rms_norm_eps": 1e-06,
  "tie_word_embeddings": false,
  "torch_dtype": "bfloat16",
  "transformers_version": "4.29.2",
  "use_cache": true,
  "vocab_size": 125696
}

tokenization_baichuan

这个代码文件里定义了一个名为BaichuanTokenizer的分词器类,它继承自PreTrainedTokenizer,用于处理自然语言文本,将文本转换为模型可以理解的形式,或者将模型的输出转换回人类可理解的文本。

类变量:

vocab_files_names,pretrained_vocab_files_map,max_model_input_sizes:存储词汇文件名,预训练词汇文件映射和最大模型输入大小的字典。model_input_names:存储模型输入名称的列表。

分词和转换相关的方法: _tokenize:将文本分词。 _convert_token_to_id 和 _convert_id_to_token:分别用于将token转换为id和将id转换回token。convert_tokens_to_string:将token序列转换回字符串。

BaichuanPreTrainedModel

class BaichuanPreTrainedModel(PreTrainedModel):  # 定义一个名为 "BaichuanPreTrainedModel" 的类,继承了 "PreTrainedModel" 类
    config_class = BaichuanConfig  # 定义一个类属性,将BaichuanConfig赋值给config_class,用于指定该模型的配置类是什么

    base_model_prefix = "model"  # 定义一个类属性,其值是字符串 "model",通常用于标识模型的主要子模块或基础模型

    supports_gradient_checkpointing = True  # 定义一个类属性,标识该模型支持梯度检查点功能,可以节省显存但会使计算稍慢

    _no_split_modules = ["BaichuanLayer"]  # 定义一个类属性,包含模型内不应该被拆分为子模型的模块名称

    _keys_to_ignore_on_load_unexpected = [r"decoder\.version"]  # 定义一个类属性,列出在加载模型时应该忽略的意外键名

    def _init_weights(self, module):  # 定义一个方法用于初始化模型权重
        std = self.config.initializer_range  # 从模型的配置中获取权重初始化的标准差
        
        if isinstance(module, torch.nn.Linear):  # 如果传入的模块是线性层
            module.weight.data.normal_(mean=0.0, std=std)  # 使用正态分布初始化线性层的权重,均值为0,标准差为std
            if module.bias is not None:  # 如果线性层有偏置
                module.bias.data.zero_()  # 使用0来初始化偏置

        elif isinstance(module, torch.nn.Embedding):  # 如果传入的模块是嵌入层
            module.weight.data.normal_(mean=0.0, std=std)  # 使用正态分布初始化嵌入层的权重,均值为0,标准差为std
            if module.padding_idx is not None:  # 如果嵌入层有填充索引
                module.weight.data[module.padding_idx].zero_()  # 将填充索引对应的嵌入向量初始化为0

    def _set_gradient_checkpointing(self, module, value=False):  # 定义一个方法用于设置模块的梯度检查点
        if isinstance(module, BaichuanModel):  # 如果传入的模块是BaichuanModel类型
            module.gradient_checkpointing = value  # 设置模块的梯度检查点属性为给定的value值

BaichuanModel

初始化

    def __init__(self, config: BaichuanConfig):  # 定义类的构造函数,接受一个名为 "config" 的参数,该参数是一个BaichuanConfig对象
        super().__init__(config)  # 调用父类 (BaichuanPreTrainedModel) 的构造函数并传入config参数
        self.padding_idx = config.pad_token_id  # 从config中获取pad_token_id并设置为类的属性
        self.vocab_size = config.vocab_size  # 从config中获取词汇表的大小并设置为类的属性
        self.n_head = config.num_attention_heads  # 从config中获取注意力头的数量并设置为类的属性
        self.embed_tokens = torch.nn.Embedding(  # 创建一个嵌入层
            config.vocab_size, config.hidden_size, self.padding_idx  # 指定词汇表大小、嵌入大小和填充索引
        )
        self.layers = torch.nn.ModuleList(  # 创建一个模块列表
            [BaichuanLayer(config) for _ in range(config.num_hidden_layers)]  # 根据隐藏层的数量多次实例化BaichuanLayer模块
        )
        self.norm = RMSNorm(config.hidden_size, epsilon=config.rms_norm_eps)  # 创建一个RMSNorm正则化层

        self.gradient_checkpointing = config.gradient_checkpointing  # 设置梯度检查点属性
        self.post_init()  # 调用post_init方法,可能在父类中定义,用于进一步的初始化
        self.max_cache_pos = config.model_max_length  # 设置模型的最大缓存位置
        self.first_run = True  # 设置一个标志,表示模型是否是第一次运行
        self.alibi_mask = None  # 初始化一个alibi_mask属性,值为None

    def get_input_embeddings(self):  # 定义一个方法用于获取输入的嵌入
        return self.embed_tokens  # 返回嵌入层

    def set_input_embeddings(self, value):  # 定义一个方法用于设置输入的嵌入
        self.embed_tokens = value  # 将传入的嵌入赋值给嵌入层属性

    def get_alibi_mask(self, tensor, seq_length_with_past):  # 定义一个方法用于获取alibi mask
        if self.training:  # 如果模型处于训练模式
            slopes = torch.Tensor(_get_interleave(self.n_head))  # 获取交错值
            position_point = (
                torch.arange(seq_length_with_past) - seq_length_with_past + 1  # 计算位置点
            )
            position_point = (
                position_point.unsqueeze(0)
                .unsqueeze(0)
                .expand(self.n_head, seq_length_with_past, -1)  # 调整位置点的形状并扩展
            )
            diag = torch.diag(position_point[0])  # 获取对角线
            position_point = position_point - diag.unsqueeze(0).unsqueeze(0).transpose(
                -1, -2
            )  # 调整位置点的形状
            alibi = slopes.unsqueeze(1).unsqueeze(1) * position_point  # 计算alibi值
            mask = _buffered_future_mask(  # 获取未来的mask
                tensor, seq_length_with_past, alibi, self.n_head
            )
        else:  # 如果模型处于评估模式
            if self.first_run:  # 如果是第一次运行
                self.first_run = False  # 设置标志为False
                self.register_buffer(  # 注册一个缓冲区
                    "future_mask",
                    _gen_alibi_mask(tensor, self.n_head, self.max_cache_pos).to(
                        tensor
                    ),
                    persistent=False,  # 使缓冲区不持久
                )
            if seq_length_with_past > self.max_cache_pos:  # 如果当前序列长度超过最大缓存位置
                self.max_cache_pos = seq_length_with_past  # 更新最大缓存位置
                self.register_buffer(  # 再次注册一个缓冲区
                    "future_mask",
                    _gen_alibi_mask(tensor, self.n_head, self.max_cache_pos).to(
                        tensor
                    ),
                    persistent=False,  # 使缓冲区不持久
                )
            mask = self.future_mask[
                : self.n_head, :seq_length_with_past, :seq_length_with_past
            ]  # 获取未来的mask
        return mask  # 返回mask

forward

    def forward(
        self,
        input_ids: torch.LongTensor = None,
        attention_mask: Optional[torch.Tensor] = None,
        past_key_values: Optional[List[torch.FloatTensor]] = None,
        inputs_embeds: Optional[torch.FloatTensor] = None,
        use_cache: Optional[bool] = False,
        output_attentions: Optional[bool] = False,
        output_hidden_states: Optional[bool] = False,
        return_dict: Optional[bool] = True,
    ) -> Union[Tuple, BaseModelOutputWithPast]:
        if input_ids is not None and inputs_embeds is not None:
            raise ValueError(
                "You cannot provide both input_ids and inputs_embeds simultaneously"
            )
        elif input_ids is not None:
            batch_size, seq_length = input_ids.shape
        elif inputs_embeds is not None:
            batch_size, seq_length, _ = inputs_embeds.shape
        else:
            raise ValueError("You need to provide input_ids or inputs_embeds")

        return_dict = (
            return_dict if return_dict is not None else self.config.use_return_dict
        )

        seq_length_with_past = seq_length

        if past_key_values is not None:
            past_key_values_length = past_key_values[0][0].shape[2]
            seq_length_with_past = seq_length_with_past + past_key_values_length

        if inputs_embeds is None:
            inputs_embeds = self.embed_tokens(input_ids)

        if self.training:
            if (
                self.alibi_mask is None
                or self.alibi_mask.shape[-1] != seq_length_with_past
            ):
                self.alibi_mask = self.get_alibi_mask(
                    inputs_embeds, seq_length_with_past
                )
            alibi_mask = self.alibi_mask
        else:
            alibi_mask = self.get_alibi_mask(inputs_embeds, seq_length_with_past)

        if attention_mask is not None:
            if len(attention_mask.shape) == 2:
                expanded_mask = attention_mask.to(alibi_mask.dtype)
                expanded_mask = torch.tril(
                    torch.gt(expanded_mask[:, :, None] * expanded_mask[:, None, :], 0)
                ) * torch.eq(expanded_mask[:, :, None] - expanded_mask[:, None, :], 0)
            else:
                expanded_mask = attention_mask
            bsz = inputs_embeds.size(0)
            src_len, tgt_len = alibi_mask.size()[-2:]
            expanded_mask = (
                expanded_mask.unsqueeze(1)
                .expand(bsz, 1, src_len, tgt_len)
                .to(alibi_mask.dtype)
            )
            inverted_mask = 1.0 - expanded_mask
            inverted_mask = inverted_mask.masked_fill(
                inverted_mask.to(torch.bool), torch.finfo(alibi_mask.dtype).min
            )
            attention_mask = inverted_mask + alibi_mask.unsqueeze(0)
        else:
            attention_mask = alibi_mask

        hidden_states = inputs_embeds

        if self.gradient_checkpointing and self.training:
            if use_cache:
                logger.warning_once(
                    "`use_cache=True` is incompatible with gradient checkpointing. Setting `use_cache=False`..."
                )
                use_cache = False

        # decoder layers
        all_hidden_states = () if output_hidden_states else None
        all_self_attns = () if output_attentions else None
        next_decoder_cache = () if use_cache else None

        for idx, decoder_layer in enumerate(self.layers):
            if output_hidden_states:
                all_hidden_states += (hidden_states,)

            past_key_value = (
                past_key_values[idx] if past_key_values is not None else None
            )

            if self.gradient_checkpointing and self.training:

                def create_custom_forward(module):
                    def custom_forward(*inputs):
                        # None for past_key_value
                        return module(*inputs, output_attentions, None)

                    return custom_forward

                layer_outputs = torch.utils.checkpoint.checkpoint(
                    create_custom_forward(decoder_layer),
                    hidden_states,
                    attention_mask,
                    None,
                )
            else:
                layer_outputs = decoder_layer(
                    hidden_states,
                    attention_mask=attention_mask,
                    past_key_value=past_key_value,
                    output_attentions=output_attentions,
                    use_cache=use_cache,
                )

            hidden_states = layer_outputs[0]

            if use_cache:
                next_decoder_cache += (layer_outputs[2 if output_attentions else 1],)

            if output_attentions:
                all_self_attns += (layer_outputs[1],)

        hidden_states = self.norm(hidden_states)

        # add hidden states from the last decoder layer
        if output_hidden_states:
            all_hidden_states += (hidden_states,)

        next_cache = next_decoder_cache if use_cache else None
        if not return_dict:
            return tuple(
                v
                for v in [hidden_states, next_cache, all_hidden_states, all_self_attns]
                if v is not None
            )
        return BaseModelOutputWithPast(
            last_hidden_state=hidden_states,
            past_key_values=next_cache,
            hidden_states=all_hidden_states,
            attentions=all_self_attns,
        )

BaichuanLayer

class BaichuanLayer(torch.nn.Module):  # 定义一个名为 "BaichuanLayer" 的 PyTorch 模型类,它继承了torch.nn.Module

    def __init__(self, config: BaichuanConfig):  # 构造函数接收一个名为config的BaichuanConfig类型参数
        super().__init__()  # 调用父类(torch.nn.Module)的初始化方法
        self.hidden_size = config.hidden_size  # 从config中提取hidden_size并赋值给类变量self.hidden_size
        self.self_attn = BaichuanAttention(config=config)  # 用config初始化BaichuanAttention对象,并赋值给self.self_attn
        self.mlp = MLP(  # 初始化一个MLP对象,并赋值给self.mlp
            hidden_size=self.hidden_size,  # MLP的hidden_size参数等于我们之前设置的self.hidden_size
            intermediate_size=config.intermediate_size,  # 从config中提取intermediate_size作为MLP的参数
            hidden_act=config.hidden_act,  # 从config中提取hidden_act作为MLP的参数
        )
        self.input_layernorm = RMSNorm(config.hidden_size, epsilon=config.rms_norm_eps)  # 初始化一个RMSNorm对象,并赋值给self.input_layernorm
        self.post_attention_layernorm = RMSNorm(  # 初始化另一个RMSNorm对象,并赋值给self.post_attention_layernorm
            config.hidden_size, epsilon=config.rms_norm_eps
        )

    def forward(  # 定义前向传播函数
        self,
        hidden_states: torch.Tensor,  # 输入参数为一个tensor,代表隐藏状态
        attention_mask: Optional[torch.Tensor] = None,  # 可选的attention_mask参数,默认为None
        past_key_value: Optional[Tuple[torch.Tensor]] = None,  # 可选的past_key_value参数,默认为None
        output_attentions: Optional[bool] = False,  # 可选的output_attentions参数,默认为False
        use_cache: Optional[bool] = False,  # 可选的use_cache参数,默认为False
    ) -> Tuple[  # 函数的返回类型为一个Tuple
        torch.FloatTensor, Optional[Tuple[torch.FloatTensor, torch.FloatTensor]]
    ]:
        residual = hidden_states  # 将输入的hidden_states保存为residual以便后面使用

        hidden_states = self.input_layernorm(hidden_states)  # 将hidden_states通过self.input_layernorm进行处理

        # Self Attention部分
        hidden_states, self_attn_weights, present_key_value = self.self_attn(  # 使用self.self_attn处理hidden_states,并返回三个结果
            hidden_states=hidden_states,
            attention_mask=attention_mask,
            past_key_value=past_key_value,
            output_attentions=output_attentions,
            use_cache=use_cache,
        )
        hidden_states = residual + hidden_states  # 将处理后的hidden_states与原始的residual进行加和,实现residual connection

        # Fully Connected部分
        residual = hidden_states  # 更新residual为处理后的hidden_states
        hidden_states = self.post_attention_layernorm(hidden_states)  # 将hidden_states通过self.post_attention_layernorm进行处理
        hidden_states = self.mlp(hidden_states)  # 将hidden_states通过self.mlp进行处理
        hidden_states = residual + hidden_states  # 再次使用residual connection

        outputs = (hidden_states,)  # 将处理后的hidden_states放入outputs tuple中

        if use_cache:  # 如果use_cache为True
            outputs += (present_key_value,)  # 将present_key_value也加入到outputs tuple中

        return outputs  # 返回outputs tuple

BaichuanAttention

Positional Embeddings 用的是 RoPE(Rotary Positional Embedding)(7B)与 ALiBi(13B)。

# 定义 BaichuanAttention 类,是一个 PyTorch 神经网络模块
class BaichuanAttention(torch.nn.Module):
    
    # 构造函数,接收 BaichuanConfig 配置类实例
    def __init__(self, config: BaichuanConfig):
        super().__init__()  # 调用父类的构造函数
        self.config = config  # 保存传入的配置
        self.hidden_size = config.hidden_size  # 从配置中取出隐藏层大小
        self.num_heads = config.num_attention_heads  # 从配置中取出注意力头数
        self.head_dim = self.hidden_size // self.num_heads  # 计算每个注意力头的维度
        self.max_position_embeddings = config.model_max_length  # 从配置中取出模型的最大长度
        
        # 确保 hidden_size 可以被 num_heads 整除
        if (self.head_dim * self.num_heads) != self.hidden_size:
            raise ValueError(
                f"hidden_size {self.hidden_size} is not divisible by num_heads {self.num_heads}"
            )
        
        # 定义一个线性层,用于获取查询、键和值
        self.W_pack = torch.nn.Linear(
            self.hidden_size, 3 * self.hidden_size, bias=False
        )
        
        # 定义一个输出线性层
        self.o_proj = torch.nn.Linear(
            self.num_heads * self.head_dim, self.hidden_size, bias=False
        )
        self.rotary_emb = RotaryEmbedding(self.head_dim, max_position_embeddings=self.max_position_embeddings)
        self.cos, self.sin = None, None
    
    # 辅助函数,用于重新整形张量以适应注意力计算
    def _shape(self, tensor: torch.Tensor, seq_len: int, bsz: int):
        return (
            tensor.view(bsz, seq_len, self.num_heads, self.head_dim)
            .transpose(1, 2)
            .contiguous()
        )

    # 前向传播函数
    def forward(
        self,
        hidden_states: torch.Tensor,
        attention_mask: Optional[torch.Tensor] = None,
        past_key_value: Optional[Tuple[torch.Tensor]] = None,
        output_attentions: bool = False,
        use_cache: bool = False,
    ) -> Tuple[torch.Tensor, Optional[torch.Tensor], Optional[Tuple[torch.Tensor]]]:
        
        bsz, q_len, _ = hidden_states.size()  # 获取输入的批次大小、序列长度和隐藏层大小

        # 使用线性变换获取查询、键和值
        proj = self.W_pack(hidden_states)
        proj = (
            proj.unflatten(-1, (3, self.hidden_size))
            .unsqueeze(0)
            .transpose(0, -2)
            .squeeze(-2)
        )
        query_states = (
            proj[0].view(bsz, q_len, self.num_heads, self.head_dim).transpose(1, 2)
        )
        key_states = (
            proj[1].view(bsz, q_len, self.num_heads, self.head_dim).transpose(1, 2)
        )
        value_states = (
            proj[2].view(bsz, q_len, self.num_heads, self.head_dim).transpose(1, 2)
        )

        # 如果给定了过去的键值对,则连接它们
        if past_key_value is not None:
            key_states = torch.cat([past_key_value[0], key_states], dim=2)
            value_states = torch.cat([past_key_value[1], value_states], dim=2)

        # 根据 use_cache 的值来决定是否保存键和值
        past_key_value = (key_states, value_states) if use_cache else None

        # 判断是否有 xformers,并根据条件使用不同的注意力计算方式
        if xops is not None and self.training:
            attn_weights = None
            with torch.backends.cuda.sdp_kernel(enable_flash=True, enable_math=True, enable_mem_efficient=True):
                attn_output = F.scaled_dot_product_attention(query_states, key_states, value_states, attn_mask=attention_mask)
            attn_output = attn_output.transpose(1, 2)
        else:
            attn_weights = torch.matmul(
                query_states, key_states.transpose(2, 3)
            ) / math.sqrt(self.head_dim)

            # 对注意力权重应用掩码
            if attention_mask is not None:
                if q_len == 1:  # 缓存中的推断
                    if len(attention_mask.size()) == 4:
                        attention_mask = attention_mask[:, :, -1:, :]
                    else:
                        attention_mask = attention_mask[:, -1:, :]
                attn_weights = attn_weights + attention_mask
                attn_weights = torch.max(
                    attn_weights, torch.tensor(torch.finfo(attn_weights.dtype).min)
                )
            
            # 获取注意力权重并计算注意力输出
            attn_weights = torch.nn.functional.softmax(attn_weights, dim=-1)  # 对最后一个维度进行 softmax,得到注意力权重
            attn_output = torch.matmul(attn_weights, value_states)  # 使用权重对 value 进行加权求和得到注意力输出

            attn_output = attn_output.transpose(1, 2)  # 调换维度,以便之后的处理

        attn_output = attn_output.reshape(bsz, q_len, self.hidden_size)  # 重新整形,以得到最终输出的形状
        attn_output = self.o_proj(attn_output)  # 通过输出的线性层

        # 如果不输出注意力权重,则设置为 None
        if not output_attentions:
            attn_weights = None

        # 返回注意力输出,注意力权重(如果需要)和过去的键值对(如果使用缓存)
        return attn_output, attn_weights, past_key_value

ROPE

\theta_i 的选择上,同样沿用了 Sinusoidal 位置编码的方案,即 \theta_i=10000^{-2i/d},它可以带来一定的远程衰减性。

baichuan-7B 源码中通过以下方式(下面代码注释中会引用,简称 rotation公式)实现旋转位置编码(代码中 d=128):

$$\begin{bmatrix} q_0\\ \vdots \\ q_{d/2-1}\\ q_{d/2}\\ \vdots\\ q_{d-1}\\ \end{bmatrix} \otimes \begin{bmatrix} cosn\theta_0\\ \vdots \\ cosn\theta_{d-2}\\ cosn\theta_0\\ \vdots \\ cosn\theta_{d-2}\\ \end{bmatrix} + \begin{bmatrix} -q_{d/2}\\ \vdots \\ -q_{d-1}\\ q_{0}\\ \vdots\\ q_{d/2-1}\\ \end{bmatrix} \otimes \begin{bmatrix} sinn\theta_0\\ \vdots \\ sinn\theta_{d-2}\\ sinn\theta_0\\ \vdots \\ sinn\theta_{d-2}\\ \end{bmatrix}$$


源代码(只含关键逻辑):

class RotaryEmbedding(torch.nn.Module):
    def __init__(self, dim, max_position_embeddings=2048, base=10000, device=None):
        # 初始化 theta, inv_freq.shape: [dim/2]
        inv_freq = 1. / (base ** (torch.arange(0, dim, 2).float() / dim))
        self.register_buffer("inv_freq", inv_freq)

        # Build here to make `torch.jit.trace` work.
        self.max_seq_len_cached = max_position_embeddings
        # t.shape: [max_position_embeddings]
        t = torch.arange(self.max_seq_len_cached, device=self.inv_freq.device, dtype=self.inv_freq.dtype)
       
        # freqs.shape: [max_position_embeddings, dim/2]
        freqs = torch.einsum("i,j->ij", t, self.inv_freq)
       
        # 根据 rotation公式,[max_position_embeddings, dim/2] -> [max_position_embeddings, dim] 
        emb = torch.cat((freqs, freqs), dim=-1)
        self.register_buffer("cos_cached", emb.cos()[None, None, :, :], persistent=False)
        self.register_buffer("sin_cached", emb.sin()[None, None, :, :], persistent=False)

    def forward(self, x, seq_dim=1, seq_len=None): 
         return (
            # [1, 1, seq_len, dim]
            self.cos_cached[:, :, :seq_len, ...].to(dtype=x.dtype),
            self.sin_cached[:, :, :seq_len, ...].to(dtype=x.dtype),
        )


def rotate_half(x):    
    x1 = x[..., : x.shape[-1] // 2]
    x2 = x[..., x.shape[-1] // 2:]
    # 根据 rotation公式, 前半部分与后半部分互换,且后半部分取负
    return torch.cat((-x2, x1), dim=-1)


def apply_rotary_pos_emb_index(q, k, cos, sin, position_id):
    cos = cos.squeeze(1).squeeze(0)  # [seq_len, dim]
    sin = sin.squeeze(1).squeeze(0)  # [seq_len, dim]
    # 融入 token 的位置信息
    cos = cos[position_ids].unsqueeze(1)  # [bs, 1, seq_len, dim]
    sin = sin[position_ids].unsqueeze(1)  # [bs, 1, seq_len, dim]  

    # 计算旋转位置编码,参考 rotation公式
    q_embed = (q * cos) + (rotate_half(q) * sin)
    k_embed = (k * cos) + (rotate_half(k) * sin)
    return q_embed, k_embed


class Attention(torch.nn.Module):
    def __init__(self, hidden_size, num_attention_heads, position_encoding_2d=True):
        self.hidden_size = hidden_size
        self.num_heads = config.num_attention_heads
        self.head_dim = self.hidden_size // self.num_heads
        self.max_position_embeddings = config.max_position_embeddings

        # 初始化 RotaryEmbedding
        self.rotary_emb = RotaryEmbedding(self.head_dim, max_position_embeddings=self.max_position_embeddings)
        )

    def forward(self, position_ids):
        proj = self.W_pack(hidden_states)
        proj = proj.unflatten(-1, (3, self.hidden_size)).unsqueeze(0).transpose(0, -2).squeeze(-2)
         # batch_size x source_len x hidden_size
        query_states = proj[0].view(bsz, q_len, self.num_heads, self.head_dim).transpose(1, 2) 
        # batch_size x target_len x head_size
        key_states = proj[1].view(bsz, q_len, self.num_heads, self.head_dim).transpose(1, 2)  
        # batch_size x source_len x hidden_size
        value_states = proj[2].view(bsz, q_len, self.num_heads, self.head_dim).transpose(1, 2)  

        kv_seq_len = key_states.shape[-2]

         # 获取 cos, sin
        cos, sin = self.rotary_emb(value_states, seq_len=kv_seq_len)

        # 计算旋转位置编码
        query_states, key_states = apply_rotary_pos_emb(query_states, key_states, cos, sin, positio

MLP

Activation 函数用的是 SwiGLU,这是一个结合了 Swish 与 GLU 的优点的东西,所以我觉得这个是有意义的,不过参数的选择要经过一定的实验才成。

class MLP(nn.Module):
    def __init__(
            self,
            hidden_size: int,
            intermediate_size: int,
            hidden_act: str,
    ):
        super().__init__()
        self.gate_proj = nn.Linear(hidden_size, intermediate_size, bias=False)
        self.down_proj = nn.Linear(intermediate_size, hidden_size, bias=False)
        self.up_proj = nn.Linear(hidden_size, intermediate_size, bias=False)
        self.act_fn = ACT2FN[hidden_act]

    def forward(self, x):
        return self.down_proj(self.act_fn(self.gate_proj(x)) * self.up_proj(x))

RMSNorm

RMSNorm(Root Mean Square Layer Normalization),是一般 LayerNorm 的一种变体,可以在梯度下降时令损失更加平滑

与 LayerNorm 相比,RMSNorm 的主要区别在于去掉了减去均值的部分(re-centering),只保留方差部分(re-scaling),从归一化的表达式上可以直观地看出。

一般的 LayerNorm:

\bar{a_i} = \frac{a_i - \mu}{\sigma}g_i \quad\quad \mu = \frac{1}{n}\sum_{i=1}^{n}a_i \quad\quad \sigma=\sqrt{\frac{1}{n}\sum_{i=1}^{n}(a_i-\mu)^2}

RMSNorm:

\bar{a_i} = \frac{a_i}{RMS(a)}g_i \quad\quad RMS(a) = \sqrt{\frac{1}{n}\sum_{i=1}^{n}a_i^2}

源代码:

class RMSNorm(nn.Module):
    def __init__(self, hidden_size, eps=1e-6):
        """
        RMSNorm is equivalent to T5LayerNorm
        """
        super().__init__()
        self.weight = nn.Parameter(torch.ones(hidden_size))
        self.variance_epsilon = eps

    def forward(self, hidden_states):
        variance = hidden_states.to(torch.float32).pow(2).mean(-1, keepdim=True)
        hidden_states = hidden_states * torch.rsqrt(variance + self.variance_epsilon)

        # convert into half-precision if necessary
        if self.weight.dtype in [torch.float16, torch.bfloat16]:
            hidden_states = hidden_states.to(self.weight.dtype)

        return self.weight * hidden_states

NormHead

用一个 normalize 后的 weight 替换了 Transformer 里的 linera。 但是它里边用了个方法:nn.init.kaiming_uniform_(self.weight, a=math.sqrt(5)) ,这个权重的初始化方法居然是何恺明大神的方法。下面是具体的实现。

计算因子:计算一个因子a,该因子取决于激活函数的类型。对于ReLU激活(以及其变种,如Leaky ReLU),a通常取为sqrt(5)(这是PyTorch中的默认值)。 计算标准差:计算权重张量的标准差为std = gain * sqrt(2 / fan_in),其中gain是根据激活函数类型计算的增益值,fan_in是权重张量的输入单元数(即权重张量的形状的第一个维度)。 均匀分布初始化:使用均匀分布U(-bound, bound)初始化权重张量,其中bound = std * sqrt(3)。

BaichuanForCausalLM中使用

self.lm_head = NormHead(config.hidden_size, config.vocab_size, bias=False)


class NormHead(nn.Module):
    def __init__(self, hidden_size, vocab_size, bias=False):
        super().__init__()
        self.weight = nn.Parameter(torch.empty((vocab_size, hidden_size)))
        nn.init.kaiming_uniform_(self.weight, a=math.sqrt(5))
        self.first_flag = True

    def forward(self, hidden_states):
        if self.training:
            norm_weight = nn.functional.normalize(self.weight)
            self.first_flag = True
        elif self.first_flag:
            self.first_flag = False
            self.weight.data = nn.functional.normalize(self.weight)
            norm_weight = self.weight
        else:
            norm_weight = self.weight
        return nn.functional.linear(hidden_states, norm_weight)

就是hidden_state进output前对其向量进行归一化

为了稳定训练并提高模型性能,将输出嵌入(也称为“头”)进行归一化。在实验中,NormHead有两个优势。首先,在初步实验中发现,头的范数容易不稳定。稀有标记嵌入的范数在训练过程中会变小,这会干扰训练动力学。NormHead可以让训练稳定。其次,发现语义信息主要由嵌入的余弦相似性编码,而不是L2距离。因为当前的线性分类器计算logits是通过点积完成的,这是一种L2距离和余弦相似性的组合。NormHead缓解了计算中的L2距离干扰。

Max-z loss

softmax_normalizer = shift_logits.max(-1).values ** 2
z_loss = self.config.z_loss_weight * softmax_normalizer.mean()

求最大值平方和的根号与原本的ce loss相加,目的是为了抑制最大值使其更稳定。在训练过程中,发现LLM的logits值可能非常大。绝对logits值在softmax函数不太重要,因为它只取决于它们之间的相对值。在大逻辑值在推理过程中造成了问题,因为重复惩罚的常见实现直接应用到logits值上。以这种方式收缩非常大的logits值可能会显著改变softmax后的概率,使模型对重复惩罚超参数的选择敏感。受NormSoftmax和PaLM的辅助Z损失的启发,添加了一个Max-z损失来规范化logits值:

\iota_{max-z} = 2e^{-4} * z^{2}

其中 z 是最大 logits值。 这有帮助稳定训练。

fine-tune

定义训练超参数

比如模型存储的本地路径、数据存储的本地路径、采用的优化器、上下文最大长度、是否使用Lora算法等等。注意,这里的模型路径"baichuan-inc/Baichuan2-7B-Base"是使用HuggingFace远程模型来完成的,如果模型已经下载到本地,就需要调整路径。

# @dataclass 是一个Python装饰器,用于自动生成初始化、比较等特殊方法。它使得数据类定义变得简洁。
@dataclass
# 这行定义了一个名为 `ModelArguments` 的类。
class ModelArguments:
    # 定义了一个可选的字符串类型的类属性 `model_name_or_path`,其默认值为 `"baichuan-inc/Baichuan2-7B-Base"`。
    model_name_or_path: Optional[str] = field(default="baichuan-inc/Baichuan2-7B-Base")

# 同第一行,用于自动生成特殊方法。
@dataclass
# 定义了一个名为 `DataArguments` 的类。
class DataArguments:
    # 定义了一个字符串类型的类属性 `data_path`,其默认值为 `None`。还为该字段添加了一些元数据描述。
    data_path: str = field(
        default=None, metadata={"help": "Path to the training data."}
    )

# 同上,用于自动生成特殊方法。
@dataclass
# 定义了一个名为 `TrainingArguments` 的类,该类继承自 `transformers.TrainingArguments`。
class TrainingArguments(transformers.TrainingArguments):
    # 定义了一个可选的字符串类型的类属性 `cache_dir`,其默认值为 `None`。
    cache_dir: Optional[str] = field(default=None)
    # 定义了一个字符串类型的类属性 `optim`,其默认值为 `"adamw_torch"`。
    optim: str = field(default="adamw_torch")
    # 定义了一个整型属性 `model_max_length`,其默认值为 `512`,并为它提供了描述。
    model_max_length: int = field(
        default=512,
        metadata={
            "help": "Maximum sequence length. Sequences will be right padded (and possibly truncated)."
        },
    )
    # 定义了一个布尔值属性 `use_lora`,默认为 `False`。
    use_lora: bool = field(default=False)

定义数据结构

在介绍SupervisedDataset之前,先对官方代码所采用的数据集有一个直观了解。 训练数据放在data/belle_chat_ramdon_10k.json,该样例数据是从multiturn_chat_0.8M采样出 1 万条,并且做了格式转换。主要是展示多轮数据怎么训练,不保证效果。下图展示了两条数据的格式。多轮对话有两个角色,分别是用户和GPT,在处理数据时需要对其作不同处理。在此处还有更加详细的介绍。

input_ids:前面拼上各自的标识符(<reserved_106>、<reserved_107> )之后拼接各自内容对应的 token ids。

labels:对于 human 的内容,其标识符(如果没记错的话是<reserved_106>)对应位置是 </s> (id=2),其他位置是 -100,不计算这些tokens的 loss 和梯度;对于 assistant 的内容,其标识符(如果没记错的话是<reserved_107>)对应位置是 -100,其他位置和 input_ids 一致。

class SupervisedDataset(Dataset):
    """Dataset for supervised fine-tuning."""

    # 定义了 `SupervisedDataset` 类的初始化方法,并接收一系列参数。
    def __init__(
        self,
        data_path,
        tokenizer,
        model_max_length,
        user_tokens=[195],  # [195]所代表的符号是<reserved_106>, 表示 human 输入
        assistant_tokens=[196],   # [196]所代表的符号是<reserved_107>, 表示 assistant 输出
    ):
        # 调用父类(Dataset)的初始化方法。
        super(SupervisedDataset, self).__init__()
        # 读取由 `data_path` 参数指定的JSON文件,并将其内容赋值给 `self.data`。
        self.data = json.load(open(data_path))
        # 这几行将传入的参数赋值给相应的类属性。
        self.tokenizer = tokenizer
        self.model_max_length = model_max_length
        self.user_tokens = user_tokens
        self.assistant_tokens = assistant_tokens
        self.ignore_index = -100
        # 对第一个数据进行预处理,并打印它的输入。
        item = self.preprocessing(self.data[0])
        print("input:", self.tokenizer.decode(item["input_ids"]))
        labels = []
        # 对预处理后的标签进行解码,并打印解码后的内容。
        for id_ in item["labels"]:
            if id_ == -100:
                continue

            labels.append(id_)
        print("label:", self.tokenizer.decode(labels))

    # 定义魔法方法,返回数据集的大小。
    def __len__(self):
        return len(self.data)

    # 定义预处理方法,对单个示例进行预处理。
    def preprocessing(self, example):
        # 初始化输入ID的空列表。
        input_ids = []
        # 初始化标签的空列表。
        labels = []

        # 遍历每个对话中的消息。
        for message in example["conversations"]:
            # 获取消息的发送者。
            from_ = message["from"]
            # 获取消息的内容。
            value = message["value"]
            # 使用tokenizer对消息内容进行编码。
            # 由于tokenizer并不是逐字编码,因此通常len(value_ids)≠len(value)
            value_ids = self.tokenizer.encode(value)

            # 如果消息来自人类用户。
            if from_ == "human":
                # 在输入ID中添加用户特定的token和消息token。
                input_ids += self.user_tokens + value_ids
                # 在标签中添加结束符和忽略标签。
                labels += [self.tokenizer.eos_token_id] + [self.ignore_index] * len(
                    value_ids
                )
            else:
                # 如果消息来自助手,添加助手特定的token和消息token。
                input_ids += self.assistant_tokens + value_ids
                # 在标签中添加忽略标签和消息token。
                labels += [self.ignore_index] + value_ids

        # input_ids[:20]  =  [195, 91529,  2365, 50020, 92415, 8945, 3337, 6824, 13449,   66,    5,  196, 2015, 65, 2835, 11024, 13172, 2665, 70, 5]
        # labels[:20]     =  [2,    -100,  -100,  -100,  -100, -100, -100, -100,  -100, -100, -100, -100, 2015, 65, 2835, 11024, 13172, 2665, 70, 5]
        #                    人类发言开始的地方             
        # 在输入ID和标签的末尾都追加结束符token。
        input_ids.append(self.tokenizer.eos_token_id)
        labels.append(self.tokenizer.eos_token_id)

        # 对输入ID和标签进行截断。
        input_ids = input_ids[: self.model_max_length]
        labels = labels[: self.model_max_length]

        # 为输入ID和标签填充token。
        # 如果文本长度不够512, 需要用'<unk>'(id=0)来补全
        input_ids += [self.tokenizer.pad_token_id] * (
            self.model_max_length - len(input_ids)
        )
        labels += [self.ignore_index] * (self.model_max_length - len(labels))

        # 转换输入ID和标签为PyTorch的LongTensor。
        input_ids = torch.LongTensor(input_ids)
        labels = torch.LongTensor(labels)

        # 创建注意力掩码。
        # input_ids中出现0(即'<unk>')的位置均为False, 其余True
        attention_mask = input_ids.ne(self.tokenizer.pad_token_id)

        # 返回预处理后的结果。
        return {
            "input_ids": input_ids,
            "labels": labels,
            "attention_mask": attention_mask,
        }

    # 定义魔法方法,允许使用索引从数据集中获取示例。
    def __getitem__(self, idx) -> Dict[str, torch.Tensor]:
        return self.preprocessing(self.data[idx])


定义训练函数

def train():
    # HfArgumentParser可以将类对象中的实例属性转换成转换为解析参数
    # 必须注意的是,这里的类对象必须是通过@dataclass()创建的类对象
    # 并且通过HfArgumentParser创建的解析参数,都是可选参数
    parser = transformers.HfArgumentParser((ModelArguments, DataArguments, TrainingArguments))
    model_args, data_args, training_args = parser.parse_args_into_dataclasses()

    # 自动加载适合进行有条件生成任务(如对话生成)的预训练语言模型,即给定输入的内容输出(有最大长度限制)
    model = transformers.AutoModelForCausalLM.from_pretrained(
        model_args.model_name_or_path,
        trust_remote_code=True, # 对从远程位置下载的代码持有高度的信任,因此将禁用某些安全检查,以加快加载过程。
        cache_dir=training_args.cache_dir,
    )
    tokenizer = transformers.AutoTokenizer.from_pretrained(
        model_args.model_name_or_path,
        use_fast=False,
        trust_remote_code=True,
        model_max_length=training_args.model_max_length,
        cache_dir=training_args.cache_dir,
    )
    if training_args.use_lora:
        from peft import LoraConfig, TaskType, get_peft_model
        peft_config = LoraConfig(
            # 因果语言建模任务(Causal Language Modeling), 模型试图预测给定上下文中的下一个单词
            # 该上下文通常包括在当前单词之前的所有单词. 这种建模方法遵循因果原则, 即当前单词只受到其前面单词的影响, 而不受后面单词的影响
            # 代表模型有GPT2, Bloom, OPT, GPT-Neo, GPT-J, LLaMA, ChatGLM
            task_type=TaskType.CAUSAL_LM, 
            target_modules=["W_pack"],
            inference_mode=False,              # 是否在推理模式下使用Peft模型
            r=4,                               # LoRA低秩矩阵的维数。关于秩的选择,通常,使用4,8,16即可
            lora_alpha=32,                     # LoRA低秩矩阵的缩放系数,为一个常数超参,调整alpha与调整学习率类似
            lora_dropout=0.1,                  # LoRA 层的丢弃(dropout)率,取值范围为[0, 1)。
        )
        model.enable_input_require_grads()
        model = get_peft_model(model, peft_config)
        model.print_trainable_parameters()
        

    dataset = SupervisedDataset(data_args.data_path, tokenizer, training_args.model_max_length)
    trainer = transformers.Trainer(model=model, args=training_args, train_dataset=dataset, tokenizer=tokenizer)
    trainer.train()
    trainer.save_state() # 保存 Trainer state
    trainer.save_model(output_dir=training_args.output_dir) # 保存模型,使得接下来可以采用 from_pretrained() 方法来加载

if __name__ == "__main__":
    train()

执行程序

在命令行执行程序,需要确定一些超参数的话,可以参考下面。由于我是单机单卡(3090)进行训练,所以没有采用deepspeed。

python fine-tune.py  \
    --report_to "none" \
    --data_path "data/belle_chat_ramdon_10k.json" \
    --model_name_or_path "baichuan-inc/Baichuan2-7B-Base" \
    --output_dir "output" \
    --model_max_length 512 \
    --num_train_epochs 4 \
    --per_device_train_batch_size 16 \
    --gradient_accumulation_steps 1 \
    --save_strategy epoch \
    --learning_rate 2e-5 \
    --lr_scheduler_type constant \
    --adam_beta1 0.9 \
    --adam_beta2 0.98 \
    --adam_epsilon 1e-8 \
    --max_grad_norm 1.0 \
    --weight_decay 1e-4 \
    --warmup_ratio 0.0 \
    --logging_steps 1 \
    --gradient_checkpointing True \
    --deepspeed ds_config.json \
    --bf16 True \
    --tf32 True

如果需要单机多卡,则使用如下代码。注意--model_name_or_path此处需要换成自己的模型和文件的路径。

hostfile=""
deepspeed --hostfile=$hostfile fine-tune.py  \
    --report_to "none" \
    --data_path "data/belle_chat_ramdon_10k.json" \
    --model_name_or_path "/data/hhy/LLM/models/Baichuan2-7B-Chat" \
    --output_dir "output" \
    --model_max_length 512 \
    --num_train_epochs 4 \
    --per_device_train_batch_size 16 \
    --gradient_accumulation_steps 1 \
    --save_strategy epoch \
    --learning_rate 2e-5 \
    --lr_scheduler_type constant \
    --adam_beta1 0.9 \
    --adam_beta2 0.98 \
    --adam_epsilon 1e-8 \
    --max_grad_norm 1.0 \
    --weight_decay 1e-4 \
    --warmup_ratio 0.0 \
    --logging_steps 1 \
    --gradient_checkpointing True \
    --deepspeed ds_config.json \
    --bf16 True \
    --tf32 True


编辑于 2024-03-21 15:39・IP 属地北京