-
1 基本映射方法
- 设计设备映射时,可以让Accelerate库来处理设备映射的计算
- 通过设置
device_map
为支持的选项之一("auto"、 "balanced"、 "balanced_low_0"、 "sequential");或者如果想更精确地控制每一层应该去哪里,也可以自己创建一个设备映射
"auto" 和 "balanced" | 在所有可用的GPU上均匀分配模型 |
"balanced_low_0" | 在除了第一个GPU之外的所有GPU上均匀分配模型,并且只有在其他GPU放不下时,才在GPU 0上放置内容 当你需要在生成 Transformers 模型的输出时使用GPU 0进行一些处理时,这个选项非常有用 |
"sequential" | 尽可能在GPU 0上放置内容,然后移动到GPU 1,依此类推(如果不需要,就不会使用最后的GPU) |
2 max_memory
- 在infer_auto_device_map中,通过使用
max_memory
参数来限制每个GPU上使用的内存 - 设置
max_memory
时,你应该传递一个包含GPU标识符(例如0、1等)和“cpu”键的字典- 值可以是一个整数(以字节为单位)或一个带单位的数字字符串,如"10GiB"或"10GB"
from accelerate import infer_auto_device_mapdevice_map = infer_auto_device_map(my_model, max_memory={0: "10GiB", 1: "10GiB", "cpu": "30GiB"})
3 完全自行设计设备映射
- 如果选择完全自行设计设备映射,它应该是一个字典,键是模型的模块名称,值是一个有效的设备标识符(例如GPU的一个整数)或“cpu”用于CPU卸载,“disk”用于磁盘卸载
- 键需要覆盖整个模型
- 例如,如果你的模型有两个块(block1和block2),每个块包含三个线性层(linear1、linear2和linear3),一个有效的设备映射可以是:
device_map = {"block1": 0, "block2": 1}
device_map = {"block1": 0,"block2.linear1": 0, "block2.linear2": 1, "block2.linear3": 1}
下面这种映射不是有效的,因为它没有覆盖模型的每个参数
device_map = {"block1": 0, "block2.linear1": 1, "block2.linear2": 1}
- 为了最有效率,确保设备映射以顺序方式将参数放在GPU上,以避免在GPU之间进行多次数据传输。
- 例如,不要将第一个权重放在GPU 0上,然后将权重放在GPU 1上,最后的权重再放回GPU 0