Skip to content

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

@gabrielmbmb

Description

@gabrielmbmb

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.

Metadata

Metadata

Assignees

Labels

enhancementNew feature or request

Type

No type

Projects

Status

No status

Relationships

None yet

Development

No branches or pull requests

Issue actions