PyTorch_Profiling: drive site-specific settings from a single local.env

MLExamples/PyTorch_Profiling/.gitignore

@@ -0,0 +1,9 @@
+# Site configs hold per-user paths and allocations.
+local.env
+local.env.*
+!local.env.example
+
+# Job output.
+logs/
+*.out
+*.err

MLExamples/PyTorch_Profiling/README_ROCM_NIGHTLY_TESTING.md

@@ -3,63 +3,84 @@
 This guide walks you through testing a **nightly ROCm build** together with
 PyTorch on an AMD GPU. You will:
 
-1. Build a self-contained Python virtual environment with nightly ROCm, PyTorch,
-   and the ROCm profilers.
-2. Create a `setup_rocm.sh` script that activates that environment.
+1. Describe your cluster in a single `local.env` file.
+2. Build a self-contained Python virtual environment with nightly ROCm, PyTorch,
+   and the ROCm profilers, exposed as an Lmod module.
 3. Run the CIFAR-100 training workload through each profiling tool using the
    provided SLURM scripts, and check that everything works end to end.
 
 The workload itself is the same `train_cifar_100.py` used throughout this
 directory — a small vision model trained on CIFAR-100. It is intentionally
 short so that a nightly build can be validated quickly.
 
-> This guide targets an AMD MI300A GPU (`gfx942`) on a SLURM cluster. For a
-> different GPU, change the architecture (`device-gfx942` / `--arch`) and the
-> SLURM `--partition` in the scripts accordingly.
+The scripts run unmodified on any SLURM cluster and on any supported GPU:
+everything machine-specific lives in `local.env`, and no script needs editing.
 
 ---
 
-## Step 1 — Build the virtual environment
+## Step 1 — Describe your site in `local.env`
 
-Follow [`ROCM_PYTORCH_PIP_VENV_SETUP.md`](./ROCM_PYTORCH_PIP_VENV_SETUP.md) to
-create the `rocm-pytorch-pip` venv. In short, it:
+```bash
+cp local.env.example local.env
+$EDITOR local.env
+```
+
+`env.sh` sources it, or `$SITE_ENV` if set, and derives every path from it.
+
+| Key | Meaning |
+|-----|---------|
+| `PROJECT` / `PARTITION` | SLURM account and partition, exported as `SBATCH_ACCOUNT` / `SBATCH_PARTITION` |
+| `VENV_BASE` | Shared directory holding the venvs, modulefiles and examples |
+| `GPU_ARCH` / `ROOFLINE_ARCH` | Wheel extra `device-<arch>`, and the roofline extractor's counter set |
+| `ROCM_VERSION` / `ROCM_INDEX_URL` | Nightly build to validate, and the wheel index |
+| `PYTHON_MODULE` / `BASE_PYTHON` | Base interpreter used to create the venv |
+| `LMOD_INIT` | Lmod init script, for sites where `module` is undefined in non-login shells |
+| `PROXY` | Outbound proxy, for sites whose nodes have no direct internet |
+| `MIOPEN_LOCAL_DB` / `MIOPEN_TMP_BASE` | Keep MIOpen's SQLite databases on node-local storage |
+
+Two worked examples ship in `local.env.example`: OLCF Frontier (MI250X, Lustre,
+proxied compute nodes) and AAC6 (MI300A, NFS, direct internet).
 
-- creates a venv under `${VENV_BASE}/venvs/rocm-pytorch-pip` (`VENV_BASE`
-  defaults to your home directory),
-- installs nightly ROCm + PyTorch + profilers from the multi-arch nightly index,
-- installs `transformers` (required by the training script),
-- runs `rocm-sdk init` to extract the development headers and device code.
+`VENV_BASE` must be visible from **both** the login and the compute nodes, since
+the login node builds the venv and the jobs read it. Prefer the fastest shared
+filesystem available; node-local paths such as `/tmp` or `/dev/shm` do not work
+however fast they are.
 
-> **Tip:** Point `VENV_BASE` at node-local/fast storage (not a shared NFS home).
-> The venv holds hundreds of MB of ROCm/PyTorch libraries, and loading them from
-> NFS on every job noticeably slows startup.
+To validate a different nightly later, change `ROCM_VERSION` and rebuild.
 
-To test a **specific nightly**, set the ROCm version pin in that guide, e.g.:
+## Step 2 — Build the environment
+
+Run this **on a login node** (one with internet access):
 
 ```bash
-ROCM_VERSION=7.15.0a20260721
+bash install_rocm_pytorch.sh
 ```
 
-Change this value to the nightly date you want to validate.
+It is idempotent, and it:
 
