尧图网站建设 
ColossalAI 一维张量并行1D Tensor Parallelism原理与 Shardformer 实战指南【免费下载链接】ColossalAIMaking large AI models cheaper, faster and more accessible项目地址: https://gitcode.com/GitHub_Trending/co/ColossalAI本文以 ColossalAI 官方中文文档《1D 张量并行》为骨架系统讲解一维张量并行的数学原理列并行/行并行切分方式、前向与后向的通信模式、理论效率边界并结合仓库源码深入说明其在 ColossalAI 中经Shardformer的实现细节、参数配置与运行方式。读完本文你将理解 1D 张量并行为什么省显存、为什么需要通信、何时需要 all-reduce并能在自己基于 HuggingFace Transformers 的模型上正确启用它。本文对应文档1D 张量并行作者 Zhengda Bian, Yongbin Li其实现与配套示例同时可参考 Shardformer 文档 与 shardformer 示例目录。1 引言为什么需要张量并行大语言模型训练面临的首要瓶颈是单卡显存。以 LLaMA-2 70B、OPT 175B 这类 Transformer 模型为例仅参数本身就需要数百 GB 内存远超单张 GPU 的容量。为此业界普遍采用模型并行Model Parallelism把模型切分到多台设备上每台设备只保存并计算一部分权重。**张量并行Tensor Parallelism**是模型并行的一种重要形式它把单个层内的权重张量按一定维度切分到多个设备上从而成比例地降低每台设备的参数内存负荷。Megatron-LM 在Efficient Large-Scale Language Model Training on GPU Clusters Using Megatron-LM中给出了一种高效的一维张量并行实现这也是 ColossalAI 中 1D 张量并行的理论来源。所谓一维是指仅沿张量的一个维度进行切分区别于 2D、2.5D、3D 张量并行沿两个或多个维度切分参见同目录下 2D_tensor_parallel.md、2p5D_tensor_parallel.md 与 3D_tensor_parallel.md。它实现简单、通信开销可控是最基础也最常用的一种张量切分策略。2 核心原理从线性层的 GEMM 说起一个线性层本质上是矩阵乘GEMM。设输入为 $X$权重为 $A$线性层前向计算 $Y XA$。假设使用 2 个处理器rank可以将 $A$ 沿列切分为 $[A_1 ~~ A_2]$每个处理器分别持有 $A_1$、$A_2$ 并计算 $Y_i XA_i$最终拼接得到$$ [Y_1 ~ Y_2] [XA_1 ~ XA_2] $$由于每个处理器上的 $Y_i$ 只是最终输出 $Y$ 的一个切片这种权重按列切分、输出按列分布的方式称为列并行column-parallel。它的关键收益是权重矩阵被分到了不同设备上单卡无需保存完整权重。2.1 列并行线性层当第一个线性层采用列并行后紧随其后的第二个线性层 $ZYB$ 需要处理 $Y$ 是按列分散的这一问题。若第二个线性层仍按列切分 $B$就必须先通过通信把分散的 $Y$ 重新聚合为完整张量这会引入额外通信。Megatron-LM 的做法是把 $B$ 按行切分为$$ \left[\begin{matrix} B_1 \ B_2 \end{matrix} \right] $$这种权重按行切分的方式称为行并行row-parallel。此时$$ Z [Y_1 ~ Y_2] \left[\begin{matrix} B_1 \ B_2 \end{matrix} \right] $$的计算被分解为每个处理器先在本地计算 $Y_iB_i$注意 $Y_i$ 的列数与 $B_i$ 的行数正好匹配随后通过一次all-reduce将各处理器的部分结果相加得到完整输出 $ZY_1B_1Y_2B_2$。这一列并行 → 行并行的成对设计非常关键上一层的列切分恰好与下一层的行切分首尾咬合因此中间隐藏状态不需要全部落盘与复制前向只需在行并行层的输出处做一次 all-reduce即可得到完整结果。2.2 后向传播输入梯度的 all-reduce张量并行不仅影响前向也重塑了后向传播中的通信模式。以列并行线性层为例记损失对 $Y_i$ 的梯度为 $\dot{Y_i}$由于每个处理器只持有权重分片 $A_i$其本地只能算出部分输入梯度$$ \dot{X_i}\dot{Y_i}A_i^T $$而完整的输入梯度应当是 $\dot{X}\dot{Y}A^T\dot{Y_1}A_1^T\dot{Y_2}A_2^T$。因此列并行线性层在后向传播中需要对输入梯度执行一次 all-reduce把所有处理器的部分梯度累加为完整梯度才能向上层继续传播。对称地前向中做 all-reduce 的行并行层其权重梯度与偏置梯度的计算方式也有讲究——详见第 4 节源码分析。整体上1D 张量并行的通信开销集中体现在这两类跨设备规约操作上。3 理论效率分析给定 $P$ 个处理器文档给出 1D 张量并行基于环形通信算法 ring algorithm的理论成本如下计算内存参数内存activations通信带宽通信时延$O(1/P)$$O(1/P)$$O(1)$$O(2(P-1)/P)$$O(2(P-1))$对该表做几点解读计算 $O(1/P)$权重分片后每台设备只做全量矩阵乘的 $1/P$ 计算量理论上单步吞吐不因并行而损失。参数内存 $O(1/P)$这是张量并行最直接的收益——权重、以及由优化器状态带来的内存占用随 $P$ 反比例下降。配合Shardformer只保留本设备负责的参数分片这一设计见第 5 节内存节省效果显著。激活内存 $O(1)$注意仅靠 1D 张量并行不会降低激活值内存激活仍是复制/整存的表中系数为 1。这正是 Shardformer 文档 中**序列并行Sequence Parallelism**存在的原因——它沿序列维度进一步切分激活可与 1D 张量并行叠加使用以削减激活占用。通信带宽 $O(2(P-1)/P)$ 与时延 $O(2(P-1))$前向行并行输出处的 all-reduce与后向列并行输入梯度处的 all-reduce各发生一次规约。使用 NCCL 的环形 all-reduce 时一次规约由 reduce-scatter 与 all-gather 构成。这也解释了为何工程上普遍将张量并行度限制在单机 8 卡等较小范围——随着 $P$ 增大时延项 $O(2(P-1))$ 与每卡通信占比都会上升。4 在 ColossalAI 中使用由 Shardformer 提供在 ColossalAI 的最新版本中1D 张量并行统一由Shardformer功能实现不需要用户手工改写模型。Shardformer为 HuggingFace Transformers 生态的主流模型自动注入张量并行 / 流水线并行策略用户只需几行代码即可让 Transformers 模型进入可并行训练状态同时可按需叠加 Flash Attention、Fused LayerNorm、序列并行等优化。其设计与原理详见 Shardformer 文档。Shardformer支持 1D 张量并行的模型包括 Llama V1/V2、OPT、BLOOM、ChatGLM 2、BERT、GPT 2、T5、ViT、Whisper、Falcon、Blip2 等完整兼容性矩阵见 shardformer.md。4.1 方式一通过 HybridParallelPlugin 启动推荐把 1D 张量并行与其他并行策略、混合精度、ZeRO 等组合训练时最推荐的方式是用HybridParallelPlugin初始化的Booster来启动。例如在 examples/language/bert 微调示例中plugin HybridParallelPlugin( tp_size2, # 张量并行度即把模型切分到几个 GPU pp_size1, num_microbatchesNone, microbatch_size1, enable_all_optimizationTrue, zero_stage1, precisionfp16, initial_scale1, )随后用Booster封装模型与优化器即可Shardformer的封装由插件自动完成。此处tp_size即对应本文讨论的张量并行处理器数量 $P$设置tp_size1表示关闭张量并行。更多插件参数见 Booster 插件文档。4.2 方式二直接调用 Shardformer API不依赖Booster时也可以手动调用ShardformerAPI注意此时流水线并行无法正常工作因此仅适用于纯张量并行等场景。参考官方收敛性测试脚本 colossalai/shardformer/examples/convergence_benchmark.py 中的核心片段import torch.distributed as dist from colossalai.shardformer import ShardConfig, ShardFormer # 1. 为张量并行创建 process group tp_group dist.new_group(backendnccl) # 2. 创建 Shardformer 配置 shard_config ShardConfig( tensor_parallel_process_grouptp_group, enable_tensor_parallelismTrue, enable_all_optimizationTrue, ) # 3. 用该配置实例化 Shardformer shard_former ShardFormer(shard_configshard_config) # 4. 对 HuggingFace 模型执行封装返回并行化后的模型与共享参数 model, shared_params shard_former.optimize(model)封装完成后模型即可按常规方式进行前向/后向与优化器更新model(**batch)、loss.backward()。脚本中这一套流程对 GLUE 的 MRPC 等任务做了收敛性校验--target_f1参数可用作验证Shardformer改造正确性的参考。4.3 ShardConfig 关键参数说明ShardConfig的字段定义于 colossalai/shardformer/init.py各参数含义综合 colossalai/shardformer/README.md 与 Shardformer 文档参数含义默认值tensor_parallel_process_group张量并行使用的通信进程组使用张量并行时需提供None全局进程组pipeline_stage_manager流水线并行使用的 stage manager不使用流水线时保持NoneNoneenable_tensor_parallelism是否启用张量并行Trueenable_fused_normalization是否启用融合 LayerNorm需安装apexFalseenable_flash_attention是否启用 Flash Attention需安装flash_attnFalseenable_jit_fused是否启用 JIT 融合算子Falseenable_sequence_parallelism是否启用序列并行仅与 1D 张量并行配合沿序列维切分激活Falseenable_sequence_overlap是否将序列并行的通信与计算重叠需先开启enable_sequence_parallelismFalseenable_all_optimization一键开启上述全部优化融合归一化、Flash Attention、JIT 融合算子、序列并行及重叠Falseextra_kwargs传递给 Shardformer 的额外关键字参数{}需要说明的使用约束启用 Flash Attention 前需pip install flash_attnxFormers 的cutlass_op可作为 Flash Attention 的补充优化。使用GPT2ForSequenceClassification、ViTForImageClassification等分类模型时请确保标签总数为张量并行度的整数倍否则Shardformer无法正确切分 classifier 层可在 transformers 的 config 中加入虚拟标签解决。依赖apex的融合算子与 FlashAttention 等需与本地 CUDA 环境匹配启用前请确认依赖安装正确参见 requirements/requirements.txt。4.4 运行示例多卡启动可直接使用torchrun例如对收敛性基准测试# 进入 colossalai/shardformer/examples 目录后 torchrun --standalone --nproc_per_node 4 convergence_benchmark.py \ --model bert --pretrain bert-base-uncased --task mrpc --target_f1 0.86其中--nproc_per_node应不小于张量并行度 $P$Shardformer会自动把 BERT 中的nn.Linear、nn.Embedding等模块替换为分布式的 1D 并行模块。5 源码级解析Linear1D_Col 与 Linear1D_Row 的实现文档描述的列并行/行并行语义在仓库中由两个分布式模块实现路径为 colossalai/shardformer/layer/linear.py5.1Linear1D_Col列并行层对应数学上 $YXA$、权重沿列切分为 $[A_1, \dots, A_P]$ 的情形。其前向逻辑Linear1D_Col.forward只在本地执行F.linear(input_, weight)得到输出切片 $Y_i$前向无需通信而真正的规约发生在后向代码走linear_with_async_comm(..., async_grad_allreduceTrue)实现在 colossalai/shardformer/layer/_operation.py 的LinearWithAsyncCommunication类中。其后向通过dist.all_reduce(grad_input, async_opTrue)对输入梯度做异步 all-reduce恰好对应文档 2.2 节的推导 $\dot{X}\dot{Y_1}A_1^T\dot{Y_2}A_2^T$由于后端为异步集合通信需要CUDA_DEVICE_MAX_CONNECTIONS1来保证 all-reduce 被优先调度该环境变量由 shardformer 封装时设置。Linear1D_Col还暴露了若干工程参数gather_outputFalse默认前向不额外 all-gather 输出保留 $Y_i$ 分片便于喂给下一个行并行层从而省去一次冗余通信gather_outputTrue前向末尾调用gather_forward_split_backward把各卡输出聚合为完整 $Y$用于必须输出完整张量的场景如最后的 LM Head权重与偏置分别通过shard_rowwise/shard_colwise来自 colossalai/tensor/d_tensor转为分布式张量各卡只保留自己负责的参数分片from_native_module中带约束检查当out_features % tp_size ! 0时抛出ValueError提示用户该层输出维度必须能被张量并行度整除若out_features tp_size则退化为普通层。5.2Linear1D_Row行并行层对应权重 $B$ 按行切分为 $[B_1; \dots; B_P]$ 的情形。其前向Linear1D_Row.forward当输入已经是按特征维切分的分片parallel_inputTrue即前一列并行层输出 $Y_i$时本地执行F.linear(input_, weight)得 $Y_iB_i$再通过reduce_forward对输出做一次all-reduce得到完整 $Z\sum_i Y_iB_i$——这与文档公式完全一致权重经shard_colwise切分由于行并行层的偏置 $b$ 在语义上作用于完整输出其初始化为全量偏置并在reset_parameters中从 rank 0 通过dist.broadcast广播复制到各卡保证各卡偏置一致反向传播中行并行层则对输入梯度做 all-gather 类的聚合split_forward_gather_backward形成与前向对称的通信模式。顺带一提linear.py中还提供了VocabParallelLMHead1D对词表维度做可整除化的列并行 LM Head与PaddingLMHead用于词表大小不能被 $P$ 整除的语言模型输出层场景这正是 Shardformer 兼容矩阵中多类生成式模型的切分基础。5.3 与序列并行的组合从 Shardformer 文档 可以看到序列并行是Shardformer面向 1D 张量并行的配套优化它把非张量并行区域如 LayerNorm、Dropout、激活等沿序列维度切分从而把文档效率表中 $O(1)$ 的激活内存降下来。其要点包括列并行层在序列并行下前向需 all-gather 输入、后向做 reduce-scatter行并行层前向做 reduce-scatter、后向做 all-gatherNCCL 的 all-reduce 由 reduce-scatter all-gather 组成两者开销等价故序列并行 张量并行不会引入额外通信开销列并行层反向求梯度需要完整输入但实现中可通过enable_sequence_overlap将输入 all-gather 与梯度计算重叠同样不引入额外通信延迟。相关集合通信原语_GatherForwardReduceScatterBackward、_ReduceScatterForwardGatherBackward、环形实现_ring_as_gather/_ring_as_reducescatter均实现在 colossalai/shardformer/layer/_operation.py。6 总结与下一步1D 张量并行把单个线性层/注意力/MLP 内的权重沿一个维度切分到 $P$ 个设备将参数内存降为原来的 $1/P$代价是前向与后向各引入一次基于环形的 all-reduce 通信。在 ColossalAI 中你无需手工切分模型直接通过HybridParallelPlugin推荐或ShardFormerShardConfig即可在任意受支持的 HuggingFace 模型上启用 1D 张量并行并可叠加 Flash Attention、序列并行、ZeRO 等能力。若希望进一步了解与 1D 并行配套的更高维并行与进阶调度可继续阅读仓库内同目录文档2D 张量并行、2.5D 张量并行、3D 张量并行、流水线并行 与 序列并行。相关可运行示例集中在 colossalai/shardformer/examples收敛性与性能基准与 examples/language/bert经HybridParallelPlugin微调 BERT。【免费下载链接】ColossalAIMaking large AI models cheaper, faster and more accessible项目地址: https://gitcode.com/GitHub_Trending/co/ColossalAI创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考 ColossalAI 一维张量并行（1D Tensor Parallelism）原理与 Shardformer 实战指南
ColossalAI 一维张量并行1D Tensor Parallelism原理与 Shardformer 实战指南【免费下载链接】ColossalAIMaking large AI models cheaper, faster and more accessible项目地址: https://gitcode.com/GitHub_Trending/co/ColossalAI本文以 ColossalAI 官方中文文档《1D 张量并行》为骨架系统讲解一维张量并行的数学原理列并行/行并行切分方式、前向与后向的通信模式、理论效率边界并结合仓库源码深入说明其在 ColossalAI 中经Shardformer的实现细节、参数配置与运行方式。读完本文你将理解 1D 张量并行为什么省显存、为什么需要通信、何时需要 all-reduce并能在自己基于 HuggingFace Transformers 的模型上正确启用它。本文对应文档1D 张量并行作者 Zhengda Bian, Yongbin Li其实现与配套示例同时可参考 Shardformer 文档 与 shardformer 示例目录。1 引言为什么需要张量并行大语言模型训练面临的首要瓶颈是单卡显存。以 LLaMA-2 70B、OPT 175B 这类 Transformer 模型为例仅参数本身就需要数百 GB 内存远超单张 GPU 的容量。为此业界普遍采用模型并行Model Parallelism把模型切分到多台设备上每台设备只保存并计算一部分权重。**张量并行Tensor Parallelism**是模型并行的一种重要形式它把单个层内的权重张量按一定维度切分到多个设备上从而成比例地降低每台设备的参数内存负荷。Megatron-LM 在Efficient Large-Scale Language Model Training on GPU Clusters Using Megatron-LM中给出了一种高效的一维张量并行实现这也是 ColossalAI 中 1D 张量并行的理论来源。所谓一维是指仅沿张量的一个维度进行切分区别于 2D、2.5D、3D 张量并行沿两个或多个维度切分参见同目录下 2D_tensor_parallel.md、2p5D_tensor_parallel.md 与 3D_tensor_parallel.md。它实现简单、通信开销可控是最基础也最常用的一种张量切分策略。2 核心原理从线性层的 GEMM 说起一个线性层本质上是矩阵乘GEMM。设输入为 $X$权重为 $A$线性层前向计算 $Y XA$。假设使用 2 个处理器rank可以将 $A$ 沿列切分为 $[A_1 ~~ A_2]$每个处理器分别持有 $A_1$、$A_2$ 并计算 $Y_i XA_i$最终拼接得到$$ [Y_1 ~ Y_2] [XA_1 ~ XA_2] $$由于每个处理器上的 $Y_i$ 只是最终输出 $Y$ 的一个切片这种权重按列切分、输出按列分布的方式称为列并行column-parallel。它的关键收益是权重矩阵被分到了不同设备上单卡无需保存完整权重。2.1 列并行线性层当第一个线性层采用列并行后紧随其后的第二个线性层 $ZYB$ 需要处理 $Y$ 是按列分散的这一问题。若第二个线性层仍按列切分 $B$就必须先通过通信把分散的 $Y$ 重新聚合为完整张量这会引入额外通信。Megatron-LM 的做法是把 $B$ 按行切分为$$ \left[\begin{matrix} B_1 \ B_2 \end{matrix} \right] $$这种权重按行切分的方式称为行并行row-parallel。此时$$ Z [Y_1 ~ Y_2] \left[\begin{matrix} B_1 \ B_2 \end{matrix} \right] $$的计算被分解为每个处理器先在本地计算 $Y_iB_i$注意 $Y_i$ 的列数与 $B_i$ 的行数正好匹配随后通过一次all-reduce将各处理器的部分结果相加得到完整输出 $ZY_1B_1Y_2B_2$。这一列并行 → 行并行的成对设计非常关键上一层的列切分恰好与下一层的行切分首尾咬合因此中间隐藏状态不需要全部落盘与复制前向只需在行并行层的输出处做一次 all-reduce即可得到完整结果。2.2 后向传播输入梯度的 all-reduce张量并行不仅影响前向也重塑了后向传播中的通信模式。以列并行线性层为例记损失对 $Y_i$ 的梯度为 $\dot{Y_i}$由于每个处理器只持有权重分片 $A_i$其本地只能算出部分输入梯度$$ \dot{X_i}\dot{Y_i}A_i^T $$而完整的输入梯度应当是 $\dot{X}\dot{Y}A^T\dot{Y_1}A_1^T\dot{Y_2}A_2^T$。因此列并行线性层在后向传播中需要对输入梯度执行一次 all-reduce把所有处理器的部分梯度累加为完整梯度才能向上层继续传播。对称地前向中做 all-reduce 的行并行层其权重梯度与偏置梯度的计算方式也有讲究——详见第 4 节源码分析。整体上1D 张量并行的通信开销集中体现在这两类跨设备规约操作上。3 理论效率分析给定 $P$ 个处理器文档给出 1D 张量并行基于环形通信算法 ring algorithm的理论成本如下计算内存参数内存activations通信带宽通信时延$O(1/P)$$O(1/P)$$O(1)$$O(2(P-1)/P)$$O(2(P-1))$对该表做几点解读计算 $O(1/P)$权重分片后每台设备只做全量矩阵乘的 $1/P$ 计算量理论上单步吞吐不因并行而损失。参数内存 $O(1/P)$这是张量并行最直接的收益——权重、以及由优化器状态带来的内存占用随 $P$ 反比例下降。配合Shardformer只保留本设备负责的参数分片这一设计见第 5 节内存节省效果显著。激活内存 $O(1)$注意仅靠 1D 张量并行不会降低激活值内存激活仍是复制/整存的表中系数为 1。这正是 Shardformer 文档 中**序列并行Sequence Parallelism**存在的原因——它沿序列维度进一步切分激活可与 1D 张量并行叠加使用以削减激活占用。通信带宽 $O(2(P-1)/P)$ 与时延 $O(2(P-1))$前向行并行输出处的 all-reduce与后向列并行输入梯度处的 all-reduce各发生一次规约。使用 NCCL 的环形 all-reduce 时一次规约由 reduce-scatter 与 all-gather 构成。这也解释了为何工程上普遍将张量并行度限制在单机 8 卡等较小范围——随着 $P$ 增大时延项 $O(2(P-1))$ 与每卡通信占比都会上升。4 在 ColossalAI 中使用由 Shardformer 提供在 ColossalAI 的最新版本中1D 张量并行统一由Shardformer功能实现不需要用户手工改写模型。Shardformer为 HuggingFace Transformers 生态的主流模型自动注入张量并行 / 流水线并行策略用户只需几行代码即可让 Transformers 模型进入可并行训练状态同时可按需叠加 Flash Attention、Fused LayerNorm、序列并行等优化。其设计与原理详见 Shardformer 文档。Shardformer支持 1D 张量并行的模型包括 Llama V1/V2、OPT、BLOOM、ChatGLM 2、BERT、GPT 2、T5、ViT、Whisper、Falcon、Blip2 等完整兼容性矩阵见 shardformer.md。4.1 方式一通过 HybridParallelPlugin 启动推荐把 1D 张量并行与其他并行策略、混合精度、ZeRO 等组合训练时最推荐的方式是用HybridParallelPlugin初始化的Booster来启动。例如在 examples/language/bert 微调示例中plugin HybridParallelPlugin( tp_size2, # 张量并行度即把模型切分到几个 GPU pp_size1, num_microbatchesNone, microbatch_size1, enable_all_optimizationTrue, zero_stage1, precisionfp16, initial_scale1, )随后用Booster封装模型与优化器即可Shardformer的封装由插件自动完成。此处tp_size即对应本文讨论的张量并行处理器数量 $P$设置tp_size1表示关闭张量并行。更多插件参数见 Booster 插件文档。4.2 方式二直接调用 Shardformer API不依赖Booster时也可以手动调用ShardformerAPI注意此时流水线并行无法正常工作因此仅适用于纯张量并行等场景。参考官方收敛性测试脚本 colossalai/shardformer/examples/convergence_benchmark.py 中的核心片段import torch.distributed as dist from colossalai.shardformer import ShardConfig, ShardFormer # 1. 为张量并行创建 process group tp_group dist.new_group(backendnccl) # 2. 创建 Shardformer 配置 shard_config ShardConfig( tensor_parallel_process_grouptp_group, enable_tensor_parallelismTrue, enable_all_optimizationTrue, ) # 3. 用该配置实例化 Shardformer shard_former ShardFormer(shard_configshard_config) # 4. 对 HuggingFace 模型执行封装返回并行化后的模型与共享参数 model, shared_params shard_former.optimize(model)封装完成后模型即可按常规方式进行前向/后向与优化器更新model(**batch)、loss.backward()。脚本中这一套流程对 GLUE 的 MRPC 等任务做了收敛性校验--target_f1参数可用作验证Shardformer改造正确性的参考。4.3 ShardConfig 关键参数说明ShardConfig的字段定义于 colossalai/shardformer/init.py各参数含义综合 colossalai/shardformer/README.md 与 Shardformer 文档参数含义默认值tensor_parallel_process_group张量并行使用的通信进程组使用张量并行时需提供None全局进程组pipeline_stage_manager流水线并行使用的 stage manager不使用流水线时保持NoneNoneenable_tensor_parallelism是否启用张量并行Trueenable_fused_normalization是否启用融合 LayerNorm需安装apexFalseenable_flash_attention是否启用 Flash Attention需安装flash_attnFalseenable_jit_fused是否启用 JIT 融合算子Falseenable_sequence_parallelism是否启用序列并行仅与 1D 张量并行配合沿序列维切分激活Falseenable_sequence_overlap是否将序列并行的通信与计算重叠需先开启enable_sequence_parallelismFalseenable_all_optimization一键开启上述全部优化融合归一化、Flash Attention、JIT 融合算子、序列并行及重叠Falseextra_kwargs传递给 Shardformer 的额外关键字参数{}需要说明的使用约束启用 Flash Attention 前需pip install flash_attnxFormers 的cutlass_op可作为 Flash Attention 的补充优化。使用GPT2ForSequenceClassification、ViTForImageClassification等分类模型时请确保标签总数为张量并行度的整数倍否则Shardformer无法正确切分 classifier 层可在 transformers 的 config 中加入虚拟标签解决。依赖apex的融合算子与 FlashAttention 等需与本地 CUDA 环境匹配启用前请确认依赖安装正确参见 requirements/requirements.txt。4.4 运行示例多卡启动可直接使用torchrun例如对收敛性基准测试# 进入 colossalai/shardformer/examples 目录后 torchrun --standalone --nproc_per_node 4 convergence_benchmark.py \ --model bert --pretrain bert-base-uncased --task mrpc --target_f1 0.86其中--nproc_per_node应不小于张量并行度 $P$Shardformer会自动把 BERT 中的nn.Linear、nn.Embedding等模块替换为分布式的 1D 并行模块。5 源码级解析Linear1D_Col 与 Linear1D_Row 的实现文档描述的列并行/行并行语义在仓库中由两个分布式模块实现路径为 colossalai/shardformer/layer/linear.py5.1Linear1D_Col列并行层对应数学上 $YXA$、权重沿列切分为 $[A_1, \dots, A_P]$ 的情形。其前向逻辑Linear1D_Col.forward只在本地执行F.linear(input_, weight)得到输出切片 $Y_i$前向无需通信而真正的规约发生在后向代码走linear_with_async_comm(..., async_grad_allreduceTrue)实现在 colossalai/shardformer/layer/_operation.py 的LinearWithAsyncCommunication类中。其后向通过dist.all_reduce(grad_input, async_opTrue)对输入梯度做异步 all-reduce恰好对应文档 2.2 节的推导 $\dot{X}\dot{Y_1}A_1^T\dot{Y_2}A_2^T$由于后端为异步集合通信需要CUDA_DEVICE_MAX_CONNECTIONS1来保证 all-reduce 被优先调度该环境变量由 shardformer 封装时设置。Linear1D_Col还暴露了若干工程参数gather_outputFalse默认前向不额外 all-gather 输出保留 $Y_i$ 分片便于喂给下一个行并行层从而省去一次冗余通信gather_outputTrue前向末尾调用gather_forward_split_backward把各卡输出聚合为完整 $Y$用于必须输出完整张量的场景如最后的 LM Head权重与偏置分别通过shard_rowwise/shard_colwise来自 colossalai/tensor/d_tensor转为分布式张量各卡只保留自己负责的参数分片from_native_module中带约束检查当out_features % tp_size ! 0时抛出ValueError提示用户该层输出维度必须能被张量并行度整除若out_features tp_size则退化为普通层。5.2Linear1D_Row行并行层对应权重 $B$ 按行切分为 $[B_1; \dots; B_P]$ 的情形。其前向Linear1D_Row.forward当输入已经是按特征维切分的分片parallel_inputTrue即前一列并行层输出 $Y_i$时本地执行F.linear(input_, weight)得 $Y_iB_i$再通过reduce_forward对输出做一次all-reduce得到完整 $Z\sum_i Y_iB_i$——这与文档公式完全一致权重经shard_colwise切分由于行并行层的偏置 $b$ 在语义上作用于完整输出其初始化为全量偏置并在reset_parameters中从 rank 0 通过dist.broadcast广播复制到各卡保证各卡偏置一致反向传播中行并行层则对输入梯度做 all-gather 类的聚合split_forward_gather_backward形成与前向对称的通信模式。顺带一提linear.py中还提供了VocabParallelLMHead1D对词表维度做可整除化的列并行 LM Head与PaddingLMHead用于词表大小不能被 $P$ 整除的语言模型输出层场景这正是 Shardformer 兼容矩阵中多类生成式模型的切分基础。5.3 与序列并行的组合从 Shardformer 文档 可以看到序列并行是Shardformer面向 1D 张量并行的配套优化它把非张量并行区域如 LayerNorm、Dropout、激活等沿序列维度切分从而把文档效率表中 $O(1)$ 的激活内存降下来。其要点包括列并行层在序列并行下前向需 all-gather 输入、后向做 reduce-scatter行并行层前向做 reduce-scatter、后向做 all-gatherNCCL 的 all-reduce 由 reduce-scatter all-gather 组成两者开销等价故序列并行 张量并行不会引入额外通信开销列并行层反向求梯度需要完整输入但实现中可通过enable_sequence_overlap将输入 all-gather 与梯度计算重叠同样不引入额外通信延迟。相关集合通信原语_GatherForwardReduceScatterBackward、_ReduceScatterForwardGatherBackward、环形实现_ring_as_gather/_ring_as_reducescatter均实现在 colossalai/shardformer/layer/_operation.py。6 总结与下一步1D 张量并行把单个线性层/注意力/MLP 内的权重沿一个维度切分到 $P$ 个设备将参数内存降为原来的 $1/P$代价是前向与后向各引入一次基于环形的 all-reduce 通信。在 ColossalAI 中你无需手工切分模型直接通过HybridParallelPlugin推荐或ShardFormerShardConfig即可在任意受支持的 HuggingFace 模型上启用 1D 张量并行并可叠加 Flash Attention、序列并行、ZeRO 等能力。若希望进一步了解与 1D 并行配套的更高维并行与进阶调度可继续阅读仓库内同目录文档2D 张量并行、2.5D 张量并行、3D 张量并行、流水线并行 与 序列并行。相关可运行示例集中在 colossalai/shardformer/examples收敛性与性能基准与 examples/language/bert经HybridParallelPlugin微调 BERT。【免费下载链接】ColossalAIMaking large AI models cheaper, faster and more accessible项目地址: https://gitcode.com/GitHub_Trending/co/ColossalAI创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考