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。

[chen2021673]

这里在一直轮询所有 metadata.json，但进程同步语义不应该绑定到文件是否出现上。我看已经有SynchronizeCheckpointRanks 函数了，应该可以通过改造复用来实现更可靠的判断。

[chen2021673]

这里如果继续用轮询检测 metadata file 的话可能有刷屏风险。

[chen2021673]

Checkpoint::Save 和 Checkpoint::Load 没有地方使用了，相应函数和测试都可以删掉。

[chen2021673]

只有一个 static 方法，不需要套一个 class 吧。

[chen2021673]

这个字段不需要吧，真正用的是 storage.data_offset

[chen2021673]

建议把 model.ckpt、optimizer.ckpt、metadata.json 等字符串都提出为常量，不要在各个地方硬编码。

[chen2021673]

这里应该 CHECK 或 LOG(WARNING)，tensors 不在 ShardedStateDict 里是不正常的。

[chen2021673]

nn 是基础设施，checkpoint 是上层功能，nn 层不应该依赖 checkpoint 层。PyTorch 里也是 DCP 依赖 nn.Module，不是反过来。建议把 shard 描述类型下沉到基础层（infini_train/include/shard_spec.h，namespace infini_train），或让 nn 层只暴露一个不含 checkpoint 语义的轻量结构。

[chen2021673]

Transformer 层不应该 include lora 层，这里引用是为了拿 kParamLoraBName ，应该在 LoRAColumnParallelLinear::ShardedStateDict 里完成。

[chen2021673]

我们现在都是硬编码dp_rank = 0，但 Megatron / PyTorch 是 replica_id 挂在 ShardedTensor 上，由数据自己决定谁写，这里可以加个TODO。

[chen2021673]

这里能不能构造时就把 global layer_number 传进来，而不是 Remap 一遍