拆解EnergonAI的张量并行实现Linear1D与词表并行的2个关键技巧【免费下载链接】EnergonAILarge-scale model inference.项目地址: https://gitcode.com/gh_mirrors/en/EnergonAIEnergonAI 是一个面向大模型推理的多 GPU 推理引擎其核心能力之一就是张量并行Tensor Parallelism把单个 Transformer 层的权重切分到多张 GPU 上协同计算从而让 66B 甚至 175B 参数的大模型跑进有限显存。本文将带你拆解 EnergonAI 张量并行实现中的 2 个关键技巧——Linear1D 列行切分与词表并行新手也能轻松看懂 1️⃣ 为什么大模型推理必须懂张量并行单卡显存有限而大模型的权重、KV 缓存越算越大。常见的解法有两条路线并行方式切分粒度通信特点适用场景张量并行TP每一层内部切分每次前向有多次通信节点内多卡、低延迟推理流水线并行PP按层切分仅在层间传激活值跨节点、超长层数EnergonAI 的做法是TP PP 双轨并行节点内用张量并行压低单卡显存节点间用流水线并行分摊深度。每个推理工作进程Worker在启动时就声明了自己的并行拓扑见 energonai/worker.py第 27–28 行colossalai.launch({parallel: {tensor: {mode: 1d, size: tp_world_size}, pipeline: pp_world_size}}, ...) 注意mode: 1d——这就是 EnergonAI 张量并行的技术底座ColossalAI 的 1D 张量并行后面提到的Linear1D、MLP1D、Block1D等命名正源于此。2️⃣ 技巧一Linear1D 列行切分通信量砍半Transformer 每层的核心是一个 MLP两层全连接hidden → intermediate → hidden。EnergonAI 的 energonai/model/mlp.py 中MLP1D是这样组合的子层使用的并行层关键参数作用dense_1Linear1D_Col列并行gather_outputFalse每张卡只算输出的 1/TP结果保持切分不合并dense_2Linear1D_Row行并行parallel_inputTrue接收切分输入内部做 all-reduce 得到完整结果这正是业界经典的列并行 行并行组合拳妙处在于两个技巧点中间激活全程保持切分状态gather_outputFalse让第一层算完后不急着 all-gather切分的数据直接喂给行并行层每层只付出 1 次 all-reduce 的通信代价all-reduce 被藏进了行并行层内部避免了两次完整张量聚合。注意力模块 energonai/model/attention.py 采用同款套路融合 QKV 投影用Linear1D_Col一次算出 3×hidden 的切分输出输出投影用Linear1D_Row。3️⃣ 技巧二词表并行Embedding 与 LM Head 也能切除了 MLP/Attention还有两个大块头容易撑爆单卡显存——词嵌入矩阵和语言模型头LM Head形状都是[词表大小, hidden]。以 OPT-175B 为例50272 × 12288 ≈ 12.3 亿参数。EnergonAI 的解法是按词表行切分词表并行词嵌入energonai/model/embedding.py 中的Embedding1D支持vocab_parallel开关。开启后使用VocabParallelEmbedding1D每张卡只持有 1/TP 的词表行查表后自动聚合模型拿到的是完整词向量LM Headenergonai/model/downstream.py 中的LMHead1D开启词表并行后使用VocabParallelClassifier1D(gather_outputTrue)每张卡只计算自己那部分词表上的 logits最后聚合为完整词表分布供采样使用。这两个开关最终由模型工厂统一配置。在 energonai/model/model_factory.py第 341–354 行中opt_175B配置就显式开启了vocab_parallelTrue与fused_qkvTrue配套的load_175b函数energonai/utils/checkpointing_opt.py负责把 175B 分片检查点正确加载到各张量并行 Rank 上。✅4️⃣ 幕后加速自研 CUDA 内核与 KV 缓存张量并行的效率还取决于单卡计算速度EnergonAI 在 energonai/kernel/cuda_native/ 下提供了一批自研内核GEMM 加速mlp_gemm直连 cuBLAScsrc/linear_wrapper.cppFP16 下批量完成 MLP 矩阵乘算法自动调优linear_func.py 的find_algo()会实测 30 多种 GEMM 算法组合自动选出最优算法注意力融合scale_mask_softmax、transpose_pad_fusion等内核把缩放、掩码、Softmax、转置合并减少显存搬运KV 缓存MLP1D、MultiHeadAttention1D内置past_cache自回归解码时只计算新 token 的中间层逐 token 推理显著提速。5️⃣ 快速上手3 步验证张量并行效果获取代码git clone https://gitcode.com/gh_mirrors/en/EnergonAI运行分布式测试执行 tests/run_standalone_tests.sh其中 tests/test_engine/test_tp.py 会启动 2 卡张量并行 1 流水线用内置的 BoringModel 校验多卡输出与单卡完全一致跑真实大模型参考 examples/opt/ 下的服务脚本配置tp_world_size启动 OPT 系列模型推理服务体验多卡吞吐提升。 总结EnergonAI 的张量并行实现可以浓缩为 2 句话算子层用Linear1D_ColLinear1D_Row列行切分覆盖 MLP 与 Attention中间激活保持切分通信代价最小化词表层用VocabParallelEmbedding1DVocabParallelClassifier1D把 Embedding 与 LM Head 按词表切分彻底释放显存瓶颈。配合自研 CUDA 内核与 KV 缓存这套组合拳让 EnergonAI 成为部署 66B / 175B 级大模型推理的实用选择。想要深入源码细节建议从energonai/model/目录下的 1D 系列模块读起 【免费下载链接】EnergonAILarge-scale model inference.项目地址: https://gitcode.com/gh_mirrors/en/EnergonAI创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考 拆解EnergonAI的张量并行实现：Linear1D与词表并行的2个关键技巧
拆解EnergonAI的张量并行实现Linear1D与词表并行的2个关键技巧【免费下载链接】EnergonAILarge-scale model inference.项目地址: https://gitcode.com/gh_mirrors/en/EnergonAIEnergonAI 是一个面向大模型推理的多 GPU 推理引擎其核心能力之一就是张量并行Tensor Parallelism把单个 Transformer 层的权重切分到多张 GPU 上协同计算从而让 66B 甚至 175B 参数的大模型跑进有限显存。本文将带你拆解 EnergonAI 张量并行实现中的 2 个关键技巧——Linear1D 列行切分与词表并行新手也能轻松看懂 1️⃣ 为什么大模型推理必须懂张量并行单卡显存有限而大模型的权重、KV 缓存越算越大。常见的解法有两条路线并行方式切分粒度通信特点适用场景张量并行TP每一层内部切分每次前向有多次通信节点内多卡、低延迟推理流水线并行PP按层切分仅在层间传激活值跨节点、超长层数EnergonAI 的做法是TP PP 双轨并行节点内用张量并行压低单卡显存节点间用流水线并行分摊深度。每个推理工作进程Worker在启动时就声明了自己的并行拓扑见 energonai/worker.py第 27–28 行colossalai.launch({parallel: {tensor: {mode: 1d, size: tp_world_size}, pipeline: pp_world_size}}, ...) 注意mode: 1d——这就是 EnergonAI 张量并行的技术底座ColossalAI 的 1D 张量并行后面提到的Linear1D、MLP1D、Block1D等命名正源于此。2️⃣ 技巧一Linear1D 列行切分通信量砍半Transformer 每层的核心是一个 MLP两层全连接hidden → intermediate → hidden。EnergonAI 的 energonai/model/mlp.py 中MLP1D是这样组合的子层使用的并行层关键参数作用dense_1Linear1D_Col列并行gather_outputFalse每张卡只算输出的 1/TP结果保持切分不合并dense_2Linear1D_Row行并行parallel_inputTrue接收切分输入内部做 all-reduce 得到完整结果这正是业界经典的列并行 行并行组合拳妙处在于两个技巧点中间激活全程保持切分状态gather_outputFalse让第一层算完后不急着 all-gather切分的数据直接喂给行并行层每层只付出 1 次 all-reduce 的通信代价all-reduce 被藏进了行并行层内部避免了两次完整张量聚合。注意力模块 energonai/model/attention.py 采用同款套路融合 QKV 投影用Linear1D_Col一次算出 3×hidden 的切分输出输出投影用Linear1D_Row。3️⃣ 技巧二词表并行Embedding 与 LM Head 也能切除了 MLP/Attention还有两个大块头容易撑爆单卡显存——词嵌入矩阵和语言模型头LM Head形状都是[词表大小, hidden]。以 OPT-175B 为例50272 × 12288 ≈ 12.3 亿参数。EnergonAI 的解法是按词表行切分词表并行词嵌入energonai/model/embedding.py 中的Embedding1D支持vocab_parallel开关。开启后使用VocabParallelEmbedding1D每张卡只持有 1/TP 的词表行查表后自动聚合模型拿到的是完整词向量LM Headenergonai/model/downstream.py 中的LMHead1D开启词表并行后使用VocabParallelClassifier1D(gather_outputTrue)每张卡只计算自己那部分词表上的 logits最后聚合为完整词表分布供采样使用。这两个开关最终由模型工厂统一配置。在 energonai/model/model_factory.py第 341–354 行中opt_175B配置就显式开启了vocab_parallelTrue与fused_qkvTrue配套的load_175b函数energonai/utils/checkpointing_opt.py负责把 175B 分片检查点正确加载到各张量并行 Rank 上。✅4️⃣ 幕后加速自研 CUDA 内核与 KV 缓存张量并行的效率还取决于单卡计算速度EnergonAI 在 energonai/kernel/cuda_native/ 下提供了一批自研内核GEMM 加速mlp_gemm直连 cuBLAScsrc/linear_wrapper.cppFP16 下批量完成 MLP 矩阵乘算法自动调优linear_func.py 的find_algo()会实测 30 多种 GEMM 算法组合自动选出最优算法注意力融合scale_mask_softmax、transpose_pad_fusion等内核把缩放、掩码、Softmax、转置合并减少显存搬运KV 缓存MLP1D、MultiHeadAttention1D内置past_cache自回归解码时只计算新 token 的中间层逐 token 推理显著提速。5️⃣ 快速上手3 步验证张量并行效果获取代码git clone https://gitcode.com/gh_mirrors/en/EnergonAI运行分布式测试执行 tests/run_standalone_tests.sh其中 tests/test_engine/test_tp.py 会启动 2 卡张量并行 1 流水线用内置的 BoringModel 校验多卡输出与单卡完全一致跑真实大模型参考 examples/opt/ 下的服务脚本配置tp_world_size启动 OPT 系列模型推理服务体验多卡吞吐提升。 总结EnergonAI 的张量并行实现可以浓缩为 2 句话算子层用Linear1D_ColLinear1D_Row列行切分覆盖 MLP 与 Attention中间激活保持切分通信代价最小化词表层用VocabParallelEmbedding1DVocabParallelClassifier1D把 Embedding 与 LM Head 按词表切分彻底释放显存瓶颈。配合自研 CUDA 内核与 KV 缓存这套组合拳让 EnergonAI 成为部署 66B / 175B 级大模型推理的实用选择。想要深入源码细节建议从energonai/model/目录下的 1D 系列模块读起 【免费下载链接】EnergonAILarge-scale model inference.项目地址: https://gitcode.com/gh_mirrors/en/EnergonAI创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考