加速文档

模型内存估算器

Hugging Face's logo
加入 Hugging Face 社区

并获得增强型文档体验

开始使用

模型内存估算器

在探索可用于您机器的潜在模型时,一个非常困难的方面是了解您当前的显卡可以容纳多大的模型(例如将模型加载到 CUDA 上)。

为了帮助解决这个问题,Accelerate 通过 accelerate estimate-memory 提供了一个 CLI 界面。本教程将帮助您逐步了解如何使用它,以及可以期待什么,最后会链接到托管在 Hub 上的交互式演示,您甚至可以将这些结果直接发布到模型仓库中!

目前,我们支持搜索可以在 timmtransformers 中使用的模型。

此 API 将在 meta 设备上将模型加载到内存中,因此我们实际上并没有将模型的全部权重下载并加载到内存中,也不需要这样做。因此,我们可以轻松地测量 80 亿参数模型(或更多),而无需担心您的 CPU 是否能够处理!

Gradio 演示

以下是与上面描述内容相关的几个 Gradio 演示。第一个是官方的 Hugging Face 内存估算空间,直接使用 Accelerate

社区成员进一步发展了这个想法,允许您直接过滤模型,并查看在给定 GPU 限制和 LoRA 配置的情况下是否可以运行特定的 LLM。要体验它,请访问 此处 了解更多详细信息。

命令

使用 accelerate estimate-memory 时,您需要传入要使用的模型的名称,可能还需要传入该模型使用的框架(如果无法自动找到),以及您希望模型以何种数据类型加载。

例如,以下是如何计算 bert-base-cased 的内存占用。

accelerate estimate-memory bert-base-cased

这将下载 bert-based-casedconfig.json,在 meta 设备上加载模型,并报告它将使用多少空间。

加载 bert-base-cased 的内存使用情况

数据类型 最大层 总大小 使用 Adam 进行训练
float32 84.95 MB 418.18 MB 1.61 GB
float16 42.47 MB 206.59 MB 826.36 MB
int8 21.24 MB 103.29 MB 413.18 MB
int4 10.62 MB 51.65 MB 206.59 MB

默认情况下,它将返回所有支持的数据类型(int4float32),但如果您只对特定数据类型感兴趣,可以对其进行过滤。

特定库

如果无法自动确定源库(例如在 bert-base-cased 的情况下),则可以传入库名称。

accelerate estimate-memory HuggingFaceM4/idefics-80b-instruct --library_name transformers

加载 HuggingFaceM4/idefics-80b-instruct 的内存使用情况

数据类型 最大层 总大小 使用 Adam 进行训练
float32 3.02 GB 297.12 GB 1.16 TB
float16 1.51 GB 148.56 GB 594.24 GB
int8 772.52 MB 74.28 GB 297.12 GB
int4 386.26 MB 37.14 GB 148.56 GB
accelerate estimate-memory timm/resnet50.a1_in1k --library_name timm

加载 timm/resnet50.a1_in1k 的内存使用情况

数据类型 最大层 总大小 使用 Adam 进行训练
float32 9.0 MB 97.7 MB 390.78 MB
float16 4.5 MB 48.85 MB 195.39 MB
int8 2.25 MB 24.42 MB 97.7 MB
int4 1.12 MB 12.21 MB 48.85 MB

特定数据类型

如前所述,虽然我们默认返回 int4float32,但可以使用任何数据类型,包括 float32float16int8int4

为此,在指定 --dtypes 后传入它们。

accelerate estimate-memory bert-base-cased --dtypes float32 float16

加载 bert-base-cased 的内存使用情况

数据类型 最大层 总大小 使用 Adam 进行训练
float32 84.95 MB 413.18 MB 1.61 GB
float16 42.47 MB 206.59 MB 826.36 MB

此计算器的注意事项

此计算器将告诉您纯粹加载模型所需的内存量,包括执行推理。

此计算与实际值误差在几个百分比之内,因此它非常准确地反映了实际所需的内存量。例如,在 CUDA 上以全精度加载 bert-base-cased 实际上需要 413.68 MB,而计算器估计为 413.18 MB

执行推理时,您预计会增加最多 20%,如 EleutherAI 所述。我们将研究如何更准确地估计这些值,并在完成研究后更新此计算器。

< > 在 GitHub 上更新