类ChatGPT项目的部署与微调(上):从LLaMA到Alpaca、BELLE

前言  

近期,除了研究ChatGPT背后的各种技术细节 不断看论文(至少100篇,100篇目录见此:ChatGPT相关技术必读论文100篇),还开始研究一系列开源模型(包括各自对应的模型架构、训练方法、训练数据、本地私有化部署、硬件配置要求、微调等细节) 

本文一开始是作为此文《ChatGPT技术原理解析:从RL之PPO算法、RLHF到GPT4、instructGPT》的第4部分,但随着研究深入 为避免该文篇幅又过长,将把『第4部分 开源项目』抽取出来 独立成本文,然后不断续写本文到最终2万字左右(3.22日7000余字)

毕竟我上半年的目标之一,便是把ChatGPT涉及的所有一切关键技术细节,以及相关的开源项目都研究的透透的,故过程中会不断产出一篇篇新文章出来,比如:

  1. 微积分和概率统计极简入门
  2. 一文通透优化算法
  3. 强化学习极简入门:通俗理解MDP、DP MC TC和Q学习、策略梯度、PPO
  4. ChatGPT技术原理解析(本系列核心主体,也是同类解读里最清晰、全面、细致的一篇)
  5. ChatGPT相关技术论文100篇
  6. 类ChatGPT项目的部署与微调:从LLaMA到Alpaca、BELLEChatLLaMA和ColossalChat从ChatGLM-6b到ChatDoctor
  7. CV多模态模型发展史(23年4月发布),详述GPT4背后多模态的能力起源与发展历史,包括但不限于DTER、DDPM、Vision Transformer、CLIP、Swin Transformer、DALL·E 2、Stable Diffusion、BEiT-3、Visual ChatGPT、GPT4等

第一部分 Colossal-AI、PaLM-rlhf-pytorch、Open-Assistant等

 虽说GPT3在2020年就出来了,但OpenAI并未开源,所以直到一年半后以后才有国内外各个团队比如DeepMind等陆续复现出来,这些大厂的复现代码我们自然无法窥知一二,毕竟人家也未开源出来

再到后来基于GPT3的InstructGPT、基于GPT3.5ChatGPT初版(GPT3.5的参数规模也尚无准确定论)、GPT4均未开源,OpenAI不再open,好在Meta等公司或研究者开源出了一系列类ChatGPT项目,本部分针对其中部分做下简要推荐(根据发布顺序排序)

1.1 基于Colossal-AI低成本实现类ChatGPT迷你版的训练过程

2.15,很多朋友在GitHub上发现了一个基于Colossal-AI低成本实现类ChatGPT迷你版训练过程的开源项目(基于OPT + RLHF + PPO),虽是类似GPT3的开源项目OPT与RLHF的结合,但可以增进我们对ChatGPT的理解,该项目有几个不错的特点

  1. 很多同学一看到DL,便会想到大数据,而数据量一大,还用CPU处理的话很可能训练一个小任务都得半天,而如果用GPU跑,可能一两分钟就出来了。于此,在深度学习大火的那几年,特别是AlphaGo出来的16年起,我司七月在线便分别为VIP、AI系统大课、在职提升大课、求职/论文/申博/留学1V1辅导提供GPU云平台进行实战训练

    但如果想训练那种千亿参数规模的开源模型,就不只是有GPU就完事了,比如1750亿参数规模这种得用64张AI 100(即便经过一系列内存开销上的优化,也得至少32张AI 100,单张AI 100售价10万以上,且现在还经常没货),这样的硬件要求是大部分个人是无法具备的,所以该开源项目提供了单GPU、独立4/8-GPUs 的版本
  2. 如下代码所示,启动简单
    from chatgpt.nn import GPTActor, GPTCritic, RewardModel
    from chatgpt.trainer import PPOTrainer
    from chatgpt.trainer.strategies import ColossalAIStrategy
    
    strategy = ColossalAIStrategy(stage=3, placement_policy='cuda')
    
    with strategy.model_init_context():
        actor = GPTActor().cuda()
        critic = GPTCritic().cuda()
        initial_model = deepcopy(actor).cuda()
        reward_model = RewardModel(deepcopy(critic.model)).cuda()
    
    trainer = PPOTrainer(strategy, actor, critic, reward_model, initial_model, ...)
    trainer.fit(prompts)
  3. 训练过程明确清晰,如下图(由于此文已经详细介绍过ChatGPT的训练步骤,故不再赘述)

    ​此外,据钟博士在我所维护的『Machine Learning读书会群』里所说,Colossal-AI的并行效率确实不错,是新加坡的一个初创团队推出的,但目前尚没有团队采用Colossal-AI框架来做主训练框架训练175b级别的超大模型,可以再了解下Meta家训练OPT用的Metaseq

1.2 TRL包:类似ChatGPT训练阶段三的PPO方式微调语言模型

通过《ChatGPT技术原理解析》一文,我们已经知道了ChatGPT的三阶段训练过程,其中,阶段三的本质其实就是通过PPO的方式去微调LM

GitHub上有个TRL(Transformer Reinforcement Learning,基于Hugging Face开发的Transformer库),便是通过PPO的方式去微调LM,需要的数据便是三元组「query, response, reward」,具体如下图所示

