LLM 之device_map参数解析

LLM大模型在Transformer做推理是，如何根据GPU的显卡做多卡加载模型呢？通常会使用device_map做参数，device_map是AutoModel.from_pretrained() 方法中的一个重要参数，它用于指定模型的各个部件应加载到哪个具体的计算设备上，以实现资源的有效分配和利用。这个参数在进行模型并行或分布式训练时特别有用。

device_map参数主要有如下4种选项：”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 会直接爆显存了（直接别用了）

示例代码：

model = AutoModel.from_pretrained(model_path, trust_remote_code=True, device_map=device)

def split_model(model_name):   #根据卡的数量和模型的层数做均分

    device_map = {}

    world_size = torch.cuda.device_count()

    world_size = 2 #使用两张卡，序号从6开始

    print(world_size)

    num_layers = {

        ‘InternVL2-1B’: 24, ‘InternVL2-2B’: 24, ‘InternVL2-4B’: 32, ‘InternVL2-8B’: 32,

        ‘InternVL2-26B’: 48, ‘InternVL2-40B’: 60, ‘InternVL2-Llama3-76B’: 80}[model_name]

    # Since the first GPU will be used for ViT, treat it as half a GPU.

    num_layers_per_gpu = math.ceil(num_layers / (world_size – 0.5))

    num_layers_per_gpu = [num_layers_per_gpu] * world_size

    num_layers_per_gpu[0] = math.ceil(num_layers_per_gpu[0] * 0.5)

    print(num_layers_per_gpu)

    layer_cnt = 0

    for i, num_layer in enumerate(num_layers_per_gpu):

        i += 6

        for j in range(num_layer):

            device_map[f’language_model.model.layers.{layer_cnt}’] = i

            # print(f’layer_cnt: {layer_cnt}—-{i}’)

            layer_cnt += 1

    print(i)  # i=6,7

    first_gpu = i – 1 #把其他配置都使用第一张卡

    device_map[‘vision_model’] = first_gpu

    device_map[‘mlp1’] = first_gpu

    device_map[‘language_model.model.tok_embeddings’] = first_gpu

    device_map[‘language_model.model.embed_tokens’] = first_gpu

    device_map[‘language_model.output’] = first_gpu

    device_map[‘language_model.model.norm’] = first_gpu

    device_map[‘language_model.lm_head’] = first_gpu

    device_map[f’language_model.model.layers.{num_layers – 1}’] = first_gpu

   

    return device_map

device = split_model(“InternVL2-40B“)

model = AutoModel.from_pretrained(model_path, trust_remote_code=True, device_map=device)

如果是用多卡做推理加载模型，需要考虑模型大小以及每张卡的内存情况