ColossalAI Booster Checkpoint 全解基于 Booster 统一封装模型、优化器与学习率调度器的保存与恢复【免费下载链接】ColossalAIMaking large AI models cheaper, faster and more accessible项目地址: https://gitcode.com/GitHub_Trending/co/ColossalAI本文基于 ColossalAI 官方文档 docs/source/zh-Hans/basics/booster_checkpoint.md 编写结合仓库源码与真实训练示例展开。作为 ColossalAI 新版训练入口Booster的配套能力Checkpoint 模块把模型权重、优化器状态、学习率调度器状态统一封装为简洁一致的 API让用户在任意分布式插件Gemini、Hybrid Parallel、TorchDDP、TorchFSDP、LowLevelZero 等之上都能以相同的方式保存与恢复训练现场并能直接与 Hugging Face 生态的预训练权重与分片格式互通。一、导读与适用场景在 ColossalAI 中Booster见 colossalai/booster/booster.py是封装模型、优化器、损失函数、DataLoader 与学习率调度器的高层训练 API可先阅读前置教程 Booster API 了解booster.boost的用法。与之配套ColossalAI 提供了一整套Booster Checkpoint能力核心价值在于统一入口所有对象都先经过Booster.boost封装再由booster.save_model / load_model / save_optimizer / load_optimizer / save_lr_scheduler / load_lr_scheduler六个方法统一读写无需关心底层是单卡、ZeRO 还是混合并行多格式自适应load_model会自动探测 checkpoint 是单文件、Hugging Face 风格分片目录还是 index 文件并按对应方式加载生态兼容分片保存的模型权重与 huggingface/transformers 的pytorch_model.bin.index.json分片格式兼容用户可以直接用from_pretrained加载反之也能从本地 Hugging Face 目录恢复大模型权重。读完本文你将掌握如何使用六个 checkpoint 方法完整保存与恢复一次分布式训练如何借助LazyInitContextbooster.load在单卡放不下的情况下加载 Hugging Face 预训练大模型以及分片 checkpoint 在磁盘上的实际目录结构与其底层加载原理。二、整体架构一切 Checkpoint 都从checkpoint_io出发从源码看Booster 的 checkpoint 能力并非自己实现 IO而是组合一个CheckpointIO实例。在 booster.py 中if self.plugin is not None and self.plugin.control_checkpoint_io(): self.checkpoint_io self.plugin.get_checkpoint_io() else: self.checkpoint_io GeneralCheckpointIO()也就是说当使用HybridParallelPlugin张量并行 / 流水线并行 / ZeRO 组合等插件时插件会返回专门的HybridParallelCheckpointIO见 colossalai/checkpoint_io/hybrid_parallel_checkpoint_io.py它在保存时会自动把张量并行分片参数gather成完整权重、识别流水线阶段的参数前缀等其它场景下使用GeneralCheckpointIO见 colossalai/checkpoint_io/general_checkpoint_io.py。所有的 IO 类都继承自抽象基类CheckpointIO见 colossalai/checkpoint_io/checkpoint_io_base.py其中定义了模型类save_model / load_model及对应的 sharded / unsharded 子方法优化器类save_optimizer / load_optimizer及对应子方法学习率调度器类save_lr_scheduler / load_lr_schedulerLoRA 扩展save_lora_as_pretrained需安装peft。模块出口在 colossalai/checkpoint_io/init.py对外暴露CheckpointIO、GeneralCheckpointIO、HybridParallelCheckpointIO、CheckpointIndexFile、MoECheckpointIO。因此即便不走Booster也可以直接构造对应的checkpoint_io对象进行读写Booster 只是为你自动选择了合适的实现。三、模型 Checkpointsave_model / load_model3.1 使用方法与约束模型在保存/加载前必须经过Booster.boost封装。两个方法的签名定义于 colossalai/booster/booster.pydef save_model( self, model, # Union[nn.Module, ModelWrapper]已被 Booster 封装的模型 checkpoint: str, # checkpoint 路径必须是本地路径 shard: bool False, # 是否以分片方式保存 gather_dtensor: bool True, prefix: Optional[str] None, size_per_shard: int 1024, use_safetensors: bool False, use_async: bool False, ) - None def load_model( self, model, checkpoint: str, strict: bool True, low_cpu_mem_mode: bool True, num_threads: int 1, ) - None关于checkpoint路径语义原文档明确指出若shardFalsecheckpoint是一个文件若shardTruecheckpoint是一个文件夹分片 checkpoint 会包含 index json 与多个分片文件。shardTrue时 checkpoint 以分片方式保存。当单个 checkpoint 文件太大、无法放进单文件时这个模式非常实用。分片格式与 huggingface/transformers 兼容所以用户可以使用from_pretrained从分片 checkpoint 加载模型。3.2 save_model 参数详解含默认值save_model的完整参数行为可从 checkpoint_io_base.py 与 booster.py 的 docstring 归纳如下参数默认值说明model必填被 Booster 封装的模型或其ModelWrapper。checkpoint必填本地路径。shardFalse时为文件路径shardTrue时必须为目录路径。shardFalse是否分片。为True时输出一个目录目录结构与 Hugging Face transformers 分片权重一致含.index.json为False时输出单一文件。gather_dtensorTrue是否将分布式张量如张量并行、ZeRO 场景下被切分的参数收集gather到首个设备再保存。分布式张量无法直接加载离线转换才能还原为完整权重。prefixNone加在参数 / buffer 名前的前缀用于拼出 state_dict 中的 key分片场景下也会体现在权重文件命名上如pytorch_model.prefix.bin。size_per_shard1024每个分片文件的最大体积单位MB仅shardTrue时生效。use_safetensorsFalse是否使用 safetensors 格式保存。为True时权重文件后缀为.safetensorsindex 为model.safetensors.index.json。use_asyncFalse是否异步保存模型的 state_dict。注意源码中use_async仅在use_safetensorsTrue时真正生效否则会打日志警告并自动把use_safetensors置为True。文件名常量集中定义在 colossalai/checkpoint_io/utils.py默认权重文件pytorch_model.bin、safetensors 权重文件model.safetensors、对应 index 文件pytorch_model.bin.index.json/model.safetensors.index.json。3.3 分片保存的原理与磁盘结构当shardTrue时GeneralCheckpointIO.save_model内部会把模型state_dict拆成不超过size_per_shard的若干块用shard_model_checkpoint将完整 state_dict 切分成多个 shard见 utils.py并以“边拼边校验”的方式保证每个 shard 不超过上限见StateDictSharder.append_param单参数大于上限时也会单独成片按 index 生成pytorch_model-00001.bin、pytorch_model-00002.bin… 这类 5 位序号的分片文件维护一个CheckpointIndexFile见 colossalai/checkpoint_io/index_file.py记录weight_map参数名 → 分片文件名与metadata最后写出pytorch_model.bin.index.json若模型是 Hugging FacePreTrainedModel还会额外调用save_config_file把config.json/generation_config.json一并保存到目录中见 utils.py这正是分片目录能直接from_pretrained的原因之一。因此典型的分片目录结构为./checkpoints/ ├── pytorch_model.bin.index.json # 记录每个参数落在哪个分片 总大小 ├── config.json # Hugging Face 模型配置自动附带 ├── pytorch_model-00001.bin ├── pytorch_model-00002.bin └── ...3.4 load_model 的自动格式探测load_model会“自动检测 checkpoint 格式并以相应方式加载”其探测逻辑在 checkpoint_io_base.py实质由工具函数has_index_file完成见 utils.py若checkpoint是一个文件名以.index.json结尾的文件判定为分片 index进入load_sharded_model若checkpoint是一个目录且其中恰好存在唯一一个*.index.*json文件判定为分片目录若目录中无 index 文件则依次尝试目录下的model.safetensors、pytorch_model.bin再退化为把传入路径当单文件加载。也就是说用户加载时完全不用关心当初是用哪种方式保存的。此外load_model还提供两个对超大模型友好的参数low_cpu_mem_modeTrue以“低 CPU 内存模式”逐片加载避免把整个权重一次性读进内存low_cpu_mem_modeFalse时可通过num_threads启用多线程 RAM cache 来加速加载。strictTrue默认则要求 checkpoint 中的 key 与模型 state_dict 完全匹配否则抛出带缺失 key 信息的RuntimeError。3.5 从 Hugging Face 加载超大预训练模型的最佳实践原文档针对“模型太大、单设备无法直接from_pretrained”给出了推荐方案其核心思路是先下载权重到本地在 Lazy Initialization 环境下初始化模型避免初始化时占满内存封装后再用booster.load_model从本地路径加载。原文伪代码如下from colossalai.lazy import LazyInitContext from huggingface_hub import snapshot_download # Initialize model under lazy init context init_ctx LazyInitContext(default_deviceget_current_device) with init_ctx: model LlamaForCausalLM(config) # Wrap the model through Booster.boost model, optimizer, _, _, _ booster.boost(model, optimizer) # download huggingface pretrained model to local directory. model_dir snapshot_download(repo_idlysandre/arxiv-nlp) # load model using booster.load booster.load(model, model_dir)其中LazyInitContext属于 ColossalAI 的懒惰初始化能力模型的参数在with块内并不会被真实分配/初始化而是登记一张“待初始化清单”从而大幅节省初始化超大模型时的峰值内存。配合booster.load_model文档中的booster.load在封装完成后、真正前向之前把本地权重写回即可在显存受限的情况下装载大模型。相关实现的入口见 colossalai/lazy/init.py 及其 lazy_init 实现更多细节可参考文档中“懒惰初始化”专题Lazy Initialization若该路径存在于当前仓库时。注意上面伪代码中的booster.load即本文对应的booster.load_model当前仓库版本统一命名为load_model。四、优化器 Checkpointsave_optimizer / load_optimizer优化器的动量、方差等状态直接决定断点续训能否精确衔接因此优化器状态需要与模型分开单独存取。同样优化器在保存/加载前必须先被 Booster 封装此时可能是被 ZeRO 等插件改写过的OptimizerWrapper。4.1 API 签名与要点def save_optimizer( self, optimizer: Optimizer, # 被 Booster 封装后的优化器 checkpoint: str, shard: bool False, gather_dtensor: bool True, prefix: Optional[str] None, size_per_shard: int 1024, use_async: bool False, ) - None def load_optimizer( self, optimizer: Optimizer, checkpoint: str, low_cpu_mem_mode: bool True, num_threads: int 1, ) - None从源码与文档说明可以确认的要点与模型类似shardFalse时checkpoint是单文件shardTrue时是一个目录目录中会生成index 文件 param group 文件 多个状态分片文件优化器状态保存不支持 safetensorsCheckpointIO.save_optimizer的 docstring 明确说明因此优化器分片文件后缀固定为.bin加载时同样会探测 index 文件。特别地若传入的是一个目录却找不到 index 文件会抛出ValueError: Cannot find index file in ...见 checkpoint_io_base.py。4.2 分片优化器目录的完整结构在 general_checkpoint_io.py 的 docstring 中给出了分片优化器保存时生成的文件清单一个 index 文件pytorch_optim.bin.index.json维护优化器状态与分片文件名的映射一个 group 文件pytorch_optim_group.bin记录param_groups的信息注意 optimizer 的param_groups里params存储的是张量对象无法直接torch.save因此单独用一个文件保存索引化之后的 group 结构多个状态分片文件pytorch_optim-00001.bin、pytorch_optim-00002.bin…真正保存各参数的 state 张量如 Adam 的exp_avg、exp_avg_sq、step。目录结构示意./optim_ckpt/ ├── pytorch_optim.bin.index.json ├── pytorch_optim_group.bin ├── pytorch_optim-00001.bin ├── pytorch_optim-00002.bin └── ...底层在切分时使用shard_optimizer_checkpoint见 utils.py只切分state_dict[state]param_groups单独处理每个参数状态可能包含多个张量如 Adam 每个参数有step/exp_avg/exp_avg_sq会把这些张量视为一个整体一起放入某个分片。加载侧则由load_param_groups_into_optimizer重建参数 id 与张量的映射id_map随后load_states_into_optimizer把各状态张量 cast 到当前参数所在的设备与 dtype最后调用sharded_optimizer_loading_epilogue完成 PyTorch 侧的收尾_patch_step_function/_hook_for_profile等见 utils.py。4.3 一个容易踩的坑因为优化器状态里有“每个参数 → 状态”的强绑定关系加载分片优化器时要求当前优化器的 param_groups 数量与结构必须和保存时一致见 utils.py 对 group 数量与参数个数的校验。因此在load_optimizer之前你需要用与训练时完全相同的配置重新构建模型和优化器同样的学习率分组、同样顺序的层而不是用一个随机初始化的优化器去硬载。五、学习率调度器 Checkpoint学习率调度器同样支持一行保存 / 恢复def save_lr_scheduler(self, lr_scheduler: LRScheduler, checkpoint: str) - None def load_lr_scheduler(self, lr_scheduler: LRScheduler, checkpoint: str) - None调度器在保存/加载前必须被 Booster 封装与模型、优化器不同checkpoint只允许是本地文件路径不支持分片目录底层实现非常简单且标准save_lr_scheduler直接torch.save(lr_scheduler.state_dict(), checkpoint)load_lr_scheduler先torch.load再load_state_dict见 checkpoint_io_base.py。由于调度器状态极小实际工程中通常直接与模型、优化器并列保存为三个独立文件。六、从一次完整训练看三者的配套使用仓库中的多个官方示例都遵循“模型 优化器 调度器”三者成组保存、断点续训时三者一起恢复的模式。以 examples/tutorial/new_api/cifar_resnet/train.py 为例同样写法也出现在 examples/images/resnet/train.py、examples/tutorial/new_api/cifar_vit/train.py# 恢复断点续训 booster.load_model(model, f{args.checkpoint}/model_{args.resume}.pth) booster.load_optimizer(optimizer, f{args.checkpoint}/optimizer_{args.resume}.pth) booster.load_lr_scheduler(lr_scheduler, f{args.checkpoint}/lr_scheduler_{args.resume}.pth) # 周期性保存 booster.save_model(model, f{args.checkpoint}/model_{epoch 1}.pth) booster.save_optimizer(optimizer, f{args.checkpoint}/optimizer_{epoch 1}.pth) booster.save_lr_scheduler(lr_scheduler, f{args.checkpoint}/lr_scheduler_{epoch 1}.pth)推荐的最佳实践可总结为启动时先booster.load_model恢复权重若只想做推理/评测这一步即可需要继续训练时务必按相同配置重建模型与优化器再依次load_optimizer、load_lr_scheduler训练中以(model, optimizer, lr_scheduler)三件套按 epoch / step 周期保存必要时把 dataloader 的迭代位置、随机种子一并记录构成完整的“训练现场”。而针对超大模型的发布/续训则推荐把模型与优化器都用shardTrue保存为目录例如 examples/tutorial/new_api/cifar_resnet 中的booster.save_model(model, args.output_path, shardTrue)或 examples/language/opt/opt_train_demo.py、examples/images/vit/vit_train_demo.py 里对输出模型同样采用shardTrue的写法。七、Checkpoint 设计背后的统一化考量原文档的最后一节将读者引向社区讨论A Unified Checkpoint System Design其背后的设计动机正是从架构层面解答“为什么六个 API 长这样”封装即边界boost之后模型 / 优化器都可能是带分布式状态的 wrapperZeRO 分片、张量并行切分、流水线阶段切分……。checkpoint API 面向的是 wrapper 而非裸对象因此底层的CheckpointIO才有机会在保存时把分片gather成完整权重、在加载时把完整权重重新正确地切回各 rank加载宽容、保存严格保存时需显式声明是否分片决定路径是文件还是目录加载时则完全由has_index_file自动判定降低出错概率通用与专精分层GeneralCheckpointIO处理最普通的单机/非并行场景HybridParallelCheckpointIO见 colossalai/checkpoint_io/hybrid_parallel_checkpoint_io.py面向混合并行会额外处理张量并行参数合并、pipeline stage 前缀、pad tensor 的裁剪/还原to_padded_tensor/to_unpadded_tensorMoE 场景则由 colossalai/checkpoint_io/moe_checkpoint.py 的MoECheckpointIO负责。值得说明的是当前实现的限制是分布式张量dtensorcheckpoint 无法直接加载必须先用离线 CLI 转成聚合gathered后的完整张量 checkpoint 才能读取这一点在CheckpointIO.load_model、load_state_dict等多处均有显式断言。规划分片恢复方案时需要把它纳入考量。八、总结Booster Checkpoint 通过save_model / load_model / save_optimizer / load_optimizer / save_lr_scheduler / load_lr_scheduler六个方法把分布式训练中最繁琐的“各种并行状态如何存取”收敛为一套规则简单的 API对象保存入口加载入口分片支持与生态的兼容性模型权重save_modelload_model自动探测格式支持格式与 HF transformers 分片一致分片目录可被from_pretrained直接读取反之亦可加载本地 HF 权重优化器状态save_optimizerload_optimizer支持index group shards不支持 safetensors仅.bin学习率调度器save_lr_schedulerload_lr_scheduler不支持单文件PyTorch 标准state_dict对于显存受限的 Hugging Face 大模型加载场景组合使用LazyInitContext惰性初始化、snapshot_download本地化权重与booster.load_model三步是目前仓库推荐的标准路径。在实际落地时建议把模型、优化器、调度器视作一个整体按固定节奏保存并保持“重建模型/优化器的配置与保存时一致”的纪律即可在任意 ColossalAI 并行插件之上获得可靠的断点续训与模型发布能力。【免费下载链接】ColossalAIMaking large AI models cheaper, faster and more accessible项目地址: https://gitcode.com/GitHub_Trending/co/ColossalAI创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考