RFC: 支持典型网络（Qwen/DS）算子级并行与自定义算子

需求背景 & 价值

大模型训练需要对算子的输入/输出张量进行切分，支持 Qwen、DeepSeek 等典型网络涉及的分布式算子。Qwen/DS 等典型网络中涉及一些专门算子（如 mhc / dsa / gdn / engram），这些算子在单卡上能正常运行，但在分布式训练中需要框架自动完成输入/输出张量的切分与重组，保证切分后计算结果与单卡数学等价。

同时，对于框架尚未内置实现的分布式算子，需要提供用户自定义注册机制，降低接入门槛——用户只需描述算子的 layout 推导逻辑，框架自动完成 DTensor 的提取与包装。

核心价值：

• 覆盖 Qwen/DS 等典型网络的算子级分布式并行需求；
• 提供统一的 DTensor 分发管道，所有算子无需感知分布式，框架自动完成 layout 推导和 DTensor 包装；
• 提供用户自定义分布式算子的注册机制（DFunction + DistributedOp / custom_shard），支持框架内建算子和用户扩展算子共存。

功能描述

1. 算子级并行基础设施

提供三层架构支持算子级分布式并行：

• 声明式模块切分（shard_module）：用户描述模块的输入/输出/参数布局，框架自动注册 forward hooks 完成 DTensor 输入→本地计算→DTensor 输出的全流程包装；
• 函数级轻量装饰器（custom_shard）：将任意原生函数包装为 DTensor 感知函数，适用于快速原型和简单自定义逻辑；
• 自定义自动求导函数（DFunction）：允许用户以 local tensor 视角编写正反向计算，框架自动处理 DTensor 提取与包装，适用于需要自定义反向传播逻辑的场景。

2. 分布式算子实现体系

基于 DistributedOp 抽象基类，为框架内建算子和用户自定义算子提供统一的开发范式（三阶段模式）：

• preprocess(args, kwargs)：参数规范化 + to_local + 构建 cache_values；
• infer_layout(cache_values)：基于输入布局推导输出布局，所有校验集中在此；
• get_expand_impl(func, infer_result, cache_values)：返回可选的扩展实现（闭包），用于注入额外分布式逻辑（如 bias scaling、partial 归约）。

3. 双通道算子注册

