一个脚本

CUDA Version: 12.4

# torch

一个脚本直接干完

echo y | conda install cuda=11.8.0 -c nvidia
pip install torch torchvision torchaudio --index-url <https://download.pytorch.org/whl/cu118>
pip install liger-kernel
pip install flash-attn deepspeed transformers --extra-index-url <https://download.pytorch.org/whl/cu118>
nvcc --version
pip install vllm --extra-index-url <https://download.pytorch.org/whl/cu118>

cuda

echo y | conda install cuda=11.8.0 -c nvidia
echo y | conda install cuda=12.2.0 -c nvidia
echo y | conda install pytorch==2.1.2 torchvision==0.16.2 torchaudio==2.1.2 pytorch-cuda=11.8 -c pytorch -c nvidia

torch

Previous PyTorch Versions | PyTorch

pip uninstall torch torchvision torchaudio

# cuda 12.1
pip install torch torchvision torchaudio --index-url <https://download.pytorch.org/whl/cu121>

# cuda 11.8
pip install torch torchvision torchaudio --index-url <https://download.pytorch.org/whl/cu118>

flash attn

For CUDA 11, we only compile for CUDA 11.8, and for CUDA 12 we only compile for CUDA 12.2

尽量选11.8

git clone [email protected]:Dao-AILab/flash-attention.git
git checkout tags/v2.3.6
python setup.py install

# A100 可以额外装一下 layer_norm
cd csrc/layer_norm && pip install .

deepspeed

pip install deepspeed

transformers