Feat/checkpoint resharding

[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 变化；支持 QKV 非连续分段布局、词表尾部 padding 和保存/目标 dtype 转换。
• 模型与 Adam m/v 共用参数分片信息，训练状态和 LR scheduler 状态随 checkpoint 一并恢复。
• 保存阶段仅由 dp_rank=0 的 TP/PP ranks 写入 shard，并在所有 rank metadata 就绪后原子发布全局 metadata。

当前限制

• 仅支持 checkpoint version 3。
• 每个张量目前只支持一个 fragmented axis；QKV 等非连续布局通过 segments 表达。
• 暂不支持 DistributedOptimizer / ZeRO optimizer state 的保存和恢复。
• optimizer resharding 依赖稳定的 named parameters。

测试

• TP 2 -> 4、TP 4 -> 2 的 overlap planning。
• 非均匀 shard 的显式 global offset。
• QKV segmented layout 在 TP 变化后的重组。
• PP stage 变化时只加载目标 stage 所需参数。
• vocabulary padding、BF16 -> FP32 dtype 转换。
• model/optimizer state metadata 和 checkpoint serialization round trip。