[JYMiracle305]

背景

Pipeline Parallel 在 stage 间接收 activation 和 gradient 时，会固定按照 float32 创建接收张量。
当训练启用 bfloat16 等非 FP32 dtype 时，发送端和接收端的张量 dtype 不一致，可能导致 PP 通信失败或程序崩溃。

修改内容

• 在 PP activation 通信前发送张量 dtype 元数据。
• 接收端根据 dtype 元数据创建对应类型的 activation 张量。
• 在 ISend 的 autograd context 中保存输入张量 dtype。
• 反向传播接收 gradient 时，按照原始输入 dtype 创建接收张量。
• 增加 tensor 数量和 dtype 范围检查，避免无效通信数据被使用。

[chen2021673]

这个检查可以提前到 ISend::Forward() 的入口处，在forward 之前就先检查

[JYMiracle305]

把这个检查提到了 ISend::Forward() 入口处

[chen2021673]

这里潜在的默认行为是 IRecv::Backward send 的类型由 grad_outputs 中各个 Tensor 自身的 dtype 决定。但比如LinearBackwardInput 的 grad_outputs 可能为 FP32，但用 BF16 计算，就会变成

前向 activation：BF16
        ↓
下游 Linear backward 产生梯度：FP32
        ↓
IRecv::Backward 原样发送：FP32
        ↓
上游 ISend::Backward 按保存的前向类型接收：BF16

如果要解决这个问题就得在 IRecv::Backward 和 ISend::Backward 配套，都先发送梯度的类型，再发送梯度。但这样其实增加了通信成本。可以先这样过渡，留个 fixme 后续对照 Megatron 支持配置 pipeline_dtype 或类似的功能。

[JYMiracle305]

已加FIXME