Amazon SageMaker HyperPod 存储设计与实践（二）

fullstate_save_policy = FullStateDictConfig(offload_to_cpu=True, rank0_only=True)

with FSDP.state_dict_type(
        model, StateDictType.FULL_STATE_DICT, fullstate_save_policy):

    # On All Ranks
    model_state = model.state_dict()
    
    # On Rank0 Only
    torch.save(model_state, <save_path>)
    
# On All Ranks
optim_state = FSDP.optim_state_dict(model, optimizer)
# On Rank0 Only
torch.save(optim_state, <save_path>)

from torch.distributed.checkpoint.default_planner import DefaultSavePlanner
import torch.distributed._shard.checkpoint as dist_cp

# On All Ranks of a shard group (Node)

distributed_writer = dist_cp.FileSystemWriter(<save_path>)

with FSDP.state_dict_type(model, StateDictType.SHARDED_STATE_DICT):
    state_dict = {}
    state_dict["model"] = model.state_dict()
    state_dict["optim"] = FSDP.optim_state_dict(model, optim)
    
    dist_cp.save_state_dict(
            state_dict=state_dict,
            storage_writer=distributed_writer,
            planner=DefaultSavePlanner(),
        )

当模型训练分片策略使用 HYBRID_SHARD 时，由于单个 shard group 持有完整的模型副本，因此仅需单个 shard group 进行写入，torch.distributed 会通过 NCCL 集合操作对于同时跨越多个 Rank 间的交互进行协调 ，具体参考 FSDP。

结果与分析

针对训练过程中模型参数及优化器的状态保存，以下表格汇总了不同分片策略及其在不同的存储方案上的检查点（checkpoint）写入性能：

	单分片存储占用(GB)	分片数量	Checkpoint至NVMe耗时(Sec)	NVMe转存S3耗时(Sec) *	NVMe链路累计耗时(Sec)	Checkpoint至Mountpoint for S3耗时(Sec)	Checkpoint至FSx for Lustre耗时(Sec) **
Model (Full)	29.93	1	114	81.54	195.54	107	129
Optimizer (Full)	59.85	1	260	176.31	436.31	225	254
Model (Sharded)	3.74	8	38	12.79	50.79	***	29
Optimizer (Sharded)	7.48	8	75	23.94	98.94	***	63

*计算节点的实例 NVMe 存储转存至 S3 使用高性能 S3 存取工具 s5cmd

**FSx for Lustre 共预置 60TiB 存储容量，单位吞吐能力 250 MB/s/TiB，累计吞吐能力 15000 MB/s

***Mountpoint for S3 暂不支持完整的 POSIX 接口，无法支持在分片存储时需要的 rename 操作

基于以上数据，该场景中对于 8B 参数量模型训练时的检查点（checkpoint）保存，其检查点（checkpoint）的存储消耗累计达到了 90GB （其中的 Optimizer 占用约 60GB）。同时，评估了其在 FULL_SHARD 及 HYBRID_SHARD 两种训练策略下的检查点（checkpoint）分片策略，在计算实例的本地 NVMe 存储，挂载 Mountpoint for S3 及 FSX for Lustre 存储上的性能。可以看出，使用分片形式进行模型检查点（checkpoint）存储，由于多个 GPU 进程可以进行并发写操作，因此整体的时间消耗小于写入一个完整的检查点（checkpoint）。另外，可以观察到，在本测试场景中，FSx for Lustre 对于检查点（checkpoint）分片的存储形式可以有极大的性能提升。因此，对于一般的训练场景的检查点（checkpoint）存储选型，有如下最佳实践建议：

1. 需要根据实际需要来选取合适的检查点（checkpoint）保存频率及保存范围，如训练代码的容错策略，或模型评估频率等
2. 对于较小参数量模型（如 8B 及以下）的检查点（checkpoint）存储场景，建议直接使用 Mountpoint for S3 进行全量的检查点（checkpoint）存储。优势有：
   • 高通用性的检查点（checkpoint）存储形式，具备极低的代码复杂度；且可以直接以文件形式直接写入 S3 对象存储，进行持久化保存
   • 检查点（checkpoint）写入的性能及成本有较好的平衡，且 S3 容量无需预置
3. 对于中大参数量模型（如 70B 及以上）且集群规模较大，需要相对高频进行检查点（checkpoint）存储的场景，建议使用 FSx for Lustre 结合分片并发检查点（checkpoint）存储。优势有：
   • 当使用分片存储策略时，一个模型副本中有多个 GPU 进程并发执行分片的存储，性能提升的同时，减少了对于单计算节点的 CPU RAM 依赖
   • 在 FSx for Lustre 预置一定规模的存储容量（如测试示例中的 60TiB）时，其具备较高的并发能力（如测试示例中的 60 个可并发的 250MB/s 的存储单元）。分片进行检查点（checkpoint）写入时间有大幅度的降低，避免了 GPU 集群的等待开销

除去以上介绍的基于 torch.save() 以及 torch.distributed 的不同存储形式，主流并行框架也内置了不同的检查点（checkpoint）存储策略，如 deepspeed 中的 save_checkpoint 实现方式，或 megatron.core 中的 dist_checkpointing 实现方式等。建议针对实际的训练任务及训练方案选型，以及根据不同方案或方案组合的性能测试结果进行实际生产场景的选型。

方案总结

本文在前文 HyperPod 存储设计与实践（一）的基础上进一步探讨了在 SageMaker HyperPod 中实现 Mountpoint for S3 支持的方法，包括 client 安装、创建挂载和检查挂载状态的脚本等。通过进行 Lifecycle 配置，可以在集群创建时自动挂载 Mountpoint for S3。文中还阐述了在 HyperPod 的 Self-healing 特性下，通过 Lifecycle 脚本实现自动存储挂载的必要性，以确保在计算节点发生故障及自动的算力节点替换后，使得训练任务可以无缝继续。最后，文章基于实际训练场景检查点存储的性能测试，对不同规模的模型训练场景，给出了存储选择的最佳实践建议。对于小规模模型训练，建议使用 Mountpoint for S3 进行检查点全量存储；对于大规模模型高频检查点场景，建议使用 FSx for Lustre 分片存储检查点,以提高写入性能。

本篇作者