尧图网页设计 
DeepSpeed 自动张量并行AutoTP推理指南无需注入策略即可为 HuggingFace 模型启用多卡切分【免费下载链接】DeepSpeedDeepSpeed is a deep learning optimization library that makes distributed training and inference easy, efficient, and effective.项目地址: https://gitcode.com/GitHub_Trending/de/DeepSpeed本文聚焦 DeepSpeed 在推理场景下的 Automatic Tensor Parallelism自动张量并行AutoTP能力当模型尚未获得专用 kernel injection内核注入支持、用户也未手写 injection policy注入策略时如何让模型在多个 GPU 上自动按张量并行切分并跑通、加速推理。文章以 DeepSpeed 官方教程 docs/_tutorials/automatic-tensor-parallelism.md 为主体骨架并结合 engine.py、auto_tp.py 等源码讲解其判定与切分原理。读完本文你将掌握AutoTP 与传统注入策略的差异、最小可用调用代码、性能验证脚本与启动命令、当前支持的模型范围以及如何在源码层面定位与排查问题。背景张量并行为什么需要注入策略张量并行Tensor ParallelismTP会把一个 Transformer 层中的大型权重矩阵按行或按列切分到多个 GPU 上每张卡只持有权重的一部分。切分本身并不难难的是哪些层需要列切分、哪些层需要行切分以及在什么位置插入跨卡的 all-reduce 通信来合并部分结果。在 DeepSpeed-Inference 的早期版本中要绕过未支持内核注入的模型用户必须通过injection_policy手动指出 Transformer Encoder/Decoder 层中的两个关键线性层attention 输出的 GeMM 层例如 T5 的SelfAttention.o、EncDecAttention.o整个 layer 输出的 GeMM 层例如 T5 的DenseReluDense.wo。这两个位置对应列切分→all-reduce→行切分的衔接点前一层按列切开后各部分结果必须通过一次 all-reduce 合并才能作为下一层的完整输入。DeepSpeed 只有知道这两处位置才能把所需的跨模型并行 rank 的通信逻辑正确地插入计算图。而 AutoTP 的目标正是把这一手工过程自动化——在运行时自动推断注入策略并应用用户不再需要填写任何层名。从源码结构看InferenceEngine 的初始化逻辑明确支持三种模式见 engine.py用户显式指定 injection policy走self.injection_dict分支内核注入replace_with_kernel_injectTrue整层替换为 DeepSpeed 专用推理 kernel自动张量并行当既没有提供 policy、也没有启用内核注入、且tp_size 1时自动触发。第三种模式即本文主题其入口是 engine.py 中对AutoTP.tp_parser(model)的调用与逐模块的 policy 应用。使用方法AutoTP 最小示例AutoTP 对用户完全透明。启用条件非常明确不启用 kernel injection、不提供 injection policy、并且mp_size或新参数tp_size大于 1。此时模型无需做任何改动即可完成张量并行推理。教程中给出的新方法是automatic-tensor-parallelism.md# --------------------------------------- # New automatic tensor parallelism method # --------------------------------------- import os import torch import transformers import deepspeed local_rank int(os.getenv(LOCAL_RANK, 0)) world_size int(os.getenv(WORLD_SIZE, 1)) # create the model pipeline pipe transformers.pipeline(tasktext2text-generation, modelgoogle/t5-v1_1-small, devicelocal_rank) # Initialize the DeepSpeed-Inference engine pipe.model deepspeed.init_inference( pipe.model, mp_sizeworld_size, dtypetorch.float ) output pipe(Input String)其中几个要点LOCAL_RANK、WORLD_SIZE由 DeepSpeed launcher或torchrun/deepspeed启动器注入环境变量本示例默认单卡也可运行world_size1时即不切分若在启动器之外直接执行需要保证每个进程的local_rank与device一一对应deepspeed.init_inference的完整签名与四种合法调用方式纯 config / 纯 kwargs / 两者混用等定义于 deepspeed/init.py内部会把 config dict 与 kwargs 合并后构造成DeepSpeedInferenceConfig。相比旧方法的本质区别为了体现 AutoTP 带来的简化教程同时给出了旧方法代码automatic-tensor-parallelism.md# ---------------------------------- # Previous tensor parallelism method # ---------------------------------- import os import torch import transformers import deepspeed from transformers.models.t5.modeling_t5 import T5Block local_rank int(os.getenv(LOCAL_RANK, 0)) world_size int(os.getenv(WORLD_SIZE, 1)) # create the model pipeline pipe transformers.pipeline(tasktext2text-generation, modelgoogle/t5-v1_1-small, devicelocal_rank) # Initialize the DeepSpeed-Inference engine pipe.model deepspeed.init_inference( pipe.model, mp_sizeworld_size, dtypetorch.float, injection_policy{T5Block: (SelfAttention.o, EncDecAttention.o, DenseReluDense.wo)} ) output pipe(Input String)对比可见新方法唯一少掉的是injection_policy{T5Block: (SelfAttention.o, ...)}。这两个版本实际走的是 engine 初始化里完全不同的两个分支——旧方法进入用户显式 policy分支新方法则进入AutoTP 自动解析分支。对于 AutoTP 已支持的模型policy 会在运行时确定并自动应用。同时兼容新式配置 API教程代码沿用了历史参数mp_size。在 inference/config.py 中mp_size已标记为deprecated推荐迁移到结构化参数tensor_parallel其内部模型为DeepSpeedTPConfig字段tp_size默认 1、tp_grain_size默认 64、可选mpu/tp_group。等效写法pipe.model deepspeed.init_inference( pipe.model, tensor_parallel{tp_size: world_size}, dtypetorch.float16, )其他常用的可选参数同样来自DeepSpeedInferenceConfig参数默认值说明replace_with_kernel_injectFalse置为True走专用推理 kernel 注入会禁用AutoTP 自动路径dtypetorch.float16支持torch.half/torch.int8/torch.float亦可传字符串由DtypeEnum转换tensor_parallel{tp_size: 1}张量并行规模tp_size 1且无 policy、无 kernel inject 时才触发 AutoTPkeep_module_on_hostFalse加载超大 checkpoint 时先把参数留在主机端、避免设备 OOM仅对有注入策略与 AutoTP 的模型生效max_out_tokens1024推理引擎一次能处理的最大 token 数含输入与输出超长生成需调大min_out_tokens1运行时对显存压力的最低保障无法满足时报错而非段错误实测脚本与启动命令AutoTP 的效果可以用文本生成推理测试脚本验证该脚本位于 DeepSpeed 官方配套的DeepSpeedExamples示例库inference/huggingface/text-generation/inference-test.py脚本会统计每 token 延迟、带宽、吞吐量以及显存占用便于横向对比。基线运行不启用 DeepSpeed、不做张量并行并加--test_performance收集性能数据deepspeed --num_gpus num_gpus DeepSpeedExamples/inference/huggingface/text-generation/inference-test.py --name model --batch_size batch_size --test_performance开启张量并行对兼容模型追加--ds_inference标志deepspeed --num_gpus num_gpus DeepSpeedExamples/inference/huggingface/text-generation/inference-test.py --name model --batch_size batch_size --test_performance --ds_inference其中--num_gpus即并行度需要与脚本内部解析到的WORLD_SIZE保持一致。注意该测试脚本本身位于 DeepSpeed 之外的示例仓库运行时需先将其克隆到本地并把命令中的路径替换为实际路径当前仓库内与之同类的自测代码可以参考 tests/unit 下的 AutoTP 相关用例。性能对比AutoTP 的实际收益教程在 V100 SXM2 32GB 上分别对 T5 11B 与 OPT 13B 做了 1/2/4 GPU 的对比实验并附有延迟与吞吐曲线图。T5 11B 推理性能对比T5-11B 在不同 GPU 规模下的每 token 延迟随批大小变化1/2/4 个 GPU 三条曲线批量越大、GPU 越多延迟收益越明显吞吐曲线同样显示 4 GPU 配置的 TFLOPS 显著高于 1/2 GPU配合曲线教程还给出了 T5-11B 的显存与吞吐数据表TestMemory Allocated per GPUMax Batch SizeMax Throughput per GPUNo TP or 1 GPU21.06 GB649.29 TFLOPS2 GPU TP10.56 GB32013.04 TFLOPS4 GPU TP5.31 GB76814.04 TFLOPSOPT 13B 推理性能对比OPT-13B 的单卡吞吐随批大小的变化曲线Baseline 与 2/4 GPU 对比其显存与吞吐汇总如下TestMemory Allocated per GPUMax Batch SizeMax Throughput per GPUNo TP23.94 GB21.65 TFlops2 GPU TP12.23 GB204.61 TFlops4 GPU TP6.36 GB564.90 TFlops需要强调以上数字是教程在**特定软硬件环境V100 SXM2 32GB**下采集的结果实际表现会随 GPU 型号、驱动、CUDA/版本、批大小与模型实现而变化请以本机复测为准。整体趋势可供参考张量并行通过摊薄每卡显存占用换取了更大的可行批大小从而在相同或更低的单卡显存下获得更高的吞吐。支持与不支持的模型范围教程给出经测试可用的模型家族如下未列出不代表必然不可用仅表示尚未测试albert、arctic、baichuan、bert、bigbird_pegasus、bloom、camembert、chatglm2、chatglm3、codegen、codellama、deberta_v2、electra、ernie、esm、falcon、glm、gpt-j、gpt-neo、gpt-neox、longt5、luke、llama、llama2、m2m_100、marian、mistral、mixtral、mpt、mvp、nezha、openai、opt、pegasus、perceiver、phi、plbart、qwen、qwen2、qwen2-moe、qwen2.5、qwen3、reformer、roberta、roformer、splinter、starcode、t5、xglm、xlm_roberta、yoso、yuan。以下模型当前不受 AutoTP 支持但其中部分仍可搭配 DeepSpeed 的其他能力使用例如 Bloom 可用内核注入deberta、flaubert、fsmt、gpt2、led、longformer、xlm、xlnet。这份名单与源码中的硬编码判定保持一致auto_tp.py 中的AutoTP.supported()内部维护了一个unsupported名单[deberta, flaubert, fsmt, gpt2, led, longformer, xlm, xlnet]通过从模型类名正则提取 model 类型后做大小写不敏感匹配来判定。若模型既不被 AutoTP 支持、又无可用内核容器 policytp_parser会抛出断言错误并提示AutoTP not supported for model. Please provide policy.。源码级原理AutoTP 如何在运行时自动找到那两个线性层AutoTP 的实现集中在 deepspeed/module_inject/auto_tp.py核心是AutoTP类的静态方法tp_parser见 auto_tp.py。它自动完成旧版本里用户手工填写injection_policy的全部工作大致分四步收集模块列表通过get_module_list遍历ModuleList等容器拿到模型里所有可替换的层模块如各 Decoder Block可行性预检调用AutoTP.supported()排除不支持家族并用AutoTP.kernel_supported()判断是否本可走内核注入仅在二者都不满足时报错逐层定位 all-reduce 线性层对每一层递归下钻出线性算子并借助 LayerNorm或ln前后相邻层、以及各模型家族特有命名来识别需要行切分 all-reduce的层——从源码可看到大量按命名启发式归纳的规则例如通用规则out_proj、o_proj、down_projMoE 的 down 投影被归入 all-reduce 层GPT-NeoXattention.denseFalconself_attention.denseMixtralMoE 的w2其前向为w2*act(w1*w3)需行切分;ChatGLMself_attention.dense、dense_4h_to_hPhiself_attn.dense组装 policy 列表把层类型 → 该层需要 all-reduce 的 Gemm 名字集合组织为policy_list交由 engine.py 逐项执行_apply_injection_policy。真正的切分替换逻辑在AutoTP._replaceauto_tp.py及其行/列两个替换产物中列并行LinearLayer权重沿输出维切开每卡计算完整输出的一个分片通过 gather 或后续 all-reduce 恢复行并行LinearAllreduce权重沿输入维切开前向末尾对 partial 结果执行跨mp_group的 all-reduce这正是旧版本需要手工指定 attention 输出层与 layer 输出层的原因。其通信实现可见 layers.py 中两类的forward以及all_reduce在 TP group 上的实际调用。此外AutoTP还负责融合 QKV 权重切分对query_key_value等融合层用ReplaceWithTensorSlicing.strided_copynum_splits3按 Q/K/V 头边界交错切分auto_tp.pyEmbedding 切分_slice_embedding把词表维按 rank 切开auto_tp.py并在参数上记录通用 checkpoint 恢复元信息超参同步对每层持有的n_heads、hidden_size、num_key_value_heads等维度属性按同名策略同步缩小update_mp_params避免后续模块使用未切分的逻辑维度高层 config 化切分较新版本还支持通过partition_config/AutoTPConfigTPLayerSpec显式提供正则式切分规则_replace_with_config并支持从 HuggingFace 模型自带的tp_plan自动提取切分方案。正因这些逻辑集中在 AutoTP 内部用户侧的 API 才能保持无 policy的极简形态。训练侧的 AutoTP 使用方式配合 ZeRO 与deepspeed.initialize支持 preset/partition_config/tp_plan可参考同主题的另一篇教程 docs/_tutorials/autotp-training.md。常见问题与适用边界何时会走 AutoTP三个条件同时满足injection_policy未提供、replace_with_kernel_injectFalse默认、tp_size 1。若显式给了 policy 却同时开启 kernel injectengine 会直接断言报错二者不可混用。AutoTP 与内核注入的关系内核注入把整层替换为 DeepSpeed 专用推理 kernel性能往往更好但只覆盖少数被支持架构AutoTP 不替换 kernel而是把普通 PyTorch 层做张量切分 插入通信因此能覆盖远更广的 HuggingFace 模型家族。对于 AutoTP 不支持但仍可内核注入的模型如名单中的 Bloom提示信息会引导用户改用replace_with_kernel_injectTrue。结果正确性由什么保证从实现看前向的 all-reduce行并行末尾合并与权重/维度的同步切分共同保证各 rank 上部分和能正确合并Loading类auto_tp.py负责把完整权重按当前 rank 的切分规则拷贝进替换层。不均匀切分源码注释与训练教程均指出hidden/vocabulary 维度不必被tp_size整除AutoTP 支持不均匀分片并通过每 rank 各自的形状信息保存/恢复。性能数据的前提文中所列图表与表格均来自原教程在 V100 SXM2 32GB、特定软件栈下的测量仅作量级参考。进一步阅读原教程文档docs/_tutorials/automatic-tensor-parallelism.md推理引擎三分支初始化逻辑deepspeed/inference/engine.pyAutoTP 实现parser / 切分 / 权重加载deepspeed/module_inject/auto_tp.py行/列并行层与通信实现deepspeed/module_inject/layers.py推理配置项定义deepspeed/inference/config.py训练侧 AutoTP与 ZeRO 结合docs/_tutorials/autotp-training.md仓库内相关自测用例目录 tests/unit/v1/autotp【免费下载链接】DeepSpeedDeepSpeed is a deep learning optimization library that makes distributed training and inference easy, efficient, and effective.项目地址: https://gitcode.com/GitHub_Trending/de/DeepSpeed创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考 DeepSpeed 自动张量并行（AutoTP）推理指南：无需注入策略即可为 HuggingFace 模型启用多卡切分
DeepSpeed 自动张量并行AutoTP推理指南无需注入策略即可为 HuggingFace 模型启用多卡切分【免费下载链接】DeepSpeedDeepSpeed is a deep learning optimization library that makes distributed training and inference easy, efficient, and effective.项目地址: https://gitcode.com/GitHub_Trending/de/DeepSpeed本文聚焦 DeepSpeed 在推理场景下的 Automatic Tensor Parallelism自动张量并行AutoTP能力当模型尚未获得专用 kernel injection内核注入支持、用户也未手写 injection policy注入策略时如何让模型在多个 GPU 上自动按张量并行切分并跑通、加速推理。文章以 DeepSpeed 官方教程 docs/_tutorials/automatic-tensor-parallelism.md 为主体骨架并结合 engine.py、auto_tp.py 等源码讲解其判定与切分原理。读完本文你将掌握AutoTP 与传统注入策略的差异、最小可用调用代码、性能验证脚本与启动命令、当前支持的模型范围以及如何在源码层面定位与排查问题。背景张量并行为什么需要注入策略张量并行Tensor ParallelismTP会把一个 Transformer 层中的大型权重矩阵按行或按列切分到多个 GPU 上每张卡只持有权重的一部分。切分本身并不难难的是哪些层需要列切分、哪些层需要行切分以及在什么位置插入跨卡的 all-reduce 通信来合并部分结果。在 DeepSpeed-Inference 的早期版本中要绕过未支持内核注入的模型用户必须通过injection_policy手动指出 Transformer Encoder/Decoder 层中的两个关键线性层attention 输出的 GeMM 层例如 T5 的SelfAttention.o、EncDecAttention.o整个 layer 输出的 GeMM 层例如 T5 的DenseReluDense.wo。这两个位置对应列切分→all-reduce→行切分的衔接点前一层按列切开后各部分结果必须通过一次 all-reduce 合并才能作为下一层的完整输入。DeepSpeed 只有知道这两处位置才能把所需的跨模型并行 rank 的通信逻辑正确地插入计算图。而 AutoTP 的目标正是把这一手工过程自动化——在运行时自动推断注入策略并应用用户不再需要填写任何层名。从源码结构看InferenceEngine 的初始化逻辑明确支持三种模式见 engine.py用户显式指定 injection policy走self.injection_dict分支内核注入replace_with_kernel_injectTrue整层替换为 DeepSpeed 专用推理 kernel自动张量并行当既没有提供 policy、也没有启用内核注入、且tp_size 1时自动触发。第三种模式即本文主题其入口是 engine.py 中对AutoTP.tp_parser(model)的调用与逐模块的 policy 应用。使用方法AutoTP 最小示例AutoTP 对用户完全透明。启用条件非常明确不启用 kernel injection、不提供 injection policy、并且mp_size或新参数tp_size大于 1。此时模型无需做任何改动即可完成张量并行推理。教程中给出的新方法是automatic-tensor-parallelism.md# --------------------------------------- # New automatic tensor parallelism method # --------------------------------------- import os import torch import transformers import deepspeed local_rank int(os.getenv(LOCAL_RANK, 0)) world_size int(os.getenv(WORLD_SIZE, 1)) # create the model pipeline pipe transformers.pipeline(tasktext2text-generation, modelgoogle/t5-v1_1-small, devicelocal_rank) # Initialize the DeepSpeed-Inference engine pipe.model deepspeed.init_inference( pipe.model, mp_sizeworld_size, dtypetorch.float ) output pipe(Input String)其中几个要点LOCAL_RANK、WORLD_SIZE由 DeepSpeed launcher或torchrun/deepspeed启动器注入环境变量本示例默认单卡也可运行world_size1时即不切分若在启动器之外直接执行需要保证每个进程的local_rank与device一一对应deepspeed.init_inference的完整签名与四种合法调用方式纯 config / 纯 kwargs / 两者混用等定义于 deepspeed/init.py内部会把 config dict 与 kwargs 合并后构造成DeepSpeedInferenceConfig。相比旧方法的本质区别为了体现 AutoTP 带来的简化教程同时给出了旧方法代码automatic-tensor-parallelism.md# ---------------------------------- # Previous tensor parallelism method # ---------------------------------- import os import torch import transformers import deepspeed from transformers.models.t5.modeling_t5 import T5Block local_rank int(os.getenv(LOCAL_RANK, 0)) world_size int(os.getenv(WORLD_SIZE, 1)) # create the model pipeline pipe transformers.pipeline(tasktext2text-generation, modelgoogle/t5-v1_1-small, devicelocal_rank) # Initialize the DeepSpeed-Inference engine pipe.model deepspeed.init_inference( pipe.model, mp_sizeworld_size, dtypetorch.float, injection_policy{T5Block: (SelfAttention.o, EncDecAttention.o, DenseReluDense.wo)} ) output pipe(Input String)对比可见新方法唯一少掉的是injection_policy{T5Block: (SelfAttention.o, ...)}。这两个版本实际走的是 engine 初始化里完全不同的两个分支——旧方法进入用户显式 policy分支新方法则进入AutoTP 自动解析分支。对于 AutoTP 已支持的模型policy 会在运行时确定并自动应用。同时兼容新式配置 API教程代码沿用了历史参数mp_size。在 inference/config.py 中mp_size已标记为deprecated推荐迁移到结构化参数tensor_parallel其内部模型为DeepSpeedTPConfig字段tp_size默认 1、tp_grain_size默认 64、可选mpu/tp_group。等效写法pipe.model deepspeed.init_inference( pipe.model, tensor_parallel{tp_size: world_size}, dtypetorch.float16, )其他常用的可选参数同样来自DeepSpeedInferenceConfig参数默认值说明replace_with_kernel_injectFalse置为True走专用推理 kernel 注入会禁用AutoTP 自动路径dtypetorch.float16支持torch.half/torch.int8/torch.float亦可传字符串由DtypeEnum转换tensor_parallel{tp_size: 1}张量并行规模tp_size 1且无 policy、无 kernel inject 时才触发 AutoTPkeep_module_on_hostFalse加载超大 checkpoint 时先把参数留在主机端、避免设备 OOM仅对有注入策略与 AutoTP 的模型生效max_out_tokens1024推理引擎一次能处理的最大 token 数含输入与输出超长生成需调大min_out_tokens1运行时对显存压力的最低保障无法满足时报错而非段错误实测脚本与启动命令AutoTP 的效果可以用文本生成推理测试脚本验证该脚本位于 DeepSpeed 官方配套的DeepSpeedExamples示例库inference/huggingface/text-generation/inference-test.py脚本会统计每 token 延迟、带宽、吞吐量以及显存占用便于横向对比。基线运行不启用 DeepSpeed、不做张量并行并加--test_performance收集性能数据deepspeed --num_gpus num_gpus DeepSpeedExamples/inference/huggingface/text-generation/inference-test.py --name model --batch_size batch_size --test_performance开启张量并行对兼容模型追加--ds_inference标志deepspeed --num_gpus num_gpus DeepSpeedExamples/inference/huggingface/text-generation/inference-test.py --name model --batch_size batch_size --test_performance --ds_inference其中--num_gpus即并行度需要与脚本内部解析到的WORLD_SIZE保持一致。注意该测试脚本本身位于 DeepSpeed 之外的示例仓库运行时需先将其克隆到本地并把命令中的路径替换为实际路径当前仓库内与之同类的自测代码可以参考 tests/unit 下的 AutoTP 相关用例。性能对比AutoTP 的实际收益教程在 V100 SXM2 32GB 上分别对 T5 11B 与 OPT 13B 做了 1/2/4 GPU 的对比实验并附有延迟与吞吐曲线图。T5 11B 推理性能对比T5-11B 在不同 GPU 规模下的每 token 延迟随批大小变化1/2/4 个 GPU 三条曲线批量越大、GPU 越多延迟收益越明显吞吐曲线同样显示 4 GPU 配置的 TFLOPS 显著高于 1/2 GPU配合曲线教程还给出了 T5-11B 的显存与吞吐数据表TestMemory Allocated per GPUMax Batch SizeMax Throughput per GPUNo TP or 1 GPU21.06 GB649.29 TFLOPS2 GPU TP10.56 GB32013.04 TFLOPS4 GPU TP5.31 GB76814.04 TFLOPSOPT 13B 推理性能对比OPT-13B 的单卡吞吐随批大小的变化曲线Baseline 与 2/4 GPU 对比其显存与吞吐汇总如下TestMemory Allocated per GPUMax Batch SizeMax Throughput per GPUNo TP23.94 GB21.65 TFlops2 GPU TP12.23 GB204.61 TFlops4 GPU TP6.36 GB564.90 TFlops需要强调以上数字是教程在**特定软硬件环境V100 SXM2 32GB**下采集的结果实际表现会随 GPU 型号、驱动、CUDA/版本、批大小与模型实现而变化请以本机复测为准。整体趋势可供参考张量并行通过摊薄每卡显存占用换取了更大的可行批大小从而在相同或更低的单卡显存下获得更高的吞吐。支持与不支持的模型范围教程给出经测试可用的模型家族如下未列出不代表必然不可用仅表示尚未测试albert、arctic、baichuan、bert、bigbird_pegasus、bloom、camembert、chatglm2、chatglm3、codegen、codellama、deberta_v2、electra、ernie、esm、falcon、glm、gpt-j、gpt-neo、gpt-neox、longt5、luke、llama、llama2、m2m_100、marian、mistral、mixtral、mpt、mvp、nezha、openai、opt、pegasus、perceiver、phi、plbart、qwen、qwen2、qwen2-moe、qwen2.5、qwen3、reformer、roberta、roformer、splinter、starcode、t5、xglm、xlm_roberta、yoso、yuan。以下模型当前不受 AutoTP 支持但其中部分仍可搭配 DeepSpeed 的其他能力使用例如 Bloom 可用内核注入deberta、flaubert、fsmt、gpt2、led、longformer、xlm、xlnet。这份名单与源码中的硬编码判定保持一致auto_tp.py 中的AutoTP.supported()内部维护了一个unsupported名单[deberta, flaubert, fsmt, gpt2, led, longformer, xlm, xlnet]通过从模型类名正则提取 model 类型后做大小写不敏感匹配来判定。若模型既不被 AutoTP 支持、又无可用内核容器 policytp_parser会抛出断言错误并提示AutoTP not supported for model. Please provide policy.。源码级原理AutoTP 如何在运行时自动找到那两个线性层AutoTP 的实现集中在 deepspeed/module_inject/auto_tp.py核心是AutoTP类的静态方法tp_parser见 auto_tp.py。它自动完成旧版本里用户手工填写injection_policy的全部工作大致分四步收集模块列表通过get_module_list遍历ModuleList等容器拿到模型里所有可替换的层模块如各 Decoder Block可行性预检调用AutoTP.supported()排除不支持家族并用AutoTP.kernel_supported()判断是否本可走内核注入仅在二者都不满足时报错逐层定位 all-reduce 线性层对每一层递归下钻出线性算子并借助 LayerNorm或ln前后相邻层、以及各模型家族特有命名来识别需要行切分 all-reduce的层——从源码可看到大量按命名启发式归纳的规则例如通用规则out_proj、o_proj、down_projMoE 的 down 投影被归入 all-reduce 层GPT-NeoXattention.denseFalconself_attention.denseMixtralMoE 的w2其前向为w2*act(w1*w3)需行切分;ChatGLMself_attention.dense、dense_4h_to_hPhiself_attn.dense组装 policy 列表把层类型 → 该层需要 all-reduce 的 Gemm 名字集合组织为policy_list交由 engine.py 逐项执行_apply_injection_policy。真正的切分替换逻辑在AutoTP._replaceauto_tp.py及其行/列两个替换产物中列并行LinearLayer权重沿输出维切开每卡计算完整输出的一个分片通过 gather 或后续 all-reduce 恢复行并行LinearAllreduce权重沿输入维切开前向末尾对 partial 结果执行跨mp_group的 all-reduce这正是旧版本需要手工指定 attention 输出层与 layer 输出层的原因。其通信实现可见 layers.py 中两类的forward以及all_reduce在 TP group 上的实际调用。此外AutoTP还负责融合 QKV 权重切分对query_key_value等融合层用ReplaceWithTensorSlicing.strided_copynum_splits3按 Q/K/V 头边界交错切分auto_tp.pyEmbedding 切分_slice_embedding把词表维按 rank 切开auto_tp.py并在参数上记录通用 checkpoint 恢复元信息超参同步对每层持有的n_heads、hidden_size、num_key_value_heads等维度属性按同名策略同步缩小update_mp_params避免后续模块使用未切分的逻辑维度高层 config 化切分较新版本还支持通过partition_config/AutoTPConfigTPLayerSpec显式提供正则式切分规则_replace_with_config并支持从 HuggingFace 模型自带的tp_plan自动提取切分方案。正因这些逻辑集中在 AutoTP 内部用户侧的 API 才能保持无 policy的极简形态。训练侧的 AutoTP 使用方式配合 ZeRO 与deepspeed.initialize支持 preset/partition_config/tp_plan可参考同主题的另一篇教程 docs/_tutorials/autotp-training.md。常见问题与适用边界何时会走 AutoTP三个条件同时满足injection_policy未提供、replace_with_kernel_injectFalse默认、tp_size 1。若显式给了 policy 却同时开启 kernel injectengine 会直接断言报错二者不可混用。AutoTP 与内核注入的关系内核注入把整层替换为 DeepSpeed 专用推理 kernel性能往往更好但只覆盖少数被支持架构AutoTP 不替换 kernel而是把普通 PyTorch 层做张量切分 插入通信因此能覆盖远更广的 HuggingFace 模型家族。对于 AutoTP 不支持但仍可内核注入的模型如名单中的 Bloom提示信息会引导用户改用replace_with_kernel_injectTrue。结果正确性由什么保证从实现看前向的 all-reduce行并行末尾合并与权重/维度的同步切分共同保证各 rank 上部分和能正确合并Loading类auto_tp.py负责把完整权重按当前 rank 的切分规则拷贝进替换层。不均匀切分源码注释与训练教程均指出hidden/vocabulary 维度不必被tp_size整除AutoTP 支持不均匀分片并通过每 rank 各自的形状信息保存/恢复。性能数据的前提文中所列图表与表格均来自原教程在 V100 SXM2 32GB、特定软件栈下的测量仅作量级参考。进一步阅读原教程文档docs/_tutorials/automatic-tensor-parallelism.md推理引擎三分支初始化逻辑deepspeed/inference/engine.pyAutoTP 实现parser / 切分 / 权重加载deepspeed/module_inject/auto_tp.py行/列并行层与通信实现deepspeed/module_inject/layers.py推理配置项定义deepspeed/inference/config.py训练侧 AutoTP与 ZeRO 结合docs/_tutorials/autotp-training.md仓库内相关自测用例目录 tests/unit/v1/autotp【免费下载链接】DeepSpeedDeepSpeed is a deep learning optimization library that makes distributed training and inference easy, efficient, and effective.项目地址: https://gitcode.com/GitHub_Trending/de/DeepSpeed创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考