LLM-text_generation_interfence

今天来看看，Huggingface 中的一个项目 text-generation-interfence 。

它的特点在 github 中的很显眼的位置，就不在这里赘述了。其中如果有一些不太清楚的名词，可以查看 最后一部分——扩展资料，或许有你想要的东西。

前置知识——模型加载

不同的 transformer 训练不同的 model，同理，不同的 model 需要被不同的 transformer 加载：

例如 ChatGLM 和 RWKV，通过代码看，ChatGLM 是通过常规的 transformer 来加载模型的。

RWKV 是基于 transformer 的变种，是将 transformer 中的 self-attention 替换为 Position Encoding 和 TimeMix。

self-attention 的含义，可以查看 LLM 100 关键词； Position Encoding 和  TimeMix 的含义，可以查看这个  链接

总体来说，不同的模型，由不同的 transformer 加载。

################# CHATGLM 加载模型的方法 ##################
>>> from transformers import AutoTokenizer, AutoModel
>>> tokenizer = AutoTokenizer.from_pretrained("THUDM/chatglm-6b", trust_remote_code=True)
>>> model = AutoModel.from_pretrained("THUDM/chatglm-6b", trust_remote_code=True).half().cuda()
>>> model = model.eval()
>>> response, history = model.chat(tokenizer, "你好", history=[])
>>> print(response)
你好 !我是人工智能助手 ChatGLM-6B,很高兴见到你,欢迎问我任何问题。
>>> response, history = model.chat(tokenizer, "晚上睡不着应该怎么办", history=history)
>>> print(response)

################# RWKV 加载模型的方法  ##################
from rwkv.model import RWKV                         # pip install rwkv
model = RWKV(model='/fsx/BlinkDL/HF-MODEL/rwkv-4-pile-1b5/RWKV-4-Pile-1B5-20220903-8040', strategy='cuda fp16')

out, state = model.forward([187, 510, 1563, 310, 247], None)   # use 20B_tokenizer.json
print(out.detach().cpu().numpy())                   # get logits
out, state = model.forward([187, 510], None)
out, state = model.forward([1563], state)           # RNN has state (use deepcopy if you want to clone it)
out, state = model.forward([310, 247], state)
print(out.detach().cpu().numpy())  

项目架构

整个项目由三部分组成：

• launcher
• router
• serve

Launcher、Router和Server（Python gRPC服务）都是服务的组成部分，它们各自承担不同的职责，共同提供一个完整的文本生成推理服务。以下是它们之间的关系：

• Launcher：这是服务的启动器，它负责启动和运行服务。它可能会启动 Router，并设置好所有的路由规则。然后，它会监听指定的地址和端口，等待并处理来自客户端的连接。当接收到一个连接时，它会将连接转发给Router 进行处理。
• Router：这是服务的中间件，它的主要职责是路由和调度请求。当客户端发送一个请求时，Router 会接收这个请求，然后根据请求的内容和当前的系统状态，决定将请求路由到哪个处理器进行处理。这个处理器可能是Server 中的一个 gRPC 方法。Router 的目的是有效地管理和调度系统资源，提高系统的并发处理能力和响应速度。
• Server（Python gRPC服务）：这是服务的核心部分，它实现了文本生成推理的主要逻辑。它提供了一些 gRPC 方法，如 Info、Health、ServiceDiscovery、ClearCache、FilterBatch、Prefill 和 Decode，这些方法用于处理客户端的请求，执行文本生成的推理任务，并返回结果。这个服务可能运行在一个单独的服务器上，独立于Launcher 和 Router。

launcher 启动器

顾名思义，launcher 启动器，就是负责启动的程序，主要做以下工作：(在 launcher/src/main.rs 中)

1. 通过 serve 的命令下载模型，代码中执行的函数为： download_convert_model(&args, running.clone())?;
2. 启动 serve ，代码中执行的函数为： spawn_shards(...)
3. 启动 router，代码中执行的函数为：spawn_webserver(args, shutdown.clone(), &shutdown_receiver)?;

所以，router 和 serve 负责主要的逻辑处理与模型调用。在项目中有一个架构图，可以更加直观的认识到它们之间的关系，其架构如下图所示：

router 路由

可以看到 router 这个 webserver 负责接收请求，然后放在 buffer 中，等收集到一定量的数据后，一个 batch 一个 batch 的以 rpc 的方式发送给 serve 的去处理。

对外暴露的 url 很少同时也很精简，只有四个：

1. /generate : 一次性生成所有回答的 token
2. /generate_stream ：流式的生成所回答的 token (就类似于 chatgpt 一样，一个字一个字的显现)
3. /metrics ： 获取该服务的 metrics 信息。
4. /info ：获取模型的相关信息

serve

在图中，也可以看到，在每个卡上都启动了一个 serve，被叫做 shard，这也是 launcher 的作用之一，通过参数来决定 serve 启动的情况。

在 serve 端的代码，有两个命令行启动脚本（serve/text_generation_server/cli.py）：

# 下载模型权重的方法
@app.command()
def download_weights(
...
)
...

# 启动 serve 服务的方法
@app.command()
def serve(
...
)
...

其实内部逻辑也很简单，稍微处理一下数据后，直接调用 model 的接口来处理。

