Fix for the LM_HEAD issue #475
Merged
Conversation
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
tgaddair
requested changes
May 20, 2024
There was a problem hiding this comment.
This was a good catch! In terms of making this more general, I think we can safely move this into the BatchLoraWeights.load function which stores off unique segment_indices per layer.
One thing we would just need to change is to (1) plumb through the batch.prefill_head_indices here and (2) plumb through the layer name (which is k) here.
Also, looks like there are a couple merge conflicts that need to be cleaned up.
|
Sounds good. Will make those changes and do a basic testing round with a few different model+LoRA combinations. |
Refactored the fix to be in `BatchLoraWeights.load` method and added plumbing for it in `generate_token` in flash_causal_lm.py so that all models get this fix.
1. `AdaterBatchData.from_meta` calls in tests and `causal_lm.py` 2. `BatchLoraWeights.load` calls in tests
tgaddair
approved these changes
May 23, 2024
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.
This suggestion is invalid because no changes were made to the code.
Suggestions cannot be applied while the pull request is closed.
Suggestions cannot be applied while viewing a subset of changes.
Only one suggestion per line can be applied in a batch.
Add this suggestion to a batch that can be applied as a single commit.
Applying suggestions on deleted lines is not supported.
You must change the existing code in this line in order to create a valid suggestion.
Outdated suggestions cannot be applied.
This suggestion has been applied or marked resolved.
Suggestions cannot be applied from pending reviews.
Suggestions cannot be applied on multi-line comments.
Suggestions cannot be applied while the pull request is queued to merge.
Suggestion cannot be applied right now. Please check back later.
Fix for the LM_HEAD issue
Root Cause. The error is caused by incorrect segments passed to the
lora_b_sgmvkernel during the prefill stage. This happens because we do not want to forward all the tokens in the prompt throughlm_headand associated adapters. The goal is to save compute and memory by not having to forward the tokens that are not involved in generating the next token. Only the last token is needed for this purpose. However, doing this changes the shape of the net batch size (batch size times number of tokens) that the lora kernel sees, but the segment start and end tensors are not changed. These tensors are used to slice/segment the batch in the kernel. Hence, keeping them unchanged is incorrect. Additionally, the resulting error from this problem not reported correctly since the kernel has a catch-call condition that reports a generickernel not found for the dtypemessage. My guess is that it's an out-of-bounds memory access, but confirming this would require changing the kernel. This is out of the scope of this investigation.Description of the Fix. The fix simply goes over the adapters for
lm_headand adjusts the segment start and end tensors to correctly point to the segments in the batch. For now, this is done only during the prefill stage and when a LoRA adapter for thelm_headis present.Tests. TBD
Fixes #163
Before submitting
to it if that's the case.
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.