|
|
#!/usr/bin/env bash |
|
|
set -eo pipefail |
|
|
|
|
|
CONDA_ROOT="/lustre/fsw/portfolios/nvr/projects/nvr_lpr_llm/users/yongganf/anaconda" |
|
|
ENV_NAME="nemotron_flash" |
|
|
PYTHON_VERSION="3.10.12" |
|
|
|
|
|
source "${CONDA_ROOT}/etc/profile.d/conda.sh" |
|
|
|
|
|
if conda env list | awk '{print $1}' | grep -qx "${ENV_NAME}"; then |
|
|
echo "Conda env '${ENV_NAME}' already exists, skipping creation." |
|
|
else |
|
|
echo "Creating conda env '${ENV_NAME}' with Python ${PYTHON_VERSION}..." |
|
|
conda create -n "${ENV_NAME}" "python=${PYTHON_VERSION}" -y |
|
|
fi |
|
|
conda activate "${ENV_NAME}" |
|
|
|
|
|
echo "Installing CUDA toolkit 12.8..." |
|
|
conda install -c nvidia cuda-toolkit=12.8 -y |
|
|
|
|
|
export CUDA_HOME |
|
|
CUDA_HOME="$(dirname "$(dirname "$(which nvcc)")")" |
|
|
echo "CUDA_HOME set to ${CUDA_HOME}" |
|
|
|
|
|
echo "Installing PyTorch and core Python packages..." |
|
|
pip install torch==2.8.0 transformers==4.56.2 psutil |
|
|
|
|
|
echo "Installing causal-conv1d from source..." |
|
|
pip install --no-build-isolation causal-conv1d |
|
|
|
|
|
echo "Installing flash-attention from source..." |
|
|
pip install --no-build-isolation flash-attn==2.7.3 |
|
|
|
|
|
echo "Installing mamba-ssm from source..." |
|
|
pip install --no-build-isolation mamba-ssm |
|
|
|
|
|
echo "Installing flash-linear-attention..." |
|
|
pip install --no-build-isolation flash-linear-attention |
|
|
|
|
|
echo "Done! Environment '${ENV_NAME}' is ready." |
|
|
|
|
|
|