扫描二维码关注公众号,回复: 14621355 查看本文章
  1. Rollout:语言模型根据query生成response
  2. 评估:怎么评估模型针对特定query生成response的质量呢,我们可以使用a function、model、human feedback或它们的某种组合进行评估,然后为每个query/response对产生一个标量值,说白了 就是奖励模型有了,那就直接打分
  3. 优化:在优化步骤中,「query/response pairs」用于计算序列中标记的对数概率,且比较下面这两个模型输出之间的 KL 散度用作额外的奖励信号
    \rightarrow  经过训练的模型(即上图中的Active model)
    \rightarrow  基线模型(即上图中的Reference model),通常是PPO微调之前的模型(比如这里的GPT2,或者instructGPT里的SFT)
    最终,使得Active model生成的响应不会偏离基线模型Reference model太远

示例代码如下

# imports
import torch
from transformers import AutoTokenizer
from trl import PPOTrainer, PPOConfig, AutoModelForCausalLMWithValueHead, create_reference_model
from trl.core import respond_to_batch

# get models
model = AutoModelForCausalLMWithValueHead.from_pretrained('gpt2')
model_ref = create_reference_model(model)

tokenizer = AutoTokenizer.from_pretrained('gpt2')

# initialize trainer
ppo_config = PPOConfig(
    batch_size=1,
)

# encode a query
query_txt = "This morning I went to the "
query_tensor = tokenizer.encode(query_txt, return_tensors="pt")

# get model response
response_tensor  = respond_to_batch(model_ref, query_tensor)

# create a ppo trainer
ppo_trainer = PPOTrainer(ppo_config, model, model_ref, tokenizer)

# define a reward for response
# (this could be any reward such as human feedback or output from another model)
reward = [torch.tensor(1.0)]

# train model for one step with ppo
train_stats = ppo_trainer.step([query_tensor[0]], [response_tensor[0]], reward)

第二部分 LLaMA到Stanford Alpaca(微调LLaMA)、BELLE(微调BLOOMZ-7B)

2.1 Meta发布LLaMA((7B 13B 33B 65B):参数少但多数任务的效果好于GPT3

一直致力于LLM模型研究的国外TOP 3大厂除了OpenAI、Google,便是Meta(原来的Facebook)

Meta曾第一个发布了基于LLM的聊天机器人——BlenderBot 3,但输出不够安全,很快下线;再后来,Meta发布一个专门为科学研究设计的模型Galactica,但用户期望过高,发布三天后又下线

23年2.24日,Meta通过论文《LLaMA: Open and Efficient Foundation Language Models》发布了自家的大型语言模型LLaMA(这是解读之一),有多个参数规模的版本(7B 13B 33B 65B)

LLaMA只使用公开的数据(总计1.4T即1,400GB的token,其中CommonCrawl的数据占比67%,C4数据占比15%,Github Wikipedia Books这三项数据均各自占比4.5%,ArXiv占比2.5%,StackExchange占比2%),论文中提到

When training a 65B-parameter model, our code processes around 380 tokens/sec/GPU on 2048 A100 GPU with 80GB of RAM.

This means that training over our dataset containing 1.4T tokens takes approximately 21 days

且试图证明小模型在足够多的的数据上训练后,也能达到甚至超过大模型的效果

  • 比如13B参数的版本在多项基准上测试的效果好于2020年的参数规模达175B的GPT-3
  • 而对于65B参数的LLaMA,则可与DeepMind的Chinchilla(70B参数)和谷歌的PaLM(540B参数)旗鼓相当
  • 且Meta还尝试使用了论文「Scaling Instruction-Finetuned Language Models」中介绍的指令微调方法,由此产生的模型LLaMA-I,在MMLU(Massive Multitask Language Understanding,大型多任务语言理解)上要优于Google的指令微调模型Flan-PaLM-cont(62B)

2.2 代码级解读:LLaMA的模型架构——RMSNorm/SwiGLU/RoPE/Transformer

2.2.1 项目环境依赖:torch、fairscale、fire、sentencepiece

此项目给出的环境依赖有4个:

  1. torch
  2. fairscale,fairscale是用来做GPU分布的,一般是当使用DDP仍然遇到超显存的问题时使用fairscale
  3. fire,fire是一个命令行工具,用或者不用他都可以
  4. sentencepiece,sentencepiece是用于tokenizer的工具包
    from sentencepiece import SentencePieceProcessor
    from logging import getLogger
    from typing import List
    import os
    
    
    logger = getLogger()
    
    
    class Tokenizer:
        def __init__(self, model_path: str):
            # reload tokenizer
            assert os.path.isfile(model_path), model_path
            self.sp_model = SentencePieceProcessor(model_file=model_path)
            logger.info(f"Reloaded SentencePiece model from {model_path}")
    
            # BOS / EOS token IDs
            self.n_words: int = self.sp_model.vocab_size()
            self.bos_id: int = self.sp_model.bos_id()
            self.eos_id: int = self.sp_model.eos_id()
            self.pad_id: int = self.sp_model.pad_id()
            logger.info(
                f"#words: {self.n_words} - BOS ID: {self.bos_id} - EOS ID: {self.eos_id}"
            )
            assert self.sp_model.vocab_size() == self.sp_model.get_piece_size()
    
        def encode(self, s: str, bos: bool, eos: bool) -> List[int]:
            assert type(s) is str
            t = self.sp_model.encode(s)
            if bos:
                t = [self.bos_id] + t
            if eos:
                t = t + [self.eos_id]
            return t
    
        def decode(self, t: List[int]) -> str:
            return self.sp_model.decode(t)

2.2.2 RMSNorm:对每个Transformer子层的输入进行归一化

为了提高训练的稳定性,对每个transformer子层的输入进行归一化,而不是对输出进行归一化,且使用由Zhang和Sennrich(2019)提出的RMSNorm(Root Mean Square Layer Normalization)归一化函数
RMS Norm是一般LayerNorm的一种变体,可以在梯度下降时令损失更加平滑
与layerNorm相比,RMS Norm的主要区别在于去掉了减去均值的部分(re-centering),只保留方差部分(re-scaling),从归一化的表达式上可以直观地看出
既然提到了,咱们就不千篇一律的泛泛而谈,而是具体说明下

  • 一般的LN:

    \overline{a}_i = \frac {a_i- \mu} \sigma g_i

    其中

    \mu = \frac 1 n \sum_{i=1}^na_i

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

  • RMS Norm:

    \overline{a}_i = \frac {a_i} {RMS(a)} g_i

    其中

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

至于RMS Norm为什么有用,需要求梯度进行分析,感兴趣的同学可以阅读RMS Norm的论文

class RMSNorm(torch.nn.Module):
    def __init__(self, dim: int, eps: float = 1e-6):
        super().__init__()
        // eps防止取倒数之后分母为0
        self.eps = eps
        self.weight = nn.Parameter(torch.ones(dim))

    // x是输入
    def _norm(self, x):
        // torch.rsqrt是开平方并取倒数
        // x.pow(2)是平方
        / mean(-1)是在最后一个维度(即hidden特征维度)上取平均
        return x * torch.rsqrt(x.pow(2).mean(-1, keepdim=True) + self.eps)

    def forward(self, x):
        output = self._norm(x.float()).type_as(x)
        // weight是末尾乘的可训练参数,即gi
        return output * self.weight

2.2.3 SwiGLU替代ReLU

用Shazeer(2020)提出的SwiGLU替代ReLU,在维度上使用的维度是2/3*4d,而不是PaLM中的4d

LLaMA采用SwiGLU替换了原有的ReLU,具体是采用SwiGLU的FNN,在论文中以如下公式进行表述:

FFN_{swiGLU}(x, W, V, W_2) = (Swish_1(xW)\otimes xV)W_2

其中

Swish_\beta(x) = x\sigma(\beta x))

