LLamafactory异步API高效调用实现与速度对比

背景

一开始处理各家闭源大模型的API调用时，基本都是用同步方式。但说实话，这玩意的响应真跟上了发条似的，一个接一个排队，慢得让人抓狂。后来一想，反正现在大模型服务动不动就是几十上百个请求并行，为什么不把调用搞成异步？干活的节奏瞬间就能起飞。

简介

这里写的代码，最大的亮点是：它完全兼容 LLaMAFactory 的原始数据集导入方式，不需要你额外改造或者折腾数据格式。用起来非常顺手。推理速度相比同步API调用——完全不是一个量级，吞吐量的差异，简直可以用“碾压”来形容。

整个方案的核心思路是：
- 直接用 LLaMAFactory 原生方法加载数据集，省去中间数据格式转换的麻烦；
- 封装了一个叫 AsyncAPICall

看看核心代码文件 async_call_api.py：

# pip install langchain langchain_openai

import os
import sys
import json
import asyncio

import fire
from tqdm import tqdm
from dataclasses import dataclass
from aiolimiter import AsyncLimiter
from typing import List
import pandas as pd
from langchain_openai import ChatOpenAI
from dotenv import load_dotenv

from llamafactory.hparams import get_train_args
from llamafactory.extras.constants import IGNORE_INDEX
from llamafactory.data.loader import _get_merged_dataset

load_dotenv()


class AsyncLLM:
    def __init__(
        self,
        model: str = "gpt-3.5-turbo",
        base_url: str = "http://localhost:{}/v1/".format(
            os.environ.get("API_PORT", 8000)
        ),
        api_key: str = "{}".format(os.environ.get("API_KEY", "0")),
        num_per_second: int = 6,
        **kwargs,
    ):
        self.model = model
        self.base_url = base_url
        self.api_key = api_key
        self.num_per_second = num_per_second

        self.limiter = AsyncLimiter(self.num_per_second, 1)

        self.llm = ChatOpenAI(
            model=self.model, base_url=self.base_url, api_key=self.api_key, **kwargs
        )

    async def __call__(self, text):
        # 限速
        async with self.limiter:
            return await self.llm.ainvoke([text])


llm = AsyncLLM(
    base_url="http://localhost:{}/v1/".format(os.environ.get("API_PORT", 8000)),
    api_key="{}".format(os.environ.get("API_KEY", "0")),
    num_per_second=10,
)
llms = [llm]


@dataclass
class AsyncAPICall:
    uid: str = "0"

    @staticmethod
    async def _run_task_with_progress(task, pbar):
        result = await task
        pbar.update(1)
        return result

    @staticmethod
    def async_run(
        llms: List[AsyncLLM],
        data: List[str],
        keyword: str = "",
        output_dir: str = "output",
        chunk_size=500,
    ) -> List[str]:

        async def infer_chunk(llms: List[AsyncLLM], data: List):
            results = [llms[i % len(llms)](text) for i, text in enumerate(data)]
            with tqdm(total=len(results)) as pbar:
                results = await asyncio.gather(
                    *[
                        AsyncAPICall._run_task_with_progress(task, pbar)
                        for task in results
                    ]
                )
            return results

        idx = 0
        all_df = []
        file_exist_skip = False
        user_confirm = False

        while idx < len(data):
            file_path = os.path.join(output_dir, "tmp", f"{idx}.csv.temp")

            if os.path.exists(file_path):
                if not user_confirm:
                    while True:
                        user_response = input(
                            f"Find {file_path} file already exists. Do you want to skip them forever?ny or Y to skip, n or N to rerun to overwrite: "
                        )
                        if user_response.lower() == "y":
                            user_confirm = True
                            file_exist_skip = True
                            break
                        elif user_response.lower() == "n":
                            user_confirm = True
                            file_exist_skip = False
                            break

                if file_exist_skip:
                    tmp_df = pd.read_csv(file_path)
                    all_df.append(tmp_df)
                    idx += chunk_size
                    continue

            tmp_data = data[idx : idx + chunk_size]
            loop = asyncio.get_event_loop()
            tmp_result = loop.run_until_complete(infer_chunk(llms=llms, data=tmp_data))
            tmp_result = [item.content for item in tmp_result]

            tmp_df = pd.DataFrame({"infer": tmp_result})

            if not os.path.exists(p := os.path.dirname(file_path)):
                os.makedirs(p, exist_ok=True)

            tmp_df.to_csv(file_path, index=False)
            all_df.append(tmp_df)
            idx += chunk_size

        all_df = pd.concat(all_df)
        return all_df["infer"]


