[Bug Fix] OnlineDPOTrainer with vLLM Server Mode

[YangKai0616]

What does this PR do?

This PR:

1. [XPU] Enable vllm_mode="server" for OnlineDPOTrainer on XPU;
2. [General] Correctly parse the return value of VLLMClient.generate;
3. [General] Correctly clean up vLLM communicator between multiple UT cases. (Code format reference link)

Root Cause:

• Issue 2 is caused by incorrect handling of the return value from vllm_client.generate in _generate_vllm_server.
  Without this PR, the test would fail with the following error:

    def _generate_vllm(self, prompts, images=None):
        eos_token_id = self.eos_token_id
        pad_token_id = self.pad_token_id
    
        # Generate completion_ids and prompt_ids based on mode
        if self.vllm_mode == "server":
            completion_ids, prompt_ids = self._generate_vllm_server(prompts, images)
        elif self.vllm_mode == "colocate":
            completion_ids, prompt_ids = self._generate_vllm_colocate(prompts, images)
    
        # Shared padding, masking, and tensor conversion logic
        max_prompt_length = max(len(ids) for ids in prompt_ids)
        prompt_mask = [[0] * (max_prompt_length - len(ids)) + [1] * len(ids) for ids in prompt_ids]
        prompt_ids = [[pad_token_id] * (max_prompt_length - len(ids)) + ids for ids in prompt_ids]
        max_tokens = self.generation_config.max_tokens
        completion_mask = [[1] * len(ids) + [0] * (max_tokens - len(ids)) for ids in completion_ids]
        completion_ids = [
            ids + [eos_token_id] if ids[-1] != eos_token_id and len(ids) < max_tokens else ids
            for ids in completion_ids
        ]
        completion_ids = [ids + [pad_token_id] * (max_tokens - len(ids)) for ids in completion_ids]
    
        # Convert to tensors
        prompt_ids = torch.tensor(prompt_ids, device=self.accelerator.device)
        prompt_mask = torch.tensor(prompt_mask, device=self.accelerator.device)
>       completion_ids = torch.tensor(completion_ids, device=self.accelerator.device)
E       ValueError: too many dimensions 'str'

trl/trainer/online_dpo_trainer.py:714: ValueError

• Issue 3 occurred because when running multiple test cases with vLLM server mode, the vLLM communicator was not properly cleaned up between test runs. The VLLMClient.init_communicator() registers cleanup via atexit, which only runs on program exit, not between individual pytest test cases.
  Without this PR, the first case of config_name would pass, but the second case would fail during initialization with the following error:

(EngineCore_DP0 pid=237629) ERROR 11-07 17:40:03 [core.py:780] Invocation of collective_rpc method failed
(EngineCore_DP0 pid=237629) ERROR 11-07 17:40:03 [core.py:780] Traceback (most recent call last):
(EngineCore_DP0 pid=237629) ERROR 11-07 17:40:03 [core.py:780]   File "/usr/local/lib/python3.10/dist-packages/vllm/v1/engine/core.py", line 777, in _handle_client_request
(EngineCore_DP0 pid=237629) ERROR 11-07 17:40:03 [core.py:780]     result = method(*self._convert_msgspec_args(method, args))
(EngineCore_DP0 pid=237629) ERROR 11-07 17:40:03 [core.py:780]   File "/usr/local/lib/python3.10/dist-packages/vllm/v1/engine/core.py", line 416, in collective_rpc
(EngineCore_DP0 pid=237629) ERROR 11-07 17:40:03 [core.py:780]     return self.model_executor.collective_rpc(method, timeout, args,
(EngineCore_DP0 pid=237629) ERROR 11-07 17:40:03 [core.py:780]   File "/usr/local/lib/python3.10/dist-packages/vllm/executor/uniproc_executor.py", line 83, in collective_rpc
(EngineCore_DP0 pid=237629) ERROR 11-07 17:40:03 [core.py:780]     return [run_method(self.driver_worker, method, args, kwargs)]
(EngineCore_DP0 pid=237629) ERROR 11-07 17:40:03 [core.py:780]   File "/usr/local/lib/python3.10/dist-packages/vllm/utils/__init__.py", line 3122, in run_method
(EngineCore_DP0 pid=237629) ERROR 11-07 17:40:03 [core.py:780]     return func(*args, **kwargs)
(EngineCore_DP0 pid=237629) ERROR 11-07 17:40:03 [core.py:780]   File "/workspace/tests/test_yk/trl/trl/scripts/vllm_serve.py", line 110, in init_communicator
(EngineCore_DP0 pid=237629) ERROR 11-07 17:40:03 [core.py:780]     raise RuntimeError("Weight update group already initialized. Call close_communicator first.")
(EngineCore_DP0 pid=237629) ERROR 11-07 17:40:03 [core.py:780] RuntimeError: Weight update group already initialized. Call close_communicator first.

[YangKai0616]

Hi @qgallouedec , please help review. Thanks!

[YangKai0616]

@qgallouedec , if there is anything I need do, please feel free to contact me.

[YangKai0616]

@kashif , could you help review. Thanks!

[qgallouedec]

Thanks for the PR @YangKai0616, we will probably need some time to review this as the bandwidth is limited on our side, and Online DPO isn't a top priority trainer, please let this PR opened in the meantime. cc @kashif if you've some insights

[kashif]

thanks @YangKai0616 reviewing it now

[HuggingFaceDocBuilderDev]

The docs for this PR live here. All of your documentation changes will be reflected on that endpoint. The docs are available until 30 days after the last update.

[kashif]

CI is failing for another reason