对应论文见:Ramachandran et al., 2017
代码实现上:可以通过调用torch内置方法F.silu()实现,会在下文的FFN部分介绍

2.2.4 位置编码:旋转位置嵌入(RoPE)

在位置编码上,删除了绝对位置嵌入,而在网络的每一层增加了苏剑林等人(2021)提出的旋转位置嵌入(RoPE),其思想是采用绝对位置编码的形式,实现相对位置编码

  • RoPE主要借助了复数的思想,为了引入复数,首先假设了在加入位置信息之前,原有的编码向量是二维行向量q_mk_n,其中mn是绝对位置,现在需要构造一个变换,将mn引入到q_mk_n中,即寻找变换: 

    \tilde {q_m} = f(q, m), \tilde{k_n} = f(k, n)

    考虑到Attention的核心计算是内积:

    Attention(Q, K,V) = softmax(\frac {QK^T} {\sqrt{d_k}})V

    所以,寻求的这个f(*)变换,应该具有特性:\langle f(q, m), f(k, n) \rangle = g(q, k, m-n)
  • 这里直接说结论,寻求的变换就是q_me^{im\theta},也就是给q_m乘以e^{im\theta},相应地,k_n乘以e^{in\theta}
    做了这样一个变换之后,根据复数的特性,有:

    \langle q_m, k_n \rangle = Re[q_mk^*_n]

    也就是,如果把二维向量看做复数,那么它们的内积,等于一个复数乘以另一个复数的共轭,得到的结果再取实部,代入上面的变换,也就有:

    \langle q_me^{im\theta}, k_ne^{in\theta} \rangle = Re[(q_me^{im\theta}) (k_ne^{in\theta})^*] =Re[q_mk_n^*e^{i(m-n)\theta}]

    这样一来,内积的结果就只依赖于(m-n),也就是相对位置了
    换言之,经过这样一番操作,通过给Embedding添加绝对位置信息,可以使得两个token的编码,经过内积变换(self-attn)之后,得到结果是受它们位置的差值,即相对位置影响的
  • 于是对于任意的位置为m的二维向量[x, y],把它看做复数,乘以e^{im\theta},而根据欧拉公式,有:

    e^{im\theta}=\cos{m\theta}+i\sin{m\theta}

    于是上述的相乘变换也就变成了:

    (x+iy)e^{im\theta}=(x\cos{m\theta}-y\sin{m\theta})+i(x\sin{m\theta}+y\cos{m\theta})

    把上述式子写成矩阵形式:

    而这个变换的几何意义,就是在二维坐标系下,对向量(q_0, q_1)进行了旋转,因而这种位置编码方法,被称为旋转位置编码
  • 根据刚才的结论,结合内积的线性叠加性,可以将结论推广到高维的情形。可以理解为,每两个维度一组,进行了上述的“旋转”操作,然后再拼接在一起:

    由于矩阵的稀疏性,会造成计算上的浪费,所以在计算时采用逐位相乘再相加的方式进行:

    其中\otimes为矩阵逐位相乘操作