• YAML 驱动注册：生产环境推荐，通过 core/shard/ops/yaml/*.yaml 配置文件将算子名称映射到 DistributedOp 子类，OpDispatcher 启动时自动加载并实例化；
• 编程式注册：DistributedOp.__init__ 自动调用 register_distributed_op()，无需任何配置文件，适用于自定义算子快速接入；
• 额外功能：对于框架内建算子，若 HyperParallel 没有实现其分布式版本，支持用户通过 DFunction + DistributedOp 或 custom_shard 自行注册分布式算子；也支持用户注册完全自定义的算子。

设计方案

1. 整体架构

┌─────────────────────────────────────────────────────────┐
│  用户入口                                                │
│  shard_module()  │  custom_shard()  │  DFunction.apply()  │
├─────────────────────────────────────────────────────────┤
│  分发管道 (OpDispatcher + DTensor)                        │
│  DTensor.__torch_function__ / __fallback__ → dispatch()   │
│  YAML加载 → DistributedOp实例化 → LayoutCache缓存         │
├─────────────────────────────────────────────────────────┤
│  算子实现 (DistributedOp 子类)                             │
│  preprocess() → infer_layout() → get_expand_impl()        │
│  覆盖 MatMul、ElementWise、Reduce、Attention 等 100+ 算子  │
└─────────────────────────────────────────────────────────┘

2. DTensor 分发流程

框架通过 DTensor 机制实现算子级并行的透明分发：

用户调用 op(x, y)，其中 x/y 是 DTensor
  │
  ▼
DTensor.__torch_function__ (Torch) / __fallback__ (MindSpore) 拦截
  │
  ▼
OpDispatcher.dispatch(op_call, args, kwargs)
  │
  ├─ 旁路算子（白名单）→ 直接调用本地实现
  ├─ 随机算子          → _dispatch_random_op（处理分布式 RNG）
  └─ 布局推导算子       → 查找 YAML/注册表 → infer_layout → 本地计算 → 包装 DTensor

分发结果通过 LayoutCacheManager 按 op_name + LayoutCacheKey 缓存，避免重复的 layout 推导。

3. DistributedOp 三阶段模式

每个分布式算子的 DistributedOp 子类推荐按三阶段实现（参考 Issue #94）：

阶段 1 — preprocess：参数规范化 + to_local + 构建 cache_values。仅做参数处理，不含校验逻辑。

def preprocess(self, args, kwargs):
    args, _ = _normalize_xxx_args(*args, **kwargs)
    local_args = (args[0].to_local(), args[1].to_local())
    local_kwargs = {}
    cache_values = [args[0].layout, args[1].layout]
    return local_args, local_kwargs, cache_values

阶段 2 — infer_layout：基于 cache_values 推导输出 layout。所有校验（类型、范围、分片约束）集中在此完成，必须手动调用 _check_partial_inputs()。

def infer_layout(self, cache_values):
    layout = cache_values[0]
    self._check_partial_inputs([layout])
    # ... 校验 & 推导 ...
    return ((output_layout,), None)

阶段 3 — get_expand_impl：返回扩展实现闭包或 None。None 表示使用原始函数，不需要扩展。

4. YAML 注册

# core/shard/ops/yaml/matmul_ops.yaml
MatMul:
  dist_op_name: _matmul_dist_op
  distributed_op_class: MatMulDistributedOp
  distributed_op_file: parallel_matmul

OpDispatcher 启动时自动加载所有 YAML 文件并实例化对应 DistributedOp 子类，实例化过程中触发 register_distributed_op() 完成注册。

5. DFunction — 用户自定义分布式正反向

DFunction 允许用户以 local tensor 视角编写自定义分布式 autograd 函数（参考 Issue #91)：

from hyper_parallel import DFunction
from hyper_parallel.core.shard.ops.parallel_ops import DistributedOp

# 步骤 1：注册 DistributedOp（描述 layout 推导逻辑）
class MyAddDistOp(DistributedOp):
    def __init__(self):
        super().__init__("MyAdd")
    def infer_layout(self, layouts, extra_args=None):
        return layouts[0]  # 逐元素：输出 layout = 输入 layout

MyAddDistOp()  # 实例化即注册

# 步骤 2：实现 DFunction 子类（操作 local tensor）
class MyAdd(DFunction):
    _op_name = "MyAdd"

    @staticmethod
    def forward(ctx, x, y):
        return x + y

    @staticmethod
    def backward(ctx, grad):
        return grad, grad

# 自动路由：
result = MyAdd.apply(x, y)               # 纯 Tensor → 单设备路径
result = MyAdd.apply(dtensor_x, dtensor_y)  # DTensor → 分布式路径

分发机制：

• 无 DTensor 输入 → super().apply() → 平台 autograd 单设备路径；
• 有 DTensor 输入 → _OP_DISPATCHER.dispatch() → 提取 local tensors → 调用 infer_layout() → 调用本地函数 → 包装为 DTensor。

6. custom_shard — 函数级轻量包装

对于不需要自定义反向的简单场景，提供更轻量的接入方式（参考 Issue #4)：

from hyper_parallel import custom_shard

@custom_shard(
    out_layouts=(out_layout,),
    in_layouts=(x_layout, y_layout),
    redistribute_inputs=True,
)
def my_distributed_func(x, y):
    out = ops.matmul(x, y)
    out = all_reduce(out, "mp")
    return out

框架自动完成：DTensor → redistribute（如需要）→ to_local → 调用函数 → from_local 包装为 DTensor。

对外 API

shard_module

from hyper_parallel import shard_module, Layout

layout = Layout((dp, tp), ("dp", "tp"))

sharding_plan = {
    "forward": {
        "input": (layout("dp", "tp"),),
        "output": (layout("dp", "None"),),
    },
    "parameter": {"weight": layout("None", "tp")},
}

model = shard_module(model, sharding_plan)

custom_shard

from hyper_parallel import custom_shard

@custom_shard(out_layouts=(Layout(...),), in_layouts=(Layout(...),))
def my_func(x, y):
    return some_operation(x, y)

DFunction

from hyper_parallel import DFunction

class MyCustomFunc(DFunction):
    _op_name = "MyCustomOp"

    @staticmethod
    def forward(ctx, x, y):
        return custom_compute(x, y)

    @staticmethod
    def backward(ctx, grad):
        return compute_grad(grad)

编程式注册

from hyper_parallel.core.shard.ops.parallel_ops_register import register_distributed_op

class MyOpDistributedOp(DistributedOp):
    def infer_layout(self, layouts, extra_args=None):
        return (layouts[0],)

register_distributed_op("MyOp", MyOpDistributedOp("MyOp"))

YAML 注册

# custom_ops.yaml（放入 core/shard/ops/yaml/ 目录）
MyOp:
  dist_op_name: _my_op_dist
  distributed_op_class: MyOpDistributedOp
  distributed_op_file: parallel_my_op

使用约束

1. 算子各输入的切分形态需匹配：如 add 算子，在不涉及广播的情况下，两个输入的 DTensor 状态需一致（即相同 mesh_shape 和 placements）；

2. 训练场景需结合 fully_shard 特性：反向的权重梯度通信都由 fully_shard 处理，若未启用 fully_shard 可能出现精度问题；建议在 shard_module 之后再调用 fully_shard；

3. Partial 状态消除：切分中若涉及 Partial 状态（如 MatMul contract dim 被分片后的 partial sum），需通过 cell 的配置触发重排来消除，或手动调用 result.reduce_partial()；

4. DFunction 约束：

   • 子类必须设置 _op_name，且与注册的 DistributedOp 的 op_name 一致；
   • forward 和 backward 内部必须操作 local tensor，不得递归调用 DFunction.apply；
   • 当前仅支持动态图，暂不支持静态图；

5. custom_shard 约束：

   • 第一个入参只支持函数，不支持 Cell；
   • 输出切分策略必须指定，输入切分策略可不指定（不指定时使用输入 DTensor 的切分策略）；
   • 函数体内调用通信算子时，需调用可自动微分的版本；

6. 算子分发兼容性：所有内置算子的 DistributedOp 实现需保证旧分发流（_with_layout_infer）和新分发流（_dispatch_new）的行为一致，逐步迁移至新模式。

测试设计

单元测试

用例 ID	描述	期望
OP-UT-01	基础算子（MatMul / Add / Reduce）单卡 vs 多卡正反向对齐	分布式计算结果与单卡数学等价
OP-UT-02	不同 placement 组合（Shard(0) + Shard(1)）的广播行为	框架正确推断 layout
OP-UT-03	不支持 placement 时抛出 ValueError	错误信息清晰可读
OP-UT-04	preprocess 路由正确性：新旧分发流等价	_dispatch_new 与 _with_layout_infer 输出一致
OP-UT-05	DFunction 单卡调用（纯 Tensor）	正向输出正确
OP-UT-06	DFunction 多卡调用（DTensor 输入）	输出为 DTensor，值与单卡对齐
OP-UT-07	含 ctx.save_for_backward 的 DFunction 正反向	梯度与单卡参考对齐
OP-UT-08	custom_shard 不指定输入 layout 时行为	使用输入 DTensor 的 layout

自定义算子注册测试

用例 ID	描述	期望
REG-01	YAML 注册后 get_distributed_op("Xxx") 能获取	注册成功
REG-02	编程式注册后 DTensor 输入自动路由到分布式路径	dispatch 正确路由
REG-03	用户自定义算子（DFunction + DistributedOp）端到端训练一步	loss/grad 与单卡对齐

回归测试

• 现有 51 个 UT（tests/ut/core/shard/ops/）无回归；
• 现有分布式 ST 测试全部通过。

规格 & 约束

• 规格：Q2 典型网络（Qwen/DS）所涉及的分布式算子全部覆盖；
• 额外功能：
  1. 支持用户自定义算子（通过 DFunction + DistributedOp 或 custom_shard）；
  2. 对于框架内建算子，若 HyperParallel 没有实现其分布式版本，支持用户自行注册分布式算子；
• 性能：NA（本期不设定性能指标）；
• 约束：
  1. 算子各输入的切分形态需匹配；
  2. 训练场景需结合 fully_shard 特性一起使用，否则可能出现精度问题；
  3. 切分中若涉及 Partial 状态，需通过 cell 的配置触发重排来消除，或手动消除。

参考

• DFunction 自定义算子：Issue #91
• 分布式算子三阶段模式：Issue #94