[Chore] Minor simplification for non-PP path #24810
Conversation
Signed-off-by: Woosuk Kwon <woosuk@thinkingmachines.ai>
WoosukKwon
requested review from
robertgshaw2-redhat,
njhill,
ywang96,
comaniac and
alexm-redhat
as code owners
![gemini-code-assist[bot]](https://avatars.githubusercontent.com/in/956858?s=60&v=4){kind=small}
There was a problem hiding this comment.
Code Review
This pull request provides a nice simplification to the execute_model method in GPUModelRunner. By moving the broadcast_pp_output flag to the __init__ method and restructuring the logic to handle the common (non-broadcast) and rare (broadcast) paths separately, the code becomes much clearer and easier to follow. This refactoring also correctly handles kv_connector_output for pooling models and appears to fix a latent bug where logits were not updated on non-last pipeline parallel ranks after being broadcast. The changes are well-contained and improve both readability and correctness. I have no further suggestions.
| model_output_broadcast_data = { | ||
| "logits": logits.contiguous(), | ||
| } if logits is not None else {} | ||
| else: |
There was a problem hiding this comment.
just a thought, could we put the else logic in a different function to be even less intrusive to the common path?
Signed-off-by: Woosuk Kwon <woosuk@thinkingmachines.ai>
Signed-off-by: Woosuk Kwon <woosuk@thinkingmachines.ai> Signed-off-by: bbartels <benjamin@bartels.dev>
Signed-off-by: Woosuk Kwon <woosuk@thinkingmachines.ai>
Currently,
broadcast_pp_outputcomplicates the logic in the common non-PP path.This PR simplifies the logic a bit.