赞
踩
模型卸载是优化大型模型训练和使用的一种常见技术。这种技术允许我们将模型参数或计算过程的部分内容从主内存(通常是显存)"卸载"到其他存储设备(如CPU的RAM),以节省主内存空间并提高训练效率。
例如,假设我们有一台配备了RTX2080Ti显卡的计算机,拥有8GB的显存和32GB的RAM。如果我们需要在GPU上加载一个需要10GB显存的7B LLaMA模型,就会出现内存溢出的错误。此时,我们可以利用模型卸载技术,将最多6GB的参数加载到8GB的显存上,然后将剩余的参数保存在32GB的RAM中,使模型能够正常运行。除了RAM,我们也可以选择将模型卸载到SSD硬盘中。
我们知道,模型中的大量矩阵运算是在GPU中完成的,而GPU在显存中操作这些数据。然而,RAM的价格远低于显存,而SSD的价格又远低于RAM。因此,模型卸载是一种非常有效的成本降低技术。
模型卸载的步骤通常包括将模型分割成多个部分,每个部分可以单独加载到GPU的显存中进行计算,而其他部分则保留在其他设备上,需要时再从其他存储设备上加载。
以CPU作为卸载空间为例,模型卸载工具会将模型参数、状态和梯度加载到CPU中。在前向传播过程中,它会一次将一层(或多层)加载到GPU上进行训练,然后在CPU上激活,再复制到GPU上进行反向传播。完成后,所有的参数都会使用CPU中存储的梯度进行更新。
FairScale的模型卸载方法位于其experimental包中,可以通过导入该包来使用。
from fairscale.experimental.nn.offload import OffloadModel
这个包中
Copyright © 2003-2013 www.wpsshop.cn 版权所有,并保留所有权利。