Get set up

#!/bin/bash
# Behind the Secrets of LLMs, HPC lab
# 00_setup/setup_env.sh: sets up your Python environment once, and switches it on later.
#
# The first time it installs packages, which takes 5 to 10 minutes. Do that inside
# a small interactive job, not on the login node:
#
#     srun --account=p_lv_llmsecrets --ntasks=1 --cpus-per-task=4 --mem-per-cpu=2G \
#          --time=00:45:00 --pty bash -l
#     source 00_setup/setup_env.sh
#
# After that, in every new shell and in every batch job, it only takes a second:
#
#     source 00_setup/setup_env.sh
#
# Use "source", not "bash", so that the venv stays active in your shell.
#
# What it does:
#   1. creates (or reuses) a workspace, which is big and fast storage for venv, data and models
#   2. loads a Python module
#   3. creates (or reuses) a venv with PyTorch and transformers inside the workspace
#   4. points the pip, Hugging Face and tmp caches into the workspace instead of $HOME
#   5. downloads the tiny Shakespeare text for exercise 4

# ---- settings you may change ------------------------------------------------
WS_NAME="${WS_NAME:-llm-lab}"     # workspace name
WS_FS="${WS_FS:-horse}"           # file system, horse works on all clusters
WS_DAYS="${WS_DAYS:-30}"          # lifetime in days, extend it later with ws_extend
VENV_NAME="${VENV_NAME:-venv-llm}"

# Python module. If it does not exist on your cluster, look with: module spider Python
PY_MODULES="release/24.10 GCCcore/13.3.0 Python/3.12.3"

# PyTorch built for CUDA 12.x runs with the drivers on Alpha (A100) and Capella (H100).
# A plain "pip install torch" may pick a CUDA 13 build that needs a newer driver.
# If nvidia-smi shows CUDA 13.0 or newer you can use .../whl/cu130 instead.
TORCH_INDEX="https://download.pytorch.org/whl/cu128"
# -----------------------------------------------------------------------------

# 1) workspace. If this shell (or the shell that submitted the job) already has it, reuse it.
if [ -z "$WS_DIR" ] || [ ! -d "$WS_DIR" ]; then
    # ws_allocate prints the path. For a workspace that already exists it just reuses it.
    WS_DIR="$(ws_allocate -F "$WS_FS" "$WS_NAME" "$WS_DAYS" 2>/dev/null)"
fi
if [ -z "$WS_DIR" ] || [ ! -d "$WS_DIR" ]; then
    echo "Could not create the workspace '$WS_NAME' on '$WS_FS'. Check with: ws_list -l" >&2
    return 1 2>/dev/null || exit 1
fi
export WS_DIR
echo "workspace : $WS_DIR"

# 2) modules
module purge > /dev/null 2>&1
if ! module load $PY_MODULES; then
    echo "'module load $PY_MODULES' did not work. Run 'module spider Python' and change PY_MODULES." >&2
    return 1 2>/dev/null || exit 1
fi

# 3) caches go into the workspace, because the $HOME quota is small
export PIP_CACHE_DIR="$WS_DIR/.cache/pip"
export HF_HOME="$WS_DIR/.cache/huggingface"
export TMPDIR="$WS_DIR/tmp"
mkdir -p "$PIP_CACHE_DIR" "$HF_HOME" "$TMPDIR" "$WS_DIR/data" "$WS_DIR/runs"

# 4) venv
if [ ! -f "$WS_DIR/$VENV_NAME/bin/activate" ]; then
    echo "creating the venv $WS_DIR/$VENV_NAME (only the first time, this takes a few minutes)"
    python -m venv "$WS_DIR/$VENV_NAME"
    source "$WS_DIR/$VENV_NAME/bin/activate"
    pip install --quiet --upgrade pip
    pip install --quiet torch --index-url "$TORCH_INDEX"
    pip install --quiet numpy matplotlib tqdm transformers accelerate ipykernel
    # the same environment as a kernel in JupyterHub
    python -m ipykernel install --user --name "$VENV_NAME" --display-name "LLM Lab ($VENV_NAME)"
else
    source "$WS_DIR/$VENV_NAME/bin/activate"
fi
echo "python    : $(which python) ($(python --version 2>&1))"

# 5) the training text for exercise 4
if [ ! -s "$WS_DIR/data/input.txt" ]; then
    curl -sSL -o "$WS_DIR/data/input.txt" \
        https://raw.githubusercontent.com/karpathy/char-rnn/master/data/tinyshakespeare/input.txt \
        && echo "downloaded tiny Shakespeare to $WS_DIR/data/input.txt"
fi

echo "environment ready"

