526互联

利用 device_map、torch.dtype、bitsandbytes 压缩模型参数控制使用设备

发布时间 2023-09-08 09:54:18作者: bringlu

device_map

以下内容参考 Huggingface Accelerate文档：超大模型推理方法

在 HuggingFace 中有个重要的关键字是 device_map，它可以简单控制模型层部署在哪些硬件上。

设置参数 device_map="auto"，Accelerate会自动检测在哪个设备放置模型的哪层参数（自动根据你的硬件资源分配模型参数）。其规则如下：

  • 首先充分利用GPU上的显存资源
  • 如果GPU上资源不够了，那么就将权重存储到内存
  • 如果内存还不够用了，将会使用内存映射的技术，将剩余的参数存储到硬盘上

设置参数 no_split_module_classes=["GPTJBlock"] 表示，模型中的 GPTJBlock 模块不会被切分移到不同的设备上。如果有些层包含残差操作，那么请做这样的设置。

可以通过 model.hf_device_map 来观察这个 device_map 的具体内容。

也可以自定义 device_map，也可以显示的设置每一层与设备的对应关系，然后使用如下代码加载 checkpoint。

model = load_checkpoint_and_dispatch(model, "sharded-gpt-j-6B", device_map=my_device_map)

设计 device_map

可以使用其他的 device map 映射方式，通过设置 device_map 参数（例如 "auto", "balanced", "balanced_low_0", "sequential"），或者手工设置这个 device map 字典。

读者可以操控模型在meta设备上的所有层（计算 device_map）。

当读者没有足够GPU显存来加载完整的模型（因为都会按照先占满GPU显存，再占满CPU/内存资源，最后占据硬盘的顺序来完成模型加载），上面所有的选项得到的层和设备对应结果将会相同。

当读者有足够的GPU资源来加载模型，那么上面4个选项得到的结果会有所不同。

  • "auto" 和 "balanced" 将会在所有的GPU上平衡切分模型，那么可以计算批尺寸大于 \(1\) 的输入
  • "balanced_low_0" 会在除了第一个GPU上的其它GPU上平衡划分模型，并且在第一个 GPU 上占据较少资源。这个选项符合需要在第一个 GPU 上进行额外操作的需求，例如需要在第一个 GPU 执行 generate 函数（迭代过程）。
  • "sequential" 按照GPU的顺序分配模型分片，从 GPU 0 开始，直到最后的 GPU（那么最后的 GPU 往往不会被占满，和 "balanced_low_0" 的区别就是第一个还是最后一个，以及非均衡填充）

这里 "auto" 和 "balanced" 会得到相同的结果，但是未来 "auto" 模式可能会被改变，主要是有可能发现更高效的分配策略。"balanced" 参数的功能则保持稳定。

还有需要注意的地方是，可以通过设置参数 max_memory 来限制每个GPU的显存使用数额（在函数 infer_auto_device_map() 函数设置）。当设置了参数 max_memory，需要构建一个包含 GPU 索引的字典（例如 \(0, 1\) 等），并且 cpu 这个 key 是指 CPU 离线加载的最大内存（可以设置占用内存，不然加载后电脑会卡）。字典的 value 值，可以是整数（单位是字节），也可以是字符串，例如 "10GiB" 或 "10GB"（表示最大占用10GB的显存）。

下面是一个例子。表示第 \(0\) 号和第 \(1\) 号 GPU 上最多提供 10GB 的显存、以及不超过 30GB 的内存给模型权重加载使用。

from accelerate import infer_auto_device_map

device_map = infer_auto_device_map(my_model, max_memory={0: "10GiB", 1: "10GiB", "cpu": "30GiB"})

当 PyTorch 加载模型时，他会先加载 CUDA 内核，这个就占据了 1-2GB 的显存（根据 GPU 的不同会略有区别）。因此能够使用的 GPU 显存要小于实际标定显存。可以使用代码 torch.ones(1).cuda() 来看看你的 GPU 上的 CUDA kernel 占用显存大小。

因此可以通过待 max_memory 参数的存储空间映射，来防止 out-of-memory 错误出现。

另外，如果有这样的需求，一些操作的输出需要在GPU上运行（例如 generate 函数来生成文本），或者将输入放到 GPU 上，那么在某个 GPU 上需要留一些显存（Accelerate 会将输出返回作为一些设备的输入）。又如果需要优化最大的批尺寸以及有很多的 GPU，那么可以在第一个的 GPU 留下足够的显存。例如在 8x80 A100 GPU 上运行 BLOOM-176B 的理想内存设置如下：

max_memory = {0: "30GIB", 1: "46GIB", 2: "46GIB", 3: "46GIB", 4: "46GIB", 5: "46GIB", 6: "46GIB", 7: "46GIB"}

除了 GPU 0 上留了足够的显存，在其他 \(7\) 个 GPU 上也留了将近 \(50\%\) 的显存。

