【大模型运行漫长的开始】 关于多GPU使用 device_map

我不听！直接用！赞啦！

device_map = “balanced_low_0”

直接使用from_pretrained(model_path, device_map = "balanced_low_0")，就能成功运行！

详细版本（苦难的开始）

1、device_map的使用

1.1 device_map是什么？

它是一个用于指定模型中各个部分所在设备的映射表，它可以简单控制模型层部署在哪些硬件上。

在深度学习中，模型通常由多个层和参数组成，这些层和参数可以在不同的设备上进行计算和存储。device_map的作用就是告诉模型的各个部分应该在哪个设备上进行计算和存储。

1.2 device_map使用方式：简单版

self.engine = AutoModelForSeq2SeqLM.from_pretrained(config.model_path, device_map = "auto", max_memory={0: "44GiB", 1: "44GiB", 2: "44GiB", 3: "44GiB",4: "44GiB",5: "44GiB"})

device_map 参数有"auto", “balanced”, “balanced_low_0”, “sequential”

• “auto” 和 “balanced” 将会在所有的GPU上平衡切分模型。主要是有可能发现更高效的分配策略。“balanced” 参数的功能则保持稳定。（个人不推荐使用）
• “balanced_low_0” 会在除了第一个GPU上的其它GPU上平衡划分模型，并且在第一个 GPU 上占据较少资源。这个选项符合需要在第一个 GPU 上进行额外操作的需求，例如需要在第一个 GPU 执行 generate 函数（迭代过程）。（墙裂推荐使用）
• “sequential” 按照GPU的顺序分配模型分片，从 GPU 0 开始，直到最后的 GPU（那么最后的 GPU 往往不会被占满，和 “balanced_low_0” 的区别就是第一个还是最后一个，以及非均衡填充），但是我在实际使用当中GPU 0 会直接爆显存了（直接别用了）

1.3 device_map使用方式：手工版

如果你有特殊的需求，需要使用指定的显卡那就需要设置这个需要device map 字典。

例如，假设我们有一个模型，其中包含了一个卷积层和一个全连接层，同时还有一些参数。具体的device_map可以如下所示：

device_map = {'shared': 1, # 在GPU 1
'decoder.embed_tokens': 1, # 在GPU 1
'encoder.embed_tokens': 1, # 在GPU 1
'encoder.block.0': 1, # 在GPU 1
'encoder.block.1': 1, # 在GPU 1
'encoder.block.2': 1, # 在GPU 1
.....
'encoder.final_layer_norm': 3, # 在GPU 3
'encoder.dropout': 3,  # 在GPU 3
'decoder.block.0': 3,  # 在GPU 3
'decoder.block.1': 4,  # 在GPU 4
'decoder.block.2': 4, 
.....
}

在加载模型时，我们可以使用device_map来指定各个部分所在的设备。

在模型的前向传播和反向传播过程中，每个部分都会被正确地分配到指定的设备上进行计算。实际使用时，需要根据具体的硬件和框架来指定正确的设备名称。

1.4 手工版的注意事项

怎么固定GPU的显存？

使用下面的例子，就可以固定现存和显卡数量来运行模型

max_memory={0: "44GiB", 1: "44GiB", 2: "44GiB", 3: "44GiB",4: "44GiB",5: "44GiB"}
self.engine = AutoModelForSeq2SeqLM.from_pretrained(model, device_map = "auto", max_memory=max_memory)

怎么查看模型的device_map构成？

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

如果device_map为None，这意味着没有指定设备映射，因此代码默认在单一设备上（可能是主GPU或CPU）执行模型。可以先device_map = "auto"来查看。

为什么手工写的device_map会报错？

如果报错RuntimeError: Expected all tensors to be on the same device, but found at least two devices, cuda:1 and cuda:0! (when checking argument for argument mat2 in method wrapper_mm)

原因是当进行涉及多个张量的操作时，如加法、乘法或更复杂的神经网络层计算，所有这些张量都必须在同一设备上。

依旧不太明白的可以看一下这个

2、Accelerate载入大模型（超慢）

2.1 Accelerate是啥

Accelerate可以在无需大幅修改代码的情况下完成并行化。同时还支持DeepSpeed的多种ZeRO策略，简直不要太爽。

代码效率高，支持无论是单机单卡还是多机多卡适配同一套代码。

允许在单个实例上训练更大的数据集：Accelerate 还可以使 DataLoaders 更高效。这是通过自定义采样器实现的，它可以在训练期间自动将部分批次发送到不同的设备，从而允许每个设备只需要储存数据的一部分，而不是一次将数据复制四份存入内存。

支持DeepSpeed：无需更改代码，只用配置文件即可对DeepSpeed开箱即用。

2.2 使用方式

1. 使用load_checkpoint_and_dispatch()加载检查点到空模型。
2. 设置device_map="auto"会让Accelerate自动判断把模型的每个层放在哪里:
3. 通过no_split_module_classes参数可以指定某些层不被分割（比如包含残差连接的Block等模块）

cuda_list = '0,1'.split(',')
memory = '10GiB'
max_memory = {int(cuda):memory for cuda in cuda_list}
no_split_modules = 你的huggingface模型._no_split_modules
# 第一步 下面三行代码可以知道每个权重的形状大小 然后就可估算出加载预训练权重需要的内存大小
# 因此下步为了适配有限内存和显存 我们需要合理的切分模型 进行CPU和GPU计算操作
engine_config = AutoModelForSeq2SeqLM.from_pretrained(model_name)  # multi GPU
with init_empty_weights():
    engine = AutoModelForSeq2SeqLM.from_config(engine_config)  # 加载到meta设备中，不需要耗时，不需要消耗内存和显存

# 第二步 自动划分每个层的设备 会尽可能的使用GPU和CPU 需要手动添加offload到磁盘 
device_map = infer_auto_device_map(self.engine, max_memory=max_memory)  # 自动划分每个层的设备

# 第三步 当前的模型并行时简单的顺序执行 同一时刻只有一个在执行 不对模型的参数进行检测
load_checkpoint_and_dispatch(engine, model_name, device_map="balanced_low_0", no_split_module_classes=no_split_module_classes, offload_folder=None, dtype='float16', offload_state_dict=True)# 加载权重
self.engine = dispatch_model(self.engine, device_map="auto")  # 并分配到具体的设备上

2.2 但不知道为什么这个巨慢

2.3 官方教程网站

官方教程：Accelerate
Handling big models for inference
Distributed Inference with 🤗 Accelerate
Launching your 🤗 Accelerate scripts

小技巧

怎么在pycharm上使用accelerate launch 命令

在教程中说，使用accelerate的话要这么做

但是我想用pycharm来实现这一步骤，那就这么做来源是Debugging in Pycharm #535

具体操作如下