【LLM教程】405B大模型如何炼成?

• Preview

• 1. Introduction

• 2. 公式推导

• 2.1 前向计算

• 2.2 损失函数计算

• 2.3 反向梯度推导

• 3. 手撕TP

• 3.1 常规线性层计算

• 3.2 手撕Col-Wise TP

• 3.3 手撕Row-Wise TP

• 4. 手撕分布式TP

• 5. Llama-3 TP层分析

• 5.1 MLP层分析

• 5.2 Attention层分析

• 6. 总结

• Reference

Preview

开个新坑，主要从代码角度来解析主流分布式训练技术(TP/DP/PP/CP…)

1. Introduction

随着模型参数规模增大(Llama-3.1-405B存储810GB(fp16))，如何高效训练模型成为当前的难题。一个简要的手段是将模型进行切分，存储在不同的计算设备中进行分布式计算。

如下图：一个140GB模型无法在一个40GB的显存上加载，那么我们可以切分模型成4个35GB模型，分别用4个40G的模型进行加载：

Megatron-LM提出一种张量并行(Tensor Parallelism, 下文简称TP)， 能够对参数矩阵(Tensor)进行切块分布在其他设备上，从而达到切割模型的作用。对于参数矩阵的切分策略我们通常有横切(Row-Wise)和竖切(Col-Wise)，而不同的切法将产生不同的计算复杂度。

为了方便理解，不依赖DeepSpeed和Megatron-LM代码基础下，我们将重点讲解：

1. Row-Wise和Col-wise切分的前向和反向梯度推导

2. 最小化TP代码实现

3. 基于torch.distributed实现TP版本的线性层

4. 分析TP版本Llama3实现

2. 公式推导

2.1 前向计算

对于参数矩阵, 其中为行数， 为列数，对于线性层我们有如下计算：

其中为输入矩阵， 为数据条目, 为了方便计算我们设定, 设定GPU数量为2个

下图所示左上角为常规的线性层前向计算。

我们对于可以进行列切割(Col-Wise 图右)和行切割(row-wise 图左下)

1. 列切割 ， 表示列切割索引，我们将不同的参数矩阵在不同GPU上加载，同时输入要传输至不同的GPU中。此时前向计算为：

列切割下,各卡结果 拼接得到的 为前向计算结果

2. 行切割, , 注意到输入需对应进行切割, 并放置在不同的GPU上，此时前向计算为:

   行切割下，各卡相加后为前向计算结果

这里三种前向计算的结果是一样的

我们简要分析通信量：

|
| Row-wise | Col-wise |
| — | — | — |
| 通信量 | 44 | 42 |
| 计算量 | 400 | 720 |

a. Row-wise:

1. 通信量: 输入X(2x3) x 2 + 输出Y(2x10) x 2 = 12+40 = 42

2. 计算量: 乘法(2x3x3x10) x 2=360 + 额外的加法 2x10x2=40 -> 400

b. Col-wise：

1. 通信量: 输入X(2x6) x 2+ 输出Y(2x5) x 2 = 24+20 = 44

2. 计算量：乘法(2x6x6x5)x2=720

注：提供计算思路，反应通信和计算数据量

2.2 损失函数计算

我们使用均方差损失来计算误差，给定标签, 损失计算如下:

其中，差值为

2.3 反向梯度推导

2.3.1 Col-Wise TP

前向计算

反向计算

其中由于 之间独立，所以可以分别计算误差，从而求出梯度

2.3.2 Row-Wise TP

前向计算

反向计算

3. 手撕TP

我们使用标准线性层、Col-Wise TP和Row-Wise TP进行反向传播，可以得到相同的梯度。

3.1 常规线性层计算

准备数据

import torch   bs = 2    row = dim = 4   col = out_dim = dim * 2   x = torch.arange(bs * dim, dtype = torch.float32).reshape(bs, dim)    w = torch.arange(dim * out_dim, dtype = torch.float32).reshape(dim, out_dim)   y_label = torch.randn(bs, out_dim)   y_pred = torch.randn(bs, out_dim)   

前向计算和梯度计算

y_pred = x @ w   delta_y = y_pred - y_label   delta_w = x.t() @ delta_y / out_dim    print(delta_w)   

得到

tensor([[151.8721, 162.4626, 173.4557, 185.0397, 195.8915, 207.5966, 218.1116,            229.2710],           [203.9184, 217.9152, 232.3963, 247.5970, 261.8807, 277.3813, 291.0889,            305.9322],           [255.9646, 273.3677, 291.3369, 310.1544, 327.8700, 347.1660, 364.0662,            382.5934],           [308.0109, 328.8203, 350.2776, 372.7117, 393.8593, 416.9507, 437.0434,            459.2547]])   

