[Core] Avoid one broadcast op when propagating metadata #4844
Conversation
njhill
force-pushed
the
single-broadcast
branch
4 times, most recently
from
1e6ccaa to
2cb1a6f
Compare
|
@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.
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.
@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).