大模型从CPU迁移至CUDA耗时久？训练多次恢复模型优化方案问询

针对大模型多次恢复到原始状态的场景，以下几个方案可以有效降低耗时：

1. 提前在GPU保留原始模型副本（优先推荐，GPU内存足够时） #

如果GPU内存能同时容纳原始模型和后续的小模型，直接在GPU上预存一份转换好的原始模型，每次恢复时直接在GPU内操作，避免跨设备数据传输的巨大开销：

# 初始化时一次性完成GPU转换，只做一次
self.gpu_origin_model = origin_model.to(device=device, dtype=torch.bfloat16)
# 后续每次恢复时，直接在GPU内复制模型
self.model = deepcopy(self.gpu_origin_model)
# 无需再调用torch.cuda.empty_cache()，PyTorch会自动管理显存

这样每次恢复的耗时会从10秒级降到毫秒级，完全避免了CPU到GPU的跨设备传输。

2. 移除不必要的torch.cuda.empty_cache()调用 #

torch.cuda.empty_cache()会强制扫描并释放未使用的显存，本身有1秒左右的开销，而PyTorch的显存管理器会自动回收未被引用的显存。只要没有出现显存碎片导致的OOM问题，完全可以去掉该调用，直接简化流程：

# CPU存原始模型时的优化版
del self.model
self.model = copy.deepcopy(origin_model).to(device=device, dtype=torch.bfloat16)
# 去掉empty_cache()，节省1秒开销

3. 预存GPU版模型的状态字典，原地加载 #

提前将原始模型转换到GPU并保存其状态字典，后续恢复时直接加载到现有模型实例，无需重新创建模型和跨设备转换：

# 初始化时预存GPU版状态字典
self.gpu_origin_state = origin_model.to(device=device, dtype=torch.bfloat16).state_dict()
# 后续每次恢复时，直接加载状态字典
self.model.load_state_dict(self.gpu_origin_state)

这种方式比重新创建模型并转换设备快得多，因为只是替换模型参数，无需重新初始化模型结构。

4. 预保存GPU版模型文件，快速加载 #

如果需要持久化保存，可以提前将转换好的GPU版模型保存为文件，后续加载时直接指定设备，避免转换步骤：

# 仅执行一次：预转换并保存GPU版模型
gpu_origin_model = origin_model.to(device=device, dtype=torch.bfloat16)
torch.save(gpu_origin_model.state_dict(), "gpu_origin_model.pt")

# 后续每次恢复时
del self.model
self.model = origin_model.__class__()  # 先创建空模型结构
self.model.load_state_dict(torch.load("gpu_origin_model.pt", map_location=device))
self.model.to(device=device, dtype=torch.bfloat16)  # 此处to操作几乎无开销，参数已在GPU上

加载GPU预存的状态字典时，无需再做大量数据的设备转换，耗时会大幅降低。

5. GPU内存不足时的折中方案 #

如果GPU内存无法同时容纳原始模型和小模型，可以使用懒加载+内存映射减少加载开销：

# 加载时使用内存映射，避免一次性加载全部数据
self.model = torch.load("gpu_origin_model.pt", map_location=device, mmap_mode='r')

这种方式会将模型参数以内存映射的方式加载，减少内存占用的同时，也能加快加载速度。

内容的提问来源于stack exchange，提问作者ZhiWei Li