import ray from functools import cached_property from vllm import LLM from vllm.worker.worker import Worker class MyLLM(LLM): def __init__(self, *args, **kwargs): import os os.environ.pop("CUDA_VISIBLE_DEVICES", None) super().__init__(*args, **kwargs) class Runner: model_path: str = "/checkpoints/Qwen2.5-32B" data_path: str = "/data/math-train_7500.jsonl" num_gpus: int = 8 tensor_parallel_size: int = 4 num_prompts_per_iter: int = 128 num_rollouts_per_iter: int = 64 dtype: str = "bfloat16" gpu_memory_utilization: float = 0.9 def run(self): from vllm import SamplingParams sampling_params = SamplingParams(n=self.num_rollouts_per_iter) for batch in self.debug_data: futures = [] num_prompts_per_worker = len(batch) // len(self.workers) for indx, vllm_engine in enumerate(self.vllm_engines): prompts = batch[indx * num_prompts_per_worker: (indx + 1) * num_prompts_per_worker] futures.append( vllm_engine.generate.remote( prompts, sampling_params=sampling_params ) ) results = ray.get(futures) import pdb; pdb.set_trace() @cached_property def workers(self): from ray.util.placement_group import PlacementGroupSchedulingStrategy pg = ray.util.placement_group([{"GPU": 1, "CPU": 0}] * self.num_gpus) ray.get(pg.ready()) workers = [] for i in range(self.num_gpus // self.tensor_parallel_size): bundle_ids = list( range(i * self.tensor_parallel_size, (i + 1) * self.tensor_parallel_size) ) workers.append( ray.remote(MyLLM).options( num_gpus=0, num_cpus=0, scheduling_strategy=PlacementGroupSchedulingStrategy( placement_group=pg, placement_group_capture_child_tasks=True, placement_group_bundle_index=bundle_ids ), runtime_env={ "env_vars": { "VLLM_RAY_PER_WORKER_GPUS": str(0.1), "VLLM_RAY_BUNDLE_INDICES": ",".join(map(str, bundle_ids)), } }, ).remote( self.model_path, enforce_eager=True, tensor_parallel_size=self.tensor_parallel_size, distributed_executor_backend="ray", gpu_memory_utilization=self.gpu_memory_utilization, enable_sleep_mode=True ) ) return workers @cached_property def tokenizer(self): from transformers import AutoTokenizer return AutoTokenizer.from_pretrained(self.model_path) @cached_property def debug_data(self): import json from torch.utils.data import DataLoader data = [] with open(self.data_path, "r") as f: lines = f.read().splitlines() for line in lines: question = json.loads(line)["question"] data.append( self.tokenizer.apply_chat_template( [{"role": "user", "content": question}], add_generation_prompt=True, tokenize=False, ) ) return DataLoader(data, batch_size=self.num_prompts_per_iter) if __name__ == "__main__": runner = Runner() runner.run()