#!/bin/bash
# The conda route, an alternative to setup_env.sh for people who prefer conda.
# Builds a conda environment with PyTorch, transformers, pandas and evaluate in its
# own workspace. Run it once, from this folder, on an Alpha login node:
#
#   cd ~/behind-llms-lab/hpc_lab/00_setup/conda
#   mkdir -p logs && sbatch create_env.sbatch
#   cat logs/create-env-<jobid>.out          # the last line should say OK
#
# Running it again builds the environment from scratch.
#
# ZIH recommends the venv from setup_env.sh. Anaconda's licence now covers universities
# of TU Dresden's size and TU Dresden has no central licence, so this script installs
# only from conda-forge. If in doubt, use the venv.
#
#SBATCH --job-name=create-env
#SBATCH --account=p_lv_llmsecrets
#SBATCH --nodes=1
#SBATCH --ntasks=1
#SBATCH --cpus-per-task=6
#SBATCH --gres=gpu:1                  # Alpha takes only GPU jobs, and the last line tests CUDA
#SBATCH --mem=32G
#SBATCH --time=01:00:00
#SBATCH --output=logs/%x-%j.out
#SBATCH --error=logs/%x-%j.err

WS_NAME=llm-conda

module purge
module load Anaconda3                                 # if this fails: module spider Anaconda3
source $EBROOTANACONDA3/etc/profile.d/conda.sh        # ZIH asks for this instead of "conda init"

# The environment lives in its own workspace, not in the small $HOME
ws_allocate -F horse $WS_NAME 30 > /dev/null
ENV=$(ws_find -F horse $WS_NAME)/env

conda create -y --prefix $ENV --override-channels -c conda-forge python=3.11 pip
conda activate $ENV

# pip on ZIH systems refuses to run outside a venv, and a conda env does not count as one
export PIP_REQUIRE_VIRTUALENV=false
# PyTorch for CUDA 12.x fits the drivers on Alpha and Capella (see setup_env.sh)
pip install --no-cache-dir torch --index-url https://download.pytorch.org/whl/cu128
pip install --no-cache-dir transformers pandas evaluate

python -c "import torch, transformers, pandas, evaluate; print('OK | CUDA:', torch.cuda.is_available())"

#!/bin/bash
# The conda route: a job that runs a Python script inside the conda environment.
# Run create_env.sbatch once first. Then, from this folder:
#
#   mkdir -p logs && sbatch job.sbatch
#   cat logs/conda-job-<jobid>.out           # ends with ALL GOOD
#
# For your own code, copy this file and change the last line.
#
#SBATCH --job-name=conda-job
#SBATCH --account=p_lv_llmsecrets
#SBATCH --nodes=1
#SBATCH --ntasks=1
#SBATCH --cpus-per-task=6
#SBATCH --gres=gpu:1
#SBATCH --mem=32G
#SBATCH --time=00:30:00
#SBATCH --output=logs/%x-%j.out
#SBATCH --error=logs/%x-%j.err

WS_NAME=llm-conda

module purge
module load Anaconda3
source $EBROOTANACONDA3/etc/profile.d/conda.sh
conda activate $(ws_find -F horse $WS_NAME)/env
export HF_HOME=$(ws_find -F horse $WS_NAME)/hf-cache   # downloaded models go to the workspace, not $HOME

export SRUN_CPUS_PER_TASK=$SLURM_CPUS_PER_TASK
srun python -u check_env.py

"""The conda route: check that the environment works on a GPU node.

Run it with job.sbatch. It tests the GPU, transformers, pandas and evaluate,
and ends with ALL GOOD. The first run downloads a small model (about 250 MB).
"""
import time

import evaluate
import pandas as pd
import torch
import transformers
from transformers import pipeline

# 1. Versions
print(f"torch {torch.__version__} | transformers {transformers.__version__} | "
      f"pandas {pd.__version__} | evaluate {evaluate.__version__}")

# 2. GPU
assert torch.cuda.is_available(), "No GPU visible. Is '#SBATCH --gres=gpu:1' in the job script?"
print("GPU:", torch.cuda.get_device_name(0))
x = torch.randn(4096, 4096, device="cuda")
x @ x  # warm-up
torch.cuda.synchronize()
t = time.time()
x @ x
torch.cuda.synchronize()
print(f"4096x4096 matmul: {(time.time() - t) * 1000:.1f} ms")

# 3. transformers: a small sentiment model on the GPU
clf = pipeline(
    "sentiment-analysis",
    model="distilbert/distilbert-base-uncased-finetuned-sst-2-english",
    device=0,
)
texts = [
    "The cluster finally runs my jobs.",
    "Waiting in the queue all day is awful.",
    "This environment setup was painless.",
    "The error messages were confusing.",
]
labels = [1, 0, 1, 0]  # 1 = positive, 0 = negative
preds = clf(texts)

# 4. pandas
df = pd.DataFrame({
    "text": texts,
    "pred": [p["label"] for p in preds],
    "score": [round(p["score"], 3) for p in preds],
})
print(df.to_string(index=False))

# 5. evaluate
accuracy = evaluate.load("accuracy")
pred_ids = [1 if p["label"] == "POSITIVE" else 0 for p in preds]
print("accuracy:", accuracy.compute(predictions=pred_ids, references=labels)["accuracy"])

print("ALL GOOD")