-## Step 2 — Verify `setup_rocm.sh`
+- creates the venv under `${VENV_BASE}/venvs/rocm-pytorch-pip` and installs
+  nightly ROCm + PyTorch + profilers from the multi-arch nightly index, plus
+  `transformers` (required by the training script),
+- runs `rocm-sdk init` to extract the development headers and device code,
+- generates an Lmod modulefile at
+  `${VENV_BASE}/modulefiles/rocm-pytorch-pip/${ROCM_VERSION}.lua`,
+- pre-stages everything the jobs need, so compute nodes need no internet:
+  `rooflineExtractor` and its requirements, the isolated
+  `rocprof-compute analyze` venv (`numpy==1.26.4`), and the CIFAR-100 dataset.
 
-The SLURM scripts in each sub-directory activate the environment by sourcing
-`../setup_rocm.sh` — i.e. the `setup_rocm.sh` shipped in **this**
-(`PyTorch_Profiling/`) directory. It is already provided; just verify (and edit
-if needed) that its `VENV` points at the venv you built in Step 1. Its full
-contents, and a GPU-node sanity check, are covered in steps 5-7 of
-[`ROCM_PYTORCH_PIP_VENV_SETUP.md`](./ROCM_PYTORCH_PIP_VENV_SETUP.md).
+[`ROCM_PYTORCH_PIP_VENV_SETUP.md`](./ROCM_PYTORCH_PIP_VENV_SETUP.md) walks
+through the same venv build by hand, for adapting it or debugging a failure.
 
-Once verified, a quick check that the nightly build runs GPU kernels through
-PyTorch:
+Confirm the nightly runs GPU kernels through PyTorch before profiling:
 
 ```bash
 source setup_rocm.sh
 srun -n1 --gpus=1 python3 -c "import torch; print('torch', torch.__version__); \
 x = torch.ones(4, device='cuda:0'); print('device ok:', (x+1).sum().item())"
 ```
 
+`setup_rocm.sh` loads the generated module when it exists and otherwise
+activates the venv directly; the SLURM scripts source it as `../setup_rocm.sh`.
 If you see `device ok:`, you are ready to profile.
 
 ## Step 3 — Run the SLURM scripts
@@ -69,8 +90,13 @@ Each sub-directory contains a single-process SLURM script that sources
 under one tool. All of them use a single GPU and a short run
 (`--batch-size 32 --max-steps 5`) so a nightly can be checked quickly.
 
-**Submit each script from its own directory** (the scripts use
-`SLURM_SUBMIT_DIR` to locate themselves):
+`run_all.sh` submits the whole suite from a login node, each script from its own
+directory, chaining the analyze job after its profile job with
+`--dependency=afterok`. It prints one `<tool> <jobid>` line per submission:
+
+```bash
+bash run_all.sh
+```
 
 | Tool | Directory | Script | What it produces |
 |------|-----------|--------|------------------|
@@ -82,14 +108,14 @@ under one tool. All of them use a single GPU and a short run
 | ROCm Systems Profiler | `rocm-systems-profiler/` | `slurm_single_process.sh` | Sampling profile + trace under `rocprofsys-python3-output/`. |
 | Roofline Extractor | `roofline-extractor/` | `slurm_single_process.sh` | Per-kernel roofline analysis + interactive HTML plot under `output/`. |
 
-Example (baseline sanity check first, then a profiler):
+To submit one tool on its own, do it **from that tool's directory** (the scripts
+use `SLURM_SUBMIT_DIR` to locate themselves), having sourced `env.sh` so the
+account and partition reach `sbatch`:
 
 ```bash
+source env.sh
 cd no-profiling
 sbatch slurm_single_process_noprofile.sh
-
-cd ../rocprofv3
-sbatch slurm_single_process_kernels.sh
 ```
 
 Check job status and output:
@@ -102,10 +128,10 @@ squeue --me
 ### Analyzing the results
 
 > **Note — running `rocprof-compute analyze`:** run it only **after** its
