Accelerate 文档

模型内存估算器

Hugging Face's logo
加入 Hugging Face 社区

并获得增强的文档体验

开始使用

模型内存估算器

当探索可能在你的机器上使用的模型时,一个非常困难的方面是知道多大的模型适应你当前显卡的内存(例如将模型加载到 CUDA 上)。

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

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

此 API 会将模型加载到 meta 设备上的内存中,因此我们实际上并没有下载和加载模型的完整权重到内存中,也不需要这样做。因此,测量 80 亿参数模型(或更多)是完全没问题的,而无需担心你的 CPU 是否能够处理它!

Gradio 演示

以下是一些与上述内容相关的 gradio 演示。第一个是官方 Hugging Face 内存估算 Space,直接使用 Accelerate

一位社区成员采用了这个想法并进一步扩展了它,允许你直接过滤模型,并查看在给定的 GPU 约束和 LoRA 配置下,你是否可以运行特定的 LLM。要体验它,请查看这里了解更多详情。

命令

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

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

accelerate estimate-memory bert-base-cased

这将下载 bert-based-casedconfig.json,将模型加载到 meta 设备上,并报告它将使用的空间量

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

dtype 最大层 总大小 使用 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

默认情况下,它将返回所有支持的 dtype (int4float32),但如果你对特定的 dtype 感兴趣,可以进行过滤。

特定库

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

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

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

dtype 最大层 总大小 使用 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 的内存使用情况

dtype 最大层 总大小 使用 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 的任何 dtype。

要这样做,在指定 --dtypes 后传入它们

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

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

dtype 最大层 总大小 使用 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 上更新