[Feature]: hyper-parallel PyTorch 侧接入 MC2 融合张量并行 Linear Style

🚀 功能描述

0. 基本信息

项目	内容
特性名称	HyperParallel PyTorch/Ascend MC2 融合张量并行 Linear Style
开发分支	feat/tp-mc2-fused-parallel
适用后端	PyTorch + Ascend（torch_npu MC2 kernel）
已验证设备	Ascend NPU（本地 UT + 2 卡 MC2 ST）
对应实现	hyper-parallel!1096（已合入 master）
对齐上游	mindformers!8397
当前阶段	Phase 1：TP + Sequence Parallel 下 Column/Row Linear 的 MC2 融合；默认关闭

---

1. 需求背景

HyperParallel 在 PyTorch / Ascend 张量并行路径上，Column / Row Parallel Linear 的集合通信与 nn.Linear 矩阵乘仍是分离调度（DTensor redistribute / Partial→AllReduce-RS）。MindFormers PyNative 侧已具备 MC2（MatMul + Communication）融合能力，但 HyperParallel Torch 侧缺少对齐实现：

能力	MindFormers PyNative	HyperParallel Torch（合入前）	用户影响
Column SP：AllGather + MatMul 融合	MC2ColwiseParallel + all_gather_matmul	通信与 nn.Linear 分离	TP+SP MLP 前向无法吃到 Ascend MC2 通算流水
Row SP：MatMul + ReduceScatter 融合	MC2RowwiseParallel + matmul_reduce_scatter	Partial→AllReduce/RS 分离路径	同上
可训练自定义反向	✅	❌	无法在训练图中直接调用 torch_npu MC2 前向核

在 TP + Sequence Parallel 场景下无法吃到 Ascend MC2 通算流水，独立通信算子调度开销偏高。

---

2. 本期目标与非目标

2.1 本期目标

1. 新增可训练 MC2Linear：前向 AG+MM / MM+RS；Column 反向默认融合 MRS（满足 n_local≥256）。
2. 新增 MC2ColwiseParallel / MC2RowwiseParallel：要求序列维切分 I/O，apply 时替换为 MC2Linear。
3. 公开导出上述 API；examples/torch/llama3/parallelize.py 增加 enable_mc2（默认 False，且需 SP）。
4. 补齐 UT + Ascend ST，覆盖 Colwise/Rowwise 前向+反向精度、MLP 链、fp16、Shard(1)。
5. 优化 tensor_parallel ST 调度与分布式 launcher：父进程不提前加载 torch/mindspore，缩短墙钟。

2.2 本期非目标

• MindSpore 后端同款 MC2 Style。
• MoE / Attention 路径的 MC2 融合。
• 正式 E2E 吞吐/加速比报告与默认开启 MC2。
• 放开 Ascend MC2 kernel 之外的 dtype / world_size / 收缩维约束。
• 无 Sequence Parallel 场景下的 MC2 启用。

---

3. 改动范围

3.1 MC2 产品能力

• MC2Linear + 融合 autograd（mc2.py）
• MC2ColwiseParallel / MC2RowwiseParallel（mc2_style.py）
• 公开导出；示例 enable_mc2 开关

3.2 测试与基础设施

• tensor_parallel ST：前向+反向合并同一次 torchrun；去重 4 卡 Linear launcher；端口错开；MC2 ST shape 保证 n_local≥256
• 新增 tests/common/distributed_launcher.py
• 新增 UT：tests/ut/core/tensor_parallel/test_mc2_style.py
• 新增 ST：tests/torch/tensor_parallel/test_mc2_distributed.py

---

4. 对外接口

# 融合 Linear（可训练）
from hyper_parallel.core.tensor_parallel.mc2 import MC2Linear

# Parallel Style（apply 时替换为 MC2Linear）
from hyper_parallel.core.tensor_parallel import (
    MC2ColwiseParallel,
    MC2RowwiseParallel,
)

# 示例开关（需 Sequence Parallel）
# examples/torch/llama3/parallelize.py: enable_mc2=True

是否影响现有 API

• 默认行为不变：现有 ColwiseParallel / RowwiseParallel 保持；MC2 为新增 Style，需显式选用或打开 enable_mc2
• 新增公开导出：MC2Linear、MC2ColwiseParallel、MC2RowwiseParallel
• 测试基础设施：distributed_launcher 与 ST 调度调整不影响产品 API

---

5. 当前支持矩阵 / 关键约束

能力	状态	限制
PyTorch + Ascend MC2 Column/Row Linear	支持	需 Sequence Parallel；enable_mc2 默认关
dtype	支持	仅 float16 / bfloat16
收缩维 k	支持	k ∈ [256, 65535)；Column 反向时为 n_local
world_size	支持	{2, 4, 8}
Column 反向融合 MRS	支持	需 n_local≥256，否则不可走融合反向
MindSpore 后端 MC2 Style	不支持	本期非目标
MoE / Attention MC2	不支持	本期非目标
无 SP 的 MC2	不支持	Style 要求序列维切分 I/O
正式 E2E 吞吐报告	未覆盖	留后续

---

6. 风险与限制

6.1 Ascend MC2 kernel 硬约束

MC2 路径直接依赖 CANN / torch_npu 融合核。dtype、收缩维、world_size、Column 反向 n_local≥256 任一不满足时，不能静默回退到“看起来能跑但未融合”的路径，必须在 Style/Linear 入口明确失败或保持默认关闭，避免用户误以为已获得 MC2 收益。

6.2 仅覆盖 Linear + SP，能力面偏窄

