[JYMiracle305]

背景

现有 distributed checkpoint 与保存时的 TP/PP 拓扑绑定，恢复训练时要求使用相同的并行配置。本 PR 引入基于全局张量坐标的 checkpoint resharding，使 checkpoint 可以在不同 TP/PP 配置之间恢复。

设计文档：Checkpoint Resharding 设计

主要修改

• 新增 ShardedTensor / ShardedStateDict，描述张量的全局形状、本地分片、全局偏移和切分方式。
• 新增 SavePlanner，统一规划模型参数和 Adam optimizer state 的本地写入布局，并生成可用于 reshard 的全局 metadata。
• 新增 LoadPlanner，根据源 checkpoint 与当前目标拓扑的分片坐标计算重叠区间。
• 新增 IndexedRegionLoadStrategy，按 metadata 中的文件和 offset 直接读取所需区域，在加载阶段完成重组，无需预先生成中间 checkpoint。
• 支持 TP 扩大、缩小、PP stage 变化以及 VPP 拓扑变化；VPP 加载使用全局 layer key，不依赖保存时的 local layer index。
• 支持 packed QKV 和 packed SwiGLU 等非连续分段布局，通过 segments 保持目标张量中的局部排列及其对应的全局区间；覆盖模型权重、Adam m/v 和 packed LoRA-B。
• 支持词表尾部 padding 和保存/目标 dtype 转换。
• 模型与 Adam m/v 共用参数分片信息，训练状态和 LR scheduler 状态随 checkpoint 一并恢复。
• 保存阶段当前仅由 dp_rank=0 的 TP/PP ranks 写入 shard，并在所有 rank metadata 就绪后原子发布全局 metadata；后续将由 shard 的 replica_id 选择 writer。
• trainer_state.json 保存训练语义、模型配置和运行时并行配置；metadata.json 仅保存 checkpoint version/format 以及 tensor shard 的 shape、offset、segments 和文件位置。
• metadata 加载严格校验 dtype、整数数组和 segment 连续性；未知 dtype、非法数字以及 segment gap/overlap 会直接报错，不再默认回退为 FP32。

当前限制

• 仅支持本 PR 引入的 distributed checkpoint version 3；加载时必须存在 metadata.json，不兼容旧 distributed checkpoint 文件。
• 每个张量目前只支持一个 fragmented axis；QKV 等非连续布局通过 segments 表达。
• 暂不支持 DistributedOptimizer / ZeRO optimizer state 的保存和恢复。
• optimizer resharding 依赖稳定的 named parameters。
• GPT-2 tied weight 暂不支持从 PP>1 checkpoint 安全恢复到 PP=1。保存端的 wte.weight 和 lm_head.weight 可能是独立张量，而 PP=1 目标端会将两个 key 映射到同一 Tensor；在实现 canonical tied key 或跨 stage weight tying 前应避免该拓扑转换。

测试

• TP 2 -> 4、TP 4 -> 2 的 overlap planning。
• 非均匀 shard 的显式 global offset。
• QKV segmented layout 在 TP 变化后的重组。
• packed SwiGLU gate/up segment metadata，以及 TP 2 -> 1、TP 1 -> 2 的实际文件读取和重组。
• packed LoRA-B 的 SwiGLU TP shard/restore。
• PP stage 变化时只加载目标 stage 所需参数。
• VPP checkpoint 使用全局 layer key 完成跨 VPP 拓扑加载。
• vocabulary padding、BF16 -> FP32 dtype 转换。
• model/optimizer state metadata 和 checkpoint serialization round trip。
• optimizer moment 使用实际 state dtype，同时复用模型 shard metadata。
• 拒绝未知 checkpoint dtype 和 malformed metadata 整数数组。
• metadata.json 不写入 iteration、parallel_config 和 model_config；TrainerState 保存并恢复 VPP 配置。

[Chamberlain0w0]

对于 PR 描述中 “当前限制” 的部分，建议在代码中对应位置加上 TODO/FIXME