相关博客
【自然语言处理】【大模型】MPT模型结构源码解析(单机版)
【自然语言处理】【大模型】ChatGLM-6B模型结构代码解析(单机版)
【自然语言处理】【大模型】BLOOM模型结构源码解析(单机版)
【自然语言处理】【大模型】极低资源微调大模型方法LoRA以及BLOOM-LORA实现代码
【深度学习】【分布式训练】Collective通信操作及Pytorch示例
【自然语言处理】【大模型】Chinchilla：训练计算利用率最优的大语言模型
【自然语言处理】【大模型】大语言模型BLOOM推理工具测试
【自然语言处理】【大模型】GLM-130B：一个开源双语预训练语言模型
【自然语言处理】【大模型】用于大型Transformer的8-bit矩阵乘法介绍
【自然语言处理】【大模型】BLOOM：一个176B参数且可开放获取的多语言模型
【自然语言处理】【ChatGPT系列】FLAN：微调语言模型是Zero-Shot学习器
【自然语言处理】【ChatGPT系列】ChatGPT的智能来自哪里？
【自然语言处理】【ChatGPT系列】大模型的涌现能力

一、MPT介绍

​ MPT是由MosaicML团队开源出来了英文预训练大模型，共开源了4个模型：MPT-7B Base、MPT-7B-StoryWriter-65k+、MPT-7B-Instruct、MPT-7B-Chat。

1. 模型结构与预训练

• MPT-7B Base是一个GPT风格的decoder-only Transformer模型，参数量为6.7B，使用了1T tokens进行预训练；
• 使用了440张A100x40G的GPU训练了9.5天；
• 使用ALiBi而不是位置编码，除了可以提供外推长度外，也能提高对loss spike的适应性；
• 优化器使用Lion而不是AdamW，能提供更加稳定的更新幅度且减少优化器缓存占用的一半；
• 通过使用FlashAttention、低精度LayerNorm层和FasterTransformer实现更快的推理；

2. 使用的数据

3. 效果

4. 不同版本的模型

​ MPT-7B Base是预训练语言模型，也是其他模型的基础；

​ MPT-7B-StoryWriter-65k+则是用长度为65k+的书籍语料微调的模型，得益于ALiBi的外推能力，其能够外推至84k的tokens；

​ MPT-7B-Instruct是一个经过指令微调的模型；

​ MPT-7B-Chat是一个对话模型，微调数据包括：ShareGPT-Vicuna、HC3、Alpaca、Helpful and Harmless以及Evol-Instruct；

二、模型结构

1. Layer Norm

​ 虽然MPT在推理时，使用的是低精度LayerNorm。但是代码实现中，也提供了RMSNorm。因此，这里会简单介绍其实现的各类Norm。代码位于https://huggingface.co/mosaicml/mpt-7b/blob/main/norm.py。

1.1 标准Layer Norm

​ 在Transformer中LayerNorm是对**(batch_size, seq_length, hidden_size)**中的hidden_size维度进行normalize。具体来说，给定一个向量$\text{x}$，则normalize的过程为：
$$\text{y}=\frac{\text{x}-\text{E}(\text{x})}{\sqrt{\text{Var}(\text{x})+ϵ}}\ast \gamma +\beta$$
其中，$\text{E}(\text{x})$表示向量$\text{x}$的期望，$\text{Var}(\text{x})$是向量$\text{x}$的方差，$ϵ$是为了防止分母为0的偏置项，$\gamma$和$\beta$是两个可学习参数。

