[core][bugfix] configure env var during import vllm (#12209)
Signed-off-by: youkaichao <youkaichao@gmail.com>
This commit is contained in:
@@ -19,7 +19,7 @@ from ray.util.placement_group import placement_group
|
||||
from ray.util.scheduling_strategies import PlacementGroupSchedulingStrategy
|
||||
from transformers import AutoModelForCausalLM
|
||||
|
||||
from vllm import LLM, SamplingParams, configure_as_vllm_process
|
||||
from vllm import LLM, SamplingParams
|
||||
from vllm.utils import get_ip, get_open_port
|
||||
from vllm.worker.worker import Worker
|
||||
|
||||
@@ -98,12 +98,7 @@ class MyLLM(LLM):
|
||||
"""
|
||||
Start the training process, here we use huggingface transformers
|
||||
as an example to hold a model on GPU 0.
|
||||
|
||||
It is important for all the processes outside of vLLM to call
|
||||
`configure_as_vllm_process` to set some common environment variables
|
||||
the same as vLLM workers.
|
||||
"""
|
||||
configure_as_vllm_process()
|
||||
|
||||
train_model = AutoModelForCausalLM.from_pretrained("facebook/opt-125m")
|
||||
train_model.to("cuda:0")
|
||||
|
||||
Reference in New Issue
Block a user