如果自定义 device_map，那么字典中的 key 必须是模型中的模块名，且 value 是合法的设备索引（例如对于 GPU，是从 \(0\) 开始的整数）或 "cpu"（内存）或 "disk"（硬盘）。并且 key 要覆盖所有的模型模块名，满足这些规则，你可以随心所欲的定义你的 device map。例如你的模型有两个模块（就叫 block1 和 block2 吧），每个模块包含三个线性层（称为 linear1，linear2 和 linear3 吧，还真是有顺序）。一个合法的 device map 如下：

device_map = {"block1": 0, "block2": 1}

另外一个合法的 device map 如下：

device_map = {"block1": 0, "block2.linear1": 0, "block2.linear2": 1, "block2.linear3": 1}

相反，下面的device map就不是合法，因为它的key没有覆盖模型的所有模块名。

device_map = {"block1": 0, "block2.linear1": 1, "block2.linear2": 1}

为了提升效率，最好是你的 device map 是按照 GPU 顺序，序贯的配置参数（例如不要在 GPU 0 上加载第一个全在，在 GPU 1 上再加载，然后权重又回到 GPU 0，交叉来回会降低模型推理速度）。主要是减少数据的 GPU 之间切换次数，提高效率。

一些注意事项

  • infer_auto_device_map() （或在load_checkpoint_and_dispatch()中设置device_map="auto"）会最大化使用 GPU 显存和 CPU 内存。当 PyTorch 会高效的管理 GPU 显存（并且会释放不适用的显存），但是管理内存不怎么高效。因此，使用 auto 模式，在 CPU 内存管理上不是很高效。一个好的解决方法，就是将部分在内存上的模块移到硬盘上。
  • infer_auto_device_map() （或在load_checkpoint_and_dispatch()中设置device_map="auto"）是按照 GPU、CPU 和硬盘的顺序分配模型模块（防止循环操作），因此如果你的第一个层需要的 GPU 显存空间大于 GPU 显存时，有可能在 CPU/硬盘上出先奇怪的东西（第一个层不要太大，不然会发生奇怪的事情）。
  • load_checkpoint_and_dispatch() 和 load_checkpoint_in_model() 不会对 checkpoint 中的 state_dict 和模型进行检测（后面会修正这个），如果加载的 checkpoint 和 model 不匹配，又会发生奇怪的事情（报错啦！）。
  • 没有做模型并行，意思是你的模型被分割到多个 GPU ，然后被顺序执行。换言之，一次只有一个 GPU 在运行，然后等待其他 GPU 运行完毕。

torch_dtype

这个形参可以设置模型中全部 Linear 层的数据格式，可以使用如下把 \(32\) 位的模型线性层参数转换为 \(16\) 位的模型参数

model = AutoModelForCausalLM.from_pretrained("./Llama-2-7b-hf", torch_dtype=torch.float16)

但是除线性层之外的所有参数仍为 \(32\) 位浮点数。

bitsandbytes (load_in_8bit / load_in_4bit)

主要参考 HuggingFace 的文档 Quantize ? Transformers models

这个是 bitsandbytes 包中特有的内容，必须安装了这个包才能使用这个形参。

具体作用是使用 \(8\) 位或 \(4\) 位来读入模型参数。具体算法可以参考 LLM.int8() 这篇论文。

自 0.39.0 版发布以来，可以利用 FP4 数据类型，使用 \(4\) 位量化加载任何支持 device_map 的模型。

如果你想量化自己的 PyTorch 模型，请查看 Accelerate 库的文档。

以下是使用 bitsandbytes 集成可以做的事情：

正常用法

如果你的模型支持使用 HF 的 Accelerate 导入，并包含线性层，那么可以在调用 from_pretrained() 函数时使用 load_in_8bit 或者 load_in_4bit。

from transformers import AutoModelForCausalLM

model_8bit = AutoModelForCausalLM.from_pretrained("facebook/opt-350m", load_in_8bit=True)
model_4bit = AutoModelForCausalLM.from_pretrained("facebook/opt-350m", load_in_4bit=True)

默认情况下其他的模块（例如 torch.nn.LayerNorm）会被转化为 torch.float16，但是其实你也可以使用上文中提及的 torch_dtype 强行改成 \(32\) 位。

import torch
from transformers import AutoModelForCausalLM

model_8bit = AutoModelForCausalLM.from_pretrained("facebook/opt-350m", load_in_8bit=True, torch_dtype=torch.float32)
model_8bit.model.decoder.layers[-1].final_layer_norm.weight.dtype

可以检查你的模型的内存占用量 footprint 使用 get_memory_footprint 方法.

print(model.get_memory_footprint())

4/8 位浮点数

需要三个最新的包：bitsandbytes、accelerate、transformers。

目前无法将 \(4\) 位的模型上传到 HF 的 Hub 里，\(8\) 位的模型如果使用最新的包可以上传到 Hub 中。训练 \(4/8\) 位的模型目前仍不支持。用 \(4/8\) 可以训练额外的参数。（截止至 2023-9-7）