[Bug]: MindSpore PipelineStage 使用 grad_position=-1 原地修改共享 kwargs，导致 PP interleave + recompute 反向错位

Checklist

• 已检索现有 issue，未发现相同问题。
• 已阅读相关文档和实现。
• 已提供最小复现，并保留了完整错误日志；下文仅摘录关键错误与验证证据。

🐛 Describe the bug

问题概述

MindSpore 后端的 PipelineStageBase._grad_position_from_requires_grad() 在“所有 positional Tensor 都需要梯度”时返回 -1。但 forward_and_gradfn() 对 grad_position=-1 的公开语义是：将所有 positional 和 keyword Tensor 都作为输入梯度目标，并原地设置 _requires_grad=True。

PipelineStage 实际只会通过 PP backward-send 路由回传 positional stage input（接收到的 activation）梯度，并没有 kwargs 的 backward-send 路由。因此，这里把 -1 当作“所有 positional Tensor”的快捷表示，会意外修改调用方持有的共享 kwargs Tensor。

在 interleaved pipeline + full recompute 场景中，这个副作用会改变重计算前后输入 Tensor 的梯度状态，进一步造成重计算 placeholder 顺序错位，最终在与根因无关的下游算子中表现为 shape mismatch。

触发条件与前因后果

已复现的场景为 MindFormers PyNative DeepSeek-V4：

• TP=2、EP=2、PP=4；
• interleave num=2，同一 rank 持有两个 virtual stage；
• 1F1B interleaved schedule；
• full recompute；
• input_ids、mask 等模型输入通过 schedule.run(**inputs) 作为 kwargs 传入并在 virtual stage 间复用。

具体过程如下：

1. rank 0 上的 stage 0 首次前向时，模型输入全部位于 kwargs，positional args 为空。旧逻辑得到空的 requires_grad_indices，返回 []，因此 stage 0 不会修改 kwargs。
2. 同一 rank 随后执行 virtual stage 4。此时 positional args 是 P2P 接收到的 activation，且这些 Tensor 均为 _requires_grad=True。旧逻辑因此返回 -1。
3. forward_and_gradfn(..., grad_position=-1, **kwargs) 会递归遍历 kwargs，并原地把共享的 input_ids._requires_grad 从 False 改为 True。
4. stage 0 进入重计算时再次使用同一份 microbatch kwargs。重计算看到的输入元数据已经与首次前向不同，保存/恢复的 Tensor placeholder 列表从该位置开始整体错位。
5. 错位最终在 mHC 路径的 Mul 中暴露为不相关的 shape mismatch，容易被误判为 Hyper-Connections 或 MindSpore 算子问题。

关闭 interleave 后能够运行，是因为 rank 0 不再同时持有后续 non-first virtual stage；没有 stage 4 在 stage 0 重计算前原地修改同一份 kwargs，触发链被切断。

根因代码

旧的 stage 侧选择逻辑：

if not requires_grad_indices:
    return []
if len(requires_grad_indices) == len(tensor_indices):
    return -1
return tuple(requires_grad_indices)

而 forward_and_gradfn() 的 -1 分支会处理 args 和 kwargs：

if grad_position == -1:
    flatten_inputs, _ = tree_flatten(inputs, tensors_only_leaf=True)
    for inp in flatten_inputs:
        inp._requires_grad = True
    flatten_kwargs, _ = tree_flatten(kwargs, tensors_only_leaf=True)
    for kwarg in flatten_kwargs:
        kwarg._requires_grad = True

因此根因不是 forward_and_gradfn(-1) 的实现错误，而是 PipelineStage 使用了语义范围更大的 -1 哨兵。

最小复现

下面的代码展示 stage 选择逻辑与 forward_and_gradfn 组合后的副作用：

import os

os.environ["HYPER_PARALLEL_PLATFORM"] = "mindspore"

import numpy as np
import mindspore as ms
from mindspore import Tensor, nn, ops

from hyper_parallel.platform.mindspore.pipeline_parallel.backward import forward_and_gradfn
from hyper_parallel.platform.mindspore.pipeline_parallel.stage import PipelineStageBase


class Net(nn.Cell):
    def construct(self, activation, input_ids=None):
        return activation * ops.cast(input_ids, activation.dtype)


ms.set_context(mode=ms.PYNATIVE_MODE)
activation = Tensor(np.array([2.0], np.float32))
activation._requires_grad = True
input_ids = Tensor(np.array([3], np.int32))

grad_position = PipelineStageBase._grad_position_from_requires_grad([activation])
assert grad_position == -1  # 修复前

forward_and_gradfn(Net(), activation, grad_position=grad_position, input_ids=input_ids)
assert not input_ids._requires_grad  # 修复前失败：实际为 True

这里最后一个断言表达的是 PipelineStage 所需的隔离性，不是要改变 forward_and_gradfn(-1) 的公开语义。

定位证据

在重计算输入和 placeholder pack/unpack 处增加 Python 打点后，能够稳定观察到：

