[Feature]: Co-locating NPU support for GRPO training with trl

[Switchsyj]

🚀 The feature, motivation and pitch

Hello, as I notice that trl (vllm-ascend) now supports to inference on a independent node during training. Does it allow to inference alongside other workloads on the same GPU as training (i.e. TP or PP)?

I think this would be useful as single-node/NPU inference could be wasteful, reducing inference throughput and training efficiency.

I think the core thing is how to update grouped parameters, adapting to NPU (but I am not very familiar with that). If there's any possibilities, I would like to help.

See the issue implemented on GPUs: huggingface/trl#3162

Alternatives

No response

Additional context

To be specifically, I guess the difference should lie in init and broadcast process:
init:

pg = StatelessProcessGroup.create(host=self.host, port=self.group_port, rank=self.rank, world_size=world_size) --> replace with corresponding  group manager?
self.pynccl_comm = PyNcclCommunicator(pg, device="cuda:0")  --> replace with NPUCommunicator

broadcast:

self.pynccl_comm.broadcast(weights, src=self.rank, stream=torch.cuda.current_stream())
self.pynccl_comm.group.barrier() --> replace with some NPU broadcast method?

[Yikun]

@ji-huazhong any comments?

[ji-huazhong]

Hi @Switchsyj. I sincerely apologize for the delayed response and the slow progress on this feature. I've been swamped with urgent work tasks in the past few weeks, which prevented me from dedicating time to this development. I plan to thoroughly research it in the coming days. My goal is to complete the development by next week. I'll provide updates here regularly and ensure timely communication moving forward. Thank you for your patience and understanding.

[Switchsyj]

Thanks @ji-huazhong , now I have made some efforts to achieve this by modifying some code in grpo_trainer.py and vllm_client.py in trl, and tested it with 7B size models (but still can not run with 32B series). It might be helpful to you : )

In grpo_trainer.py:

            if self.accelerator.is_main_process:
                # backend = "hccl" if is_npu_available() else "nccl" 
                self.vllm_client = VLLMColocationClient(
                    args, model, accelerator
                )

And in vllm_client.py:

