Repository: huggingface/trl
Pull request #3394: 🧑‍🤝‍🧑 Co-Locating vLLM w/ training to for higher throughput and GPU utilization
Author: login: toslali-ibm
Created: 2025-04-30T14:06:38Z | Merged: 2025-05-01T23:17:26Z

# What does this PR do?
Enables colocating vLLM with training in each GPU to improve utilization and throughput.

Fixes #3064 and #3113
Addresses: #3195, #2971, #2922, #2887 etc.

## Enabler:
vLLM (version >0.7.3) introduced support for an [external launcher](https://github.com/vllm-project/vllm/pull/12071), allowing vLLM processes to run alongside other workloads on the same GPU.

## Benefits:
- Faster Inference: Speeds up GRPO training by reducing inference latency via parallel prompt processing (each vLLM works on their device's batch)
- Better GPU Efficiency: Frees up GPU resources by removing the need for a dedicated vLLM server. Multiple vLLM instances can now share GPUs with training jobs (reducing GPU idle time)
- Supports TP + DP
- Ray-less solution


## Testing vllm colocation

Run it w/ the following: 
`VLLM_USE_V1=0 ACCELERATE_LOG_LEVEL=info CUDA_VISIBLE_DEVICES=0,1,2,3,4,5,6,7 accelerate launch --config_file recipes/accelerate_configs/zero3.yaml --num_processes=8 -m open_r1.grpo --config config_tpcoloc.yaml`
* change `vllm_colocation` in the config to the sharding you would like. 
* E.g., If `vllm_colocation=1`, model is not sharded, each GPU holds a full copy of the model. 
* `vllm_colocation=2`, model is sharded by two, and groups: [0,1], [2,3], [4,5], [6,7]. 

<details>
  <summary>Click to view config.yaml</summary>

  ```yaml
# Model arguments
model_name_or_path: Qwen/Qwen2.5-Math-1.5B
model_revision: main
torch_dtype: bfloat16
attn_implementation: flash_attention_2

# Data training arguments
dataset_name: DigitalLearningGmbH/MATH-lighteval
dataset_config: default
dataset_prompt_column: problem
system_prompt: "You are a helpful AI Assistant, designed to provided well-reasoned and detailed responses. You FIRST think about the reasoning process as an internal monologue and then provide the user with the answer. The reasoning process MUST BE enclosed within <think> and </think> tags."

# GRPO trainer config
bf16: true
use_vllm: true
vllm_colocation: 2
vllm_gpu_memory_utilization: 0.3
vllm_max_model_len: 2048
do_eval: false
gradient_accumulation_steps: 1
gradient_checkpointing: true
gradient_checkpointing_kwargs:
use_reentrant: false
learning_rate: 3.0e-06
log_completions: false
log_level: info
logging_first_step: true
logging_steps: 1
logging_strategy: steps
lr_scheduler_type: cosine
max_prompt_length: 512
max_completion_length: 1024
max_steps: 50
num_generations: 8
num_train_epochs: 1
overwrite_output_dir: true
# per_device_eval_batch_size: 16
per_device_train_batch_size: 16
push_to_hub: false
report_to:
- wandb
reward_funcs:
- accuracy
- format
reward_weights:
- 1.0
- 1.0
save_strategy: steps
save_steps: 100
save_total_limit: 1
seed: 42
warmup_ratio: 0.1

  ```
</details>


## Sanity check
* GRPO results `Qwen/Qwen2.5-Math-1.5B` on `DigitalLearningGmbH/MATH-lighteval` dataset (as shown above) using both plain TRL (w/ vLLM server) and colocated TRL (w/ TP =1,TP =2, and TP =4); The rewards are identical.

![TRL-PR-coloc](https://github.com/user-attachments/assets/cf1bd7c7-5cfd-415f-94f5-54f1a5b0bd99)


## Before submitting
- [ ] This PR fixes a typo or improves the docs (you can dismiss the other checks if that's the case).
- [x] Did you read the [contributor guideline](https://github.com/huggingface/trl/blob/main/CONTRIBUTING.md#create-a-pull-request),
      Pull Request section?
- [x] Was this discussed/approved via a GitHub issue? Please add a link
      to it if that's the case.
- [ ] Did you make sure to update the documentation with your changes? Here are the
      [documentation guidelines](https://github.com/huggingface/trl/tree/main/docs).
- [ ] Did you write any new necessary tests?


## Who can review?

Anyone in the community is free to review the PR once the tests have passed. Feel free to tag
members/contributors who may be interested in your PR.


New version of #3162 (incorporated @qgallouedec 's comments)


CC @fabianlim 

Linked issue #3064: Enable External Launcher Support for vLLM in TRL for Efficient GRPO Training
### Feature request

vLLM has introduced support for an [external launcher](https://github.com/vllm-project/vllm/pull/12071), enabling vLLM processes to be co-located with other processes, such as training. By running multiple vLLM instances alongside the training process, we can improve inference process, reducing the time required for GRPO training. I propose adding an option in TRL to spawn vLLM processes per GPU using its external launcher.

### Motivation

Efficient GRPO relies heavily on fast and scalable inference. Currently, inference and training processes are executed separately, introducing bottlenecks that slow down training. The ideal is to enable multiple vLLM instances in the training process as done by the others like [OpenRLHF](https://github.com/OpenRLHF/OpenRLHF) and [VERL](https://github.com/volcengine/verl).

With vLLM's newly introduced external launcher ([PR #12071](https://github.com/vllm-project/vllm/pull/12071)), it is now possible to co-locate vLLM instances with training processes, allowing spawn vLLM instances to run per GPU. This reduces inference latency, leading to shorter training durations.

By integrating vLLM’s external launcher into TRL, we can enhance distributed inference efficiency and accelerate GRPO training, making large-scale reinforcement learning more practical and scalable.

### Your contribution

Modify GRPO_trainer to enable initialization of vllm via external launcher - if TRL flag (such as self.args.external_launcher) is provided. We are considering doing a RAY-less version, in which case the changes can be quite minimal. 

Diff:
diff --git a/docs/source/grpo_trainer.md b/docs/source/grpo_trainer.md
index 96d88a090ef..a5a8fdd549e 100644
--- a/docs/source/grpo_trainer.md
+++ b/docs/source/grpo_trainer.md
@@ -170,26 +170,59 @@ A higher value means more tokens are clipped, which constrains how much the poli
 
 ### Speed up training with vLLM-powered generation
 
-Generation is often the main bottleneck that makes training slow with online methods. To accelerate generation, you can use [vLLM](https://github.com/vllm-project/vllm), a library that enables fast generation. To enable it, first install the package with
-
+Generation is often the main bottleneck when training with online methods. To accelerate generation, you can use [vLLM](https://github.com/vllm-project/vllm), a high-throughput, low-latency inference engine for LLMs. To enable it, first install the package with
 ```shell
 pip install trl[vllm]
 ```
 
-Then, start the vLLM server with the desired model:
+We support two ways of using vLLM during training: **server mode** and **colocate mode**.
 
-```bash
-trl vllm-serve --model <model_name>
-```
+#### 🔌 Option 1: Server mode
+
+In this mode, vLLM runs in a separate process (and using separate GPUs) and communicates with the trainer via HTTP. This is ideal if you have dedicated GPUs for inference.
+
+1. **Start the vLLM server**:
+   ```bash
+   trl vllm-serve --model <model_name>
+   ```
 
-Then, pass `use_vllm=True` in the training arguments and run the training script:
+2. **Enable server mode in your training script**:
+   ```python
+   from trl import GRPOConfig
+
+   training_args = GRPOConfig(
+       ...,
+       use_vllm=True,
+       vllm_mode="server",  # default value, can be omitted
+   )
+   ```
+
+<Tip warning={true}>
+
+Make sure that the server is using different GPUs than the trainer, otherwise you may run into NCCL errors. You can specify the GPUs to use with the `CUDA_VISIBLE_DEVICES` environment variable.
+
+</Tip>
+
+#### 🧩 Option 2: Colocate mode
+
+In this mode, vLLM runs inside the trainer process and shares GPU memory with the training model. This avoids launching a separate server and can improve GPU utilization, but may lead to memory contention on the training GPUs.
 
 ```python
 from trl import GRPOConfig
 
-training_args = GRPOConfig(..., use_vllm=True)
+training_args = GRPOConfig(
+    ...,
+    use_vllm=True,
+    vllm_mode="colocate",
+)
 ```
 
+<Tip>
+
+Depending on the model size and the overall GPU memory requirements for training, you may need to adjust the `vllm_gpu_memory_utilization` parameter in [`GRPOConfig`] to avoid underutilization or out-of-memory errors.
+
+</Tip>
+
 For more information, see [Speeding up training with vLLM](speeding_up_training#vllm-for-fast-generation-in-online-methods).
 
 ### GRPO at scale: train a 70B+ Model on multiple nodes
diff --git a/trl/trainer/grpo_config.py b/trl/trainer/grpo_config.py
index b449165a5f1..48923232c82 100644
--- a/trl/trainer/grpo_config.py
+++ b/trl/trainer/grpo_config.py
@@ -85,17 +85,39 @@ class GRPOConfig(TrainingArguments):
         > Parameters that control generation acceleration powered by vLLM
 
         use_vllm (`bool`, *optional*, defaults to `False`):
-            Whether to use vLLM for generating completions. If set to `True`, ensure that a GPU is kept unused for
-            training, as vLLM will require one for generation. vLLM must be installed (`pip install vllm`).
+            Whether to use vLLM for generating completions. If set to `True`, the trainer will use vLLM for generation
+            instead of the default model.generate(). Requires `vllm` to be installed.
+        vllm_mode (`str`, *optional*, defaults to `"server"`):
+            Mode to use for vLLM integration when `use_vllm` is set to `True`. Must be one of `"server"` or
+            `"colocate"`.
+
+            - `"server"`: The trainer will send generation requests to a separate vLLM server. Make sure a TRL vLLM
+              server is running (start with `trl vllm-serve`).
+            - `"colocate"`: vLLM will run in the same process and share the training GPUs. This avoids the need for a
+              separate server but may cause resource contention with training.
+        vllm_guided_decoding_regex (`str` or `None`, *optional*, defaults to `None`):
+            Regex for vLLM guided decoding. If `None` (default), guided decoding is disabled.
+
+        > Parameters that control the vLLM server (only used when `vllm_mode` is `"server"`)
+
         vllm_server_host (`str`, *optional*, defaults to `"0.0.0.0"`):
             Host of the vLLM server to connect to.
         vllm_server_port (`int`, *optional*, defaults to `8000`):
             Port of the vLLM server to connect to.
-        vllm_server_timeout (`float`, *optional*, defaults to `120.0`):
+        vllm_server_timeout (`float`, *optional*, defaults to `240.0`):
             Total timeout duration in seconds to wait for the vLLM server to be up. If the server is not up after the
             timeout, a `ConnectionError` is raised.
-        vllm_guided_decoding_regex (`str` or `None`, *optional*, defaults to `None`):
-            Regex for vLLM guided decoding. If `None` (default), guided decoding is disabled.
+
+        > Parameters that control colocated vLLM execution (only used when `vllm_mode` is `"colocate"`)
+
+        vllm_gpu_memory_utilization (`float`, *optional*, defaults to `0.3`):
+            Control the GPU memory utilization for vLLM. This setting only applies when `vllm_mode` is set to
+            `"colocate"`. If you are using `vllm_mode="server"`, this parameter must be passed separately when
+            launching the vLLM server via the `--vllm_gpu_memory_utilization` flag.
+        vllm_tensor_parallel_size (`int`, *optional*, defaults to `1`):
+            Control the tensor parallel size for vLLM. This setting only applies when `vllm_mode` is set to
+            `"colocate"`. If you are using `vllm_mode="server"`, this parameter must be passed separately when
+            launching the vLLM server via the `--vllm_tensor_parallel_size` flag.
 
         > Parameters that control the training
 
@@ -272,10 +294,26 @@ class GRPOConfig(TrainingArguments):
     use_vllm: bool = field(
         default=False,
         metadata={
-            "help": "Whether to use vLLM for generating completions. If set to `True`, ensure that a vLLM server is "
-            "running. To run the server, install vLLM (`pip install vllm`) and run `trl vllm-serve`."
+            "help": "Whether to use vLLM for generating completions. If set to `True`, the trainer will use vLLM for "
+            "generation instead of the default model.generate(). Requires `vllm` to be installed."
         },
     )
+    vllm_mode: str = field(
+        default="server",
+        metadata={
+            "help": "Mode to use for vLLM integration when `use_vllm` is set to `True`. Must be one of `server` or "
+            "`'colocate'`. `'server'`: The trainer will send generation requests to a separate vLLM server. Make sure a "
+            "TRL vLLM server is running (start with `trl vllm-serve`). `'colocate'`: vLLM will run in the same "
+            "process and share the training GPUs. This avoids the need for a separate server but may cause resource "
+            "contention with training."
+        },
+    )
+    vllm_guided_decoding_regex: Optional[str] = field(
+        default=None,
+        metadata={"help": "Regex for vLLM guided decoding. If `None` (default), guided decoding is disabled."},
+    )
+
+    # Parameters that control the vLLM server (only used when `vllm_mode` is `"server"`)
     vllm_server_host: str = field(
         default="0.0.0.0",
         metadata={"help": "Host of the vLLM server to connect to."},
@@ -291,9 +329,23 @@ class GRPOConfig(TrainingArguments):
             "after the timeout, a `ConnectionError` is raised."
         },
     )
-    vllm_guided_decoding_regex: Optional[str] = field(
-        default=None,
-        metadata={"help": "Regex for vLLM guided decoding. If `None` (default), guided decoding is disabled."},
+
+    # Parameters that control colocated vLLM execution (only used when `vllm_mode` is `"colocate"`)
+    vllm_gpu_memory_utilization: float = field(
+        default=0.3,
+        metadata={
+            "help": "Control the GPU memory utilization for vLLM. This setting only applies when `vllm_mode` is set "
+            "to `'colocate'`. If you are using `vllm_mode='server'`, this parameter must be passed separately when "
+            "launching the vLLM server via the `--vllm_gpu_memory_utilization` flag."
+        },
+    )
+    vllm_tensor_parallel_size: int = field(
+        default=1,
+        metadata={
+            "help": "Control the tensor parallel size for vLLM. This setting only applies when `vllm_mode` is set "
+            "to `'colocate'`. If you are using `vllm_mode='server'`, this parameter must be passed separately when "
+            "launching the vLLM server via the `--vllm_tensor_parallel_size` flag."
+        },
     )
 
     # Parameters that control the training
diff --git a/trl/trainer/grpo_trainer.py b/trl/trainer/grpo_trainer.py
index f839ebc6b12..927472c910b 100644
--- a/trl/trainer/grpo_trainer.py
+++ b/trl/trainer/grpo_trainer.py
@@ -69,6 +69,10 @@
 if is_liger_kernel_available():
     from liger_kernel.chunked_loss import LigerFusedLinearGRPOLoss
 
+if is_vllm_available():
+    from vllm import LLM, SamplingParams
+    from vllm.sampling_params import GuidedDecodingParams
+
 if is_wandb_available():
     import wandb
 
@@ -471,6 +475,9 @@ def data_collator(features):  # No data collation is needed in GRPO
         self.min_p = args.min_p
         self.repetition_penalty = args.repetition_penalty
         self.use_vllm = args.use_vllm
+        self.vllm_mode = args.vllm_mode
+        self.vllm_gpu_memory_utilization = args.vllm_gpu_memory_utilization  # only applies to colocation mode
+        self.vllm_tensor_parallel_size = args.vllm_tensor_parallel_size  # only applies to colocation mode
         self.use_liger_loss = args.use_liger_loss
         self.loss_type = args.loss_type
         self.scale_rewards = args.scale_rewards
@@ -509,14 +516,15 @@ def data_collator(features):  # No data collation is needed in GRPO
         # This acts as a flag to indicate that the warning has already been issued.
         model.warnings_issued["estimate_tokens"] = True
 
-        # redirect the model.module forward to the model forward to ensure pre-forward hooks are called
-        self._forward_redirection = _ForwardRedirection()
         if self.use_liger_loss:
             if not is_liger_kernel_available():
                 raise ImportError(
                     "Liger is required to use `liger_loss` as the GRPO loss. Run `pip install liger-kernel`."
                 )
 
+            # Redirect the model.module forward to the model forward to ensure pre-forward hooks are called
+            self._forward_redirection = _ForwardRedirection()
+
             self.liger_grpo_loss = LigerFusedLinearGRPOLoss(
                 beta=self.beta,
                 epsilon_low=self.epsilon_low,
@@ -617,12 +625,44 @@ def data_collator(features):  # No data collation is needed in GRPO
                     "`pip install vllm` to use it."
                 )
 
-            if self.accelerator.is_main_process:
+            if self.vllm_mode == "server" and self.accelerator.is_main_process:
                 self.vllm_client = VLLMClient(
                     args.vllm_server_host, args.vllm_server_port, connection_timeout=args.vllm_server_timeout
                 )
                 self.vllm_client.init_communicator()
 
+            elif self.vllm_mode == "colocate":
+                # Make sure vllm_tensor_parallel_size group size evenly divides the world size - each group should have
+                # the same number of ranks
+                if not self.accelerator.num_processes % self.vllm_tensor_parallel_size == 0:
+                    raise ValueError(
+                        f"vllm_tensor_parallel_size ({self.vllm_tensor_parallel_size}) must divide world size "
+                        f"({self.accelerator.num_processes}) evenly."
+                    )
+
+                if self.vllm_tensor_parallel_size > 1:
+                    # Create subgroups of ranks for TP, each group with `vllm_tensor_parallel_size` ranks.
+                    # For example, if world_size=8 and vllm_tensor_parallel_size=2 → groups: [0,1], [2,3], [4,5], [6,7]
+                    self.tp_group, _ = torch.distributed.new_subgroups_by_enumeration(
+                        [
+                            list(range(i * self.vllm_tensor_parallel_size, (i + 1) * self.vllm_tensor_parallel_size))
+                            for i in range(self.accelerator.num_processes // self.vllm_tensor_parallel_size)
+                        ]
+                    )
+
+                self.llm = LLM(
+                    model=model.name_or_path,
+                    tensor_parallel_size=args.vllm_tensor_parallel_size,
+                    gpu_memory_utilization=self.vllm_gpu_memory_utilization,
+                    max_num_seqs=self.args.per_device_train_batch_size
+                    * self.vllm_tensor_parallel_size
+                    * self.args.gradient_accumulation_steps,
+                    max_model_len=self.max_prompt_length + self.max_completion_length,
+                    distributed_executor_backend="external_launcher",
+                    # Feed identical seed for tp groups to ensure sampling results are the same across workers
+                    seed=self.accelerator.process_index // self.vllm_tensor_parallel_size,
+                )
+
             # vLLM specific sampling arguments
             self.guided_decoding_regex = args.vllm_guided_decoding_regex
 
@@ -825,6 +865,34 @@ def _get_per_token_logps(self, model, input_ids, attention_mask, logits_to_keep,
             all_logps.append(logps)
         return torch.cat(all_logps, dim=0)
 
+    def _sync_fsdp_params_to_vllm(self, module: nn.Module, prefix: str = "", visited=None):
+        """Memory-efficient post-order traversal of FSDP modules to extract full parameters and sync with vLLM."""
+        if visited is None:
+            visited = set()
+
+        for child_name, child_module in module.named_children():
+            child_prefix = f"{prefix}.{child_name}" if prefix else child_name
+            self._sync_fsdp_params_to_vllm(
+                child_module, prefix=child_prefix, visited=visited
+            )  # recurse into the child
+
+        if isinstance(module, FSDP):
+            with FSDP.summon_full_params(module, recurse=False, writeback=False):
+                for param_name, param in module.named_parameters():
+                    full_name = f"{prefix}.{param_name}" if prefix else param_name
+                    for extra in ("_fsdp_wrapped_module.", "_checkpoint_wrapped_module."):
+                        full_name = full_name.replace(extra, "")
+
+                    if full_name in visited:
+                        continue  # skip FSDP subtrees already traversed
+                    visited.add(full_name)
+
+                    if self.vllm_mode == "server" and self.accelerator.is_main_process:
+                        self.vllm_client.update_named_param(full_name, param.data)
+                    elif self.vllm_mode == "colocate":
+                        llm_model = self.llm.llm_engine.model_executor.driver_worker.model_runner.model
+                        llm_model.load_weights([(full_name, param.data)])
+
     @profiling_decorator
     def _move_model_to_vllm(self):
         # For DeepSpeed ZeRO-3 and FSDP, we need to gather all parameters before operations
@@ -837,54 +905,20 @@ def _move_model_to_vllm(self):
         else:
             gather_if_zero3 = nullcontext
 
-        def post_order_fsdp_processing(module: nn.Module, prefix: str = "", visited=None):
-            """memory-efficient module gather"""
-            extra_names = ["_fsdp_wrapped_module.", "_checkpoint_wrapped_module."]
-            if visited is None:
-                visited = set()
-
-            for child_name, child_module in module.named_children():
-                if prefix == "":
-                    child_prefix = child_name
-                else:
-                    child_prefix = f"{prefix}.{child_name}"
-
-                # Recurse into the child
-                post_order_fsdp_processing(child_module, prefix=child_prefix, visited=visited)
-
-            if isinstance(module, FSDP):
-                with FSDP.summon_full_params(module, recurse=False, writeback=False):
-                    for param_name, param in module.named_parameters():
-                        if prefix == "":
-                            full_name = param_name
-                        else:
-                            full_name = f"{prefix}.{param_name}"
-
-                        for extra in extra_names:
-                            full_name = full_name.replace(extra, "")
-
-                        if full_name in visited:
-                            # skip FSDP subtrees already traversed
-                            continue
-
-                        visited.add(full_name)
-
-                        if self.accelerator.is_main_process:
-                            self.vllm_client.update_named_param(full_name, param.data)
-
         if is_peft_model(self.model):
-            # With PEFT and FSDP/DeepSpeed ZeRO Stage 3, we must gather the full model at once before merging, as merging
-            # adapters in a sharded manner is not supported.
+            # With PEFT and FSDP/DeepSpeed ZeRO Stage 3, we must gather the full model at once before merging, as
+            # merging adapters in a sharded manner is not supported.
+            # TODO: does this work with FSDP?
             with gather_if_zero3(list(self.model.parameters())):
-                if self.is_fsdp_enabled:
+                self.model.merge_adapter()
+
+                # Update vLLM weights while parameters are gathered
+                if self.is_fsdp_enabled:  # note if using FSDP, gather_if_zero3 is nullcontext
                     # Update vLLM weights while parameters are gathered
                     # For PEFT with FSDP we need to use the memory efficient post-order traversal
-                    self.model.merge_adapter()
-                    post_order_fsdp_processing(self.model)
-                    self.model.unmerge_adapter()
+                    self._sync_fsdp_params_to_vllm(self.model)
                 else:
-                    # DeepSpeed ZeRO-3 with PEFT (not using FSDP)
-                    self.model.merge_adapter()
+                    # DeepSpeed ZeRO-3 with PEFT
                     for name, param in self.model.named_parameters():
                         # When using PEFT, we need to recover the original parameter name and discard some parameters
                         name = name.removeprefix("base_model.model.").replace(".base_layer", "")
@@ -895,24 +929,32 @@ def post_order_fsdp_processing(module: nn.Module, prefix: str = "", visited=None
                             continue
                         name = name.replace("modules_to_save.default.", "")
 
-                        if self.accelerator.is_main_process:
+                        if self.vllm_mode == "server" and self.accelerator.is_main_process:
                             self.vllm_client.update_named_param(name, param.data)
-                    self.model.unmerge_adapter()
+                        elif self.vllm_mode == "colocate":
+                            llm_model = self.llm.llm_engine.model_executor.driver_worker.model_runner.model
+                            llm_model.load_weights([(name, param.data)])
+                # Unmerge adapters while parameters are still gathered
+                self.model.unmerge_adapter()
+                # Parameters will automatically be repartitioned when exiting the context
         else:
-            # For non-PEFT models
+            # For non-PEFT models, simply gather (if needed) and update each parameter individually.
             if self.is_fsdp_enabled:
-                # Use memory-efficient post-order traversal for FSDP
-                post_order_fsdp_processing(self.model)
-            elif zero_stage_3:
-                # For DeepSpeed ZeRO-3, simply gather and update each parameter individually
+                self._sync_fsdp_params_to_vllm(self.model)  # use memory-efficient post-order traversal for FSDP
+            else:
                 for name, param in self.model.named_parameters():
                     with gather_if_zero3([param]):
-                        if self.accelerator.is_main_process:
+                        if self.vllm_mode == "server" and self.accelerator.is_main_process:
                             self.vllm_client.update_named_param(name, param.data)
+                        elif self.vllm_mode == "colocate":
+                            llm_model = self.llm.llm_engine.model_executor.driver_worker.model_runner.model
+                            llm_model.load_weights([(name, param.data)])
 
-        # Reset cache on main process
-        if self.accelerator.is_main_process:
+        # Reset cache on vLLM
+        if self.vllm_mode == "server" and self.accelerator.is_main_process:
             self.vllm_client.reset_prefix_cache()
+        elif self.vllm_mode == "colocate":
+            self.llm.reset_prefix_cache()
 
     @profiling_decorator
     def _prepare_inputs(
@@ -966,40 +1008,78 @@ def _generate_and_score_completions(
 
         # Generate completions using either vLLM or regular generation
         if self.use_vllm:
-            # First, have main process load weights if needed
+            # First, update the vLLM weights if needed
             if self.state.global_step != self._last_loaded_step:
                 self._move_model_to_vllm()
                 self._last_loaded_step = self.state.global_step
 
             # Generate completions using vLLM: gather all prompts and use them in a single call in the main process
-            all_prompts_text = gather_object(prompts_text)
-            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 = all_prompts_text[:: self.num_generations]
+            if self.vllm_mode == "server":
+                all_prompts_text = gather_object(prompts_text)
+                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 = all_prompts_text[:: self.num_generations]
+                    with profiling_context(self, "vLLM.generate"):
+                        completion_ids = self.vllm_client.generate(
+                            prompts=ordered_set_of_prompts,
+                            n=self.num_generations,
+                            repetition_penalty=self.repetition_penalty,
+                            temperature=self.temperature,
+                            top_p=self.top_p,
+                            top_k=-1 if self.top_k is None else self.top_k,
+                            min_p=0.0 if self.min_p is None else self.min_p,
+                            max_tokens=self.max_completion_length,
+                            guided_decoding_regex=self.guided_decoding_regex,
+                        )
+                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]
+
+            # Generate completions using colocated vLLM instances: each device holds vLLM copy and work on their own batch of prompts
+            elif self.vllm_mode == "colocate":
+                if self.guided_decoding_regex:
+                    guided_decoding = GuidedDecodingParams(backend="outlines", regex=self.guided_decoding_regex)
+                else:
+                    guided_decoding = None
+                sampling_params = SamplingParams(
+                    n=1,  # vLLM on each GPU generates only 1 in colocate mode
+                    repetition_penalty=self.repetition_penalty,
+                    temperature=self.temperature,
+                    top_p=self.top_p,
+                    top_k=-1 if self.top_k is None else self.top_k,
+                    min_p=0.0 if self.min_p is None else self.min_p,
+                    max_tokens=self.max_completion_length,
+                    guided_decoding=guided_decoding,
+                )
+
+                if self.vllm_tensor_parallel_size > 1:
+                    # Gather prompts from all ranks in the TP group and flatten.
+                    # Each rank starts with its own prompts; after gathering, all ranks see the full group set.
+                    orig_size = len(prompts_text)
+                    gathered_prompts = [None for _ in range(self.vllm_tensor_parallel_size)]
+                    torch.distributed.all_gather_object(gathered_prompts, prompts_text, group=self.tp_group)
+                    prompts_text = [p for sublist in gathered_prompts for p in sublist]
+
                 with profiling_context(self, "vLLM.generate"):
-                    completion_ids = self.vllm_client.generate(
-                        prompts=ordered_set_of_prompts,
-                        n=self.num_generations,
-                        repetition_penalty=self.repetition_penalty,
-                        temperature=self.temperature,
-                        top_p=self.top_p,
-                        top_k=-1 if self.top_k is None else self.top_k,
-                        min_p=0.0 if self.min_p is None else self.min_p,
-                        max_tokens=self.max_completion_length,
-                        guided_decoding_regex=self.guided_decoding_regex,
-                    )
-            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]
+                    all_outputs = self.llm.generate(prompts_text, sampling_params=sampling_params, use_tqdm=False)
+
+                completion_ids = [output.token_ids for outputs in all_outputs for output in outputs.outputs]
+
+                if self.vllm_tensor_parallel_size > 1:
+                    # Slice completions for this rank within its TP group.
+                    # Each rank generates all outputs — we keep only our share.
+                    local_rank_in_group = torch.distributed.get_rank(group=self.tp_group)
+                    tp_slice = slice(local_rank_in_group * orig_size, (local_rank_in_group + 1) * orig_size)
+                    completion_ids = completion_ids[tp_slice]
 
             # Pad the completions, and concatenate them with the prompts
             completion_ids = [torch.tensor(ids, device=device) for ids in completion_ids]


Reviews:
user: login: toslali-ibm
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-04-30T15:59:53Z
user: login: toslali-ibm
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-04-30T16:05:11Z
user: login: toslali-ibm
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-04-30T16:06:20Z
user: login: qgallouedec
state: COMMENTED
author_association: MEMBER
submitted_at: 2025-05-01T19:00:42Z
user: login: toslali-ibm
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-05-01T19:03:44Z
user: login: qgallouedec
state: COMMENTED
author_association: MEMBER
submitted_at: 2025-05-01T20:58:04Z
user: login: qgallouedec
state: COMMENTED
author_association: MEMBER
submitted_at: 2025-05-01T20:59:37Z
user: login: qgallouedec
body: some changements not related to this PR, but that I did while merge main to your branch
state: COMMENTED
author_association: MEMBER
submitted_at: 2025-05-01T22:02:58Z
user: login: qgallouedec
body: Wonderful work @toslali-ibm 
Can you please take a final look before I merge?
state: APPROVED
author_association: MEMBER
submitted_at: 2025-05-01T22:59:17Z

Review comments:
diff_hunk: @@ -903,7 +935,10 @@ def _move_model_to_vllm(self):
                         continue
                     name = name.replace("modules_to_save.default.", "")
 
-                    if self.accelerator.is_main_process:
+                    if self.args.vllm_colocation:
+                        llm_model = self.llm.llm_engine.model_executor.driver_worker.model_runner.model
path: trl/trainer/grpo_trainer.py
user: login: toslali-ibm
body: Move this outside the loop.
created_at: 2025-04-30T15:59:53Z
author_association: CONTRIBUTOR
diff_hunk: @@ -617,7 +619,37 @@ def data_collator(features):  # No data collation is needed in GRPO
                     "`pip install vllm` to use it."
                 )
 
-            if self.accelerator.is_main_process:
+            if self.args.vllm_colocation:
+                # Ensure vllm_colocation TP value is valid (at least 1)
+                assert self.args.vllm_colocation >= 1, "vllm_colocation must be greater than 0"
+                # Make sure vllm_colocation TP group size evenly divides the world size - each group should have the same number of ranks
+                assert self.accelerator.num_processes % self.args.vllm_colocation == 0, (
+                    f"TP size of vllm_colocation ({self.args.vllm_colocation}) must divide world size "
+                    f"({self.accelerator.num_processes}) evenly."
+                )
+
+                if self.args.vllm_colocation > 1:
+                    # Create subgroups of ranks for TP, each group with `vllm_colocation` ranks.
+                    # For example, if world_size=8 and vllm_colocation=2 → groups: [0,1], [2,3], [4,5], [6,7]
+                    self.tp_group, _  = torch.distributed.new_subgroups_by_enumeration(
+                        [
+                            list(range(i*self.args.vllm_colocation, (i+1) * self.args.vllm_colocation)) 
+                            for i in range(self.accelerator.num_processes // self.args.vllm_colocation)
+                        ]
+                    )
+
+                self.llm = LLM(
+                    model=model.name_or_path,
+                    # device=f"{device_type}:{self.accelerator.process_index}", # ToDo: we do not need to set the device
+                    tensor_parallel_size=args.vllm_colocation, 
+                    gpu_memory_utilization=self.args.vllm_gpu_memory_utilization,
+                    # max_num_seqs=self.args.per_device_train_batch_size * self.args.vllm_colocation, # ToDo: this should be multiplied by gradient_accumulation_steps 
+                    max_model_len=self.args.vllm_max_model_len,
path: trl/trainer/grpo_trainer.py
user: login: toslali-ibm
body: set this to prompt + generation length
created_at: 2025-04-30T16:05:11Z
author_association: CONTRIBUTOR
diff_hunk: @@ -617,7 +619,37 @@ def data_collator(features):  # No data collation is needed in GRPO
                     "`pip install vllm` to use it."
                 )
 
-            if self.accelerator.is_main_process:
+            if self.args.vllm_colocation:
+                # Ensure vllm_colocation TP value is valid (at least 1)
+                assert self.args.vllm_colocation >= 1, "vllm_colocation must be greater than 0"
+                # Make sure vllm_colocation TP group size evenly divides the world size - each group should have the same number of ranks
+                assert self.accelerator.num_processes % self.args.vllm_colocation == 0, (
+                    f"TP size of vllm_colocation ({self.args.vllm_colocation}) must divide world size "
+                    f"({self.accelerator.num_processes}) evenly."
+                )
+
+                if self.args.vllm_colocation > 1:
+                    # Create subgroups of ranks for TP, each group with `vllm_colocation` ranks.
+                    # For example, if world_size=8 and vllm_colocation=2 → groups: [0,1], [2,3], [4,5], [6,7]
+                    self.tp_group, _  = torch.distributed.new_subgroups_by_enumeration(
+                        [
+                            list(range(i*self.args.vllm_colocation, (i+1) * self.args.vllm_colocation)) 
+                            for i in range(self.accelerator.num_processes // self.args.vllm_colocation)
+                        ]
+                    )
+
+                self.llm = LLM(
+                    model=model.name_or_path,
+                    # device=f"{device_type}:{self.accelerator.process_index}", # ToDo: we do not need to set the device
+                    tensor_parallel_size=args.vllm_colocation, 
+                    gpu_memory_utilization=self.args.vllm_gpu_memory_utilization,
+                    # max_num_seqs=self.args.per_device_train_batch_size * self.args.vllm_colocation, # ToDo: this should be multiplied by gradient_accumulation_steps 
path: trl/trainer/grpo_trainer.py
user: login: toslali-ibm
body: self.args.per_device_train_batch_size * self.args.vllm_colocation * gradient_accumulation_steps
created_at: 2025-04-30T16:06:20Z
author_association: CONTRIBUTOR
diff_hunk: @@ -978,34 +1025,70 @@ def _generate_and_score_completions(
                 self._last_loaded_step = self.state.global_step
 
             # Generate completions using vLLM: gather all prompts and use them in a single call in the main process
-            all_prompts_text = gather_object(prompts_text)
-            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 = all_prompts_text[:: self.num_generations]
-                with profiling_context(self, "vLLM.generate"):
-                    completion_ids = self.vllm_client.generate(
-                        prompts=ordered_set_of_prompts,
-                        n=self.num_generations,
-                        repetition_penalty=self.repetition_penalty,
-                        temperature=self.temperature,
-                        top_p=self.top_p,
-                        top_k=-1 if self.top_k is None else self.top_k,
-                        min_p=0.0 if self.min_p is None else self.min_p,
-                        max_tokens=self.max_completion_length,
-                        guided_decoding_regex=self.guided_decoding_regex,
-                    )
-            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 self.vllm_mode == "server":
+                all_prompts_text = gather_object(prompts_text)
+                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 = all_prompts_text[:: self.num_generations]
+                    with profiling_context(self, "vLLM.generate"):
+                        completion_ids = self.vllm_client.generate(
+                            prompts=ordered_set_of_prompts,
+                            n=self.num_generations,
+                            repetition_penalty=self.repetition_penalty,
+                            temperature=self.temperature,
+                            top_p=self.top_p,
+                            top_k=-1 if self.top_k is None else self.top_k,
+                            min_p=0.0 if self.min_p is None else self.min_p,
+                            max_tokens=self.max_completion_length,
+                            guided_decoding_regex=self.guided_decoding_regex,
+                        )
+                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]
+
+            # Generate completions using colocated vLLM instances: each device holds vLLM copy and work on their own batch of prompts
+            elif self.vllm_mode == "colocation":
+                if self.guided_decoding_regex:
+                    guided_decoding = GuidedDecodingParams(backend="outlines", regex=self.guided_decoding_regex)
+                else:
+                    guided_decoding = None
+                sampling_params = SamplingParams(
+                    n=1,  # vLLM on each GPU generates only 1 in vllm_colocation mode
+                    repetition_penalty=self.repetition_penalty,
+                    temperature=self.temperature,
+                    top_p=self.top_p,
+                    top_k=-1 if self.top_k is None else self.top_k,
+                    min_p=0.0 if self.min_p is None else self.min_p,
+                    max_tokens=self.max_completion_length,
+                    guided_decoding=guided_decoding,
+                )
+
+                if self.vllm_tensor_parallel_size > 1:
+                    # Gather prompts from all ranks in the TP group and flatten.
+                    # Each rank starts with its own prompts; after gathering, all ranks see the full group set.
+                    orig_size = len(prompts_text)
+                    gathered_prompts = [None for _ in range(self.vllm_tensor_parallel_size)]
+                    torch.distributed.all_gather_object(gathered_prompts, prompts_text, group=self.tp_group)
+                    prompts_text = [p for sublist in gathered_prompts for p in sublist]
+
+                all_outputs = self.llm.generate(prompts_text, sampling_params=sampling_params, use_tqdm=False)
path: trl/trainer/grpo_trainer.py
user: login: qgallouedec
body: In a multi-GPUsetup  (2 GPUs) with TP=2, if each rank is given a separate subset of prompts—e.g., rank 0 gets `["a", "b"]` and rank 1 gets `["c", "d"]`. Does each rank independently call:

```python
llm.generate(["a", "b", "c", "d"])
```

It seems like duplicated call, but is it coordinated such that each rank only processes its subset of prompts? In other words, if the full prompt list is passed on each rank, does vLLM handle this duplication internally to avoid redundant work?

created_at: 2025-05-01T19:00:42Z
author_association: MEMBER
diff_hunk: @@ -978,34 +1025,70 @@ def _generate_and_score_completions(
                 self._last_loaded_step = self.state.global_step
 
             # Generate completions using vLLM: gather all prompts and use them in a single call in the main process
-            all_prompts_text = gather_object(prompts_text)
-            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 = all_prompts_text[:: self.num_generations]
-                with profiling_context(self, "vLLM.generate"):
-                    completion_ids = self.vllm_client.generate(
-                        prompts=ordered_set_of_prompts,
-                        n=self.num_generations,
-                        repetition_penalty=self.repetition_penalty,
-                        temperature=self.temperature,
-                        top_p=self.top_p,
-                        top_k=-1 if self.top_k is None else self.top_k,
-                        min_p=0.0 if self.min_p is None else self.min_p,
-                        max_tokens=self.max_completion_length,
-                        guided_decoding_regex=self.guided_decoding_regex,
-                    )
-            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 self.vllm_mode == "server":
+                all_prompts_text = gather_object(prompts_text)
+                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 = all_prompts_text[:: self.num_generations]
+                    with profiling_context(self, "vLLM.generate"):
+                        completion_ids = self.vllm_client.generate(
+                            prompts=ordered_set_of_prompts,
+                            n=self.num_generations,
+                            repetition_penalty=self.repetition_penalty,
+                            temperature=self.temperature,
+                            top_p=self.top_p,
+                            top_k=-1 if self.top_k is None else self.top_k,
+                            min_p=0.0 if self.min_p is None else self.min_p,
+                            max_tokens=self.max_completion_length,
+                            guided_decoding_regex=self.guided_decoding_regex,
+                        )
+                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]
+
+            # Generate completions using colocated vLLM instances: each device holds vLLM copy and work on their own batch of prompts
+            elif self.vllm_mode == "colocation":
+                if self.guided_decoding_regex:
+                    guided_decoding = GuidedDecodingParams(backend="outlines", regex=self.guided_decoding_regex)
+                else:
+                    guided_decoding = None
+                sampling_params = SamplingParams(
+                    n=1,  # vLLM on each GPU generates only 1 in vllm_colocation mode
+                    repetition_penalty=self.repetition_penalty,
+                    temperature=self.temperature,
+                    top_p=self.top_p,
+                    top_k=-1 if self.top_k is None else self.top_k,
+                    min_p=0.0 if self.min_p is None else self.min_p,
+                    max_tokens=self.max_completion_length,
+                    guided_decoding=guided_decoding,
+                )
+
+                if self.vllm_tensor_parallel_size > 1:
+                    # Gather prompts from all ranks in the TP group and flatten.
+                    # Each rank starts with its own prompts; after gathering, all ranks see the full group set.
+                    orig_size = len(prompts_text)
+                    gathered_prompts = [None for _ in range(self.vllm_tensor_parallel_size)]
+                    torch.distributed.all_gather_object(gathered_prompts, prompts_text, group=self.tp_group)
+                    prompts_text = [p for sublist in gathered_prompts for p in sublist]
+
+                all_outputs = self.llm.generate(prompts_text, sampling_params=sampling_params, use_tqdm=False)
path: trl/trainer/grpo_trainer.py
user: login: toslali-ibm
body: Yes. When using TP along with external_launcher, we need to make sure that all participating shards receive the same prompts -- and vLLM internally handles it. 

So if TP = 2 and GPU = 2, then all workers get the `["a", "b", "c", "d"]`

So if TP = 1 and GPU = 2, then first worker get the `["a", "b"]` and second worker get the `["c", "d"]`
created_at: 2025-05-01T19:03:44Z
author_association: CONTRIBUTOR
diff_hunk: @@ -617,12 +624,44 @@ def data_collator(features):  # No data collation is needed in GRPO
                     "`pip install vllm` to use it."
                 )
 
-            if self.accelerator.is_main_process:
+            if self.vllm_mode == "server" and self.accelerator.is_main_process:
                 self.vllm_client = VLLMClient(
                     args.vllm_server_host, args.vllm_server_port, connection_timeout=args.vllm_server_timeout
                 )
                 self.vllm_client.init_communicator()
 
+            elif self.vllm_mode == "colocate":
+                # Make sure vllm_tensor_parallel_size group size evenly divides the world size - each group should have
+                # the same number of ranks
+                if not self.accelerator.num_processes % self.vllm_tensor_parallel_size == 0:
+                    raise ValueError(
+                        f"vllm_tensor_parallel_size ({self.vllm_tensor_parallel_size}) must divide world size "
+                        f"({self.accelerator.num_processes}) evenly."
+                    )
+
+                if self.vllm_tensor_parallel_size > 1:
+                    # Create subgroups of ranks for TP, each group with `vllm_tensor_parallel_size` ranks.
+                    # For example, if world_size=8 and vllm_tensor_parallel_size=2 → groups: [0,1], [2,3], [4,5], [6,7]
+                    self.tp_group, _ = torch.distributed.new_subgroups_by_enumeration(
+                        [
+                            list(range(i * self.vllm_tensor_parallel_size, (i + 1) * self.vllm_tensor_parallel_size))
+                            for i in range(self.accelerator.num_processes // self.vllm_tensor_parallel_size)
+                        ]
+                    )
+
+                self.llm = LLM(
+                    model=model.name_or_path,
+                    tensor_parallel_size=args.vllm_tensor_parallel_size,
+                    gpu_memory_utilization=self.vllm_gpu_memory_utilization,
+                    max_num_seqs=self.args.per_device_train_batch_size
+                    * self.vllm_tensor_parallel_size
+                    * self.args.gradient_accumulation_steps,
+                    max_model_len=self.max_prompt_length + self.max_completion_length,
+                    distributed_executor_backend="external_launcher",
+                    # Feed identical seed for tp groups to ensure ...
path: trl/trainer/grpo_trainer.py
user: login: qgallouedec
body: @toslali-ibm I wasn't sure how to motivate this, can you complete this comment?
created_at: 2025-05-01T20:58:04Z
author_association: MEMBER
diff_hunk: @@ -617,12 +624,44 @@ def data_collator(features):  # No data collation is needed in GRPO
                     "`pip install vllm` to use it."
                 )
 
-            if self.accelerator.is_main_process:
+            if self.vllm_mode == "server" and self.accelerator.is_main_process:
                 self.vllm_client = VLLMClient(
                     args.vllm_server_host, args.vllm_server_port, connection_timeout=args.vllm_server_timeout
                 )
                 self.vllm_client.init_communicator()
 
+            elif self.vllm_mode == "colocate":
+                # Make sure vllm_tensor_parallel_size group size evenly divides the world size - each group should have
+                # the same number of ranks
+                if not self.accelerator.num_processes % self.vllm_tensor_parallel_size == 0:
+                    raise ValueError(
+                        f"vllm_tensor_parallel_size ({self.vllm_tensor_parallel_size}) must divide world size "
+                        f"({self.accelerator.num_processes}) evenly."
+                    )
+
+                if self.vllm_tensor_parallel_size > 1:
+                    # Create subgroups of ranks for TP, each group with `vllm_tensor_parallel_size` ranks.
+                    # For example, if world_size=8 and vllm_tensor_parallel_size=2 → groups: [0,1], [2,3], [4,5], [6,7]
+                    self.tp_group, _ = torch.distributed.new_subgroups_by_enumeration(
+                        [
+                            list(range(i * self.vllm_tensor_parallel_size, (i + 1) * self.vllm_tensor_parallel_size))
+                            for i in range(self.accelerator.num_processes // self.vllm_tensor_parallel_size)
+                        ]
+                    )
+
+                self.llm = LLM(
+                    model=model.name_or_path,
+                    tensor_parallel_size=args.vllm_tensor_parallel_size,
+                    gpu_memory_utilization=self.vllm_gpu_memory_utilization,
+                    max_num_seqs=self.args.per_device_train_batch_size
+                    * self.vllm_tensor_parallel_size
+                    * self.args.gradient_accumulation_steps,
+                    max_model_len=self.max_prompt_length + self.max_completion_length,
+                    distributed_executor_backend="external_launcher",
+                    # Feed identical seed for tp groups to ensure ...
path: trl/trainer/grpo_trainer.py
user: login: qgallouedec
body: btw, is `os.getenv("RANK", "0")` the same as `self.accelerator.process_index`? if so I'd use the later
created_at: 2025-05-01T20:59:36Z
author_association: MEMBER
diff_hunk: @@ -509,14 +516,15 @@ def data_collator(features):  # No data collation is needed in GRPO
         # This acts as a flag to indicate that the warning has already been issued.
         model.warnings_issued["estimate_tokens"] = True
 
-        # redirect the model.module forward to the model forward to ensure pre-forward hooks are called
-        self._forward_redirection = _ForwardRedirection()
         if self.use_liger_loss:
             if not is_liger_kernel_available():
                 raise ImportError(
                     "Liger is required to use `liger_loss` as the GRPO loss. Run `pip install liger-kernel`."
                 )
 
+            # Redirect the model.module forward to the model forward to ensure pre-forward hooks are called
+            self._forward_redirection = _ForwardRedirection()
+
path: trl/trainer/grpo_trainer.py
user: login: qgallouedec
body: not related to this PR
created_at: 2025-05-01T22:01:28Z
author_association: MEMBER
diff_hunk: @@ -825,6 +865,34 @@ def _get_per_token_logps(self, model, input_ids, attention_mask, logits_to_keep,
             all_logps.append(logps)
         return torch.cat(all_logps, dim=0)
 
+    def _sync_fsdp_params_to_vllm(self, module: nn.Module, prefix: str = "", visited=None):
+        """Memory-efficient post-order traversal of FSDP modules to extract full parameters and sync with vLLM."""
+        if visited is None:
+            visited = set()
+
+        for child_name, child_module in module.named_children():
+            child_prefix = f"{prefix}.{child_name}" if prefix else child_name
+            self._sync_fsdp_params_to_vllm(
+                child_module, prefix=child_prefix, visited=visited
+            )  # recurse into the child
+
+        if isinstance(module, FSDP):
+            with FSDP.summon_full_params(module, recurse=False, writeback=False):
+                for param_name, param in module.named_parameters():
+                    full_name = f"{prefix}.{param_name}" if prefix else param_name
+                    for extra in ("_fsdp_wrapped_module.", "_checkpoint_wrapped_module."):
+                        full_name = full_name.replace(extra, "")
+
+                    if full_name in visited:
+                        continue  # skip FSDP subtrees already traversed
+                    visited.add(full_name)
+
+                    if self.vllm_mode == "server" and self.accelerator.is_main_process:
+                        self.vllm_client.update_named_param(full_name, param.data)
+                    elif self.vllm_mode == "colocate":
+                        llm_model = self.llm.llm_engine.model_executor.driver_worker.model_runner.model
+                        llm_model.load_weights([(full_name, param.data)])
+
path: trl/trainer/grpo_trainer.py
user: login: qgallouedec
body: not related to this PR
created_at: 2025-05-01T22:01:47Z
author_association: MEMBER
diff_hunk: @@ -837,54 +905,20 @@ def _move_model_to_vllm(self):
         else:
             gather_if_zero3 = nullcontext
 
-        def post_order_fsdp_processing(module: nn.Module, prefix: str = "", visited=None):
-            """memory-efficient module gather"""
-            extra_names = ["_fsdp_wrapped_module.", "_checkpoint_wrapped_module."]
-            if visited is None:
-                visited = set()
-
-            for child_name, child_module in module.named_children():
-                if prefix == "":
-                    child_prefix = child_name
-                else:
-                    child_prefix = f"{prefix}.{child_name}"
-
-                # Recurse into the child
-                post_order_fsdp_processing(child_module, prefix=child_prefix, visited=visited)
-
-            if isinstance(module, FSDP):
-                with FSDP.summon_full_params(module, recurse=False, writeback=False):
-                    for param_name, param in module.named_parameters():
-                        if prefix == "":
-                            full_name = param_name
-                        else:
-                            full_name = f"{prefix}.{param_name}"
-
-                        for extra in extra_names:
-                            full_name = full_name.replace(extra, "")
-
-                        if full_name in visited:
-                            # skip FSDP subtrees already traversed
-                            continue
-
-                        visited.add(full_name)
-
-                        if self.accelerator.is_main_process:
-                            self.vllm_client.update_named_param(full_name, param.data)
-
         if is_peft_model(self.model):
-            # With PEFT and FSDP/DeepSpeed ZeRO Stage 3, we must gather the full model at once before merging, as merging
-            # adapters in a sharded manner is not supported.
+            # With PEFT and FSDP/DeepSpeed ZeRO Stage 3, we must gather the full model at once before merging, as
+            # merging adapters in a sharded manner is not supported.
+            # TODO: does this work with FSDP?
             with gather_if_zero3(list(self.model.parameters())):
-                if self.is_fsdp_enabled:
+                self.model.merge_adapter()
+
+                # Update vLLM weights while parameters are gathered
+                if self.is_fsdp_enabled:  # note if using FSDP, gather_if_zero3 is nullcontext
                     # Update vLLM weights while parameters are gathered
                     # For PEFT with FSDP we need to use the memory efficient post-order traversal
-                    self.model.merge_adapter()
-                    post_order_fsdp_processing(self.model)
-                    self.model.unmerge_adapter()
+                    self._sync_fsdp_params_to_vllm(self.model)
                 else:
-                    # DeepSpeed ZeRO-3 with PEFT (not using FSDP)
-                    self.model.merge_adapter()
+                    # DeepSpeed ZeRO-3 with PEFT
path: trl/trainer/grpo_trainer.py
user: login: qgallouedec
body: not related to this PR
created_at: 2025-05-01T22:02:12Z
author_association: MEMBER

PR comments:
user: login: qgallouedec
created_at: 2025-05-01T18:49:16Z
body: @toslali-ibm I've updated your PR by changing the logic a bit, so I'll let you have a look, test it out, and tell me what you think.
author_association: MEMBER
user: login: toslali-ibm
created_at: 2025-05-01T18:50:37Z
body: > @toslali-ibm I've updated your PR by changing the logic a bit, so I'll let you have a look, test it out, and tell me what you think.

I think it looks good. Let me run a sanity-check experiment. 
author_association: CONTRIBUTOR
user: login: toslali-ibm
created_at: 2025-05-01T19:07:22Z
body: I am getting an error from the current version of the code

```Traceback (most recent call last):
  File "<frozen runpy>", line 198, in _run_module_as_main
  File "<frozen runpy>", line 88, in _run_code
  File "/workspace/open-r1/src/open_r1/grpo.py", line 179, in <module>
    main(script_args, training_args, model_args)
  File "/workspace/open-r1/src/open_r1/grpo.py", line 133, in main
    train_result = trainer.train(resume_from_checkpoint=checkpoint)
                   ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/usr/local/lib/python3.12/dist-packages/transformers/trainer.py", line 2238, in train
    return inner_training_loop(
           ^^^^^^^^^^^^^^^^^^^^
  File "/usr/local/lib/python3.12/dist-packages/transformers/trainer.py", line 2553, in _inner_training_loop
    tr_loss_step = self.training_step(model, inputs, num_items_in_batch)
                   ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/usr/local/lib/python3.12/dist-packages/transformers/trainer.py", line 3724, in training_step
    inputs = self._prepare_inputs(inputs)
             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/workspace/trl/trl/extras/profiling.py", line 87, in wrapper
    return func(self, *args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/workspace/trl/trl/trainer/grpo_trainer.py", line 991, in _prepare_inputs
    accumulated_local_batch = self._generate_and_score_completions(accumulated_local_batch)
                              ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/workspace/trl/trl/trainer/grpo_trainer.py", line 1094, in _generate_and_score_completions
    completion_ids = [torch.tensor(ids, device=device) for ids in completion_ids]```
author_association: CONTRIBUTOR
user: login: qgallouedec
created_at: 2025-05-01T19:10:42Z
body: Any idea why?
author_association: MEMBER
user: login: toslali-ibm
created_at: 2025-05-01T19:33:07Z
body: > Any idea why?

There was a mismatch between config and trainer (colocate vs. colocation). I fixed that, now there is another error I am debugging: 

```
Traceback (most recent call last):
  File "<frozen runpy>", line 198, in _run_module_as_main
  File "<frozen runpy>", line 88, in _run_code
  File "/workspace/open-r1/src/open_r1/grpo.py", line 179, in <module>
    main(script_args, training_args, model_args)
  File "/workspace/open-r1/src/open_r1/grpo.py", line 133, in main
    train_result = trainer.train(resume_from_checkpoint=checkpoint)
                   ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/usr/local/lib/python3.12/dist-packages/transformers/trainer.py", line 2238, in train
    return inner_training_loop(
           ^^^^^^^^^^^^^^^^^^^^
  File "/usr/local/lib/python3.12/dist-packages/transformers/trainer.py", line 2553, in _inner_training_loop
    tr_loss_step = self.training_step(model, inputs, num_items_in_batch)
                   ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/usr/local/lib/python3.12/dist-packages/transformers/trainer.py", line 3724, in training_step
    inputs = self._prepare_inputs(inputs)
             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/workspace/trl/trl/extras/profiling.py", line 87, in wrapper
    return func(self, *args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/workspace/trl/trl/trainer/grpo_trainer.py", line 991, in _prepare_inputs
    accumulated_local_batch = self._generate_and_score_completions(accumulated_local_batch)
                              ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/workspace/trl/trl/trainer/grpo_trainer.py", line 1024, in _generate_and_score_completions
    self._move_model_to_vllm()
  File "/workspace/trl/trl/extras/profiling.py", line 87, in wrapper
    return func(self, *args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/workspace/trl/trl/trainer/grpo_trainer.py", line 961, in _move_model_to_vllm
    llm_model = self.llm.llm_engine.model_executor.driver_worker.model
                ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/usr/local/lib/python3.12/dist-packages/vllm/worker/worker_base.py", line 624, in __getattr__
    return getattr(self.worker, attr)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^
AttributeError: 'Worker' object has no attribute 'model'
```
author_association: CONTRIBUTOR
user: login: toslali-ibm
created_at: 2025-05-01T20:15:35Z
body: Okay... I think all fixed now. I was able to run a quick training. I am now running sanity-check experiment for TP =1, 2, 4 and will report rewardds.
author_association: CONTRIBUTOR
user: login: toslali-ibm
created_at: 2025-05-01T20:47:01Z
body: @qgallouedec , the sanity experiment looks good - please see the figure below.
Labels: Refactored (current version of the code) vs. vLLM coloc (before Quentin's refactor) vs. vLLM server


![W B Chart 5_1_2025, 4_46_24 PM](https://github.com/user-attachments/assets/1460b47b-6c99-4e50-a9cc-a38d637aab7b)

author_association: CONTRIBUTOR
user: login: qgallouedec
created_at: 2025-05-01T20:57:28Z
body: Nice! Trying to run on my side.
author_association: MEMBER
user: login: HuggingFaceDocBuilderDev
created_at: 2025-05-01T22:07:38Z
body: The docs for this PR live [here](https://moon-ci-docs.huggingface.co/docs/trl/pr_3394). All of your documentation changes will be reflected on that endpoint. The docs are available until 30 days after the last update.
author_association: NONE
user: login: toslali-ibm
created_at: 2025-05-01T23:14:00Z
body: > Wonderful work @toslali-ibm Can you please take a final look before I merge?

Everything looks great—my sanity experiment ran successfully on the latest version. Thanks so much for the solid help, @qgallouedec ! :)
author_association: CONTRIBUTOR
user: login: mayanks43
created_at: 2025-05-03T12:42:19Z
body: Very cool improvement!
author_association: NONE