Skip to content

[DLC bug] Image v1.0-hf-5.3.0-pt-2.9.0-tr-gpu-py312 fails to work with flash-attn-2 #6347

Description

@vidhig

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.

Metadata

Metadata

Assignees

No one assigned

    Labels

    No labels
    No labels

    Type

    No type

    Fields

    No fields configured for issues without a type.

    Projects

    No projects

    Milestone

    No milestone

    Relationships

    None yet

    Development

    No branches or pull requests

    Issue actions