本期只覆盖 TP+SP 下的 Column/Row Linear。Attention、MoE、无 SP 拓扑仍走原分离通信路径。若示例或文档表述成“全模型 MC2”，会造成预期偏差；enable_mc2 必须继续默认关闭，并由调用方显式打开。

6.3 数值与调度差异

融合核与“通信 + nn.Linear”分离路径在内核实现、rounding、流水排布上可能存在可接受范围内的数值差。验收应以与参考路径的精度阈值/相对误差为准，而不是 bit-wise 一致。ST 需覆盖前向+反向，避免只验前向。

6.4 测试基础设施变更的连带风险

父进程 launcher 轻量化、ST 合并 torchrun、端口错开等改动会影响 tensor_parallel 相关用例调度。风险是“MC2 本身正确，但并行起测/导入路径回归”。验收需同时覆盖：

• MC2 功能 ST；
• 导入 parallel_case / distributed_launcher / shard-ops suite 时父进程不加载 torch/mindspore。

6.5 与 MindFormers 对齐漂移

HyperParallel 侧实现需持续对齐 mindformers PyNative MC2 语义。上游 kernel ABI、反向融合条件或 Style 约束若变更，本仓库可能出现行为分叉；后续变更应同步对照 mindformers!8397 及后续修复。

---

7. 验证设计与当前结果

7.1 UT 覆盖

tests/ut/core/tensor_parallel/test_mc2_style.py：

• 参数保留与 mode 配置；
• 序列维布局约束（非法布局 fail-closed）；
• apply 后模块替换为 MC2Linear；
• Column 反向走融合 MRS 等路径选择。

7.2 ST 覆盖（Ascend）

tests/torch/tensor_parallel/test_mc2_distributed.py：

Level	内容
level0	Colwise / Rowwise 前向+反向精度（各 2 卡，同一波）
level1	MLP 链前向+反向；fp16 Colwise；Shard(1) Colwise

补充约束：MC2 ST 默认 shape 保证 n_local≥256，满足融合反向条件。

7.3 基础设施验证

• 导入 parallel_case / distributed_launcher / shard-ops suite 时，父进程不加载 torch / mindspore；
• tensor_parallel ST 调度调整后，相关用例可稳定并行启动（端口错开，减少冲突重试）。

7.4 本地复现

export HYPER_PARALLEL_PLATFORM=torch
pytest tests/ut/core/tensor_parallel/test_mc2_style.py -q

# ST（示例 2 卡）
export ASCEND_RT_VISIBLE_DEVICES=4,5
pytest tests/torch/tensor_parallel/test_mc2_distributed.py -v

7.5 当前结果

• 本地 UT / 2 卡 MC2 ST 通过；
• launcher 轻量化抽检通过；
• 正式多机 E2E 吞吐对比报告未覆盖（见非目标）。

---

8. 验收标准

8.1 功能验收

• MC2Linear、MC2ColwiseParallel、MC2RowwiseParallel 可从公开路径导入。
• 在 TP+SP 合法配置下，Column/Row Style apply 后实际替换为 MC2Linear。
• Colwise / Rowwise 前向+反向与参考路径精度达标（按 ST 阈值）。
• MLP 链、fp16 Colwise、Shard(1) Colwise ST 通过。
• enable_mc2=False（默认）时，现有非 MC2 TP 行为不变。

8.2 兼容性验收

• 未启用 MC2 时，ColwiseParallel / RowwiseParallel 路径行为不变。
• 新增 API 为增量导出，不破坏既有 Parallel Style 名称与默认并行策略。
• 测试 launcher 变更不导致非 MC2 tensor_parallel / shard-ops 用例大面积不可启动。

8.3 明确报错 / fail-closed

以下场景必须明确失败或拒绝启用，而不是静默退回“未融合但宣称 MC2”：

• 无 Sequence Parallel 却启用 MC2 Style；
• dtype 非 float16/bfloat16；
• 收缩维 / n_local 不满足 MC2 kernel 或融合反向约束；
• world_size 不在 {2, 4, 8}。

8.4 验收方法

• 参考示例中 enable_mc2 的使用方式，自己构建一个MC2的TP切分方式
• 能力边界（仅 Linear、不含 MoE/Attention/MindSpore）在 issue/PR 描述中可见。

---

9. 现有替代方案

继续使用现有 ColwiseParallel / RowwiseParallel：通信与 MatMul 分离调度，功能可用，但在 Ascend TP+SP 下无法走 MC2 融合核，通信调度开销更高。MindFormers PyNative 侧已有同类能力，本需求为 HyperParallel Torch 路径对齐补齐。

与 DFX 相关性

• 可观测性：无新增强制日志；示例通过 enable_mc2 可选启用
• 可靠性：UT + Ascend ST（前向+反向）
• 性能：降低独立通信算子调度开销；正式吞吐报告留后续
• 兼容性：默认关闭，不改变现有 TP 默认行为

补充信息

• 对齐上游：mindformers!8397
• 实现 PR：hyper-parallel!1096

schema_version: 1
source: gitcode
gitcode_repo: mindspore/hyper-parallel
gitcode_issue: 318
source_url: https://gitcode.com/mindspore/hyper-parallel/issues/318

Start with the implementation areas named in the issue: mc2.py, mc2_style.py, public exports, and examples/torch/llama3/parallelize.py. Run tests/ut/core/tensor_parallel/test_mc2_style.py first, then the Ascend tests/torch/tensor_parallel/test_mc2_distributed.py suite. Done means the supported TP+SP cases, listed precision checks, fail-closed constraints, and default-disabled behavior pass without regressing existing tensor-parallel tests.