为 Step 分配资源

当处理在具有丰富资源（CPU 和 GPU）的分布式环境中执行的复杂 pipelines 时，有时需要在 Step 之间明智地分配这些资源。这就是为什么 distilabel 允许为每个 Step 指定 replicas、cpus 和 gpus 的数量。让我们看一个例子

from distilabel.pipeline import Pipeline
from distilabel.models import vLLM
from distilabel.steps import StepResources
from distilabel.steps.tasks import PrometheusEval


with Pipeline(name="resources") as pipeline:
    ...

    prometheus = PrometheusEval(
        llm=vLLM(
            model="prometheus-eval/prometheus-7b-v2.0",
            chat_template="[INST] {{ messages[0]['content'] }}\\n{{ messages[1]['content'] }}[/INST]",
        ),
        resources=StepResources(replicas=2, cpus=1, gpus=1)
        mode="absolute",
        rubric="factual-validity",
        reference=False,
        num_generations=1,
        group_generations=False,
    )

在上面的例子中，我们正在创建一个 PrometheusEval 任务（记住 Task 是 Step）它将使用 vLLM 来服务 prometheus-eval/prometheus-7b-v2.0 模型。此任务是资源密集型的，因为它需要 LLM，而 LLM 又需要 GPU 才能快速运行。考虑到这一点，我们使用 StepResources 类指定了任务所需的 resources，并且我们定义了每个任务副本需要 1 个 GPU 和 1 个 CPU。此外，我们定义了我们需要 2 个副本，即我们将运行两个任务实例，以便整个数据集的计算运行得更快。此外，StepResources 使用 RuntimeParametersMixin，因此我们还可以在运行 pipeline 时为每个步骤指定资源

...

if __name__ == "__main__":
    pipeline.run(
        parameters={
            prometheus.name: {"resources": {"replicas": 2, "cpus": 1, "gpus": 1}}
        }
    )

就是这样！当运行 pipeline 时，distilabel 将在具有指定资源的节点中创建任务。