超大模型加载转换Trick

文章预览

原文：https://zhuanlan.zhihu.com/p/698950172 在深度学习领域，大模型的训练和推理通常需要消耗大量的计算和内存。如何高效地加载和使用大模型是一个相当关键的问题。在这篇博客中，我将分享一些关于更快加载大模型和减少内存的技巧. 问题分析 假设现在我们有一个236B 超大模型的原始权重的  checkpoint.pth  文件, 比如 DeepSeek Chat V2, 以BF16 格式存储, 一个标准的加载流程如下 import torch state_dict = torch . load(checkpoint_file) my_model = BigModelClass( ... ) my_model . load_state_dict(state_dict) 在这段代码的中,  my_model = BigModelClass(...)  会初始化一个模型,  torch.load(checkpoint_file) 函数会将模型权重从磁盘加载到内存中。然后， my_model.load_state_dict(state_dict) 函数会将权重从内存加载到模型的参数中。这两个步骤都可能会消耗大量的时间和内存。理想情况下, 一个236B BF16格式的 ………………………………