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 函数使用 AllReduce 实现同步，应该可以直接调用，通过 AllReduce 来同步。WaitForWriterManifests 和 WaitForGlobalMetadata 应该可以直接删掉。

[chen2021673]

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

[JYMiracle305]

不再使用轮训的方式没有这个风险了

[chen2021673]

关于 Checkpoint::Save/Load 与 LoadDistributedCheckpoint 的分层，建议调整一下。

现在 Checkpoint::Save/Load 已经没有调用方了（只剩 test_trainer_state.cc、test_lr_scheduler_state.cc、test_checkpoint_serialization.cc 在调），实际入口变成了 Checkpoint::SaveSharded 和 LoadDistributedCheckpoint。这样不太规范，建议改成：

• Checkpoint::Save/Load —— 格式层，唯一入口。负责分片布局、metadata、rank 协同。TP=1 是 world_size=1 的退化情形，不是另一条代码路径。这也正是本 PR 的目的：TP=1 和 TP=8 的 ckpt 可以互读，那么入口自然也该是同一个，不需要 LoadDistributedCheckpoint 这个名字。
• SaveCheckpoint/ResumeFromCheckpoint —— 策略层。负责 iter_xxxxxxx 目录、架构校验等。

相应地 reshard.h/reshard.cc 可以整体并入 Checkpoint::Load —— Load 本身就 reshard-aware 之后，不存在第二个 reshard 入口。

[chen2021673]

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

[JYMiracle305]

已修改

[chen2021673]

这个函数没有被使用，可以删掉。

[JYMiracle305]

done

[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 ，但也不能挪到其他地方，因为 lora 本身不知道 qkv 是不是 packed 的。这里要不加个 fixme，之后把 packed QKV 本身抽成专门的 sharding abstraction。

[chen2021673]

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

[chen2021673]

这里能不能构造时就把 global layer_number 传进来，TransformerModel::ShardedStateDict() 中直接按 layer 生成，而不是 Remap 一遍。

[chen2021673]

这个函数其实是为了把 private 的 SaveTrainerState 方法暴露出来。
如果分布式 checkpoint 变成默认的 Checkpoint::Save/Load，那这些 wrapper 基本都可以删掉。顶层只有Checkpoint::Save/Load，内部直接 SaveTrainerState 即可。下同。

[chen2021673]

同上，不要硬编码。

[chen2021673]

TrainerState 和 metadata.json 的部分信息是重复的。建议

// trainer_state.json：

// 训练语义 / common state

global_step
consumed_train_samples

model config:
    n_layer
    n_head
    n_kv_head
    n_embd
    original_vocab_size
    padded_vocab_size

runtime config:
    tp_size
    pp_size
    dp_size
    sp_size
    vpp_size


// metadata.json：

// checkpoint format/version

tensor A:
    key
    dtype
    global_shape

    shard 0:
        global_offset
        local_shape
        segments
        writer/file
        file_offset
        byte_size

    shard 1:
        ...

也就是说 metadata 存

{
  "version": 3,
  "format": "infinitrain_sharded",
  "tensors": [...]
}

就够了。

[chen2021673]

同上，这个函数需要修改。

[chen2021673]

不应该 由 TP/PP 是否变化来控制逻辑，应该 load 时先根据当前运行拓扑重新构造 target sharded state dict，由 saved layout + target layout 驱动是否 load，参考 Megatron https://github.com/NVIDIA/Megatron-LM/blob/main/megatron/training/checkpointing.py?#L1838-L1908 。
另外这里如果暂时不支持的 VPP 话可以加一个抛出错误，同时说明 TODO。