Cannot use sequence length > 1024 with fast_multihead_attention
Nobody has claimed this yet.
- Dominant language
- Python
- Stars
- 9k
- Forks
- 1.5k
- Avg merge
- 2d 4h
- Merged PRs (30d)
- 3
Description
On running the func_test_multihead_attn.py
as follows:
python func_test_multihead_attn.py --seq-length 1088 --num-seqs-start 1 --num-seqs-stop 1
It throws an error:
python: apex/contrib/csrc/multihead_attn/self_multihead_attn_cuda.cu:153:
std::vector<at::Tensor> multihead_attn::self::cublas_gemmex::fwd_cuda(bool, bool, int, const at::Tensor&, const at:: Tensor&, const at::Tensor&, const uint8_t*, float):
Assertion `softmax_success' failed.
Aborted
This happens with any length > 1024 (works till length 1024). In the above example I am using 1 sequence per batch just to make sure its not related to OOM.
Any help in this matter is highly appreciated!
Contributor guide
No contributing guide indexed for this repository
First steps
- Read the whole issue, then the project's contributing guide.
- Comment on the issue to say you are picking it up — it saves two people doing the same work.
- Fork the repository and make your change on a branch.
- Open a pull request that references the issue number.
Research direction
Reproduce the failure with func_test_multihead_attn.py using --seq-length 1088, then inspect apex/contrib/csrc/multihead_attn/self_multihead_attn_cuda.cu around line 153 and the fwd_cuda path. Determine why softmax_success fails above length 1024; done means the reported command works for sequence lengths greater than 1024 without the assertion.
Written by the indexing model from the issue text.
Assessment
- Tech stack
- cpp, python, pytorch
- Domain
- machine-learning, performance
- Issue type
- Bug
- Difficulty
- 4/5
- Estimated time
- 3-5 days
- Activity status
- Stale
- Clarity
- Needs clarification
- Newbie friendliness
- 35/100