File size: 2,634 Bytes
c7a6fe6 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 | #!/bin/bash
USE_MEGATRON=${USE_MEGATRON:-1}
USE_SGLANG=${USE_SGLANG:-1}
export MAX_JOBS=32
echo "0. Install uv (The fast package installer)"
pip install uv
echo "1. install inference frameworks and pytorch they need"
if [ $USE_SGLANG -eq 1 ]; then
# --system is needed if not running inside an active virtual environment
uv pip install --system "sglang[all]==0.5.2" --no-cache-dir && uv pip install --system torch-memory-saver --no-cache-dir
fi
uv pip install --system --no-cache-dir "vllm==0.11.0"
echo "2. install basic packages"
uv pip install --system "transformers[hf_xet]>=4.51.0" accelerate datasets peft hf-transfer \
"numpy<2.0.0" "pyarrow>=15.0.0" pandas "tensordict>=0.8.0,<=0.10.0,!=0.9.0" torchdata \
ray[default] codetiming hydra-core pylatexenc qwen-vl-utils wandb dill pybind11 liger-kernel mathruler \
pytest py-spy pre-commit ruff tensorboard
echo "pyext is lack of maintainace and cannot work with python 3.12."
echo "if you need it for prime code rewarding, please install using patched fork:"
echo "uv pip install --system git+https://github.com/ShaohonChen/PyExt.git@py311support"
uv pip install --system "nvidia-ml-py>=12.560.30" "fastapi[standard]>=0.115.0" "optree>=0.13.0" "pydantic>=2.9" "grpcio>=1.62.1"
echo "3. install FlashAttention and FlashInfer"
# Install flash-attn-2.8.1 (cxx11abi=False)
# uv can install directly from the file after wget
wget -nv https://github.com/Dao-AILab/flash-attention/releases/download/v2.8.1/flash_attn-2.8.1+cu12torch2.8cxx11abiFALSE-cp312-cp312-linux_x86_64.whl && \
uv pip install --system --no-cache-dir flash_attn-2.8.1+cu12torch2.8cxx11abiFALSE-cp312-cp312-linux_x86_64.whl
uv pip install --system --no-cache-dir flashinfer-python==0.3.1
if [ $USE_MEGATRON -eq 1 ]; then
echo "4. install TransformerEngine and Megatron"
echo "Notice that TransformerEngine installation can take very long time, please be patient"
uv pip install --system "onnxscript==0.3.1"
# Keeping no-deps here as per original script logic
NVTE_FRAMEWORK=pytorch uv pip install --system --no-deps git+https://github.com/NVIDIA/TransformerEngine.git@v2.6
uv pip install --system --no-deps git+https://github.com/NVIDIA/Megatron-LM.git@core_v0.13.1
fi
echo "5. May need to fix opencv"
uv pip install --system opencv-python
uv pip install --system opencv-fixer && \
python -c "from opencv_fixer import AutoFix; AutoFix()"
if [ $USE_MEGATRON -eq 1 ]; then
echo "6. Install cudnn python package (avoid being overridden)"
uv pip install --system nvidia-cudnn-cu12==9.10.2.21
fi
echo "Successfully installed all packages" |