Bottleneck in GRPO training

[ZYM66]

Feature request

The default code only deploys the generation model on a single GPU, which is slow and inefficient, as other processes have to wait for VLLM completion.

GPU7 deployed the vLLM model

source code below:

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."
        )
    # vLLM is not compatible with accelerate. So we need to patch it to make sure we can (1) place the vLLM
    # model on the desired device (world_size_patch) and (2) avoid a test that is not designed for our
    # setting (profiling_patch).
    world_size_patch = patch("torch.distributed.get_world_size", return_value=1)
    profiling_patch = patch(
        "vllm.worker.worker.Worker._assert_memory_footprint_increased_during_profiling", return_value=None
    )
    with world_size_patch, profiling_patch:
        self.llm = LLM(
            model=model.name_or_path,
            device=vllm_device,
            gpu_memory_utilization=self.args.vllm_gpu_memory_utilization,
            dtype=self.args.vllm_dtype,
            # Automatic Prefix Caching caches the KV cache of existing queries, so that a new query can
            # directly reuse the KV cache if it shares the same prefix with one of the existing queries.
            # This is particularly useful here because we generate completions from the same prompts.
            enable_prefix_caching=True,
            max_model_len=self.args.vllm_max_model_len,
        )

Motivation

One GPU inference， too slow!

Your contribution

[cuiyuhao1996]

same issue!

[GuodongFan]

same!

[weizhepei]

Is it possible to specify the tensor_parallel_size for vllm in GRPO? (as suggested in this issue)

[qgallouedec]

The default code only deploys the generation model on a single GPU

I temporarily solved the problem in #2901

These are different things right?

So currently it's done this way because vLLM is quite hard to hack to make it work on selected devices. Just to have it on a single device, we need two very hacky patches.

We're very opened to contributions on this point.

[zygi]

For a custom training pipeline I implemented a hack that allows spawning vLLM instances through python/torch multiprocessing , have them bound to select devices, and support performant weight copying from the training process to vllm. Just wanted to share it here in case someone wants to adapt it to trl: https://gist.github.com/zygi/2155c6c3fefdea69237baee16efdd7e5

[SeiunSky0131]

For a custom training pipeline I implemented a hack that allows spawning vLLM instances through python/torch multiprocessing , have them bound to select devices, and support performant weight copying from the training process to vllm. Just wanted to share it here in case someone wants to adapt it to trl: https://gist.github.com/zygi/2155c6c3fefdea69237baee16efdd7e5

This is a great work! However, I am struggling with adapting it to trl. So, just want to make it happen, I want to check whether my understanding is on the right path. Currently, I think there are two places that we need to change in grpo_trainer.py.

In line 441: change to allow multiple vllm devices

        if self.use_vllm:
            if not is_vllm_available():
                raise ImportError(
                    "vLLM is not available and `use_vllm` is set to True. Please install vLLM with "
                    "`pip install vllm` to use it."
                )

            if self.accelerator.is_main_process:
                vllm_device = self.args.vllm_device
                device_type = PartialState().default_device.type
                device_module = getattr(torch, device_type)
                if vllm_device == "auto":
                    if device_module.device_count() == 1:
                        vllm_device = f"{device_type}:0"  # particular case when training with onyl 1 device: share it
                    else:
                        vllm_device = f"{device_type}:{self.accelerator.num_processes}"  # take the next GPU idx
                # Check that the requested device is available
                if (
                    vllm_device.split(":")[0] == f"{device_type}"
                    and int(vllm_device.split(":")[1]) >= device_module.device_count()
                ):       <--------- overwrite this line to support multiple vLLM devices
                    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 {device_module.device_count() - 1}`."
                    )
                # Check that the requested device is not also used for training
                if vllm_device in {f"{device_type}:{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."
                    )
                # vLLM is not compatible with accelerate. So we need to patch it to make sure we can (1) place the vLLM
                # model on the desired device (world_size_patch) and (2) avoid a test that is not designed for our
                # setting (profiling_patch).
                world_size_patch = patch("torch.distributed.get_world_size", return_value=1)
                profiling_patch = patch(
                    "vllm.worker.worker.Worker._assert_memory_footprint_increased_during_profiling", return_value=None
                )

In line 656: change to allow move models to multiple vLLM devices

def _move_model_to_vllm(self):
        with unwrap_model_for_generation(
            self.model_wrapped, 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.         <----- change this line to a list of vllm models, to ensure we get access to all vllm models.
                llm_model.load_weights(state_dict.items())          <----- use the reload_weight function provided by zygi to reload weights to all vllm models.
            # 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()

In line 727 of function _generate_and_score_completions, change to use multiple vLLM engines for generation

            if self.accelerator.is_main_process:
                # Since 'prompts' contains 'num_generations' duplicates, we first take unique prompts, and generate
                # num_generations outputs for each one. This is faster than generating outputs for each duplicate
                # prompt individually.
                ordered_set_of_prompts = list(dict.fromkeys(all_prompts_text))
                with profiling_context(self, "vLLM.generate"):
                    all_outputs = self.llm.generate(
                        ordered_set_of_prompts, sampling_params=self.sampling_params, use_tqdm=False
                    )     <----- change this line to use multiple vLLM engine to generate completions. We have to make them completely random, i.e. not duplicated across vLLM engines.
                completion_ids = []
                for outputs in all_outputs:
                    for output in outputs.outputs:
                        completion_ids.append(output.token_ids)
            else:
                completion_ids = [None] * len(all_prompts_text)
            # Broadcast the completions from the main process to all processes, ensuring each process receives its
            # corresponding slice.
            completion_ids = broadcast_object_list(completion_ids, from_process=0)
            process_slice = slice(
                self.accelerator.process_index * len(prompts),
                (self.accelerator.process_index + 1) * len(prompts),
            )
            completion_ids = completion_ids[process_slice]

If I miss some places that need changes as well, please help me point them out. Also, I would really appreciate it if someone could do the adaption !!!! Many thanks in advance!