This will not be necessary once we merge:

• Align use of vllm_max_model_length in RLOOTrainer #4702

I have removed it after the merge.

I am also thinking how to improve this.

I think we should pass logprobs_mode in the RLOOTrainer as well, after having read the motivation of the PR that introduced this:

• 🌡️ Have vLLM return processed (temperature scaled) log probs #4163

I am reverting this condition.

Done.

I like the idea of having a class dedicated to generation for vLLM now that the amount of code related to generation is starting to saturate the online trainers.

However I'm a bit annoyed by this back-reference. It's not an ideal design choice in my opinion. In this state, I even have trouble seeing how to properly test this.

Conceptually, the generator doesn't need to have a trainer to generation. It only needs the generation parameters, probably the accelerator, and an method to update the weights.

Can we think of a alternative design, something like

class VLLMGeneration:
    def __init__(self, model_id, accelerator, mode="server", ...):
        ...

    def sync_weights(self, model):
        ...

    def generate(self, prompts, temperature, ...):
        ...

Thanks for you insightful review! 🤗

Yes, I totally agree with you that the back-reference to the trainer is not ideal.

Indeed, this PR is a preliminary step in the complete refactoring and I was planning to change that on a later PR. This was my original plan:

1. PR 1 (current): Extract the vLLM logic into a separate class (prove the concept, remove duplication)
2. PR 2: Refine the interface to use explicit parameters instead of trainer reference
3. PR 3: Add a proper protocol/interface

However, I could fix this now in this PR if you think it is better.

I have seen that the trainer instance is used by the rollout_func. This will need further refactoring to separate (but coordinate) rollout_func and vllm functionalities. 🤔

See related PR:

• Refactor vLLM generation [2/N]: Decouple rollout_func and vLLM functionalities #4712

I have seen that the trainer instance is used by the rollout_func. This will need further refactoring to separate (but coordinate) rollout_func and vllm functionalities. 🤔

Yes, indeed. The design of rollout_func is really not ideal. We implemented it that way initially because we had to move quickly for the release of OpenEnv, but in my opinion it's really something we should completely rethink and refactor.

Indeed, this PR is a preliminary step in the complete refactoring and I was planning to change that on a later PR. This was my original plan:

I'm fine with this multi-stage PR for this. It will probably be easier to review this way, and for you to write

Additionally, I see that the trainer instance is also used by the profiling_context function. This will need further refactoring as well.

I created a dedicated PR to decouple profiling from trainer:

• Refactor vLLM generation [3/N]: Decouple profiling from trainer #4717

what do you think of removing args and have instead

• enable_sleep_mode
• gpu_memory_utilization
• group_port
• guided_decoding_regex
• max_model_length
• mode
• model_impl
• server_base_url
• server_host
• server_timeout
• tensor_parallel_size

in the argument of this function?

you may want to add max_num_seqs also

Thanks, @qgallouedec.

I had something similar in mind in my original plan:

• PR 2: Refine the interface
  • I had in mind to create something like a VLLMGenerationConfig class with all present (and open for future) VLLM parameters.

If you prefer, I can do all the refactoring in a single PR?

@qgallouedec I think I addressed your suggestion above.

not necessarily for this PR, but since we add a generation submodule, we could have vllm_client in this submodule as well

nit

instead of

llm_kwargs = {
    "model": model.name_or_path,
    ...
}
self.llm = LLM(**llm_kwargs)

we can also simply do

self.llm = LLM(
    model=model.name_or_path,
    ...,
)