[基础能力4] loss_parallel / 分布式 CE：当前差异梳理与补齐计划

一、背景与目标

📎 关联： changzherui1/hyper-parallel#6 张量并行 TP · changzherui1/hyper-parallel#1 接口梳理总览 · mindspore/hyper-parallel#270 TP redistribute 异步

本 Issue 目标： 在保持 Hyper 现有 loss_parallel 架构（上下文管理器 + _OP_DISPATCHER 路由 + 双栈 kernel）不变的前提下，根据下文当前差异逐项补齐 loss_parallel / 分布式 Cross-Entropy 能力。

背景：

• TP 训练下 lm_head 常用 ColwiseParallel，logits 在 vocab 维为 Shard(-1)，每个 rank 仅持有 V/tp 切片。直接 F.cross_entropy 语义错误，必须在 TP 维做分布式 softmax（all-reduce max / sumexp / gather）。
• PyTorch 提供 with loss_parallel(): F.cross_entropy(...)，经 DTensor custom handler 拦截 _log_softmax / nll_loss。
• TorchTitan 生产路径不用上下文，而用 components/loss.py 的 _LossParallelCrossEntropy（plain tensor + tp_group）+ 外层 ChunkedLossWrapper（seq 分块降显存）。
• Hyper 已有 loss_parallel 上下文 + DistributedCrossEntropyFunction（Torch/MindSpore 双栈），核心 CE 算法与 Titan _LossParallelCrossEntropy 同类，但在 mesh/labels/reduction 语义、框架集成、ChunkedLoss 上与 PyTorch / Titan 存在下述差异。

分析依据（源码）：

侧	路径
PyTorch	torch/distributed/tensor/parallel/loss.py
HyperParallel	hyper_parallel/core/tensor_parallel/loss_parallel.py、platform/*/loss_parallel_ops.py、_ce_op_registry.py
TorchTitan	torchtitan/components/loss.py（CrossEntropyLoss、_LossParallelCrossEntropy、ChunkedLossWrapper）

图例： ✅ Hyper 已有 · ⚠️ 部分具备 · ❌ Hyper 缺失 · 🔶 Hyper 扩展

---

二、范围与约束

2.1 模块范围

子模块	PyTorch	Hyper	TorchTitan
上下文 API	loss_parallel()	loss_parallel(mesh, strict)	—（不用上下文）
分布式 CE kernel	ATen custom handlers	DistributedCrossEntropyFunction	_LossParallelCrossEntropy
直调 API	—（隐含在 handler 内）	—（缺失）	_LossParallelCrossEntropy.apply(...)
框架 loss 组件	用户自行包上下文	—（缺失）	CrossEntropyLoss + ChunkedLossWrapper
op 路由 / 防护	DTensor dispatcher	_OP_DISPATCHER + _ce_op_registry	N/A

2.2 约束

• 保留 with loss_parallel() 上下文 + _OP_DISPATCHER 拦截。
• 保留 Torch + MindSpore 双栈 platform/*/loss_parallel_ops.py。
• ChunkedLoss 作为独立框架层模块，不并入 loss_parallel.py 内核。
• TP parallelize_module / lm_head 编排见 #6；TP 边界通信重叠见 #270（与本 Issue 正交，不影响 CE 正确性）。

2.3 三端现状对比

                    PyTorch              Hyper（现状）           Titan（生产）
                    ───────              ───────────           ─────────────
触发方式            with loss_parallel   with loss_parallel    cross_entropy_loss() 自动
内核                ATen handlers        DistributedCEFn       _LossParallelCrossEntropy
显存优化            —                    —                     ChunkedLossWrapper
双后端              Torch only           Torch + MS            Torch only

2.4 算子拦截机制（与 PyTorch 差异）

层次	Hyper	PyTorch
DTensor 算子	DTensorBase.__torch_function__ → _OP_DISPATCHER	DTensor.__torch_dispatch__
CE 路由	_ce_op_registry + 上下文 loss_parallel()	DTensor custom handler on _log_softmax / nll_loss
内核	DistributedCrossEntropyFunction（双栈）	ATen handlers

说明： Hyper 不在 DTensor 上实现全套 __torch_dispatch__；CE 语义由 _OP_DISPATCHER 在上下文激活时路由到 DistributedCrossEntropyFunction。与 #270 的 ACT 异步 wait 无关——CE 路径走融合 kernel，不依赖 redistribute 重叠。

---

三、当前差异、影响与补齐方向

3.1 Hyper 已具备（本 Issue 不改动）

能力	说明
TP vocab 分片 CE 算法	Forward：MAX + SUM all-reduce + gather；Backward：融合、0 次 collective
不均匀 vocab 切分	ceil(V/tp) chunk 语义
ignore_index=-100	与 PyTorch / Titan 相同
分解算子防护	禁止上下文内直接调 log_softmax/nll_loss
loss_parallel(mesh, strict)	🔶 Hyper 扩展

3.2 差异总表

