-
-
Notifications
You must be signed in to change notification settings - Fork 3.8k
New issue
Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.
By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.
Already on GitHub? Sign in to your account
[Core] Avoid one broadcast op when propagating metadata #4844
Conversation
1e6ccaa
to
2cb1a6f
Compare
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
QQ: can you compare the overhead with other approach at tp 8?
@rkooo567 @youkaichao I haven't tried with tp=8 yet but for TP=2 with llama-2-7b this change gives ~4-5% benefit to end-to-end latency. I expect this will be additive to the benefits from #4894, and I think we can further reduce the number of broadcast ops too. I didn't compare with #4757 if that's what you mean since I expect they will be equivalent from a perf benefit pov - since the improvement comes from avoiding the additional sizing broadcast op done by @youkaichao I did also make an update on top of this to send lists rather than dicts. The measured benefit of that is negligible but I do still think it's worthwhile. To avoid complicating this PR I've pushed it to a separate one for now: njhill#2 |
# Register call site for sizing broadcast buffer. | ||
self.prepare_input_tensors_callsite_id = register_broadcast_callsite() |
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
This code is too intrusive. I thought you can extract callsite information automatically, from the function's call stack.
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
@youkaichao yes this is possible but a little hacky imo and there's a non-negligible cost to interrogating the stack (including file access I think).
IMHO this is minimally intrusive, especially since: it's not a public API, it's not used from many places, no other class needs to be updated/kept in sync per call site as is the case with the original approach.
In fact we could just remove this so that we aren't sizing per call site - empirically the latency difference due to the tensor size differences in question here isn't measurable and so could be premature optimization in any case. Especially if we combine the two call sites into one.
Alternative approach to #4757.
Does not include the changes to only broadcast values or compress the tensor metadata. For the former, it would be simplest to just change the API to broadcast lists rather than dicts imo (update: proposed changes for that added in njhill#2).