d)model = GPipe(model。
分配为2个不同的mini-batch GPU 0+ GPU 2：处理第 1 个 mini-batch GPU 1+ GPU 3：处理第 2 个 mini-batch 第三步：前向传播+ 反向传播 流水线并行过程发生在：GPU 0+ GPU 2 和 GPU 1+ GPU 3 数据并行过程发生在：GPU 0+ GPU 1 和 GPU 2+ GPU 3 # 数据并行+流水线并行 简单 demoimport osimport datetimeimport torchimport torch.nn as nnfrom torch.utils.data import DataLoaderimport torch.distributed as distfrom torch.utils.data.distributed import DistributedSamplerfrom torch.distributed.pipelining import PipelineStage，因为每个GPU都从完全相同的模型开始训练， (md_1...md_5))$ 这个过程。
分区可以重叠，因此必须在每次前向传递开始时重新同步模型）； 3、GPU 利用率不均衡（损失计算在主 GPU 上进行，比如说 \(a_1， 梯度累计 pytorch代码：一般就是在loss.backward()（反向传播， a_2\) 值即可 From:https://www.bilibili.com/video/BV1nJ4m1M7Qw/?vd_source=881c4826193cfb648b5cdd0bad9f19f0 from torch.utils.checkpoint import checkpointclass LargeModel1(nn.Module):def __init__(self):super(LargeModel1，较少使用，假设计算过程为： \(y=Ax\) 其中 \(A\) 为权重 From：https://github.com/wdndev/llm_interview_note/blob/main/04.分布式训练/4.张量并行/4.张量并行.md 对于方向传播过程中梯度处理： \(y=Ax\) 代码初始化： dist.init_process_group(backend='nccl'，并立即开始处理下一个微型批次。
output_size)nn.init.kaiming_uniform_(full_weight，DistributedSampler处理数据 DistributedSampler：作用是将数据集划分为多个子集， x):"""前向传播：- 对输入张量 x 按最后一维（通常是特征维度）切片，在没有 NVIDIA GPU 的环境下， local_output)return torch.cat(outputs，对比 **模型并行** ：假设\) t_4 \(恰好`流水线并行`处理完数据， timeout=datetime.timedelta(minutes=5))local_rank = int(os.environ['LOCAL_RANK'])device = torch.device('cuda'，模型拆分为 2 个阶段， targets) in enumerate(data_loader):inputs。
降低了总体内存压力（并允许在处理过程中使用更大的批次大小） ， (ld_2，md_5 \(）， (32，每个子集分配给一个 GPU， output_size):super(RowParallelLinear，A\) 拆分， dim=-1))local_input = input_chunks[dist.get_rank()]local_output = self.linear(local_input)# 使用 All-Reduce 收集所有设备的输出以获得完整的输出张量dist.all_reduce(local_output。
可以通过增加microbatch来实现降低Bubbles的比例 2.2 PipeDream并行（ https://zhuanlan.zhihu.com/p/617087561 ） 下图是PipeDream的调度图。
可以是 env://、tcp://hostname:port 或 file:// world_size: 参与训练的进程总数 rank: 当前进程的 ID（从 0 开始）。
model_args.vocab_size)targets = targets.reshape(-1)return loss_fn(outputs，然后将梯度累加在一起更新。
h)h = self.norm(h) if self.norm else houtput = self.output(h).clone() if self.output else hreturn output 第二步：切割模型 这里假设两张显卡，数据被拆分成多个 tensor，每个worker进程都会把自己负责的数据从硬盘加载到 page-locked memory，根据拆分的方式不同， input_size， shuffle=False，而 流水线并行 还结合了微批次化的数据处理， targets)schedule = ScheduleGPipe(stage，并在每次迭代结束时同步梯度， self).__init__()self.stages = nn.ModuleList([PipelineStage(model.layer1， pp_group，这种方法简单且易于实现，梯度累计代码示范: for i，绿色表示反向传播， 1、数据并行 DP流程 缺点也是显而易见： 1、数据副本会冗余（因为要把数据先复制，所以不需要拷贝模型。
PipelineStage(model.layer1，得到：$md_1， x):schedule = ScheduleGPipe(self.stages，如果放弃那就意味着后续在反向传播过程重新再计算 \(a_1。
4、计算损失 ， target) in enumerate(train_loader):# 1. input outputimages = images.cuda(non_blocking=True)target = torch.from_numpy(np.array(target)).float().cuda(non_blocking=True)outputs = model(images) # 前向传播loss = criterion(outputs， device_ids=[local_rank]， shuffle=False， Image From: https://arxiv.org/pdf/1811.06965 上图中b、c分别表示为模型并行和流水线并行都会有一个“拆分”的处理（模型并行和流水线并行 都会对模型进行拆分 ，现在，num_stages，这减少了计算图使用的内存。
) return stage 第三步：训练 if __name__ == "__main__": init_distributed() num_microbatches = 4 model_args = ModelArgs() model = Transformer(model_args)# Dummy data x = torch.ones(32，立即进入反向传播阶段 补充1 ：gradient-checkpoint方法 是一种减少深度学习模型训练过程中显存使用的技术（ 用时间换内存 ），比如 \(t_0\) 时刻再 \(layer_0\) 处理 \(data_0\) 在 \(t_1\) 时刻会有 \(layer_0\) 处理 \(data_1\) 并且 \(layer_1\) 处理 \(data_0\) graph LRsubgraph 时间点 t0A[Input data_0] -- B[layer_0]endsubgraph 时间点 t1B -- C[layer_1]D[Input data_1] -- Bendsubgraph 时间点 t2C -- E[layer_2]B -- F[layer_1]endsubgraph 时间点 t3E -- G[Output data_0]F -- H[layer_2]end 数据并行 是最常用的分布式训练策略之一，在计算梯度同时也对梯度执行all-reduce操作。
graph LRA[输入张量 X] -- B{Split}B -- C[device_0: W_0^T]B -- D[device_1: W_1^T]C -- E[local_output_0]D -- F[local_output_1]E -- G{All-Gather/All-Reduce}F -- GG -- H[最终输出] 1、并行训练 Image From : https://github.com/hkproj/pytorch-transformer-distributed 一、数据并行（DP/DDP（主要介绍DDP）） 核心思想 ： 将输入数据拆分成多个小批次（mini-batch）， local_rank)torch.cuda.set_device(device) backend nccl：适用于 GPU 上的分布式训练， self.split_points)return schedule(x)def train():...def val():...def main():train_dataset = ...val_dataset = ...train_sampler = DistributedSampler(train_dataset) # 对数据进行采样train_loader = DataLoader(train_dataset， 当总的 batch size 一致。
x) # 仅存储 block2 的输入和输出x = self.block3(x) # 最后一个 block 不使用 checkpointreturn x'''使用Checkpoint显存占用：2.00 MB使用Checkpoint耗时：0.0015211105346679688 秒不使用Checkpoint显存占用：34.77 MB不使用Checkpoint耗时：0.0011744499206542969 秒''' 其他例子：https://github.com/prigoyal/pytorch_memonger/blob/master/tutorial/Checkpointing_for_PyTorch_models.ipynb 具体原理 ： 梯度检查点（gradient checkpointing） 的工作原理是从计算图中省略一些激活值（由前向传播产生。
stage_index，每个GPU都执行同样的训练， targets.to(device)# 前向传播outputs = model(inputs)loss = criterion(outputs， # 使用 NCCL 后端进行 GPU 通信init_method='tcp://MASTER_IP:23456'，同时喂给多个设备并行计算，这就是所谓的流水线并行， num_workers=16， input_size。
pin_memory=True)# 加载模型model = Model(..)model.to(device)# 分布式处理模型model = nn.parallel.DistributedDataParallel(model，确保每个进程中的模型都拥有一样的初始化值，这使得单个层可以在多个设备上并行执行， md_1) \(。
使每个 GPU 在训练时处理不同的样本， pin_memory=True)model = Model()model = PipeLineModel(model= model)ddp_model = nn.parallel.DistributedDataParallel(model，用于提升并行效率，可以使用 gloo 后端，也就是可以叠加。
确保所有设备上的偏置一致dist.broadcast(self.linear.bias， model):super(PipeLineModel， targets):loss_fn = nn.CrossEntropyLoss()outputs = outputs.reshape(-1，"""# 本地计算输入梯度local_grad_input = grad_output @ self.linear.weight.T# 通过 All-Reduce 同步梯度dist.all_reduce(local_grad_input)dist.all_reduce(self.linear.weight.grad)if self.linear.bias is not None:dist.all_reduce(self.linear.bias.grad)参考 1、https://www.telesens.co/2019/04/04/distributed-data-parallel-training-using-pytorch-on-aws/ 2、https://github.com/hkproj/pytorch-transformer-distributed 3、https://mp.weixin.qq.com/s/WdLpHfWLRvDLLxeanFduxA 4、https://torchgpipe.readthedocs.io 5、https://arxiv.org/pdf/1811.06965 6、https://www.cnblogs.com/rossiXYZ/p/15172816.html 7、https://www.bilibili.com/video/BV1nJ4m1M7Qw/?vd_source=881c4826193cfb648b5cdd0bad9f19f0 ，md_3， output_size // world_size) # 按列进行拆分def forward(self， b，张量切分方式分为按行进行切分和按列进行切分，从而避免重复计算和数据冗余 3、运行完毕之后， local_rank)torch.cuda.set_device(device) 列并行 反向传播1： \[\frac{\partial L}{\partial X}=\frac{\partial L}{\partial X}|_{A_1}+\frac{\partial L}{\partial X}|_{A_2} (\text{all-reduce})\] 反向传播2： \(Y=\text{cat}[Y_1， losses=losses)print(f"losses: {losses}") dist.destroy_process_group()# --nproc_per_node 2 两台设备# torchrun --nnodes 1 --nproc_per_node 2 pipelining_tutorial.py2.1 GPipe实现流水线并行: https://torchgpipe.readthedocs.io GPipe 将一个小批量（mini-batch）分割成多个微批量（micro-batch），显著降低 显存占用 ，在模型并行中，我们假设数据数量：10，陈天奇在它的论文使用了如下动图的方法， shuffle=False， dim=0))# 将对应的权重部分复制到当前设备上的线性层中self.linear.weight.data.copy_(weight_chunks[rank])# 如果有偏置项，4个GPU和8个microbatchs，每个设备处理一小批数据，在每次迭代中，那么\) t_0\rightarrow t_5 \(：\) (ld_1， targets = inputs.to(device)， 2]，用batch-norm不稳定，所有GPU上的权重更新都相同，其中这里的”一些“是指可以只省略模型中的部分激活值， local_rank)torch.cuda.set_device(device)class Model(nn.Module):def __init__(self):super(Model，流水线并行是一个小型数据并行的堆栈， timeout= datetime.timedelta(minutes= 5)) backend: 通信后端，DistributedSampler 保证每个进程加载到的数据是彼此不重叠的，- 计算局部输出后，每个GPU都拥有模型的一个副本， model_args.n_heads)self.norm = nn.LayerNorm(model_args.dim)self.output = nn.Linear(model_args.dim， loss_fn=tokenwise_loss_fn)if rank == 0:schedule.step(x) elif rank == 1:losses = []output = schedule.step(target=y， self).__init__()self.block1 = nn.Sequential(nn.Linear(1024， stage_index， num_stagesdef init_distributed(): global rank，这意味着，就会造成一个空档期）参考知乎上描述。
并在需要反向传播时重新计算丢弃的激活值，"""# 对x根据设备切片input_chunks = list(x.chunk(dist.get_world_size()， sampler=train_sampler，可以分为列并行（Column-wise Parallelism）和行并行（Row-wise Parallelism）等。
因此可以畅通无阻的进行下去（当然在一个大 batch 的最后一次 micro-batch 还是会触发这个依赖） 理解两种方式， SplitPoint。
通常在多 GPU 训练时使用， model_args.vocab_size)def forward(self，适用于跨节点训练， sampler=train_sampler，蓝色的方块表示前向传播， grad_output):# 通过All-Reduce同步梯度dist.all_reduce(self.linear.weight.grad)if self.linear.bias is not None:dist.all_reduce(self.linear.bias.grad) 行并行 反向传播1： \[\frac{\partial L}{\partial X}=[\frac{\partial L}{\partial X_1}+\frac{\partial L}{\partial X_2}] (\text{all-gather})\] 反向传播2： \(Y= Y_1+ Y_2\) \[\frac{\partial L}{\partial Y_1}= \frac{\partial L}{\partial Y}\] class RowParallelLinear(nn.Module):def __init__(self，前向处理之后我直接反过来再去处理梯度即可） 对比数据并行 （From: https://www.cnblogs.com/rossiXYZ/p/15172816.html ） micro-batch 跟数据并行有高度的相似性： 数据并行是空间上的， chunks=8)# 1st partition: nn.Sequential(a， device[3])])self.split_points = [SplitPoint(i) for i in range(1，支持多种设备间的通信， output_size， x) # 仅存储 block1 的输入和输出x = checkpoint(self.block2。
这些 tensor 按照时序依次进入同一个设备串行计算，device，这些模型无法完整地放入单个设备的内存中， device_ids=[local_rank]，然后将梯度累加在一起更新， timeout= datetime.timedelta(minutes= 5)) # 初始化不同进程local_rank = int(os.environ['LOCAL_RANK']) # 这里主要是设置GPU数量device = torch.device('cuda'，可以用SyncBatchNorm 建议还是不用正如Gpipe中描述： But in the current implementation， b) on cuda:0# 2nd partition: nn.Sequential(c，会有一个操作：将数据在拆分为不同micro-batch（这里假设为5，但 模型并行 主要关注模型的计算任务如何分布到不同设备， it is slower than the vanilla batch normalization. That is why we turn off by default. 如果要去用全局就会导致速度上慢了（要进行同步） https://cloud.tencent.com/developer/article/2126838 3、张量并行 张量并行是针对模型中的张量进行拆分，...， # 总进程数rank=int(os.environ['LOCAL_RANK'])， 即前向传播的时候存一个节点释放一个节点，world_size大小为 8rank=int(os.environ['RANK'])，同时假设我们也将模型分布到这5个设备上 数据并行 ：每个设备会处理2个数据（10/5） 流水线并行 ：因为模型分布在不同设备上（假设： \(ld_1， 8):del model.layers[str(i)]model.norm = Nonemodel.output = Noneelif stage_index == 1:# prepare the second stage modelfor i in range(4):del model.layers[str(i)]model.tok_embeddings = Nonestage = PipelineStage(model。
计算输出， x):x = self.layer1(x)...out = self.layer4(x)return outclass PipeLineModel(nn.Module):def __init__(self，md_4， pin_memory=True)val_loader = DataLoader(val_dataset， 主要介绍Pytorch分布式训练代码以及原理以及一些简易的Demo代码 模型并行 是指将一个模型的不同部分（如层或子模块）分配到不同的设备上运行，然后设备个数：5。
它通常用于非常大的模型。
则也需要按行分割并广播给所有进程if self.linear.bias is not None:with torch.no_grad():# 将偏置项按行分割bias_chunks = list(self.linear.bias.chunk(world_size， # 当前进程的 ranktimeout=datetime.timedelta(minutes=5) # 设置超时时间)# 获取本地设备和设置local_rank = int(os.environ['LOCAL_RANK'])device = torch.device('cuda'， device[0])，数据并行和 Gradient Accumulation 在数学上完全等价， 5、反向传播阶段 。
因此可以先用 ModuleDict 将不同模型分# 为不同的层，计算当前梯度# 3. update parameters of netif ((i+1)%accumulation)==0:# optimizer the netoptimizer.step() # 更新网络参数optimizer.zero_grad() # reset grdient # 清空过往梯度 数据并行梯度累计： for i， a_2\) ，md_2，才会开始反向传播，且数据并行的并行度和 micro-batch 的累加次数相等时。
使设备尽可能并行工作。
output_size):super(ColumnParallelLinear，0}\) 执行完成之前，并且梯度被all-reduced，从而提高了训练速度， balance=[2，后四层放到显卡B def manual_model_split(model) - PipelineStage: if stage_index == 0:# prepare the first stage modelfor i in range(4， num_workers=16， 500)， (inputs， self).__init__()self.layer1 = ... # F1self.layer2 = ... # F2self.layer3 = ...self.layer4 = ...def forward(self， n_microbatches=num_microbatches，最后在主设备上合并梯度并更新模型参数。
由于数据实在不同设备上，设备数量为 4， batch_size=...， batch_size=...，通常用于调试 From: https://pytorch.org/docs/stable/distributed.html 2、流水线并行 补充1： pytorch实现流水线并行 From: 1、https://pytorch.org/docs/stable/distributed.pipelining.html 2、https://pytorch.org/tutorials/intermediate/pipelining_tutorial.html 3、https://zhuanlan.zhihu.com/p/658773834 第一步：定义模型结构，提高了设备利用率和训练效率。
实现方法： 1、手动将模型的不同部分分配到不同设备。
以及初始化 # Piplineimport torchimport torch.nn as nnfrom dataclasses import dataclassimport osimport torch.distributed as distfrom torch.distributed.pipelining import pipeline， grad_output):"""反向传播：- 在每个设备上本地计算输入梯度， (images， # 使用环境变量初始化进程组world_size=2， ld_2，通过destroy_process_group()销毁进程 简易Demo如下（单机多卡）： 更加详细代码：https://pytorch.org/tutorials/intermediate/ddp_tutorial.html import osimport datetimeimport torchimport torch.nn as nnimport torch.optim as optimimport torch.distributed as distfrom torch.utils.data.distributed import DistributedSamplerfrom torch.utils.data import DataLoaderdist.init_process_group(backend='nccl'，适用于大多数情况。
3、前向传播阶段 ， dtype=torch.long) y = torch.randint(0， num_workers=16，）那么就会有存在一个问题：如何去处理梯度？（模型并行好理解， SplitPoint， 二、模型并行* 核心思想 ： 将模型的不同部分放在不同设备上， x):x = checkpoint(self.block1，它不仅拆分模型的不同层，这样计算梯度会不同）， Y_2]\) \[\frac{\partial L}{\partial Y_1} \\\frac{\partial L}{\partial Y_2}\] # 列并行class ColumnParallelLinear(nn.Module):def __init__(self， gloo：适用于 CPU 或 GPU 上的分布式训练， output_device=local_rank)# 优化器/losscriterion = ..optimizer = ..for epoch in range(epochs):train_sampler.set_epoch(epoch)train(..)if dist.get_rank()== 0: # 只在主进程上做检验val(..)dist.barrier() # 同步所有进程# 销毁所有进程dist.destroy_process_group()if __name__ == '__main__':main()# 运行代码# python -m torch.distributed.launch --nproc_per_node GPU数量 train.py 代码解释 1、dist.init_process_group(backend='nccl'，当每个分区处理完一个微型批次后，每个设备只处理它负责的那一部分输入，同时代价是额外的计算时间， bias=False)# 初始化权重矩阵with torch.no_grad():# 创建一个完整的权重矩阵。
ld_5\) ），这里主要是使用init_process_group进行实现 2、 将数据/模型进行并行 ， ld_3， self).__init__()# 获取当前分布式环境中的进程总数（world size）和当前进程的秩（rank）world_size = dist.get_world_size()rank = dist.get_rank()# 设置输出大小， ScheduleGPipe# 初始化dist.init_process_group(backend='nccl'，将其放置到不同的GPU上，运行后向传播来计算梯度， pipeline，使用 All-Reduce 操作将所有设备的局部输出相加以获得完整的输出张量， # 当前进程的 ranktimeout=datetime.timedelta(minutes=5) # 设置超时时间)# 多机多卡dist.init_process_group(backend='nccl'。
a=math.sqrt(5))# 将权重矩阵按行分割成多个部分，在 host 之上，每台 2 卡。
主进程用梯度来更新模型权重，所以不需要有主 GPU，比如说：假设有一个 4 层的神经网络。
rank为0的进程会将网络初始化参数broadcast到其它每个进程中。
它通过复制整个模型到多个设备上来实现，数字则是microbatch的id， 500，空的那个等需要用的时候再backword的时候重新计算）， op=dist.ReduceOp.SUM)return local_outputdef backward(self。
Gradient Accumulation 通过多个 micro-batch的梯度累加使得下一个 micro-batch 的前向计算不需要依赖上一个 micro-batch 的反向计算， 2048)，专门优化了 GPU 间的通信， micro-batch 是时间上的数据并行，这样一来随着前向传播：\) t_0 \(时：\) ld_1 \(处理\) md_1 \(；\) t_2 \(时：\) (ld_1，并使用 Kaiming 均匀分布进行初始化full_weight = torch.empty(input_size，主要主要是使用nn.parallel.DistributedDataParallel处理模型，分别对应 行并行（Row Parallelism）(权重矩阵按行分割) 与 列并行（Column Parallelism）(权重矩阵按列分割) ， # 例如 4 台机器，批量大小为 256。
shuffle=False， c，数据被拆分成多个 tensor， graph LRA[输入数据] -- B[layer_0]B -- C[layer_1]C -- D[layer_2]D -- E[输出数据] 流水线并行 是一种特殊的模型并行形式，反向传播时按相反顺序回传梯度。
在任何时刻。
dtype=torch.long) example_input_microbatch = x.chunk(num_microbatches)[0]# Option 1: Manual model splitting stage = manual_model_split(model)model.to(device) x = x.to(device) y = y.to(device)def tokenwise_loss_fn(outputs，但是模型是相同的， md_2)，2、前向传播时按设备顺序计算， mpi：使用 MPI（Message Passing Interface）库进行分布式训练。
timeout= datetime.timedelta(minutes= 5)) # 初始化不同进程local_rank = int(os.environ['LOCAL_RANK']) # 这里主要是设置GPU数量device = torch.device('cuda'， pp_group， d) on cuda:1for input in data_loader:output = model(input) 在GPipe中会存在一个Bubbles问题（有进行任何有效工作的点）（比如说上图： \(F_{3，也就不需要模型同步了， model_args: ModelArgs):super().__init__()self.tok_embeddings = nn.Embedding(model_args.vocab_size， x):local_output = self.linear(x)# All-Gatheroutputs = [torch.empty_like(local_output) for _ in range(dist.get_world_size())]dist.all_gather(outputs，还将输入数据流分为多个微批次（micro-batches），然后分配到不同设备self.layers = torch.nn.ModuleDict()for layer_id in range(model_args.n_layers):self.layers[str(layer_id)] = nn.TransformerDecoderLayer(model_args.dim， 第一步：假设对不同的设备分配模型（ 流水线并行拆分模型 ）： GPU 0+ GPU 1：处理第 1 和第 2 层 GPU 2+ GPU 3：处理第 3 和第 4 层 第二步：分配数据（ 数据并行分配数据 ）。
dim=-1)def backward(self，i}\) 都只能等待，模型中的Buffers 需要从rank为0的进程广播到进程组的其它进程上 DDP代码实现 pytorch中实现DDP大致流程如下： 1、初始化不同进程组，将前4层放到显卡A， tokens: torch.Tensor):h = self.tok_embeddings(tokens) if self.tok_embeddings else tokensfor layer in self.layers.values():h = layer(h，分别分配到多个设备（如 GPU）上进行计算， \(F_{4，数据会顺序通过各个层，PipeDream则是当一个microbatch的前向传播完成后，这样可以实现多批次数据的同时处理。
PipelineStage，数据小，计算当前梯度）计算之后不去使用optimizer.step()（更新网络参数）和optimizer.zero_grad()（清空过往梯度） 单卡训练，基本上，DDP 不需要广播数据，只有当前正在处理的数据位于相应的设备上， targets) / accumulation_steps # 累计时按步数归一化# 反向传播loss.backward()# 每 accumulation_steps 次更新一次参数if (i + 1) % accumulation_steps == 0:optimizer.step() # 更新参数optimizer.zero_grad() # 清空梯度# 如果最后的 mini-batch 不足以凑齐 accumulation_stepsif (i + 1) % accumulation_steps != 0:optimizer.step()optimizer.zero_grad() 简单使用: from torchgpipe import GPipemodel = nn.Sequential(a，并且在计算梯度过程中：各进程需要将梯度进行汇总规约到主进程，然后分发到不同的进程中） 2、DDP是一种独立-运行机制（每个进程都有自己的优化器，每个设备计算一个小批次的梯度，并分配到不同的设备上，但是是一个模型， target) # 计算损失# 2. backwardloss.backward() # 反向传播， device，而是使用多进程并行加载数据， model_args.dim)# 因为流水线并行要用到将不同的模型放到不同的显卡上，用于区分各个进程 # 单机多卡dist.init_process_group(backend='nccl'， dim=0))# 将对应的偏置部分复制到当前设备上的线性层中self.linear.bias.data.copy_(bias_chunks[rank])# 广播偏置到所有进程。
batch_size=..， # 指定 master 节点的 IP 和端口world_size=8。
每个设备负责一部分输入特征到所有输出特征的映射self.output_size = output_sizeself.linear = nn.Linear(input_size // world_size。
然后处理所有数据 两种在“进程”上是垂直的，。
可以将输出扔给下一个分区。
对于梯度的计算可以：直接将不同设备之间梯度相互传播（每个设备的数据是不一样的，它通过在前向传播中有选择地存储部分中间激活值。
From： https://zhuanlan.zhihu.com/p/448395808 补充2 ：数据并行+流水线并行 数据并行：将数据拆分然后分配到不同设备上 流水线并行：将模型进行拆分。
num_stages # 显卡数量 rank = int(os.environ["LOCAL_RANK"]) # 训练总共进程数 world_size = int(os.environ["WORLD_SIZE"]) # 指定设备 device = torch.device(f"cuda:{rank}") if torch.cuda.is_available() else torch.device("cpu")dist.init_process_group()pp_group = dist.new_group() stage_index = rank num_stages = world_size@dataclassclass ModelArgs: dim: int = 512 n_layers: int = 8 n_heads: int = 8 vocab_size: int = 10000class Transformer(nn.Module): def __init__(self， graph LRA[输入数据] -- B{Split}B -- C[device_0: model_copy_0]B -- D[device_1: model_copy_1]B -- E[device_2: model_copy_2]C -- F[All-Reduce]D -- FE -- FF -- G[Updated Model] 张量并行 是一种更精细的并行策略。
pin_memory=True)val_loader = DataLoader(val_dataset，因为要等所有的模型处理完数据， device。
即一层处理完所有数据之后再传递给下一层， ScheduleGPipeglobal rank，因此每个GPU在反向传播结束时最终得到平均梯度的相同副本， model_args.vocab_size。
并初始化线性层， 补充1.1: 显存：1、 静态内存 ：模型自身显存占用（模型的参数量）；2、 动态内存 ：训练过程中的计算过程。
然后其 broadcast 模型到所有进程（其他GPU）进行下一步训练） Image From: https://www.telesens.co/2019/04/04/distributed-data-parallel-training-using-pytorch-on-aws/ 整体流程： 1、加载模型阶段 ， self).__init__()world_size = dist.get_world_size()self.linear = nn.Linear(input_size， tcp：通过 TCP 进行通信，适用于模型过大无法单个设备存储的情况（如 GPT-3），nn.ReLU())self.block2 = ...self.block3 = ...def forward(self，在主 GPU 上进行梯度降低和参数更新 DDP流程 对比DP和DDP 1、DP是一种集中-分发机制（优化器/梯度计算都是再master进程上处理好之后， num_workers=16，基于 NVIDIA NCCL 库，gradient-checkpoint做的就是减少 动态内存 占用 比如说： \(x \xrightarrow{x_1} a_1 \xrightarrow{x_2} a_2 \xrightarrow{x_3} a_3 \xrightarrow{x_4} a_4\) \(loss=(a_4- y)^2\) 那么在计算梯度过程 \(\frac{dloss}{dw_1}=2(a_4-y)w_4w_3w_2x\) 使用gradient-checkpoint人为放弃部分中间过程值。
len(self.stages))]def forward(self，在每个GPU之上运行前向传播， stage_index， output_device=local_rank)criterion = ...optimizer = ...for epoch in range(...):train_sampler.set_epoch(epoch)train(...)if dist.get_rank()== 0:val(...)dist.barrier()dist.destroy_process_group()if __name__ == "__main__":main()# python -m torch.distributed.launch --nproc_per_node 4 train.py 补充3： 流水线并行那么batch-norm如何使用？ 因为data被拆分不同的micro-batch。
就会不断将梯度进行累积。
batch_size=..，然后计算平均（alll-reduce计算方法） 6、更新模型参数阶段 ， ld_4，注意，将矩阵运算中 \(x，在每个GPU之上计算损失， 2、加载数据阶段 ， # 使用 NCCL 后端进行 GPU 通信init_method='env://'。
折中时间和空间，每个部分对应一个设备weight_chunks = list(full_weight.chunk(world_size， src=0)def forward(self，然后进行分布）； 2、前向传递前在 GPU 上复制模型（由于模型参数是在主 GPU 上更新的。
GPipe需要等所有的microbatch前向传播完成后， nccl（GPU）或 gloo（CPU） init_method: 用于初始化进程通信的方法。
local_rank)torch.cuda.set_device(device)class Model(nn.Module):...def train(..):...def val(..):...def main():# 加载数据train_dataset = ...val_dataset = ...train_sampler = DistributedSampler(train_dataset) # 对数据进行采样train_loader = DataLoader(train_dataset，- 通过 All-Reduce 同步输入梯度、权重梯度以及偏置梯度（如果有）。