-> profile job has finished (the counter database must exist). Because it needs
-> `numpy==1.26.4` (vs the shared venv's `numpy>=2.0`), the analysis script uses
-> its own isolated venv (`~/venvs/rocprof-compute-analyze`) and never touches
-> the shared venv.
+> profile job has finished (the counter database must exist); `run_all.sh`
+> chains it for you. Because it needs `numpy==1.26.4` (vs the shared venv's
+> `numpy>=2.0`), the analysis script uses its own isolated venv
+> (`${VENV_BASE}/venvs/rocprof-compute-analyze`) and never touches the shared venv.
 
 - **ROCm Compute Profiler:** submit the companion analysis job from
   `rocm-compute-profiler/` (it locates the workload and runs the analysis for you):
@@ -126,9 +152,19 @@ sbatch slurm_single_process_analyze.sh
 
 ## Notes
 
-- If a job fails to start, check the SLURM `--partition` and time limits in the
-  script headers match your cluster.
+- **Account and partition are not in the `#SBATCH` headers.** SLURM parses those
+  directives before the script runs, so they cannot expand variables. `env.sh`
+  exports `SBATCH_ACCOUNT` and `SBATCH_PARTITION` instead, which `sbatch` honours
+  at submit time. Source `env.sh` (or use `run_all.sh`) before submitting by hand.
+- **MIOpen on NFS/Lustre.** MIOpen's SQLite perf and kernel databases need real
+  POSIX file locking, and `conv2d` fails with
+  `RuntimeError: miopenStatusInternalError` without it. `setup_rocm.sh` points
+  `MIOPEN_USER_DB_PATH` and `MIOPEN_CUSTOM_CACHE_DIR` at node-local storage; set
+  `MIOPEN_LOCAL_DB=0` if your shared filesystem locks correctly.
+- **Pre-fetch the dataset with `download_only_nogpus.py`.**
+  `train_cifar_100.py --download-only` calls `dist.init_process_group("nccl", ...)`
+  even in download-only mode, which hangs on a GPU-less login node.
+- If a job fails to start, check the time limits in the script headers and the
+  `PARTITION` in `local.env` against your cluster's limits.
 - The scripts derive a per-job rendezvous port from the SLURM job ID, so
   multiple jobs can share a node without port collisions.
-- To validate a different nightly, rebuild the venv (Step 1) with a new
-  `ROCM_VERSION` and re-run the scripts.

MLExamples/PyTorch_Profiling/ROCM_PYTORCH_PIP_VENV_SETUP.md

@@ -1,54 +1,60 @@
-# Creating the `rocm-pytorch-pip` venv (nightly ROCm + PyTorch, MI300A / gfx942)
-
-This guide walks you through building a Python virtual environment with ROCm,
-PyTorch, and the ROCm profiling tools, so you can train and profile a model on
-an AMD GPU. Follow the steps in order.
+# Creating the `rocm-pytorch-pip` venv (nightly ROCm + PyTorch)
 
+This guide builds a Python virtual environment holding ROCm, PyTorch, and the
+ROCm profiling tools, so you can train and profile a model on an AMD GPU.
 Everything installs from pip into one self-contained venv, so PyTorch and the
-profilers use the same ROCm. Each step below is a single command block you can
-copy and run.
+profilers use the same ROCm.
+
+> `install_rocm_pytorch.sh` performs every step below and pre-stages the test
+> jobs. Follow the steps by hand only to adapt the build or debug a failure; see
+> [`README_ROCM_NIGHTLY_TESTING.md`](./README_ROCM_NIGHTLY_TESTING.md) for the
+> scripted path.
 
-> This guide targets an AMD MI300A GPU (`gfx942`). If you have a different GPU,
-> change `device-gfx942` to your architecture.
+The commands use the settings from `local.env`, so create one first and let
+`env.sh` load it. Nothing below is specific to a cluster or a GPU:
+
+```bash
+cp local.env.example local.env
+$EDITOR local.env          # VENV_BASE, GPU_ARCH, ROCM_VERSION at minimum
+source env.sh
+```
 
 ---
 
 ## 1. Create and activate the venv
 
-First choose where the venv should live. Set `VENV_BASE` to the directory that
-will hold the `venvs` folder (defaults to your home directory). Keep this shell
-open for the remaining steps, which reuse the variable.
-
-> **Tip:** Prefer node-local/fast storage (e.g. a local NVMe scratch path) over
-> a shared NFS home directory. The venv holds hundreds of MB of ROCm/PyTorch
-> shared libraries, and loading them from NFS on every job noticeably slows
-> startup.
+`VENV_BASE` holds the `venvs` folder. It must be visible from **both** the login
+node that builds the venv and the compute nodes that run the jobs, so use the
+fastest *shared* filesystem available (a parallel filesystem where there is one).
+Node-local paths such as `/tmp` or `/dev/shm` cannot work here, however fast they
+are: the compute node would not see what the login node wrote.
 
 ```bash
-VENV_BASE=~
 mkdir -p "${VENV_BASE}/venvs"
-python -m venv "${VENV_BASE}/venvs/rocm-pytorch-pip"
-source "${VENV_BASE}/venvs/rocm-pytorch-pip/bin/activate"
+python3 -m venv "${VENV}"
+source "${VENV}/bin/activate"
 ```
 
 ## 2. Install ROCm + PyTorch from the nightly multi-arch index
 
 ```bash
-# Pin the nightly ROCm version once and reuse it everywhere below.
-ROCM_VERSION=7.15.0a20260721
-
-pip install --index-url https://rocm.nightlies.amd.com/whl-multi-arch/ \
-    "rocm[profiler,devel,libraries,device-gfx942]==${ROCM_VERSION}" \
-    "torch[device-gfx942]" \
-    "torchvision[device-gfx942]"
+pip install --index-url "${ROCM_INDEX_URL}" \
+    "rocm[profiler,devel,libraries,device-${GPU_ARCH}]==${ROCM_VERSION}" \
+    "torch[device-${GPU_ARCH}]" \
+    "torchvision[device-${GPU_ARCH}]"
 ```
 
 The `rocm[...]` extras pull in the pieces this workflow needs:
 - `profiler`  — the ROCm profilers: `rocprof-compute`, `rocprofv3`, and
   `rocprof-sys` (bundled `_rocm_profiler`)
 - `devel`     — development package (headers/device code, extracted in step 4)
 - `libraries` — math libraries (hipBLAS, rocBLAS, ...)
-- `device-gfx942` — the GPU-arch kernels for MI300A
+- `device-${GPU_ARCH}` — the GPU-arch kernels (`gfx942` for MI300A, `gfx90a` for
+  MI250X); see the index for the architectures a given nightly ships
+
+To validate a different nightly, change `ROCM_VERSION` in `local.env` and
+rebuild. Available versions are listed at the index URL itself,
+<https://rocm.nightlies.amd.com/whl-multi-arch/>.
 
 ## 3. Install `transformers` (required by the training script)
 
@@ -62,7 +68,7 @@ required.
 ## 4. Extract development headers and device code
 
 ```bash
-"${VENV_BASE}/venvs/rocm-pytorch-pip/bin/rocm-sdk" init
+"${VENV}/bin/rocm-sdk" init
 ```
 
 `rocm-sdk init` unpacks the `devel` payload (headers, LLVM device bitcode) into
@@ -71,44 +77,24 @@ and the paths `setup_rocm.sh` points at next.
 
 ---
 
-## 5. Verify `setup_rocm.sh`
+## 5. Activate with `setup_rocm.sh`
 
-The repo already ships `setup_rocm.sh` in `MLExamples/PyTorch_Profiling/` (the
-SLURM scripts source it as `../setup_rocm.sh`). It activates the venv and points
-the ROCm environment at the extracted `_rocm_sdk_devel` tree. It defaults
-`VENV_BASE` to your home directory; if you used a different `VENV_BASE` in
-step 1, export it before sourcing (or edit the default here). Its contents are:
+The repo ships `setup_rocm.sh` in `MLExamples/PyTorch_Profiling/` (the SLURM
+scripts source it as `../setup_rocm.sh`). It reads `local.env` through `env.sh`,
+activates the venv, points the ROCm environment at the extracted
+`_rocm_sdk_devel` tree, and keeps MIOpen's databases on node-local storage.
+There is nothing in it to edit:
 
 ```bash
-#!/usr/bin/env bash
-# Source this to activate the ROCm venv and set ROCm env vars:
-#   source setup_rocm.sh
-VENV_BASE="${VENV_BASE:-$HOME}"
-VENV="$VENV_BASE/venvs/rocm-pytorch-pip"
-source "$VENV/bin/activate"
-DEVEL="$(python3 -c 'import site; print(site.getsitepackages()[0])')/_rocm_sdk_devel"
-export ROCM_PATH="$DEVEL"
-export HIP_PATH="$DEVEL"
-export HIP_DEVICE_LIB_PATH="$DEVEL/lib/llvm/amdgcn/bitcode"
-export PATH="$DEVEL/bin:$PATH"
-export LD_LIBRARY_PATH="$DEVEL/lib:$DEVEL/lib/rocm_sysdeps/lib:$LD_LIBRARY_PATH"
-echo "ROCm venv active: $VENV"
-```
-
-## 6. Re-source to pick up the ROCm env vars
-
-If the venv is already active from step 1, deactivate and source the script so
-the `ROCM_PATH` / `LD_LIBRARY_PATH` exports take effect (run from
-`MLExamples/PyTorch_Profiling/`):
-
-```bash
-deactivate
+deactivate                 # if the venv is still active from step 1
 source setup_rocm.sh
 ```
 
----
+Where `install_rocm_pytorch.sh` has generated an Lmod modulefile under
+`${VENV_BASE}/modulefiles`, `setup_rocm.sh` loads that module instead of
+activating the venv directly. Either way the resulting environment is the same.
 
-## 7. Verify (on a GPU node)
+## 6. Verify (on a GPU node)
 
 ```bash
 source setup_rocm.sh
@@ -119,7 +105,7 @@ x = torch.ones(4, device='cuda:0'); print('device ok:', (x+1).sum().item())"
 Expected output resembles:
 
 ```
-ROCm venv active: /.../rocm-pytorch-pip
+setup_rocm.sh: ROCm venv active: /.../rocm-pytorch-pip
 torch 2.12.0+rocm7.15.0a20260721
 device ok: 8.0
 ```