【vLLM 学习】Distributed_vllm distributed代码分析-优快云博客

vLLM 是一款专为大语言模型推理加速而设计的框架，实现了 KV 缓存内存几乎零浪费，解决了内存管理瓶颈问题。

更多 vLLM 中文文档及教程可访问 →https://vllm.hyper.ai/

源码 examples/offline_inference/distributed.py

# SPDX-License-Identifier: Apache-2.0

"""
此示例显示了如何使用 ray 数据进行离线批处理推断
在多节点群集上分布。
在 https://docs.ray.io/en/latest/data/data.html 中了解有关 ray 数据的更多信息
"""

from typing import Any

import numpy as np
import ray
from packaging.version import Version
from ray.util.scheduling_strategies import PlacementGroupSchedulingStrategy

from vllm import LLM, SamplingParams

assert Version(ray.__version__) >= Version(
    "2.22.0"), "Ray version must be at least 2.22.0"

# 创建一个采样参数对象。
sampling_params = SamplingParams(temperature=0.8, top_p=0.95)

# 每个实例设置张量并行性。
tensor_parallel_size = 1

# 设置数量的实例。每个实例都将使用 Tensor_parallel_size GPU。
num_instances = 1


# 创建一个以进行批处理推理的类。
class LLMPredictor:

    def __init__(self):
        # Create an LLM.
        # 创建一个 LLM。
        self.llm = LLM(model="meta-llama/Llama-2-7b-chat-hf",
                       tensor_parallel_size=tensor_parallel_size)

    def __call__(self, batch: dict[str, np.ndarray]) -> dict[str, list]:
        # 从提示中生成文本。
        # 输出是包含提示的 RequestOutput 对象的列表，
        # 生成的文本和其他信息。
        outputs = self.llm.generate(batch["text"], sampling_params)
        prompt: list[str] = []
        generated_text: list[str] = []
        for output in outputs:
            prompt.append(output.prompt)
            generated_text.append(' '.join([o.text for o in output.outputs]))
        return {
            "prompt": prompt,
            "generated_text": generated_text,
        }


# 从 S3读取一个文本文件。 ray 数据支持读取多个文件
# 来自云存储 (例如 JSONL，Parquet，CSV，二进制格式)。
ds = ray.data.read_text("s3://anonymous@air-example-data/prompts.txt")



# 对于 tensor_parallel_size> 1，我们需要为 vLLM 创建放置组
# 要使用。每个演员都必须拥有自己的安置小组。
def scheduling_strategy_fn():
    # One bundle per tensor parallel worker
    # 每张张量的平行工人一捆
    pg = ray.util.placement_group(
        [{
            "GPU": 1,
            "CPU": 1
        }] * tensor_parallel_size,
        strategy="STRICT_PACK",
    )
    return dict(scheduling_strategy=PlacementGroupSchedulingStrategy(
        pg, placement_group_capture_child_tasks=True))


resources_kwarg: dict[str, Any] = {}
if tensor_parallel_size == 1:
    # 对于 tensor_parallel_size == 1，我们只是设置 num_gpus = 1。
    resources_kwarg["num_gpus"] = 1
else:
    # 否则，我们必须设置 num_gpus = 0并提供
    # 一个将创建一个安置组的函数
    # 每个实例。
    resources_kwarg["num_gpus"] = 0
    resources_kwarg["ray_remote_args_fn"] = scheduling_strategy_fn

# 应用所有输入数据的批处理推断。
ds = ds.map_batches(
    LLMPredictor,
    # Set the concurrency to the number of LLM instances.
    # 将并发设置为 LLM 实例的数量。
    concurrency=num_instances,
    # Specify the batch size for inference.
    # 指定推理的批次大小。
    batch_size=32,
    **resources_kwarg,
)

# 窥视前10个结果。
# 注意:这是用于本地测试和调试。对于生产用例，
# 应该写出完整的结果，如下所示。
outputs = ds.take(limit=10)
for output in outputs:
    prompt = output["prompt"]
    generated_text = output["generated_text"]
    print(f"Prompt: {prompt!r}, Generated text: {generated_text!r}")

# ds.write_parquet("s3://<your-output-bucket>")
# 将推理输出数据作为 parquet 文件输出到 S3。
# 多个文件将写入输出目标，
# 每个任务将分别编写一个或多个文件。
#
# ds.write_parquet ( "s3:// <your-out-output-bucket>")