| | #!/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." |
| |
|
| |
|
