ITADN

Runtime check for the FlashAttention version rejects local version identifiers when it shouldn't

#3334OpenlupreCSC 创建于 11 天前
bugattention
L
lupreCSCcommented
**Describe the bug** The runtime check for the FlashAttention version rejects local version identifiers when it shouldn't (i.e., an installation with version `2.8.3+local_version` is rejected by the check in https://github.com/NVIDIA/TransformerEngine/blob/8260f49660cbadb78bc52c90449428c51625469d/transformer_engine/pytorch/attention/dot_product_attention/backends.py#L102-L106 ). As local version identifiers leave the API compatible, this is too restrictive. **Expected behavior** An installation of FlashAttention with a compatible public version identifier (e.g. `2.8.3`) should be accepted whether it also provides a local version identifier or not. **Suggested fix** Make the max version check use an exclusive upper bound, i.e. - `FlashAttentionUtils.max_version` set to the next higher incompatible version (currently `FlashAttentionUtils.max_version = PkgVersion("2.8.4")`) - The version check in `transformer_engine/pytorch/attention/dot_product_attention/backends.py` then would use `<` rather than `<=` to compare to `fa_utils.max_version`: `fa_utils.version_required <= fa_utils.version < fa_utils.max_version`. This would admit local version identifiers for compatible releases (as indicated by the public version identifier) while still rejecting all incompatible versions of FlashAttention (e.g. `2.8.4`, `2.9`, etc)
1 条评论