精度转换与存储瘦身实战指南)
人工智能大模型AI 技能/插件分布式训练微调深度学习【免费下载链接】ml-engineeringMachine Learning Engineering Open Book项目地址https://gitcode.com/gh_mirrors/ml/ml-engineering点击查看免费下载本指南聚焦 training/checkpoints 目录下两件核心工具torch-checkpoint-convert-to-bf16将 fp32 权重检查点批量转换为 bf16并同步转换 safetensors 及索引文件与torch-checkpoint-shrink.py修复因「tensor 的底层 storage 大于视图 view」而白白浪费磁盘空间的检查点。在大模型训练中检查点动辄数 TB转换与瘦身直接决定存储成本与恢复速度。读完本文你将掌握这两个工具的使用方式、底层原理与适用边界并能在自己的训练工程中直接复用之。为什么需要处理检查点存储成本与保存策略的现实约束在进入工具细节之前先明确这类脚本存在的工程背景。训练过程中检查点保存的频率与体积是一对直接矛盾保存越频繁故障恢复时丢失的训练步数越少但保存动作本身会拖慢训练体积越大磁盘与备份云存储的成本越高甚至导致训练因磁盘写满而崩溃。training/fault-tolerance/README.md 记录了 BLOOM-176B 训练的实测在 GPFS over NVMe 文件系统上384 个进程并发写入一个 2.3TB 检查点只需 40 秒每 3 小时保存一次、训练约 3 个月累计 720 个检查点仅保存动作就消耗了约 8 小时约占训练总时长的 0.37%若 IO 慢 5 倍这一开销将升至 2%。而存储章节也专门给出了「找出谁的大检查点在吞噬磁盘」的排查思路可见检查点体积管理是集群运维的常态痛点。在此背景下降低单个检查点的体积就成了直接收益fp32 → bf16的转换可将权重体积直接减半而 storage 瘦身则能剔除那些因保存时机不当而携带的冗余底层存储。两者合在一起往往能显著缓解磁盘压力并加快检查点读写速度。精度背景fp32 到 bf16为什么可以安全减半要理解转换脚本的价值先回顾 training/dtype.md 中的精度知识。bf16 与 fp32 的关键差异在于尾数位fp32 为 1 位符号 8 位指数 23 位尾数bf16 为 1 位符号 8 位指数 7 位尾数。两者的指数位同为 8 位意味着 bf16 完整保留了 fp32 的动态范围最大可表示数级一致只是把尾数从 23 位压缩到 7 位、丢失部分精度——这正是训练 LLM 时 bf16 比 fp16 更稳定的根本原因fp16 只有 5 位指数易上溢。因此fp32 权重转 bf16 一般安全范围不变仅精度下降适合推理或继续微调而反过来把 fp16 训练的模型放进 bf16 通常可行但会有精度损失最好先微调再使用用 bf16 预训练模型直接跑 fp16 则通常失败因为 fp16 最大只能表示约 64k极易溢出详见 training/dtype.md。对权重而非梯度、优化器状态这类「已收敛数值」而言转成 bf16 保存几乎不损失可用性却能换来存储与 IO 的双倍收益这正是仓库提供转换脚本的动机。工具一torch-checkpoint-convert-to-bf16 —— 一键将 fp32 检查点转为 bf16用法该工具是一个纯 bash 脚本使用前提是已安装 PyTorch若检查点包含 safetensors 文件还需要安装safetensors包。运行时需先cd到检查点目录再执行cd checkpoint bash training/checkpoints/torch-checkpoint-convert-to-bf16它会新建一个名为bf16的子目录并在其中生成与原始检查点等价的 bf16 版本原始文件保持不动。完整脚本位于 training/checkpoints/torch-checkpoint-convert-to-bf16处理流程共分四步第一步创建目标目录并复制元数据文件target_dirbf16 mkdir -p $target_dir cp *json *model $target_dir脚本把bf16作为固定目标目录名并复制所有*.json与*model文件如config.json、generation_config.json、tokenizer 相关文件等。注释中也提示了更粗暴的做法cp * $target_dir可根据实际检查点布局调整——若你的检查点还包含*.tokenizer、*.txt等文件需要自行补充cp规则这是「应易于适配其他类似场景」的落点。第二步转换 torch 的*.bin文件python -c import torch, sys; [torch.save({k:v.to(torch.bfloat16) for k,v in torch.load(f).items()}, f{sys.argv[1]}/{f}) for f in sys.argv[2:]] $target_dir *bin核心动作是一行 Python对每个*.bin文件执行torch.load把 state_dict 中的每个张量v做v.to(torch.bfloat16)再以原文件名写回目标目录。要注意的细节只转换了 state_dict 顶层值为 Tensor 的键{k:v.to(...) ...}仅作用于一层。如果你的检查点里嵌套了optimizer_state、module等子字典这套转换不会深入嵌套——脚本注释写明这是「for similar use cases」的模板需按需扩展成递归转换默认torch.load会按保存时的设备映射加载脚本未指定map_location因此推荐在 GPU 环境运行若在无 GPU 的机器上执行可能需要对torch.load加map_locationcpu命令末尾的*bin是 shell 通配符展开后的文件列表若目录中没有*.binpython将收到空参数列表此步静默无事发生。第三步转换 safetensors 文件if compgen -G *.safetensors /dev/null; then cd $target_dir python -c import re, sys, torch; from safetensors.torch import save_file; [save_file(torch.load(f), re.sub(r.*?(model.*?)\.bin,r\1.safetensors,f), metadata{format: pt}) for f in sys.argv[1:]] *bin if test -e pytorch_model.bin.index.json; then cp pytorch_model.bin.index.json model.safetensors.index.json perl -pi -e s|pytorch_||; s|\.bin|.safetensors| model.safetensors.index.json fi cd - /dev/null fi这段逻辑需要仔细拆解因为它涉及 torch 与 safetensors 两种格式的桥接compgen -G *.safetensors检测原始检查点是否包含 safetensors 文件。只有在原检查点存在 safetensors 时才进入该分支进入bf16子目录后用torch.load读取刚刚转换好的*.bin文件再通过safetensors.torch.save_file以model.*.safetensors命名写出。注意这里读的是bf16目录下的*bin也就是第二步的产物因此最终的.safetensors权重本身就是 bf16 精度metadata{format: pt}标记这是 PyTorch 格式的张量保证 Hugging Face transformers 等库能正确加载若存在分片索引pytorch_model.bin.index.json则复制为model.safetensors.index.json并用perl将索引内的文件名从pytorch_...bin批量改写为...safetensors使 transformers 能通过索引找到新的分片文件。第四步输出结果echo the dir $target_dir now contains a copy of the original checkpoint with bf16 weights完成后bf16/目录即为一份可直接加载的 bf16 检查点。若你的框架加载时偏好 safetensors如 transformers 默认优先.safetensors此目录已经同时具备两种格式的产物可按需选择。使用注意与扩展点不要重复执行脚本会把bf16固定为目标目录且不做存在性判断重复运行会在已有目录上继续写入cp与torch.save会覆盖同名文件但多分片场景可能有残留建议执行前确认目标目录为空或改名CPU 环境适配可在第二步的 Python 片段中为torch.load增加map_locationcpu再对v.to(torch.bfloat16)前加v.float()之类显式转换以应对跨 dtype 保存嵌套结构适配若 state_dict 存在嵌套将第二步的单层字典推导改为递归函数即可这与同仓库 torch-checkpoint-shrink.py 的递归处理方式可互相参考。工具二torch-checkpoint-shrink.py —— 剔除检查点中的冗余底层存储问题根源storage 大于 viewPyTorch 中一个 Tensor 由**底层存储storage与视图view**组成。视图只是对底层存储的一段「窗口」通过view、slice、expand、transpose等操作产生。正常情况下切片的底层存储与切片等大但某些框架或保存路径例如分布式优化器状态、算子中间产物会在保存时把「视图很小、底层存储很大」的张量原样写盘导致检查点文件体积远大于逻辑内容——这正是脚本注释中 some reason stored tensors with storage larger than their view at the moment of saving 所指的问题。用法# 处理检查点目录下全部 *.pt 文件 python training/checkpoints/torch-checkpoint-shrink.py --checkpoint_dir ./checkpoints/global_step10 # 只处理匹配多个模式的 pt 文件务必加引号避免 shell 展开 python training/checkpoints/torch-checkpoint-shrink.py --checkpoint_dir ./checkpoints/global_step10 --patterns layer*pt zero*pt # 开启调试输出每个文件的处理明细 python training/checkpoints/torch-checkpoint-shrink.py --checkpoint_dir ./checkpoints/global_step10 -d命令行参数由 torch-checkpoint-shrink.py 中的argparse定义含义如下参数类型默认值说明--checkpoint_dirstr必填目标检查点目录例如path/checkpoints/global_step10目录不存在或没有*.pt文件会抛出FileNotFoundError--patternsstr可多个*.pt一个或多个文件名通配模式使用fnmatch逐文件匹配只匹配文件名部分不含目录默认匹配所有*.pt-d/--debugflag关开启逐文件调试输出打印每个键的路径与每个文件的转换前后体积核心处理逻辑脚本先用glob收集目录下所有*.pt文件get_pt_files再按 patterns 过滤然后对每个文件执行 shrink_pt_filetorch.load(f, map_locationdevice)其中device torch.device(cpu)保证在任意机器上都能加载递归遍历 state_dict 的所有值shrink_dict_values对字典类型的值递归深入对torch.is_tensor(v)为真的叶子张量执行d[k] v.clone()clone()会为当前视图重新分配一块恰好等于视图大小的存储丢弃原有的大 storagetorch.save(sd, f)原地覆盖保存即原地修改原文件脚本不生成新文件请务必先备份或确认磁盘空间足够因为加载与保存可能同时存在新旧两份数据统计处理前后体积最终打印汇总checkpoint_shrinkProcessing zero checkpoint ./checkpoints/global_step10 - ./checkpoints/global_step10/zero_pp_rank_0_mp_rank_0_optim_states.pt Done. Before 1234.00MiB, after 1000.00MiB, saved 234.00MiBdebug模式下还会逐文件打印before/after/saved三个 MiB 数值便于定位哪些文件浪费最多。适用场景与边界该脚本只处理*.ptPyTorch 原生格式不处理*.bin或*.safetensors。若目标是 bf16 检查点瘦身可先跑本脚本再跑转换脚本或反之视目录内格式而定只做「存储裁剪」不做精度转换对已经是 bf16 的*.pt同样有效patterns的参数里fnmatch只匹配文件名若想只处理某个子目录下的文件需要自行调整glob模式原地覆盖意味着一旦出问题难以回滚建议先复制一份或对单文件试跑用--patterns限定-d观察效果。结合仓库的整体运维建议两个脚本配合仓库其他章节可形成一套完整的检查点生命周期管理保存策略按 training/fault-tolerance/README.md 的建议先实测单次保存耗时再权衡频率保留最近 2 份本地检查点以支持快速恢复其余异步离线上云体积优化训练期间定期用torch-checkpoint-shrink.py清理 optimizer state 类*.pt文件的冗余 storage训练结束后用torch-checkpoint-convert-to-bf16产出 bf16 副本用于推理或微调两者都能在存储层面直接减负空间监控结合 storage/README.md 中的磁盘使用排查手段配合du -ahd1 | sort -rh定位大检查点并考虑用定时任务离线上传、定期清理过期检查点。总结training/checkpoints/README.md 下的两个工具构成了检查点运维的最小实用集转换解决「精度与存储」问题fp32→bf16 体积减半、范围不变瘦身解决「存储与视图不匹配」的隐性浪费。它们都是刻意写成的小型模板脚本——转换脚本的注释明言 should be easily adaptable to other similar use cases瘦身脚本则用递归结构天然支持任意嵌套 state_dict。读者可以按本文拆解的逻辑将其扩展到 safetensors 直接转换、嵌套 state_dict 递归 bf16、或接入自己的检查点后处理流水线配合仓库的 fault-tolerance 与 storage 章节形成完整的训练资产闭环。赞分享人工智能大模型AI 技能/插件分布式训练微调深度学习【免费下载链接】ml-engineeringMachine Learning Engineering Open Book项目地址https://gitcode.com/gh_mirrors/ml/ml-engineering点击查看免费下载相关推荐Machine Learning Engineering Open Book用 Tiny 模型、Tokenizer 与数据集加速调试与开发Machine Learning Engineering Open Book用 Tiny 模型、Tokenizer 与数据集加速调试与开发 导读 本篇文章基人工智能大模型AI 技能/插件分布式训练微调深度学习TorchTitan 检查点系统实战DCP 存读、Seed Checkpoint 创建与 HuggingFace/Torch 格式转换全解TorchTitan 检查点系统实战DCP 存读、Seed Checkpoint 创建与 HuggingFace/Torch 格式转换全解 TorchTita人工智能大模型预训练分布式训练强化学习Fizzy API 集成指南从认证、缓存到分页与文件上传的完整实践Fizzy API 集成指南从认证、缓存到分页与文件上传的完整实践 本篇技术指南以 Fizzy 开源仓库的 API 文档 https://link.gitco人工智能大模型AI 技能/插件分布式训练微调深度学习上一篇3分钟解决HTML转Markdown难题Turndown浏览器端极速集成指南下一篇StaticScript 开发环境搭建指南Docker 多版本 CI 镜像如何帮你快速搞定 Node 与 LLVM 环境创作声明:本文部分内容由AI辅助生成(AIGC),仅供参考