在 Sagemaker 端点上部署 LLM - CUDA 内存不足

Question

我正在尝试使用自定义脚本将 LLM 部署到 Sagemaker Endpoint，但遇到错误： CUDA 内存不足。尝试分配 20.00 MiB。 GPU 1 的总容量为 22.20 GiB，其中 13.12 MiB 是免费的。进程 13234 有 2.25 GiB 内存正在使用。进程 13238 有 3.82 GiB 内存正在使用。进程 13236 正在使用 8.06 GiB 内存。进程 13239 正在使用 8.06 GiB 内存。在已分配的内存中，PyTorch 分配了 6.93 GiB，PyTorch 保留了 49.59 MiB，但未分配。如果保留但未分配的内存很大，请尝试设置 max_split_size_mb 以避免碎片。请参阅内存管理和 PYTORCH_CUDA_ALLOC_CONF 的文档：400

**问题是我能够得到几个 ping 的响应，然后我开始收到此错误。我正在清除预测函数中的缓存，但我无法理解 3 件事：

1. 为什么在端点多次成功响应后我收到此错误？
2. 这四个进程占用空间的原因是什么？
3. 很明显，我该如何解决这个错误？**

Answer 1

Cuda 内存不足

• 更新 PYTORCH_CUDA_ALLOC_CONF max_split_size_mb。但这可能是最没有帮助的选择。 os.environ["PYTORCH_CUDA_ALLOC_CONF"] = "max_split_size_mb:在此处输入大小"

• 另一个论坛的一位用户提到需要安装以下软件包：transformers==4.28.1、sentpiece==0.1.97、accelerate==0.18.0、bitsandbytes==0.37.2 和 torch 1.13.1 。但我认为内存不足问题主要需要通过内存管理来处理。新版本中的软件包问题可能是暂时的，但肯定会得到解决。

• 仅与训练相关 - 在训练视觉模型时，图像可能无法仅适合 GPU，因此您应该调整其大小并从 GPU 内存中释放它们。

• 仅与训练相关 - 将训练批量大小减少至 1

• 垃圾收集 gc.collect()

• 清空缓存 torch.cuda.empty_cache()

• 增加系统 RAM/更大的计算实例

真正对我有帮助的是通过定义可用于存储模型的GPU的max_内存来将LLM分布在GPU上。这意味着我的 GPU 没有被 LLM 完全预订。这是一个三步过程：

• 在 GPU 上加载没有权重/空模型的模型。当推理集 no_grad 以避免任何计算权重更新，即使没有 权重将被更新。另外，设置设备映射以修复最大内存 加载模型权重可以采取。

• 在CPU上加载模型权重

• 将各层的权重加载到GPU，执行/计算 完成后，权重将从 GPU 中移除。

 使用 torch.no_grad()：
    
            使用 init_empty_weights()：
                old_prediction_model = AutoModelForCausalLM.from_pretrained(
                    模型目录，
                    torch_dtype=torch.bfloat16,
                    量化配置=量化配置
                    ）
            模型 = load_checkpoint_and_dispatch(
            old_prediction_model, offload_folder="/offload_folder_name_or_location",
                checkpoint=model_dir, device_map=infer_auto_device_map(old_prediction_model, max_memory={0: "10GiB"}),
                dtype=torch.bfloat16
        ）

注意：除此之外，内存不足错误的另一个重要原因是将变量留在 GPU 上，即因为执行是在 GPU 上进行的，并且在此过程中您尝试创建模型推理或评估的列表，GPU当你做出新的推论时，记忆会继续被填满。为了避免这种情况，每次推理时都将变量从 GPU 移至 CPU 内存