Move flash_attn assert from __init__ into calling func
#32
by rogerxfeng8 - opened
When enabling phi-3-small on non-cuda devices, flash_attn package is not available. The assert of flash_attn in init will force the exit. The patch changes the assert into warning, so that we can use customized implementation of flash attention in users' own modeling.
Mostly LGTM. One ask: can you move the assert warning into _apply_dense_attn, where (flash_attn_varlen_kvpacked_func is being called ?
Otherwise might be hard to understand that the issue is that flash attention is not available.
rogerxfeng8 changed pull request title from Change the assert to warning in __init__ to Move flash_attn assert from __init__ into calling func
nguyenbh changed pull request status to merged