#	差异点	PyTorch	Titan	Hyper	会导致的问题	补齐方向
1	顶层 API 导出	tensor.parallel 公开	内置于 components/loss	未进 hyper_parallel.__all__	用户找不到 API；误用未包上下文的 F.cross_entropy	加入 __all__ + 文档
2	多维 mesh (CP+TP)	自动找 TP 维；labels 须 DTensor	主路径 1D TP；local_map 支持多维	要求 1D mesh；mesh_dim=0 写死	CP+TP 训练 CE 报错或算错	_find_tp_mesh_dim；DTensor labels 校验；非 TP 维 Partial 输出
3	直调 Function API	—	_LossParallelCrossEntropy.apply	有 kernel、无公开直调	无法复用 Titan 式编排；必须包上下文	公开 loss_parallel_cross_entropy(...)
4	reduction 语义	none/sum 多维 layout 完整；mean 仅 1D	训练默认 sum + /global_valid_tokens	mean/sum/none 有；loss 非 DTensor	多维 mean 可能不正确；与 Titan token 归一化路径不同	补 placement 规则；CrossEntropyLoss 组件
5	labels 类型	多维须 DTensor	plain 或 DTensor	仅 plain tensor	CP 分片 labels 与 logits 可能不一致	推导并校验 target placements
6	class weight	支持	不支持	支持	—	文档说明三端差异
7	label_smoothing	不支持	不支持	不支持	—	文档声明；可选 P2
8	框架 CrossEntropyLoss	用户包上下文	有	无	Trainer 无法自动走 loss parallel	新增 loss 组件
9	ChunkedLossWrapper	无	生产默认	无	长 seq + 大 vocab 峰值显存 O(B×L×V)	新增 ChunkedLoss
10	Shard logits 未包上下文	行为未定义	N/A	主动报错	—	保持并补充测试
11	nn.CrossEntropyLoss 模块	经 ATen 子算子	走 cross_entropy_loss	依赖 op 名注册	部分调用路径可能未拦截	扩展 _ce_op_registry
12	推理 gather logits	—	—	gather_tensor_parallel_logits	训练/推理 CE 路径分离	P2 文档或统一

3.3 按场景影响

用户场景	Hyper 现状	风险
1D tp_mesh + with loss_parallel()	✅	低
CP+TP 2D mesh，logits/labels 分片	❌/⚠️	高：报错或 silent 错误
使用 Titan 式 CrossEntropyLoss + ChunkedLossWrapper	❌ 无对应组件	高：无组件；显存模型不同
使用 PyTorch 式 with loss_parallel() + 2D mesh	⚠️	中：行为与 PyTorch 不一致
MindSpore 后端 TP CE	🔶	中：需 MS 精度/通信测试
reduction="mean" + 2D mesh	PyTorch 会拒绝；Hyper 可能错	中

---

四、补齐计划（P0 → P1 → P2）

原则： 内核已有 DistributedCrossEntropyFunction，优先补 API + 多维语义 + 框架集成，再补 ChunkedLoss。

4.1 优先级总览

优先级	目标	项数
P0	消除 CP+TP 等阻塞性差异	4
P1	补全 Titan 训练路径相关差异	4
P2	显存路径与文档	3

4.2 P0

#	任务	Hyper 现状	补齐方向	工作量
1	顶层导出 loss_parallel	子包 only	hyper_parallel.__all__	小
2	自动识别 TP mesh 维	mesh_dim=0 写死	_find_tp_mesh_dim(placements, class_dim)	中
3	多维 mesh + DTensor labels	plain labels	推导 target placements；多维 plain labels 报错	中
4	reduction="sum" 非 TP 维 Partial	仅 TP all-reduce	补 DTensor loss 输出或文档化 reduce 契约	中

4.3 P1

#	任务	补齐方向	工作量
5	loss_parallel_cross_entropy()	公开包装 DistributedCrossEntropyFunction	小～中
6	CrossEntropyLoss 组件	检测 Shard(-1)；sum + global_valid_tokens	中
7	扩展 CE op 注册	nn.CrossEntropyLoss、MindSpore 变体	小
8	reduction="mean" 多维约束	2D mesh NotImplementedError（与 PyTorch 相同限制）	小

4.4 P2

#	任务	补齐方向	工作量
9	ChunkedLossWrapper	seq 分块 lm_head + FSDP 编排	大
10	Trainer 接线	LLMTrainer 可选 ChunkedLoss	中
11	文档 + 示例	tp_cp_example、三端差异说明	小

4.5 里程碑

M1  CP+TP CE 可用     P0 #1 #2 #3
M2  框架 loss 路径     P0 #4 + P1 #5 #6
M3  ChunkedLoss 显存   P2 #9 #10
M4  文档与 op 覆盖     P1 #7 #8 + P2 #11

4.6 不在本 Issue 范围

项	原因
label_smoothing	三端均未支持
移除上下文、仅保留 Function API	破坏现有用法
TP parallelize_module	属 #6
推理 gather_tensor_parallel_logits 重写	推理侧单独立项

---

五、测试计划

类型	内容
UT	mesh 维识别、reduction、labels 校验、op 注册
数值对比	与 PyTorch loss_parallel（1D/2D）；与 Titan _LossParallelCrossEntropy
ST	NPU 2-card TP；CP+TP 组合
回归	tests/torch/loss_parallel、tests/mindspore/st/loss_parallel

---

附录：三端 API 对照

API	PyTorch	Hyper	Titan
上下文	loss_parallel()	loss_parallel(mesh, strict)	—
分布式 CE	ATen handlers	distributed_cross_entropy	_LossParallelCrossEntropy.apply
框架入口	—	（缺失）	CrossEntropyLoss
显存优化	—	（缺失）	ChunkedLossWrapper
调试	—	is_loss_parallel_active()	—

---

关联：changzherui1/hyper-parallel#1 · changzherui1/hyper-parallel#6 · mindspore/hyper-parallel#270

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

Start by reading hyper_parallel/core/tensor_parallel/loss_parallel.py, platform/*/loss_parallel_ops.py, and _ce_op_registry.py, then review the listed PyTorch and TorchTitan loss implementations. Run tests/torch/loss_parallel and tests/mindspore/st/loss_parallel to establish current behavior. Done means the prioritized mesh, labels, API, reduction, framework-loss, and registration gaps are implemented with numerical and distributed regression coverage.