nvcr.io/nvidia/pytorch
https://catalog.ngc.nvidia.com/orgs/nvidia/containers/pytorch
nvidia/cuda
と違って、CUDAドライバーだけでなく、PythonやPyTorchまで入ったイメージ
JupyterLabなどデータサイエンス系のパッケージが一式入っている（なのでサイズが大きい）
コンテナに含む主なライブラリのバージョンはリリースノートにある
https://docs.nvidia.com/deeplearning/frameworks/pytorch-release-notes/index.html
PyTorchはnightlyビルドっぽい（24.08タグはtorch 2.5.0（PyPIにまだない）だった）
flash-attnを入れてくれている（バージョン固定されるが、試しやすいかも）
flash-attn
が使っていて知った
https://github.com/Dao-AILab/flash-attention/blob/3cea2fb6ee54fb7e1aad9db6ac6c9331184b8647/training/Dockerfile#L5