Server 对外暴露了一下接口：（这里说的对外，指的是 router ）

1. Info : 返回 model 信息
2. Health : 检查 serve 的健康状况
3. ServiceDiscovery : 服务发现，实现也很简单，将所有的 serve 的地址发送出去
4. ClearCache : 清除 cache 中的数据 （cache 的功能再看）
5. FilterBatch
6. Prefill
7. Decode

cache 中的存储单位是 batch （在 router 中提过，router 就是一个 batch 一个 batch 来传的。）

内部接口的含义

再然后，就剩下最重要的三个功能：FilterBatch、Prefill、Decode

FilterBatch 流程如下：（使用场景还不太清楚）

先从 cache 中以 batch_id 获取特定的 batch 再从 batch 中过滤出我们想要留下的 request_ids（这里的 request_id 指的是 客户端发送的请求 id ） 过滤后，再将 batch 放回 cache 中。

Prefill 的主要功能是：

1. 从 router 接收 batch ,然后根据模型给的 from_pb 方法整理一下 batch 中的信息 并且 通过 tokenizer 来将相应的词转化成词向量。（from_pb 方法之后在说）
2. 将 整理后的 batch 信息，通过 model 的 generate_token 方法，生成新的 token (也就是预测的词)，同时也会返回 next_batch。（generate_token 方法之后在说）
3. 将 next_batch 存放到 cache 中。
4. 返回消息。

Decode 的功能也很简单，主要功能是：

1. 通过 request 传入的 batch.id 从 cache 中获取 batch
2. 将这些 batch 通过 model 的 generate_token 方法，生成新的 token，同时会返回 next_batch。
3. 将 next_batch 存放到 cache 中。
4. 返回消息。

主要是第一步，从 缓存中获取 batch，这样有两个好处：第一，request 不需要传输历史的 信息，上下文都在 cache 中；第二，cache 中缓存的是 词向量 的信息，所以，在每次预测词的时候，只需要将传入的 信息 通过词嵌入 转化成 词向量，其他的信息就不需要再做转化了，减少了大量的计算工作。

扩展资料

flash-attation

Flash-Attention 也是 Huggingface 中比较出名的一个项目，它是一种优化技术，使用了一种特殊的计算方法，可以在不影响模型精度的情况下，显著提高计算速度。这使得模型可以更快地处理大量数据。

优化的过程如下：

在传统的 Transformer 模型中，注意力机制需要计算一个大小为 (n^2) 的矩阵，其中 n 是序列的长度。这个过程在计算上是非常昂贵的，特别是对于长序列。Flash-Attention 通过使用一种叫做 "softmax decomposition" 的技术，将这个复杂的计算过程分解为几个更简单的步骤。这种方法可以显著提高计算速度，而不会影响模型的精度。

soft decomposition 在 100 words 中介绍一下吧，在这个引用一下。

不仅如此，在传统的 Transformer 模型中，注意力机制需要存储一个大小为 (n^2) 的矩阵，这在内存上是非常昂贵的，特别是对于长序列。Flash-Attention 还使用了 “checkpointing” 来减少内存的使用。

同理， checkpointing 在 100 words 中也介绍一下，在这里引用。

bitsandbytes

"bitsandbytes" 是一个库，它提供了一些用于 PyTorch 的 8 位 CUDA 函数。这个库的主要目标是提供一种高效的方式来处理和存储大型模型的权重，特别是在进行大规模训练或推理时。

"bitsandbytes" 的一个关键特性是它支持 4 位和 8 位的量化。量化 是一种减少模型大小和计算需求的技术，它通过将模型的权重从 32 位浮点数（通常用于训练神经网络的数据类型）减少到更小的位数（如 8 位或 4 位）来实现。这种方法可以显著减少模型的内存占用，从而使模型能够在更小的硬件上运行，或者处理更大的数据集。

"bitsandbytes" 库的另一个重要特性是它支持在 GPU 上进行计算，这使得它可以充分利用 GPU 的并行计算能力，从而提高计算速度。

总的来说，"bitsandbytes" 是一个强大的工具，它可以帮助研究人员和开发人员更有效地处理和存储大型模型的权重，从而使这些模型能够在更小的硬件上运行，或者处理更大的数据集。

Safetensors

"Safetensors" 是一个用于保存和加载张量的库，它支持最常见的框架（包括 PyTorch、TensorFlow、JAX、PaddlePaddle 和 NumPy）。这个库的主要目标是提供一种安全和高效的方式来处理张量，特别是在进行模型保存和加载时。

"Safetensors" 的一个重要特性是它可以安全地加载文件。这是因为它不使用 Python 的 pickle 库，pickle 库在处理文件时存在安全风险，可能会执行恶意代码。相反，"Safetensors" 使用了一种安全的方法来处理文件，从而避免了这种风险。

此外，"Safetensors" 还支持懒加载（lazy loading），这意味着它可以在需要时才加载部分张量，而不是一次性加载整个张量。这种特性可以提高处理大型模型的效率，特别是在分布式环境中。

import torch
from safetensors.torch import load_file, save_file

weights = {"embeddings": torch.zeros((10, 100))}
save_file(weights, "model.safetensors")
weights2 = load_file("model.safetensors")

原文地址：

LLM-text_generation_interfence - 知乎