huggingface 笔记：device_map

• 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_map

device_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