distilabel
distilabel copied to clipboard
[FEATURE] Use `Step.resources` to set `tensor_parallel_size` and `pipeline_parallel_size` in `vLLM`
Is your feature request related to a problem? Please describe.
To use 8 GPUs and two nodes for vLLM
right now we need to define tensor_parallel_size
and pipeline_parallel_size
in class instantiation, and also match these values when defining the step resources.
with Pipeline() as pipeline:
step = MyTask(vLLM(..., extra_kwargs={"tensor_parallel_size": 8, "pipeline_parallel_size": 2})
if __name__ == "__main__":
pipeline.run(parameters={step.name: {"resources": {"gpus": 8, "replicas": 2}}}
This is annoying and cumbersome, and could be automated.
Describe the solution you'd like
If user define resources for a Task
using vLLM
and tensor_parallel_size
and pipeline_parallel_size
haven't been set, then set values for those using the Step.resources
.