distilabel icon indicating copy to clipboard operation
distilabel copied to clipboard

[FEATURE] Use `Step.resources` to set `tensor_parallel_size` and `pipeline_parallel_size` in `vLLM`

Open gabrielmbmb opened this issue 6 months ago • 0 comments

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.

gabrielmbmb avatar Aug 19 '24 16:08 gabrielmbmb