generated from fastai/nbdev_template
-
Notifications
You must be signed in to change notification settings - Fork 2.1k
Open
Labels
Description
trl/trl/trainer/grpo_trainer.py
Lines 439 to 461 in e5ae703
| if self.accelerator.is_main_process: | |
| vllm_device = self.args.vllm_device | |
| if vllm_device == "auto": | |
| if torch.cuda.device_count() == 1: | |
| vllm_device = "cuda:0" # particular case when training with onyl 1 GPU: share it | |
| else: | |
| vllm_device = f"cuda:{self.accelerator.num_processes}" # take the next GPU idx | |
| # Check that the requested device is available | |
| if vllm_device.split(":")[0] == "cuda" and int(vllm_device.split(":")[1]) >= torch.cuda.device_count(): | |
| raise ValueError( | |
| f"The requested device for vllm ({vllm_device}) is not available. You are likely using vLLM " | |
| "without restricting the number of GPUs for training. Set the `--num_processes` argument to a " | |
| "value lower than the number of GPUs available on your machine—typically, reducing it by one " | |
| f"is sufficient. In your case: `--num_processes {torch.cuda.device_count() - 1}`." | |
| ) | |
| # Check that the requested device is not also used for training | |
| if vllm_device in {f"cuda:{idx}" for idx in range(self.accelerator.num_processes)}: | |
| warnings.warn( | |
| f"The requested device {vllm_device} is also being used for training. For higher throughput " | |
| "and to avoid out-of-memory errors, it is recommended to use a dedicated device for vLLM. " | |
| "If this is intentional, you may ignore this warning but should adjust " | |
| "`vllm_gpu_memory_utilization` accordingly." | |
| ) |
In the current GRPO implementation, VLLM can only run on a single GPU, which becomes a performance bottleneck. For example, in an 8-GPU setup, the remaining 7 GPUs have to wait for 1 GPU to complete inference, and it also can't accommodate larger models.
How can we enable VLLM to run on multiple GPUs? The only concern is that we need to figure out a way to update the parameters across multiple GPUs each time the model is reloaded:
trl/trl/trainer/grpo_trainer.py
Lines 624 to 653 in e5ae703
| def _move_model_to_vllm(self): | |
| with unwrap_model_for_generation( | |
| self.model, self.accelerator, gather_deepspeed3_params=self.args.ds3_gather_for_generation | |
| ) as unwrapped_model: | |
| if is_compiled_module(unwrapped_model): | |
| unwrapped_model = unwrapped_model._orig_mod | |
| if is_peft_model(unwrapped_model): | |
| unwrapped_model.merge_adapter() | |
| state_dict = unwrapped_model.state_dict() | |
| # Remove base_model and base_layer prefixes | |
| state_dict = { | |
| k.removeprefix("base_model.model.").replace(".base_layer", ""): v for k, v in state_dict.items() | |
| } | |
| # Remove values with adapter prefix (example: "_lora") | |
| state_dict = {k: v for k, v in state_dict.items() if unwrapped_model.prefix not in k} | |
| # When module to save, remove its prefix and discard the original module | |
| state_dict = { | |
| k.replace("modules_to_save.default.", ""): v | |
| for k, v in state_dict.items() | |
| if "original_module" not in k | |
| } | |
| else: | |
| state_dict = unwrapped_model.state_dict() | |
| if self.accelerator.is_main_process: | |
| llm_model = self.llm.llm_engine.model_executor.driver_worker.model_runner.model | |
| llm_model.load_weights(state_dict.items()) | |
| # Unmerge the adapter to restore the model to its original state. | |
| # This must be done after loading weights to ensure they correspond to the merged state. | |
| if is_peft_model(unwrapped_model): | |
| unwrapped_model.unmerge_adapter() |
Hasuer, SeiunSky0131, edwardzjl, iamansinha, muupan and 3 more