[RFC]: distributed_checkpoint 去 platform 抽象，收敛为 PyTorch-only-hyper-parallel-AtomGit

已开启

[RFC]: distributed_checkpoint 去 platform 抽象，收敛为 PyTorch-only #387

zhangbuxue创建于 17 天前

zhangbuxue成员

17 天前创建

动机（Motivation）

hyper_parallel/core/distributed_checkpoint/（下称 DCP）目前通过 get_platform() 抽象层同时支持 PyTorch 与 MindSpore 两个后端，代价是：

间接调用散布全目录：15 个源文件中有 11 个存在 platform.xxx() 调用，共 100+ 处。读代码时无法直接看出一次 platform.detach() / platform.broadcast_async() 究竟落到哪个框架的哪个 API。
双分支重复实现：filesystem_storage.py 中 _fetch_tensor_file / _open_checkpoint_files / _broadcast_batch_bytes 各带一套 PlatformType.PYTORCH 分支与 MindSpore 对侧实现；offline_transform.py 的类型校验与 checkpoint 格式推断同样双份。
MindSpore 侧已无维护：tests/mindspore/st/distributed_checkpoint/ 的离线转换 ST 早已失效（MindSpore Tensor 无 .detach 方法），在本次改动之前就无法通过。

抽象层带来的可读性与维护成本已显著超过其价值。DCP 收敛为 PyTorch-only 后，可直接使用 torch 语义，并顺带清理掉因双后端而存在的冗余结构。

目标设计

① 去 platform，直调 torch

get_platform() / platform.xxx() 全部替换为对应 torch 调用：

原调用	替换为
`platform.get_rank()` / `get_world_size()` / `barrier()`	`dist.get_rank()` / `dist.get_world_size()` / `dist.barrier()`
`platform.broadcast_async(t, src, g)`	`dist.broadcast(t, src, g, async_op=True)`
`platform.new_group(ranks)` / `destroy_process_group(g)`	`dist.new_group(ranks=list(ranks))` / `dist.destroy_process_group(g)`
`platform.all_gather_object(...)`	`dist.all_gather_object(...)`
`platform.detach(t)` / `platform.is_tensor(o)`	`t.detach()` / `isinstance(o, torch.Tensor)`
`platform.empty(...)` / `platform.new_tensor(...)`	`torch.empty(...)`
`platform.Tensor` / `platform.dtype`	`torch.Tensor` / `torch.dtype`

torch 无单一对应调用的（checkpoint 文件读写、通信组缓存查询、设备流同步、fused copy、dtype 字符串解析），在 DCP 目录内自行封装，不再穿透到 platform 层。

删除全部 PlatformType 分支及 MindSpore 对侧实现：_fetch_ms_tensor_file、_MS_FILES_KEPT、_broadcast_batch_bytes、src_platform="mindspore"。

② util.py 按职责拆分

原 util.py 达 960 行，混杂多种关注点。拆为（并顺带把 util.py 更名为 utils.py、ragged_utils.py 更名为 ragged.py，使目录里只剩一个公共工具模块）：

模块	职责
`checkpoint_io.py`	单个 checkpoint 文件的读写（safetensors / pickle）
`broadcast.py`	副本分片广播：通信组建立与销毁、发送、小分片批量合并
`utils.py`	路径、分片几何、state dict 遍历、区间求交、分阶段计时、统一 logger

③ 删除 DCP 不依赖的模块

经全仓库依赖排查，以下模块在 DCP 内部零使用，仓库内引用方仅为其自身 UT 与 __init__.py 的 re-export：

loader.py / saver.py：单文件 safetensors 读写包装（各约 40 行）
layout.py：layout 采集/落盘/跨 rank 汇总（153 行）
reshard.py 的 ReshardHandler 及其辅助函数（约 300 行）
ensure_broadcast_groups：_build_broadcast_groups 的薄封装，引入时即无生产调用方，加载路径走 broadcast_groups_for_load

其中 infer_intersection（计算两个分片区间的重叠）是 reshard.py 中唯一被 DCP 依赖的函数——standard_planner.py 在构造 read item 时用它计算本地分片与已存分片的交集——移入 utils.py，与产出其输入的 chunk_to_area 相邻。

需要设计评审DFX建议

替代方案比较

