【PyTorch 奇淫技巧】Async Checkpoint Save

文章预览

↑ 点击 蓝字  关注极市平台 作者丨GiantPandaCV 来源丨GiantPandaCV 编辑丨极市平台 极市导读   本文介绍了PyTorch 2.4的新特性——异步Checkpoint保存功能，该功能通过将模型状态的保存过程移到CPU线程来减少训练中断，大幅提升了大型模型训练时Checkpoint保存的速度。   >> 加入极市CV技术交流群，走在计算机视觉的最前沿 PyTorch Async Checkpoint Save PyTorch博客资料：https://pytorch.org/blog/reducing-checkpointing-times/ PyTorch实现和使用Demo：https://github.com/pytorch/pytorch/blob/main/torch/distributed/checkpoint/state_dict_saver.py 功能介绍 在PyTorch 2.4之后，我们可以尝试使用PyTorch开发的异步Checkpoint保存功能，这个功能是和IBM联合开发的，在7B的大模型训练中，Checkpoint保存的时间从平均 148.8 秒缩短至 6.3 秒，快了 23.62 倍。这可以转化为以下两种好处： 在继续鲁棒的保存Checkpoint的同时，在每给 ………………………………