PyTorch 实现《响应长度感知和序列调度：一个 LLM 赋能的 LLM 推理管道》

论文简析

论文聚焦

聚焦于响应长度感知和序列调度，以提高 LLM 批推理的效率。注意是批推理，而不是提高单次推理的效率。提高批推理效率，主要聚焦于在一次推理过程中同时处理多个查询请求时的整体性能。它强调的是在单位时间内能够处理的查询数量，即提高系统的吞吐量。提高批推理效率的意义在于，在实际应用场景中，往往会有大量的用户查询需要处理，如果能够同时对多个查询进行批处理，就能更充分地利用计算资源，降低总体的处理时间和成本。

本项目中提出的技术，通过将预计响应长度相近的查询分组，减少冗余计算，提高计算资源的利用率。例如，在处理多个查询时，如果不进行调度，可能会出现一些查询很快完成，而另一些查询还需要很长时间，导致计算资源闲置。通过序列调度，将预计响应长度相近的查询放在一起处理，可以使计算资源更均衡地被利用，从而提高整体的推理效率。

核心思想

该项目发现大语言模型（LLMs）具备提前感知其生成响应长度的能力。利用这一特性，提出了名为 "序列调度" 的技术，通过将预计响应长度相近的查询分组，显著减少冗余计算，在不影响性能的前提下，使推理吞吐量提升了 86%。

主要模块及功能

generate.py 定义了 Creator 类，负责生成文本。支持单样本流式生成（stream）、批量生成（batch）和分组生成（group）三种策略。包含 MLP 类，是一个简单的多层感知机模型。
utils.py 提供了多种工具函数，如设置随机种子（set_seed）、时间测量（timeit）、数据描述（describe）、数据加载和保存（jload, jdump）等。定义了 EvalDataset 类，用于处理评估数据集。实现了 Predictor 类，根据不同策略预测响应长度。包含 schedule 函数，用于对预测的长度进行排序和批处理。
train_lora.py 用于对模型进行 LoRA（Low-Rank Adaptation）微调训练。
eval.py 评估响应长度感知模块的性能，计算误差和准确率。
benchmark.py 对不同的推理策略进行基准测试，比较吞吐量、平均长度、有效令牌比率和失败率等指标。
construct.py 构建用于指令微调的训练数据集。
lenpred.py 对训练数据集进行多次推理，收集长度信息。
sample.py 从原始数据集中随机采样生成训练集和验证集。

存放项目所需的数据集，包括原始对话数据、训练集、验证集以及对应的长度信息文件。

用于存放模型的检查点和 LoRA 权重文件。

包含项目文档中使用的图片，如序列调度流程、提前感知示例等。

其他文件

README.md 项目的详细说明文档，包含项目介绍、安装步骤、使用方法和实验结果等信息。
requirements.txt 列出了项目所需的 Python 包。
train.sh 用于执行模型的指令微调训练的脚本。

论文复现

1. 准备工作

先下载两个模型权重文件（准备科学上网工具， 37G 大小）

huggingface-cli download lmsys/vicuna-7b-v1.5 --local-dir ./vicuna-7b

huggingface-cli download --resume-download huggyllama/llama-7b --local-dir ./llama-7b --local-dir-use-symlinks False

再安装 python 依赖，注意修改版本 transformers==4.28.0 peft==0.4.0，指定 python 为 3.9

pip install -r requirements.txt

2. 修改代码

a. 去除 Nvida GPU 依赖

# 去除Nvida GPU依赖
class Creator:
    def __init__(
        self,
        model_name,
        conv_template="vicuna_v1.1",
        device=None,  # 修改默认值为 None
        num_gpus=0,   # 修改默认值为 0
        load_8bit=False,
        debug=False,
        lora_path=None,
    ):
        # 自动检测设备
        if device is None:
            device = "cuda" if torch.cuda.is_available() else "cpu"
        if device == "cuda" and num_gpus == 0:
            num_gpus = torch.cuda.device_count() if torch.cuda.is_available() else 0

        self.model, self.tokenizer = load_model(
            model_name, device, num_gpus, load_8bit, debug
        )
// ... existing code ...

utils.py 中的 timeit 函数，它试图调用 torch.cuda.synchronize()，但在 CPU 环境下不应该这样做

def timeit(start_time=None):
    if start_time is None:
        return time.time()
    else:
        if torch.cuda.is_available():
            torch.cuda.synchronize()
        return time.time() - start_time

3. 序列调度

Vanilla 策略（无调度）：

CUDA_VISIBLE_DEVICES=0 python -m src.benchmark --num-data 1024

序列调度策略（带 VBS 和 FCR）：

CUDA_VISIBLE_DEVICES=0 python -m src.benchmark --num-data 1024 --strategy seqsch --vbs --fcr --lora-path ./ckpts/vicuna-response-length-perception-module

序列调度策略（不带 VBS 和 FCR）：

CUDA_VISIBLE_DEVICES=0 python -m src.benchmark --num-data 1024 --strategy seqsch --lora-path ./ckpts/vicuna-response-length-perception-module

通过以上步骤，用户可以复现论文中的实验结果，并验证序列调度技术对 LLM 推理效率的提升效果。

小结

这个批量推理实现是工业级提升批量推理的常用算法之一。也是为什么火山引擎和阿里云等云服务商批量推理比单个推理单价低 50%以上。

需要 GPU 服务器

如果一开始就使用 GPU 服务器就不会有这么多卡点，Nvidia GPU 还是不能少。

PyTorch 实现《响应长度感知和序列调度：一个 LLM 赋能的 LLM 推理管道》 ​

相关资料： ​

论文简析 ​

论文聚焦 ​

核心思想 ​

主要模块及功能 ​

论文复现 ​

1. 准备工作 ​

2. 修改代码 ​

3. 序列调度 ​

小结 ​