input_ids expected: Tensor(shape=[1, 1024], dtype=Int32, requires_grad=False)
input_ids actual:   Tensor(shape=[1, 1024], dtype=Int32, requires_grad=True)

placeholder index 67 expected: Tensor(shape=[1024, 32], dtype=Float32)
placeholder index 67 actual:   Tensor(shape=[1024], dtype=Int32)

随后报错：

ValueError: For 'Mul', input1.shape and input2.shape need to broadcast,
but got input1.shape = [2048, 512] and input2.shape = [1024, 1, 4, 512]

将 stage 的返回值改为显式 positional tuple 后，input_ids 不再发生状态变化，placeholder mismatch 和后续 Mul shape mismatch 同时消失。该因果关系已通过最小 UT、4 卡 ST 和原始 8 卡训练场景交叉验证，并非仅根据最终异常推测。

Expected behavior

• PipelineStage 只把需要 PP backward-send 的 positional stage inputs 作为输入梯度目标；
• 不应原地修改调用方传入的 kwargs Tensor 梯度状态；
• interleaved pipeline 下，多个 virtual stage 复用 kwargs 不应改变其他 stage 的重计算输入签名；
• forward_and_gradfn(-1) 的公开 args+kwargs 求梯度语义保持不变；
• kwargs 参与的计算仍应正常产生模型参数梯度。

Additional context

可选修复方案

1. 修改 MindSpore 或 forward_and_gradfn(-1) 的公共语义

   • 例如只修改浮点 kwargs、只收集原本已经需要梯度的 kwargs，或在前向后恢复 _requires_grad。
   • 缺点：会破坏 forward_and_gradfn(-1) 已有的 args+kwargs 求梯度语义和对应测试；同时把 PipelineStage 的策略泄漏到通用 autograd API，影响范围过大。

2. 在 MindFormers 侧规避

   • 例如关闭 PP interleave、每个 virtual stage 克隆 kwargs，或把 input_ids 从 kwargs 改为其他传递方式。
   • 缺点：只能绕开当前触发链，其他 Hyper 上层仍可能遇到相同问题；克隆还会引入额外生命周期、显存和维护成本。

3. 为 forward_and_gradfn 增加独立的 kwargs 梯度选择 API

   • 例如增加 grad_keyword_names，允许 positional 和 keyword 梯度目标分别指定。
   • 优点：可完整支持真正需要 kwargs 输入梯度的场景。
   • 缺点：涉及公共 API、梯度返回结构和兼容性设计；当前 PipelineStage 也没有 kwargs 的 PP backward-send 路由，作为本问题的修复过重。

4. Hyper PipelineStage 始终传入显式 positional indices（当前采用）

   • _grad_position_from_requires_grad() 统一返回 tuple(requires_grad_indices)；没有目标时为 ()，全部 positional Tensor 都需要梯度时为 (0, 1, ...)，不再返回 -1。
   • 优点：修改范围仅限 MindSpore PipelineStage；与现有 PP backward-send 路由一致；保留 forward_and_gradfn(-1) 公共语义；权重梯度目标不变。
   • 兼容性：如果未来 PipelineSchedule 需要跨 stage 回传 kwargs 输入梯度，应采用方案 3 并补齐明确的通信路由，而不是依赖旧逻辑中仅在“全部 positional Tensor 都需要梯度”时偶然出现的行为。

当前核心修改为：

requires_grad_indices = [
    i for i in tensor_indices
    if composite_args[i]._requires_grad
]
return tuple(requires_grad_indices)

并增加以下回归覆盖：

• CPU UT：non-first stage 的 positional activation dx 正确，Int32 input_ids 保持 _requires_grad=False；
• 保留并通过现有 forward_and_gradfn(-1) kwargs 梯度测试，确认公共 API 未回归；
• 4 卡 PP2×TP2×VPP2 ST：所有 virtual stage 复用共享 Int32 kwarg，断言其梯度状态不被污染，同时保留输出与权重梯度精度对齐；
• 8 卡 DeepSeek-V4：TP2、EP2、PP4、interleave=2、full recompute、seq_length=1024，完成 1 step：

loss: 11.783281
load_balancing_loss: 1.128077
mtp_1_loss: 3.541678
grad_norm: 7.811133

Environment info

• hyper-parallel base commit: a7a9b2b4fad6dfea6d73a255795d7ced5b261c00
• MindSpore: 2.10.0
• MindFormers ms/master: 650abf288
• Execution mode: PyNative
• Hardware: 8 × Ascend 910B2
• Architecture: aarch64

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

Start with PipelineStageBase._grad_position_from_requires_grad() in hyper_parallel.platform.mindspore.pipeline_parallel.stage and the grad_position handling in pipeline_parallel.backward.forward_and_gradfn(). Run the CPU regression covering positional activation and shared Int32 kwargs, then the existing kwargs-gradient test. Done means explicit positional indices avoid kwargs mutation while the public -1 behavior and distributed regressions remain covered.