方案	优点	缺点
A. 保留 platform 抽象层，仅删 MindSpore 实现	改动面小；未来若恢复 MindSpore 支持成本低	间接层仍在，可读性问题未解决；留下一个只有单一实现的抽象，是典型的"预留扩展点"反模式
B. 直调 torch + 局部封装（本方案）	调用点语义直白；删除双分支后 `filesystem_storage.py` 减少约 90 行；DCP 目录自洽	若未来恢复 MindSpore 支持需重新引入抽象
C. 直调 torch，不做局部封装	最直接	`get_created_group` / `synchronize` / `copy_each` 等无 torch 单一对应，会在多个调用点重复实现

选 B。判断依据是 MindSpore 侧已无维护需求，且 core/dtensor/ 已有 torch-only 的先例（pylint 插件的 TORCH_ONLY_CORE_PARTS 已包含该目录）。

对现有代码的影响范围

39 个文件，+1796 / −3285（净减 1489 行）。

新增：broadcast.py、checkpoint_io.py、tests/ut/core/distributed_checkpoint/test_broadcast.py
删除：loader.py、saver.py、layout.py、reshard.py 及其 UT/ST（共 9 个文件）
core/dtensor/ 零改动——DCP 仅引用其原生的 Layout 与 infer_slice_area_by_layout
pylint 插件 scripts/pylint_hyperparallel.py：将 distributed_checkpoint/ 加入 TORCH_ONLY_CORE_PARTS，与既有的 core/dtensor/ 同列

向后兼容性

hyper_parallel.core.distributed_checkpoint.__all__ 移除 8 个导出：save_checkpoint、load_checkpoint、get_current_layout、get_global_layout、save_layout、load_layout、combine_layout、ReshardHandler。

已确认这 8 个名字在仓库内（hyper_parallel/、tests/、scripts/、examples/）均无实际调用方，仅存在于 __init__.py 的 re-export 与其自身 UT 中。若外部用户有依赖，需在 release note 中说明。

save / async_save / load / get_optim_state_dict / set_optim_state_dict 及全部 Planner / Storage 接口不变。

顺带修复的缺陷

StandardLoadPlanner.acquire_tensor 的两条分支中，只有非 ragged 分支做了 detach；ragged 分支经 get_ragged_box_tensor 返回的是带 grad_fn 的 view。实测对这种 view 做 copy_ 会抛：

RuntimeError: a view of a leaf Variable that requires grad is being used in an in-place operation.

此前未暴露，是因为 filesystem_storage._validate_and_copy_tensor 有一层兜底 detach。本次将两条分支统一为先选出本地分片、再统一 detach + narrow，并把「返回值必须 detach」写入 LoadPlanner.acquire_tensor 抽象接口契约，兜底层随之移除——责任从调用方收回到实现方。

性能影响

无。所有替换均为同语义直调（platform.get_rank() 本就是 dist.get_rank() 的转发）。广播的批量合并阈值 DEFAULT_BROADCAST_BATCH_BYTES 及在途上限 _MAX_BROADCASTS_IN_FLIGHT 行为不变；_open_checkpoint_files 在 torch 路径上的保留文件数仍为 8。

已知遗留

api.py 与 async_persist.py 之间存在 pylint R0401 cyclic-import。该循环在本次改动前即存在，由维护者有意通过函数内 import 在运行时打断（async_persist 中两处 from ...api import _save_impl，带 # pylint: disable=import-outside-toplevel）。本 RFC 不处理；彻底消除需将 storage 协同（FileType / gather_all_results_from_storage）与 _save_impl 各自拆出独立模块，属于另一项重构。

完整的反馈期限

一周。

CC List

（待补充）

其他补充说明

验证情况

pylint（.pylintrc + scripts/pylint_hyperparallel.py）：改动涉及的 35 个文件 0 告警（除上述改动前即存在的 R0401）
UT：tests/ut/core/distributed_checkpoint + tests/ut/core/dtensor 591 passed / 98 skipped
全量 tests/ut：与同机 upstream master 基线逐条比对，5 个失败两侧完全相同且均在 tests/ut/platform/**，零新增失败
ST：Ascend 910B3 双环境（torch 2.12/CANN 9.1、torch 2.9/CANN 9.2）level0 + level1 全绿，各 31 passed / 0 failed，明细见关联 PR 的 Test Plan