在上一篇文章「LLaMA 大模型并行计算原理图解：从数据并行到张量并行的演进与对比」中我们详细介绍了张量平行计算的原理，并通过图示展示了其计算过程。在接下来的这篇文章当中，我们将详细给大家介绍张量平行在 fairscale 库中的实现过程，以及如何在 CPU 上来模拟整个运行过程。
下面，我们依旧按着先学会用，再探究为什么的原则来进行叙述。
首先，我们来介绍一下 fairscale 库。
关键词：LLaMA、张量并行、ColumnParallelLinear、RowParallelLinear、FairScale
1. fairscale 并行计算库#
FairScale 是 Meta 于2021年 [1] 提出的一个用于高性能、大规模训练的 PyTorch 扩展库。从名字可以看出，Fair Scale，“公平的规模” 确实有那么一点并行的味道。
1.1 fairscale 库介绍#
FairScale 是在 PyTorch 的基础功能之上进行了扩展，同时引入了新的最先进的模型扩展技术。FairScale 通过可组合的模块和易于使用的 API，提供了当前最前沿的分布式训练方法。这些 API 是研究人员在资源受限的情况下尝试扩展模型时的重要工具之一。
FairScale 的设计理念包括以下几个方面：
-
可用性（Usability） —— 用户应该能够在认知负担最小的情况下理解并使用 FairScale 的 API。
-
模块化（Modularity） —— 用户应该能够将多个 FairScale API 无缝整合到自己的训练循环中。
-
性能（Performance） —— FairScale 的 API 在扩展能力和效率方面提供最优性能。
以上描述摘自 FairScale 项目介绍。
1.2 安装 fairscale 库#
在介绍完 FairScale 后我们再来看如何安装这个库。首先，如果是在 CPU 上安装的话，可以直接使用 pip install fairscal==0.4.13 进行安装。这里安装的是最新的 0.4.13 版本，你也可以根据自己的实际情况调整。
本来按道理，如果是在 GPU 上安装，上面的 pip 命令应该也没问题，结果我们在安装过程中反复提示如下报错：
could not find a version that satisfies the requirement setuptools>=40.6.2. 但是检查后发现当前 setuptools 版本为 68.2.2 ，不存在上诉问题。又经过一番排查，才在官网看到了解决方案。所以，无论你是在 CPU 或者 GPU 环境下安装，若果出现类似上面错误，都可以按照下面的方法解决。
方案一：安装时加上--no-build-isolation
pip install fairscal==0.4.13 --no-build-isolation方案二：克隆或下载官方仓库 [1] 后，进入仓库目录输入python setup.py install 进行安装
当然，我们正式使用这个库肯定是在 GPU 上使用，CPU 上仅供学习，模拟。
1.3 CPU 上模拟运行#
在完成 fairscale 库的安装以后，我们便可以通过如下代码完成在 CPU 上的模拟运行，完整示例代码可参见 Code/C04_Parallel/parallel_linear_cpu.py 文件。
① 导入模块
1 import os
2 import torch
3 import torch.distributed as dist
4 import torch.multiprocessing as mp
5 from fairscale.nn.model_parallel.initialize import initialize_model_parallel
6 from fairscale.nn.model_parallel import ColumnParallelLinear
7 from fairscale.nn.model_parallel import RowParallelLinear在上述代码中，第3行是后续将使用分布式方式来运行整个示例代码。第4行是后续将使用 spawn 模块来实现 CPU 山的多进程模拟。第5~7行是载入并行计算时需要用到模块。
② 设置环境变量
1 def run(rank, world_size):
2 print(f"Rank: {rank} initializing model parallel with size {world_size}")
3 os.environ['MASTER_ADDR'] = 'localhost'
4 os.environ['MASTER_PORT'] = '12355'
5 os.environ['RANK'] = str(rank)
6 os.environ['WORLD_SIZE'] = str(world_size)
7 dist.init_process_group(backend="gloo", rank=rank, world_size=world_size)
8 initialize_model_parallel(world_size)在上述代码中，第1行是定义一个函数 run()，它在每个进程中运行；其中 rank 表示当前进程的全局编号，取值范围是 [0, world_size - 1]，world_size 是总共参与分布式训练的进程数。第2行则是打印出在当前进程中自己的编号，在后续输出结果中我们可以看到。第3行是设置主节点地址为 'localhost'，适用于单机多进程，多机训练时应设为主节点的 IP。第4行是设置进程间通信使用的端口，所有进程都必须设置为同一个端口。第5行是设置当前进程的 rank（编号）。
这里需要注意的一点是，虽然这些子进程都运行在同一台主机上，但后续mp.spawn() 会复制当前 Python 进程的代码和环境，然后为每个进程传入一个独立的 rank 参数。在每个子进程中，os.environ['RANK'] = str(rank) 是在子进程自己的内存空间中执行的赋值操作，即os.environ 实际是当前进程的一个本地字典，它并不共享（或覆盖）其他进程的环境变量。这和平时我们理解的 OS 里的环境变量有点区别。
第7行初始化 PyTorch 的分布式进程组，让所有进程进入通信准备状态，其中 backend="gloo" 表示使用 CPU 通信后端（可跨平台），若使用 GPU，则通常用 "nccl"。第8行是初始化模型并行所需的辅助结构（例如张量并行的组划分等），意思是把全局通信组再切分成多个模型并行子组，准备 ColumnParallelLinear 等模块使用。
下面便是关于两个参数的相关介绍：
| 选项 | 特性 | 描述 |
|---|---|---|
| gloo | 类型 | 通用通信后端，Facebook 的 CPU/GPU 通信库 |
| 支持设备 | CPU / GPU 都支持 | |
| 通信方式 | 使用 TCP/IP 或共享内存 | |
| 性能 | 相比 NCCL，在 GPU 上慢很多 | |
| 推荐场景 | 本地 CPU 模拟测试、没有 GPU 的机器、调试环境、小规模训练 | |
| nccl | 类型 | 高性能 GPU 通信库（NVIDIA 出品） |
| 支持设备 | GPU 专用 | |
| 通信方式 | 使用 CUDA-aware 的通信（P2P, RDMA） | |
| 性能 | 非常快，尤其是在多卡（单机/多机）训练时 | |
| 推荐场景 | 多 GPU 的训练，包括单机多卡和多机多卡 |
③ 模型定义
在完成上面的准备工作以后，我们继续在 run() 函数里定义一个简单的两层神经网络，完成此次并行过程，示例代码如下：
1 batch_size,input_dim = 2, 3
2 hidden_dim,output_dim = 4, 5
3 x = torch.randn(batch_size, input_dim)
4 layer1 = ColumnParallelLinear(in_features=input_dim,
5 out_features=hidden_dim, gather_output=False)
6 layer2 = RowParallelLinear(in_features=hidden_dim,
7 input_is_parallel=True, out_features=output_dim)
8 o1 = layer1(x)
9 o2 = layer2(o1)
10 print(f"Rank: {rank} output shape: {o2.shape}")
11 print(f"Rank: {rank} output shape:\n {o2}")
12 dist.destroy_process_group()在上述代码中，第4~5行便是实例化一个 ColumnParallelLinear 类对象，gather_output=False 表示不讲最后分布在各个进程（GPU）上的结果拼接到一起。第6~7行则是实例化一个 RowParallelLinear 类对象，input_is_parallel=True 表示指定输入需要分布在各个进程上，也就是和上面的 gather_output=False 配合使用。第8~11行便是计算整个前向传播过程，并输出最后的结果。第12行是销毁整个进程组回收资源。
同时，这里需要注意的一点是，因为模型并行的组件通常跨进程工作，所以不能用单进程的 Sequential 封装，即不能通过一下方式来封装：
net = nn.Sequential(ColumnParallelLinear(...), RowParallelLinear(...))④ 结果分析
在完成上述准备工作以后，我们可以通过如下代码运行：
1 def main():
2 world_size = 2 # 模拟2卡
3 mp.spawn(run, args=(world_size,), nprocs=world_size, join=True)
4
5 if __name__ == "__main__":
6 main()上述代码运行结束后，便会输出类似如下结果：
1 当前进程 PID: 21544，父进程 PPID: 21535
2 Rank: 0 initializing model parallel with size 2
3 当前进程 PID: 21545，父进程 PPID: 21535
4 Rank: 1 initializing model parallel with size 2
5 > initializing model parallel with size 2
6 > initializing ddp with size 1
7 > initializing pipeline with size 1
8 Rank: 0 output shape: torch.Size([2, 5])
9 Rank: 1 output shape: torch.Size([2, 5])
10 Rank: 0 output shape:
11 tensor([[-0.8471, -0.4116, 0.1244, 0.0042, 1.1281],
12 [ 1.0552, 1.6335, 1.5514, -0.1004, -2.1235]], grad_fn=<AddBackward0>)
13 Rank: 1 output shape:
14 tensor([[-0.8471, -0.4116, 0.1244, 0.0042, 1.1281],
15 [ 1.0552, 1.6335, 1.5514, -0.1004, -2.1235]], grad_fn=<AddBackward0>)根据前4行的输出结果我们可以看出，两个进程的 PID 分别是 21544 和 21545，且它们的父进程都是 21535。同时，根据11~15行的输出结果可以看出，在两个进程上的输出结果都是一样，而这也是刚好符合我们预期的。
注意，这里我们需要明白一件事，由于程序是并行所以其实函数 run() 在每个设备上都会执行一次，并且内部依靠 rank 来区分计算逻辑。
1.4 GPU 上真实运行#
在有了前面的介绍以后，在 GPU 上的运行过程就变得非常简单了，只需要修改对应的环境配置即可，如下所示：
上述代码运行结束以后便可以得到类似如下结果：
完整示例代码可参见 Code/C04_Parallel/parallel_linear_gpu.py 文件。
2. 从零实现 ColumnParallelLinear#
在介绍完如何使用 ColumnParallelLinear 和 RowParallelLinear 以后，我们再来看详细看看两者内部在 FairScale 框架中的实现过程，也即践行我们的“先学会怎么用，在探究为什么”这一理念。完整注释标注代码可参见 Code/C04_Parallel/parallel.py 文件。
为了让大家能够更加清楚的弄懂整个实现过程，我们先来给大家介绍一下 stride 这个概念。
2.1 stride 概念介绍#
学习过卷积神经网络的朋友对于 stride 这个概念一定不会陌生，即每次卷积操作以后移动的距离。当然，不清楚的也没关系，因为两者虽然都叫 stride 但是还是有不少的差别。
根据我们在上一篇文章中的介绍，张量并行的计算过程并不涉及到 stride 这个东西，但是为什么在代码实现的时候有多出来了呢？
一句话总结，FairScale 库在实现 ColumnParallelLinear 和 RowParallelLinear 这两种张量并行的时候考虑了更加通用的场景，可以让用户更加灵活地去划分权重矩阵到不同的 GPU 上进行计算，所以引入了 stride 这样一个概念。
下面，我们以 ColumnParallelLinear 为例，来看看 stride 到底是什么样的含义。
如图2所示，最上方为 ColumnParallelLinear 计算时初始化的权重矩阵，形状为 [output_size, in_feature] ；左右两块分别用来示例权重矩阵在两块显卡（world_size = 2）上，stride 分别为 1 和 2 时，各个设备上划分得到的权重块。
根据图2中的计算过程，首先会根据 output_size 和 world_size 计算得到 output_size_per_partion，即得到每个 GPU 设备上权重块的大小，例如此处计算得到的结果便是 4 ；然后再根据 stride 计算得到 per_partition_per_stride_size，即得到在划分每个设备上的权重块的时候应该以多大的间隔进行，例如这里左右两边分别是 $4//1=4$ 和 $4//2=2$。
可以看出，这里尤其需要注意的是，两步计算过程都要保证能够整除。
在计算得到 per_partition_per_stride_size 后，便可以把权重矩阵进行划分得到 weight_list，最终根据 weight_list[rank::world_size] 索引得到每块 GPU 上对应的权重块。
根据最后得到的结果可以看出，图2中左右两侧最终得到的每块设备上权重块的大小都是 4 ，但是区别在于索引得到了原权重矩阵中不同的行。当然，这中设计只是为了满足其它更为复杂的场景，在本文所要实现的 ColumnParallelLinear 和 RowParallelLinear 中，stride 均为 1，即直接按照每块设备权重块的大小在原始权重均值上连续划分即可。
说到这里你是否会产生一个疑问？
我们在上一篇文章当中似乎隐约提过，对于每块设备上的权重块我们可以直接在各个设备上进行初始化，而非先初始化一个整块然后划分到各个设备上，因为前者看起来效率更高。
为什么呢？这个问题我们放到稍后进行回答。
2.2 ColumnParallelLinear 代码实现#
在介绍完 stride 的概念以后，我们再来看下如何实现整个 ColumnParallelLinear 部分。
为了让大家能够更加清晰地理解整个实现过程，我们先通过一个图来看看整个代码的层次结构，如图3所示。
如图3所示，我们将会定义 ColumnParallelLinear 这个类，并且在初始化方法 __init__() 中来完成各个设备上矩阵块的划分，然后通过 forward() 来完成前向传播计算过程。最终，self.weight 中记录的便是每个设备上对应的矩阵块，get_master_weight() 返回的便是整个权重矩阵。
① 初始化方法
下面，我们带着大家来仔细看看初始化方法中的实现过程，如下所示：
1 class ColumnParallelLinear(torch.nn.Module):
2 def __init__(self, in_features: int, out_features: int,
3 bias: bool = True, gather_output: bool = True,
4 init_method: Callable[[torch.Tensor], torch.Tensor] = init.xavier_normal_,
5 stride: int = 1, keep_master_weight_for_test: bool = False) -> None:
6 super(ColumnParallelLinear, self).__init__()
7 self.in_features = in_features
8 self.out_features = out_features
9 self.gather_output = gather_output
10 world_size = get_model_parallel_world_size()
11 self.output_size_per_partition = divide_and_check_no_remainder(out_features, world_size)
12 self.weight = nn.Parameter(torch.Tensor(self.output_size_per_partition, self.in_features))
13 if bias:
14 self.bias = nn.Parameter(torch.Tensor(self.output_size_per_partition))
15 with torch.no_grad():
16 self.bias.zero_()
17 else:
18 self.register_parameter("bias", None)
19 self.master_weight = _initialize_affine_weight( self.weight,
20 self.out_features, self.in_features, self.output_size_per_partition,
21 0, init_method, stride=stride,
22 return_master_weight=keep_master_weight_for_test)在上述代码中，第2~5行是实例化类时所接收的参数，其中 init_method 用于后续初始化完整的权重矩阵，gather_output 用于指定是否将前向传播后每个设备上的结果合并到一起；keep_master_weight_for_test 用于指定第19行代码中是否返回结果，为 False 时 self.master_weight 将为 None，其主要用于模型开发时结果分析。
第10行是得到当前主机上可并行的设备数量 world_size。第11行是计算得到每个 GPU 设备上权重块的大小，同时会判断能否整除，不能则会报错。第12行是初始化每块设备上权重块的大小所以形状是 [output_size_per_partition, in_features]，相当于是一个占位符的作用，因为后面在第19行 _initialize_affine_weight() 函数中会被重新赋值。第13~18行则是初始化偏置。第19~22行则是先初始化一个完整的权重矩阵，然后再划分到各个设备上，下面我们继续来看这部分代码。
1 def _initialize_affine_weight( weight: torch.Tensor, out_features: int,
2 in_features: int, per_partition_size: int, partition_dim: int,
3 init_method: Callable[[torch.Tensor], torch.Tensor],
4 stride: int = 1,return_master_weight: bool = False, ):
5
6 world_size = get_model_parallel_world_size()
7 if world_size == 1:
8 init_method(weight)
9 if return_master_weight:
10 return weight
11 return None
12 master_weight = torch.empty(out_features, in_features, dtype=weight.dtype, requires_grad=False)
13 init_method(master_weight)
14 per_partition_per_stride_size = divide_and_check_no_remainder(per_partition_size, stride)
15 weight_list = torch.split(master_weight, per_partition_per_stride_size, dim=partition_dim)
16 rank = get_model_parallel_rank()
17 my_weight_list = weight_list[rank::world_size]
18 with torch.no_grad():
19 torch.cat(my_weight_list, dim=partition_dim, out=weight)
20 if return_master_weight:
21 return master_weight
22 return None在上述代码，第1行 weight 便是 ColumnParallelLinear 初始化方法中传入的 self.weight。第7~10行是判断，如果当前仅在存在一个设备，则直接重新初始化 weight 即可。这里需要注意的是，因为 init_method 默认是 in-place 操作 init.xavier_normal_，所以 init_method(weight) 执行完毕以后就会修改原始 self.weight 中的值。
第12~13行是初始化一个完整的权重矩阵 master_weight，并对其进行重新初始化。再次注意，这里我们需要明白一件事，由于程序是并行所以其实整个程序在每个设备上都会执行一次，并且依靠 rank 来区分计算逻辑，也就是说在每个设备上都会初始化一个 master_weight，这一点我们通过后面的图示进行说明。
第14行是计算得到在划分每个设备上的权重块的时候应该以多大的间隔进行。第15行则是根据指定的间隔对 master_weight 进行划分，对应图2中所示的过程。第16行是得到当前设备设备（进程）对应的 rank，并根据17行来划分得到自己对应部分的权重块。第18~19行是将索引得到的权重拼接到一起。
② 前向传播
进一步，实现完初始化操作后开始完成前向传播计算过程，实现代码如下：
1 def forward(self, input_: torch.Tensor) -> torch.Tensor:
2 input_parallel = copy_to_model_parallel_region(input_)
3 output_parallel = F.linear(input_parallel, self.weight, self.bias)
4 if self.gather_output:
5 output = gather_from_model_parallel_region(output_parallel)
6 else:
7 output = output_parallel
8 return output在上述代码中，第2行是将原始完整输入拷贝到各个设备上，使各个设备使用同一份输入。不过这里可能会问，明明每个进程接收的 input_ 是一样的，为什么还需要 copy_to_model_parallel_region(input_)？那是因为即使你在主进程构造了相同的 input_ 张量，如果你不调用 copy_to_model_parallel_region(input_)，那么在多进程并行中，每个进程上的 input_ 实际上仍然是各自独立的一份内存副本，并非同步的共享张量。第3行则是完整的输入，同每个设备上对应的权重块进行线性变换的过程。第4~8行则是判断是否将计算完成后分布在各个设备上的输出拼接到一起并返回。
③得到完整矩阵
最后，我们还可以通过 get_master_weight() 方法来返回得到参与线性计算的 master_weight 矩阵，代码如下：
1 def get_master_weight(self) -> torch.Tensor:
2 return gather_from_model_parallel_region(
3 self.weight.data.transpose(0, 1)).transpose_(0, 1)这里尤其需要 注意一点，get_master_weight() 函数返回的 master_weight 才是真正各个设备上权重块拼接后的结果，并不是每个类对象中的全局 self.master_weight 。
2.3 图解 ColumnParallelLinear 并行过程#
虽然整个 ColumnParallelLinear 的实现过程已经介绍完了，但是我们相信大家一定会弄错一个地方，下面我们再通过一个图示来给大家回顾一下整个代码在运行时变量的变化过程，完整示例代码见 Code/C04_Parallel/ColumnParallelLinear_cpu.py 文件。
首先，我们给出程序代码，如下所示：
1 def run(rank, world_size=2):
2 ......
3 dist.init_process_group(backend="gloo", rank=rank, world_size=world_size)
4 initialize_model_parallel(world_size)
5 batch_size,input_dim,output_dim = 2,3,8
6 x = torch.randn(batch_size, input_dim)
7 layer = ColumnParallelLinear( in_features=input_dim,
8 out_features=output_dim, gather_output=False,
9 keep_master_weight_for_test=True)
10 output = layer(x)
11 print(f"Rank:{rank} Output:\n{output}")
12 print(f"Rank:{rank} master_weight:\n{layer.master_weight}")
13 print(f"Rank:{rank} get_master_weight:\n{layer.get_master_weight()}")
14 print(f"Rank:{rank} weight:\n{layer.weight}")
15 dist.destroy_process_group()在上述代码中，我们指定了 world_size=2，并且输出我们需要验证和分析部分的变量。同时，我们还可以通过如下图所示的过程来表示 run() 在并行时的情况。
如图4所示为双卡并行时 ColumnParallelLinear 中各个变量的变化情况。从上往下看，并行时函数 run() 会分别在两个进程上执行实例化一个 layer 类对象，然后各自初始得到一个 weight。进一步，各自再初始化得到一个 master_weight，并且根据 rank 计算得到对应的划分块，接着从各自的 master_weight 取出这部分权重作为自己设备上对应的权重块 weight。所以，最后通过 layer.master_weight 得到的便是各个进程上一开始初始化的 master_weight，每个进程都不一样；通过 layer.get_master_weight() 得到的便是所有设备上权重块拼接后的结果，每个进程返回的结果都是一样的；通过 layer.weight 得到的便是每个设备上对应的权重块。
上述代码运行结束后的结果类似如下所示：
> initializing model parallel with size 2
> initializing ddp with size 1
> initializing pipeline with size 1
[Rank:0] Output:
tensor([[ 1.1377, -0.2127, -0.1992, 1.9902],
[ 0.0053, -0.4526, -0.6691, -0.2070]], grad_fn=<AddmmBackward0>)
[Rank:1] Output:
tensor([[-0.5462, -0.4056, 0.5707, -0.0984],
[ 0.0651, 0.3634, -0.2874, -0.1482]], grad_fn=<AddmmBackward0>)
Rank:0 master_weight:
tensor([[-0.6482, 0.0033, -0.4027],
[ 0.1520, 0.6575, -0.2176],
[-0.4115, 0.6603, -0.0949],
[-0.6823, 0.5604, -1.0497],
[ 0.5418, 0.0205, -0.2759],
[-0.0383, 0.2029, -0.1678],
[ 0.2387, -0.0888, 0.1351],
[ 0.2006, -0.0577, -0.1250]])
Rank:1 master_weight:
tensor([[ 0.0858, 0.1115, -0.4627],
[ 0.3607, -0.2240, 0.3769],
[ 0.0109, 0.2240, 0.7165],
[ 0.0143, 0.5922, 0.2749],
[-0.8927, -0.4710, -0.5375],
[-0.1299, -0.1249, 0.3380],
[-0.4457, 0.4655, -0.3996],
[-0.0582, -0.2470, -0.3050]])
Rank:0 get_master_weight:
tensor([[-0.6482, 0.0033, -0.4027],
[ 0.1520, 0.6575, -0.2176],
[-0.4115, 0.6603, -0.0949],
[-0.6823, 0.5604, -1.0497],
[-0.8927, -0.4710, -0.5375],
[-0.1299, -0.1249, 0.3380],
[-0.4457, 0.4655, -0.3996],
[-0.0582, -0.2470, -0.3050]])
Rank:1 get_master_weight:
tensor([[-0.6482, 0.0033, -0.4027],
[ 0.1520, 0.6575, -0.2176],
[-0.4115, 0.6603, -0.0949],
[-0.6823, 0.5604, -1.0497],
[-0.8927, -0.4710, -0.5375],
[-0.1299, -0.1249, 0.3380],
[-0.4457, 0.4655, -0.3996],
[-0.0582, -0.2470, -0.3050]])
Rank:0 weight:
Parameter containing:
tensor([[-0.6482, 0.0033, -0.4027],
[ 0.1520, 0.6575, -0.2176],
[-0.4115, 0.6603, -0.0949],
[-0.6823, 0.5604, -1.0497]], requires_grad=True)
Rank:1 weight:
Parameter containing:
tensor([[-0.8927, -0.4710, -0.5375],
[-0.1299, -0.1249, 0.3380],
[-0.4457, 0.4655, -0.3996],
[-0.0582, -0.2470, -0.3050]], requires_grad=True)介绍到这里，相信大家对于 ColumnParallelLinear 内部的实现流程已经比较清楚了，下面我们需要来回答前面提出的问题：
为什么非先初始化一个整块 master_weight 然后划分到各个设备上，而不是仅在各个设备上初始自己对应的权重块，也就是上面对应的 layer.weight？
2.4 全局 master_weight 作用#
对于为什么需要全局 master_weight 有下面几点原因：
先造一个完整的大矩阵再切，可以保证“切出来的每一片”在逻辑上、位置上、随机数上都严格对应全局矩阵，并且方便做 checkpoint、debug 和各种跳跃式切分。并且，由于代码实现是并行各自切自己的，所以还没有通信上的开销。
虽然每个进程都构造了 master_weight，但它仅用于初始化自己的 self.weight 变量，master_weight 本身不会参与训练和计算图构建，也不会保留下来（除非 keep_master_weight_for_test=True），而整个过程是为了保证每个 rank 的 self.weight 是从相同初始化逻辑中切出来的 —— 以确保整个模型行为一致性。
如果你只初始化局部权重块，虽然节省了那点微不足道的内存／计算，但是会丢失一致性、可重组性以及对复杂切分（stride>1）的天然支持。当 stride>1，切分并不是简单的连续区间，而是跳跃式地取若干小块，如果你只给每卡初始化它自己那几段，那么在不先生成全局 master_weight 的前提下，要正确地“插入”这些跳跃段就非常复杂——得搞清楚它们在全局矩阵的原始排布。
3. 从零实现 RowParallelLinear#
在介绍完 ColumnParallelLinear 以后，我们再来看 RowParallelLinear 的实现过程就非常容易了，我们下面简单地来看一下。
3.1 RowParallelLinear 代码实现#
对于 RowParallelLinear 的实现过程来说，其整个代码层次结构同图3类似，所以我们下面直接来看代码的实现过程。
① 初始化方法
下面，首先来看看初始化方法中的实现过程，如下所示：
1 class RowParallelLinear(torch.nn.Module):
2 def __init__( self, in_features: int, out_features: int,
3 bias: bool = True, input_is_parallel: bool = False,
4 init_method: Callable[[torch.Tensor], torch.Tensor] = init.xavier_normal_,
5 stride: int = 1, keep_master_weight_for_test: bool = False,):
6 super(RowParallelLinear, self).__init__()
7 self.in_features = in_features
8 self.out_features = out_features
9 self.input_is_parallel = input_is_parallel
10 world_size = get_model_parallel_world_size()
11 self.input_size_per_partition = divide_and_check_no_remainder(in_features, world_size)
12 self.weight = nn.Parameter(torch.Tensor(self.out_features, self.input_size_per_partition))
13 if bias:
14 self.bias = nn.Parameter(torch.Tensor(self.out_features))
15 with torch.no_grad():
16 self.bias.zero_()
17 else:
18 self.register_parameter("bias", None)
19 self.master_weight = _initialize_affine_weight( self.weight,
20 self.out_features, self.in_features, self.input_size_per_partition,
21 1, init_method, stride=stride,return_master_weight=keep_master_weight_for_test)上述代码基本上和 ColumnParallelLinear 中的一致，就不再赘述。唯一的区别在于第11~12行中，RowParallelLinear 中是对输入维度进行划分，所以 input_size 是需要能够被 world_size 整除的，最后 self.weight 的形状为 [out_features, input_size_per_partition]。
② 前向传播
进一步，实现完初始化操作后开始完成前向传播计算过程，实现代码如下：
1 def forward(self, input_: torch.Tensor) -> torch.Tensor:
2 if self.input_is_parallel:
3 input_parallel = input_
4 else:
5 input_parallel = scatter_to_model_parallel_region(input_)
6 output_parallel = F.linear(input_parallel, self.weight)
7 output_ = reduce_from_model_parallel_region(output_parallel)
8 if self.bias is not None:
9 output = output_ + self.bias
10 else:
11 output = output_
12 return output在上述代码中，第2~5行用来判断输入 input_ 此时是否已经分布在各个设备上，例如前一步刚刚是计算结束的 ColumnParallelLinear 那 input_is_parallel 就应该设置为 True，否则就要将输入进行划分并广播到各个设备上。第5~6行是各个设备分别各自完成对应的线性组合计算然后相加得到最后的输出 output_ 结果。
3.2 图解 RowParallelLinear 并行过程#
同样，这里我们也通过一张图来简单地看一下在 RowParallelLinear 中整个代码在运行时权重矩阵的变化过程，整个代码在运行时变量的变化过程，完整示例代码见 Code/C04_Parallel/RowParallelLinear_cpu.py 文件。
首先，我们给出程序代码，如下所示：
1 def run(rank, world_size = 2):
2 ......
3 dist.init_process_group(backend="gloo", rank=rank, world_size=world_size)
4 initialize_model_parallel(world_size)
5 batch_size,input_dim,output_dim = 2,4,8
6 x = torch.randn(batch_size, input_dim)
7 layer = RowParallelLinear( in_features=input_dim, out_features=output_dim,
8 input_is_parallel=False, keep_master_weight_for_test=True)
9 output = layer(x)
10 print(f"Rank:{rank} Output:\n{output}")
11 print(f"Rank:{rank} master_weight:\n{layer.master_weight}")
12 print(f"Rank:{rank} get_master_weight:\n{layer.get_master_weight()}")
13 print(f"Rank:{rank} weight:\n{layer.weight}")
14 dist.destroy_process_group()在上述代码中，第8行中我们指定了 input_is_parallel = False，也就表示当前输入 x 并没有分配到各个设备上，需要通过我们上面介绍的 input_parallel = scatter_to_model_parallel_region(input_) 进行广播。 同时，我们还可以通过如下图所示的过程来表示 run() 在并行时的情况。
如图5所示便是双卡并行时 RowParallelLinear 中权重参数的变化过程，整体上同图4中的过程一致，这里就不再赘述。
最后，上述 run() 函数运行结束以后的结果为：
> initializing model parallel with size 2
> initializing ddp with size 1
> initializing pipeline with size 1
Rank:1 Output:
tensor([[-1.5984, 0.3450, -0.8908, -1.0466, 0.1380, -0.0195, -0.6476, 1.5980],
[ 0.7189, 0.0848, 0.2930, -0.2203, -1.7449, -0.4163, -0.6861, -0.8015]],
grad_fn=<AddBackward0>)Rank:0 Output:
tensor([[-1.5984, 0.3450, -0.8908, -1.0466, 0.1380, -0.0195, -0.6476, 1.5980],
[ 0.7189, 0.0848, 0.2930, -0.2203, -1.7449, -0.4163, -0.6861, -0.8015]],
grad_fn=<AddBackward0>)
Rank:1 master_weight:
tensor([[ 0.4043, -0.3316, 0.4244, -0.1243],
[-0.2400, -0.0837, 0.2050, 0.0367],
[-0.2023, -0.1222, 0.0777, 0.0610],
[ 0.1911, -0.3903, -0.2756, -0.2109],
[-0.0623, -0.5191, -0.9562, -0.9136],
[-0.6861, -0.2524, -0.1756, -0.0309],
[ 0.2642, -0.2509, -0.1439, -0.4757],
[ 0.3479, -0.1346, -0.2158, -0.0095]])
Rank:0 master_weight:
tensor([[-0.3933, -0.0278, 0.5974, -0.0525],
[ 0.0853, -0.3322, 0.2918, -0.1643],
[-0.2847, 0.1345, -0.3409, 0.2798],
[-0.2567, 0.3180, 0.0633, 0.0461],
[ 0.2794, 0.1202, -0.3799, -0.4058],
[-0.0569, -0.1793, -0.7418, -0.7882],
[-0.1109, -0.4893, -0.0777, -0.1438],
[ 0.4086, -0.5037, -0.2567, -0.6275]])
Rank:0 get_master_weight:
tensor([[-0.3933, -0.0278, 0.4244, -0.1243],
[ 0.0853, -0.3322, 0.2050, 0.0367],
[-0.2847, 0.1345, 0.0777, 0.0610],
[-0.2567, 0.3180, -0.2756, -0.2109],
[ 0.2794, 0.1202, -0.9562, -0.9136],
[-0.0569, -0.1793, -0.1756, -0.0309],
[-0.1109, -0.4893, -0.1439, -0.4757],
[ 0.4086, -0.5037, -0.2158, -0.0095]])
Rank:1 get_master_weight:
tensor([[-0.3933, -0.0278, 0.4244, -0.1243],
[ 0.0853, -0.3322, 0.2050, 0.0367],
[-0.2847, 0.1345, 0.0777, 0.0610],
[-0.2567, 0.3180, -0.2756, -0.2109],
[ 0.2794, 0.1202, -0.9562, -0.9136],
[-0.0569, -0.1793, -0.1756, -0.0309],
[-0.1109, -0.4893, -0.1439, -0.4757],
[ 0.4086, -0.5037, -0.2158, -0.0095]])
Rank:1 weight:
Parameter containing:
tensor([[ 0.4244, -0.1243],
[ 0.2050, 0.0367],
[ 0.0777, 0.0610],
[-0.2756, -0.2109],
[-0.9562, -0.9136],
[-0.1756, -0.0309],
[-0.1439, -0.4757],
[-0.2158, -0.0095]], requires_grad=True)
Rank:0 weight:
Parameter containing:
tensor([[-0.3933, -0.0278],
[ 0.0853, -0.3322],
[-0.2847, 0.1345],
[-0.2567, 0.3180],
[ 0.2794, 0.1202],
[-0.0569, -0.1793],
[-0.1109, -0.4893],
[ 0.4086, -0.5037]], requires_grad=True)到此，对于 ColumnParallelLinear 和 RowParallelLinear 这两个张量并行的实现过程就介绍完了，感谢您的阅读！
引用#
[1] https://github.com/facebookresearch/fairscale