新媒网 · 跨境数据社区

当大语言模型（LLM）的浪潮席卷全球，深耕跨境领域的你，是否也在为那些动辄千亿参数的“大家伙”如何高效部署、顺畅推理而头疼？以备受关注的Llama 3 70B为例，光是模型本身的半精度（FP16）权重就需要大约140 GB的显存，而更高阶的Llama 4 Scout 109B更是高达218 GB。更别提在推理过程中，为了支持超长上下文窗口和多用户并发，还需要巨大的键值缓存（KV Cache）。比如，一个支持128k上下文窗口、单用户（批处理大小为1）的Llama 3 70B模型，其KV缓存就可能占用约40 GB显存，并且会随着用户数量线性增长。在实际生产环境中，如果试图将如此庞大的模型完全载入GPU显存，几乎百分之百会遇到“显存不足”（OOM）的窘境。新媒网跨境获悉，这正是许多技术团队在规模化部署大模型时面临的“拦路虎”。

传统的GPU和CPU内存是相互独立的，数据在两者之间传输需要耗费大量时间和资源。然而，英伟达（NVIDIA）的Grace Blackwell和Grace Hopper架构彻底改变了这一局面。它们通过NVLink-C2C这种高速互联技术，构建了一个高达900 GB/s的内存一致性连接。这比PCIe Gen 5的带宽还要快7倍！正是NVLink-C2C带来的内存一致性，让CPU和GPU能够共享一个统一的内存地址空间（如图1所示）。这意味着它们可以无缝访问和操作同一份数据，不再需要耗时的数据显式传输或冗余内存复制。这种创新设计，让即便是超大规模的数据集和模型，也能在GPU和CPU之间高效流动，轻松突破传统GPU显存的局限。

图1. NVLink-C2C内存一致性与地址转换服务

得益于NVLink-C2C的高带宽连接和统一内存架构，英伟达Grace Hopper和Grace Blackwell平台在LLM微调、KV缓存卸载、推理以及科学计算等领域展现出惊人的效率提升。当GPU自身显存不足时，模型能够快速地将数据转移并利用CPU内存，确保各项任务顺利进行。

举个例子，在搭载统一内存架构的英伟达GH200 Grace Hopper 超级芯片平台®上，Llama 3 70B模型不仅可以使用芯片自带的96 GB高带宽GPU显存，还能无缝调用连接到CPU的480 GB LPDDR内存，整个过程无需任何显式的数据传输。这大大扩展了可用内存总量，使得处理那些原本对单一GPU来说过于庞大的模型和数据集成为可能。

实战演练：告别显存不足的烦恼

接下来，我们通过一个实战演练，看看如何用Llama 3 70B模型和GH200超级芯片，来展示统一内存的巨大威力。

1. 前期准备工作

要开始实操，我们首先需要配置好环境，并获取Llama 3 70B模型的访问权限。请注意，以下代码示例专为英伟达Grace Hopper GH200超级芯片平台设计，以充分展示其统一内存架构的优势。当然，这些技术同样适用于英伟达Grace Blackwell系列系统。

步骤很简单，请跟着我一步步操作：

• 申请模型访问权限： 前往知名AI模型社区平台Hugging Face的Llama 3 70B模型页面，提交访问申请。
• 生成访问令牌： 申请获批后，在你的Hugging Face账户设置中生成一个访问令牌。这个令牌将用于程序化访问模型的身份验证。
• 安装必要的软件包： 在你可以与模型交互之前，请确保安装了必需的Python库。在GH200机器上打开Jupyter Notebook，运行以下命令：

#安装huggingface和cuda相关包
!pip install --upgrade huggingface_hub
!pip install transformers
!pip install nvidia-cuda-runtime-cu12

• 登录Hugging Face： 安装完软件包后，使用你生成的令牌登录Hugging Face。huggingface_hub库提供了一个便捷的方法：

#使用生成的令牌登录Hugging Face
from huggingface_hub import login
login("输入你的令牌")

2. 挑战：Llama 3 70B模型在GH200上会遇到什么？

