Software 2.0

估算加载模型所需的GPU 内存

admin -

在上一篇文章中，我们讨论了在 pytorch 中使用 GPU 的基础知识。在本文中，我们将理解如何估算给定模型需要多少 GPU 内存。

这将帮助我们根据模型估算使用哪个 GPU，或者根据 GPU 估算使用哪个模型。这还将帮助我们防止 OOM（cuda out of memory 错误），当 GPU RAM 上没有可用内存时会发生这种错误

1、内存不等于计算

但在我们深入细节之前，最重要的理解是内存 ≠ 计算。

我们谈论 GPU 内存，并认为计算性能来自那里。内存影响性能（因为 GPU 上的操作是内存受限的），但内存不是计算。

如今大多数 GPU 的计算性能以 T-FLOPs（万亿次浮点运算）报告，FLOPs 指的是每秒浮点运算次数。

内存指的是高带宽内存，它在 GPU RAM 上存储所有模型和数据。

这在 GPU 性能中起着重要作用，因为大多数 GPU 操作是内存受限而非计算受限的，这意味着内存操作通常是性能瓶颈。

我们今天在 colab notebook 中使用的 GPU（Nvidia — T4）报告的 FP32 计算性能为 8.1 Tflops，FP16 计算约为 65 Tflops，具有 16 GB 的 HBM（GPU RAM），内存带宽约为 320 GB/sec。

Nvidia Tesla-T4 规格：来源

2、在 GPU 上进行推理的过程涉及什么？

GPU 上的模型-数据生命周期

将模型权重移动到 GPU RAM：我们在 GPU 上加载模型，这涉及将模型权重加载到 GPU RAM 上。- 将数据移动到 GPU RAM：当我们进行批量预测时，我们将转换后的数据移动到 GPU RAM 上。- 将模型权重和数据从 RAM 移动到 SM 缓存：- 计算在 GPU 上的 SM 中发生。- 将计算数据从 SM 缓存移回 GPU RAM，然后卸载到 CPU。

3、计算机内存基础

位（bit）是计算机中的基本内存单元，可以表示二进制数（0/1）。单个字节由 4 位组成。

1 字节 = 4 位

我们习惯处理 GB 和 MB，它们基本上是

  • 1 千字节 = 1024 字节* MB（兆字节）：1024 KB，约 ~1e6 字节。
  • GB（千兆字节）：1024 MB，接近 ~1e9 字节。

4、ML 中的基本浮点数据类型

单精度：FP32

  • 用于大多数机器学习模型
  • 每个单元使用 4 字节表示。

半精度：FP16

  • 用于量化和混合精度模型
  • 每个单元用 2 字节表示。
  • 与 FP32 相比范围有限。

Bfloat16

  • 由 Google 开发，具有 FP32 的范围，但以较低的表达能力为代价。
  • 内存占用与 FP16 类似，每个单元 2 字节。

5、模型参数意味着什么？

（参数化）模型通常以参数数量定义。我们将使用最新的 Qwen-3.5 4B 模型。

这些参数是单独的权重元素的数量，每个都可以被认为是一个可调旋钮，我们在训练过程中调整它以使其很好地拟合训练数据。

现在，4B 模型意味着它大约有 40 亿个活跃参数（大约 40 亿个可调旋钮）

每个参数都有某种数值精度，这取决于所使用的模型权重的数据类型。我们将在以后的博客中讨论这个。

大多数机器学习模型通常以全精度（FP32）或 Bfloat16（BF16）训练和提供。

这意味着每个参数由 32 位整数表示，即每个参数 4 字节。

因此，模型在全精度下的内存占用为

40 亿参数 × 4 字节/参数 ~ 加载全精度模型需要 16 GB 内存。

在 BF16 中，同一模型将需要

40 亿参数 × 2 字节/参数 ~ 加载模型需要 8 GB 内存。

6、加载模型

在加载模型之前，我们可以检查 GPU 内存占用。

RAM 有空闲内存

7、将模型加载到 GPU 内存

import torch
from transformers import AutoModelForCausalLM, AutoTokenizer

device = torch.cuda.current_device()


# Set the exact model ID from Hugging Face
model_id = "Qwen/Qwen3.5-4B"

# 1. Load the Tokenizer
tokenizer = AutoTokenizer.from_pretrained(model_id)
>>> 


# 2. Load the Model
model = AutoModelForCausalLM.from_pretrained(
    model_id,
    device_map="cuda:0",
    torch_dtype="auto" 
)

现在如果我们再次检查内存占用

GPU 大约占用了 8 GB 内存。这大致对应（4B 模型，BF16 精度）→ 4e9 × 2 ~ 8e9（8GB）内存。

要使用模型进行预测，我们使用分词器对输入文本进行分词，并将其移动到 GPU

# 3. Prepare your prompt using the Chat Template
messages = [
    {"role": "system", "content": "You are a helpful AI assistant."},
    {"role": "user", "content": "Explain gpu computing in one simple paragraph."}
]

# This applies Qwen's specific formatting tags to the prompt
text = tokenizer.apply_chat_template(
    messages,
    tokenize=False,
    add_generation_prompt=True
)

# Convert the text into tensor tokens and move them to the same device as the model
model_inputs = tokenizer([text], return_tensors="pt").to(model.device)

print(model_inputs)

分词后的文本，移动到 cuda

# 4. Generate the output
generated_ids = model.generate(
    **model_inputs,
    max_new_tokens=512, 
)
预测的模型输出是输入 + 输出 token 的组合。

接下来，我们根据分词器的 id-token 映射将这些 id 转换为文本

# 5. Decode the output
# We strip the length of the input prompt from the generated IDs so it only prints the new text
generated_ids = [
    output_ids[len(input_ids):] for input_ids, output_ids in zip(model_inputs.input_ids, generated_ids)
]

response = tokenizer.batch_decode(generated_ids, skip_special_tokens=True)[0]

print(response)
最终模型输出。

谢谢，在下一篇中，我们将看看本地训练和预测模型的内存要求。

原文链接: Understands LLMs : How do I estimate GPU memory requirement for loading a model?

汇智网翻译整理，转载请标明出处