背景

在分布式训练场景中，存在若干需要用户自定义正反向计算逻辑的场景，使得现有算子级并行框架难以直接处理：

1. 自定义 autograd 函数与分布式并行框架的割裂：用户通过 torch.autograd.Function 或 MindSpore 等效类自定义正反向时，并行框架对正向插入的 layout 推导、DTensor 包装等额外操作对用户不透明，导致自定义反向无法感知并正确处理分布式张量。
2. 缺少跨平台统一抽象：用户在 PyTorch 和 MindSpore 后端上编写的自定义 autograd 函数无法复用，需为不同后端分别适配。
3. 无法与 DTensor dispatch 系统衔接：现有用户自定义函数无法接入 DistributedOp 的 layout 推导和 DTensor 输入/输出包装流程，导致多卡场景下无法直接使用分布式张量作为输入输出。

为解决上述问题，我们引入 DFunction 机制，允许用户以 local tensor 视角编写自定义分布式 autograd 函数，框架自动完成 DTensor 的提取与包装。

设计方案

DFunction 是 platform.Function 的子类（在 PyTorch 上继承 torch.autograd.Function，在 MindSpore 上继承对应的 _Function）。用户子类实现 forward / backward 静态方法，操作的是 local tensor；当输入包含 DTensor 时，apply() 自动路由到 _OP_DISPATCHER.dispatch() 完成 layout 推导和 DTensor 包装，用户感知不到多卡和单卡的区别。

如图所示：

• 单卡路径：DFunction.apply(local_x, local_y) → super().apply() → 平台 autograd 机制 → forward(ctx, local_x, local_y)
• 多卡路径：DFunction.apply(dtensor_x, dtensor_y) → _OP_DISPATCHER.dispatch() → 提取 local tensors → forward(ctx, local_x, local_y) → DTensor.from_local(output, mesh, placements)

layout 推导逻辑由用户配套实现的 DistributedOp 子类提供，通过 _op_name 字符串与 DFunction 子类关联。

对外 API

class DFunction(platform.Function):
    _op_name: str = None  # 与注册的 DistributedOp 的 op_name 对应

    @staticmethod
    def forward(ctx, *args, **kwargs) -> Tensor: ...

    @staticmethod
    def backward(ctx, *grad_outputs) -> ...: ...

    @classmethod
    def apply(cls, *args, **kwargs) -> Tensor | DTensor: ...

配套的 DistributedOp（已有接口，无需新增）：

class DistributedOp:
    def __init__(self, op_name: str): ...
    def preprocess(self, args: tuple, kwargs: dict) -> None | tuple: ...
    def infer_layout(self, cache_values) -> Layout | tuple: ...
    def get_expand_impl(self, func, infer_result, cache_values) -> None | Callable: ...

使用样例：

from hyper_parallel import init_device_mesh, DFunction
from hyper_parallel.core.dtensor.dtensor import distribute_tensor
from hyper_parallel.core.dtensor.placement_types import Shard, Replicate
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):
        ctx.save_for_backward(x, y)
        return x + y

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

# 单卡调用（plain tensor）
result = MyAdd.apply(x_local, y_local)

# 多卡调用（DTensor，自动走分布式路径）
mesh = init_device_mesh("npu", (2, 4), mesh_dim_names=("dp", "tp"))
x_dist = distribute_tensor(x, mesh, (Shard(0), Replicate()))
y_dist = distribute_tensor(y, mesh, (Shard(0), Replicate()))
result_dist = MyAdd.apply(x_dist, y_dist)  # 返回 DTensor

使用约束

1. DFunction 子类必须设置 _op_name，且与注册的 DistributedOp 实例的 op_name 完全一致；当输入包含 DTensor 时，若未设置 _op_name 则抛出 ValueError；
2. forward 和 backward 内部必须操作 local tensor，不得递归调用 DFunction.apply；
3. 非 Tensor 位置参数在 legacy dispatch 路径（未实现 preprocess）中不会被传入 forward，需改用 kwargs 或实现 preprocess 走新 dispatch 路径；
4. 当 get_expand_impl 返回值为可调用对象且输出具有 partial 状态时，需在使用前调用 result.reduce_partial()；
5. 当前仅支持动态图，暂不支持静态图；

测试设计

1. 普通用例（不涉及自定义反向）

   1. 构造单卡用例，验证 DFunction.apply 正向结果正确；
   2. 构造并行用例（DTensor 输入），验证输出类型为 DTensor 且值与单卡对齐；
   3. 分别跑一个训练 step，观察正向 loss / 反向 grad 是否完全对齐；

2. 涉及自定义反向的用例

   1. 实现包含 ctx.save_for_backward / ctx.saved_tensors 的自定义正反向；
   2. 单卡 vs 多卡反向梯度精度对齐验证；
   3. 参考用例：

      class LinearFunction(DFunction):
          _op_name = "TestLinear"
      
          @staticmethod
          def forward(ctx, x, weight):
              output = platform.matmul(x, weight)
              ctx.save_for_backward(x, weight)
              return output
      
          @staticmethod
          def backward(ctx, grad_output):
              x, weight = ctx.saved_tensors
              grad_input = platform.matmul(grad_output, weight.t())
              grad_weight = platform.matmul(x.t(), grad_output)
              return grad_input, grad_weight
      
      mesh = init_device_mesh("npu", (1, 8), mesh_dim_names=("dp", "tp"))
      in_layouts = (layout("None", "None"), layout("None", "tp"))
      out_layout = (layout("None", "tp"),)
      # 单卡与多卡前向/反向结果完全对齐
      

3. row-parallel 场景（含 get_expand_impl bias scaling）

   1. 构造输入在 TP 维度 Shard 的 row-parallel linear；
   2. 验证 bias 在每个 rank 上被正确缩放（bias / tp_size）；
   3. 反向梯度与单卡参考值对齐；