def _cast_if_autocast_enabled(tensor):"""自动转换tensor的数据类型，用来实现后续的低精度LayerNorm。"""if torch.is_autocast_enabled():if tensor.device.type == 'cuda':dtype = torch.get_autocast_gpu_dtype()elif tensor.device.type == 'cpu':dtype = torch.get_autocast_cpu_dtype()else:raise NotImplementedError()return tensor.to(dtype=dtype)return tensorclass LPLayerNorm(torch.nn.LayerNorm):"""低精度LayerNorm，也输入和LayerNorm的参数转换为低精度"""def __init__(self, normalized_shape, eps=1e-05, elementwise_affine=True, device=None, dtype=None):super().__init__(normalized_shape=normalized_shape, eps=eps, elementwise_affine=elementwise_affine, device=device, dtype=dtype)def forward(self, x):module_device = x.device# 将输入和LayerNorm的参数都转换为低精度downcast_x = _cast_if_autocast_enabled(x)downcast_weight = _cast_if_autocast_enabled(self.weight) if self.weight is not None else self.weightdowncast_bias = _cast_if_autocast_enabled(self.bias) if self.bias is not None else self.biaswith torch.autocast(enabled=False, device_type=module_device.type):return torch.nn.functional.layer_norm(downcast_x, self.normalized_shape, downcast_weight, downcast_bias, self.eps)

1.2 RMS Norm

​ 相比于LayerNorm，RMS Norm主要是去掉了减均值的部分。
$$\begin{array}{rl}\text{y}& =\frac{\text{x}}{\text{RMS}(\text{x})}\\ \text{RMS}(\text{x})& =\sqrt{\frac{1}{n}\sum _{i=1}^n{x}_i^2}\end{array}$$

def rms_norm(x, weight=None, eps=1e-05):"""RMS的实现"""output = x / torch.rsqrt(x.pow(2).mean(-1, keepdim=True) + eps)if weight is not None:return output * weightreturn outputclass RMSNorm(torch.nn.Module):def __init__(self, normalized_shape, eps=1e-05, weight=True, dtype=None, device=None):super().__init__()self.eps = epsif weight:self.weight = torch.nn.Parameter(torch.ones(normalized_shape, dtype=dtype, device=device))else:self.register_parameter('weight', None)def forward(self, x):return rms_norm(x.float(), self.weight, self.eps).to(dtype=x.dtype)class LPRMSNorm(RMSNorm):"""低精度RMS Norm的实现"""def __init__(self, normalized_shape, eps=1e-05, weight=True, dtype=None, device=None):super().__init__(normalized_shape=normalized_shape, eps=eps, weight=weight, dtype=dtype, device=device)def forward(self, x):downcast_x = _cast_if_autocast_enabled(x)downcast_weight = _cast_if_autocast_enabled(self.weight) if self.weight is not None else self.weightwith torch.autocast(enabled=False, device_type=x.device.type):return rms_norm(downcast_x, downcast_weight, self.eps).to(dtype=x.dtype)
NORM_CLASS_REGISTRY = {'layernorm': torch.nn.LayerNorm, 'low_precision_layernorm': LPLayerNorm, 'rmsnorm': RMSNorm, 'low_precision_rmsnorm': LPRMSNorm}

2. 自注意力

​ 虽然MPT中仍然使用多头注意力，但是其同时实现了"多头注意力"和"多Query注意力"，并且提供了纯torch、flash_attn和triton_flash_attn的实现。这里会介绍所有这些实现，代码位于https://huggingface.co/mosaicml/mpt-7b/blob/main/attention.py。

2.1 三个版本的Self-Attention实现

​ 自注意力机制的原理如下：
$$\begin{array}{rl}Q& =W_{q}X\\ K& =W_{k}X\\ V& =W_{v}X\\ \text{Attention}(Q,K,V)& =\text{softmax}(\frac{Q{K}^T}{\sqrt{d_k}})V\end{array}$$
其中，$X$是输入， $W_q$、$W_k$和$W_v$分别是query、key、value的投影矩阵。

纯torch版本

