多张卡部署一个codellama实例

本文主要是介绍多张卡部署一个codellama实例，希望对大家解决编程问题提供一定的参考价值，需要的开发者们随着小编来一起学习吧！

模型推理过程

使用Transformers框架进行文本生成类任务会经过以下步骤：

加载预训练模型和tokenizer
主要涉及到网络传输（下载模型参数）、解压缩以及模型参数的初始化，这些过程通常是在 CPU 上执行
文本编码
对输入文本进行分词、转换为 token ID，并最终将 token ID 转换为 PyTorch 张量。这个过程不涉及模型参数的加载，也不需要进行模型的推理或训练，因此不需要 GPU 资源。
生成文本
将模型参数需要加载到 GPU 的显存中，把编码后的输入文本作为输入得到输出项，一般会经过以下几个步骤：

模型推理: 将编码后的输入文本输入到预训练的语言模型中。模型会根据输入文本的上下文信息，预测下一个 token 的概率分布。这个过程通常是在模型的前向传播中完成的。
采样下一个 token: 根据模型预测的概率分布，从候选 token 中采样一个作为下一个 token。常见的采样方法包括贪婪采样、随机采样和核采样等。
更新输入文本: 将新采样的 token 添加到输入文本末尾，并丢弃最早的 token，形成新的输入文本序列。
重复步骤 2-4: 重复进行模型推理和采样，直到生成了指定长度的文本或满足停止条件。

解码
将生成的 token IDs 解码成文本字符串（人类可读的文本），即生成的文本。

以下是一个文本生成模型的代码示例：

from transformers import AutoModelForCausalLM, AutoTokenizer# 加载预训练模型和 tokenizer
model_name = "gpt2"
model = AutoModelForCausalLM.from_pretrained(model_name, trust_remote_code=True)
tokenizer = AutoTokenizer.from_pretrained(model_name)# 从模型中生成文本
input_text = "Once upon a time,"
input_ids = tokenizer(input_text, return_tensors="pt").input_ids
outputs = model.generate(input_ids)# 将生成的文本解码为人类可读的文本
generated_text = tokenizer.decode(outputs[0], skip_special_tokens=True)
print(generated_text)

模型Layer

在深度学习中，模型通常由多个层（或称为模块）组成，每个层执行特定的功能，并且它们以一种有序的方式连接在一起，构成了整个模型。不同类型的模型可能由不同种类的层组成，每个层都有其特定的功能和作用。

以下是一些常见的模型层及其功能的示例：

输入层（Input Layer）：输入层负责接收原始输入数据，并将其转换为模型可以处理的张量格式。在自然语言处理任务中，输入层通常由词嵌入层组成，用于将文本数据编码成词嵌入向量。
卷积层（Convolutional Layer）：卷积层主要用于处理图像数据，通过卷积操作提取图像的特征。在卷积神经网络（CNN）中，通常会包含多个卷积层，每个卷积层会对输入进行不同尺寸和数量的卷积操作。
循环层（Recurrent Layer）：循环层用于处理序列数据，具有记忆功能，能够捕捉序列数据中的时序信息。循环神经网络（RNN）和长短期记忆网络（LSTM）是常见的循环层结构。
注意力层（Attention Layer）：注意力层用于处理序列数据或序列-序列数据，能够动态地计算输入序列中各个位置的重要性，并将注意力权重应用于相应的位置。Transformer 模型中的自注意力机制就是一种常见的注意力层。
全连接层（Fully Connected Layer）：全连接层通常位于神经网络的最后几层，用于将模型提取的特征映射到最终的输出空间。在分类任务中，全连接层通常用于将特征向量映射到类别概率分布。
输出层（Output Layer）：输出层负责生成模型的最终输出，通常根据任务的不同，输出层可能采用不同的激活函数和损失函数。在分类任务中，输出层通常采用 softmax 激活函数生成类别概率分布。

可以使用以下代码查看某个模型的Layer分布

for name, module in model.named_modules():print(f"模块名称: {name}, 模块对象: {module}")

将Layer分配到不同的显卡上进行计算。每个显卡只负责计算模型的一部分，可以解决单个显卡内存不足的问题。

Layer分散到多卡推理

import torch
from transformers import AutoModelForCausalLM, AutoTokenizer, LlamaForCausalLM
from accelerate import dispatch_model
from accelerate.utils import get_balanced_memory, infer_auto_device_mapmodel_name = "codellama/CodeLlama-7b-Instruct-hf"
tokenizer = AutoTokenizer.from_pretrained(model_name, trust_remote_code=True)
model = AutoModelForCausalLM.from_pretrained(model_name, trust_remote_code=True)no_split_module_classes = LlamaForCausalLM._no_split_modules
max_memory = {1: '5GiB', 2: '5GiB', 3: '7GiB'}
device_map = infer_auto_device_map(model, max_memory=max_memory, no_split_module_classes=no_split_module_classes)
print(device_map)
model = dispatch_model(model, device_map=device_map, offload_dir="tmp")input_text = "def function("
input_ids = tokenizer(input_text, return_tensors="pt").input_ids.cuda()# 生成文本
outputs = model.generate(input_ids, max_length=100)# 将生成的文本解码为人类可读的文本
generated_texts = [tokenizer.decode(output, skip_special_tokens=True) for output in outputs]
print(generated_texts)