class VLLMColocationClient:
    """
    A client class to interact with vLLM processes colocated with the training process.

    This client bypasses remote communication and directly interacts with the in-process vLLM engine.
    It supports weight updates and text generation functionalities similar to `VLLMClient`, but is optimized
    for scenarios where vLLM is running in the same process or node as training.

    Args:
        args (`GRPOConfig`): Configuration object containing vLLM parameters.
        model (`transformers.PreTrainedModel`): The model being used.
        vllm_device (`torch.device` or `str`): Device on which the model is loaded (e.g., "cuda:0").
    """

    def __init__(self, args: GRPOConfig, model, accelerator):
        self.args: GRPOConfig = args
        self.model = model
        # self.vllm_device = vllm_device
        self.vllm_device = accelerator.device
        self.world_size = accelerator.num_processes
        self.process_index = accelerator.process_index
        set_seed(42)

        if self.args.vllm_colocation_tp:
            # Ensure TP value is valid (at least 1)
            assert self.args.vllm_colocation_tp >= 1, "vllm_colocation_tp must be greater than 0"

            # Make sure TP group size evenly divides the world size
            # This ensures each group has the same number of ranks
            assert self.world_size % self.args.vllm_colocation_tp == 0, (
                f"TP size of vllm_colocation_tp ({self.args.vllm_colocation_tp}) must divide world size "
                f"({self.world_size}) evenly."
            )

        # Create subgroups of ranks for TP, each group with `vllm_colocation_tp` ranks.
        # For example, if world_size=8 and vllm_colocation_tp=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_tp, (i+1) * self.args.vllm_colocation_tp)) 
                for i in range(self.world_size // self.args.vllm_colocation_tp)
            ]
        )
        self.llm = LLM(
            model=self.model.name_or_path,
            device=self.vllm_device,
            gpu_memory_utilization=self.args.vllm_gpu_memory_utilization,
            dtype=self.args.vllm_dtype,
            enable_prefix_caching=self.args.vllm_enable_prefix_caching,
            max_model_len=self.args.vllm_max_model_len,
            tensor_parallel_size=args.vllm_colocation_tp,
            distributed_executor_backend="external_launcher",
            max_num_seqs=self.args.per_device_train_batch_size * self.args.vllm_colocation_tp,
            enable_sleep_mode=True,
            # model="/workspace/user_code/yuanjun_ws/pretrained_models/QwQ-32B",
            # tensor_parallel_size=2,
            # distributed_executor_backend="external_launcher",
            # device=self.vllm_device,
        )
        
    def update_named_param(self, name: str, weights: torch.Tensor):
        """
        Updates a specific named parameter in the model.

        Args:
            name (`str`):
                Name of the layer whose weights are being updated.
            weights (`torch.Tensor`):
                Tensor containing the updated weights.
        """
        torch.cuda.empty_cache()
        self.llm.wake_up()
        llm_model = self.llm.llm_engine.model_executor.driver_worker.model_runner.model
        llm_model.load_weights([(name,weights)])

    def generate(
        self,
        prompts: list[str],
        n: int = 1,
        repetition_penalty: float = 1.0,
        temperature: float = 1.0,
        top_p: float = 1.0,
        top_k: int = -1,
        min_p: float = 0.0,
        max_tokens: int = 16,
        guided_decoding_regex: Optional[str] = None,
    ) -> list[list[str]]:
        """
        Generates model completions for the provided prompts.

        Args:
            prompts (`list[str]`):
                List of text prompts for which the model will generate completions.
            n (`int`, *optional*, defaults to `1`):
                Number of completions to generate for each prompt.
            repetition_penalty (`float`, *optional*, defaults to `1.0`):
                Parameter for repetition penalty. 1.0 means no penalty.
            temperature (`float`, *optional*, defaults to `1.0`):
                Temperature parameter for sampling. Higher values increase diversity.
            top_p (`float`, *optional*, defaults to `1.0`):
                Top-p sampling parameter.`1.0` means no truncation.
            top_k (`int`, *optional*, defaults to `-1`):
                Top-k sampling parameter. `-1` means no truncation.
            min_p (`float`, *optional*, defaults to `0.0`):
                Minimum probability for sampling.
            max_tokens (`int`, *optional*, defaults to `16`):
                Maximum number of tokens to generate for each prompt.
            guided_decoding_regex (`str` or `None`, *optional*, defaults to `None`):
                Regular expression to guide the decoding process.

        Returns:
            `list[list[int]]`:
                List of lists of token IDs representing the model-generated completions for each prompt.
        """
        torch.cuda.empty_cache()
        self.llm.wake_up()
        # Guided decoding, if enabled
        if guided_decoding_regex is not None:
            guided_decoding = GuidedDecodingParams(backend="outlines", regex=guided_decoding_regex)
        else:
            guided_decoding = None
        
        if self.args.vllm_colocation_tp:
            # 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)
            gathered_prompts = [None for _ in range(self.args.vllm_colocation_tp)]
            torch.distributed.all_gather_object(gathered_prompts, prompts, group=self.tp_group)
            prompts = [p for sublist in gathered_prompts for p in sublist]

        sampling_params = SamplingParams(
            n=1, # vLLM on each device generates only 1 in vllm_colocation mode
            repetition_penalty=repetition_penalty,
            temperature=temperature,
            top_p=top_p,
            top_k=top_k,
            min_p=min_p,
            max_tokens=max_tokens,
            guided_decoding=guided_decoding,
        )
        
        all_outputs = self.llm.generate(
            prompts, sampling_params=sampling_params, use_tqdm=False
        )
        completion_ids = [output.token_ids for outputs in all_outputs for output in outputs.outputs]

        if self.args.vllm_colocation_tp:
            # 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)
            print(f"*** output completions: {tp_slice} ***")
            completion_ids = completion_ids[tp_slice]
        try:
            assert any(len(x) > 0 for x in completion_ids)
        except:
            print(f"*** VLLM client {self.vllm_device} generate {completion_ids} ***")

        self.llm.sleep(level=2)
        return completion_ids

    def reset_prefix_cache(self):
        """
        Resets the prefix cache for the model.
        """
        self.llm.reset_prefix_cache()