为在多个 Task 之间共享而服务 LLM

通常情况下，我们希望在 Pipeline 中为多个 Task 使用相同的 LLM。为了避免多次加载 LLM（次数与 Task 的数量相同）并避免浪费资源，建议使用 text-generation-inference 或 vLLM 等解决方案来服务模型，然后使用与 AsyncLLM 兼容的客户端（如 InferenceEndpointsLLM 或 OpenAILLM）分别与服务器通信。

使用 text-generation-inference 服务 LLM

model=meta-llama/Meta-Llama-3-8B-Instruct
volume=$PWD/data # share a volume with the Docker container to avoid downloading weights every run

docker run --gpus all --shm-size 1g -p 8080:80 -v $volume:/data \
    -e HUGGING_FACE_HUB_TOKEN=<secret> \
    ghcr.io/huggingface/text-generation-inference:2.0.4 \
    --model-id $model

注意

上面的 bash 命令是从官方文档 text-generation-inference 中复制粘贴过来的。请参考官方文档以获取更多信息。

然后我们可以使用 InferenceEndpointsLLM，并将 base_url=http://localhost:8080 （指向我们的 TGI 本地部署）

from distilabel.models import InferenceEndpointsLLM
from distilabel.pipeline import Pipeline
from distilabel.steps import LoadDataFromDicts
from distilabel.steps.tasks import TextGeneration, UltraFeedback

with Pipeline(name="serving-llm") as pipeline:
    load_data = LoadDataFromDicts(
        data=[{"instruction": "Write a poem about the sun and moon."}]
    )

    # `base_url` points to the address of the `TGI` serving the LLM
    llm = InferenceEndpointsLLM(base_url="http://192.168.1.138:8080")

    text_generation = TextGeneration(
        llm=llm,
        num_generations=3,
        group_generations=True,
        output_mappings={"generation": "generations"},
    )

    ultrafeedback = UltraFeedback(aspect="overall-rating", llm=llm)

    load_data >> text_generation >> ultrafeedback

使用 vLLM 服务 LLM

docker run --gpus all \
    -v ~/.cache/huggingface:/root/.cache/huggingface \
    --env "HUGGING_FACE_HUB_TOKEN=<secret>" \
    -p 8000:8000 \
    --ipc=host \
    vllm/vllm-openai:latest \
    --model meta-llama/Meta-Llama-3-8B-Instruct

注意

上面的 bash 命令是从官方文档 vLLM 中复制粘贴过来的。请参考官方文档以获取更多信息。

然后我们可以使用 OpenAILLM，并将 base_url=http://localhost:8000 （指向我们的 vLLM 本地部署）

from distilabel.models import OpenAILLM
from distilabel.pipeline import Pipeline
from distilabel.steps import LoadDataFromDicts
from distilabel.steps.tasks import TextGeneration, UltraFeedback

with Pipeline(name="serving-llm") as pipeline:
    load_data = LoadDataFromDicts(
        data=[{"instruction": "Write a poem about the sun and moon."}]
    )

    # `base_url` points to the address of the `vLLM` serving the LLM
    llm = OpenAILLM(base_url="http://192.168.1.138:8000", model="")

    text_generation = TextGeneration(
        llm=llm,
        num_generations=3,
        group_generations=True,
        output_mappings={"generation": "generations"},
    )

    ultrafeedback = UltraFeedback(aspect="overall-rating", llm=llm)

    load_data >> text_generation >> ultrafeedback