3.2 手撕Col-Wise TP

前向计算

w_col = w   col_dim = out_dim // 2   w_col_1 = w_col[:, :col_dim]   w_col_2 = w_col[:, col_dim:]   y_1 = x @ w_col_1   y_2 = x @ w_col_2   

反向计算

y_1_delta = y_1 - y_label[:, col_dim:]   y_2_delta = y_2 - y_label[:, :col_dim]   grad_col_1 = x.t() @ y_1_delta   grad_col_2 = x.t() @ y_2_delta    grad_col = torch.cat( (grad_col_1, grad_col_2), dim = 1) / out_dim   

计算结果为:

tensor([[151.8915, 163.5966, 174.1116, 185.2710, 195.8721, 206.4626, 217.4557,            229.0397],           [203.8807, 219.3813, 233.0889, 247.9322, 261.9184, 275.9152, 290.3964,            305.5970],           [255.8700, 275.1660, 292.0662, 310.5934, 327.9646, 345.3677, 363.3370,            382.1544],           [307.8593, 330.9507, 351.0434, 373.2547, 394.0109, 414.8203, 436.2776,            458.7117]])   

3.3 手撕Row-Wise TP

前向计算

row_dim = dim // 2   w_row_1 = w_row[:row_dim, :]   w_row_2 = w_row[row_dim:, :]   x_col_1 = x[:, :row_dim]   x_col_2 = x[:, row_dim:]   y_1 = x[:, :row_dim] @ w_row_1   y_2 = x[:, row_dim:] @ w_row_2   

反向计算

delta_y = (y_1 + y_2 - y_label)   grad_row_1 = x_col_1.t() @ delta_y    grad_row_2 = x_col_2.t() @ delta_y    grad_row = torch.cat((grad_row_1,grad_row_2), dim = 0) / out_dim   print(grad_row)   

输出为

tensor([[151.8915, 163.5966, 174.1116, 185.2710, 195.8721, 206.4626, 217.4557,            229.0397],           [203.8807, 219.3813, 233.0889, 247.9322, 261.9184, 275.9152, 290.3964,            305.5970],           [255.8700, 275.1660, 292.0662, 310.5934, 327.9646, 345.3677, 363.3370,            382.1544],           [307.8593, 330.9507, 351.0434, 373.2547, 394.0109, 414.8203, 436.2776,            458.7117]])   

4. 手撕分布式TP

笔者实现了基于torch.distributed的手动求导版本的Col-Wise的TP

1. 将权重scatter_w()到每个GPU上，self.w是部分权重

2. 计算过程forward,backward和update均不需要与其他gpu进行交互

3. 当模型训练好时，通过gather_w可以聚合所有权重，使用torch.cat合并成完整权重

class ColWiseParallel():       def __init__(self, dim=6, dim_out=10, n_gpus=2, device_id=0) -> None:           self.dim = dim           self.dim_out = dim_out           self.n_gpus = n_gpus           self.col_per_device = dim_out // n_gpus           self.w = torch.zeros(dim, self.col_per_device).to(device_id)           dprint(self.w.shape)           self.lr = 0.01          def scatter_w(self, device_id):           if is_main_device():               w_all = torch.randn(self.dim, self.dim_out).to(device_id)               w_list = torch.split(w_all, self.col_per_device, dim=1)               w_list = list(w_list)           else:               w_list = None           dist.scatter(tensor=self.w, scatter_list=w_list, src=0)          def forward(self, x):           y = x @ self.w           return y          def backward(self, x, error):           grad = x.t() @ error * (1/self.dim_out)           return grad          def update(self, grad):           self.w -= self.lr * grad          def gather_w(self, device_id):           if is_main_device():               w_list = [torch.zeros_like(self.w) for _ in range(self.n_gpus)]           else:               w_list = None           dist.gather(self.w, w_list, dst=0)           return w_list   

调用代码如下：

model = ColWiseParallel(dim=6, dim_out=10, n_gpus=n_gpus, device_id=rank)   model.scatter_w(rank)   x = torch.randn(bs, in_dim).to(rank)   y = model.forward(x)   e = y - y_label   grad = model.backward(x, e)   model.update(grad)   w_dist = model.gather_w(rank)   

5. Llama-3 TP层分析

实际上只要了解Row-Wise和Col-Wise TP后，就可以自由组合不同的并行。

5.1 MLP层分析

MLP里存在两个矩阵WA(up)和WB(down)，涉及到大规模的矩阵乘法计算，我们分析不同的TP策略：

• Row-Wise First：在遇到GeLU前需要gather所有卡上的激活者，增加通信消耗

• Col-Wise Firse：可以独立计算GeLU，不需要再Gather，而后再用Row-Wise的TP

  由于TP不改变计算量，切分的原则是尽可能减少gather操作，减少通信消耗

