Accelerate 文档

模型内存估算器

Hugging Face's logo
加入 Hugging Face 社区

并获得增强的文档体验

开始使用

模型内存估算器

在探索潜在模型时,一个非常困难的方面是了解多大的模型才能*装入*您当前显卡(例如将模型加载到 CUDA 上)的内存中。

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

目前我们支持搜索可在 `timm` 和 `transformers` 中使用的模型。

此 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-cased` 的 `config.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

默认情况下,它将返回所有支持的数据类型(`int4` 到 `float32`),但如果您对特定的数据类型感兴趣,可以进行筛选。

特定库

如果无法自动确定源库(如 `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

特定数据类型

如前所述,虽然我们默认返回从 `int4` 到 `float32` 的数据类型,但可以使用 `float32`、`float16`、`int8` 和 `int4` 中的任何一种。

要这样做,只需在指定 `--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

使用此计算器的注意事项

此计算器会告诉您纯粹加载模型所需的内存,而*不是*执行推理所需的内存。

这个计算结果与实际值的误差在几个百分点以内,所以它能很好地反映出实际需要多少内存。例如,加载 `bert-base-cased` 在 CUDA 上以全精度加载时实际需要 `413.68 MB`,而计算器估算为 `413.18 MB`。

根据 EleutherAI 的发现,在执行推理时,您可以预期会额外增加高达 20% 的内存。我们将进行研究以找到对这些值更准确的估算,并在完成后更新此计算器。

< > 在 GitHub 上更新