大模型参数量和占用显存

Question

由于设备条件有限，在调用大模型的时候有一个困惑，就是模型的参数量和我的显存有什么关系，比如使用ollama调用模型时，一个4b的模型要占用多少显存？只能通过实际调用才能知道吗？还是有计算公式？

Answer 1

阿里嘎多学长整理AIGC生成，因移动端显示问题导致当前答案未能完全显示，请使用PC端查看更加详细的解答过程

大模型参数量和占用显存的关系

大模型的参数量和占用显存的关系主要取决于模型的架构和参数类型。一般来说，模型的参数量越大，占用显存的量也越大。

计算公式

虽然没有一个通用的公式可以直接计算出模型的参数量和占用显存的关系，但我们可以根据模型的架构和参数类型进行一些估算。

• 参数量: 模型的参数量通常取决于模型的层数、神经元数量和激活函数类型。例如，一个常见的LSTM模型，其参数量可以估算为：

  参数量 = (层数 * 神经元数量 * (输入维度 + 1)) + (层数 * 神经元数量 * 神经元数量)
  

• 占用显存: 占用显存的量取决于模型的参数类型和数据类型。例如，浮点数参数通常占用更多显存空间，而整数参数则占用较少。

  占用显存 = 参数量 * 数据类型的字节数
  

  例如，一个浮点数参数占用4字节的显存空间，而一个整数参数占用4字节的显存空间。

实际调用

虽然有计算公式，但实际调用时仍然需要通过实际调用来确定模型的参数量和占用显存的量。因为不同的模型架构和参数类型可能会有不同的计算公式。

示例代码

以下是使用PyTorch和TensorFlow来计算模型参数量和占用显存的示例代码：

import torch
import tensorflow as tf

# PyTorch示例
model = torch.nn.LSTM(input_size=10, hidden_size=20, num_layers=2)
print("参数量:", sum(p.numel() for p in model.parameters()))
print("占用显存:", model.parameters().__sizeof__() / (1024 * 1024), "MB")

# TensorFlow示例
model = tf.keras.models.Sequential([
    tf.keras.layers.LSTM(20, input_shape=(10, 10)),
    tf.keras.layers.Dense(10)
])
print("参数量:", sum(p.size for p in model.trainable_variables))
print("占用显存:", model.trainable_variables().__sizeof__() / (1024 * 1024), "MB")

以上示例代码计算了LSTM模型的参数量和占用显存的量。