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

动机（Motivation）

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

1. 间接调用散布全目录：15 个源文件中有 11 个存在 platform.xxx() 调用，共 100+ 处。读代码时无法直接看出一次 platform.detach() / platform.broadcast_async() 究竟落到哪个框架的哪个 API。
2. 双分支重复实现：filesystem_storage.py 中 _fetch_tensor_file / _open_checkpoint_files / _broadcast_batch_bytes 各带一套 PlatformType.PYTORCH 分支与 MindSpore 对侧实现；offline_transform.py 的类型校验与 checkpoint 格式推断同样双份。
3. 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 各自拆出独立模块，属于另一项重构。

相关的RFCs和API

对外接口变更

• 移除：hyper_parallel.core.distributed_checkpoint 的 8 个 __all__ 导出（见上）
• 契约新增：LoadPlanner.acquire_tensor 的返回值必须为 detached view（自定义 Planner 实现方需注意）
• convert_full_checkpoint_to_dcp 的 src_platform 取值从 Literal["torch", "huggingface", "mindspore"] 收窄为 Literal["torch", "huggingface"]

依赖关系

本 RFC 不依赖其他 RFC。与 core/dtensor/ 收敛为 torch-only 的既有决策方向一致（该目录已在 pylint 插件中登记为 torch-only）。

文档同步更新：docs/guide/distributed_checkpoint.md 的模块表、示例与正文中残留的 platform.* 提法。

完整的反馈期限

一周。

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

MindSpore 相关 ST（tests/mindspore/st/distributed_checkpoint/）随实现一并删除。

Before submitting a new issue...

• Make sure you already searched for previous RFCs.

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

Start by reviewing hyper_parallel/core/distributed_checkpoint/ and the listed filesystem_storage.py, offline_transform.py, util.py, ragged_utils.py, and init.py changes, then inspect the related unit tests and docs/guide/distributed_checkpoint.md. Validate the proposed torch-only structure, removals, API compatibility, and detached acquire_tensor contract against the stated tests and pylint checks. Done means the refactor and documentation updates pass the listed unit, ST, and lint validation without new failures.