ARTICLE DETAIL

资讯详情

深耕郑州网站建设与运营推广的一线实战洞察。

DeepSpeed AutoTP 自定义切分模式实战:以 Pythia 6.9B 为例的 `partition_config` 深度解析

DeepSpeed AutoTP 自定义切分模式实战:以 Pythia 6.9B 为例的 `partition_config` 深度解析 示例工程【免费下载链接】DeepSpeedExamplesExample models using DeepSpeed项目地址https://gitcode.com/gh_mirrors/de/DeepSpeedExamples点击查看免费下载本指南基于 DeepSpeedExamples 仓库中的 custom_patterns 示例完整讲解如何为不受 AutoTP 内置规格覆盖的模型以使用 fusedquery_key_value投影的 Pythia 6.9B / GPT-NeoX 为例编写自定义张量并行切分规则。你将掌握tensor_parallel.partition_config的完整配置语法、column / row 两种切分类型的选择逻辑、显式shape的用法以及如何配合 ZeRO-2 与 DP-rank 数据采样器跑通一个真实的多卡训练脚本。背景什么时候需要自定义切分模式AutoTP 的目标是让张量并行Tensor ParallelTP的切分工作自动化它识别模型参数中典型的命名与形状模式自动施加正确的切分方式。对于 Llama 这类受支持的架构只需要在 DeepSpeed 配置中开启tensor_parallel.autotp_size无需任何手工切分规则参见 basic_example。但当模型使用了非标准参数定义——例如 PythiaGPT-NeoX 架构把 Q、K、V 三个投影融合成一个query_key_value权重此时 AutoTP 无法确定融合张量内部如何切分就需要通过partition_config显式声明切分模式。这正是 custom_patterns 示例存在的意义它把「按层指定切分规则」的能力开放给用户同时用一个 10 条文本的内存数据集 DP-rank 随机采样器演示了完整的 TP DP 训练流程。核心配置partition_config完整解析示例中的完整 DeepSpeed 配置如下节选自 autotp_custom_patterns.pyds_config { zero_optimization: {stage: 2}, tensor_parallel: { autotp_size: args.tp_size, partition_config: { use_default_specs: False, layer_specs: [ { patterns: [.*(self_attention|attention)\\.query_key_value\\.weight$], partition_type: column, shape: ((q_size, kv_size, kv_size), -1), partition_dim: 0, }, { patterns: [.*(self_attention|attention)\\.dense\\.weight$], partition_type: row, }, { patterns: [.*mlp\\.dense_h_to_4h\\.weight$], partition_type: column, }, { patterns: [.*mlp\\.dense_4h_to_h\\.weight$], partition_type: row, }, ], }, }, data_parallel_size: args.dp_size, }各字段含义与要点配置项作用说明autotp_size张量并行组大小与data_parallel_size的乘积必须等于 world sizepartition_config.use_default_specs是否叠加 AutoTP 内置默认规格设为False表示仅使用我们自定义的layer_specs设为True则内置规则与自定义规则合并partition_config.layer_specs自定义切分规则列表每条规则由一个或多个patterns正则表达式匹配参数名并指定切分方式patterns参数名匹配正则例如.*(self_attention\|attention)\.query_key_value\.weight$匹配query_key_value.weight结尾的融合投影权重partition_type切分类型column或row决定参数沿哪个维度切分详见下文shape显式提供参数形状仅在 AutoTP 无法从参数本身推断出内部结构时必需例如融合的 Q/K/V 权重partition_dim切分维度在显式shape的前提下指定沿第几个维度切开融合张量column 与 row 切分的直觉column列切分把权重矩阵沿输出维度out_features切分适用于「前向投影」类算子例如 attention 的 QKV 投影、MLP 的dense_h_to_4h升维投影。每个 TP rank 持有一列权重产出的是部分和通常需要后续的 all-reduce或通过分区列式实现来聚合。row行切分把权重矩阵沿输入维度in_features切分适用于「聚合」类算子例如 attention 的dense输出投影、MLP 的dense_4h_to_h降维投影。每个 TP rank 持有完整输出通道的一部分部分和可在本地直接累加。示例中的四条规则恰好构成 GPT-NeoX 一个 Transformer 层的完整切分方案query_key_value.weight融合的 Q/K/V 投影column切分需要显式shapeattention 的dense.weight输出投影row切分MLP 的dense_h_to_4h.weight升维column切分MLP 的dense_4h_to_h.weight降维row切分。为什么 fused QKV 需要显式shapePythia 6.9B 使用 fusedquery_key_value投影即单个权重同时承载 Q、K、V 三份投影。AutoTP 无法仅凭权重张量本身判断「哪一段是 Q、哪一段是 K/V」因此需要用户提供shape说明内部布局shape: ((q_size, kv_size, kv_size), -1), partition_dim: 0,其中q_size num_heads * head_dimkv_size kv_heads * head_dim。示例脚本从模型配置中动态推导这些数值见 autotp_custom_patterns.pynum_heads model.config.num_attention_heads kv_heads getattr(model.config, num_kv_heads, None) if kv_heads is None: kv_heads getattr(model.config, num_key_value_heads, None) if kv_heads is None: kv_heads num_heads head_dim getattr(model.config, head_dim, None) if head_dim is None: head_dim model.config.hidden_size // num_heads q_size num_heads * head_dim kv_size kv_heads * head_dim这组推导是可配置驱动的shape只是 AutoTP 切分融合张量的依据实际数值完全来自模型config因此换用其他 GPT-NeoX 系模型时脚本无需改动。MQA / GQA 模型的自动回退脚本还处理了 Multi-Query AttentionMQA场景。当检测到kv_heads % tp_size ! 0KV 头无法被 TP 均分时会打印提示并把 QKV 切分方式自动切换为 rowuses_mqa bool(getattr(model.config, multi_query, False)) if kv_heads % args.tp_size ! 0: uses_mqa True qkv_spec { patterns: [.*(self_attention|attention)\\.query_key_value\\.weight$], partition_type: row if uses_mqa else column, shape: ((q_size, kv_size, kv_size), -1), partition_dim: 0, }这是use_default_specs机制之外、脚本层面提供的第二层防御当融合 QKV 中 K/V 头数无法按 TP size 整除时KV 部分不可切分整体退化为 row 切分来保证正确性。完整训练脚本TP DP 分组与数据采样手工构建 TP / DP 进程组示例没有依赖外部 MPUmodel parallel unit库而是在脚本内通过dist.new_group手工构造 TP 与 DP 组见 autotp_custom_patterns.pytp_rank rank % tp_size dp_rank rank // tp_size for dp_idx in range(dp_size): tp_ranks list(range(dp_idx * tp_size, (dp_idx 1) * tp_size)) group dist.new_group(tp_ranks) # TP 组同一 dp 块内的连续 tp_size 个 rank for tp_idx in range(tp_size): dp_ranks [tp_idx dp_idx * tp_size for dp_idx in range(dp_size)] group dist.new_group(dp_ranks) # DP 组跨 dp 块、同 tp_idx 的 rank组的编号逻辑rank 号 dp_rank * tp_size tp_rankTP 组内 rank 相邻DP 组内 rank 跨块取同列。随后用最小化的ModelParallelUnitdataclass仅暴露get_data_parallel_group/get_model_parallel_group等接口传入deepspeed.initializempu ModelParallelUnit(tp_group, dp_group, args.tp_size, args.dp_size, tp_rank, dp_rank) engine, _, _, _ deepspeed.initialize(modelmodel, optimizeroptimizer, configds_config, mpumpu)DP-rank 随机采样器每个 DP rank 看到不同样本与 basic_example 中「TP 组内广播同一份合成 token」的做法不同custom_patterns 使用一个真正的文本数据集并通过DPRandomSampler保证每个 DP rank 看到不同的样本而同一 DP 组内的 TP rank 共享相同的数据顺序class DPRandomSampler(Sampler[int]): def __iter__(self) - Iterable[int]: g torch.Generator() g.manual_seed(self.seed self.epoch) indices torch.randperm(len(self.data_source), generatorg).tolist() return iter(indices[self.dp_rank :: self.dp_size])实现要点以seed epoch作为随机源保证每个 epoch 的采样序列稳定可复现对打乱的索引按dp_rank步进dp_size切片每个 DP rank 拿到互不重叠的一份数据被pad_sequence填充到 batch 内等长labels用IGNORE_INDEX -100填充以屏蔽 padding 位置的 loss见collate_batchautotp_custom_patterns.py。训练循环在数据集耗尽时调用sampler.set_epoch(step)重新洗牌让 10 条文本的迷你数据集也能支撑任意步数的迭代for step in range(args.num_steps): try: batch next(data_iter) except StopIteration: sampler.set_epoch(step) data_iter iter(dataloader) batch next(data_iter) ... outputs engine(input_idsinput_ids, attention_maskattention_mask, labelslabels) engine.backward(outputs.loss) engine.step()命令行参数一览脚本通过 argparse 暴露了完整的可调参数autotp_custom_patterns.py参数默认值说明--model_nameEleutherAI/pythia-6.9b要加载的 Hugging Face 因果 LM 模型--tp_size4张量并行组大小--dp_size2数据并行组大小需满足tp_size * dp_size world_size--zero_stage2ZeRO 优化阶段示例默认 stage 2--batch_size1每 GPU 微批大小--seq_length512序列长度--num_steps20训练步数--learning_rate2e-6AdamW 学习率--precisionbf16可选bf16/fp16/fp32对应启用 DeepSpeed 的bf16或fp16配置--trust_remote_code关闭允许从 Hub 加载带自定义代码的模型ChatGLM 等会自动启用--seed42采样器随机种子如何运行选择满足tp_size * dp_size world_size的卡数。8 卡、TP4 / DP2 的标准运行方式deepspeed --num_gpus 8 autotp_custom_patterns.py \ --model_name EleutherAI/pythia-6.9b \ --tp_size 4 \ --dp_size 2 \ --seq_length 512 \ --num_steps 20torchrun同样可用脚本通过deepspeed.init_distributed()完成分布式初始化并由启动器注入LOCAL_RANK。每 5 步 rank 0 会打印当前 loss方便快速验证配置是否生效。与仓库中其他 AutoTP 示例的对照basic_example/autotp_example.py最小化 AutoTP ZeRO-2 流程仅需tensor_parallel.autotp_size使用 TP 组内广播的合成 token不涉及任何手工切分规则——适合验证 AutoTP 对 Llama 等受支持架构的自动识别。basic_example/autotp_memory_compare.py对比「配置驱动」路径与「传统路径」的内存占用传统路径需在deepspeed.initialize前手动调用deepspeed.tp_model_init(...)与set_autotp_mode(trainingTrue)custom_patterns 与 basic_example 均采用新的配置驱动方式无需任何初始化辅助调用。hf_integration在 Hugging FaceTrainer中通过configs/ds_config.json的tensor_parallel.autotp_size启用 AutoTPrun.sh提供zero1tp/zero2tp/tp等模式切换。三者覆盖了「受支持模型开箱即用 → 自定义切分模式 → 与 HF Trainer 集成」的完整路径custom_patterns 正是其中唯一需要手工编写partition_config的进阶场景。常见问题与注意点tp_size * dp_size必须等于 world sizebuild_tp_dp_groups会在不满足时直接抛出ValueError这是保证 TP/DP 组编号正确性的前提。use_default_specs的选择如果模型大部分结构符合内置规格、只有个别融合参数特殊可保留True让内置规则兜底Pythia 这类与内置规格差异较大的模型建议设为False完全接管。正则匹配要精确锚定patterns中的表达式以weight结尾并锚定行尾$避免误匹配 bias、embedding 等其他参数。MQA 场景的注意当 K/V 头数不能被 TP size 整除时脚本会自动退化为 row 切分并打印提示属预期行为。精度配置--precision bf16对应 DeepSpeedbf16.enabledTruefp16同理fp32时不额外写精度字段。通过本示例的配置与脚本你可以把任意「参数命名/结构非标准」的 Hugging Face 模型接入 DeepSpeed AutoTP 训练只需按层列出patterns、选定partition_type必要时用shape显式描述融合张量布局即可获得与标准模型一致的 TP 切分能力。赞分享示例工程【免费下载链接】DeepSpeedExamplesExample models using DeepSpeed项目地址https://gitcode.com/gh_mirrors/de/DeepSpeedExamples点击查看免费下载相关推荐pandoc LaTeX 宏解析实战以 \newcommand 自定义命令为例深入 latex_macros 扩展pandoc LaTeX 宏解析实战以 \newcommand 自定义命令为例深入 latex_macros 扩展 导读 本文以 pandoc 官方命令测试用文档开发工具CLI为 webpack 接入自定义 JavaScript 解析器以 acorn、oxc、meriyah 为例的 ParseFunction 适配实战为 webpack 接入自定义 JavaScript 解析器以 acorn、oxc、meriyah 为例的 ParseFunction 适配实战 本文以仓库中前端构建开发工具蓝鲸PaaS路线图前瞻AI开发、云原生与可观测性的下一步演进蓝鲸PaaS路线图前瞻AI开发、云原生与可观测性的下一步演进 蓝鲸智云 PaaS 平台blueking paas是一个开放式的开发平台帮助开发者快速创建后端云原生微服务前端企业应用开发者门户上一篇DLSS Swapper深度解析游戏画质升级的终极解决方案下一篇Firstmate Secondmate 父通道Parent Channel设计详解与无模型现场验证实录创作声明:本文部分内容由AI辅助生成(AIGC),仅供参考
返回列表