The image generated for transformers v5 when used with a sagemaker Huggingface estimator class is failing with the following error:
import flash_attn_2_cuda as flash_attn_gpu
ImportError: libcudart.so.12: cannot open shared object file: No such file or directory
seems like this is happening because the flash attention 2 version wheel that was baked into the image was using cuda 12.x vs the actual cuda kernel available within the image is cuda 13.0.
To ensure compatibility either cuda runtime should be updated before the flash attention version is installed. In addition, there is a new flash-attn version (2.8.3.post1) which fixes some vulnerabilities so we can upgrade to that too.
The image generated for transformers v5 when used with a sagemaker Huggingface estimator class is failing with the following error:
seems like this is happening because the flash attention 2 version wheel that was baked into the image was using cuda 12.x vs the actual cuda kernel available within the image is cuda 13.0.
To ensure compatibility either cuda runtime should be updated before the flash attention version is installed. In addition, there is a new flash-attn version (2.8.3.post1) which fixes some vulnerabilities so we can upgrade to that too.