原理理解了,接下来可以代码实现旋转位置编码

def precompute_freqs_cis(dim: int, end: int, theta: float = 10000.0):
    // 首先torch.arange创建了一个tensor,[ 0 , 2 , 4 , . . . , 60 , 62 ] [0, 2, 4, ..., 60, 62][0,2,4,...,60,62]
    // 然后统一除以64,把它变成分数,然后整体作为基础角度的指数,它的shape是(32)
    freqs = 1.0 / (theta ** (torch.arange(0, dim, 2)[: (dim // 2)].float() / dim))

    // t比较容易理解,也就是绝对位置信息,它的shape是(1024)
    t = torch.arange(end, device=freqs.device)

    // torch.outer是把一个向量的转置乘以另一个向量:torch.outer(a, b) = a^T * b
    // 于是根据torch.outer运算,我们得到了一个shape为(1024, 32)的tensor。其意义也就是将每一个绝对位置,分配到对应的角度,相乘
    // 直观理解一下,就是每一个绝对位置上,都有32个角度
    // 为什么是这样的呢,回顾计算的公式,对于旋转矩阵,每两个元素为一组,它们乘以的角度是同一个θ,所以这个(1024, 32)
    // 在后续的过程中,就可以reshape成(512, 64),并且在64的那个维度上,每两个是相同的
    freqs = torch.outer(t, freqs).float()

    // torch.polar(abs, angle)利用一个绝对数值和一个角度值,从而在极坐标下构造一个复数张量
    // 即abs∗cos(angle)+abs∗sin(angle)j
    // torch.polar(torch.tensor([1], dtype=torch.float64), torch.tensor([np.pi / 2], dtype=torch.float64))
    // # tensor([6.1232e-17+1.j], dtype=torch.complex128)

    // freqs_cis其实就是需要计算出来的mθ,也就是跟绝对位置相关的旋转的角度,在极坐标下对应的复数tensor
    // 这一步就是在生成我们需要的位置信息
    // 直观理解一下,像是在复平面内,以原点为中心,转了1024组,每一组64个的单位向量,它的shape是(1024, 64)
    freqs_cis = torch.polar(torch.ones_like(freqs), freqs)  # complex64
    return freqs_cis


// 第二个函数reshape_for_broadcast,是把freqs_cis变成和输入的tensor相同的形状
def reshape_for_broadcast(freqs_cis: torch.Tensor, x: torch.Tensor):
    ndim = x.ndim
    assert 0 <= 1 < ndim
    assert freqs_cis.shape == (x.shape[1], x.shape[-1])

    // 这个方法的作用是为了把freqs_cis变成和输入的tensor相同的形状
    // 需要注意的是,这里的freqs_cis并不是precompute_freqs_cis生成的形状为(1024, 64)的那个tensor
    // 而是根据输入的绝对位置,在(1024, 64)的tensor中,截取了长度为当前seq_len的一部分
    // 代码在Transformer类的forward方法中freqs_cis = self.freqs_cis[start_pos : start_pos + seqlen]
    // 也就是说,假如当前输入的序列长度是512,那么截取出来的这个新的freqs_cis,形状就是(512, 64)
    // reshape之后,形状就变成了(1, 512, 1, 32),也就是在每一个位置上,都对应有32个角度
    // 根据上面torch.polar的介绍,当我们固定绝对值(也就是向量的模长)时,角度就可以在笛卡尔坐标系下唯一确定一个复数
    // 这样一来也就是32个复数,即64个特征维度,所以就可以对应的将它融合到每个attention head的64个特征中去了
    shape = [d if i == 1 or i == ndim - 1 else 1 for i, d in enumerate(x.shape)]
    return freqs_cis.view(*shape)


// apply_rotary_emb方法,这个方法其实就是把位置信息添加到原有的编码结果上,在multi-head attention阶段调用
def apply_rotary_emb(
    xq: torch.Tensor,
    xk: torch.Tensor,
    freqs_cis: torch.Tensor,
) -> Tuple[torch.Tensor, torch.Tensor]:

    // torch.view_as_complex是把一个tensor转为复数形式
    // 比如torch.view_as_complex(torch.Tensor([[1, 2], [3, 4], [5, 6]]))
    // # tensor([1.+2.j, 3.+4.j, 5.+6.j])
    
    // 假设输入x_q的尺寸就是(2, 512, 12, 64)
    // 那么这一句操作的reshape,就是把它变成(2, 512, 12, -1, 2),也就是(2, 512, 12, 32, 2)。x_k同理,略
    // 紧接着把它变成复数形式,也就是变成了(2, 512, 12, 32)的形状。
    xq_ = torch.view_as_complex(xq.float().reshape(*xq.shape[:-1], -1, 2))
    xk_ = torch.view_as_complex(xk.float().reshape(*xk.shape[:-1], -1, 2))

    // 然后进入到上面的第二个函数reshape_for_broadcast
    freqs_cis = reshape_for_broadcast(freqs_cis, xq_)

    // torch.view_as_real是把复数tensor变回实数
    // torch.view_as_real(torch.view_as_complex(torch.Tensor([[1, 2], [3, 4], [5, 6]])))
    // # tensor([[1., 2.],
    // #         [3., 4.],
    // #         [5., 6.]])
    // reshape之后,就是将位置信息融入query和key中
    // 这一步将二者相乘得到的复数tensor,重新转换为实数形式,得到的shape为(2, 512, 12, 32, 2)
    // 然后再flatten成(2, 512, 12, 64),这样一来,就变回了和最开始x_q相同的形状,也就完成了将位置信息融入到x_q的这一操作,x_k同理
    xq_out = torch.view_as_real(xq_ * freqs_cis).flatten(3)
    xk_out = torch.view_as_real(xk_ * freqs_cis).flatten(3)
    return xq_out.type_as(xq), xk_out.type_as(xk)

引用此文的介绍,再着重解释下precompute_freqs_cis的作用

  • 假设
    batch_size为2
    seq_len固定为512
    attention_head的数量为12
    每个attention_head的维度为64,那么,对于输入到multi-head attn中的输入x_q的尺寸就是
    (2, 512, 12, 64)

而freqs_cis其实就是需要计算出来的m\theta也就是跟绝对位置相关的旋转的角度,在极坐标下对应的复数tensor

而precompute_freqs_cis就是提前将这些旋转角度对应的tensor给创建出来,并可以重复利用。因为确定了序列的最大长度,所以这个tensor是固定死的。根据后续的数据流我们可以发现,在调用该函数时,传入的两个参数分别是attention_head的维度,以及最大长度的两倍,具象地,也就是64和1024

2.2.4 Transform架构的实现:Attention计算、SA、FFN

LLaMA和GPT一样,都是基于Transformer这个架构,通常,我们在构建transformer时,是按Block构建的,每个transformer Block包含SA和FFN两部分,然后再通过堆叠block的形式,构建起整个transformer网络,LLaMA也是这样做的

回顾一下Attention计算的总体过程是:

  1. 输入x,分别经过三个Linear得到x_q, x_k, x_v
  2. 在 x_q 和x_k中加入旋转位置编码
  3. 缓存 x_q 和 x_k 
  4. 计算softmax(\frac {QK^T} {\sqrt{d_k}})V

其中有一个细节就是缓存机制,它设计的目的是在generate时减少token的重复计算。简单解释一下,就是在计算第n个token特征的时候,需要用到第1,...,n-1个token,即每次生成时,需要知道前面所有的过往信息,如果每次都从头算的话,那就会造成极大的浪费,所以就没算一个位置的信息,就把它缓存下来

接下来,我们来看下代码实现,首先是SA部分:

class Attention(nn.Module):
    def __init__(self, args: ModelArgs):
        super().__init__()

        self.n_local_heads = args.n_heads // fs_init.get_model_parallel_world_size()
        self.head_dim = args.dim // args.n_heads

        self.wq = ColumnParallelLinear(
            args.dim,
            args.n_heads * self.head_dim,
            bias=False,
            gather_output=False,
            init_method=lambda x: x,
        )
        self.wk = ColumnParallelLinear(
            args.dim,
            args.n_heads * self.head_dim,
            bias=False,
            gather_output=False,
            init_method=lambda x: x,
        )
        self.wv = ColumnParallelLinear(
            args.dim,
            args.n_heads * self.head_dim,
            bias=False,
            gather_output=False,
            init_method=lambda x: x,
        )
        self.wo = RowParallelLinear(
            args.n_heads * self.head_dim,
            args.dim,
            bias=False,
            input_is_parallel=True,
            init_method=lambda x: x,
        )

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

    def forward(self, x: torch.Tensor, start_pos: int, freqs_cis: torch.Tensor, mask: Optional[torch.Tensor]):
        bsz, seqlen, _ = x.shape
        xq, xk, xv = self.wq(x), self.wk(x), self.wv(x)

        xq = xq.view(bsz, seqlen, self.n_local_heads, self.head_dim)
        xk = xk.view(bsz, seqlen, self.n_local_heads, self.head_dim)
        xv = xv.view(bsz, seqlen, self.n_local_heads, self.head_dim)

        xq, xk = apply_rotary_emb(xq, xk, freqs_cis=freqs_cis)

        self.cache_k = self.cache_k.to(xq)
        self.cache_v = self.cache_v.to(xq)

        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]

        xq = xq.transpose(1, 2)
        keys = keys.transpose(1, 2)
        values = values.transpose(1, 2)
        scores = torch.matmul(xq, keys.transpose(2, 3)) / math.sqrt(self.head_dim)
        if mask is not None:
            scores = scores + mask  # (bs, n_local_heads, slen, cache_len + slen)
        scores = F.softmax(scores.float(), dim=-1).type_as(xq)
        output = torch.matmul(scores, values)  # (bs, n_local_heads, slen, head_dim)
        output = output.transpose(
            1, 2
        ).contiguous().view(bsz, seqlen, -1)

        return self.wo(output)

然后是FFN部分,需要注意的点就是采用的激活函数,以及激活函数的位置

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

这里与常见模型中的FFN做一下简单的对比

  • BART中的FFN,用的是fc->act->fc,用了两层全连接
  • GPT中的FFN,用的是conv1D->act->conv1D,也是只用了两层
  • 而LLaMA中的FFN采用了三个全连接层以实现FFNSwiGLU,即
    FFN_{swiGLU}(x, W, V, W_2) = (Swish_1(xW)\otimes xV)W_2

然后将SA和FFN这两部分拼在一起就是一个transformer block

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

最后利用torch的module list将transformer block进行堆叠,拼上最前头的embedding部分,就是一个完整的transformer decoder结构了

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

        // 输入是token,先做token embedding,然后添加位置信息
        h = self.tok_embeddings(tokens)
        self.freqs_cis = self.freqs_cis.to(h.device)
        freqs_cis = self.freqs_cis[start_pos : start_pos + seqlen]

        // 对于decoder模型,为了防止标签泄漏,需要mask,所以做了一个上三角的mask矩阵
        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)

        // 接下来就是逐层的计算transformer
        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()

接着看下生成过程,如下:

  1. 对prompts进行tokenize,得到token ids;
  2. 计算当前batch的最大长度total_len,用来创建输入的token tensor,最大长度不能超过前文所述缓存的大小;
  3. 从当前batch中,最短的一个prompt的位置,作为生成的开始位置,开始生成;
  4. 输入的token tensor传入transformer模型,计算logits,得到形状为(batch_size, hidden_size)的logits(transformer最后一层的输出);
  5. softmax+top_p采样,得到当前预测的token,并更新当前位置,准备预测下一个token;
  6. 解码得到生成的文本

代码如下

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

2.3 LLaMA的Optimizer设计、模型加速优化与微型版本

在Optimizer设计上

  • 该模型使用AdamW优化器(Loshchilov和Hutter,2017)进行训练,超参数设置为β1=0.9,β2=0.95
    此外,使用余弦学习率方式,使最终学习率等于最大学习率的10%,以及使用0.1的权重衰减和1.0的梯度剪裁,和2,000个warm up策略,使得可以根据模型的大小改变学习率和批次大小

在模型的加速优化方面

  1. 首先,使用一个高效的因果多头注意力方式的实现,灵感来自Rabe和Staats(2021)以及Dao等人(2022),这个实现可在xformers库中找到,可以有效减少内存的使用和计算
    具体原理为通过不存储注意力权重和不计算由于语言建模任务的因果性质而被掩盖的键/查询分数来实现的
  2. 其次,为了进一步提高训练效率,减少了在check point的后向传递中重新计算的激活量,在实现上,通过手动实现trasnformer层的后向函数来进行操作
    为了充分受益于这种优化,还通过如Korthikanti等人(2022)中采用的方法,进行使用模型和序列并行来减少模型的内存使用
  3. 最后,该工作还尽可能地重叠激活的计算和GPU之间在网络上的通信
    最终的优化性能效果为:当训练一个65B参数的模型时,代码在2048A100的GPU上处理大约380个token/秒/GPU,并耗费80GB的内存,这意味着对包含1.4Ttoken的数据集进行训练大约花费了21天

LLaMA发布不久后,一些研究者基于它做了不少工作

  • 一开始最小参数7B的模型也需要近30GB的GPU才能运行,但通过比特和字节库进行浮点优化,能够让模型在单个NVIDIA RTX 3060上运行
  • 之后,GitHub 上的一名研究人员甚至能够在Ryzen 7900X CPU上运行LLM的7B 版本,每秒能推断出几个单词
  • 再之后,有研究者推出了llama.cpp,无需 GPU,就能运行 LLaMA
    llama.cpp 项目实现了在MacBook上运行 LLaMA,还有开发者成功的在 4GB RAM 的树莓派上运行了 LLaMA 7B

第三部分 通过self-instruct低成本微调LLaMA:Stanford Alpaca到BELLE

3.1 Stanford Alpaca:结合英文语料通过Self Instruct方式微调LLaMA 7B

3.1.1 什么是self-instruct方式:提示GPT3的API收集数据

3月中旬,斯坦福发布Alpaca(中文名:羊驼):号称只花100美元,人人都可微调Meta家70亿参数的LLaMA大模型(即LLaMA 7B),具体做法是通过52k指令数据,然后在8个80GB A100上训练3个小时,最终性能比肩GPT-3.5(text-davinci-003)

而斯坦福团队微调LLaMA 7B所用的52K指令数据,便是通过Self-Instruct『Self-Instruct是来自华盛顿大学Yizhong Wang等22年12月通过这篇论文《SELF-INSTRUCT: Aligning Language Model with Self Generated Instructions》提出的』提示GPT3的API拿到的

​具体而言,论文中提出

  1. 人工设计175个任务,每个任务都有对应的{指令 输入 输出/实例}或{指令 输出/实例},将这175个任务数据作为种子集
  2. 然后提示模型比如GPT3对应的text-davinci-001 (原论文中没用text-davinci-003,because their newer engines are trained with the latest user data and are likely to already see the SUPERNI evaluation set,但实际应用时比如斯坦福Alpaca可以指定text-davinci-003生成指令),使用种子集作为上下文示例来生成更多新的指令
  3. 对该模型生成的指令判断是否分类任务
  4. 使用模型生成实例
  5. 对上述模型生成的数据{指令 输入 输出/实例}过滤掉低质量或相似度高的
  6. 将经过过滤和后处理的数据添加到种子池中
    一直重复上述2-6步直到种子池有足够多的数据

而斯坦福的Alpaca,就是花了不到500美元使用OpenAI API生成了5.2万个这样的示例微调LLaMA搞出来的,个人觉得可以取名为 instructLLaMA-7B,^_^

3.1.2 Alpaca-LoRA:通过PEFT库在消费级GPU上微调「基于LLaMA的Alpaca」

在神经网络模型中,模型参数通常以矩阵的形式表示。对于一个预训练好的模型,其参数矩阵已经包含了很多有用的信息。为了使模型适应特定任务,我们需要对这些参数进行微调

LoRA的核心思想是用一种低秩的方式来调整这些参数矩阵。在数学上,低秩意味着一个矩阵可以用两个较小的矩阵相乘来近似,通过论文《LORA: LOW-RANK ADAPTATION OF LARGE LANGUAGE MODELS》可知(这是解读之一)

  1. 选择目标层:首先,在预训练神经网络模型中选择要应用LoRA的目标层。这些层通常是与特定任务相关的,如自注意力机制中的查询Q和键K矩阵
  2. 初始化映射矩阵和逆映射矩阵:为目标层创建两个较小的矩阵A和B
    \rightarrow  A是映射矩阵(随机高斯分布初始化),维度上是升维
    \rightarrow  B是逆映射矩阵(用0矩阵初始化),维度上是降维
    其中,矩阵的大小由LoRA的秩(rank)和alpha值确定
  3. 参数变换:将目标层的原始参数矩阵W通过映射矩阵A和逆映射矩阵B进行变换。计算公式为:W' = W + A * B。这里W'是变换后的参数矩阵
  4. 微调模型:使用新的参数矩阵W'替换目标层的原始参数矩阵W,然后在特定任务的训练数据上对模型进行微调
  5. 梯度更新:在微调过程中,计算损失函数关于映射矩阵A和逆映射矩阵B的梯度,并使用优化算法(如Adam、SGD等)对A和B进行更新
    注意,在更新过程中,原始参数矩阵W保持不变,说白了,训练的时候固定原始PLM的参数,只训练降维矩阵A与升维矩阵B
  6. 重复更新:在训练的每个批次中,重复步骤3-5,直到达到预定的训练轮次(epoch)或满足收敛条件

总之,LoRA的详细步骤包括选择目标层、初始化映射矩阵和逆映射矩阵、进行参数变换和模型微调。在微调过程中,模型会通过更新映射矩阵U和逆映射矩阵V来学习特定任务的知识,从而提高模型在该任务上的性能。

而Huggingface公司推出的PEFT(Parameter-Efficient Fine-Tuning)库便封装了LoRA这个方法,PEFT库可以使预训练语言模型高效适应各种下游任务,而无需微调模型的所有参数,即仅微调少量(额外)模型参数,从而大大降低了计算和存储成本

Model Full Finetuning PEFT-LoRA PyTorch PEFT-LoRA DeepSpeed with CPU Offloading
bigscience/T0_3B (3B params) 47.14GB GPU / 2.96GB CPU 14.4GB GPU / 2.96GB CPU 9.8GB GPU / 17.8GB CPU
bigscience/mt0-xxl (12B params) OOM GPU 56GB GPU / 3GB CPU 22GB GPU / 52GB CPU
bigscience/bloomz-7b1 (7B params) OOM GPU 32GB GPU / 3.8GB CPU 18.1GB GPU / 35GB CPU

且PEFT库支持以下流行的方法

  1. LoRA: LORA: LOW-RANK ADAPTATION OF LARGE LANGUAGE MODELS
  2. Prefix Tuning: Prefix-Tuning: Optimizing Continuous Prompts for GenerationP-Tuning v2: Prompt Tuning Can Be Comparable to Fine-tuning Universally Across Scales and Tasks
  3. P-Tuning: GPT Understands, Too
  4. Prompt Tuning: The Power of Scale for Parameter-Efficient Prompt Tuning

Alpaca-LoRA则可以通过PEFT库实现的LoRA方法在消费级GPU微调「基于LLaMA的Alpaca」,比如项目中的这个文件finetune.py 包含了PEFT在LLaMA上的直接应用,以及一些与prompt construction和tokenization相关的代码,以下是用法示例:

python finetune.py \
    --base_model 'decapoda-research/llama-7b-hf' \
    --data_path 'yahma/alpaca-cleaned' \
    --output_dir './lora-alpaca'

我们还可以调整我们的超参数(为方便大家理解,我给每个参数都加了注释说明):

python finetune.py \                             # 运行微调脚本
    --base_model 'decapoda-research/llama-7b-hf' \  # 选择预训练的基础模型
    --data_path 'yahma/alpaca-cleaned' \            # 用于微调的数据集路径
    --output_dir './lora-alpaca' \                  # 微调后模型的输出目录
    --batch_size 128 \                              # 设置每个批次的样本数量
    --micro_batch_size 4 \                          # 设置每个小批次的样本数量
    --num_epochs 3 \                                # 设置训练的轮次(epoch)
    --learning_rate 1e-4 \                          # 设置学习速率
    --cutoff_len 512 \                              # 设置截断长度
    --val_set_size 2000 \                           # 设置验证集的大小
    --lora_r 8 \                                    # 设置LoRA方法中的秩
    --lora_alpha 16 \                               # 设置LoRA方法中的alpha值
    --lora_dropout 0.05 \                           # 设置LoRA方法中的dropout率
    --lora_target_modules '[q_proj,v_proj]' \       # 设置使用LoRA进行微调的模型模块
    --train_on_inputs                               # 指示模型在训练时使用输入文本

3.1.3 Alpaca所用的self-instruct的影响力:解决一大批模型的数据扩展问题

很快,通过下文你会发现

  1. self-instruct启发出很多「羊驼类模型」
    羊驼率先带动的self-instruct,启发后续很多人/团队也用这个方式去采集『提示ChatGPT API』的数据,比如BELLE、ChatLLaMA、ColossalChat
  2. 很多「羊驼类模型」的数据被用于微调新一批模型
    然后还有一批模型各种叠加组合比如『Alpaca/BELLE』,又用于微调一批批模型
    比如ChatDoctor 有用到Alpaca的数据进行微调,再比如有人拿BELLE数据tuning去调chatglm

一下子出来这么新的模型 似乎有点懵,没事,请看下文及下一篇文章娓娓道来..

3.2 BELLE(中文版):结合中文语料通过Self Instruct方式微调BLOOMZ-7B

Stanford Alpaca的种子任务都是英语,收集的数据也都是英文,因此训练出来的模型未对中文优化。为了提升对话模型在中文上的效果,开源中文对话大模型70 亿参数的 BELLE(Bloom-Enhanced Large Language model Engine)来了(项目地址)。

它基于Stanford Alpaca完成,但进行了中文优化,并对生成代码进行了一些修改,不仅如此,模型调优也仅使用由 GPT3.5 (默认使用模型text-davinci-003,如果想使用ChatGPT的API比如gpt-3.5-turbo模型,可通过参数控制) 生产的数据(不包含任何其他数据)。

在数据方面,该项目开源了基于Stanford Alpaca的数据收集代码,基于这段代码生成了约 100 万条中文数据,结合 Alpaca 的 5.2 万条英文数据,在 BLOOMZ-7B 模型训练得到的 checkpoint 上传在 Hugging Face

BLOOM是由HuggingFace于2022年3月中旬推出的大模型,规模最大版本的参数量达到176B(GPT-3是175B),基于从 Megatron-LM GPT-2修改而来的仅解码器 transformer 模型架构,对应的论文为《BLOOM: A 176B-Parameter Open-Access Multilingual Language Model》(翻译之一解读之一)
至于HuggingFace是著名开源工具Transformers的开发公司,很多推理工具都会支持Transformers中的模型


截至23年3月中旬,超过100B参数量且能够支持中文的开源大模型只有BLOOM和GLM-130B

该项目主要包含以下三部分内容:

  • 175 个中文种子任务,斯坦福Alpaca一样,每个任务都包含对应的指令/任务、prompt、输出
    \rightarrow  zh_seed_tasks.jsonl:样例如下
     {  "id": "seed_task_20", "name": "horror_movie_opening",
    "instruction": "你需要为一部恐怖电影写一个创意的开场场景。",
    "instances": [{"input": "","output":" 太阳已经落山,留下了一个黑暗的小镇。微风吹拂空荡的街道,让每一个冒险走出门外的人感到一阵寒意。唯一的声音是被风吹动的树叶发出的轻微沙沙声。突然,一声令人毛骨悚然的尖叫声划破了寂静,随后是玻璃破碎的声音。一所房子亮起了灯光,可以看到一个人影朝镇中心奔跑。当> 那个人影越来越靠近时,清楚地看到那是一个年轻女子,她浑身血迹斑斑。"}],
    "is_classification": false  }
  • \rightarrow  prompt_cn.txt: 生成所使用的提示语
    \rightarrow  0.5M 生成的数据
  • 生成数据及其代码
    沿用 Alpaca 的方式:
    pip install -r requirements.txt
    export OPENAI_API_KEY=YOUR_API_KEY
    python generate_instruction.py generate_instruction_following_data

    默认使用 Completion API,模型 text-davinci-003。如果想使用 Chat API 并使用 gpt-3.5-turbo 模型,可通过参数控制:
    python generate_instruction.py generate_instruction_following_data \
      --api=chat --model_name=gpt-3.5-turbo

    输出文件在 Belle.train.json,可以人工筛选后再使用
  • 基于 BLOOMZ-7B1-mt 模型和 Belle.train.json 训练模型

猜你喜欢

转载自blog.csdn.net/v_JULY_v/article/details/129709105