Add vllm colocation - #3
Conversation
| (self.accelerator.process_index + 1) * len(prompts), | ||
| ) | ||
| completion_ids = completion_ids[process_slice] | ||
| # corresponding slice (if colocated, no need for broadcasting). |
There was a problem hiding this comment.
best is to eventually move this control logic into the class
| all_prompts_text = gather_object(prompts_text) | ||
| if self.accelerator.is_main_process: | ||
| # Generate completions using vLLM: gather all prompts and use them in a single call in the main process (if colocated, work on your own batch) | ||
| all_prompts_text = prompts_text if self.args.vllm_colocation else gather_object(prompts_text) |
There was a problem hiding this comment.
control logic move into class
| ) | ||
| vllm_dtype: Optional[str] = field( | ||
| default=None, | ||
| default="auto", |
There was a problem hiding this comment.
not sure if we should change this..
There was a problem hiding this comment.
It used to be auto - now it is not used in the upstream repo, but we need it.
|
|
||
|
|
||
| class VLLMClient: | ||
| class VLLMClient(BaseVLLMClient): |
There was a problem hiding this comment.
i think to keep it simple, jsut put all clients in one file.
| # synchronize all processes after vLLM has been fully initialized. | ||
| self.accelerator.wait_for_everyone() | ||
| # synchronize all processes after vLLM has been fully initialized (if colocated, no need to wait). | ||
| if not self.args.vllm_colocation: |
There was a problem hiding this comment.
Not sure how to move this.. This can be a non-main process of a vllm_coloc False setting where vllm_client is None but should wait. We can leave this if statement here, imo...
| name = name.replace("modules_to_save.default.", "") | ||
|
|
||
| if self.accelerator.is_main_process: | ||
| if self.accelerator.is_main_process or self.args.vllm_colocation: |
| for name, param in self.model.named_parameters(): | ||
| with gather_if_zero3([param]): | ||
| if self.accelerator.is_main_process: | ||
| if self.accelerator.is_main_process or self.args.vllm_colocation: |
| # Reset cache on main process | ||
| if self.accelerator.is_main_process: | ||
| # Reset cache on main process (if colocated, reset cache on all vllms) | ||
| if self.accelerator.is_main_process or self.args.vllm_colocation: |
| guided_decoding_regex=self.guided_decoding_regex, | ||
| ) | ||
| # prompt individually (if colocated, work on your own batch). | ||
| ordered_set_of_prompts = all_prompts_text if self.args.vllm_colocation else all_prompts_text[:: self.num_generations] |
fabianlim
left a comment
fabianlim
left a comment
There was a problem hiding this comment.
Alot of "move into class" comments, but try to be creative on how to reduce code verbage, rather than just doing it by copying and pasting
|
I think moving if/else logics to the class (vllm clients) may not be really good. Because, if noncolocated and not main process, vllm client will be empty/None... So those checks have to be made in trainer as we cannot do it in vllmclient of none. |
Signed-off-by: Yu Chin Fabian Lim <flim@sg.ibm.com>
Signed-off-by: Yu Chin Fabian Lim <flim@sg.ibm.com>
2 participants
What does this PR do?
Fixes # (issue)
Before submitting
Pull Request section?
to it if that's the case.
documentation guidelines.
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.