Docker image appear to be broken due to flash_atten

[tcapelle]

Please check that this issue hasn't been reported before.

• I searched previous Bug Reports didn't find any similar reports.

Expected Behavior

It should run.

Current behaviour

Steps to reproduce

You just pull latest docker image and run it.

sudo docker run --gpus '"all"' --rm -it winglian/axolotl:main-py3.10-cu118-2.1.0

Tried on 2 different machines.

Config yaml

No response

Possible solution

Also the image should have pytorch 2.1 and has 2.0.1

Which Operating Systems are you using?

• Linux
• macOS
• Windows

Python Version

3.10

axolotl branch-commit

docker images

Acknowledgements

• My issue title is concise, descriptive, and in title casing.
• I have searched the existing issues to make sure this bug has not been reported yet.
• I am using the latest version of axolotl.
• I have provided enough information for the maintainers to reproduce and diagnose the issue.

[hamelsmu]

What is your base CUDA env? I could not repro this. Can you try the following:

• Upgrade your base CUDA (not through conda)
• Install a newer version of nvidia-docker-toolkit

[tcapelle]

In 2 different machines I get the same error. One with A100 the other with H100.

• Also, the installed PyTorch is 2.0.1+cu117 on the image instead of the mentioned 2.1

[winglian]

Once the docker images from #987 finish, this should be resolved.