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

• 示例工程

本指南基于 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 Parallel，TP）的切分工作自动化：它识别模型参数中典型的命名与形状模式，自动施加正确的切分方式。对于 Llama 这类受支持的架构，只需要在 DeepSpeed 配置中开启tensor_parallel.autotp_size，无需任何手工切分规则，参见 basic_example。

但当模型使用了非标准参数定义——例如 Pythia（GPT-NeoX 架构）把 Q、K、V 三个投影融合成一个query_key_value权重，此时 AutoTP 无法确定融合张量内部如何切分，就需要通过partition_config显式声明切分模式。这正是 custom_patterns 示例存在的意义：它把「按层指定切分规则」的能力开放给用户，同时用一个 10 条文本的内存数据集 + DP-rank 随机采样器演示了完整的 TP + DP 训练流程。

核心配置：partition_config完整解析

示例中的完整 DeepSpeed 配置如下（节选自 autotp_custom_patterns.py）：

ds_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, }

各字段含义与要点：

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 层的完整切分方案：

1. query_key_value.weight：融合的 Q/K/V 投影，column切分，需要显式shape；
2. attention 的dense.weight（输出投影）：row切分；
3. MLP 的dense_h_to_4h.weight（升维）：column切分；
4. MLP 的dense_4h_to_h.weight（降维）：row切分。

为什么 fused QKV 需要显式shape

Pythia 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_dim，kv_size = kv_heads * head_dim。示例脚本从模型配置中动态推导这些数值（见 autotp_custom_patterns.py）：

num_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 Attention（MQA）场景。当检测到kv_heads % tp_size != 0（KV 头无法被 TP 均分）时，会打印提示并把 QKV 切分方式自动切换为 row：

uses_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 进程组

示例没有依赖外部 MPU（model parallel unit）库，而是在脚本内通过dist.new_group手工构造 TP 与 DP 组（见 autotp_custom_patterns.py）：

tp_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_rank，TP 组内 rank 相邻，DP 组内 rank 跨块取同列。随后用最小化的ModelParallelUnit（dataclass，仅暴露get_data_parallel_group/get_model_parallel_group等接口）传入deepspeed.initialize：

mpu = ModelParallelUnit(tp_group, dp_group, args.tp_size, args.dp_size, tp_rank, dp_rank) engine, _, _, _ = deepspeed.initialize(model=model, optimizer=optimizer, config=ds_config, mpu=mpu)

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), generator=g).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_batch，autotp_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_ids=input_ids, attention_mask=attention_mask, labels=labels) engine.backward(outputs.loss) engine.step()

命令行参数一览

脚本通过 argparse 暴露了完整的可调参数（autotp_custom_patterns.py）：

如何运行

选择满足tp_size * dp_size = world_size的卡数。8 卡、TP=4 / DP=2 的标准运行方式：

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 20

torchrun同样可用，脚本通过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(training=True)；custom_patterns 与 basic_example 均采用新的配置驱动方式，无需任何初始化辅助调用。
• hf_integration：在 Hugging FaceTrainer中通过configs/ds_config.json的tensor_parallel.autotp_size启用 AutoTP，run.sh提供zero1tp/zero2tp/tp等模式切换。

三者覆盖了「受支持模型开箱即用 → 自定义切分模式 → 与 HF Trainer 集成」的完整路径，custom_patterns 正是其中唯一需要手工编写partition_config的进阶场景。

常见问题与注意点

1. tp_size * dp_size必须等于 world size：build_tp_dp_groups会在不满足时直接抛出ValueError，这是保证 TP/DP 组编号正确性的前提。
2. use_default_specs的选择：如果模型大部分结构符合内置规格、只有个别融合参数特殊，可保留True让内置规则兜底；Pythia 这类与内置规格差异较大的模型，建议设为False完全接管。
3. 正则匹配要精确锚定：patterns中的表达式以weight结尾并锚定行尾（$），避免误匹配 bias、embedding 等其他参数。
4. MQA 场景的注意：当 K/V 头数不能被 TP size 整除时，脚本会自动退化为 row 切分并打印提示，属预期行为。
5. 精度配置：--precision bf16对应 DeepSpeedbf16.enabled=True，fp16同理；fp32时不额外写精度字段。

通过本示例的配置与脚本，你可以把任意「参数命名/结构非标准」的 Hugging Face 模型接入 DeepSpeed AutoTP 训练：只需按层列出patterns、选定partition_type，必要时用shape显式描述融合张量布局，即可获得与标准模型一致的 TP 切分能力。

• 示例工程