def async_api_infer(
    model_name_or_path: str = "",
    eval_dataset: str = "",
    template: str = "",
    dataset_dir: str = "data",
    do_predict: bool = True,
    predict_with_generate: bool = True,
    max_samples: int = None,
    output_dir: str = "output",
    chunk_size=50,
):

    if len(sys.argv) == 1:
        model_args, data_args, training_args, finetuning_args, generating_args = (
            get_train_args(
                dict(
                    model_name_or_path=model_name_or_path,
                    dataset_dir=dataset_dir,
                    eval_dataset=eval_dataset,
                    template=template,
                    output_dir=output_dir,
                    do_predict=True,
                    predict_with_generate=True,
                    max_samples=max_samples,
                )
            )
        )
    else:
        model_args, data_args, training_args, finetuning_args, generating_args = (
            get_train_args()
        )

    dataset = _get_merged_dataset(
        data_args.eval_dataset, model_args, data_args, training_args, "sft"
    )

    labels = [item[0]["content"] for item in dataset["_response"]]
    prompts = [item[0]["content"] for item in dataset["_prompt"]]

    infers = AsyncAPICall.async_run(
        llms,
        prompts,
        chunk_size=chunk_size,
        output_dir=training_args.output_dir,
    )

    if not os.path.exists(training_args.output_dir):
        os.makedirs(training_args.output_dir, exist_ok=True)

    output_prediction_file = os.path.join(
        training_args.output_dir, "generated_predictions.jsonl"
    )

    with open(output_prediction_file, "w", encoding="utf-8") as writer:
        res: List[str] = []
        for text, pred, label in zip(prompts, infers, labels):
            res.append(
                json.dumps(
                    {"prompt": text, "predict": pred, "label": label},
                    ensure_ascii=False,
                )
            )
        writer.write("n".join(res))


if __name__ == "__main__":
    fire.Fire(async_api_infer)

调用方式也很简洁，配合一个 yaml 配置文件即可：

async.yaml

## model
model_name_or_path: qwen/Qwen2.5-7B-Instruct

### method
do_predict: true

### dataset
dataset_dir: ../data/
eval_dataset: alpaca_zh_demo
template: qwen

### output
output_dir: output
# overwrite_output_dir: true


### eval
predict_with_generate: true
max_samples: 100

用下面的命令就可以直接跑起来：

python async_call_api.py async.yaml

看一个输出的例子，帮你理解下结果格式：

{"prompt": "识别并解释给定列表中的两个科学理论：细胞理论和日心说。", "predict": "细胞理论和日心说是两个重要的科学理论，分别属于生物学和天文学领域。nn1. 细胞理论：细胞理论是生物学中的一个基本理论，它描述了所有生命体的基本结构和功能单位——细胞。细胞理论包含以下三个主要观点：nn   a. 所有生物体都是由一个或多个细胞组成的。n   b. 细胞是生命活动的基本单位，负责维持生物体的生命活动。n   c. 新细胞只能由已存在的细胞分裂产生。nn   细胞理论的提出和发展对于生物学研究具有重要意义，它揭示了生物体的基本结构和功能单位，为生物学研究提供了理论基础。nn2. 日心说：日心说是一种关于太阳系结构的理论，它描述了太阳在太阳系中的中心位置，而地球和其他行星则围绕太阳旋转。日心说最初由古希腊天文学家阿里斯塔克提出，但直到16世纪，哥白尼重新提出这一理论，并得到了进一步的发展和完善。日心说的提出碘伏了传统的地心说观念，为现代天文学的发展奠定了基础。nn这两个科学理论分别在生物学和天文学领域中具有重要的意义，它们揭示了生命体和宇宙的基本结构和运行规律。", "label": "细胞理论是生物科学的一个理论，它认为所有生命体都是由微小的基本单元——细胞所构成。这是生物学的一个基础理论，认为细胞是所有生物的基本结构和功能单位，所有的生物都是由一个或多个细胞组成，细胞只能通过细胞分裂产生新的细胞。这一理论由薛定谔、施瓦内和雪莱于1839年首次提出。nn日心说是指太阳是太阳系的中心，也就是说，行星围绕太阳旋转的理论。这个理论打破了传统的地心说观点，认为地球并不是宇宙的中心。日心说的提出者是尼古拉·哥白尼，他在16世纪初发表了他的著作《天体运行论》，阐述了太阳系行星围绕太阳运行的模型，为天文学的发展做出了巨大贡献。"},...,
}

最后一条补充：输出文件里 predict 字段就是大模型生成的口语化推理结果，方便拿它跟 label（参考答案）做交叉比对，评估模型精度。

速度对比

这里放两组硬核数据，直观展示一下“异步”到底能快多少。

异步调用速度

下面是两个异步调用的进度条（因为分块设了 chunk_size=50，所以100条数据被分成了2块）：

100%|██████████████████████████████████████████████████████████████████████████████████████████████| 50/50 [00:22<00:00, 2.27it/s]
100%|██████████████████████████████████████████████████████████████████████████████████████████████| 50/50 [00:22<00:00, 2.22it/s]

100条数据总共用时 44秒，平均每秒吞进去 2.2 条。

同步调用速度

对比之下，同步调用就简单粗暴多了，一个 for 循环愣走——

infers = []
for prompt in tqdm(prompts):
    infers.append(llm.llm.invoke(prompt))

看看它的惨状：

100%|████████████████████████████████████████████████████████████████████████████████████████████| 100/100 [06:54<00:00, 4.15s/it]

同样的100条数据，同步要耗掉整整 6分54秒，平均每条4.15秒。

一目了然，异步比同步快了大约 9.41倍。与其在同步模式下干耗着等API回来，不如直接上异步，把计算资源彻底用满。

你是一名 AI 行业编辑，请围绕下面这条热点输出一份资讯解读：
热点：LLamafactory异步API高效调用实现与速度对比要求：
1. 先用一句话解释这条热点在讲什么
2. 再总结它为什么重要
3. 说明会影响哪些 AI 产品或内容方向
4. 最后给出 3 个适合资讯站使用的标题