进一步分析Llama-3 代码

# llama3/llama/model.py   class FeedForward(nn.Module):       def __init__(           ....       ):           super().__init__()           ...           # 使用TP版本的Linear 代替 nn.Linear           self.w1 = ColumnParallelLinear(               dim, hidden_dim, ...           ) # activation branch           self.w2 = RowParallelLinear(               hidden_dim, dim, ...           ) # down           self.w3 = ColumnParallelLinear(               dim, hidden_dim, ...           ) # up          def forward(self, x):         	# 原版         	# return self.w2(F.silu(self.w1(x)) * self.w3(x))                      # 便于分析         	y1 = F.silu(self.w1(x)) * self.w3(x) # col-wise           y2 = self.w2(y)											 # row-wise           return y2   

5.2 Attention层分析

在Llama3中, wq, wk, wv都用Col-Wise来做TP，此时好处是每个GPU都能算单头注意力。Attenion计算完后可类似MLP层处理，将wo 用Row-Wise做TP。

# llama3/llama/model.py   class Attention(nn.Module):       def __init__(self, args: ModelArgs):           super().__init__()           # ...              self.wq = ColumnParallelLinear(               #....           )           self.wk = ColumnParallelLinear(               #....           )           self.wv = ColumnParallelLinear(               #....           )           self.wo = RowParallelLinear(               #...           )       #....   

6. 总结

1. 本文推导了TP下前向和反向计算，并手撕了TP实现，基于torch.distributed重写了Col-Wise TP

2. 验证Col-Wise和Row-Wise 的参数更新梯度一致

3. 分析Llama-3的MLP和Attenion TP实现，需要Col-Wise First计算

---

如何学习大模型 AI ？

由于新岗位的生产效率，要优于被取代岗位的生产效率，所以实际上整个社会的生产效率是提升的。

但是具体到个人，只能说是：

“最先掌握AI的人，将会比较晚掌握AI的人有竞争优势”。

这句话，放在计算机、互联网、移动互联网的开局时期，都是一样的道理。

我在一线互联网企业工作十余年里，指导过不少同行后辈。帮助很多人得到了学习和成长。

我意识到有很多经验和知识值得分享给大家，也可以通过我们的能力和经验解答大家在人工智能学习中的很多困惑，所以在工作繁忙的情况下还是坚持各种整理和分享。但苦于知识传播途径有限，很多互联网行业朋友无法获得正确的资料得到学习提升，故此将并将重要的AI大模型资料包括AI大模型入门学习思维导图、精品AI大模型学习书籍手册、视频教程、实战学习等录播视频免费分享出来。

😝有需要的小伙伴，可以Vx扫描下方二维码免费领取🆓

👉1.大模型入门学习思维导图👈

要学习一门新的技术，作为新手一定要先学习成长路线图，方向不对，努力白费。

对于从来没有接触过AI大模型的同学，我们帮你准备了详细的学习成长路线图&学习规划。可以说是最科学最系统的学习路线，大家跟着这个大的方向学习准没问题。（全套教程文末领取哈）

👉2.AGI大模型配套视频👈

很多朋友都不喜欢晦涩的文字，我也为大家准备了视频教程，每个章节都是当前板块的精华浓缩。

👉3.大模型实际应用报告合集👈

这套包含640份报告的合集，涵盖了AI大模型的理论研究、技术实现、行业应用等多个方面。无论您是科研人员、工程师，还是对AI大模型感兴趣的爱好者，这套报告合集都将为您提供宝贵的信息和启示。（全套教程文末领取哈）

👉4.大模型落地应用案例PPT👈

光学理论是没用的，要学会跟着一起做，要动手实操，才能将自己的所学运用到实际当中去，这时候可以搞点实战案例来学习。（全套教程文末领取哈）

👉5.大模型经典学习电子书👈

随着人工智能技术的飞速发展，AI大模型已经成为了当今科技领域的一大热点。这些大型预训练模型，如GPT-3、BERT、XLNet等，以其强大的语言理解和生成能力，正在改变我们对人工智能的认识。 那以下这些PDF籍就是非常不错的学习资源。

👉6.大模型面试题&答案👈

截至目前大模型已经超过200个，在大模型纵横的时代，不仅大模型技术越来越卷，就连大模型相关的岗位和面试也开始越来越卷了。为了让大家更容易上车大模型算法赛道，我总结了大模型常考的面试题。

这份完整版的大模型 AI 学习资料已经上传CSDN，朋友们如果需要可以微信扫描下方CSDN官方认证二维码免费领取【保证100%免费】

😝有需要的小伙伴，可以Vx扫描下方二维码免费领取🆓