fix: assistant msg cropping might crop suffix special tokens #329
Conversation
starpit
force-pushed
the
assistant-cropping
branch
from
8bb31c6 to
952fdd9
Compare
The complexity here: we need to crop the assistant messages so that they end on a block boundary. We cannot pad them out, because this is not what the model server will do when generating the assistant message. But it will cache the prefix of full blocks. Hence the need to crop. However, we cannot just crop at the end, as this would also crop off any "end of text" special tokens that the chat template adds to the end of the `self.assistant(m)` token sequence. Therefore, we need to crop just the message part. This logic tries to do all of that in a way that is agnostic to the cast template. However, the logic does currently assume that the chat template will never add special tokens *in the middle* of the given message `m`; it assumes special tokens are only ever added (if at all) to the beginning or end. Signed-off-by: Nick Mitchell <[email protected]>
starpit
force-pushed
the
assistant-cropping
branch
from
952fdd9 to
f20fc4c
Compare
|
Assume below 4 tokens per block. We will visualize a block as something like Scenario 1: AAAA and BBB as the two inner generate outputsIn this case, fortune is with us, and we only need to pad to get good inner-outer cache locality. We also send good quality info to the outer generate, because no cropping is needed.
Outer Input Note how, excluding a short suffix, Scenario 2: BBB and AAA as the two inner generate outputsIf we swap the order of presentation of the inner outputs to the outer input, then the trailing
Outer Input Scenario 3: Crop assistant outputIf instead we crop assistant output, we avoid the miss in the middle, but have cropped a special token. I.e. this one has good locality, but we send low quality input to the outer generate. Scenario 4: Crop assistant output with fully padded suffix special tokenIf instead we crop assistant output and add a fully padded assistant token block, then we get good locality, and send reasonable quality input to the outer generate. |
1 participant
The complexity here: we need to crop the assistant messages so that they end on a block boundary. We cannot pad them out, because this is not what the model server will do when generating the assistant message. But it will cache the prefix of full blocks. Hence the need to crop. However, we cannot just crop at the end, as this would also crop off any "end of text" special tokens that the chat template adds to the end of the
self.assistant(m)token sequence. Therefore, we need to crop just the message part. This logic tries to do all of that in a way that is agnostic to the cast template. However, the logic does currently assume that the chat template will never add special tokens in the middle of the given messagem; it assumes special tokens are only ever added (if at all) to the beginning or end.DO NOT MERGE
TODO: see the discussion below. We need to crop and also pad out the assistant suffix special token.