当你尝试将Llama 3 70B模型载入GPU显存时，其参数（权重）会被加载到英伟达CUDA显存中。由于Llama 3 70B的半精度（FP16）权重需要大约140 GB的显存，而GH200提供的GPU显存仅为96 GB，因此模型无法完全载入可用显存，加载过程会因“显存不足”（OOM）错误而失败。

在下面的代码单元格中，我们将演示这种现象：

import torch
from transformers import pipeline
pipe = pipeline("text-generation", model="meta-llama/Llama-3.1-70B") #尝试将模型加载到GPU内存中

运行上述命令后，你将看到类似下面的错误信息：

Error message: OutOfMemoryError: CUDA out of memory. Tried to allocate 896.00 MiB. GPU 0 has a total capacity of 95.00 GiB of which 524.06 MiB is free. Including non-PyTorch memory, this process has 86.45 GiB memory in use. Of the allocated memory 85.92 GiB is allocated by PyTorch, and 448.00 KiB is reserved by PyTorch but unallocated. If reserved but unallocated memory is large try setting PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True to avoid fragmentation. See documentation for Memory Management.

从错误信息中，我们可以清楚地看到GPU显存已经满了。你也可以通过运行!nvidia-smi命令来确认GPU的显存状态：

!nvidia-smi

运行该命令后，你应该会得到一个类似于下图的输出。该输出清楚地告诉我们，在总计97.871 GB的GPU显存中，我们已经使用了96.746 GB。

图2. nvidia-smi命令输出

为了准备下一步操作并释放GPU的内存，我们需要清除上次失败尝试留下的任何变量。在下面的命令中，请将<PID>替换为你的Python进程ID，你可以通过!nvidia-smi命令找到它。

!kill -9 <PID>

3. 突破：如何解决“显存不足”错误？

要解决这个“显存不足”的问题，我们可以利用“托管内存分配”技术，它允许GPU除了自身显存之外，还能访问CPU内存。在GH200系统上，统一内存架构使得CPU（高达480 GB）和GPU（高达144 GB）能够共享一个单一的地址空间，并透明地访问彼此的内存。通过配置RAPIDS Memory Manager (RMM) 库来使用托管内存，开发者可以分配一个同时可被GPU和CPU访问的内存区域，从而使工作负载能够突破物理GPU显存的限制，而无需手动进行数据传输。RMM库就像一个智能调度员，帮你高效管理并利用整个系统的内存资源。

import rmm
import torch
from rmm.allocators.torch import rmm_torch_allocator
from transformers import pipeline

rmm.reinitialize(managed_memory=True) #启用对CPU内存的访问
torch.cuda.memory.change_current_allocator(rmm_torch_allocator) #指示PyTorch使用RMM内存管理器来管理所有内存分配，实现统一内存

pipe = pipeline("text-generation", model="meta-llama/Llama-3.1-70B")

再次运行模型加载命令，你将惊喜地发现，由于现在我们拥有了更大的内存空间，再也不会遇到“显存不足”的错误了！现在，你可以愉快地向LLM发送提示并接收回复：

pipe("世界上最高的山峰是哪座？")

总结与展望

新媒网跨境认为，随着大语言模型规模的持续膨胀，如何将这些“巨无霸”高效加载到GPU上进行推理，已经成为一个核心挑战。通过本次实战，我们深入探讨了统一内存架构如何有效解决这一难题。它使得CPU和GPU的内存可以无缝共享，消除了显式数据传输的必要，极大地简化了在现代硬件上处理最前沿LLM的复杂性。

要深入了解如何管理CPU和GPU内存，建议查阅RAPIDS Memory Manager的官方文档。掌握这项技术，将助你更好地驾驭大模型，开启跨境业务的更多可能！

新媒网（公号: 新媒网跨境发布），是一个专业的跨境电商、游戏、支付、贸易和广告社区平台，为百万跨境人传递最新的海外淘金精准资讯情报。

本文来源：新媒网 https://nmedialink.com/posts/gh200-fix-llm-oom-in-1-min.html