背景

基于算子级并行的编程范式的整体原则是：用户仅进行单卡的编程，所有的算子切分逻辑由框架进行推导，这是一个非常远大的目标。而在实际的大模型训推过程中，存在一些场景使得算子级并行不太好处理，比如：1）当前典型的MoE训练场景，大家都倾向于选择Dropless的方案，即每个专家接收到的token数目不再进行任何的drop/padding，这要求通信的数据量是和计算结果相关，这其实就违背了算子级并行推导的基础逻辑：张量的分布式逻辑与计算结果无关；2）用户需对网络中的某个cell自定义正反向，因为并行框架只对正向算子进行并行切分，反向依赖于自动微分，如果用户自定义正反向，用户无法感知并行框架对正向插入了哪些额外的操作，因此可能导致自定义的反向无法工作；3）用户自己希望对网络中的某段逻辑进行手动并行优化，需要有介入的入口；为了解决这个问题，我们引入自定义并行的机制，开放接口给用户进行自定义的分布式逻辑处理。

设计方案

在自定义并行中，用户需提供一个自定义函数，并描述其输入/输出张量的切分形态。自定义函数是站在local tensor视角进行编程，内部可以显式调用集合通信接口，而其输入/输出通过切分形态的描述与前后的算子级并行进行衔接。

如图所示，框架将输入的global tensor转换成local tensor后进入自定义函数，自定义函数体内以local tensor视角进行编程并可以显式调用集合通信，最后框架再将自定义函数输出的local tensor转换成global tensor。

对外API

def custom_shard(
func: Callable,
out_layouts: Tuple[Layout, ...],
in_layouts: Optional[Tuple[Optional[Layout], ...]] = None,
redistribute_inputs: bool = False,
) -> Callable:

Wraps a function to handle distributed tensor conversions.

Args:
    func (Callable): The function to be wrapped.
    out_layouts (Tuple[Layout, ...]): Layouts for each output tensor.
    in_layouts (Optional[Tuple[Optional[Layout], ...]], optional):
        Layouts for each input argument. None entries indicate non-tensor inputs.
    redistribute_inputs (bool): Whether to redistribute inputs to required layouts.

Returns:
    Callable: Wrapped function that handles distributed tensors.

对redistribute_inputs入参的说明：
默认为True，即当用户对Cell输入指定的layout，与输入tensor本身的layout不一致，会去做一次重排布，但是这样会多一次判断layout的开销；如果用户自己保证配置的layout和输入tensor本身的layout完全一致，可以配置为False，框架不再去校验是否需要插入重排布，而是统一不插入。接口形式参考pytorch的local_map函数。

使用样例：

def mm_allreduce(x, w): 
      out = ops.matmul(x, w) 
      out = ms.parallel.func.all_reduce(out, “mp”)
      return out

class Net(Cell):
    def __init__(self, w_size, in_layouts, out_layouts):
        super().__init__()
        self.p_func = custom_shard(mm_allreduce, out_layouts, in_layouts)
        self.w = Parameter(initializer(w_size), "w")

    def construct(self, x):
        return self.p_func(x, self.w)

layout = Layout((2, 4), (“dp”, “mp”))
x_layout = layout(“dp”, “mp”)
y_layout = layout(“mp”, “None”)
out_layout = layout(“dp”, “None”)
net = Net([64, 64], (x_layout, y_layout), (out_layout,))
…

使用约束

1，custom_shard的第一个入参只支持函数，不支持cell；
2，函数的输出切分策略必须指定，而输入切分策略可不指定。若输入切分策略不指定，则使用输入DTensor的切分策略；
3，在用户自定义的函数体内调用通信算子时，需调用可自动微分的版本，否则反向会有问题；（比如mindspore的mint类通信算子不支持自动微分，而comm_func类通信算子支持自动微分）
4，当前仅支持动态图，尚未支持静态图；

测试设计

1，普通用例，不涉及自定义反向
1）构造一个单卡用例；
2）构造一个并行用例，将其中的某个部分采用custom_shard进行自定义并行，自定义并行逻辑能保障与单卡等价；
3）分别跑一个训练step，观察正向loss/反向grad是否能完全对齐；

2，若涉及自定义反向
1）总体步骤与上述普通用例一致；
2）自定义反向参考：

class LinearFunction(torch.autograd.Function):
    @staticmethod
    def forward(x, weight):
        output = torch.matmul(x, weight)
        return output

    @staticmethod
    def setup_context(ctx, inputs, output):
        x, weight = inputs
        ctx.save_for_backward(x, weight)

    @staticmethod
    def backward(ctx, grad_output):
        x, weight = ctx.saved_tensors
        grad_input = grad_output.mm(weight.t())
        grad_weight = x.t().mm(grad_output)

        return grad_input, grad_weight


def linear(x, weight):
    return LinearFunction.apply(x, weight)


class SimpleModel(nn.Module):
    """simple model"""
    def __init__(self, dist=False):
        super().__init__()
        self.weight = nn.Parameter(torch.ones(8, 8).npu())
        self.dist = dist
        layout = Layout((1, 8), ("dp", "tp"))
        in_layouts = (layout("None", "None"), layout("None", "tp"))
        out_layouts = (layout("None", "tp"),)
        self.wrap = custom_shard(linear, out_layouts, in_layouts)

    def forward(self, x):
        if self.dist == True:
            x = self.wrap(x, self.weight)
        else:
            x = torch.matmul(x, self.weight)
        x = torch.relu(x)
        x = torch.sum(x)
        return x