def scaled_multihead_dot_product_attention(query, # [batch_size, seq_length, (head_num*head_dim)]key, # 在Multi-Head Attention中形状同query，在Multi-Query Attention中head_num=1value, # 在Multi-Head Attention中形状同query，在Multi-Query Attention中head_num=1n_heads, # 注意力头的数量softmax_scale=None, # 注意力分数的缩放因子attn_bias=None, # 注意力分数的额外偏置项key_padding_mask=None,is_causal=False, # 是否为causal，决定了注意力mask的类型dropout_p=0.0, multiquery=False, # 是否是Mulit-Query Attentiontraining=False, needs_weights=False):# 通过rearrange函数重塑query、key、value的形状# 这里的重塑过程同时包括了拆分多头以及调整维度的顺序q = rearrange(query, 'b s (h d) -> b h s d', h=n_heads)# 在Mulit-Query Attention中, h=1，也就是多个头共享同一份key和valuek = rearrange(key, 'b s (h d) -> b h d s', h=1 if multiquery else n_heads)v = rearrange(value, 'b s (h d) -> b h s d', h=1 if multiquery else n_heads)min_val = torch.finfo(q.dtype).min# b: batch_size, s_q: seq_length, d: head_dim(b, _, s_q, d) = q.shapes_k = k.size(-1)# 若没有指定softmax的缩放因子，则是默认的因子if softmax_scale is None:softmax_scale = 1 / math.sqrt(d)# 得到未经过softmax规范化的注意力权重attn_weight = q.matmul(k) * softmax_scale# 将注意力偏置项添加到注意力分数attn_weight上if attn_bias is not None:if attn_bias.size(-1) != 1 and attn_bias.size(-1) != s_k or (attn_bias.size(-2) != 1 and attn_bias.size(-2) != s_q):raise RuntimeError(f'attn_bias (shape: {attn_bias.shape}) is expected to broadcast to shape: {attn_weight.shape}.')attn_weight = attn_weight + attn_bias# 应用key padding maskif key_padding_mask is not None:if attn_bias is not None:warnings.warn('Propogating key_padding_mask to the attention module ' + 'and applying it within the attention module can cause ' + 'unneccessary computation/memory usage. Consider integrating ' + 'into attn_bias once and passing that to each attention ' + 'module instead.')attn_weight = attn_weight.masked_fill(~key_padding_mask.view((b, 1, 1, s_k)), min_val)# 是否使用causal mask(关于causal mask可以见文章https://zhuanlan.zhihu.com/p/625911234)if is_causal:s = max(s_q, s_k)# 全1矩阵causal_mask = attn_weight.new_ones(s, s, dtype=torch.float16)causal_mask = causal_mask.tril()causal_mask = causal_mask.to(torch.bool)causal_mask = ~causal_maskcausal_mask = causal_mask[-s_q:, -s_k:]attn_weight = attn_weight.masked_fill(causal_mask.view(1, 1, s_q, s_k), min_val)# softmaxattn_weight = torch.softmax(attn_weight, dim=-1)# dropoutif dropout_p:attn_weight = torch.nn.functional.dropout(attn_weight, p=dropout_p, training=training, inplace=True)out = attn_weight.matmul(v)out = rearrange(out, 'b h s d -> b s (h d)')if needs_weights:return (out, attn_weight)return (out, None)

FlashAttention

​ 在标准的自注意力计算中，有一些中间计算结果$S$和$P$：
$$\begin{gathered} S=Q{K}^{\top } \\ P=\text{Softmax}(S) \end{gathered}$$

这些中间结果需要通过高带宽内存(HBM)进行存取，FlashAttention主要是通过减少对HBM的访问次数来优化速度。这里不展开介绍该技术的实现，其通过调用包flash_attn即可实现，在MPT的代码中对其进行了简单的封装flash_attn_fn。

TritonFlashAttention

​ 使用Triton实现的FlashAttention，也可以通过调用flash_attn来实现，在MPT代码中封装为triton_flash_attn_fn。

2.2 多头注意力

​ 多头注意力就是将多个自注意力的结果合并，如下：
$$\begin{array}{rl}{\text{head}}_i& =\text{Attention}(Q_i,K_i,V_i,A_i)\\ \text{MultiHead}(Q,K,V,A)& =\text{Concat}({\text{head}}_1,\dots ,{\text{head}}_h)W_o\end{array}$$

class MultiheadAttention(nn.Module):def __init__(self, d_model: int, n_heads: int, attn_impl: str='triton', clip_qkv: Optional[float]=None, qk_ln: bool=False, softmax_scale: Optional[float]=None, attn_pdrop: float=0.0, low_precision_layernorm: bool=False, device: Optional[str]=None):super().__init__()self.attn_impl = attn_impl #注意力的实现方式self.clip_qkv = clip_qkv # query、key、value是否进行clipself.qk_ln = qk_ln # 是否对query和key进行LayerNormself.d_model = d_modelself.n_heads = n_heads # 注意力头的数量self.softmax_scale = softmax_scale # softmax缩放因子if self.softmax_scale is None:self.softmax_scale = 1 / math.sqrt(self.d_model / self.n_heads)self.attn_dropout_p = attn_pdrop# 投影层self.Wqkv = nn.Linear(self.d_model, 3 * self.d_model, device=device)fuse_splits = (d_model, 2 * d_model)self.Wqkv._fused = (0, fuse_splits)if self.qk_ln:# 初始化query和key的LayerNormlayernorm_class = LPLayerNorm if low_precision_layernorm else nn.LayerNormself.q_ln = layernorm_class(self.d_model, device=device)self.k_ln = layernorm_class(self.d_model, device=device)if self.attn_impl == 'flash':self.attn_fn = flash_attn_fnelif self.attn_impl == 'triton':self.attn_fn = triton_flash_attn_fnelif self.attn_impl == 'torch':self.attn_fn = scaled_multihead_dot_product_attentionelse:raise ValueError(f'attn_impl={attn_impl!r} is an invalid setting.')# 输出投影层self.out_proj = nn.Linear(self.d_model, self.d_model, device=device)self.out_proj._is_residual = Truedef forward(self, x, past_key_value=None, attn_bias=None, attention_mask=None, is_causal=True, needs_weights=False):# 投影获得query、key、valueqkv = self.Wqkv(x)# 对query、key和value进行clip(MPT并没有使用该选项)if self.clip_qkv:qkv.clamp_(min=-self.clip_qkv, max=self.clip_qkv)# 将query、key和value中qkv中拆分出来(query, key, value) = qkv.chunk(3, dim=2)key_padding_mask = attention_maskif self.qk_ln:# 对query和key进行LayerNorm(MPT并没有使用该选项)dtype = query.dtypequery = self.q_ln(query).to(dtype)key = self.k_ln(key).to(dtype)# 推理时，会将前面token的key和value传递过来，这里进行合并if past_key_value is not None:if len(past_key_value) != 0:key = torch.cat([past_key_value[0], key], dim=1)value = torch.cat([past_key_value[1], value], dim=1)past_key_value = (key, value)if attn_bias is not None:attn_bias = attn_bias[:, :, -query.size(1):, -key.size(1):]# 执行自注意力(context, attn_weights) = self.attn_fn(query, key, value, self.n_heads, softmax_scale=self.softmax_scale, attn_bias=attn_bias, key_padding_mask=key_padding_mask, is_causal=is_causal, dropout_p=self.attn_dropout_p, training=self.training, needs_weights=needs_weights)return (self.out_proj(context), attn_weights, past_key_value)

2.3 多Query注意力

​ 在多头注意力中，每个头都有自己独立的query、key和value。多Query注意力中，各个头之间共享key和value。所以，多Query注意力的参数量要比多头注意力少。多Query注意力和多头注意力的实现非常类似，仅有少量的不同，这里仅展示代码中不同的部分：

class MultiQueryAttention(nn.Module):def __init__(self, d_model: int, n_heads: int, attn_impl: str='triton', clip_qkv: Optional[float]=None, qk_ln: bool=False, softmax_scale: Optional[float]=None, attn_pdrop: float=0.0, low_precision_layernorm: bool=False, device: Optional[str]=None):super().__init__()...# 计算每个头的维度self.head_dim = d_model // n_heads...# 投影层与多头注意力具有明显区别，多头注意力的输出维度为3*d_model，而这里是d_model + 2*head_dimself.Wqkv = nn.Linear(d_model, d_model + 2 * self.head_dim, device=device)...

3. 位置偏差ALiBi

​ MPT使用ALiBi的方式向模型注入位置信息。ALiBi注入位置信息的方式是在注意力分数矩阵添加一个偏差(bias)来实现的。具体来说，给定一个长度为L的输入序列， 那么每个注意力头的第i个query ${\text{q}}_i\in {\mathbb{R}}^{1\times d}(1\le i\le L)$针对前i个key $\text{K}\in {\mathbb{R}}^{i\times d}$的注意力分数为

$$\text{softmax}({\text{q}}_i{\text{K}}^{\top })$$
在使用ALiBi时，不需要向网络添加位置嵌入。仅需要在query-key点积中添加静态偏差即可。
$$\text{softmax}({\text{q}}_i{\text{K}}^{\top }+m\cdot [-(i-1),\dots ,-2,-1,0])$$
其中m是与注意力头相关的斜率(slope)，也就是超参；$[-(i-1),\dots ,-2,-1,0]$其实就是${\text{q}}_i$与各个key的相对距离。

​ 对于8个注意力头，m是等比序列：$\frac{1}{2^1},\frac{1}{2^2},\dots ,\frac{1}{2^8}$。对于16个注意力头的模型，m则是等比序列：$\frac{1}{2^{0.5}},\frac{1}{2^1},\frac{1}{2^{1.5}},\dots ,\frac{1}{2^8}$。

def gen_slopes(n_heads, alibi_bias_max=8, device=None):"""计算斜率"""# _n_heads是与n_heads接近的2的次数，例如：n_heads为5/6/7时，_n_heads为8_n_heads = 2 ** math.ceil(math.log2(n_heads))m = torch.arange(1, _n_heads + 1, dtype=torch.float32, device=device)# m是alibi_bias_max/_n_heads到alibi_bias_max的等差数列m = m.mul(alibi_bias_max / _n_heads)# 计算斜率slopes = 1.0 / torch.pow(2, m)if _n_heads != n_heads:slopes = torch.concat([slopes[1::2], slopes[::2]])[:n_heads]return slopes.view(1, n_heads, 1, 1)def build_alibi_bias(n_heads, seq_len, full=False, alibi_bias_max=8, device=None, dtype=None):"""构建alibi注意力偏差"""alibi_bias = torch.arange(1 - seq_len, 1, dtype=torch.int32, device=device).view(1, 1, 1, seq_len)if full:alibi_bias = alibi_bias - torch.arange(1 - seq_len, 1, dtype=torch.int32, device=device).view(1, 1, seq_len, 1)alibi_bias = alibi_bias.abs().mul(-1)slopes = gen_slopes(n_heads, alibi_bias_max, device=device)alibi_bias = alibi_bias * slopesreturn alibi_bias.to(dtype=dtype)def build_attn_bias(attn_impl, attn_bias, n_heads, seq_len, causal=False, alibi=False, alibi_bias_max=8):"""该函数只是对`build_alibi_bias`进行了封装，由于flash版本的注意力不支持attn_bias所以返回none"""if attn_impl == 'flash':return Noneelif attn_impl in ['torch', 'triton']:if alibi:(device, dtype) = (attn_bias.device, attn_bias.dtype)attn_bias = attn_bias.add(build_alibi_bias(n_heads, seq_len, full=not causal, alibi_bias_max=alibi_bias_max, device=device, dtype=dtype))return attn_biaselse:raise ValueError(f'attn_impl={attn_impl!r} is an invalid setting.')

4. MLP层

$$\text{MLP}(X)=\text{GELU}(X{W}_1)W_2$$

class MPTMLP(nn.Module):def __init__(self, d_model: int, expansion_ratio: int, device: Optional[str]=None):super().__init__()self.up_proj = nn.Linear(d_model, expansion_ratio * d_model, device=device)self.act = nn.GELU(approximate='none')self.down_proj = nn.Linear(expansion_ratio * d_model, d_model, device=device)self.down_proj._is_residual = Truedef forward(self, x):return self.down_proj(self.act(self.up_proj(x)))

5. MPTBlock层

class MPTBlock(nn.Module):def __init__(self, d_model: int, n_heads: int, expansion_ratio: int, attn_config: Dict={'attn_type': 'multihead_attention', 'attn_pdrop': 0.0, 'attn_impl': 'triton', 'qk_ln': False, 'clip_qkv': None, 'softmax_scale': None, 'prefix_lm': False, 'attn_uses_sequence_id': False, 'alibi': False, 'alibi_bias_max': 8}, resid_pdrop: float=0.0, norm_type: str='low_precision_layernorm', device: Optional[str]=None, **kwargs):del kwargssuper().__init__()# LayerNorm的类别norm_class = NORM_CLASS_REGISTRY[norm_type.lower()]# 注意力的类别attn_class = ATTN_CLASS_REGISTRY[attn_config['attn_type']]# LayerNormself.norm_1 = norm_class(d_model, device=device)# 注意力self.attn = attn_class(attn_impl=attn_config['attn_impl'], clip_qkv=attn_config['clip_qkv'], qk_ln=attn_config['qk_ln'], softmax_scale=attn_config['softmax_scale'], attn_pdrop=attn_config['attn_pdrop'], d_model=d_model, n_heads=n_heads, device=device)# LayerNormself.norm_2 = norm_class(d_model, device=device)# MLP层self.ffn = MPTMLP(d_model=d_model, expansion_ratio=expansion_ratio, device=device)self.resid_attn_dropout = nn.Dropout(resid_pdrop)self.resid_ffn_dropout = nn.Dropout(resid_pdrop)def forward(self, x: torch.Tensor, past_key_value: Optional[Tuple[torch.Tensor]]=None, attn_bias: Optional[torch.Tensor]=None, attention_mask: Optional[torch.ByteTensor]=None, is_causal: bool=True) -> Tuple[torch.Tensor, Optional[Tuple[torch.Tensor]]]:a = self.norm_1(x)(b, _, past_key_value) = self.attn(a, past_key_value=past_key_value, attn_bias=attn_bias, attention_mask=attention_mask, is_causal=is_causal)x = x + self.resid_attn_dropout(b)m = self.norm_2(x)n = self.ffn(m)x = x + self.resid_ffn_dropout(n)return (x, past_key_value)

6. MPTModel

​ MPTModel的代码比较多，这里就不贴完整代码了，介绍一下主要的结构和值得关注的点。

Shrink Embedding Gradient

​ Shrink Embedding Gradient技术来自于GLM-130B，用于稳定预训练，防止出现loss峰值。MPT在代码中也支持该技术，但并没有使用

# 当self.embedding_fraction小于1时，执行该段代码
# x接收到的梯度缩减为self.embedding_fraction
x_shrunk = x * self.embedding_fraction + x.detach() * (1 - self.embedding_fraction)
assert isinstance(self.emb_drop, nn.Module)
# embedding dropout
x = self.emb_drop(x_shrunk)

模型结构相关的代码

# embedding
tok_emb = self.wte(input_ids)
x = tok_emb
# embedding dropout
x = self.emb_drop(x)
# 计算alibi的注意力偏差
(attn_bias, attention_mask) = self._attn_bias(device=x.device, dtype=x.dtype, attention_mask=attention_mask, prefix_mask=prefix_mask, sequence_id=sequence_i)
# 多个block的前向传播
for (b_idx, block) in enumerate(self.blocks):past_key_value = past_key_values[b_idx] if past_key_values is not None else None(x, past_key_value) = block(x, past_key_value=past_key_value, attn_bias=attn_bias, attention_mask=attention_mask, is_causal=self.is_causal)if past_key_values is not None:past_key_values[b_idx] = past_key_value
# 最后的输出进行LayerNorm
x = self.norm_f(x)

参考资料

https://www.mosaicml.com/blog/mpt-7b#training-stability

https://huggingface.co/mosaicml/mpt-7b/tree/main