From 56fa823f2498b4fc764ff4c5c1a59202ff0febef Mon Sep 17 00:00:00 2001 From: Ashwin Aji Date: Thu, 21 May 2026 00:54:00 +0000 Subject: [PATCH 01/31] Add HydraGNN perf-analysis subagent recipe (TraceLens + Omnistat) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Adds a multi-subagent bottleneck-analysis workflow for HydraGNN on Lux: - recipes/perf-analysis/ — README + 6 agent prompt files (launcher, tracelens_analyst, tracelens_verifier, omnistat_analyst, omnistat_verifier, synthesizer) and an Omnistat user-mode config template. - examples/sbatch_train_perf_amd.sh — 2-node SLURM wrapper that injects the PyTorch profiler into HydraGNN's JSON config (under NeuralNetwork.Profile, where train_validate_test() actually reads it), starts/stops omnistat-usermode around srun, and renders the omnistat config from a @JOB_DIR@ template at submit time (configparser does not interpolate os.environ, so %(SLURM_JOB_ID)s is unreliable). - .cursor/skills/ai4science-perf-analysis/SKILL.md — orchestration guide for the main agent. - .cursor/skills/ai4science-studio/SKILL.md — captures the cross-cutting lessons surfaced by the smoke run on JOB 6762: configparser interpolation, /tmp/omni_rmsjobinfo permission collision, VM -fs.disableMmap on login node, PromQL rmsjob_info join, HydraGNN Profile-block scoping, two-phase analyst+verifier rationale, and UTC timestamp handling for omnistat-inspect. Validated end-to-end on a 2-node (16 × MI355X) HydraGNN training run: COMPLETED in 8m28s, 353 MB rank-0 trace + 1.6 MB Omnistat DB, all 6 analyst claims verified by independent re-derivation, combined_report.md generated. Co-authored-by: Cursor --- .../skills/ai4science-perf-analysis/SKILL.md | 67 +++ .cursor/skills/ai4science-studio/SKILL.md | 81 ++++ .../examples/sbatch_train_perf_amd.sh | 398 ++++++++++++++++++ material_science/models/HydraGNN/model.yaml | 4 + .../HydraGNN/recipes/perf-analysis/README.md | 115 +++++ .../recipes/perf-analysis/agents/launcher.md | 190 +++++++++ .../perf-analysis/agents/omnistat_analyst.md | 143 +++++++ .../perf-analysis/agents/omnistat_verifier.md | 92 ++++ .../perf-analysis/agents/synthesizer.md | 131 ++++++ .../perf-analysis/agents/tracelens_analyst.md | 124 ++++++ .../agents/tracelens_verifier.md | 108 +++++ .../omnistat-lux.config.template | 81 ++++ 12 files changed, 1534 insertions(+) create mode 100644 .cursor/skills/ai4science-perf-analysis/SKILL.md create mode 100755 material_science/models/HydraGNN/examples/sbatch_train_perf_amd.sh create mode 100644 material_science/models/HydraGNN/recipes/perf-analysis/README.md create mode 100644 material_science/models/HydraGNN/recipes/perf-analysis/agents/launcher.md create mode 100644 material_science/models/HydraGNN/recipes/perf-analysis/agents/omnistat_analyst.md create mode 100644 material_science/models/HydraGNN/recipes/perf-analysis/agents/omnistat_verifier.md create mode 100644 material_science/models/HydraGNN/recipes/perf-analysis/agents/synthesizer.md create mode 100644 material_science/models/HydraGNN/recipes/perf-analysis/agents/tracelens_analyst.md create mode 100644 material_science/models/HydraGNN/recipes/perf-analysis/agents/tracelens_verifier.md create mode 100644 material_science/models/HydraGNN/recipes/perf-analysis/omnistat-lux.config.template diff --git a/.cursor/skills/ai4science-perf-analysis/SKILL.md b/.cursor/skills/ai4science-perf-analysis/SKILL.md new file mode 100644 index 0000000..aee6733 --- /dev/null +++ b/.cursor/skills/ai4science-perf-analysis/SKILL.md @@ -0,0 +1,67 @@ +--- +name: ai4science-perf-analysis +description: Runs the AMD AI agents bottleneck-analysis workflow on a HydraGNN (or other ai4science) training job using TraceLens + Omnistat user-mode + paired analyst/verifier subagents. Use when the user wants to "analyze a 2-node HydraGNN run", "find bottlenecks with TraceLens/Omnistat", or "run the AMD AI agents on a job". +--- + +# AI4Science perf-analysis (multi-subagent bottleneck workflow) + +## When this skill applies + +The user wants an automated bottleneck analysis of a multi-node training/inference run using AMD's open-source observability tooling (TraceLens, Omnistat). This is **distinct** from the `ai4science-run-models` skill — that one launches a model; this one diagnoses a model after it runs. + +Default target: HydraGNN on Lux MI355X. The same pattern can extend to ORBIT-2 / StormCast / GP-MoLFormer in iteration 2. + +## Repository entry points + +- Recipe: [material_science/models/HydraGNN/recipes/perf-analysis/](../../../material_science/models/HydraGNN/recipes/perf-analysis/) +- Sbatch wrapper: [material_science/models/HydraGNN/examples/sbatch_train_perf_amd.sh](../../../material_science/models/HydraGNN/examples/sbatch_train_perf_amd.sh) +- Agent prompt files: `material_science/models/HydraGNN/recipes/perf-analysis/agents/*.md` + +## Orchestration loop (what the main agent does) + +1. **Read the recipe README** to refresh on the run topology and artifact layout. +2. **Dispatch the launcher subagent** (Task tool, `generalPurpose` or `shell`) with the prompt at `agents/launcher.md`. Wait for the job to complete and the manifest to be written. +3. **Dispatch the two analyst subagents in parallel** (one Task call per subagent in the same message): + - `tracelens_analyst` — prompt at `agents/tracelens_analyst.md` + - `omnistat_analyst` — prompt at `agents/omnistat_analyst.md` +4. **Dispatch the two verifier subagents in parallel** once analysts return: + - `tracelens_verifier` + - `omnistat_verifier` +5. **Dispatch the synthesizer subagent**, hand it both `verified_claims.json` files. +6. **Print the combined report** to the user with a short executive summary in chat. + +## Subagent contract (every subagent must) + +- Read **only** the files listed in its `## Inputs` section. +- Produce **only** the file(s) listed in its `## Outputs` section. +- Write a single line `STATUS=ok|partial|fail; reason=` to stdout as the last line. +- Never escalate to multi-node `srun`; verifiers may use **at most one** 1-node `srun -p lux -A vultr_lux -N1 --time=00:05:00` interactive probe. +- Never edit files outside `/shared/aaji/models/HydraGNN/perf-runs//`. + +## Cluster constraints (Lux) + +- Partition: `lux`. Account: `vultr_lux`. Read from [.cluster-config.yaml](../../../.cluster-config.yaml). +- Central Prometheus at `atlvrmonad01:9090` is **unreachable from compute nodes** (TCP RST, IP allowlist). Do **not** point omnistat-query at it. Always use the user-mode VictoriaMetrics that the launcher started. +- System-mode Omnistat is running on every compute node at port 8001. Don't touch it; we run our own user-mode collector alongside. +- The login node (`rad-vultr-login`) has no GPU and no MPI. All analysis runs after the job — no live profiling on the login node. + +## When something goes wrong + +- If the launcher times out waiting for `sacct` to report a terminal state, it must `scontrol show job ` and write `state=TIMEOUT_WAITING` in the manifest; the orchestrator should surface this to the user and stop. +- If TraceLens or omnistat-inspect can't be installed (network, py-version), the launcher's install step writes a clear error to stderr and exits non-zero — the orchestrator must NOT fall back to "skip analysis", it must surface the error. +- If a verifier refutes an analyst's top claim, it stays in the report (with `verdict=refuted`) so future iterations don't re-derive the same wrong conclusion. + +## Lessons captured (smoke-test on JOB 6762) + +The first 2-node end-to-end run on Lux exposed five real bugs that are now all worked around in [`sbatch_train_perf_amd.sh`](../../../material_science/models/HydraGNN/examples/sbatch_train_perf_amd.sh) and the agent prompts: + +1. **`%(SLURM_JOB_ID)s` in omnistat config is broken** — `configparser` doesn't interpolate `os.environ`. Use `@JOB_DIR@` placeholder and `sed` at submit time. +2. **`/tmp/omni_rmsjobinfo` permission collision with system-mode Omnistat** — override `job_detection_file` to a per-job path under `/shared/aaji/...`. +3. **VictoriaMetrics on login node needs `-fs.disableMmap`** to load even a 1.6 MB DB (cgroup mmap restriction). +4. **HydraGNN `Profile` block must be under `NeuralNetwork`**, not at the top level — `train_validate_test()` is called with `config["NeuralNetwork"]` so that's the scope its Profiler reads. +5. **PromQL with `jobid="..."` requires joining via `rmsjob_info`** — `rocm_*` metrics don't carry a `jobid` label directly. Verifier subagents must use: + ```promql + metric * on (instance) group_left() (max by (instance) (rmsjob_info{jobid="..."})) + ``` + +The full list (with cross-cutting context) is at the end of [.cursor/skills/ai4science-studio/SKILL.md](../ai4science-studio/SKILL.md). When fixing a bug here, check there first and propagate the lesson. diff --git a/.cursor/skills/ai4science-studio/SKILL.md b/.cursor/skills/ai4science-studio/SKILL.md index 7355375..d8f07a6 100755 --- a/.cursor/skills/ai4science-studio/SKILL.md +++ b/.cursor/skills/ai4science-studio/SKILL.md @@ -464,3 +464,84 @@ export LD_LIBRARY_PATH="/opt/venv/lib/python3.12/site-packages/torch/lib:${LD_LI ``` Only use `docker run -e` for variables you are **introducing** (e.g. `-e ORBIT2_ROOT=/orbit2`). Never use it for `LD_LIBRARY_PATH`. + +## Performance-analysis subagent workflow (TraceLens + Omnistat user-mode) + +This pattern lives in [material_science/models/HydraGNN/recipes/perf-analysis/](../../../material_science/models/HydraGNN/recipes/perf-analysis/) and is documented in detail in [.cursor/skills/ai4science-perf-analysis/SKILL.md](../ai4science-perf-analysis/SKILL.md). Lessons that apply to **any** ai4science model that wants the same workflow: + +### 1. ConfigParser does NOT interpolate `os.environ` + +Upstream Omnistat configs (`omnistat.ornl`, `omnistat.frontier`) include lines like: +```ini +victoria_datadir = /lustre/.../%(SLURM_JOB_ID)s +``` + +This **does not work** with stock omnistat-usermode because `omnistat.utils.readConfig` uses plain `configparser.ConfigParser()`, whose `BasicInterpolation` only resolves keys defined inside the config file (DEFAULT section or local section), **not** environment variables. Calling `.get()` on a key with `%(SLURM_JOB_ID)s` raises `InterpolationMissingOptionError`. + +**Fix:** Use a `@JOB_DIR@`-style placeholder in the template and run `sed` from the sbatch wrapper at submit time: +```bash +sed -e "s|@JOB_DIR@|${HG_OUTPUT_DIR}|g" omnistat-lux.config.template > $HG_OUTPUT_DIR/omnistat.config +``` + +### 2. omnistat-usermode `/tmp/omni_rmsjobinfo` permission collision + +If the cluster runs system-mode Omnistat as user `omnidc` (or any non-root, non-self user), `/tmp/omni_rmsjobinfo` will already exist and be owned by that user. omnistat-usermode's `omnistat-rms-env` step then fails with `PermissionError` on the very first `--start`, and **exporters never launch** (silent: VictoriaMetrics is up but DB stays empty). + +**Fix:** override the path in the user config: +```ini +[omnistat.collectors.rms] +job_detection_mode = file-based +job_detection_file = /shared/aaji/.../omni_rmsjobinfo +step_detection_file = /shared/aaji/.../omni_rmsjobinfo_step +``` + +Path must be on a shared FS reachable by every compute node in the allocation (NFS/Lustre/etc.) because each node's `omnistat-rms-env` runs via `srun -N $NODES --ntasks-per-node=1` and writes to that path. + +### 3. VictoriaMetrics needs `-fs.disableMmap` on the login node + +Lux login node `vm.max_map_count` is 1048576 (high), but the cgroup memory limit is 8 GB, and even the small (1.6 MB) per-job Omnistat DB cannot mmap successfully — VM panics with: +``` +FATAL: cannot mmap "/.../index.bin": cannot mmap file with size 4096 bytes; already memory mapped files: 0: no such device +``` + +**Fix:** add `-fs.disableMmap` to the VictoriaMetrics command line whenever you start VM on the login node to read a job-scoped DB: +```bash +victoria-metrics-prod -storageDataPath=$DB -fs.disableMmap -httpListenAddr=127.0.0.1:8428 ... +``` + +### 4. PromQL: filter by jobid via `rmsjob_info` join, not direct label + +`rocm_*` metrics from the omnistat collector do **not** carry a `jobid` label themselves. The omnistat-inspect query layer joins them with `rmsjob_info` per-instance: +```promql +avg(rocm_utilization_percentage * on (instance) group_left() (max by (instance) (rmsjob_info{jobid="6762"}))) +``` + +A naive `rocm_utilization_percentage{jobid="6762"}` returns 0 series even though the data is there. Verifier subagents that re-derive metrics from raw curl must use the join pattern. + +### 5. HydraGNN `Profile` block goes under `NeuralNetwork`, not the top level + +`hydragnn.train.train_validate_test()` is invoked with `config["NeuralNetwork"]` and reads `config["Profile"]` inside that scope: +```python +profiler = Profiler("./logs/" + model_with_config_name) +if "Profile" in config: + profiler.setup(config["Profile"]) +``` + +So injecting `{"Profile": {"enable": 1, "target_epoch": 1}}` at the top of `gfm_mlip.json` does **nothing** — the profiler stays disabled. The block must be inside `NeuralNetwork`: +```python +cfg.setdefault("NeuralNetwork", {})["Profile"] = {"enable": 1, "target_epoch": epoch} +``` + +(This is specific to the HydraGNN multi-dataset HPO example; other HydraGNN entrypoints may pass the full config and behave differently. Always trace where the entrypoint reads `Profile` before injecting.) + +### 6. Two-phase analyst+verifier pattern + +Every claim from the analyst subagent (TraceLens or omnistat-inspect output) is independently re-derived by a verifier subagent from raw data: +- TraceLens claims → re-derived by parsing the raw `.pt.trace.json` event stream. +- Omnistat claims → re-derived via raw `curl` PromQL against the same VictoriaMetrics endpoint. + +This catches both tool bugs and analyst hallucinations. Refuted claims stay in `verified_claims.json` (so future iterations don't re-derive the same wrong conclusion); the synthesizer drops only `verdict=refuted` from the final report. In the smoke test, all 6 claims verified within tolerance, but the structure is what catches a future flaky case. + +### 7. Datetimes from omnistat-inspect are UTC + +`job_info.json` returns `start_time`/`end_time` as ISO-8601 strings without a `Z` or offset — they are **UTC**. Verifier subagents that build PromQL `start`/`end` from these MUST use `calendar.timegm(time.strptime(s, "%Y-%m-%dT%H:%M:%S"))`, **not** `time.mktime` (which interprets local time and produces wrong epoch on most clusters). Wrong epoch → empty PromQL result series → false "no data" finding. diff --git a/material_science/models/HydraGNN/examples/sbatch_train_perf_amd.sh b/material_science/models/HydraGNN/examples/sbatch_train_perf_amd.sh new file mode 100755 index 0000000..9cddff2 --- /dev/null +++ b/material_science/models/HydraGNN/examples/sbatch_train_perf_amd.sh @@ -0,0 +1,398 @@ +#!/usr/bin/env bash +# HydraGNN 2-node training with PyTorch profiling + Omnistat user-mode telemetry. +# +# This is a thin variant of sbatch_train_amd.sh tailored for the perf-analysis +# recipe at material_science/models/HydraGNN/recipes/perf-analysis/. +# +# Differences from sbatch_train_amd.sh: +# 1. Wraps srun with omnistat-usermode --start/--stopexporters/--stopserver. +# 2. Generates a per-job copy of gfm_mlip.json with a "Profile" block injected +# so HydraGNN's built-in torch.profiler captures rank-0 of node-0 only. +# 3. Hard-codes the lux partition / vultr_lux account so this script works +# out-of-the-box on Lux. Override via SBATCH_PARTITION/SBATCH_ACCOUNT if +# submitting to a different cluster. +# 4. Defaults are tuned for a quick perf run: --nodes=2, NUM_EPOCH=2, +# MAX_NUM_BATCH=30, time=00:30:00 — enough for the wait=5/warmup=3/active=3 +# profiler schedule plus a clean epoch 0 for warm caches. +# +# Quick start: +# export AI4S_SHARED_DIR=/shared/aaji +# sbatch material_science/models/HydraGNN/examples/sbatch_train_perf_amd.sh +# +# Required: +# AI4S_SHARED_DIR — shared base path +# /shared/aaji/tools/omnistat-pr271/ (created by the launcher subagent) +# /shared/aaji/tools/victoriametrics/ (created by the launcher subagent) +# +# Optional env-var overrides (with defaults): +# HG_NUM_EPOCH=2 +# HYDRAGNN_MAX_NUM_BATCH=30 +# HG_PRECISION=fp64 +# HG_BATCH_SIZE=200 +# PROFILE_TARGET_EPOCH=1 +# OMNISTAT_USERMODE_INTERVAL=1 # seconds + +#SBATCH --job-name=hydragnn-perf +#SBATCH --partition=lux +#SBATCH --account=vultr_lux +#SBATCH --nodes=2 +#SBATCH --ntasks-per-node=8 +#SBATCH --gpus-per-node=8 +#SBATCH --cpus-per-task=16 +#SBATCH --time=00:30:00 +#SBATCH --output=hydragnn-train-%j.out +#SBATCH --error=hydragnn-train-%j.out + +set -euo pipefail + +# --------------------------------------------------------------------------- +# SCRIPT_DIR — works both at submit time and inside the SLURM spool +# --------------------------------------------------------------------------- +if [[ -n "${SLURM_JOB_ID:-}" ]]; then + _ORIG_CMD=$(scontrol show job "$SLURM_JOB_ID" | sed -n 's/.*Command=\(\S\+\).*/\1/p') + SCRIPT_DIR=$(cd "$(dirname "$_ORIG_CMD")" && pwd) +else + SCRIPT_DIR=$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd) +fi + +# --------------------------------------------------------------------------- +# Paths and configuration +# --------------------------------------------------------------------------- +HG_BASE="${AI4S_SHARED_DIR:?AI4S_SHARED_DIR must be set}/models/HydraGNN" +HG_SIF="${HG_SIF:-${AI4S_SHARED_DIR}/images/pytorch_rocm7.2.2_ubuntu24.04_py3.12_pytorch_release_2.10.0.sif}" +HG_OVERLAY="${HG_OVERLAY:-${HG_BASE}/overlays/hydragnn-overlay.img}" +HG_DATASETS="${HG_DATASETS:-ANI1x,Alexandria}" +HG_DATA_DIR="${HG_DATA_DIR:-${HG_BASE}/weights}" +HG_BATCH_SIZE="${HG_BATCH_SIZE:-200}" +HG_NUM_EPOCH="${HG_NUM_EPOCH:-2}" +HG_PRECISION="${HG_PRECISION:-fp64}" +HG_OUTPUT_DIR="${HG_OUTPUT_DIR:-${HG_BASE}/perf-runs/${SLURM_JOB_ID:-$$}}" +HG_REPO_DIR="${HG_REPO_DIR:-${HG_BASE}/code/HydraGNN}" +HYDRAGNN_MAX_NUM_BATCH="${HYDRAGNN_MAX_NUM_BATCH:-30}" +PROFILE_TARGET_EPOCH="${PROFILE_TARGET_EPOCH:-1}" + +OMNISTAT_VENV="${OMNISTAT_VENV:-/shared/aaji/tools/omnistat-pr271}" +OMNISTAT_TEMPLATE="${OMNISTAT_TEMPLATE:-${HG_BASE}/perf-runs/omnistat-lux.config}" +OMNISTAT_USERMODE_INTERVAL="${OMNISTAT_USERMODE_INTERVAL:-1}" + +NODES="${SLURM_JOB_NUM_NODES:-2}" +GPUS_PER_NODE=8 +TOTAL_RANKS=$((NODES * GPUS_PER_NODE)) + +# --------------------------------------------------------------------------- +# Validate required files +# --------------------------------------------------------------------------- +for var in HG_SIF HG_OVERLAY OMNISTAT_TEMPLATE; do + if [[ ! -e "${!var}" ]]; then + echo "ERROR: $var not found: ${!var}" >&2 + exit 2 + fi +done + +if [[ ! -x "${OMNISTAT_VENV}/bin/omnistat-usermode" ]]; then + echo "ERROR: omnistat-usermode not found at ${OMNISTAT_VENV}/bin/omnistat-usermode" >&2 + echo " Run the launcher subagent first to install (see recipes/perf-analysis/agents/launcher.md)" >&2 + exit 2 +fi + +IFS=',' read -ra DATASET_ARRAY <<< "$HG_DATASETS" +for ds in "${DATASET_ARRAY[@]}"; do + ds_path="${HG_DATA_DIR}/${ds}-v2.bp" + if [[ ! -d "$ds_path" ]]; then + echo "ERROR: Dataset not found: $ds_path" >&2 + exit 2 + fi +done + +# --------------------------------------------------------------------------- +# Clone HydraGNN source if not present (mirrors sbatch_train_amd.sh) +# --------------------------------------------------------------------------- +HG_HYDRAGNN_SHA="${HG_HYDRAGNN_SHA:-6c45f1682783e66dc89e9e23009f61716186432b}" +if [[ ! -d "${HG_REPO_DIR}/examples/multidataset_hpo_sc26" ]]; then + echo "--- Cloning HydraGNN (pinned SHA: ${HG_HYDRAGNN_SHA}) ---" + git clone https://github.com/ORNL/HydraGNN.git "$HG_REPO_DIR" + git -C "$HG_REPO_DIR" checkout "$HG_HYDRAGNN_SHA" +fi + +# --------------------------------------------------------------------------- +# Symlink datasets into the expected location +# --------------------------------------------------------------------------- +EXAMPLE_DIR="${HG_REPO_DIR}/examples/multidataset_hpo_sc26" +DATASET_DIR="${EXAMPLE_DIR}/dataset" +mkdir -p "$DATASET_DIR" + +for ds in "${DATASET_ARRAY[@]}"; do + target="${DATASET_DIR}/${ds}-v2.bp" + source="${HG_DATA_DIR}/${ds}-v2.bp" + if [[ ! -e "$target" ]]; then + ln -sfn "$source" "$target" + fi +done + +# --------------------------------------------------------------------------- +# Render the per-job omnistat config from the template (substitute @JOB_DIR@) +# This sidesteps configparser's lack of os.environ interpolation. +# --------------------------------------------------------------------------- +mkdir -p "$HG_OUTPUT_DIR" +OMNISTAT_CONFIG="${HG_OUTPUT_DIR}/omnistat.config" +sed -e "s|@JOB_DIR@|${HG_OUTPUT_DIR}|g" "$OMNISTAT_TEMPLATE" > "$OMNISTAT_CONFIG" + +# --------------------------------------------------------------------------- +# Generate the per-job profile config (inject "Profile" block) +# Done on the submit/host side (login or compute), uses python3 if available. +# --------------------------------------------------------------------------- +HG_CONFIG_OVERRIDE="${HG_OUTPUT_DIR}/gfm_mlip_profile.json" +SRC_CONFIG="${EXAMPLE_DIR}/gfm_mlip.json" + +python3 - "$SRC_CONFIG" "$HG_CONFIG_OVERRIDE" "$PROFILE_TARGET_EPOCH" <<'PYEOF' +import json, sys +src, dst, epoch = sys.argv[1], sys.argv[2], int(sys.argv[3]) +with open(src) as f: + cfg = json.load(f) +# HydraGNN's train_validate_test() is called with config["NeuralNetwork"]; +# its Profiler reads config["Profile"] from THAT scope, so the block must +# go under NeuralNetwork, not at the top level. +cfg.setdefault("NeuralNetwork", {})["Profile"] = {"enable": 1, "target_epoch": epoch} +with open(dst, "w") as f: + json.dump(cfg, f, indent=2) +print(f"Wrote {dst} with NeuralNetwork.Profile.enable=1, target_epoch={epoch}") +PYEOF + +# --------------------------------------------------------------------------- +# Print job configuration +# --------------------------------------------------------------------------- +echo "=== HydraGNN Perf Analysis Run ===" +echo " Nodes : $NODES" +echo " GPUs/node : $GPUS_PER_NODE" +echo " Total ranks : $TOTAL_RANKS" +echo " SIF : $HG_SIF" +echo " Overlay : $HG_OVERLAY" +echo " Datasets : $HG_DATASETS" +echo " Precision : $HG_PRECISION" +echo " Batch size : $HG_BATCH_SIZE" +echo " Num epochs : $HG_NUM_EPOCH" +echo " Max batches : $HYDRAGNN_MAX_NUM_BATCH" +echo " Profile epoch: $PROFILE_TARGET_EPOCH (rank-0 only)" +echo " Output dir : $HG_OUTPUT_DIR" +echo " Repo dir : $HG_REPO_DIR" +echo " Profile cfg : $HG_CONFIG_OVERRIDE" +echo " Omnistat venv: $OMNISTAT_VENV" +echo " Omnistat tmpl: $OMNISTAT_TEMPLATE" +echo " Omnistat cfg : $OMNISTAT_CONFIG (rendered)" +echo " Node(s) : ${SLURM_NODELIST:-$(hostname)}" +echo " Date : $(date)" +echo "" + +# --------------------------------------------------------------------------- +# Start Omnistat user-mode (datadir from omnistat-lux.config: under HG_OUTPUT_DIR) +# --------------------------------------------------------------------------- +echo "--- Starting Omnistat user-mode (interval=${OMNISTAT_USERMODE_INTERVAL}s) ---" +export PATH="${OMNISTAT_VENV}/bin:${PATH}" +"${OMNISTAT_VENV}/bin/omnistat-usermode" --configfile "$OMNISTAT_CONFIG" --start --interval "$OMNISTAT_USERMODE_INTERVAL" \ + 2>&1 | tee "${HG_OUTPUT_DIR}/omnistat_start.log" || { + echo "WARN: omnistat-usermode --start returned nonzero; continuing without telemetry" >&2 +} + +# Make sure we tear down even if the training fails +cleanup_omnistat() { + echo "--- Stopping Omnistat user-mode ---" + "${OMNISTAT_VENV}/bin/omnistat-usermode" --configfile "$OMNISTAT_CONFIG" --stopexporters || true + "${OMNISTAT_VENV}/bin/omnistat-usermode" --configfile "$OMNISTAT_CONFIG" --stopserver || true +} +trap cleanup_omnistat EXIT + +# --------------------------------------------------------------------------- +# Write rank script (runs inside the container on every rank) +# Mirrors sbatch_train_amd.sh but uses HG_CONFIG_OVERRIDE and gates the +# profiler to rank 0 only via PROFILE_RANK0_ONLY. +# --------------------------------------------------------------------------- +RANK_SCRIPT="${HG_OUTPUT_DIR}/hydragnn-rank-${SLURM_JOB_ID:-$$}.sh" +cat > "$RANK_SCRIPT" << 'RANKEOF' +#!/usr/bin/env bash +set -euo pipefail +source /opt/venv/bin/activate + +export PYTHONPATH="/opt/hydragnn-pkgs:${PYTHONPATH:-}" +export LD_LIBRARY_PATH="/opt/hydragnn-pkgs/adios2:/opt/ompi/lib:${LD_LIBRARY_PATH:-}" + +SCRATCH_RANK="${SCRATCH_LOCAL:-/scratch}/${USER:?}/hydragnn-${SLURM_JOB_ID:-$$}/${SLURM_PROCID:-0}" +mkdir -p "$SCRATCH_RANK" +export TMPDIR="$SCRATCH_RANK" + +export OMP_NUM_THREADS="${SLURM_CPUS_PER_TASK:-16}" +export MIOPEN_DISABLE_CACHE=1 +export MIOPEN_USER_DB_PATH="${SCRATCH_RANK}/miopen" +mkdir -p "$MIOPEN_USER_DB_PATH" +export PYTHONNOUSERSITE=1 + +export HYDRAGNN_USE_VARIABLE_GRAPH_SIZE=1 +export HYDRAGNN_AGGR_BACKEND=mpi +export HYDRAGNN_VALTEST="${HYDRAGNN_VALTEST:-0}" +export HYDRAGNN_USE_FSDP=0 +export HYDRAGNN_TRACE_LEVEL="${HYDRAGNN_TRACE_LEVEL:-1}" + +export HSA_NO_SCRATCH_RECLAIM=1 + +cd "$HG_OUTPUT_DIR" + +export HYDRAGNN_AVG_NUM_NEIGHBORS="${HYDRAGNN_AVG_NUM_NEIGHBORS:-13.735293601560318}" + +# Profiler is enabled for the whole world (the JSON config already has +# Profile.enable=1) but we want trace files from RANK 0 only so we don't +# stomp on each other in $HG_OUTPUT_DIR/logs/. We disable it in the JSON +# at runtime for non-zero ranks via a tiny monkey-patch in HydraGNN's +# Profiler class. +exec python -u -c " +import sys, os, runpy + +# Rank 0 keeps the Profile block; others zero it out so HydraGNN's Profiler +# falls through to the null context. Block lives under NeuralNetwork because +# that is the dict train_validate_test() receives. +if os.environ.get('PROFILE_RANK0_ONLY', '0') == '1' and os.environ.get('SLURM_PROCID', '0') != '0': + import json + cfg_path = os.environ['HG_CONFIG_OVERRIDE'] + with open(cfg_path) as f: + cfg = json.load(f) + cfg.setdefault('NeuralNetwork', {}).setdefault('Profile', {})['enable'] = 0 + # Write to a per-rank scratch path to avoid races with rank 0 + rank_cfg = os.path.join(os.environ.get('TMPDIR', '/tmp'), 'gfm_mlip_profile_rank.json') + with open(rank_cfg, 'w') as f: + json.dump(cfg, f) + cfg_arg = rank_cfg +else: + cfg_arg = os.environ['HG_CONFIG_OVERRIDE'] + +avg_nn = float(os.environ.get('HYDRAGNN_AVG_NUM_NEIGHBORS', '0')) +if avg_nn > 0: + import hydragnn.utils.datasets.adiosdataset as adm + _orig_init = adm.AdiosMultiDataset.__init__ + def _patched_init(self, *a, **kw): + _orig_init(self, *a, **kw) + self.avg_num_neighbors = avg_nn + adm.AdiosMultiDataset.__init__ = _patched_init + +batch_size = int(os.environ.get('HG_BATCH_SIZE', '0') or 200) +max_num_batch = int(os.environ.get('HYDRAGNN_MAX_NUM_BATCH', '0')) +num_samples_val = max_num_batch * batch_size if max_num_batch > 0 else 0 + +sys.argv = [ + os.environ['HG_EXAMPLE_DIR'] + '/gfm_mlip_all_mpnn.py', + '--inputfile=' + cfg_arg, + '--multi', + '--everyone', + '--multi_model_list=' + os.environ['HG_DATASETS'], + '--precision=' + os.environ['HG_PRECISION'], + '--batch_size=' + str(batch_size), + '--num_epoch=' + os.environ.get('HG_NUM_EPOCH', '1'), + '--startfrom=none', + '--log=hydragnn-train-' + os.environ.get('SLURM_JOB_ID','0') + '-N' + os.environ.get('SLURM_JOB_NUM_NODES','1'), +] + +if num_samples_val > 0: + sys.argv += [ + '--num_samples=' + str(num_samples_val), + '--oversampling', + '--oversampling_num_samples=' + str(num_samples_val), + ] + +runpy.run_path(sys.argv[0], run_name='__main__') +" +RANKEOF +chmod +x "$RANK_SCRIPT" + +# --------------------------------------------------------------------------- +# Multi-node network environment (read from .cluster-config.yaml) +# Identical to sbatch_train_amd.sh — kept verbatim to avoid drift. +# --------------------------------------------------------------------------- +RCCL_MULTINODE_ENVS=() +MPI_MULTINODE_ENVS=() +if [[ "$NODES" -gt 1 ]]; then + _CLUSTER_CFG="${SCRIPT_DIR}/../../../../.cluster-config.yaml" + if [[ -f "$_CLUSTER_CFG" ]]; then + _yaml_get() { grep "^ $1:" "$_CLUSTER_CFG" 2>/dev/null | sed 's/.*: *"\?\([^"]*\)"\?.*/\1/' | grep -v '^$'; } + : "${NCCL_IB_HCA:=$(_yaml_get ib_hca)}" + : "${NCCL_SOCKET_IFNAME:=$(_yaml_get mgmt_iface)}" + : "${RCCL_ANP_PLUGIN:=$(_yaml_get rccl_anp_plugin)}" + : "${LIBIONIC_PATH:=$(_yaml_get libionic_path)}" + fi + + : "${NCCL_IB_HCA:?Set NCCL_IB_HCA or configure network.ib_hca in .cluster-config.yaml}" + : "${NCCL_SOCKET_IFNAME:?Set NCCL_SOCKET_IFNAME or configure network.mgmt_iface in .cluster-config.yaml}" + : "${RCCL_ANP_PLUGIN:?Set RCCL_ANP_PLUGIN or configure network.rccl_anp_plugin in .cluster-config.yaml}" + : "${LIBIONIC_PATH:?Set LIBIONIC_PATH or configure network.libionic_path in .cluster-config.yaml}" + + MPI_MULTINODE_ENVS=( + --env OMPI_MCA_pml=ob1 + --env OMPI_MCA_btl=tcp,self + --env OMPI_MCA_btl_tcp_if_include="$NCCL_SOCKET_IFNAME" + --env MPI4PY_RC_THREADS=false + ) + + RCCL_MULTINODE_ENVS=( + --bind "${RCCL_ANP_PLUGIN}:${RCCL_ANP_PLUGIN}:ro" + --bind "${LIBIONIC_PATH}:${LIBIONIC_PATH}:ro" + --env NCCL_NET_PLUGIN="$RCCL_ANP_PLUGIN" + --env NCCL_IB_HCA="$NCCL_IB_HCA" + --env NCCL_IB_GID_INDEX=1 + --env NCCL_GDR_FLUSH_DISABLE=1 + --env RCCL_GDR_FLUSH_GPU_MEM_NO_RELAXED_ORDERING=0 + --env NCCL_GDRCOPY_ENABLE=0 + --env NCCL_IB_QPS_PER_CONNECTION=1 + --env HSA_NO_SCRATCH_RECLAIM=1 + --env NCCL_IB_TC=96 + --env NCCL_IB_FIFO_TC=192 + --env NCCL_IGNORE_CPU_AFFINITY=1 + --env NCCL_PXN_DISABLE=0 + --env NET_OPTIONAL_RECV_COMPLETION=1 + --env NCCL_IB_USE_INLINE=1 + --env NCCL_SOCKET_IFNAME="$NCCL_SOCKET_IFNAME" + --env RCCL_LL128_FORCE_ENABLE=1 + --env NCCL_IB_PCI_RELAXED_ORDERING=1 + --env NCCL_DMABUF_ENABLE=1 + --env NCCL_DEBUG="${NCCL_DEBUG:-WARN}" + ) +fi + +# --------------------------------------------------------------------------- +# Launch distributed training via srun + apptainer +# --------------------------------------------------------------------------- +echo "--- Launching training: $TOTAL_RANKS ranks across $NODES nodes ---" +echo "" + +srun --mpi=pmix \ + apptainer exec \ + --rocm \ + --overlay "${HG_OVERLAY}:ro" \ + --bind "/opt/ompi:/opt/ompi:ro" \ + --bind "${SCRATCH_LOCAL:-/scratch}:${SCRATCH_LOCAL:-/scratch}" \ + --bind "${HG_REPO_DIR}:${HG_REPO_DIR}:ro" \ + --bind "${HG_DATA_DIR}:${HG_DATA_DIR}:ro" \ + --bind "${HG_OUTPUT_DIR}:${HG_OUTPUT_DIR}" \ + --env PMIX_MCA_gds=hash \ + --env PMIX_MCA_psec=native \ + "${MPI_MULTINODE_ENVS[@]}" \ + --env SCRATCH_LOCAL="${SCRATCH_LOCAL:-/scratch}" \ + --env HG_DATASETS="$HG_DATASETS" \ + --env HG_PRECISION="$HG_PRECISION" \ + --env HG_BATCH_SIZE="${HG_BATCH_SIZE:-200}" \ + --env HG_NUM_EPOCH="${HG_NUM_EPOCH:-1}" \ + --env HYDRAGNN_MAX_NUM_BATCH="${HYDRAGNN_MAX_NUM_BATCH:-}" \ + --env HYDRAGNN_VALTEST="${HYDRAGNN_VALTEST:-0}" \ + --env HYDRAGNN_TRACE_LEVEL="${HYDRAGNN_TRACE_LEVEL:-1}" \ + --env HG_EXAMPLE_DIR="$EXAMPLE_DIR" \ + --env HG_OUTPUT_DIR="$HG_OUTPUT_DIR" \ + --env HG_CONFIG_OVERRIDE="$HG_CONFIG_OVERRIDE" \ + --env PROFILE_RANK0_ONLY=1 \ + "${RCCL_MULTINODE_ENVS[@]}" \ + "$HG_SIF" \ + bash "$RANK_SCRIPT" + +TRAIN_RC=$? + +echo "" +echo "=== HydraGNN perf-analysis training complete (rc=$TRAIN_RC) ===" +echo " Logs: ${HG_OUTPUT_DIR}/" +echo " Trace dir: ${HG_OUTPUT_DIR}/logs/" +echo " Omnistat DB: ${HG_OUTPUT_DIR}/omnistat-db/" + +exit $TRAIN_RC diff --git a/material_science/models/HydraGNN/model.yaml b/material_science/models/HydraGNN/model.yaml index 9f10928..6f84764 100644 --- a/material_science/models/HydraGNN/model.yaml +++ b/material_science/models/HydraGNN/model.yaml @@ -25,6 +25,10 @@ recipes: recipe_path: recipes/train/ script: examples/run_train.sh sbatch_script: examples/sbatch_train_amd.sh + - task: perf-analysis + description: 2-node bottleneck analysis with TraceLens + Omnistat (user-mode) via paired analyst/verifier subagents + recipe_path: recipes/perf-analysis/ + sbatch_script: examples/sbatch_train_perf_amd.sh env_vars: # Required for HPC scripts (sbatch_*_amd.sh, build_overlay_amd.sh) diff --git a/material_science/models/HydraGNN/recipes/perf-analysis/README.md b/material_science/models/HydraGNN/recipes/perf-analysis/README.md new file mode 100644 index 0000000..6bb3ee0 --- /dev/null +++ b/material_science/models/HydraGNN/recipes/perf-analysis/README.md @@ -0,0 +1,115 @@ +# HydraGNN Performance Analysis with AMD AI Agents + +Multi-subagent workflow that launches a 2-node HydraGNN training on Lux (MI355X), captures PyTorch traces and Omnistat telemetry, then runs paired analytics + verifier subagents to identify bottlenecks and propose remedies. + +> **Audience:** internal AMD performance-engineering. The output is a diagnosis of the run, not a scientific claim about HydraGNN. + +## What this recipe does + +1. **Launcher** subagent submits a 2-node training (`sbatch_train_perf_amd.sh`) on the `lux` partition with PyTorch profiling armed for one target epoch and Omnistat user-mode collecting on every node. +2. **Analytics** subagents (TraceLens + Omnistat) run after the job, in parallel, and each emits a `claims.json`. +3. **Verifier** subagents independently re-derive the top claims from raw data and optionally probe cheap remedies on a 1-node interactive `srun`. +4. **Synthesizer** merges, ranks, and writes `combined_report.md`. + +```mermaid +flowchart TD + User([User]) --> Main[Main agent: orchestrator] + Main -->|Task| Launcher[launcher: sbatch + manifest] + Launcher -->|jobid + paths| Main + Main -->|Task parallel| TLA[tracelens_analyst] + Main -->|Task parallel| OSA[omnistat_analyst] + TLA -->|claims.json| TLV[tracelens_verifier] + OSA -->|claims.json| OSV[omnistat_verifier] + TLV --> Synth[synthesizer] + OSV --> Synth + Synth -->|combined_report.md| User +``` + +## Quick start + +```bash +# 1. From the ai4science-studio repo root, ensure the cluster config is present +ls .cluster-config.yaml + +# 2. Set the shared dir and submit the perf-analysis job +export AI4S_SHARED_DIR=/shared/aaji +sbatch material_science/models/HydraGNN/examples/sbatch_train_perf_amd.sh + +# 3. Once the job completes, drive the analyst + verifier + synthesizer subagents +# via the Cursor agent (or Claude Code). The agent reads +# .cursor/skills/ai4science-perf-analysis/SKILL.md to dispatch them. +``` + +The orchestrating agent is expected to run on the **login node** (`rad-vultr-login`) with shell + filesystem access to `/shared/aaji`. Verifier subagents that need a compute node use short interactive `srun -p lux -A vultr_lux -N1 --time=00:05:00` allocations; they never grab 2-node allocations. + +## Artifacts + +After a complete run: + +``` +/shared/aaji/models/HydraGNN/perf-runs// +├── manifest.json # written by launcher +├── omnistat-db/ # VictoriaMetrics datadir +├── logs//*.pt.trace.json # PyTorch trace (rank 0 only) +├── hydragnn-train-.out # SLURM stdout +├── tracelens/ +│ ├── report.xlsx # TraceLens output +│ ├── claims.json # tracelens_analyst output +│ └── verified_claims.json # tracelens_verifier output +├── omnistat/ +│ ├── inspect_outputs/*.json # omnistat-inspect raw outputs +│ ├── claims.json # omnistat_analyst output +│ └── verified_claims.json # omnistat_verifier output +└── combined_report.md # synthesizer output +``` + +The `combined_report.md` is the deliverable: ranked bottlenecks, remedies tried with deltas, remedies proposed but not tried, and clearly-flagged "system limit reached" cases. + +## Prerequisites + +- Lux access (`vultr_lux` account) with the standard HydraGNN setup at `/shared/aaji/models/HydraGNN/` — overlay, weights, code clone. See [HydraGNN/recipes/train/](../train/). +- One-time install of Omnistat (PR #271 branch `jorda/skills`, with `origin/main` merged in), VictoriaMetrics, and TraceLens — performed lazily by the launcher subagent into `/shared/aaji/tools/`. +- The launcher writes `gfm_mlip_with_profile.json` at submit time (does not modify the upstream `gfm_mlip.json`). + +## Bottleneck taxonomy (used by analysts and verifier) + +Each claim must map to one of: + +| class | examples | +|---|---| +| `gpu_compute` | matmul/MFMA bound; low MFU; low TFLOPS vs roofline | +| `gpu_memory_hbm` | HBM bandwidth-bound; high `FETCH_SIZE`+`WRITE_SIZE` | +| `cpu_dispatch` | Python overhead between kernels; aten dispatch dominating; "GPU starvation" | +| `comm_xgmi` | intra-node allreduce / scatter-gather over XGMI | +| `comm_scaleout` | inter-node RCCL over ionic; ANP plugin throughput | +| `dataloader` | torch DataLoader; ADIOS2 read time; preprocessing | +| `host_io` | `/shared` NFS reads, MIOpen cache misses | +| `host_cpu` | OMP_NUM_THREADS, CPU saturation, NUMA effects | + +## Agent prompt files + +Per-subagent prompts (each ≤200 lines, self-contained): + +- [agents/launcher.md](agents/launcher.md) +- [agents/tracelens_analyst.md](agents/tracelens_analyst.md) +- [agents/tracelens_verifier.md](agents/tracelens_verifier.md) +- [agents/omnistat_analyst.md](agents/omnistat_analyst.md) +- [agents/omnistat_verifier.md](agents/omnistat_verifier.md) +- [agents/synthesizer.md](agents/synthesizer.md) + +The orchestrating agent dispatches each as a Task subagent (or shell script in iteration 2 with LangGraph). Each subagent's output JSON is the typed state for the next. + +## Out of scope (iteration 1) + +- LangGraph or any async backplane — subagents communicate via JSON files only. +- Multi-rank trace fusion (`TraceLens_generate_multi_rank_collective_report_pytorch`) — single-rank trace for now. +- rocprofiler hardware counters in user-mode Omnistat — disabled by default; verifier may turn on for one targeted probe. +- A/B comparative runs — recommended in `combined_report.md` but not auto-launched. +- Editing upstream HydraGNN — everything driven via JSON config + env vars. + +## Safety / etiquette + +- Verifier remedy probes use **1-node** interactive allocations only, ≤5 min each. +- Never re-runs the full 2-node job during analysis — that's a deliberate choice for the agent. +- All artifacts under `/shared/aaji/...`; large traces and DBs are not committed (see repo `.gitignore`). +- Research/perf-engineering only — no scientific claims about HydraGNN should be derived from these short runs. diff --git a/material_science/models/HydraGNN/recipes/perf-analysis/agents/launcher.md b/material_science/models/HydraGNN/recipes/perf-analysis/agents/launcher.md new file mode 100644 index 0000000..621c9c1 --- /dev/null +++ b/material_science/models/HydraGNN/recipes/perf-analysis/agents/launcher.md @@ -0,0 +1,190 @@ +# launcher subagent + +Submits a 2-node HydraGNN training on Lux with PyTorch profiling and Omnistat user-mode telemetry, waits for completion, and writes `manifest.json` for downstream subagents. + +## Inputs + +- Repo at `/home/aaji/git/ai4science-studio` on branch `aaji/perf-agents-hydragnn`. +- Cluster config at `.cluster-config.yaml` (partition `lux`, account `vultr_lux`). +- Optional env-var overrides: `HG_BATCH_SIZE`, `HYDRAGNN_MAX_NUM_BATCH`, `HG_NUM_EPOCH`, `HG_PRECISION`, `PROFILE_TARGET_EPOCH`. + +## Outputs + +- `/shared/aaji/tools/omnistat-pr271/` — Python venv with omnistat (PR #271 + main merged) and TraceLens. +- `/shared/aaji/tools/victoriametrics/victoria-metrics-prod` — VictoriaMetrics binary. +- `/shared/aaji/models/HydraGNN/perf-runs/omnistat-lux.config` — Omnistat user-mode config. +- `/shared/aaji/models/HydraGNN/perf-runs//manifest.json` — manifest schema below. +- `/shared/aaji/models/HydraGNN/perf-runs//hydragnn-train-.out` — symlinked from output dir. +- `/shared/aaji/models/HydraGNN/perf-runs//logs/` — symlinked rank-0 trace. +- `/shared/aaji/models/HydraGNN/perf-runs//omnistat-db/` — VictoriaMetrics datadir. + +## Manifest schema + +```json +{ + "jobid": "", + "submitted_at": "ISO 8601 UTC", + "completed_at": "ISO 8601 UTC", + "exit_state": "COMPLETED|FAILED|TIMEOUT|...", + "nodes": 2, + "gpus_per_node": 8, + "ranks": 16, + "nodelist": "lux-mi355x-aN,lux-mi355x-bM", + "partition": "lux", + "account": "vultr_lux", + "runtime_seconds": , + "config_used": "", + "profile_target_epoch": , + "trace_paths": [""], + "omnistat_db_path": "", + "training_log_path": "", + "perf_run_dir": "", + "hg_precision": "fp64", + "hg_batch_size": 200, + "hg_num_epoch": 2, + "hydragnn_max_num_batch": 30, + "tools_versions": { + "omnistat_commit": "", + "tracelens_commit": "", + "victoriametrics_version": "" + } +} +``` + +## Steps + +### 1. Lazy install tools (idempotent) + +```bash +set -euo pipefail +TOOLS=/shared/aaji/tools +mkdir -p "$TOOLS" + +# 1a. Omnistat: jorda/skills branch with origin/main merged in +SRC=$TOOLS/omnistat-src +if [[ ! -d $SRC/.git ]]; then + git clone https://github.com/ROCm/omnistat.git "$SRC" +fi +cd "$SRC" +git fetch --all --prune +git checkout jorda/skills 2>/dev/null || git checkout -b jorda/skills origin/jorda/skills +git reset --hard origin/jorda/skills +if ! git merge --no-edit origin/main; then + echo "ERROR: merging origin/main into jorda/skills produced conflicts; resolve manually" >&2 + exit 2 +fi +OMNISTAT_COMMIT=$(git rev-parse HEAD) + +# 1b. Venv (py3.12 to match cluster python) +VENV=$TOOLS/omnistat-pr271 +if [[ ! -d $VENV ]]; then + python3 -m venv "$VENV" +fi +"$VENV/bin/pip" install --quiet --upgrade pip +"$VENV/bin/pip" install --quiet -e "$SRC[query]" + +# 1c. TraceLens (same venv) +"$VENV/bin/pip" install --quiet "git+https://github.com/AMD-AGI/TraceLens.git" +TRACELENS_COMMIT=$("$VENV/bin/python" -c "import TraceLens, importlib.metadata; print(importlib.metadata.version('TraceLens'))") + +# 1d. VictoriaMetrics binary +VM_DIR=$TOOLS/victoriametrics +if [[ ! -x $VM_DIR/victoria-metrics-prod ]]; then + mkdir -p "$VM_DIR" + cd "$VM_DIR" + VM_VERSION=$(curl -fsSL https://api.github.com/repos/VictoriaMetrics/VictoriaMetrics/releases/latest | grep '"tag_name":' | sed -E 's/.*"([^"]+)".*/\1/') + wget -q "https://github.com/VictoriaMetrics/VictoriaMetrics/releases/download/${VM_VERSION}/victoria-metrics-linux-amd64-${VM_VERSION}.tar.gz" + tar xzf "victoria-metrics-linux-amd64-${VM_VERSION}.tar.gz" + echo "$VM_VERSION" > VERSION +fi +VM_VERSION=$(cat $VM_DIR/VERSION) +``` + +### 2. Author the Omnistat user-mode config (once, in repo `perf-runs/`) + +If `/shared/aaji/models/HydraGNN/perf-runs/omnistat-lux.config` doesn't exist, write it from the template at `/home/aaji/git/ai4science-studio/material_science/models/HydraGNN/recipes/perf-analysis/omnistat-lux.config.template`. The template uses `%(SLURM_JOB_ID)s` placeholders that omnistat-usermode resolves at runtime. + +### 3. Generate the per-job profile config + +Read the upstream `gfm_mlip.json` from `/shared/aaji/models/HydraGNN/code/HydraGNN/examples/multidataset_hpo_sc26/gfm_mlip.json` and inject the `Profile` block **inside `NeuralNetwork`** (not at the top level — `train_validate_test()` is invoked with `config["NeuralNetwork"]` as `config`, so its Profiler reads `config["Profile"]` from that scope): + +```python +cfg.setdefault("NeuralNetwork", {})["Profile"] = { + "enable": 1, + "target_epoch": int(os.environ.get("PROFILE_TARGET_EPOCH", "1")), +} +``` + +Write the result to `/shared/aaji/models/HydraGNN/perf-runs//gfm_mlip_profile.json`. Use Python (`json.load`/`json.dump`) — do NOT do this with sed. + +(The wrapper `examples/sbatch_train_perf_amd.sh` already does this; this step exists in the launcher prompt for the case where the launcher is reused for a non-HydraGNN model.) + +### 4. Submit the job + +```bash +export AI4S_SHARED_DIR=/shared/aaji +export HG_OUTPUT_DIR=/shared/aaji/models/HydraGNN/perf-runs/PENDING-$$ +mkdir -p "$HG_OUTPUT_DIR" + +cd /home/aaji/git/ai4science-studio +OUT=$(sbatch \ + --partition=lux \ + --account=vultr_lux \ + --nodes=2 \ + --ntasks-per-node=8 \ + --gpus-per-node=8 \ + --cpus-per-task=16 \ + --time=00:30:00 \ + --output="${HG_OUTPUT_DIR}/hydragnn-train-%j.out" \ + --error="${HG_OUTPUT_DIR}/hydragnn-train-%j.out" \ + material_science/models/HydraGNN/examples/sbatch_train_perf_amd.sh) +JOBID=$(echo "$OUT" | awk '{print $NF}') +echo "Submitted JOBID=$JOBID" +``` + +The sbatch wrapper inherits these env vars and exports `HG_CONFIG_OVERRIDE`, `OMNISTAT_VENV`, `OMNISTAT_CONFIG`, `OMNISTAT_USERMODE_INTERVAL`, `PROFILE_RANK0_ONLY=1` to the rank script. + +### 5. Wait for terminal state + +Poll `sacct -j $JOBID -X -n --format=State,ExitCode,Elapsed,NodeList -P` every 30 s. Treat any of `COMPLETED|FAILED|TIMEOUT|CANCELLED|NODE_FAIL|OUT_OF_MEMORY` as terminal. **Hard ceiling 45 minutes** of polling — if not terminal, write `state=TIMEOUT_WAITING` to manifest and exit 1. + +### 6. Move artifacts into final perf-run dir + +Once terminal: + +```bash +PERF_RUN=/shared/aaji/models/HydraGNN/perf-runs/${JOBID} +mv "$HG_OUTPUT_DIR" "$PERF_RUN" +# Find the rank-0 trace +TRACE=$(find "$PERF_RUN/logs" -name '*.pt.trace.json*' 2>/dev/null | head -1 || true) +# Find the omnistat DB (the sbatch wrapper writes it under the perf-run dir) +DB="$PERF_RUN/omnistat-db" +``` + +### 7. Write the manifest + +Use Python in the omnistat venv to write the manifest JSON exactly per the schema above. Include all `tools_versions` captured in step 1. + +### 8. Final stdout + +Print: +``` +STATUS=ok; reason=jobid= state= runtime=s perf_run= +``` + +## Failure modes (must surface, never silently swallow) + +| Failure | Action | +|---|---| +| Merge conflict in step 1 | Exit 2 with the message in the script. | +| `sbatch` returns nonzero | Exit 3, print sbatch stderr verbatim. | +| Job state == `FAILED` | Still write the manifest (with `exit_state=FAILED`), but **also** print the last 50 lines of the SLURM out to help downstream subagents skip cleanly. | +| No trace file found | Manifest gets `trace_paths=[]`. tracelens_analyst will gracefully no-op. | +| omnistat-db empty | Manifest still written; omnistat_analyst will detect and report it. | + +## Notes for the implementing agent + +- The launcher subagent runs on the **login node**. Do not call `srun` from inside the launcher; only `sbatch`. +- Do not import torch in the launcher's venv — it should stay lightweight. +- All paths absolute, never relative. +- Use `flock` on the install step if multiple invocations may race: `( flock -x 9; install_steps; ) 9>$TOOLS/.install.lock`. diff --git a/material_science/models/HydraGNN/recipes/perf-analysis/agents/omnistat_analyst.md b/material_science/models/HydraGNN/recipes/perf-analysis/agents/omnistat_analyst.md new file mode 100644 index 0000000..c34f6e4 --- /dev/null +++ b/material_science/models/HydraGNN/recipes/perf-analysis/agents/omnistat_analyst.md @@ -0,0 +1,143 @@ +# omnistat_analyst subagent + +Drive `omnistat-inspect` (PR #271) through the analyze-job phases on the user-mode VictoriaMetrics database, then map findings to the bottleneck taxonomy. + +## Inputs + +- `/manifest.json` +- The Omnistat DB at `manifest.omnistat_db_path`. +- The `omnistat-pr271` venv at `/shared/aaji/tools/omnistat-pr271/`. +- The VictoriaMetrics binary at `/shared/aaji/tools/victoriametrics/victoria-metrics-prod`. + +## Outputs + +- `/omnistat/inspect_outputs/db_info.json` +- `/omnistat/inspect_outputs/job_info.json` +- `/omnistat/inspect_outputs/data_check.json` +- `/omnistat/inspect_outputs/health.json` +- `/omnistat/inspect_outputs/stats_global.json` +- `/omnistat/inspect_outputs/stats__.json` — for any anomalous category, drilled to `node` and `gpu-id`/`interface-id`. +- `/omnistat/report_summary.md` +- `/omnistat/claims.json` (same schema as `tracelens/claims.json`) + +## Steps + +### 1. Start VictoriaMetrics on the DB (login node) + +Following the [load-database SKILL](https://github.com/ROCm/omnistat/blob/jorda/skills/skills/load-database/SKILL.md), with one Lux-specific addition: pass `-fs.disableMmap` so VM can open a job-scoped DB inside the login-node cgroup (without it, even a 1.6 MB DB panics with "cannot mmap file with size 4096 bytes ... no such device"). + +```bash +DB=$(jq -r .omnistat_db_path ) +[[ -d "$DB/data" ]] || { echo "STATUS=fail; reason=no_db_data"; exit 1; } + +# Pick free port +PORT=8428 +ss -lnt "sport = :$PORT" 2>/dev/null | grep -q LISTEN && PORT=8429 + +VMLOG=$PERF_RUN/omnistat/vm.log +mkdir -p "$PERF_RUN/omnistat/inspect_outputs" +nohup /shared/aaji/tools/victoriametrics/victoria-metrics-prod \ + -storageDataPath="$DB" \ + -httpListenAddr=127.0.0.1:$PORT \ + -retentionPeriod=100y \ + -search.disableCache \ + -search.latencyOffset=0 \ + -search.maxPointsPerTimeseries=90000 \ + -fs.disableMmap \ + > "$VMLOG" 2>&1 & +VM_PID=$! +echo $VM_PID > "$PERF_RUN/omnistat/vm.pid" +TSDB_URL="http://127.0.0.1:$PORT" + +# Wait up to 15 s for ready +for i in {1..15}; do + sleep 1 + curl -sf "$TSDB_URL/api/v1/status/tsdb" > /dev/null && break +done + +# Trap-style cleanup is the orchestrator's job; this subagent leaves VM running +# so the omnistat_verifier can hit the same endpoint. +``` + +### 2. Run analyze-job phases (per the SKILL) + +```bash +. /shared/aaji/tools/omnistat-pr271/bin/activate +SCRATCH=$PERF_RUN/omnistat/scratch +mkdir -p "$SCRATCH" +JOBID=$(jq -r .jobid ) + +OUTD=$PERF_RUN/omnistat/inspect_outputs + +omnistat-inspect --tsdb-url $TSDB_URL db info > $OUTD/db_info.json +omnistat-inspect --tsdb-url $TSDB_URL --scratch-dir $SCRATCH job $JOBID info > $OUTD/job_info.json +omnistat-inspect --tsdb-url $TSDB_URL --scratch-dir $SCRATCH job $JOBID data-check > $OUTD/data_check.json +omnistat-inspect --tsdb-url $TSDB_URL --scratch-dir $SCRATCH job $JOBID health > $OUTD/health.json +omnistat-inspect --tsdb-url $TSDB_URL --scratch-dir $SCRATCH job $JOBID stats --level global > $OUTD/stats_global.json +``` + +### 3. Identify anomalous categories and drill down + +For each category in `stats_global.json` (`gpu`, `host`, `network`, `xgmi`, `vendor`): + +- Compute coefficient of variation (`stddev/mean`) for utilization-like gauges. +- Flag a category as anomalous if **any** of: + - Mean GPU utilization < 60% (high cpu_dispatch / dataloader / comm risk) + - Network rx/tx rate stddev/mean > 0.3 (interface imbalance) + - VRAM stddev/mean > 0.2 (workload heterogeneity or memory leak) + - HBM throttle / power throttle counts > 0 + - XGMI traffic stddev/mean > 0.3 (unbalanced intra-node comm) + +For each flagged category, run `--level node` and either `--level gpu-id` (for `gpu`/`xgmi`) or `--level interface-id` (for `network`): + +```bash +for cat in gpu network xgmi host vendor; do + if is_anomalous "$cat"; then + omnistat-inspect --tsdb-url $TSDB_URL --scratch-dir $SCRATCH job $JOBID stats --category $cat --level node > $OUTD/stats_${cat}_node.json + case $cat in + gpu|xgmi) fine=gpu-id ;; + network) fine=interface-id ;; + *) fine="" ;; + esac + [[ -n "$fine" ]] && omnistat-inspect --tsdb-url $TSDB_URL --scratch-dir $SCRATCH job $JOBID stats --category $cat --level $fine > $OUTD/stats_${cat}_${fine//-/_}.json + fi +done +``` + +### 4. Translate findings into claims + +Concrete rules for the first iteration: + +- Mean GPU util < 50% → `class=cpu_dispatch` (or `dataloader`, lower confidence) — magnitude = `1 - mean_util`. +- HBM `vram_used_percentage` near 100% **and** `omnistat_fom` low → `class=gpu_memory_hbm`. +- Power throttling > 0 events → `class=gpu_compute` with remedy "raise power cap if available; otherwise system limit". +- Inter-node `network` rate << intra-node `xgmi` rate AND comm-heavy workload → `class=comm_scaleout` with remedy "verify ANP plugin via `NCCL_DEBUG=INFO`". +- High `network` interface CV at `interface-id` level → `class=comm_scaleout` with remedy "specific interface (e.g. ionic_3) is degraded; investigate cable/firmware". +- Hot GPU (>peak temp threshold from `gpus/mi355x.md` if present, else conservative 90°C) → `class=gpu_compute` with remedy "thermal-bound; check airflow". +- Per-node CPU mem usage CV > 0.3 → `class=host_cpu` with remedy "rebalance num_workers or check NUMA pinning". + +### 5. Emit report and claims + +`report_summary.md` should mirror Phase 1-3 of the analyze-job SKILL: discovery → data check → health → global stats → drill-down narrative. Cap at ~80 lines. + +`claims.json` schema is identical to the tracelens version. Cap at 6 claims. + +### 6. Final stdout + +``` +STATUS=ok; reason= claims; vm_port=; vm_pid= +``` + +The orchestrator (or the verifier, when finished) is responsible for stopping VM via `kill $(cat /omnistat/vm.pid)`. + +## Failure modes + +- DB empty (no metrics): write a single claim `class=host_io` `hypothesis="omnistat-usermode never collected"` and STATUS=partial. Likely cause: collector/server didn't start; check sbatch wrapper logs. +- omnistat-inspect not on PATH: STATUS=fail; the launcher should have installed it. +- VictoriaMetrics startup timeout (15 s): STATUS=fail; print last 20 lines of `vm.log`. + +## Notes + +- This subagent has read-write access to the perf-run dir. +- It must NOT modify or delete the omnistat DB. +- Never query `atlvrmonad01:9090` — it's IP-blocked and not the right database for this job anyway. diff --git a/material_science/models/HydraGNN/recipes/perf-analysis/agents/omnistat_verifier.md b/material_science/models/HydraGNN/recipes/perf-analysis/agents/omnistat_verifier.md new file mode 100644 index 0000000..e1d4cc1 --- /dev/null +++ b/material_science/models/HydraGNN/recipes/perf-analysis/agents/omnistat_verifier.md @@ -0,0 +1,92 @@ +# omnistat_verifier subagent + +Independently re-derive the top 2-3 claims from `omnistat/claims.json` by issuing raw PromQL via `curl` against the same VictoriaMetrics endpoint, at the **finest sampling step**. Optionally probe one cheap remedy on a 1-node interactive `srun`. + +## Inputs + +- `/manifest.json` +- `/omnistat/claims.json` +- VictoriaMetrics already running (PID at `/omnistat/vm.pid`, port from analyst output). +- `/omnistat/inspect_outputs/job_info.json` (for time range and sampling interval). + +## Outputs + +- `/omnistat/verified_claims.json` (same schema as `tracelens/verified_claims.json`). + +## Steps + +### 1. Re-derive top claims via raw PromQL + +Read the time range from `job_info.json` (`start_time`, `end_time`, `sampling_interval`). For each top claim, issue a `query_range` at the finest step (`max(sampling_interval, runtime/90000)`). + +Examples: + +**Important:** `rocm_*` and `omnistat_host_*` metrics do **not** carry a `jobid` label directly. Filter by joining with `rmsjob_info` per-instance: + +```promql +metric_name * on (instance) group_left() (max by (instance) (rmsjob_info{jobid="$JOBID"})) +``` + +A naive `rocm_utilization_percentage{jobid="..."}` will return 0 series even though the data is there. Use the join in every query. + +Also: `start_time`/`end_time` from `job_info.json` are **UTC**. Convert with `calendar.timegm(time.strptime(...))` (NOT `time.mktime`, which assumes local time) before passing to PromQL. + +| Claim | PromQL | +|---|---| +| "Mean GPU util = X%" | `avg(rocm_utilization_percentage * on (instance) group_left() (max by (instance) (rmsjob_info{jobid="$JOBID"})))` | +| "VRAM near 100% on N nodes" | `max(rocm_vram_used_percentage * on (instance) group_left() (max by (instance) (rmsjob_info{jobid="$JOBID"})))` per `instance` | +| "Network tx rate plateau" | `rate(omnistat_network_tx_bytes * on (instance) group_left() (max by (instance) (rmsjob_info{jobid="$JOBID"}))[60s])` peak | +| "XGMI imbalance" | per-card stddev of `rocm_xgmi_*_bytes` rate, joined as above | +| "Power throttling events" | sum increase of throttle counters over the job span, joined as above | + +```bash +PORT=$(jq -r '.[0].port' ) # or pull from analyst stdout +curl -sG "http://127.0.0.1:$PORT/api/v1/query_range" \ + --data-urlencode "query=avg_over_time(rocm_utilization_percentage{jobid=\"$JOBID\"}[1m])" \ + --data-urlencode "start=$START" \ + --data-urlencode "end=$END" \ + --data-urlencode "step=$STEP" \ + | jq '.data.result | map(.values[][1] | tonumber) | add/length' +``` + +The verifier should re-derive the actual number from the PromQL result and compare to `claim.magnitude.value`. ±20% → `verdict=verified`. Outside → `refuted`. Empty result series → `inconclusive`. + +### 2. Optional remedy probe + +Same constraints as the tracelens verifier — at most one 1-node `srun -p lux -A vultr_lux -N1 --time=00:05:00`. Telemetry-side probes are typically configuration changes (e.g. enable rocprofiler `hbm` counter set in the omnistat config and re-run for 1 minute on 1 node to see if HBM bandwidth is actually saturating). + +If a probe needs the omnistat config to change, write a temp config file based on `omnistat-lux.config` with the rocprofiler section uncommented, point `OMNISTAT_CONFIG` at it, and submit via the existing sbatch script as a tiny job with `HG_NUM_EPOCH=1 HYDRAGNN_MAX_NUM_BATCH=10 --nodes=1`. Time-budget that branch generously (10 min) since it does include the omnistat collector startup. + +If multi-node is required (e.g. ANP plugin retest), set `remedy_probe.ran=false; notes="multi-node probe deferred"`. + +### 3. Emit `verified_claims.json` + +Same shape as the tracelens version. Each claim copied through with `verdict`, `verifier_evidence`, `verifier_value`, `remedy_probe`. + +### 4. Stop VictoriaMetrics + +After all queries are done: + +```bash +kill "$(cat /omnistat/vm.pid)" 2>/dev/null || true +``` + +The synthesizer doesn't need VM running. + +### 5. Final stdout + +``` +STATUS=ok; reason=v/r/i; probe= +``` + +## Failure modes + +- VM not running: try to start it the same way the analyst did before giving up; if still down, STATUS=fail. +- Query returns no data for a key metric the claim relies on: that's an `inconclusive` verdict, not a fail. +- Network blip mid-query: retry once with 5s sleep. + +## Notes + +- Use **only** `curl` for re-derivation, **not** `omnistat-inspect`. The whole point is an independent path. +- Use the finest step that VM allows (`runtime / 90000` floor); document the actual step used in `verifier_evidence`. +- Never delete the DB or any analyst output. diff --git a/material_science/models/HydraGNN/recipes/perf-analysis/agents/synthesizer.md b/material_science/models/HydraGNN/recipes/perf-analysis/agents/synthesizer.md new file mode 100644 index 0000000..476d884 --- /dev/null +++ b/material_science/models/HydraGNN/recipes/perf-analysis/agents/synthesizer.md @@ -0,0 +1,131 @@ +# synthesizer subagent + +Merge `tracelens/verified_claims.json` and `omnistat/verified_claims.json` into a single ranked, deduplicated bottleneck report. + +## Inputs + +- `/manifest.json` +- `/tracelens/verified_claims.json` +- `/omnistat/verified_claims.json` +- (Optional) `/tracelens/report_summary.md`, `/omnistat/report_summary.md` + +## Outputs + +- `/combined_report.md` + +## Steps + +### 1. Load and filter + +```python +import json +tl = json.load(open(f"{perf_run}/tracelens/verified_claims.json")) +om = json.load(open(f"{perf_run}/omnistat/verified_claims.json")) +all_claims = [{"src":"TL", **c} for c in tl] + [{"src":"OS", **c} for c in om] +``` + +Drop `verdict == "refuted"`. Keep `inconclusive` claims but tag them so the user sees the gap. + +### 2. Deduplicate by class + topic + +Group claims by `class`. Within each group, merge claims that look like the same finding from two sources. Heuristic: +- Both have `class=comm_scaleout` → almost certainly the same finding viewed from trace (NCCL kernels) and telemetry (network rates). Merge into one entry that lists both `src` values and both magnitudes. +- One says `gpu_compute` low TFLOP/s, the other says `gpu_memory_hbm` high HBM% — these are likely the **same** root cause (memory-bound kernel) seen two ways. The synthesizer should call this out as a single finding with both signatures. + +### 3. Rank + +Score = `magnitude.value × confidence_weight × (corroborated ? 1.5 : 1.0)`. +`confidence_weight`: high=1.0, medium=0.7, low=0.4. `corroborated` = both TL and OS contributed to the merged entry. + +### 4. Tag "system limit reached" + +A claim is a system limit if any of: +- `proposed_remedy is null` +- both `verdict=verified` and `remedy_probe.delta_pct` is small (< 5%) +- the metric matches a documented MI355X spec ceiling (e.g. fp64 39 TFLOP/s, HBM 8 TB/s, ANP scale-out ~25 GB/s) + +Mark these explicitly so the report doesn't promise a fix that won't materialize. + +### 5. Write `combined_report.md` + +Layout: + +```markdown +# HydraGNN bottleneck analysis — JOBID + +**Run:** 2 nodes × 8 GPUs (MI355X), s, precision=

, batch= +**Sources:** TraceLens v on rank-0 trace, Omnistat user-mode (s sampling) +**Verdict counts:** // + +## Executive summary + +<1 paragraph, 4-5 sentences. Named bottleneck class first, magnitude, then "system-limit" or "improvable"> + +## Ranked bottlenecks + +### 1. — [score: ] [] [] + +**Hypothesis:** +**Magnitude:** +**Evidence:** +**Remedy:** — **tried?** ; **delta:** +**Verifier note:** + +### 2. ... + +## Remedies tried in this session + +| Class | Probe | Baseline | After remedy | Delta | +|---|---|---|---|---| +| ... | + +## Remedies proposed but not tried (cost/risk) + +| Class | Remedy | Why deferred | +|---|---|---| + +## Limits reached + +| Class | Spec ceiling | Observed | +|---|---|---| + +## Inconclusive claims + +| Source | Class | Hypothesis | Why inconclusive | +|---|---|---|---| + +## Next steps + +1. +2. +3. + +## Artifacts + +- `tracelens/report.xlsx` +- `tracelens/verified_claims.json` +- `omnistat/inspect_outputs/` +- `omnistat/verified_claims.json` +- `manifest.json` + +--- +Generated by `material_science/models/HydraGNN/recipes/perf-analysis/`. +``` + +### 6. Final stdout + +``` +STATUS=ok; reason=; top=; tried= +``` + +## Style rules + +- Specific numbers from evidence — never vague ("most of the time", "huge"). +- If TL and OS disagree on a number that's not refuted, show **both** with sources rather than pick one. +- When `remedy_probe.ran=false` because of multi-node requirement, say so explicitly. +- Do not editorialize about HydraGNN as a model; only about the run. + +## Failure modes + +- Both inputs missing → STATUS=fail. +- Only one input present → emit a partial report with a banner explaining which side is missing; STATUS=partial. diff --git a/material_science/models/HydraGNN/recipes/perf-analysis/agents/tracelens_analyst.md b/material_science/models/HydraGNN/recipes/perf-analysis/agents/tracelens_analyst.md new file mode 100644 index 0000000..b33572d --- /dev/null +++ b/material_science/models/HydraGNN/recipes/perf-analysis/agents/tracelens_analyst.md @@ -0,0 +1,124 @@ +# tracelens_analyst subagent + +Generate a TraceLens performance report for the rank-0 PyTorch trace and produce structured `claims.json` mapped to the bottleneck taxonomy. + +## Inputs + +- `/shared/aaji/models/HydraGNN/perf-runs//manifest.json` + +## Outputs + +- `/tracelens/report.xlsx` — the raw TraceLens Excel. +- `/tracelens/report_summary.md` — human-readable digest of the workbook (top-N kernels, GPU idle %, comm vs compute split, roofline categorization). +- `/tracelens/claims.json` — structured claims (schema below). + +## Claim schema + +Every entry: + +```json +{ + "id": "TL-001", + "class": "gpu_compute|gpu_memory_hbm|cpu_dispatch|comm_xgmi|comm_scaleout|dataloader|host_io|host_cpu", + "hypothesis": "", + "magnitude": { + "metric": "", + "value": 0.34, + "unit": "fraction|seconds|GBps|TFLOPs|count" + }, + "evidence_paths": ["::", ":"], + "evidence_excerpt": "<200 chars max>", + "confidence": "high|medium|low", + "proposed_remedy": "", + "remedy_test_command_or_null": "srun -p lux -A vultr_lux -N1 --time=00:05:00 ... | null" +} +``` + +## Steps + +### 1. Sanity check + +```bash +. /shared/aaji/tools/omnistat-pr271/bin/activate +PERF_RUN=$(jq -r .perf_run_dir ) +TRACE=$(jq -r '.trace_paths[0] // empty' ) +[[ -z "$TRACE" ]] && { echo "STATUS=partial; reason=no_trace_in_manifest"; exit 0; } +mkdir -p "$PERF_RUN/tracelens" +``` + +### 2. Run TraceLens + +```bash +cd "$PERF_RUN/tracelens" +TraceLens_generate_perf_report_pytorch \ + --profile_json_path "$TRACE" \ + --output_xlsx_path report.xlsx \ + --output_csvs_dir csvs/ \ + --enable_kernel_summary \ + --short_kernel_study 2>&1 | tee tracelens.log +``` + +CLI confirmed against TraceLens v0.1.0+ on Lux: `--output_xlsx_path` controls the workbook path; `--output_csvs_dir` writes per-sheet CSVs (handy for the verifier). `--enable_kernel_summary` adds the kernel-level summary sheet. `--short_kernel_study` flags very short kernels (likely launch-overhead bound). + +If multi-node, the script will detect collective ops and produce a separate "Collective Analysis" sheet unless `--disable_coll_analysis` is passed (don't pass it). + +### 3. Extract digest from the workbook + +Use `openpyxl` (already a TraceLens dep) in the venv: + +```python +import openpyxl, json, pathlib +wb = openpyxl.load_workbook("report.xlsx", data_only=True, read_only=True) +sheets = wb.sheetnames # expected: GPU_Timeline, Ops_Summary, Roofline, NN_Module, ... +``` + +Pull these specific signals (skip a sheet gracefully if absent): + +| Sheet | Signals to extract | +|---|---| +| `GPU_Timeline` | total_gpu_time, idle_time, kernel_time, comm_time (sum durations of NCCL/RCCL events) | +| `Ops_Summary` | top-10 ops by total time + their percent | +| `Roofline` | per-op TFLOP/s and TB/s; flag any in compute-bound vs memory-bound region | +| `NN_Module` | top-3 modules by GPU time | + +Write the digest to `report_summary.md` as a few markdown tables. + +### 4. Translate findings into claims + +Concrete rules for the first iteration (the agent may add more): + +- If `comm_time / total_step_time > 0.20` → `class=comm_xgmi` if intra-node only OR `comm_scaleout` if inter-node (detect by NCCL kernel name `AllReduce` + `nNodes>1` flag). Magnitude = the ratio. +- If `idle_time / total_step_time > 0.15` → `class=cpu_dispatch` (could also be dataloader; raise as `confidence=medium`). +- For each top-3 op: classify by Roofline region. Compute-bound + low TFLOP/s → `class=gpu_compute`. Memory-bound + high TB/s → `class=gpu_memory_hbm`. Memory-bound + low TB/s → `class=gpu_memory_hbm` with confidence=low + remedy "verify with HBM counters via Omnistat rocprofiler probe". +- If `dataloader` events appear in the timeline (look for `enumerate(DataLoader)#`) and their cumulative time > 5% → `class=dataloader`. + +### 5. Emit `claims.json` + +Cap at 6 claims. Order by `magnitude.value` (when comparable). Each claim must include a remedy that is concrete and **verifiable**. Examples: + +| Hypothesis | Remedy | +|---|---| +| "RCCL allreduce dominates 35% of step time" | "Try `RCCL_LL128_FORCE_ENABLE=0` or larger bucket size via `find_unused_parameters=False`" | +| "fp64 GEMM at 12 TFLOP/s, well below MI355X 39 TFLOP/s peak" | "Test `HG_PRECISION=bf16` for 5 batches on 1 node" | +| "DataLoader events occupy 18% of step time" | "Increase `HYDRAGNN_NUM_WORKERS` from 0 to 4" | +| "fp64 SQ FMA approaches MI355X spec ceiling" | "null — system limit; only path is bf16/fp32" | + +A null remedy is acceptable for "limit reached" findings. + +### 6. Final stdout + +``` +STATUS=ok; reason= claims, top=: +``` + +## Failure modes + +- TraceLens import error → STATUS=fail; do not attempt to "patch" the install. +- Trace too large to load (>2 GB) → write a single claim of class `cpu_dispatch` with hypothesis "trace too large to fully analyze; profiler captured >X seconds — schedule may be misconfigured". +- No NCCL events in trace → multi-node still ran but profiler missed comm; flag as `inconclusive` rather than claiming "comm is fast". + +## Notes + +- Do NOT run `omnistat-inspect`; that's the other analyst's job. If you need the omnistat DB to validate something, hand the question to the omnistat_verifier instead. +- Do NOT modify the raw trace file. +- Keep `evidence_paths` machine-parseable so the verifier can re-derive the number from the same source. diff --git a/material_science/models/HydraGNN/recipes/perf-analysis/agents/tracelens_verifier.md b/material_science/models/HydraGNN/recipes/perf-analysis/agents/tracelens_verifier.md new file mode 100644 index 0000000..a9494ae --- /dev/null +++ b/material_science/models/HydraGNN/recipes/perf-analysis/agents/tracelens_verifier.md @@ -0,0 +1,108 @@ +# tracelens_verifier subagent + +Independently re-derive the top 2-3 claims from `tracelens/claims.json` against the raw `.pt.trace.json` and (optionally) probe **one** cheap remedy on a 1-node interactive `srun`. + +## Inputs + +- `/manifest.json` +- `/tracelens/claims.json` +- The raw trace at `manifest.trace_paths[0]`. + +## Outputs + +- `/tracelens/verified_claims.json` — same schema as `claims.json` plus a `verdict` field per entry. + +## Verdict schema + +Each input claim is copied through with one new field: + +```json +{ + "...all fields from claims.json...": "...", + "verdict": "verified|refuted|inconclusive", + "verifier_evidence": "<200 chars: how you re-derived it>", + "verifier_value": , + "remedy_probe": { + "ran": true|false, + "command": "", + "baseline_seconds_per_step": , + "remedy_seconds_per_step": , + "delta_pct": , + "notes": "" + } +} +``` + +## Steps + +### 1. Re-derive top 3 claims from the raw JSON + +```python +import json, gzip +from pathlib import Path +def open_trace(p): + return gzip.open(p, 'rt') if str(p).endswith('.gz') else open(p, 'r') +with open_trace(trace) as f: + data = json.load(f) +events = data["traceEvents"] +``` + +For each top claim, do a parallel re-derivation: + +| Claim class | Re-derivation | +|---|---| +| `comm_xgmi` / `comm_scaleout` | sum `dur` for events whose `name` matches `^(nccl|rccl|ncclKernel|AllReduce|AllGather|ReduceScatter)`; divide by total span (`max(ts+dur)-min(ts)` of GPU events). Compare to claim's `magnitude.value`. | +| `gpu_compute` | filter `cat=='kernel'`, find top-N by `sum(dur)`; check the top kernel name + percent matches the claim. | +| `cpu_dispatch` | `(total_span - sum(GPU kernel dur)) / total_span` should match "idle %". | +| `dataloader` | sum `dur` of events with `name` containing `enumerate(DataLoader)`. | + +If verifier_value within ±20% of claim's value → `verdict=verified`. If outside → `verdict=refuted` with the actual number. If we can't isolate the events → `verdict=inconclusive`. + +### 2. Pick at most ONE remedy probe + +Choose the highest-impact `verified` claim with a non-null `remedy_test_command_or_null`. Constraints: + +- 1 node only (`-N1`). +- ≤5 minutes wall-time. +- Must produce a comparable number (steps/sec or seconds/step) to the original 2-node baseline. If the remedy can only be tested at scale (e.g. RCCL inter-node tweak), set `remedy_probe.ran=false; notes="remedy requires multi-node; deferred to next iteration"`. + +Example probe for the bf16 / fp64 ceiling claim — run via the existing HydraGNN sbatch script in interactive mode: + +```bash +srun -p lux -A vultr_lux -N1 --ntasks=8 --gpus-per-node=8 --cpus-per-task=16 \ + --time=00:05:00 --pty bash -c ' +export AI4S_SHARED_DIR=/shared/aaji +export HG_PRECISION=bf16 +export HYDRAGNN_MAX_NUM_BATCH=10 +export HG_NUM_EPOCH=1 +export HG_OUTPUT_DIR=/tmp/perf-probe-$$ +mkdir -p $HG_OUTPUT_DIR +bash /home/aaji/git/ai4science-studio/material_science/models/HydraGNN/examples/sbatch_train_amd.sh \ + 2>&1 | tee /tmp/probe-$SLURM_JOB_ID.out +grep -oE "[0-9.]+s/it" /tmp/probe-$SLURM_JOB_ID.out | tail -5 +' +``` + +Note: `sbatch_train_amd.sh` is designed for `sbatch`, not `srun --pty`. For probes prefer using its rank-script logic directly via a small ad-hoc wrapper rather than running the whole sbatch under srun. The verifier may inline a minimal Python invocation that bypasses sbatch and just runs the upstream `gfm_mlip_all_mpnn.py` for 10 batches inside the container with the alternate precision. + +If the verifier judges that no cheap probe is feasible, set `remedy_probe.ran=false` with a clear note and proceed. + +### 3. Emit `verified_claims.json` + +Same length and order as `claims.json`. Even refuted claims stay in the file — the synthesizer will drop them, but they should remain visible. + +### 4. Final stdout + +``` +STATUS=ok; reason=v/r/i; probe= +``` + +## Constraints + +- Never modify `claims.json` — only write `verified_claims.json`. +- Never run a 2-node `srun` or `sbatch` from this subagent. +- If the trace file is missing, write `verified_claims.json` as `[]` and `STATUS=partial; reason=no_trace`. + +## Why two phases (analyst + verifier)? + +Hallucination defense. The analyst can produce a number from a tool's report; the verifier re-derives it from raw events. Disagreement => the user gets to see both numbers and the synthesizer flags the conflict in the final report rather than silently picking one. This is the same approach the omnistat side uses with PromQL-from-curl re-derivation. diff --git a/material_science/models/HydraGNN/recipes/perf-analysis/omnistat-lux.config.template b/material_science/models/HydraGNN/recipes/perf-analysis/omnistat-lux.config.template new file mode 100644 index 0000000..c06b7ce --- /dev/null +++ b/material_science/models/HydraGNN/recipes/perf-analysis/omnistat-lux.config.template @@ -0,0 +1,81 @@ +#-- +# Omnistat user-mode config for Lux (MI355X) +# +# Adapted from HydraGNN's upstream omnistat.hydragnn-external.config. +# Differences from upstream / from system-mode default on Lux: +# - enable_amd_smi (newer SMI), not rocm_smi +# - rocprofiler counters DISABLED in v1 (turn on per-probe via verifier) +# - allowed_ips includes 0.0.0.0 so the user-mode VM scrapes +# - victoria_datadir under /shared/aaji (NFS, persists after job) +#-- + +[omnistat.collectors] + +port = 8101 +enable_rocm_smi = False +enable_amd_smi = True +enable_ras_ecc = True +enable_rms = True +enable_rocprofiler = False +enable_network = True +enable_vendor_counters = True +enable_xgmi = True +enable_cu_occupancy = False + +# Path to local ROCm install (Lux: 7.2.x) +rocm_path = /opt/rocm + +# allowed scrapers: localhost (the user-mode VM sits on the same node) + 0.0.0.0 +# wildcard for any user-mode VictoriaMetrics that might run on a sibling node. +allowed_ips = 127.0.0.1, 0.0.0.0 + +[omnistat.collectors.rms] + +host_skip = "login.*" +enable_annotations = True +# Lux note: /tmp/omni_rmsjobinfo (the upstream default) is owned by user +# `omnidc` on Lux compute nodes for the system-mode collector, so user-mode +# omnistat-rms-env hits PermissionError if it tries to write there. The +# sbatch wrapper substitutes @JOB_DIR@ with the job-scoped perf-run dir +# (NFS-mounted on every compute node) before passing the resolved config +# to omnistat-usermode --configfile. +job_detection_mode = file-based +job_detection_file = @JOB_DIR@/omni_rmsjobinfo +step_detection_file = @JOB_DIR@/omni_rmsjobinfo_step + +[omnistat.collectors.host] +enable_proc_io_stats = True + +# rocprofiler section kept as a reference for the verifier probe — disabled +# above. Enable enable_rocprofiler=True + select a profile here for HBM/FLOPS +# counters on a 1-node remedy probe only. +[omnistat.collectors.rocprofiler] +profile = hbm + +[omnistat.collectors.rocprofiler.hbm] +sampling_mode = constant +counters = ["FETCH_SIZE", "WRITE_SIZE"] + +[omnistat.query] +prometheus_url = http://localhost:8428 +system_name = lux + +#-- +# User-mode settings — VictoriaMetrics + per-job datadir +# (Like the [rms] section above, @JOB_DIR@ is substituted by the sbatch +# wrapper at runtime; ConfigParser's %(VAR)s interpolation does NOT pull +# from os.environ, so we can't rely on %(SLURM_JOB_ID)s here.) +#-- +[omnistat.usermode] + +victoria_binary = /shared/aaji/tools/victoriametrics/victoria-metrics-prod +victoria_datadir = @JOB_DIR@/omnistat-db +victoria_logfile = vic_server.log + +# No external VM; we keep everything local to the job. +external_victoria = False + +# Bind exporter to a single core to keep monitoring overhead off the data path. +exporter_corebinding = 0 + +ssh_key = ~/.ssh/id_rsa From 478532552dec25bf9921e25f88147c9276831aad Mon Sep 17 00:00:00 2001 From: Ashwin Aji Date: Thu, 21 May 2026 07:16:00 +0000 Subject: [PATCH 02/31] HydraGNN: bump to upstream main + add persistent_workers patch MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Bumps the pinned HydraGNN SHA from 6c45f16 to upstream main (2fb0bd0, 2026-05-20) for both sbatch_train_amd.sh and sbatch_train_perf_amd.sh, and rebuilds the Apptainer overlay on top of that SHA with a small, opt-in patch that exposes a HYDRAGNN_PERSISTENT_WORKERS env var. The existing HYDRAGNN_NUM_WORKERS knob is unchanged; together the two let the dataloader keep workers alive across epochs instead of forking them every epoch boundary. Why this matters: single-node MI355X probes with NW=8 + PW=1 reduce the steady-state per-iteration time from 2.35 s/it (baseline NW=0, PW=0) to ~0.25 s/it — a ~9x speedup — and eliminate the bimodal iteration-time distribution the perf-analysis subagents had previously flagged as the primary cause of >96% GPU idle time in job 6762. Build script changes: - Replace `git clone --depth=1 + git fetch --depth=1 ` with a full clone + checkout, since some remote configs reject fetch-by-sha on shallow clones. - Stage examples/patches/*.patch and apply them inside the container with `patch -p1 --dry-run` first so a stale patch fails fast. - Beef up the verify step to print `hydragnn.__file__` and confirm the patch is present in the installed source. sbatch / probe scripts: - Drop the rank-0-only `python3 -c "import hydragnn"` verify block. That import triggers `from mpi4py import MPI` which calls MPI_Init() on rank 0 alone; under `srun --mpi=pmix` this leaves ranks 1..N-1 hanging in their later MPI_Init() and the whole job deadlocks in `dist.init_process_group → _create_c10d_store`. Move the same verification logic inside the main python where all ranks align at the same MPI_Init. Skill update: - Add an "Anti-pattern: rank-0 mpi4py import before srun" section to ai4science-studio SKILL.md so a fresh session recognizes the symptom (all ranks stuck in `_create_c10d_store`) and the cure. Patch is documented in examples/patches/README.md and is opt-in: the default behaviour with HYDRAGNN_PERSISTENT_WORKERS unset matches stock upstream main. Co-authored-by: Cursor --- .cursor/skills/ai4science-studio/SKILL.md | 52 ++++++++++++++ .../HydraGNN/examples/build_overlay_amd.sh | 68 +++++++++++++++++-- .../HydraGNN/examples/patches/README.md | 40 +++++++++++ .../load_data.persistent_workers.patch | 15 ++++ .../HydraGNN/examples/sbatch_train_amd.sh | 2 +- .../examples/sbatch_train_perf_amd.sh | 23 ++++++- .../omnistat-lux.config.template | 36 ++++++++-- 7 files changed, 221 insertions(+), 15 deletions(-) create mode 100644 material_science/models/HydraGNN/examples/patches/README.md create mode 100644 material_science/models/HydraGNN/examples/patches/load_data.persistent_workers.patch diff --git a/.cursor/skills/ai4science-studio/SKILL.md b/.cursor/skills/ai4science-studio/SKILL.md index d8f07a6..4c340e6 100755 --- a/.cursor/skills/ai4science-studio/SKILL.md +++ b/.cursor/skills/ai4science-studio/SKILL.md @@ -349,6 +349,58 @@ When using `srun --mpi=pmix` with Apptainer, ranks may segfault with `PMIx ERROR --env PMIX_MCA_psec=native ``` +## Never `import` mpi4py-bearing modules from a rank-0-only python before the real `srun` + +If a model script does `from mpi4py import MPI` at import time (most do — directly, or transitively via the model's `__init__.py`), the import calls `MPI_Init()` because `mpi4py.rc.initialize=True` is the default. Under `srun --mpi=pmix`, `MPI_Init()` is a **collective rendezvous**: every rank in the SLURM step must arrive at it together before any rank can return. + +**Anti-pattern that deadlocks the job:** +```bash +# Inside rank.sh — runs once per SLURM rank, 0..N-1 +if [[ "${SLURM_PROCID:-0}" == "0" ]]; then + python3 - <<'PYEOF' # short-lived rank-0-only python + import hydragnn # ← triggers mpi4py auto-MPI_Init on rank 0 alone + print("[verify] ...") +PYEOF +fi + +exec python -u -c "..." # the real distributed python on every rank +``` + +**What goes wrong:** rank 0 enters `MPI_Init()` from the verify python and waits there. Ranks 1..N-1 reach `MPI_Init()` from the main python a moment later. PMIx counts rank 0 as already-arrived for *that* rendezvous, so the verify python returns and exits — taking rank 0's MPI participation with it. When rank 0's main python then calls `MPI_Init()` again, PMIx no longer matches it with ranks 1..N-1, and the whole job hangs in `dist.init_process_group → _create_c10d_store` with no useful error message until the 30-minute TCPStore timeout. + +**Symptoms (py-spy on a stuck rank):** +``` +_create_c10d_store (torch/distributed/rendezvous.py:200) +_env_rendezvous_handler (torch/distributed/rendezvous.py:281) +init_process_group (torch/distributed/distributed_c10d.py:1806) +setup_ddp (hydragnn/utils/distributed/distributed.py:254) +``` +…on every rank except rank 0, which has died or is in a different MPI_Init that will never match. + +**Correct patterns (pick one):** + +1. **Move the verify inside the main python**, gated on `SLURM_PROCID == 0`. Same information gets logged, but mpi4py's `MPI_Init` happens once per rank, in lockstep: + ```python + exec python -u -c " + import os + if os.environ.get('SLURM_PROCID', '0') == '0': + import inspect, hydragnn + print('[verify]', hydragnn.__file__, flush=True) + # ...rest of training driver... + " + ``` + +2. **Disable mpi4py auto-init at the top of the verify script** (only safe if the verify never actually calls any MPI op): + ```python + import mpi4py + mpi4py.rc.initialize = False + import hydragnn # now safe to import without joining a rendezvous + ``` + +3. **Print the verify info on rank 0 from the launching shell only**, never from a python that imports mpi4py-bearing modules. + +This bit us once on the HydraGNN perf-analysis recipe (see git log on `aaji/perf-agents-hydragnn`). Watch for it whenever you add `python3 -c "..."` scaffolding around an `srun` that uses `--mpi=pmix`. + --- # Docker-specific patterns diff --git a/material_science/models/HydraGNN/examples/build_overlay_amd.sh b/material_science/models/HydraGNN/examples/build_overlay_amd.sh index d005b34..a915c28 100755 --- a/material_science/models/HydraGNN/examples/build_overlay_amd.sh +++ b/material_science/models/HydraGNN/examples/build_overlay_amd.sh @@ -40,7 +40,7 @@ HG_BASE="${AI4S_SHARED_DIR:?AI4S_SHARED_DIR must be set}/models/HydraGNN" HG_SIF="${HG_SIF:-${AI4S_SHARED_DIR}/images/pytorch_rocm7.2.2_ubuntu24.04_py3.12_pytorch_release_2.10.0.sif}" OVERLAY="${HG_OVERLAY:-${HG_BASE}/overlays/hydragnn-overlay.img}" OVERLAY_SIZE_MB="${HG_OVERLAY_SIZE_MB:-4096}" -HG_HYDRAGNN_SHA="${HG_HYDRAGNN_SHA:-6c45f1682783e66dc89e9e23009f61716186432b}" +HG_HYDRAGNN_SHA="${HG_HYDRAGNN_SHA:-2fb0bd0157e3c85a74f9841887155095bd163303}" PYG_ROCM_RELEASE="${PYG_ROCM_RELEASE:-15}" PYG_PYTHON_TAG="${PYG_PYTHON_TAG:-py312}" @@ -82,6 +82,27 @@ echo "" mkdir -p "$LOCAL_TMP" "$STAGE" "$PYG_WHEELS_DIR" +# --------------------------------------------------------------------------- +# Stage local patches (applied to the HydraGNN source inside the container) +# --------------------------------------------------------------------------- +SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)" +PATCHES_SRC="${SCRIPT_DIR}/patches" +PATCHES_STAGE="${LOCAL_TMP}/patches" +mkdir -p "$PATCHES_STAGE" +if [[ -d "$PATCHES_SRC" ]]; then + shopt -s nullglob + for p in "$PATCHES_SRC"/*.patch; do + cp "$p" "$PATCHES_STAGE/" + done + shopt -u nullglob +fi +PATCH_COUNT=$(ls "$PATCHES_STAGE"/*.patch 2>/dev/null | wc -l) +echo " Local patches to apply: $PATCH_COUNT" +if [[ "$PATCH_COUNT" -gt 0 ]]; then + for p in "$PATCHES_STAGE"/*.patch; do echo " - $(basename "$p")"; done +fi +echo "" + # --------------------------------------------------------------------------- # Download and extract Looong01 PyG ROCm wheels # --------------------------------------------------------------------------- @@ -153,15 +174,42 @@ cp "$ADIOS2_INST"/lib/libadios2*.so* "$STAGE/adios2/" 2>/dev/null || true echo " adios2 installed (MPI + Python bindings)" cd / -echo "--- Step 5: HydraGNN source (main branch) ---" +echo "--- Step 5: HydraGNN source (pinned SHA: $BUILD_SHA) ---" if [[ ! -d /tmp/HydraGNN-src ]]; then - git clone -q --depth=1 https://github.com/ORNL/HydraGNN.git /tmp/HydraGNN-src - cd /tmp/HydraGNN-src && git fetch --depth=1 origin "$BUILD_SHA" && git checkout "$BUILD_SHA" + # Full clone (not --depth=1) so arbitrary SHAs can be checked out; some + # remotes refuse fetch-by-sha on shallow clones depending on + # uploadpack.allowReachableSHA1InWant config. + git clone -q https://github.com/ORNL/HydraGNN.git /tmp/HydraGNN-src + cd /tmp/HydraGNN-src && git checkout -q "$BUILD_SHA" + HEAD_SHA="$(git rev-parse HEAD)" + if [[ "$HEAD_SHA" != "$BUILD_SHA" ]]; then + echo "ERROR: checkout produced $HEAD_SHA, expected $BUILD_SHA" >&2 + exit 1 + fi cd / fi + +echo "--- Step 5b: apply local patches (BUILD_PATCHES=$BUILD_PATCHES) ---" +if [[ -d "$BUILD_PATCHES" ]]; then + cd /tmp/HydraGNN-src + shopt -s nullglob + for p in "$BUILD_PATCHES"/*.patch; do + echo " applying: $(basename "$p")" + # --check first so we fail fast with a clear message if the patch + # was authored against a different SHA than BUILD_SHA. + if ! patch -p1 --dry-run --silent < "$p"; then + echo "ERROR: patch $(basename "$p") does not apply cleanly to SHA $BUILD_SHA" >&2 + exit 1 + fi + patch -p1 < "$p" + done + shopt -u nullglob + cd / +fi + # Copy the source tree directly rather than pip-installing (avoids partial package issues) cp -r /tmp/HydraGNN-src/hydragnn "$STAGE/hydragnn" -echo " HydraGNN source copied (SHA: $BUILD_SHA)" +echo " HydraGNN source copied (SHA: $BUILD_SHA + $(ls "$BUILD_PATCHES"/*.patch 2>/dev/null | wc -l) patches)" echo "--- Step 6: strip torch / nvidia / triton / cuda from staging dir ---" for pkg in torch torchvision torchaudio torchgen functorch nvidia triton cuda sympy; do @@ -202,7 +250,14 @@ print('adios2 :', adios2.__version__, '(MPI: OK)') python3 -c "import vesin; print('vesin : ok')" python3 -c "import e3nn; print('e3nn :', e3nn.__version__)" -python3 -c "import hydragnn; print('hydragnn : ok')" +python3 -c " +import hydragnn, os, inspect +print('hydragnn : imported from', os.path.dirname(hydragnn.__file__)) +# Confirm the persistent_workers patch is present iff a patch file was applied. +src = inspect.getsource(__import__('hydragnn.preprocess.load_data', fromlist=['_'])) +patched = 'HYDRAGNN_PERSISTENT_WORKERS' in src +print('hydragnn : persistent_workers patch =', patched) +" echo "=== Overlay install complete ===" INNEREOF @@ -224,6 +279,7 @@ apptainer exec \ --env BUILD_STAGE="$STAGE" \ --env BUILD_PYG_WHEELS="$PYG_WHEELS_DIR" \ --env BUILD_SHA="$HG_HYDRAGNN_SHA" \ + --env BUILD_PATCHES="$PATCHES_STAGE" \ "$HG_SIF" \ bash "${LOCAL_TMP}/overlay_install.sh" diff --git a/material_science/models/HydraGNN/examples/patches/README.md b/material_science/models/HydraGNN/examples/patches/README.md new file mode 100644 index 0000000..91aa844 --- /dev/null +++ b/material_science/models/HydraGNN/examples/patches/README.md @@ -0,0 +1,40 @@ +# HydraGNN local patches + +Small, opt-in patches applied during overlay build (see +`examples/build_overlay_amd.sh`). Each patch is unified-format and applies +cleanly with `patch -p1` against the pinned `HG_HYDRAGNN_SHA` in the build +script. + +## load_data.persistent_workers.patch + +**Target SHA:** `2fb0bd0157e3c85a74f9841887155095bd163303` (upstream `main`, 2026-05-20) + +**Lines changed:** +4 / -0 in `hydragnn/preprocess/load_data.py` + +**What it does** + +Upstream already exposes `num_workers` as an env-var knob +(`HYDRAGNN_NUM_WORKERS`) but hard-codes `persistent_workers = False`. This +patch adds a companion opt-in env var: + +- `HYDRAGNN_PERSISTENT_WORKERS=1` and `num_workers > 0` + → `persistent_workers=True` is passed to all DataLoaders +- Otherwise the default upstream behaviour (`persistent_workers=False`) is + unchanged. + +**Why** + +When `num_workers > 0` and `persistent_workers=False`, PyTorch forks the +worker processes at the start of every epoch and re-imports the dataset +module + re-mmaps ADIOS2 BP files. This shows up in our profiler traces as a +bimodal iteration-time distribution — most iterations at the steady-state +s/it, plus a long tail clustered around epoch boundaries and after +validation hops. Setting `persistent_workers=True` keeps the workers alive +across epochs and is PyTorch's documented cure for that pattern. + +**Status** + +Ad-hoc local patch for the `ai4science-studio` perf-analysis recipe. Once +its impact is measured, the right follow-up is a small upstream PR to +`ORNL/HydraGNN` adding the `HYDRAGNN_PERSISTENT_WORKERS` env knob with the +same opt-in semantics. diff --git a/material_science/models/HydraGNN/examples/patches/load_data.persistent_workers.patch b/material_science/models/HydraGNN/examples/patches/load_data.persistent_workers.patch new file mode 100644 index 0000000..21be2bd --- /dev/null +++ b/material_science/models/HydraGNN/examples/patches/load_data.persistent_workers.patch @@ -0,0 +1,15 @@ +diff --git a/hydragnn/preprocess/load_data.py b/hydragnn/preprocess/load_data.py +index f595e2b..cabff01 100644 +--- a/hydragnn/preprocess/load_data.py ++++ b/hydragnn/preprocess/load_data.py +@@ -286,6 +286,10 @@ def create_dataloaders( + num_workers = 0 + if os.getenv("HYDRAGNN_NUM_WORKERS") is not None: + num_workers = int(os.environ["HYDRAGNN_NUM_WORKERS"]) ++ # ai4science-studio perf-analysis patch: opt-in persistent workers. ++ # Default behaviour unchanged (False) unless HYDRAGNN_PERSISTENT_WORKERS=1. ++ if os.getenv("HYDRAGNN_PERSISTENT_WORKERS") is not None and num_workers > 0: ++ persistent_workers = bool(int(os.environ["HYDRAGNN_PERSISTENT_WORKERS"])) + + use_custom_dataloader = 0 + if os.getenv("HYDRAGNN_CUSTOM_DATALOADER") is not None: diff --git a/material_science/models/HydraGNN/examples/sbatch_train_amd.sh b/material_science/models/HydraGNN/examples/sbatch_train_amd.sh index 20dfe2c..5652a76 100755 --- a/material_science/models/HydraGNN/examples/sbatch_train_amd.sh +++ b/material_science/models/HydraGNN/examples/sbatch_train_amd.sh @@ -101,7 +101,7 @@ done # --------------------------------------------------------------------------- # Clone HydraGNN source if not present # --------------------------------------------------------------------------- -HG_HYDRAGNN_SHA="${HG_HYDRAGNN_SHA:-6c45f1682783e66dc89e9e23009f61716186432b}" +HG_HYDRAGNN_SHA="${HG_HYDRAGNN_SHA:-2fb0bd0157e3c85a74f9841887155095bd163303}" if [[ ! -d "${HG_REPO_DIR}/examples/multidataset_hpo_sc26" ]]; then echo "--- Cloning HydraGNN (pinned SHA: ${HG_HYDRAGNN_SHA}) ---" git clone https://github.com/ORNL/HydraGNN.git "$HG_REPO_DIR" diff --git a/material_science/models/HydraGNN/examples/sbatch_train_perf_amd.sh b/material_science/models/HydraGNN/examples/sbatch_train_perf_amd.sh index 9cddff2..123a19d 100755 --- a/material_science/models/HydraGNN/examples/sbatch_train_perf_amd.sh +++ b/material_science/models/HydraGNN/examples/sbatch_train_perf_amd.sh @@ -107,7 +107,7 @@ done # --------------------------------------------------------------------------- # Clone HydraGNN source if not present (mirrors sbatch_train_amd.sh) # --------------------------------------------------------------------------- -HG_HYDRAGNN_SHA="${HG_HYDRAGNN_SHA:-6c45f1682783e66dc89e9e23009f61716186432b}" +HG_HYDRAGNN_SHA="${HG_HYDRAGNN_SHA:-2fb0bd0157e3c85a74f9841887155095bd163303}" if [[ ! -d "${HG_REPO_DIR}/examples/multidataset_hpo_sc26" ]]; then echo "--- Cloning HydraGNN (pinned SHA: ${HG_HYDRAGNN_SHA}) ---" git clone https://github.com/ORNL/HydraGNN.git "$HG_REPO_DIR" @@ -212,6 +212,9 @@ cat > "$RANK_SCRIPT" << 'RANKEOF' set -euo pipefail source /opt/venv/bin/activate +# The overlay's /opt/hydragnn-pkgs is rebuilt from HG_HYDRAGNN_SHA + patches/ +# (see build_overlay_amd.sh), so we import directly from there. Keep the +# bind-mounted clone available for `examples/multidataset_hpo_sc26/*.py`. export PYTHONPATH="/opt/hydragnn-pkgs:${PYTHONPATH:-}" export LD_LIBRARY_PATH="/opt/hydragnn-pkgs/adios2:/opt/ompi/lib:${LD_LIBRARY_PATH:-}" @@ -234,6 +237,11 @@ export HYDRAGNN_TRACE_LEVEL="${HYDRAGNN_TRACE_LEVEL:-1}" export HSA_NO_SCRATCH_RECLAIM=1 cd "$HG_OUTPUT_DIR" +# NOTE: do NOT `import hydragnn` from a separate rank-0-only python here. +# That import triggers `from mpi4py import MPI` which calls MPI_Init() in +# rank 0 only and leaves ranks 1..N-1 hanging in their later MPI_Init() +# because the PMIX rendezvous requires all ranks concurrently. We log the +# imported path inside the main python's first stanza instead. export HYDRAGNN_AVG_NUM_NEIGHBORS="${HYDRAGNN_AVG_NUM_NEIGHBORS:-13.735293601560318}" @@ -245,6 +253,16 @@ export HYDRAGNN_AVG_NUM_NEIGHBORS="${HYDRAGNN_AVG_NUM_NEIGHBORS:-13.735293601560 exec python -u -c " import sys, os, runpy +# Rank 0 logs which HydraGNN copy is being imported and from what SHA. +# Done inside the main python (after MPI ranks align) — see the comment +# above about why this cannot live in a separate rank-0-only python. +if os.environ.get('SLURM_PROCID', '0') == '0': + import inspect, hydragnn, subprocess + pkg_dir = os.path.dirname(hydragnn.__file__) + src = inspect.getsource(__import__('hydragnn.preprocess.load_data', fromlist=['_'])) + print('[hydragnn] imported from:', pkg_dir, flush=True) + print('[hydragnn] persistent_workers patch present:', 'HYDRAGNN_PERSISTENT_WORKERS' in src, flush=True) + # Rank 0 keeps the Profile block; others zero it out so HydraGNN's Profiler # falls through to the null context. Block lives under NeuralNetwork because # that is the dict train_validate_test() receives. @@ -381,7 +399,10 @@ srun --mpi=pmix \ --env HYDRAGNN_TRACE_LEVEL="${HYDRAGNN_TRACE_LEVEL:-1}" \ --env HG_EXAMPLE_DIR="$EXAMPLE_DIR" \ --env HG_OUTPUT_DIR="$HG_OUTPUT_DIR" \ + --env HG_REPO_DIR="$HG_REPO_DIR" \ --env HG_CONFIG_OVERRIDE="$HG_CONFIG_OVERRIDE" \ + --env HYDRAGNN_NUM_WORKERS="${HYDRAGNN_NUM_WORKERS:-}" \ + --env HYDRAGNN_PERSISTENT_WORKERS="${HYDRAGNN_PERSISTENT_WORKERS:-}" \ --env PROFILE_RANK0_ONLY=1 \ "${RCCL_MULTINODE_ENVS[@]}" \ "$HG_SIF" \ diff --git a/material_science/models/HydraGNN/recipes/perf-analysis/omnistat-lux.config.template b/material_science/models/HydraGNN/recipes/perf-analysis/omnistat-lux.config.template index c06b7ce..ab44ea3 100644 --- a/material_science/models/HydraGNN/recipes/perf-analysis/omnistat-lux.config.template +++ b/material_science/models/HydraGNN/recipes/perf-analysis/omnistat-lux.config.template @@ -16,7 +16,7 @@ enable_rocm_smi = False enable_amd_smi = True enable_ras_ecc = True enable_rms = True -enable_rocprofiler = False +enable_rocprofiler = True enable_network = True enable_vendor_counters = True enable_xgmi = True @@ -46,15 +46,37 @@ step_detection_file = @JOB_DIR@/omni_rmsjobinfo_step [omnistat.collectors.host] enable_proc_io_stats = True -# rocprofiler section kept as a reference for the verifier probe — disabled -# above. Enable enable_rocprofiler=True + select a profile here for HBM/FLOPS -# counters on a 1-node remedy probe only. +# Rocprofiler — HBM bandwidth + FP64 VALU/MFMA counters in one block. +# Counter group selection follows the user-validated pattern (one profile, +# one sub-section, all counters listed together; rocprofiler-sdk +# time-multiplexes hardware counter sets internally as needed). +# +# All eight counters below are defined for gfx950 (MI355X) in +# /opt/rocm/share/rocprofiler-sdk/counter_defs.yaml; verified before pinning. +# +# Useful derivations from these raw counters (compute in the verifier or +# downstream; rocprofiler-sdk also exposes derived counters +# MfmaFlopsF64 and TotalFlopsF64 if a downstream consumer prefers those): +# peak_fp64_GFLOP/s = (SQ_INSTS_VALU_MFMA_MOPS_F64 * 512 +# + (SQ_INSTS_VALU_FMA_F64*2 +# + SQ_INSTS_VALU_ADD_F64 +# + SQ_INSTS_VALU_MUL_F64) * 64 +# ) / wall_seconds / 1e9 +# HBM_GB/s = (FETCH_SIZE + WRITE_SIZE) / wall_seconds / 1e9 [omnistat.collectors.rocprofiler] -profile = hbm +profile = hbm_flops_f64 -[omnistat.collectors.rocprofiler.hbm] +[omnistat.collectors.rocprofiler.hbm_flops_f64] sampling_mode = constant -counters = ["FETCH_SIZE", "WRITE_SIZE"] +counters = [ + "FETCH_SIZE", + "WRITE_SIZE", + "SQ_INSTS_VALU_ADD_F64", + "SQ_INSTS_VALU_MUL_F64", + "SQ_INSTS_VALU_TRANS_F64", + "SQ_INSTS_VALU_FMA_F64", + "SQ_INSTS_VALU_MFMA_MOPS_F64" +] [omnistat.query] prometheus_url = http://localhost:8428 From 502cace08210a46c168c2dcc2f6164a67aece628 Mon Sep 17 00:00:00 2001 From: Ashwin Aji Date: Thu, 21 May 2026 18:54:03 +0000 Subject: [PATCH 03/31] HydraGNN perf-analysis: enable rocprofiler counters by default + record gfx950 limits MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The omnistat configs used by jobs 6762 and 6840 had `enable_rocprofiler = False` because the staged copy under `/shared/aaji/models/HydraGNN/perf-runs/` had drifted from the in-repo template, silently dropping all GPU hardware counters from both analyses. This change closes that gap and documents what we learned getting the counter pipeline working on Lux MI355X. Script changes (sbatch_train_perf_amd.sh): * Default OMNISTAT_TEMPLATE to the in-repo `recipes/perf-analysis/*.template` instead of a /shared staged path; the staged copy was the source of the silent-drop bug. * Parse `enable_rocprofiler` and `profile` from the chosen template at submit time and surface them in the run banner. Print a WARN if rocprofiler is off, so a future submitter catches the silent-counters case before the run finishes. Template changes (omnistat-lux.config.template): * Reduce to a single working `constant`-mode profile with 5 counters (FETCH_SIZE + 4 fp64 SQ_INSTS_VALU_*_F64). This is the empirically-validated set that fits gfx950's per-block hardware register budget AND avoids the omnistat periodic-mode bug surfaced during this work. * Add a long comment block explaining the constraints — what fails, why, and the workaround — so a future maintainer doesn't recompose the broken 7-counter or two-set version. SKILL changes (.cursor/skills/ai4science-studio/SKILL.md), four new lessons under "Performance-analysis subagent workflow": * §8: never read configs from a /shared staged copy — always from the in-repo template (the root-cause of the silent-drop above). * §9: omnistat rocprofiler-sdk extension is a separate C++ build (not a pip extra). `pip install -e '.[query]'` does NOT trigger the build even with BUILD_ROCPROFILER_SDK_EXTENSION=1; use `python setup.py build_ext --inplace` on a compute node where ROCm headers are present. * §10: gfx950 rocprofiler counter limits — FETCH_SIZE + WRITE_SIZE in one profile fails (error 38, both in TCC block), 4 x SQ_INSTS_VALU_*_F64 work in one profile. * §11: omnistat sampling_mode = periodic in 1.12.0+git.65ea9ac only fires the FIRST counter set — half the data silently lost. Use `constant` mode with one set that fits in the hardware limits. End-to-end verification: 1-node, 10-batch HydraGNN run (job 6858) with the new config collected non-zero FETCH_SIZE and SQ_INSTS_VALU_MFMA_MOPS_F64 samples on all 8 cards. Raw samples + per-card summary at /shared/aaji/models/HydraGNN/perf-runs/6858/omnistat/hardware_counters/. The rocprofiler MFMA_F64 rate triangulates with TraceLens's analytical roofline (TFLOP/s during burst windows) and Omnistat's 1-Hz utilization duty cycle — see the addendum to perf-runs/6840/combined_report.md. Co-authored-by: Cursor --- .cursor/skills/ai4science-studio/SKILL.md | 59 +++++++++++++++++++ .../examples/sbatch_train_perf_amd.sh | 16 ++++- .../omnistat-lux.config.template | 34 ++++++++--- 3 files changed, 99 insertions(+), 10 deletions(-) diff --git a/.cursor/skills/ai4science-studio/SKILL.md b/.cursor/skills/ai4science-studio/SKILL.md index 4c340e6..a5e6f54 100755 --- a/.cursor/skills/ai4science-studio/SKILL.md +++ b/.cursor/skills/ai4science-studio/SKILL.md @@ -597,3 +597,62 @@ This catches both tool bugs and analyst hallucinations. Refuted claims stay in ` ### 7. Datetimes from omnistat-inspect are UTC `job_info.json` returns `start_time`/`end_time` as ISO-8601 strings without a `Z` or offset — they are **UTC**. Verifier subagents that build PromQL `start`/`end` from these MUST use `calendar.timegm(time.strptime(s, "%Y-%m-%dT%H:%M:%S"))`, **not** `time.mktime` (which interprets local time and produces wrong epoch on most clusters). Wrong epoch → empty PromQL result series → false "no data" finding. + +### 8. Never read configs from a `/shared` "staged copy" — always from the in-repo template + +We hit this once: the omnistat config under `/shared/aaji/models/HydraGNN/perf-runs/omnistat-lux.config` was a stale copy of the in-repo `omnistat-lux.config.template`, with `enable_rocprofiler = False` left over from an early v1 decision. The in-repo template had since been updated to `enable_rocprofiler = True` with the `hbm_flops_f64` profile, but the staged copy was never refreshed, so two perf-analysis runs (jobs 6762 and 6840) silently collected zero hardware counters (no HBM bandwidth, no fp64 VALU/MFMA FLOPs). The TraceLens analytical roofline still ran fine, masking the gap until someone asked "where are the HBM counters?" + +Two compounding mistakes: +1. The sbatch script defaulted `OMNISTAT_TEMPLATE` to the `/shared` staged path, not the in-repo path. +2. There was no run-time signal that rocprofiler was off — the run banner printed `Omnistat tmpl: ` but not the rocprofiler state of that template. + +**Fix pattern for any model recipe that reads a config template**: +- Default the path to the in-repo source of truth: `${SCRIPT_DIR}/../recipes//.template`. Only override via env var for ad-hoc probes. +- After resolving the template, parse the critical knobs and **print them in the run banner**, with a `WARN` line if any expected counter group is disabled. Example: +```bash +_ROCPROF_STATE=$(awk -F'[= \t]+' '/^enable_rocprofiler/ {print $2; exit}' "$OMNISTAT_TEMPLATE") +echo " Rocprofiler : enable_rocprofiler=${_ROCPROF_STATE:-?}" +[[ "$_ROCPROF_STATE" != "True" ]] && echo " WARN: rocprofiler counters DISABLED — HBM/FLOP counters NOT collected" >&2 +``` +- Don't `exit` on it: a config-only probe that intentionally turns rocprofiler off is legitimate. Just make it impossible to be silently wrong. + +**Detection rule for any future template**: if a config file lives **both** in `/...//*.template` and on a shared filesystem (`/shared`, `/lustre`, etc.), assume the shared copy is stale until proven otherwise. Either delete the staged copy or refresh it from the repo before submitting. + +### 9. Omnistat rocprofiler-sdk extension is a separate C++ build (not a pip extra) + +`pip install -e '.[query]'` does NOT build the rocprofiler-sdk Python binding even with `BUILD_ROCPROFILER_SDK_EXTENSION=1` exported, because pip skips the build step on a re-install of an editable package. Symptom: `enable_rocprofiler = True` in the config + `omnistat-monitor` exporter dies silently with `Missing ROCProfiler-SDK extension: build with installation required` and `ModuleNotFoundError: No module named 'rocprofiler_sdk'`. + +The doc-recommended path for editable installs (`docs/installation/extensions.md` upstream) is: +```bash +cd /path/to/omnistat-src +BUILD_ROCPROFILER_SDK_EXTENSION=1 python setup.py build_ext --inplace +``` +This must run on a **compute node** (where `/opt/rocm/include/rocprofiler-sdk/` headers and `/opt/rocm/share/rocprofiler-sdk/counter_defs.yaml` are present — login nodes often have neither). Build deps: `cmake-build-extension`, `nanobind`, `cmake>=3.15`, ninja, a C++20 compiler. Output is one `.so` file named `omnistat/rocprofiler_sdk_extension.cpython-3*-x86_64-linux-gnu.so`. Verify with `python3 -c "from omnistat.rocprofiler_sdk_extension import get_samplers, initialize"`. + +Imports note: omnistat's collector imports `from omnistat.rocprofiler_sdk_extension import ...`, NOT `import rocprofiler_sdk`. A naive sanity check that imports the upstream `rocprofiler_sdk` Python package will always fail (that package is unrelated, ships with ROCm, and is not on the omnistat venv path). + +### 10. gfx950 (MI355X) rocprofiler counter register limits — multiple counters in same hardware block fail + +ROCm rocprofiler-sdk on gfx950 returns +``` +create profile failed with error code 38: Request exceeds the capabilities of the hardware to collect +``` +when a single profile asks for more counters than fit in the per-block register file. Per-block limits are not documented uniformly; we determined empirically on `lux-mi355x-b1` (ROCm 7.2.2): + +- **TCC (L2 cache) block**: **`FETCH_SIZE` AND `WRITE_SIZE` together fail.** Either one alone works (with up to 4 SQ_INSTS_VALU_*_F64 counters in the same profile). +- **SQ (scalar/vector unit) block**: ≥4 SQ_INSTS_VALU_*_F64 counters in one profile work. +- Workaround: use `sampling_mode = periodic` and `counters = [[set0], [set1]]` to time-multiplex — but see lesson 11. + +Always probe interactively first (1 node `salloc`, run `omnistat-monitor --configfile ` foreground for 8 s, grep stderr for "create profile failed") before submitting an sbatch with new counters; the omnistat user-mode launcher swallows the error from the spawned exporter and the only symptom from the outside is "Missing exporter on ". + +### 11. Omnistat `sampling_mode = periodic` only fires the first set in this build + +In `omnistat 1.12.0+git.65ea9ac` (PR #271 + main merged) on `lux-mi355x-b1` with rocprofiler-sdk freshly built, `sampling_mode = periodic` with two counter sets only fires the FIRST set. Every subsequent set returns identically zero across all cards and timestamps. Half the intended counter coverage is silently lost. + +We surfaced this with the `[[FETCH+4×SQ_F64], [WRITE+4×SQ_F64]]` config: FETCH and MFMA_MOPS samples were healthy (~50% non-zero, ~7% non-zero respectively), but WRITE and FMA/ADD/MUL were uniformly zero. Workaround: use `sampling_mode = constant` and a single counter set that fits in the per-block hardware limits — for fp64 + HBM-read on gfx950 that's: +```ini +[omnistat.collectors.rocprofiler.fp64_and_hbm_read] +sampling_mode = constant +counters = ["FETCH_SIZE", "SQ_INSTS_VALU_ADD_F64", "SQ_INSTS_VALU_MUL_F64", "SQ_INSTS_VALU_FMA_F64", "SQ_INSTS_VALU_MFMA_MOPS_F64"] +``` +Trade-off: lose `WRITE_SIZE` (HBM-write bandwidth) until the periodic-mode bug is fixed upstream. For most ML training workloads, FETCH_SIZE is the dominant traffic anyway. diff --git a/material_science/models/HydraGNN/examples/sbatch_train_perf_amd.sh b/material_science/models/HydraGNN/examples/sbatch_train_perf_amd.sh index 123a19d..f33d4a9 100755 --- a/material_science/models/HydraGNN/examples/sbatch_train_perf_amd.sh +++ b/material_science/models/HydraGNN/examples/sbatch_train_perf_amd.sh @@ -72,7 +72,12 @@ HYDRAGNN_MAX_NUM_BATCH="${HYDRAGNN_MAX_NUM_BATCH:-30}" PROFILE_TARGET_EPOCH="${PROFILE_TARGET_EPOCH:-1}" OMNISTAT_VENV="${OMNISTAT_VENV:-/shared/aaji/tools/omnistat-pr271}" -OMNISTAT_TEMPLATE="${OMNISTAT_TEMPLATE:-${HG_BASE}/perf-runs/omnistat-lux.config}" +# Read the omnistat config from the in-repo recipe (single source of truth). +# Override with OMNISTAT_TEMPLATE=/path/to/file if you need a custom probe config. +# A staged copy under ${HG_BASE}/perf-runs/ is intentionally NOT used as the +# default — it has drifted in the past from the in-repo template, silently +# disabling rocprofiler counters. See ai4science-studio SKILL.md for context. +OMNISTAT_TEMPLATE="${OMNISTAT_TEMPLATE:-${SCRIPT_DIR}/../recipes/perf-analysis/omnistat-lux.config.template}" OMNISTAT_USERMODE_INTERVAL="${OMNISTAT_USERMODE_INTERVAL:-1}" NODES="${SLURM_JOB_NUM_NODES:-2}" @@ -89,6 +94,11 @@ for var in HG_SIF HG_OVERLAY OMNISTAT_TEMPLATE; do fi done +# Surface the rocprofiler state of the chosen template so a silent +# "counters off" never happens again. Computed here, printed in the banner. +_ROCPROF_STATE=$(awk -F'[= \t]+' '/^enable_rocprofiler/ {print $2; exit}' "$OMNISTAT_TEMPLATE") +_ROCPROF_PROFILE=$(awk -F'[= \t]+' '/^profile/ {print $2; exit}' "$OMNISTAT_TEMPLATE") + if [[ ! -x "${OMNISTAT_VENV}/bin/omnistat-usermode" ]]; then echo "ERROR: omnistat-usermode not found at ${OMNISTAT_VENV}/bin/omnistat-usermode" >&2 echo " Run the launcher subagent first to install (see recipes/perf-analysis/agents/launcher.md)" >&2 @@ -179,6 +189,10 @@ echo " Profile cfg : $HG_CONFIG_OVERRIDE" echo " Omnistat venv: $OMNISTAT_VENV" echo " Omnistat tmpl: $OMNISTAT_TEMPLATE" echo " Omnistat cfg : $OMNISTAT_CONFIG (rendered)" +echo " Rocprofiler : enable_rocprofiler=${_ROCPROF_STATE:-?} profile=${_ROCPROF_PROFILE:-?}" +if [[ "$_ROCPROF_STATE" != "True" ]]; then + echo " WARN : rocprofiler counters DISABLED — HBM/FLOP counters will NOT be collected" >&2 +fi echo " Node(s) : ${SLURM_NODELIST:-$(hostname)}" echo " Date : $(date)" echo "" diff --git a/material_science/models/HydraGNN/recipes/perf-analysis/omnistat-lux.config.template b/material_science/models/HydraGNN/recipes/perf-analysis/omnistat-lux.config.template index ab44ea3..ca58c0a 100644 --- a/material_science/models/HydraGNN/recipes/perf-analysis/omnistat-lux.config.template +++ b/material_science/models/HydraGNN/recipes/perf-analysis/omnistat-lux.config.template @@ -67,16 +67,32 @@ enable_proc_io_stats = True profile = hbm_flops_f64 [omnistat.collectors.rocprofiler.hbm_flops_f64] +# MI355X (gfx950) rocprofiler counter limits — see ai4science-studio SKILL §10. +# +# Constraints surfaced on lux-mi355x-b1 / ROCm 7.2.2: +# * FETCH_SIZE + WRITE_SIZE in ONE profile fails: "error code 38: Request +# exceeds the capabilities of the hardware to collect" (both in TCC block). +# * FETCH_SIZE OR WRITE_SIZE + 4 x SQ_INSTS_VALU_*_F64 in one profile works. +# * sampling_mode = periodic in omnistat 1.12.0+git.65ea9ac (PR #271 + main) +# ONLY fires the first counter set — half the data silently lost. (See +# SKILL §11.) +# +# Workaround: use `constant` mode with a single counter set that fits the +# hardware limits and the omnistat bug. We capture HBM read traffic only +# (FETCH_SIZE) plus the 4 fp64 instruction counters needed to derive the +# fp64 FLOP rate using the formula: +# fp64_GFLOP/s = (MFMA_MOPS_F64 * 512 + (FMA_F64*2 + ADD_F64 + MUL_F64) * 64) +# / wall_seconds / 1e9 +# Drop SQ_INSTS_VALU_TRANS_F64 (not relevant for GEMM-heavy ML training). +# +# WRITE_SIZE (HBM-write bandwidth) is NOT collected in this profile until +# upstream omnistat fixes the periodic-mode bug. For most ML training +# workloads, FETCH_SIZE is the dominant component of HBM traffic anyway. +# +# All counters on one line: configparser multi-line values require every +# continuation line (including the closing `]`) to have leading whitespace. sampling_mode = constant -counters = [ - "FETCH_SIZE", - "WRITE_SIZE", - "SQ_INSTS_VALU_ADD_F64", - "SQ_INSTS_VALU_MUL_F64", - "SQ_INSTS_VALU_TRANS_F64", - "SQ_INSTS_VALU_FMA_F64", - "SQ_INSTS_VALU_MFMA_MOPS_F64" -] +counters = ["FETCH_SIZE", "SQ_INSTS_VALU_ADD_F64", "SQ_INSTS_VALU_MUL_F64", "SQ_INSTS_VALU_FMA_F64", "SQ_INSTS_VALU_MFMA_MOPS_F64"] [omnistat.query] prometheus_url = http://localhost:8428 From 7128a4dbd8211780419357c585cbdd99e055853f Mon Sep 17 00:00:00 2001 From: Ashwin Aji Date: Thu, 21 May 2026 23:22:35 +0000 Subject: [PATCH 04/31] HydraGNN perf-analysis: opt-in omnistat kernel-dispatch tracing MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Adds an OMNISTAT_KERNEL_TRACE=1 env-gate to sbatch_train_perf_amd.sh that loads omnistat's libomnistat_trace.so via ROCP_TOOL_LIBRARIES, flips enable_kernel_trace=True in the rendered config, and surfaces the state in the run banner. Hard-fails (exit 2) if the env-gate is on but the tool library is missing — same defensive pattern as the rocprofiler counter-state guard from §8/SKILL, which keeps us from re-creating the "silent zero-counters" failure mode for the kernel-trace surface. The kernel-trace library itself is built separately on a compute node inside the workload SIF; launcher.md picks up the build step (idempotent, gated on OMNISTAT_KERNEL_TRACE=1). SKILL §12 documents the three rocprofiler-sdk surfaces (device counting, kernel tracing, rocprofv3) and when to pick which. Default behavior is unchanged: enable_kernel_trace=False in the template, no apptainer-exec changes when the gate is off. Co-authored-by: Cursor --- .cursor/skills/ai4science-studio/SKILL.md | 31 +++++++++++ .../examples/sbatch_train_perf_amd.sh | 52 +++++++++++++++++++ .../recipes/perf-analysis/agents/launcher.md | 24 +++++++++ .../omnistat-lux.config.template | 14 +++++ 4 files changed, 121 insertions(+) diff --git a/.cursor/skills/ai4science-studio/SKILL.md b/.cursor/skills/ai4science-studio/SKILL.md index a5e6f54..7abae6b 100755 --- a/.cursor/skills/ai4science-studio/SKILL.md +++ b/.cursor/skills/ai4science-studio/SKILL.md @@ -656,3 +656,34 @@ sampling_mode = constant counters = ["FETCH_SIZE", "SQ_INSTS_VALU_ADD_F64", "SQ_INSTS_VALU_MUL_F64", "SQ_INSTS_VALU_FMA_F64", "SQ_INSTS_VALU_MFMA_MOPS_F64"] ``` Trade-off: lose `WRITE_SIZE` (HBM-write bandwidth) until the periodic-mode bug is fixed upstream. For most ML training workloads, FETCH_SIZE is the dominant traffic anyway. + +### 12. Three rocprofiler-sdk surfaces — pick the right one + +Omnistat exposes **two** independent rocprofiler-sdk integrations, and there's a **third** path that bypasses omnistat entirely. They sample disjoint things and can mostly co-exist; pick by the question you're answering. + +| Surface | What you get | Cardinality / cost | Build artifact | Switch in our stack | +|---|---|---|---|---| +| **Device counting** (omnistat collector, `enable_rocprofiler=True`) | Whole-card hardware counters: FETCH_SIZE, fp64 VALU/MFMA, etc. Sampled at 1 Hz, scraped by the omnistat exporter. | One time-series per (card, counter). Constant-mode is essentially free. | `omnistat/rocprofiler_sdk_extension.cpython-*.so` (Python binding, built in-place from `rocprofiler-sdk/CMakeLists.txt` with `BUILD_KERNEL_TRACE_LIB=OFF`) | `enable_rocprofiler = True` in `omnistat-lux.config.template` (default). | +| **Kernel tracing** (omnistat collector, `enable_kernel_trace=True`) | Per-kernel-name dispatch count and total duration on every card, every node. Aggregated in 1 s bins, POSTed via HTTP to the omnistat exporter. | One time-series per (card, kernel-name) — can be 1k+ unique kernel names per training step. Light overhead per dispatch (microseconds), but cardinality blows up VictoriaMetrics if left on for hours. | `build-trace/libomnistat_trace.so` (standalone C++ tool library, built from same CMake with `-DBUILD_KERNEL_TRACE_LIB=ON`) | `OMNISTAT_KERNEL_TRACE=1` env-gate in `sbatch_train_perf_amd.sh` — flips `enable_kernel_trace=True` in the rendered config and exports `ROCP_TOOL_LIBRARIES=` into the apptainer exec. | +| **`rocprofv3` standalone trace** (no omnistat) | Full kernel + memcpy + HSA API trace on a single rank. JSON / Perfetto / CSV output you read with TraceLens or chrome://tracing. | One trace file per rank, multi-MB per second of training. Heaviest of the three. | None — uses `/opt/rocm/bin/rocprofv3` directly. | Wrap a single rank's command in `rocprofv3 --kernel-trace -o ./rank0_trace -- python ...`. We don't run this in our standard sbatch yet. | + +**When to use which:** + +- **Device counting only** → multi-node, multi-hour runs. You want HBM bandwidth and fp64 FLOP rate per card across the whole job, with telemetry that survives a 2-hour wallclock without filling the DB. +- **Kernel tracing on top of device counting** → short probe runs (≤ 5 minutes, ≤ 1 node) when you need to know *which kernel* dominated time without rank-0-only kineto traces. Combine with device counting to correlate per-kernel duration against whole-card FLOP rate. +- **`rocprofv3` instead of omnistat** → deepest single-rank kernel detail (per-dispatch timing, register usage, kernel arguments). Use when omnistat's aggregate-by-name view loses the signal you need (e.g. tail latency on one specific dispatch). Don't run it across all ranks; the file size will eat the NFS share. + +**Cross-runtime gotcha:** `enable_kernel_trace=True` without the tool library loaded does *nothing* — the omnistat collector just sits with an empty endpoint. The `OMNISTAT_KERNEL_TRACE=1` switch in `sbatch_train_perf_amd.sh` enforces this with a hard `exit 2` if `libomnistat_trace.so` isn't found at the configured path; do **not** loosen that check, otherwise you re-create the exact "silent zero counters" failure mode from §8. + +**Build location:** Both omnistat artifacts (the Python extension and the kernel-trace tool library) must be built on a **compute node** inside the same SIF the workload runs in — login nodes don't have apptainer or the ROCm headers, and a host-side build will pick up the wrong libcurl / glibc. Standard build: + +```bash +salloc -p lux -A vultr_lux -N 1 --time=00:15:00 --gpus-per-node=1 \ + apptainer exec --rocm \ + /shared/aaji/images/pytorch_rocm7.2.2_ubuntu24.04_py3.12_pytorch_release_2.10.0.sif \ + bash -c 'cd /shared/aaji/tools/omnistat-src && \ + cmake -S rocprofiler-sdk/ -B build-trace/ -DBUILD_KERNEL_TRACE_LIB=ON && \ + cmake --build build-trace/ -j 8' +``` + +Verify with `ls -la /shared/aaji/tools/omnistat-src/build-trace/libomnistat_trace.so` (~MB-class file, not the tiny 4 kB stub you get when CMake fails halfway). diff --git a/material_science/models/HydraGNN/examples/sbatch_train_perf_amd.sh b/material_science/models/HydraGNN/examples/sbatch_train_perf_amd.sh index f33d4a9..d7c948f 100755 --- a/material_science/models/HydraGNN/examples/sbatch_train_perf_amd.sh +++ b/material_science/models/HydraGNN/examples/sbatch_train_perf_amd.sh @@ -80,6 +80,16 @@ OMNISTAT_VENV="${OMNISTAT_VENV:-/shared/aaji/tools/omnistat-pr271}" OMNISTAT_TEMPLATE="${OMNISTAT_TEMPLATE:-${SCRIPT_DIR}/../recipes/perf-analysis/omnistat-lux.config.template}" OMNISTAT_USERMODE_INTERVAL="${OMNISTAT_USERMODE_INTERVAL:-1}" +# Kernel-dispatch tracing (per-kernel name + duration histogram from every +# rank). OFF by default — opt in with OMNISTAT_KERNEL_TRACE=1. When enabled +# we auto-flip enable_kernel_trace=True in the rendered config and load +# libomnistat_trace.so via ROCP_TOOL_LIBRARIES inside the container. +# See material_science/models/HydraGNN/recipes/perf-analysis/agents/launcher.md +# for how to build libomnistat_trace.so on a compute node. +OMNISTAT_KERNEL_TRACE="${OMNISTAT_KERNEL_TRACE:-0}" +OMNISTAT_TRACE_LIB="${OMNISTAT_TRACE_LIB:-/shared/aaji/tools/omnistat-src/build-trace/libomnistat_trace.so}" +OMNISTAT_TRACE_ENDPOINT_PORT="${OMNISTAT_TRACE_ENDPOINT_PORT:-8001}" + NODES="${SLURM_JOB_NUM_NODES:-2}" GPUS_PER_NODE=8 TOTAL_RANKS=$((NODES * GPUS_PER_NODE)) @@ -147,6 +157,27 @@ mkdir -p "$HG_OUTPUT_DIR" OMNISTAT_CONFIG="${HG_OUTPUT_DIR}/omnistat.config" sed -e "s|@JOB_DIR@|${HG_OUTPUT_DIR}|g" "$OMNISTAT_TEMPLATE" > "$OMNISTAT_CONFIG" +# Opt-in kernel tracing: flip enable_kernel_trace=True in the rendered config +# and verify the rocprofiler-sdk tool library exists. Bail out loudly if +# the user asked for it but the .so isn't present — silent fall-through is +# exactly the "rocprofiler counters disabled" failure mode we already burned +# a day on (see SKILL.md §11). +if [[ "$OMNISTAT_KERNEL_TRACE" == "1" ]]; then + if [[ ! -f "$OMNISTAT_TRACE_LIB" ]]; then + echo "ERROR: OMNISTAT_KERNEL_TRACE=1 but tool library not found:" >&2 + echo " $OMNISTAT_TRACE_LIB" >&2 + echo " Build it on a compute node:" >&2 + echo " salloc -p lux -A vultr_lux -N 1 --time=00:15:00 --gpus-per-node=1 \\" >&2 + echo " apptainer exec --rocm \"\$HG_SIF\" bash -c '" >&2 + echo " cd /shared/aaji/tools/omnistat-src && \\" >&2 + echo " cmake -S rocprofiler-sdk/ -B build-trace/ -DBUILD_KERNEL_TRACE_LIB=ON && \\" >&2 + echo " cmake --build build-trace/ -j 8'" >&2 + exit 2 + fi + sed -i 's/^enable_kernel_trace = False/enable_kernel_trace = True/' "$OMNISTAT_CONFIG" +fi +_KTRACE_STATE=$(awk -F'[= \t]+' '/^enable_kernel_trace/ {print $2; exit}' "$OMNISTAT_CONFIG") + # --------------------------------------------------------------------------- # Generate the per-job profile config (inject "Profile" block) # Done on the submit/host side (login or compute), uses python3 if available. @@ -193,6 +224,11 @@ echo " Rocprofiler : enable_rocprofiler=${_ROCPROF_STATE:-?} profile=${_ROCPRO if [[ "$_ROCPROF_STATE" != "True" ]]; then echo " WARN : rocprofiler counters DISABLED — HBM/FLOP counters will NOT be collected" >&2 fi +echo " KernelTrace : enable_kernel_trace=${_KTRACE_STATE:-?} (opt-in via OMNISTAT_KERNEL_TRACE=1)" +if [[ "$OMNISTAT_KERNEL_TRACE" == "1" ]]; then + echo " TraceLib : $OMNISTAT_TRACE_LIB" + echo " TracePort : $OMNISTAT_TRACE_ENDPOINT_PORT" +fi echo " Node(s) : ${SLURM_NODELIST:-$(hostname)}" echo " Date : $(date)" echo "" @@ -385,6 +421,21 @@ if [[ "$NODES" -gt 1 ]]; then ) fi +# --------------------------------------------------------------------------- +# Optional: kernel-trace tool library bind + env (loaded by rocprofiler-sdk +# when ROCP_TOOL_LIBRARIES points at the .so). The library POSTs dispatch +# records to the omnistat collector on localhost:OMNISTAT_TRACE_ENDPOINT_PORT. +# --------------------------------------------------------------------------- +KTRACE_BIND_ENVS=() +if [[ "$OMNISTAT_KERNEL_TRACE" == "1" ]]; then + KTRACE_BIND_ENVS=( + --bind "${OMNISTAT_TRACE_LIB}:${OMNISTAT_TRACE_LIB}:ro" + --env ROCP_TOOL_LIBRARIES="${OMNISTAT_TRACE_LIB}" + --env OMNISTAT_TRACE_ENDPOINT_PORT="${OMNISTAT_TRACE_ENDPOINT_PORT}" + --env OMNISTAT_TRACE_LOG="${OMNISTAT_TRACE_LOG:-1}" + ) +fi + # --------------------------------------------------------------------------- # Launch distributed training via srun + apptainer # --------------------------------------------------------------------------- @@ -400,6 +451,7 @@ srun --mpi=pmix \ --bind "${HG_REPO_DIR}:${HG_REPO_DIR}:ro" \ --bind "${HG_DATA_DIR}:${HG_DATA_DIR}:ro" \ --bind "${HG_OUTPUT_DIR}:${HG_OUTPUT_DIR}" \ + "${KTRACE_BIND_ENVS[@]}" \ --env PMIX_MCA_gds=hash \ --env PMIX_MCA_psec=native \ "${MPI_MULTINODE_ENVS[@]}" \ diff --git a/material_science/models/HydraGNN/recipes/perf-analysis/agents/launcher.md b/material_science/models/HydraGNN/recipes/perf-analysis/agents/launcher.md index 621c9c1..65b66bb 100644 --- a/material_science/models/HydraGNN/recipes/perf-analysis/agents/launcher.md +++ b/material_science/models/HydraGNN/recipes/perf-analysis/agents/launcher.md @@ -100,6 +100,30 @@ fi VM_VERSION=$(cat $VM_DIR/VERSION) ``` +### 1e. (Optional) Build the omnistat kernel-trace tool library + +Only required if the run will set `OMNISTAT_KERNEL_TRACE=1` to enable per-kernel dispatch tracing. Idempotent — skip if `build-trace/libomnistat_trace.so` already exists. + +```bash +TRACE_LIB=$TOOLS/omnistat-src/build-trace/libomnistat_trace.so +if [[ "${OMNISTAT_KERNEL_TRACE:-0}" == "1" && ! -f "$TRACE_LIB" ]]; then + # Must build inside the same SIF the workload uses — login nodes lack + # apptainer and ROCm headers. C++20 + libcurl + rocprofiler-sdk needed. + salloc -p lux -A vultr_lux -N 1 --time=00:15:00 --gpus-per-node=1 \ + apptainer exec --rocm \ + "/shared/aaji/images/pytorch_rocm7.2.2_ubuntu24.04_py3.12_pytorch_release_2.10.0.sif" \ + bash -c " + set -e + cd $TOOLS/omnistat-src + cmake -S rocprofiler-sdk/ -B build-trace/ -DBUILD_KERNEL_TRACE_LIB=ON + cmake --build build-trace/ -j 8 + " + test -f "$TRACE_LIB" || { echo "ERROR: kernel-trace build did not produce $TRACE_LIB" >&2; exit 2; } +fi +``` + +The sbatch wrapper expects the `.so` at `/shared/aaji/tools/omnistat-src/build-trace/libomnistat_trace.so` by default; override with `OMNISTAT_TRACE_LIB=/path/to/lib.so`. See `ai4science-studio` SKILL §12 for when to enable kernel tracing vs the default device-counting collector. + ### 2. Author the Omnistat user-mode config (once, in repo `perf-runs/`) If `/shared/aaji/models/HydraGNN/perf-runs/omnistat-lux.config` doesn't exist, write it from the template at `/home/aaji/git/ai4science-studio/material_science/models/HydraGNN/recipes/perf-analysis/omnistat-lux.config.template`. The template uses `%(SLURM_JOB_ID)s` placeholders that omnistat-usermode resolves at runtime. diff --git a/material_science/models/HydraGNN/recipes/perf-analysis/omnistat-lux.config.template b/material_science/models/HydraGNN/recipes/perf-analysis/omnistat-lux.config.template index ca58c0a..1eee37e 100644 --- a/material_science/models/HydraGNN/recipes/perf-analysis/omnistat-lux.config.template +++ b/material_science/models/HydraGNN/recipes/perf-analysis/omnistat-lux.config.template @@ -22,6 +22,20 @@ enable_vendor_counters = True enable_xgmi = True enable_cu_occupancy = False +# Kernel-dispatch tracing collector (per-kernel name + duration histogram on +# every node). Disabled by default: it requires the standalone tool library +# `libomnistat_trace.so` to be built (see rocprofiler-sdk/README.md +# "Kernel Tracing Library") and loaded into the workload via +# ROCP_TOOL_LIBRARIES. The sbatch wrapper enables this automatically when +# OMNISTAT_KERNEL_TRACE=1 is set; otherwise it stays False. +# +# IMPORTANT: kernel tracing and the rocprofiler device-counting service +# (enable_rocprofiler above) BOTH talk to rocprofiler-sdk. They can co-exist +# on a single GPU, but if you're chasing per-kernel attribution it's usually +# cleaner to either (a) run kernel-trace alone, or (b) run device counting +# alone — see ai4science-studio SKILL §11 for the tradeoffs. +enable_kernel_trace = False + # Path to local ROCm install (Lux: 7.2.x) rocm_path = /opt/rocm From bb869d6927739b5c8c7df85ced670eb45546774f Mon Sep 17 00:00:00 2001 From: Ashwin Aji Date: Fri, 22 May 2026 02:09:31 +0000 Subject: [PATCH 05/31] HydraGNN: validate kernel-trace e2e + fix empty-string env passthrough MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Validation probe (job 7034, lux-mi355x-b1, 1 node x 8 ranks, 5 batches): - libomnistat_trace.so loaded on every rank via ROCP_TOOL_LIBRARIES - 22.5k records/rank x 8 ranks (~180k dispatches), 0 dropped - 20/20 successful HTTP flushes per rank to localhost:8101/kernel_trace - VictoriaMetrics shows 130 distinct kernel names x 8 cards x 2 metrics - Top-by-duration confirms expected shape (NCCL 4.2s; Tensile GEMMs 1.7s; PyTorch elementwise tail) for a small-batch single-node DDP run Two fixes folded in: 1. Endpoint port alignment: the kernel-trace collector registers /kernel_trace on the SAME Flask app as the other omnistat collectors, so the tool library must POST to [omnistat.collectors] port (8101 in our template), NOT the library's own default 8001. Auto-derive the port from the rendered config to prevent drift. 2. Empty-string env passthrough trap (latent, surfaced by job 7033): HydraGNN's load_data.py and train_validate_test.py both use `os.getenv(K) is not None` then `int(os.environ[K])` to detect optional knobs (HYDRAGNN_NUM_WORKERS, HYDRAGNN_PERSISTENT_WORKERS, HYDRAGNN_MAX_NUM_BATCH). The old `--env KEY="${KEY:-}"` passthrough passes empty strings, which pass the existence check and then crash on int(''). Replaced with a conditional-array pattern that only appends --env flags when the value is non-empty. Applied to both sbatch_train_amd.sh and sbatch_train_perf_amd.sh; documented as SKILL §13 with an explicit "audit every */models/*/examples/*.sh" directive for future agents. SKILL also captures the SIF build recipe (pip-cmake + .deb-with-zstd extraction for libcurl headers; apt-get download is unavailable inside the SIF) so the next agent can build libomnistat_trace.so without re-discovering these gotchas. Co-authored-by: Cursor --- .cursor/skills/ai4science-studio/SKILL.md | 37 +++++++++++++++++++ .../HydraGNN/examples/sbatch_train_amd.sh | 14 ++++++- .../examples/sbatch_train_perf_amd.sh | 28 ++++++++++++-- 3 files changed, 74 insertions(+), 5 deletions(-) diff --git a/.cursor/skills/ai4science-studio/SKILL.md b/.cursor/skills/ai4science-studio/SKILL.md index 7abae6b..5b8648d 100755 --- a/.cursor/skills/ai4science-studio/SKILL.md +++ b/.cursor/skills/ai4science-studio/SKILL.md @@ -687,3 +687,40 @@ salloc -p lux -A vultr_lux -N 1 --time=00:15:00 --gpus-per-node=1 \ ``` Verify with `ls -la /shared/aaji/tools/omnistat-src/build-trace/libomnistat_trace.so` (~MB-class file, not the tiny 4 kB stub you get when CMake fails halfway). + +**Lux SIF build prerequisites (one-time gotchas):** the `pytorch_rocm7.2.2_ubuntu24.04_py3.12_pytorch_release_2.10.0.sif` ships **without** `cmake` and **without** libcurl headers (only the `.so.4` runtime). Both omnistat artifacts need cmake; the kernel-trace library additionally needs ``. Working build recipe inside the SIF: + +1. `python -m venv /tmp/build_venv && /tmp/build_venv/bin/pip install cmake nanobind zstandard` — installs `cmake` 4.x as a wheel, plus `zstandard` for unpacking modern Ubuntu `.deb` files (which use `data.tar.zst`). +2. Fetch `libcurl4-openssl-dev_*.deb` from `archive.ubuntu.com/ubuntu/pool/main/c/curl/`, `ar x` it, decompress `data.tar.zst` with `zstandard.ZstdDecompressor` in Python, untar to a tmp prefix. +3. Pass `-DCURL_INCLUDE_DIR=/usr/include/x86_64-linux-gnu -DCURL_LIBRARY=/usr/lib/x86_64-linux-gnu/libcurl.so.4` to cmake configure. The library finds rocprofiler-sdk via `/opt/rocm/lib/cmake/rocprofiler-sdk/`. + +Verified on `lux-mi355x-b1` job 7030: 533 kB `libomnistat_trace.so`, ELF64, all rocprofiler-sdk symbols resolved, env vars `OMNISTAT_TRACE_{MAX_INTERVAL,BUFFER_SIZE,ENDPOINT_PORT,LOG}` baked in. `apt-get download` does **not** work in the SIF — apt sources aren't configured — so don't waste a node on it. + +**Endpoint port alignment:** the kernel-trace **collector** registers `/kernel_trace` on the same Flask app as the rest of the omnistat exporter, so it listens on `[omnistat.collectors] port` (8101 in our config). The kernel-trace **tool library**'s default `OMNISTAT_TRACE_ENDPOINT_PORT` is 8001, which is wrong for our setup — they must match or every dispatch record gets a connection-refused. The sbatch wrapper auto-derives the port from the rendered config; do not hard-code 8001. + +### 13. Empty-string env passthrough breaks `os.getenv() is not None` checks (HydraGNN-class trap) + +HydraGNN's `train_validate_test.py` and `preprocess/load_data.py` detect optional knobs with the idiom + +```python +if os.getenv("HYDRAGNN_NUM_WORKERS") is not None: + num_workers = int(os.environ["HYDRAGNN_NUM_WORKERS"]) +``` + +This **fails** when the variable is set to the empty string: `getenv() is not None` returns `True`, and `int("")` raises `ValueError`. Our older `--env KEY="${KEY:-}"` passthrough pattern passes empty strings whenever the wrapper's caller didn't set the var, which masquerades as "set but empty" inside the container. + +**Symptom (probe job 7033 on this repo):** all 8 ranks crash on the second line of `create_dataloaders` with `ValueError: invalid literal for int() with base 10: ''` before any kernel ever launches. + +**Fix:** build a separate array of `--env` flags that are only appended when the value is non-empty: + +```bash +OPT_ENVS=() +for k in HYDRAGNN_NUM_WORKERS HYDRAGNN_PERSISTENT_WORKERS HYDRAGNN_MAX_NUM_BATCH; do + v="${!k:-}" + [[ -n "$v" ]] && OPT_ENVS+=( --env "${k}=${v}" ) +done + +srun ... apptainer exec ... "${OPT_ENVS[@]}" "$HG_SIF" bash "$RANK_SCRIPT" +``` + +Applied to both `sbatch_train_amd.sh` and `sbatch_train_perf_amd.sh` for HydraGNN. **General lesson** for any model recipe: never use `--env KEY="${KEY:-}"` for variables that the workload reads with `os.getenv(K) is not None` — either pass a numeric default (`--env KEY="${KEY:-0}"`) or use the conditional-array pattern above. Audit every `*/models/*/examples/*.sh` against this when introducing new optional env knobs. diff --git a/material_science/models/HydraGNN/examples/sbatch_train_amd.sh b/material_science/models/HydraGNN/examples/sbatch_train_amd.sh index 5652a76..b42b200 100755 --- a/material_science/models/HydraGNN/examples/sbatch_train_amd.sh +++ b/material_science/models/HydraGNN/examples/sbatch_train_amd.sh @@ -311,6 +311,18 @@ fi echo "--- Launching training: $TOTAL_RANKS ranks across $NODES nodes ---" echo "" +# Optional env passthrough — HydraGNN's training code uses +# `os.getenv(KEY) is not None` to detect these knobs, so an empty string +# passes the check and then `int("")` blows up. Build a separate array of +# `--env` flags that are only added when the var is set non-empty. +OPT_ENVS=() +for k in HYDRAGNN_MAX_NUM_BATCH HYDRAGNN_NUM_WORKERS HYDRAGNN_PERSISTENT_WORKERS; do + v="${!k:-}" + if [[ -n "$v" ]]; then + OPT_ENVS+=( --env "${k}=${v}" ) + fi +done + srun --mpi=pmix \ apptainer exec \ --rocm \ @@ -328,11 +340,11 @@ srun --mpi=pmix \ --env HG_PRECISION="$HG_PRECISION" \ --env HG_BATCH_SIZE="${HG_BATCH_SIZE:-200}" \ --env HG_NUM_EPOCH="${HG_NUM_EPOCH:-1}" \ - --env HYDRAGNN_MAX_NUM_BATCH="${HYDRAGNN_MAX_NUM_BATCH:-}" \ --env HYDRAGNN_VALTEST="${HYDRAGNN_VALTEST:-0}" \ --env HYDRAGNN_TRACE_LEVEL="${HYDRAGNN_TRACE_LEVEL:-1}" \ --env HG_EXAMPLE_DIR="$EXAMPLE_DIR" \ --env HG_OUTPUT_DIR="$HG_OUTPUT_DIR" \ + "${OPT_ENVS[@]}" \ "${RCCL_MULTINODE_ENVS[@]}" \ "$HG_SIF" \ bash "$RANK_SCRIPT" diff --git a/material_science/models/HydraGNN/examples/sbatch_train_perf_amd.sh b/material_science/models/HydraGNN/examples/sbatch_train_perf_amd.sh index d7c948f..effef88 100755 --- a/material_science/models/HydraGNN/examples/sbatch_train_perf_amd.sh +++ b/material_science/models/HydraGNN/examples/sbatch_train_perf_amd.sh @@ -88,7 +88,13 @@ OMNISTAT_USERMODE_INTERVAL="${OMNISTAT_USERMODE_INTERVAL:-1}" # for how to build libomnistat_trace.so on a compute node. OMNISTAT_KERNEL_TRACE="${OMNISTAT_KERNEL_TRACE:-0}" OMNISTAT_TRACE_LIB="${OMNISTAT_TRACE_LIB:-/shared/aaji/tools/omnistat-src/build-trace/libomnistat_trace.so}" -OMNISTAT_TRACE_ENDPOINT_PORT="${OMNISTAT_TRACE_ENDPOINT_PORT:-8001}" +# IMPORTANT: the kernel-trace collector registers /kernel_trace on the SAME +# Flask app as the other omnistat endpoints, so the tool library must POST to +# the [omnistat.collectors] `port` (8101 here), NOT the library's own default +# of 8001. Mismatch = silent zero kernel metrics. Auto-derive from the +# template so we can never drift. +_OMNI_PORT=$(awk -F'[= \t]+' '/^[[:space:]]*port[[:space:]]*=/ {print $2; exit}' "$OMNISTAT_TEMPLATE") +OMNISTAT_TRACE_ENDPOINT_PORT="${OMNISTAT_TRACE_ENDPOINT_PORT:-${_OMNI_PORT:-8101}}" NODES="${SLURM_JOB_NUM_NODES:-2}" GPUS_PER_NODE=8 @@ -436,6 +442,22 @@ if [[ "$OMNISTAT_KERNEL_TRACE" == "1" ]]; then ) fi +# --------------------------------------------------------------------------- +# Optional env passthrough — only forward HYDRAGNN_NUM_WORKERS / +# HYDRAGNN_PERSISTENT_WORKERS / HYDRAGNN_MAX_NUM_BATCH when set NON-EMPTY. +# HydraGNN's load_data.py uses `os.getenv(KEY) is not None` to detect the +# var, so an empty string passes the check and then `int("")` blows up. +# Building a separate array keeps `apptainer exec` from receiving stale +# `--env KEY=` flags. (See SKILL §11; observed on probe job 7033.) +# --------------------------------------------------------------------------- +OPT_ENVS=() +for k in HYDRAGNN_NUM_WORKERS HYDRAGNN_PERSISTENT_WORKERS HYDRAGNN_MAX_NUM_BATCH; do + v="${!k:-}" + if [[ -n "$v" ]]; then + OPT_ENVS+=( --env "${k}=${v}" ) + fi +done + # --------------------------------------------------------------------------- # Launch distributed training via srun + apptainer # --------------------------------------------------------------------------- @@ -460,15 +482,13 @@ srun --mpi=pmix \ --env HG_PRECISION="$HG_PRECISION" \ --env HG_BATCH_SIZE="${HG_BATCH_SIZE:-200}" \ --env HG_NUM_EPOCH="${HG_NUM_EPOCH:-1}" \ - --env HYDRAGNN_MAX_NUM_BATCH="${HYDRAGNN_MAX_NUM_BATCH:-}" \ --env HYDRAGNN_VALTEST="${HYDRAGNN_VALTEST:-0}" \ --env HYDRAGNN_TRACE_LEVEL="${HYDRAGNN_TRACE_LEVEL:-1}" \ --env HG_EXAMPLE_DIR="$EXAMPLE_DIR" \ --env HG_OUTPUT_DIR="$HG_OUTPUT_DIR" \ --env HG_REPO_DIR="$HG_REPO_DIR" \ --env HG_CONFIG_OVERRIDE="$HG_CONFIG_OVERRIDE" \ - --env HYDRAGNN_NUM_WORKERS="${HYDRAGNN_NUM_WORKERS:-}" \ - --env HYDRAGNN_PERSISTENT_WORKERS="${HYDRAGNN_PERSISTENT_WORKERS:-}" \ + "${OPT_ENVS[@]}" \ --env PROFILE_RANK0_ONLY=1 \ "${RCCL_MULTINODE_ENVS[@]}" \ "$HG_SIF" \ From 3dadfb41828b0e410dc7b7204ec0c689adef4b41 Mon Sep 17 00:00:00 2001 From: Ashwin Aji Date: Fri, 22 May 2026 06:38:21 +0000 Subject: [PATCH 06/31] SKILL: capture session lessons (drain handling, srun-vs-salloc, dual-failure debug) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Adds three procedural notes that were only in chat after the kernel-trace build and validation session: §14. Cluster-state checks before queueing build/probe jobs on Lux. `sinfo -p lux,rad` + `sinfo -R` distinguishes "queued" from "queued forever" when the partition is drained for slurm maintenance. Includes the salloc-vs-srun bind-mount gotcha: `salloc + apptainer exec` doesn't auto-bind /shared/aaji/tools/ in a way that survives across hops, so for one-shot interactive builds inside the SIF prefer `srun ... apptainer exec --bind /shared/aaji/tools:/shared/aaji/tools`. §15. Dual-failure debug pattern — separate "tool wired correctly?" from "workload happy?". When a probe job fails, look for the instrumentation's startup banner before re-building. Worked example: kernel-trace job 7033 crashed in HydraGNN's create_dataloaders (workload bug §13), yet libomnistat_trace.so still printed its 0/0 exit summary on every rank — proving the tool wiring was correct all along. Saved a full rebuild cycle here. Also refreshes the perf-analysis recipe README: - Drops the stale "rocprofiler counters disabled by default" line (they have been on by default since commit a23e5c4). - Adds a "Telemetry knobs" table documenting the OMNISTAT_KERNEL_TRACE switch, the trace-lib path override, and the smoke-signal value of OMNISTAT_TRACE_LOG=1. Co-authored-by: Cursor --- .cursor/skills/ai4science-studio/SKILL.md | 28 +++++++++++++++++++ .../HydraGNN/recipes/perf-analysis/README.md | 11 +++++++- 2 files changed, 38 insertions(+), 1 deletion(-) diff --git a/.cursor/skills/ai4science-studio/SKILL.md b/.cursor/skills/ai4science-studio/SKILL.md index 5b8648d..d9b203d 100755 --- a/.cursor/skills/ai4science-studio/SKILL.md +++ b/.cursor/skills/ai4science-studio/SKILL.md @@ -724,3 +724,31 @@ srun ... apptainer exec ... "${OPT_ENVS[@]}" "$HG_SIF" bash "$RANK_SCRIPT" ``` Applied to both `sbatch_train_amd.sh` and `sbatch_train_perf_amd.sh` for HydraGNN. **General lesson** for any model recipe: never use `--env KEY="${KEY:-}"` for variables that the workload reads with `os.getenv(K) is not None` — either pass a numeric default (`--env KEY="${KEY:-0}"`) or use the conditional-array pattern above. Audit every `*/models/*/examples/*.sh` against this when introducing new optional env knobs. + +### 14. Cluster-state checks before queueing build/probe jobs on Lux + +`sinfo -p lux,rad` is the cheapest way to know whether a build alloc will ever start. Lux and rad partitions get drained for slurm maintenance roughly every few weeks — the symptom is `salloc: job NNNN queued and waiting for resources` followed by `(Nodes required for job are DOWN, DRAINED or reserved for jobs in higher priority partitions)`. The drain reasons are visible via `sinfo -R`; treat any `slurm deploy maintenance`-class reason as "wait, don't queue". Cancel pending allocs with `scancel` when you discover the drain — they don't time out on their own. + +When the cluster is up but you only need a one-shot interactive build inside the SIF, **prefer `srun --no-shell`-equivalent direct exec over `salloc + apptainer exec`**: + +```bash +srun -p lux -A vultr_lux -N 1 --time=00:15:00 --gpus-per-node=1 --cpus-per-task=8 \ + --job-name= \ + apptainer exec --rocm \ + --bind /shared/aaji/tools:/shared/aaji/tools \ + "$HG_SIF" \ + bash -c '' +``` + +Why: `salloc` doesn't auto-bind `/shared/aaji/tools/` into the apptainer namespace (only the `omnistat-src` symlink, which is bind-mounted system-wide, makes it through). Direct `srun … apptainer exec --bind …` forces the bind explicitly and writes the build artifact to its persistent NFS path in one shot. Verified during the kernel-trace library build (job 7030). + +### 15. Dual-failure debug pattern — separate "tool wired correctly?" from "workload happy?" + +When a probe job fails, ask **two** questions before retrying with a bigger run: + +1. **Did the new tool/instrumentation initialize?** Look for its banner / startup log even on a crashed run — a healthy tool prints something before the workload imports anything heavy. +2. **Did the workload reach the part the tool measures?** A crash at line 863 of `gfm_mlip_all_mpnn.py` (HydraGNN's `create_dataloaders`) is upstream of any GPU kernel launch, so any "0 dispatches captured" is consistent with both "tool broken" and "workload broken before kernels". + +**Concrete example from this session (kernel-trace probe job 7033):** all 8 ranks crashed in `int(os.environ["HYDRAGNN_NUM_WORKERS"])` (`ValueError: invalid literal for int() with base 10: ''`). The kernel-trace library *also* printed `[hostname][pid][omnistat] Trace summary: 0/0 processed records (0/0 successful flushes)` per rank — meaning the rocprofiler-sdk tool loaded, registered its callback, and only saw zero dispatches because the workload crashed first. Fixing the env-passthrough bug (§13) and re-running gave 22.5k records/rank — the kernel-trace wiring was correct all along. + +**Action:** when you see "0 records / 0 events / 0 metrics" after a failed probe, *don't* re-build the instrumentation. Read the workload error, fix that, and re-run before changing anything tool-side. This saved one full rebuild cycle here. diff --git a/material_science/models/HydraGNN/recipes/perf-analysis/README.md b/material_science/models/HydraGNN/recipes/perf-analysis/README.md index 6bb3ee0..cdea3ce 100644 --- a/material_science/models/HydraGNN/recipes/perf-analysis/README.md +++ b/material_science/models/HydraGNN/recipes/perf-analysis/README.md @@ -71,6 +71,16 @@ The `combined_report.md` is the deliverable: ranked bottlenecks, remedies tried - One-time install of Omnistat (PR #271 branch `jorda/skills`, with `origin/main` merged in), VictoriaMetrics, and TraceLens — performed lazily by the launcher subagent into `/shared/aaji/tools/`. - The launcher writes `gfm_mlip_with_profile.json` at submit time (does not modify the upstream `gfm_mlip.json`). +## Telemetry knobs (set at `sbatch` submit time) + +| Env var | Default | What it does | +|---|---|---| +| `OMNISTAT_KERNEL_TRACE` | `0` | When `1`, loads `libomnistat_trace.so` via `ROCP_TOOL_LIBRARIES` on every rank and turns on the omnistat kernel-trace collector. Adds per-kernel dispatch count + duration time series across all 8 cards on every node. Requires the library to be built once (see `agents/launcher.md` step 1e). Validated end-to-end on job 7034. | +| `OMNISTAT_TRACE_LIB` | `/shared/aaji/tools/omnistat-src/build-trace/libomnistat_trace.so` | Override only for development. The wrapper hard-fails if the path is missing when `OMNISTAT_KERNEL_TRACE=1`. | +| `OMNISTAT_TRACE_LOG` | `1` | Library prints `[host][pid][omnistat] Trace summary: N/N processed records (M/M successful flushes)` on rank exit; useful as a smoke-signal that the tool initialized even when the workload crashed. | + +`enable_rocprofiler=True` is **on** in `omnistat-lux.config.template` by default (since commit `a23e5c4`); the rendered config's state is printed in the sbatch banner so a silent "counters off" cannot recur. Kernel tracing and device-counting can co-exist on a single GPU but raise per-job VictoriaMetrics cardinality — pick by run length and the question you're asking. See `ai4science-studio` SKILL §12 for the device-counting vs kernel-trace vs `rocprofv3` decision matrix. + ## Bottleneck taxonomy (used by analysts and verifier) Each claim must map to one of: @@ -103,7 +113,6 @@ The orchestrating agent dispatches each as a Task subagent (or shell script in i - LangGraph or any async backplane — subagents communicate via JSON files only. - Multi-rank trace fusion (`TraceLens_generate_multi_rank_collective_report_pytorch`) — single-rank trace for now. -- rocprofiler hardware counters in user-mode Omnistat — disabled by default; verifier may turn on for one targeted probe. - A/B comparative runs — recommended in `combined_report.md` but not auto-launched. - Editing upstream HydraGNN — everything driven via JSON config + env vars. From 3b71da5af392ad46fdb1cf79bb38c0ea37068f5b Mon Sep 17 00:00:00 2001 From: Ashwin Aji Date: Fri, 22 May 2026 09:02:53 +0000 Subject: [PATCH 07/31] HydraGNN: add iterative sysopt-loop recipe (TraceLens+Omnistat, claude CLI) New recipe at material_science/models/HydraGNN/recipes/perf-optimizer-loop/ that drives 3-5 iterations of the existing perf-analysis pipeline, picks one lever per iter from lever_catalog.yaml, accepts/auto-reverts based on primary FOM (epoch_time_s), and writes a progression story. Subagent contract mirrors perf-analysis/: orchestrator -> lever_picker -> launcher -> tracelens+omnistat analysts (parallel) -> verifiers (parallel) -> synthesizer -> fom_extractor -> accept/revert -> next iter, ending in story_writer with story.md + foms.png. Lever catalog has 9 entries ranked by expected payoff at R2 state: batch_800, torch_compile_e3nn (rank-script patch), rccl_high_priority, tunable_op, num_workers_12, nccl_minchannels, precision_fp32 (gated on compute-bound state), kernel_trace_diag_only (gated diagnostic-only). bf16 is intentionally excluded (MACE-class equivariance requires fp32 floor). Each lever cites the ROCm blog / doc that motivates it. fom_extractor adds a TraceLens<->Omnistat 1-second-window correlation (kernel_correlation.csv) so we can attribute the dominant kernel without turning on the omnistat kernel-trace collector (which inflates VM cardinality 1000x; see SKILL #12). Async unattended execution: examples/run_optimizer_loop.sh runs pre-flight checks (cluster up, disk free, claude CLI, ANTHROPIC_API_KEY, https egress) then invokes the orchestrator via claude CLI with retry-with-backoff. Designed for a tmux session on rad-vultr-login so the loop survives ssh disconnect. sbatch_train_perf_amd.sh: extend OPT_ENVS to forward lever knobs (HG_TORCH_COMPILE, TORCH_NCCL_HIGH_PRIORITY, PYTORCH_TUNABLEOP_*, etc.) only when non-empty (avoids the empty-string trap from SKILL #13), and add HG_RANK_PRE_TRAIN_HOOK with parent-dir bind-mount so rank_script_patch levers (e.g. torch.compile) work end-to-end. parse_convergence.py: add --json-foms emitter for the FOM extractor. .gitignore: cover stray omnistat/VM logs and convergence parser outputs. Co-authored-by: Cursor --- .gitignore | 10 + .../HydraGNN/examples/parse_convergence.py | 84 ++++++ .../HydraGNN/examples/run_optimizer_loop.sh | 231 ++++++++++++++++ .../examples/sbatch_train_perf_amd.sh | 52 +++- material_science/models/HydraGNN/model.yaml | 5 + .../recipes/perf-optimizer-loop/README.md | 180 ++++++++++++ .../agents/fom_extractor.md | 258 ++++++++++++++++++ .../agents/lever_picker.md | 120 ++++++++ .../agents/orchestrator.md | 218 +++++++++++++++ .../agents/story_writer.md | 182 ++++++++++++ .../perf-optimizer-loop/lever_catalog.yaml | 182 ++++++++++++ 11 files changed, 1521 insertions(+), 1 deletion(-) create mode 100755 material_science/models/HydraGNN/examples/run_optimizer_loop.sh create mode 100644 material_science/models/HydraGNN/recipes/perf-optimizer-loop/README.md create mode 100644 material_science/models/HydraGNN/recipes/perf-optimizer-loop/agents/fom_extractor.md create mode 100644 material_science/models/HydraGNN/recipes/perf-optimizer-loop/agents/lever_picker.md create mode 100644 material_science/models/HydraGNN/recipes/perf-optimizer-loop/agents/orchestrator.md create mode 100644 material_science/models/HydraGNN/recipes/perf-optimizer-loop/agents/story_writer.md create mode 100644 material_science/models/HydraGNN/recipes/perf-optimizer-loop/lever_catalog.yaml diff --git a/.gitignore b/.gitignore index 79e9c27..666c9b5 100755 --- a/.gitignore +++ b/.gitignore @@ -45,6 +45,16 @@ secrets/ # SLURM job output *-*.out +# SLURM job artifacts named foo..out (e.g. omnistat_failed_hosts.6851.out) +*.[0-9][0-9][0-9][0-9].out + +# Misc tooling artifacts emitted into repo root by past omnistat / VM runs +exporter.log +vic_server.log + +# Convergence parser outputs (artifact, not source) +**/examples/convergence*.png +**/examples/convergence*.csv # Model outputs & generated data **/examples/outputs/ diff --git a/material_science/models/HydraGNN/examples/parse_convergence.py b/material_science/models/HydraGNN/examples/parse_convergence.py index f5fadaa..0aa558a 100755 --- a/material_science/models/HydraGNN/examples/parse_convergence.py +++ b/material_science/models/HydraGNN/examples/parse_convergence.py @@ -272,6 +272,73 @@ def plot_convergence(records: list[dict], plot_path: str, mode: str = "slurm"): print(f"Plot saved to {plot_path}") +def aggregate_foms_from_records(records: list[dict]) -> dict: + """Aggregate per-epoch records into the FOM JSON consumed by the + perf-optimizer-loop's fom_extractor subagent. + + Returns a dict with: + num_epochs_completed, mean_epoch_time_excluding_epoch_0, + epoch_times_s, final_train_loss, final_val_loss. + Any field may be None if the underlying data is absent (e.g. HYDRAGNN_VALTEST=0 + omits val/test losses; older HydraGNN may omit per-epoch wall times). + """ + epoch_summary_records = [r for r in records if r.get("source") == "epoch_summary"] + tqdm_records = [r for r in records if r.get("source") == "tqdm_final"] + + # Prefer epoch_summary wall_time_s only when it provides per-epoch values + # (>= 2 of them). HydraGNN's existing log format usually attaches the + # "train_validate_test" timer to only the FINAL epoch (cumulative whole-job + # wall), which is useless for per-epoch FOMs. In that case, fall through to + # the per-epoch tqdm_final records which always provide s/it × batches. + epoch_times: list[float] = [] + summary_walls = [float(r["wall_time_s"]) for r in epoch_summary_records + if r.get("wall_time_s") is not None] + if len(summary_walls) >= 2: + epoch_times = summary_walls + elif tqdm_records: + for r in tqdm_records: + rate = r.get("wall_time_s") + batch = r.get("batch") + if rate is not None and batch: + epoch_times.append(float(rate) * float(batch)) + elif summary_walls: + epoch_times = summary_walls # single epoch; better than nothing + + if epoch_times: + if len(epoch_times) >= 2: + mean_excluding_epoch_0 = sum(epoch_times[1:]) / float(len(epoch_times) - 1) + warning = None + else: + mean_excluding_epoch_0 = epoch_times[0] + warning = "single epoch, epoch_0 not dropped" + else: + mean_excluding_epoch_0 = None + warning = "no epoch wall_time_s present in log" + + final_train_loss = None + final_val_loss = None + if epoch_summary_records: + final_train_loss = epoch_summary_records[-1].get("train_loss") + final_val_loss = epoch_summary_records[-1].get("val_loss") + + foms = { + "num_epochs_completed": len(epoch_summary_records) if epoch_summary_records else len(tqdm_records), + "mean_epoch_time_excluding_epoch_0": mean_excluding_epoch_0, + "epoch_times_s": epoch_times if epoch_times else None, + "final_train_loss": final_train_loss, + "final_val_loss": final_val_loss, + # Note: total_batches_observed is intentionally NOT derived from the log + # here. The existing HydraGNN entrypoint emits the "Max memory allocated" + # line ONCE per job on rank 0, not per batch. The fom_extractor subagent + # computes samples_processed from manifest values + # (num_epochs * max_num_batch * batch_size * ranks) which is the + # correct denominator for throughput / energy_per_sample FOMs. + } + if warning: + foms["warning"] = warning + return foms + + def main(): parser = argparse.ArgumentParser( description="Extract convergence metrics from HydraGNN training runs." @@ -294,6 +361,13 @@ def main(): "--json", action="store_true", help="Also write a JSON summary alongside the CSV" ) + parser.add_argument( + "--json-foms", type=str, default=None, + help=("If set, write an aggregated FOM JSON to this path " + "(consumed by perf-optimizer-loop's fom_extractor subagent). " + "Includes num_epochs_completed, mean_epoch_time_excluding_epoch_0, " + "epoch_times_s, final_train_loss, final_val_loss."), + ) args = parser.parse_args() @@ -334,6 +408,16 @@ def main(): if args.plot: plot_convergence(records, args.plot, mode=mode) + if args.json_foms: + if mode != "slurm": + print("WARNING: --json-foms only supported with --log mode; skipping.", + file=sys.stderr) + else: + foms = aggregate_foms_from_records(records) + with open(args.json_foms, "w") as f: + json.dump(foms, f, indent=2) + print(f"FOMs JSON written to {args.json_foms}") + # Print summary if mode == "slurm" and records: train_values = [r["train_loss"] for r in records diff --git a/material_science/models/HydraGNN/examples/run_optimizer_loop.sh b/material_science/models/HydraGNN/examples/run_optimizer_loop.sh new file mode 100755 index 0000000..bf27417 --- /dev/null +++ b/material_science/models/HydraGNN/examples/run_optimizer_loop.sh @@ -0,0 +1,231 @@ +#!/usr/bin/env bash +# run_optimizer_loop.sh — entrypoint for the HydraGNN iterative sysopt loop. +# +# Runs pre-flight checks, sets up the loop directory, and invokes the +# orchestrator subagent via Claude Code CLI in the current tmux session. +# Survives ssh disconnect when run inside `tmux new -s hg-loop`. +# +# Usage: +# export ANTHROPIC_API_KEY=... +# tmux new -s hg-loop +# bash material_science/models/HydraGNN/examples/run_optimizer_loop.sh \ +# +# # Ctrl-b d to detach; the loop continues until completion or LOOP_ABORT. +# +# Args: +# loop-uuid Unique id for this loop (use `uuidgen`). +# n_iters_budget Max iterations (recommend 5). +# +# Abort: +# Graceful: touch /shared/aaji/models/HydraGNN/perf-runs/loop-/STOP +# Emergency: scancel + tmux kill-session -t hg-loop + +set -euo pipefail + +if [[ $# -lt 2 ]]; then + echo "Usage: $0 [--preflight-only]" >&2 + echo "" >&2 + echo " --preflight-only Run pre-flight checks only; do not invoke claude." >&2 + echo " Useful for smoke-testing the recipe." >&2 + exit 1 +fi + +LOOP_UUID="$1" +N_ITERS="$2" +PREFLIGHT_ONLY=0 +if [[ "${3:-}" == "--preflight-only" ]]; then + PREFLIGHT_ONLY=1 +fi + +# --- Resolve paths --- +SCRIPT_DIR=$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd) +REPO_ROOT=$(cd "${SCRIPT_DIR}/../../../.." && pwd) +RECIPE_DIR="${REPO_ROOT}/material_science/models/HydraGNN/recipes/perf-optimizer-loop" +ORCH_PROMPT="${RECIPE_DIR}/agents/orchestrator.md" + +# Loop-dir convention matches recipes/perf-optimizer-loop/README.md +AI4S_SHARED_DIR="${AI4S_SHARED_DIR:-/shared/aaji}" +HG_BASE="${AI4S_SHARED_DIR}/models/HydraGNN" +PERF_RUNS_DIR="${HG_BASE}/perf-runs" +LOOP_DIR="${PERF_RUNS_DIR}/loop-${LOOP_UUID}" +STATUS_FILE="${LOOP_DIR}/STATUS.txt" + +mkdir -p "$LOOP_DIR" + +# --- Helper: log to STATUS.txt with flock so multiple writers don't clobber --- +_log_status() { + local msg="$1" + local line + line="$(date -u +%Y-%m-%dT%H:%M:%SZ) $msg" + ( + flock -x 9 + echo "$line" >> "$STATUS_FILE" + ) 9>>"$STATUS_FILE" + echo "[STATUS] $line" +} + +_log_status "LOOP_START uuid=${LOOP_UUID} n_iters_budget=${N_ITERS} driver=claude-code-cli script=$0" + +# --- Pre-flight --- +PREFLIGHT_FAIL_REASON="" +PREFLIGHT_NOTES=() + +# 1. cluster up +if command -v sinfo > /dev/null 2>&1; then + _IDLE_OR_MIX=$(sinfo -p lux,rad -h -o '%t' 2>/dev/null | grep -cE '^(idle|mix|alloc)$' || true) + if [[ "${_IDLE_OR_MIX:-0}" -eq 0 ]]; then + PREFLIGHT_FAIL_REASON="cluster_down" + PREFLIGHT_NOTES+=("sinfo reports 0 usable nodes on partitions lux,rad") + fi +else + PREFLIGHT_FAIL_REASON="no_slurm" + PREFLIGHT_NOTES+=("sinfo not in PATH — must run on a SLURM login node") +fi + +# 2. disk free +if [[ -z "$PREFLIGHT_FAIL_REASON" ]]; then + _USE_PCT=$(df -P "$PERF_RUNS_DIR" | awk 'NR==2 {sub("%","",$5); print $5}') + if [[ "${_USE_PCT:-0}" -gt 95 ]]; then + PREFLIGHT_FAIL_REASON="disk_full" + PREFLIGHT_NOTES+=("/shared at ${_USE_PCT}% — refusing to start a loop that adds ~0.7GB") + fi +fi + +# 3. tools present +for _bin in \ + /shared/aaji/tools/omnistat-pr271/bin/omnistat-usermode \ + /shared/aaji/tools/victoriametrics/victoria-metrics-prod; do + if [[ -z "$PREFLIGHT_FAIL_REASON" && ! -x "$_bin" ]]; then + PREFLIGHT_FAIL_REASON="tool_missing" + PREFLIGHT_NOTES+=("missing: $_bin (run the perf-analysis launcher subagent first to install)") + fi +done + +# 4. SIF + overlay +for _f in \ + "${HG_BASE}/overlays/hydragnn-overlay.img" \ + "${AI4S_SHARED_DIR}/images/pytorch_rocm7.2.2_ubuntu24.04_py3.12_pytorch_release_2.10.0.sif"; do + if [[ -z "$PREFLIGHT_FAIL_REASON" && ! -f "$_f" ]]; then + PREFLIGHT_FAIL_REASON="image_missing" + PREFLIGHT_NOTES+=("missing: $_f") + fi +done + +# 5. claude CLI + API key + egress +if [[ -z "$PREFLIGHT_FAIL_REASON" ]]; then + if ! command -v claude > /dev/null 2>&1; then + PREFLIGHT_FAIL_REASON="claude_cli_missing" + PREFLIGHT_NOTES+=("'claude' CLI not in PATH — install per https://docs.anthropic.com/claude/docs/claude-code") + elif [[ -z "${ANTHROPIC_API_KEY:-}" ]]; then + PREFLIGHT_FAIL_REASON="no_api_key" + PREFLIGHT_NOTES+=("ANTHROPIC_API_KEY env var not set — required for unattended runs") + else + # Egress check: we only care that we can REACH api.anthropic.com — not + # that the bare endpoint accepts our key. A 401/403 response means TLS + # completed and we got a server reply, i.e. egress works (and the API + # key, if any, doesn't have read access to /v1/models — that's normal + # for many key scopes; /v1/messages is what actually matters and is + # exercised when claude actually runs). Treat 401/403 as success. + _http_code=$(curl -s -o /dev/null -w '%{http_code}' -m 5 https://api.anthropic.com/v1/models 2>/dev/null || echo "000") + case "$_http_code" in + 000|5*) PREFLIGHT_FAIL_REASON="api_egress_blocked" + PREFLIGHT_NOTES+=("cannot reach api.anthropic.com:443 (curl returned ${_http_code})") ;; + 2*|401|403) : ;; # ok: server reachable + *) PREFLIGHT_FAIL_REASON="api_egress_unexpected" + PREFLIGHT_NOTES+=("unexpected http code ${_http_code} from api.anthropic.com/v1/models") ;; + esac + fi +fi + +# 6. orchestrator prompt present +if [[ -z "$PREFLIGHT_FAIL_REASON" && ! -f "$ORCH_PROMPT" ]]; then + PREFLIGHT_FAIL_REASON="orch_prompt_missing" + PREFLIGHT_NOTES+=("orchestrator prompt not found: $ORCH_PROMPT") +fi + +if [[ -n "$PREFLIGHT_FAIL_REASON" ]]; then + _log_status "PREFLIGHT_FAIL reason=${PREFLIGHT_FAIL_REASON} notes='${PREFLIGHT_NOTES[*]}'" + echo "Pre-flight failed: ${PREFLIGHT_FAIL_REASON}" >&2 + printf ' - %s\n' "${PREFLIGHT_NOTES[@]}" >&2 + exit 2 +fi + +_DISK_FREE=$(df -h "$PERF_RUNS_DIR" | awk 'NR==2 {print $4}') +_log_status "PREFLIGHT_OK cluster=lux disk_free=${_DISK_FREE} claude_cli=ok api_egress=ok orch_prompt=ok" + +if [[ $PREFLIGHT_ONLY -eq 1 ]]; then + _log_status "PREFLIGHT_ONLY mode: exiting before invoking orchestrator" + echo "" + echo "Pre-flight complete; orchestrator NOT invoked (--preflight-only)." + echo " STATUS file: $STATUS_FILE" + exit 0 +fi + +# --- Trap SIGINT/SIGTERM and log LOOP_ABORT before exiting --- +_on_signal() { + _log_status "LOOP_ABORT reason=signal sig=$1 (driver exiting; in-flight SLURM job not cancelled)" + exit 130 +} +trap '_on_signal INT' INT +trap '_on_signal TERM' TERM + +# --- Invoke orchestrator via claude CLI with retry-with-backoff --- +# We invoke claude ONCE per attempt and let the orchestrator drive the full +# loop. If claude exits non-zero (transient API failure, network hiccup) we +# retry up to 3 times with exp backoff; on the fourth failure we log +# LOOP_ABORT and exit 1. +# +# The orchestrator persists all state to disk (foms.csv, STATUS.txt, +# do_not_retry.json) so a fresh invocation can resume from where the +# previous one left off — see orchestrator.md §0 (Resume detection). + +USER_PROMPT="Drive the HydraGNN iterative sysopt loop with these parameters: +- Loop UUID: ${LOOP_UUID} +- Loop directory: ${LOOP_DIR} +- Iterations budget: ${N_ITERS} +- Status file: ${STATUS_FILE} +- Recipe: ${RECIPE_DIR}/README.md + +Follow the steps in the orchestrator.md system prompt. Persist all state to +disk so you can resume from where you left off if I have to restart you. +On any LOOP_COMPLETE, LOOP_ABORT, or LOOP_PAUSE, exit with a single line: + STATUS=; reason=" + +_MAX_ATTEMPTS=3 +_attempt=1 +while [[ $_attempt -le $_MAX_ATTEMPTS ]]; do + _log_status "ORCH_INVOKE attempt=${_attempt} max=${_MAX_ATTEMPTS}" + set +e + claude \ + --print \ + --dangerously-skip-permissions \ + --max-turns 250 \ + --append-system-prompt "$(cat "$ORCH_PROMPT")" \ + "$USER_PROMPT" + _rc=$? + set -e + if [[ $_rc -eq 0 ]]; then + _log_status "ORCH_EXIT_OK attempt=${_attempt}" + break + fi + _log_status "ORCH_EXIT_FAIL attempt=${_attempt} rc=${_rc}" + if [[ $_attempt -eq $_MAX_ATTEMPTS ]]; then + _log_status "LOOP_ABORT reason=orchestrator_failed rc=${_rc} attempts=${_attempt}" + exit 1 + fi + _backoff=$(( 30 * (2 ** (_attempt - 1)) )) # 30s, 60s, 120s + _log_status "ORCH_BACKOFF sleep=${_backoff}s" + sleep "$_backoff" + _attempt=$(( _attempt + 1 )) +done + +_log_status "DRIVER_EXIT clean" + +echo "" +echo "=== Loop ${LOOP_UUID} driver exited cleanly ===" +echo " STATUS file : $STATUS_FILE" +echo " Loop dir : $LOOP_DIR" +echo "" +echo "Inspect the final result:" +echo " cat $LOOP_DIR/story.md" +echo " open $LOOP_DIR/foms.png" diff --git a/material_science/models/HydraGNN/examples/sbatch_train_perf_amd.sh b/material_science/models/HydraGNN/examples/sbatch_train_perf_amd.sh index effef88..2111b54 100755 --- a/material_science/models/HydraGNN/examples/sbatch_train_perf_amd.sh +++ b/material_science/models/HydraGNN/examples/sbatch_train_perf_amd.sh @@ -345,6 +345,24 @@ if avg_nn > 0: self.avg_num_neighbors = avg_nn adm.AdiosMultiDataset.__init__ = _patched_init +# Optional pre-train hook used by the perf-optimizer-loop recipe to apply a +# rank_script_patch lever (e.g. torch.compile wrapping). The orchestrator writes +# a small python file into the loop dir and exports HG_RANK_PRE_TRAIN_HOOK to +# its absolute path; the hook executes in this rank's interpreter BEFORE +# train_validate_test() runs, so it can monkey-patch hydragnn modules. +# Hook is expected to be small + idempotent; errors are logged on rank 0 but +# never abort the training (so a hook bug doesn't waste a whole 2-node sbatch). +_hook = os.environ.get('HG_RANK_PRE_TRAIN_HOOK', '') +if _hook and os.path.isfile(_hook): + if os.environ.get('SLURM_PROCID', '0') == '0': + print(f'[rank_hook] executing pre-train hook: {_hook}', flush=True) + try: + with open(_hook) as _hf: + exec(_hf.read(), {'__name__': '__hg_rank_hook__'}) + except Exception as _e: + if os.environ.get('SLURM_PROCID', '0') == '0': + print(f'[rank_hook] FAILED: {type(_e).__name__}: {_e}', flush=True) + batch_size = int(os.environ.get('HG_BATCH_SIZE', '0') or 200) max_num_batch = int(os.environ.get('HYDRAGNN_MAX_NUM_BATCH', '0')) num_samples_val = max_num_batch * batch_size if max_num_batch > 0 else 0 @@ -449,15 +467,46 @@ fi # var, so an empty string passes the check and then `int("")` blows up. # Building a separate array keeps `apptainer exec` from receiving stale # `--env KEY=` flags. (See SKILL §11; observed on probe job 7033.) +# +# The same conditional-forward pattern applies to lever knobs from the +# perf-optimizer-loop recipe (HG_TORCH_COMPILE, TORCH_NCCL_HIGH_PRIORITY, +# PYTORCH_TUNABLEOP_ENABLED, etc.) — most of these are read via plain +# `os.environ.get(K, default)` so an empty string would override the default +# back to "" rather than the desired default. Forward only when non-empty. # --------------------------------------------------------------------------- OPT_ENVS=() -for k in HYDRAGNN_NUM_WORKERS HYDRAGNN_PERSISTENT_WORKERS HYDRAGNN_MAX_NUM_BATCH; do +for k in \ + HYDRAGNN_NUM_WORKERS HYDRAGNN_PERSISTENT_WORKERS HYDRAGNN_MAX_NUM_BATCH \ + HG_TORCH_COMPILE HG_TORCH_COMPILE_MODE TORCHINDUCTOR_MAX_AUTOTUNE \ + TORCH_NCCL_HIGH_PRIORITY GPU_MAX_HW_QUEUES \ + PYTORCH_TUNABLEOP_ENABLED PYTORCH_TUNABLEOP_TUNING PYTORCH_TUNABLEOP_VERBOSE \ + NCCL_MIN_NCHANNELS; do v="${!k:-}" if [[ -n "$v" ]]; then OPT_ENVS+=( --env "${k}=${v}" ) fi done +# --------------------------------------------------------------------------- +# Optional pre-train hook (perf-optimizer-loop rank_script_patch lever). +# When HG_RANK_PRE_TRAIN_HOOK points at a host-side .py file, bind-mount its +# parent dir read-only and forward the env var. The rank script's main +# python -c block sources this file before train_validate_test(). +# --------------------------------------------------------------------------- +HOOK_BIND_ENVS=() +if [[ -n "${HG_RANK_PRE_TRAIN_HOOK:-}" ]]; then + if [[ -f "$HG_RANK_PRE_TRAIN_HOOK" ]]; then + _HOOK_DIR=$(cd "$(dirname "$HG_RANK_PRE_TRAIN_HOOK")" && pwd) + HOOK_BIND_ENVS=( + --bind "${_HOOK_DIR}:${_HOOK_DIR}:ro" + --env HG_RANK_PRE_TRAIN_HOOK="$HG_RANK_PRE_TRAIN_HOOK" + ) + echo " PreTrainHook : $HG_RANK_PRE_TRAIN_HOOK (bind-mounted ${_HOOK_DIR})" + else + echo "WARN: HG_RANK_PRE_TRAIN_HOOK=$HG_RANK_PRE_TRAIN_HOOK is set but file not found; continuing without hook" >&2 + fi +fi + # --------------------------------------------------------------------------- # Launch distributed training via srun + apptainer # --------------------------------------------------------------------------- @@ -489,6 +538,7 @@ srun --mpi=pmix \ --env HG_REPO_DIR="$HG_REPO_DIR" \ --env HG_CONFIG_OVERRIDE="$HG_CONFIG_OVERRIDE" \ "${OPT_ENVS[@]}" \ + "${HOOK_BIND_ENVS[@]}" \ --env PROFILE_RANK0_ONLY=1 \ "${RCCL_MULTINODE_ENVS[@]}" \ "$HG_SIF" \ diff --git a/material_science/models/HydraGNN/model.yaml b/material_science/models/HydraGNN/model.yaml index 6f84764..30e546c 100644 --- a/material_science/models/HydraGNN/model.yaml +++ b/material_science/models/HydraGNN/model.yaml @@ -29,6 +29,11 @@ recipes: description: 2-node bottleneck analysis with TraceLens + Omnistat (user-mode) via paired analyst/verifier subagents recipe_path: recipes/perf-analysis/ sbatch_script: examples/sbatch_train_perf_amd.sh + - task: perf-optimizer-loop + description: Iterative systems optimizer; 3-5 iter loop over the perf-analysis pipeline, picks one lever per iter, auto-revert on regression, writes story.md + foms.png + recipe_path: recipes/perf-optimizer-loop/ + script: examples/run_optimizer_loop.sh + sbatch_script: examples/sbatch_train_perf_amd.sh env_vars: # Required for HPC scripts (sbatch_*_amd.sh, build_overlay_amd.sh) diff --git a/material_science/models/HydraGNN/recipes/perf-optimizer-loop/README.md b/material_science/models/HydraGNN/recipes/perf-optimizer-loop/README.md new file mode 100644 index 0000000..8dd8aba --- /dev/null +++ b/material_science/models/HydraGNN/recipes/perf-optimizer-loop/README.md @@ -0,0 +1,180 @@ +# HydraGNN Iterative Systems Optimizer Loop + +Multi-iteration optimizer that runs HydraGNN on 2 nodes × 8 × MI355X (gfx950), picks one performance lever per iteration, runs the existing analyst/verifier/synth pipeline from [`recipes/perf-analysis/`](../perf-analysis/), records six FOMs, accepts or auto-reverts based on the primary FOM, and writes a progression story across 3-5 iterations. + +> Audience: AMD performance engineering. Output is a diagnosis + tunings, not a HydraGNN scientific claim. + +## What this recipe adds on top of `perf-analysis/` + +| Concern | `perf-analysis/` | `perf-optimizer-loop/` | +|---|---|---| +| Number of runs | one | 3-5, with auto accept/revert | +| Lever selection | manual (user reads `combined_report.md`) | LLM-driven via `lever_picker` over `lever_catalog.yaml` | +| FOM tracking | per-run claims | per-iter `foms.csv` + progression chart | +| Output | one `combined_report.md` | per-iter reports + final `story.md` + `foms.png` | +| Async | foreground only | tmux + Claude Code CLI driver for overnight runs | +| Attribution | TraceLens xlsx + Omnistat JSON | adds 1-s window TraceLens↔Omnistat correlation | + +## Orchestration loop + +```mermaid +flowchart TD + User([User: run_optimizer_loop.sh]) --> Orch[orchestrator subagent] + Orch --> Iter0[iter 0: fresh baseline] + Iter0 --> Launch[launcher subagent: sbatch+manifest] + Launch --> Analysts + subgraph Analysts [analyst phase parallel] + TLA[tracelens_analyst] + OSA[omnistat_analyst] + end + Analysts --> Verifiers + subgraph Verifiers [verifier phase parallel] + TLV[tracelens_verifier] + OSV[omnistat_verifier] + end + Verifiers --> Synth[synthesizer] + Synth --> FOM[fom_extractor: foms.json + kernel_correlation.csv] + FOM --> Accept{"primary FOM
improved?"} + Accept -->|yes or iter==0| Update[update best, append foms.csv] + Accept -->|no| Revert[record do_not_retry, restore prev best] + Update --> Done{"iter==5 or
stop criteria?"} + Revert --> Done + Done -->|no| Pick[lever_picker: combined_report + history + ROCm crib-sheet] + Pick --> NextIter[next iter: apply ONE lever] + NextIter --> Launch + Done -->|yes| StoryW[story_writer: story.md + foms.png] + StoryW --> Final([final report to user]) +``` + +## Quick start (overnight unattended via Claude Code CLI in tmux) + +```bash +ssh aaji@rad-vultr-login +export ANTHROPIC_API_KEY=... # user sets manually; never committed +cd /home/aaji/git/ai4science-studio +git checkout aaji/perf-optimizer-loop-hydragnn +tmux new -s hg-loop +bash material_science/models/HydraGNN/examples/run_optimizer_loop.sh $(uuidgen) 5 +# Ctrl-b d to detach. tmux session lives until login-node reboot. +``` + +Morning workflow: + +```bash +cat /shared/aaji/models/HydraGNN/perf-runs/loop-/STATUS.txt | tail -50 +cat /shared/aaji/models/HydraGNN/perf-runs/loop-/story.md +``` + +## Artifact layout (per-loop) + +``` +/shared/aaji/models/HydraGNN/perf-runs/ +├── loop-.json # registers loop id, baseline jobid, best jobid +├── loop-/ +│ ├── STATUS.txt # append-only event log (driver writes after every step) +│ ├── STOP # touch this file to gracefully stop the loop +│ ├── do_not_retry.json # levers that regressed primary FOM +│ ├── foms.csv # one row per iteration +│ ├── foms.png # FOM progression line chart +│ ├── story.md # final progression narrative +│ ├── iter-0-baseline.json # symlink → ..//manifest.json +│ ├── iter-1-batch_800.json # symlink → ..//manifest.json +│ └── ... +└── / # per-iter dir, written by existing perf-analysis launcher + ├── manifest.json + ├── logs/*.pt.trace.json + ├── omnistat-db/ + ├── tracelens/{report.xlsx,csvs/,verified_claims.json} + ├── omnistat/{inspect_outputs/,verified_claims.json} + ├── kernel_correlation.csv # NEW: TraceLens↔Omnistat 1-s window join + ├── foms.json # NEW: this iter's six FOMs + └── combined_report.md # synthesizer output +``` + +## Iteration shape + +Each iteration submits the existing [`sbatch_train_perf_amd.sh`](../../examples/sbatch_train_perf_amd.sh) with overrides: + +| Knob | Loop default | Why | +|---|---|---| +| `HG_NUM_EPOCH` | 3 | Steady-state past warm-up; epoch-0 timing dropped | +| `HYDRAGNN_MAX_NUM_BATCH` | 50 | 6-8 min wall; fits in `--time=00:30:00` | +| `PROFILE_TARGET_EPOCH` | 2 | Sidesteps the worker-respawn ProfilerStep#8 artifact from R2 | +| `HG_BATCH_SIZE` | 400 (baseline) | R2 best; loop may move to 800 in iter-1 | + +## Figures of Merit + +Per `agents/fom_extractor.md`: + +| FOM | Source | Role | +|---|---|---| +| `epoch_time_s` | `parse_convergence.py` over `hydragnn-train-.out`, mean over `epoch >= 1` | **PRIMARY** (accept/revert) | +| `throughput_samples_per_s` | `max_num_batch * batch_size * ranks / epoch_time` | report | +| `mfma_tflops` | PromQL: `rate(SQ_INSTS_VALU_MFMA_MOPS_F64) * join(rmsjob_info{jobid})` | report | +| `energy_J` | integral of `rocm_avg_pwr` over runtime, summed across 16 cards | report | +| `mean_power_W` | `energy_J / runtime` | report | +| `energy_per_sample_J` | `energy_J / total_samples_processed` | report | +| `final_loss` | last `train_loss` from `convergence.csv` | **CONTROL** (auto-reject if > 1.5× baseline) | + +## Levers in the catalog (ranked by expected payoff at R2 state) + +See [`lever_catalog.yaml`](lever_catalog.yaml) for the full machine-readable list. High-level ranking and provenance: + +| # | Lever id | Type | Expected payoff | Citation | +|---|---|---|---|---| +| 1 | `batch_800` | env-only | `HG_BATCH_SIZE` 400→800; R2 #1 next-lever | R2 `combined_report.md` | +| 2 | `torch_compile_e3nn` | rank-script patch | `torch.compile(model, fullgraph=False, mode='max-autotune')`; R2 #2 next-lever | [MI300X workload §torch.compile](https://rocm.docs.amd.com/en/latest/how-to/tuning-guides/mi300x/workload.html) | +| 3 | `rccl_high_priority` | env-only | `TORCH_NCCL_HIGH_PRIORITY=1`, `GPU_MAX_HW_QUEUES=2` | [MI300X workload §RCCL](https://rocm.docs.amd.com/en/latest/how-to/tuning-guides/mi300x/workload.html) | +| 4 | `tunable_op` | env-only | `PYTORCH_TUNABLEOP_ENABLED=1`, `PYTORCH_TUNABLEOP_TUNING=1` | [MI300X workload §TunableOp](https://rocm.docs.amd.com/en/latest/how-to/tuning-guides/mi300x/workload.html) | +| 5 | `num_workers_sweep` | env-only | `HYDRAGNN_NUM_WORKERS=12`, persistent_workers=1 | local lesson (R1, R2 reports) | +| 6 | `nccl_minchannels` | env-only | `NCCL_MIN_NCHANNELS=112` | [Multi-node ROCm AI setup](https://rocm.docs.amd.com/en/docs-7.0.1/how-to/rocm-for-ai/system-setup/multi-node-setup.html) | +| 7 | `precision_fp32` | env-only | `HG_PRECISION=fp32`; gfx950 fp32 MFMA peak ~2× fp64 | gfx950 spec; HydraGNN/MACE fp32 floor (bf16 breaks equivariance) | +| 8 | `kernel_trace_diag_only` | env-only, diagnostic | `OMNISTAT_KERNEL_TRACE=1`; only when 1-s TL↔OS correlation can't attribute the dominant kernel | [SKILL §12](../../../../../.cursor/skills/ai4science-studio/SKILL.md) | + +`bf16` is **deliberately excluded** — the model is fp32-floor for MACE-style equivariant features. + +## Loop control + +- **Cap:** 5 iterations. +- **Stop early:** when (a) `lever_picker` reports the catalog is exhausted (every entry tried or in `do_not_retry.json`), OR (b) the last 2 accepted iterations each improved primary FOM by < 5%. +- **Auto-revert:** an iteration that regresses `epoch_time_s` (vs current best) by any amount, OR raises `final_loss` by > 50% vs baseline, is rejected. The lever id lands in `do_not_retry.json`; the next iteration restores the previous-best config. +- **Diagnostic iterations** (`diagnostic_only=true` in the catalog) do NOT count toward the auto-revert logic, but DO count against the 5-iter budget. +- **Single variable per iteration** — the orchestrator must change exactly one lever, otherwise FOM deltas cannot be attributed. + +## Abort controls + +- **Graceful:** `touch /shared/aaji/models/HydraGNN/perf-runs/loop-/STOP`. Driver checks the flag (a) before every `sbatch`, (b) before every analyst phase, (c) before every iteration decision. Active SLURM jobs run to completion and get analyzed; no further iterations launch. +- **Emergency:** `scancel ` for the active job + `tmux kill-session -t hg-loop` (or `kill -INT `). The driver traps SIGINT and writes `LOOP_ABORT reason=signal` to STATUS.txt before exiting. + +## Disk budget + +- Per iteration footprint: ~135 MB (matches R2 shape: kineto trace 110-125 MB + TraceLens 8 MB + omnistat-db 1-3 MB + omnistat inspect ~200 KB). +- 5-iter loop projected: ~0.7 GB. `/shared` has 26 TB free (18% used); not pressured. +- No per-loop quota enforced; orchestrator logs `du -sh loop-/` after each iter to STATUS.txt for visibility. +- `OMNISTAT_KERNEL_TRACE=1` raises VictoriaMetrics cardinality ~1000× (per [SKILL §12](../../../../../.cursor/skills/ai4science-studio/SKILL.md)). The `kernel_trace_diag_only` lever is gated `diagnostic_only=true` and only fires when the 1-s correlation falls short — at most 1 of 5 iterations. + +## Agent prompt files + +- [`agents/orchestrator.md`](agents/orchestrator.md) — main loop, dispatches all other subagents, accept/revert/STATUS.txt/STOP-flag. +- [`agents/lever_picker.md`](agents/lever_picker.md) — proposes the next single lever as JSON, biased by `lever_catalog.yaml` + previous reports + ROCm crib-sheet. +- [`agents/fom_extractor.md`](agents/fom_extractor.md) — computes 6 FOMs + writes `kernel_correlation.csv` from TraceLens trace + Omnistat PromQL. +- [`agents/story_writer.md`](agents/story_writer.md) — final progression narrative + matplotlib `foms.png`. + +Plus the existing prompts inherited from [`recipes/perf-analysis/agents/`](../perf-analysis/agents/): `launcher.md`, `tracelens_analyst.md`, `tracelens_verifier.md`, `omnistat_analyst.md`, `omnistat_verifier.md`, `synthesizer.md`. + +## References (ROCm blogs and docs cited in the lever catalog) + +- [Optimizing LLM Workloads: AMD Instinct MI355X GPUs Drive Competitive Performance — ROCm Blogs (ROCm 7.0)](https://rocm.blogs.amd.com/artificial-intelligence/ROCm7-MI355X-training-performance/README.html) +- [AMD Instinct MI300X workload optimization — ROCm Documentation](https://rocm.docs.amd.com/en/latest/how-to/tuning-guides/mi300x/workload.html) +- [QuickReduce FP4 Quantization and Benchmarking on MI355 — ROCm Blogs](https://rocm.blogs.amd.com/artificial-intelligence/quick-reduce-2/README.html) +- [Accelerating ComfyUI Workflows on AMD Instinct MI355X GPUs with ROCm — ROCm Blogs](https://rocm.blogs.amd.com/artificial-intelligence/comfyui/README.html) +- [Multi-node setup for AI workloads — ROCm Documentation](https://rocm.docs.amd.com/en/docs-7.0.1/how-to/rocm-for-ai/system-setup/multi-node-setup.html) +- [Training a model with PyTorch on ROCm — ROCm Documentation](https://rocm.docs.amd.com/en/develop/how-to/rocm-for-ai/training/benchmark-docker/previous-versions/pytorch-training-v25.9.html) + +## Out of scope + +- Cross-model levers (HydraGNN only here). +- Multi-variable lever combinations in one iteration (breaks delta attribution; do as separate study after this loop converges). +- Auto-rebuild of HydraGNN overlay; only env-var and rank-script-monkey-patch changes allowed. +- Editing upstream HydraGNN source; use the rank-script monkey-patch pattern already in [`sbatch_train_perf_amd.sh`](../../examples/sbatch_train_perf_amd.sh). +- A scientific claim about HydraGNN convergence — short runs, loss is a control, not a headline. diff --git a/material_science/models/HydraGNN/recipes/perf-optimizer-loop/agents/fom_extractor.md b/material_science/models/HydraGNN/recipes/perf-optimizer-loop/agents/fom_extractor.md new file mode 100644 index 0000000..5ee453d --- /dev/null +++ b/material_science/models/HydraGNN/recipes/perf-optimizer-loop/agents/fom_extractor.md @@ -0,0 +1,258 @@ +# fom_extractor subagent + +Computes the six Figures of Merit for an iteration and writes the +TraceLens↔Omnistat 1-second-window correlation that subsequent analysts and +lever_picker use as evidence for per-kernel attribution (cheaper than the +omnistat kernel-trace collector). + +## Inputs + +- `/manifest.json` (from the existing launcher subagent) +- `/hydragnn-train-.out` (training log; epoch summaries + tqdm) +- `/logs/*.pt.trace.json` (rank-0 PyTorch trace; may be a single file) +- `/omnistat-db/` (VictoriaMetrics datadir) +- `/tracelens/csvs/` (already produced by tracelens_analyst; we re-use `ops_summary.csv` etc., but we MUST NOT depend on tracelens_analyst running successfully — fall back to raw trace if csvs missing) +- A live VictoriaMetrics endpoint started by omnistat_analyst on this perf_run, OR start our own on a free port if not running + +## Outputs + +- `/foms.json` — schema below +- `/kernel_correlation.csv` — per-1s-window, per-card, with columns: + `window_start_iso, window_start_ts, instance, gpu_id, top_kernel_name, top_kernel_busy_frac, second_kernel_name, second_kernel_busy_frac, mfma_tflops, hbm_read_GBps, mean_power_W` + +## FOM schema (`foms.json`) + +```json +{ + "iter": 1, + "jobid": "7050", + "lever": {"id": "batch_800", "env_diff": {"HG_BATCH_SIZE": "800"}}, + "runtime_s": 487.3, + "num_epochs_completed": 3, + "samples_processed": 384000, + "foms": { + "epoch_time_s": 145.3, + "throughput_samples_per_s": 11020.0, + "mfma_tflops": 0.687, + "energy_J": 1.12e6, + "mean_power_W": 4800.0, + "energy_per_sample_J": 12.4, + "final_loss": 5.21 + }, + "fom_sources": { + "epoch_time_s": "parse_convergence.py over hydragnn-train-7050.out, mean(wall_time_s) for epoch>=1", + "throughput_samples_per_s": "max_num_batch * batch_size * ranks / epoch_time_s", + "mfma_tflops": "PromQL: avg(rate(SQ_INSTS_VALU_MFMA_MOPS_F64[10s]) * on(instance) group_left() max by(instance)(rmsjob_info{jobid='7050'})) * 4", + "energy_J": "integral of rocm_avg_pwr over runtime_s, summed across 16 cards", + "mean_power_W": "energy_J / runtime_s", + "energy_per_sample_J": "energy_J / samples_processed", + "final_loss": "last epoch_summary train_loss from convergence.csv" + }, + "kernel_correlation_summary": { + "windows_examined": 280, + "windows_with_top_kernel_busy_frac_gt_0p5": 165, + "top_kernel_dominants": [["aten::mm", 78], ["aten::bmm", 52], ["triton_red_fused_*", 35]], + "median_mfma_tflops_when_aten_mm_dominant": 3.15, + "median_mfma_tflops_when_aten_bmm_dominant": 0.019, + "median_hbm_read_GBps_when_mul_dominant": 4.5, + "attribution_quality": "good" + } +} +``` + +`attribution_quality`: "good" if > 50% of windows have `top_kernel_busy_frac > 0.5`; "fair" if 25-50%; "poor" if < 25%. lever_picker uses "poor" as one of the gates for picking `kernel_trace_diag_only`. + +## Steps + +### 1. Sanity check + env + +```bash +set -euo pipefail +. /shared/aaji/tools/omnistat-pr271/bin/activate +PERF_RUN=$(jq -r .perf_run_dir ) +JOBID=$(jq -r .jobid ) +RUNTIME=$(jq -r .runtime_seconds ) +TRACE=$(jq -r '.trace_paths[0] // empty' ) +[[ -z "$TRACE" ]] && { echo "STATUS=partial; reason=no_trace; FOMs computed without correlation" >&2; } +``` + +### 2. Time-based FOMs from training log + +Use the existing parser, extended with a `--json-foms` flag (see "parse_convergence.py extension" below): + +```bash +python3 /home/aaji/git/ai4science-studio/material_science/models/HydraGNN/examples/parse_convergence.py \ + --log "$PERF_RUN/hydragnn-train-$JOBID.out" \ + --output "$PERF_RUN/convergence.csv" \ + --json-foms "$PERF_RUN/_convergence_foms.json" +``` + +`_convergence_foms.json` is consumed below; it provides `mean_epoch_time_excluding_epoch_0`, `num_epochs_completed`, `final_train_loss`, `total_batches_observed`. + +Compute: +- `epoch_time_s = mean_epoch_time_excluding_epoch_0` +- `samples_processed = num_epochs_completed * (max_num_batch * batch_size * ranks)` where the three multiplicand values come from `manifest.json`'s `hg_batch_size`, `hydragnn_max_num_batch`, `ranks`. +- `throughput_samples_per_s = samples_processed / (runtime_s)`. (Use runtime_s rather than mean_epoch_time × n_epochs because runtime includes warm-up, JIT, and finalization — the user wants real wall-time efficiency.) + +### 3. Counter-based FOMs via PromQL + +Start a local VictoriaMetrics if not already running: + +```bash +PORT=8428 +ss -lnt "sport = :$PORT" 2>/dev/null | grep -q LISTEN && PORT=8429 +if ! curl -sf "http://127.0.0.1:$PORT/api/v1/status/tsdb" > /dev/null 2>&1; then + nohup /shared/aaji/tools/victoriametrics/victoria-metrics-prod \ + -storageDataPath="$PERF_RUN/omnistat-db" \ + -httpListenAddr=127.0.0.1:$PORT \ + -retentionPeriod=100y -fs.disableMmap \ + > "$PERF_RUN/vm_for_foms.log" 2>&1 & + echo $! > "$PERF_RUN/vm_for_foms.pid" + for i in {1..20}; do sleep 1; curl -sf "http://127.0.0.1:$PORT/api/v1/status/tsdb" > /dev/null && break; done +fi +TSDB_URL="http://127.0.0.1:$PORT" +``` + +For each FOM, run the indicated PromQL via `curl ${TSDB_URL}/api/v1/query` and parse the scalar. + +| FOM | PromQL (substitute $JOBID) | +|---|---| +| `mfma_tflops` (per-card avg) | `avg(rate(SQ_INSTS_VALU_MFMA_MOPS_F64[10s]) * on(instance) group_left() max by(instance)(rmsjob_info{jobid="$JOBID"})) * 4 / 1e12` | +| `mean_power_W` (per-card mean) | `avg_over_time((rocm_avg_pwr * on(instance) group_left() max by(instance)(rmsjob_info{jobid="$JOBID"}))[${RUNTIME}s:])` | +| `energy_J` (total, all 16 cards) | `sum(avg_over_time((rocm_avg_pwr * on(instance) group_left() max by(instance)(rmsjob_info{jobid="$JOBID"}))[${RUNTIME}s:])) * ${RUNTIME}` | + +The MFMA→TFLOP/s conversion uses the standard MFMA MOPS-to-FLOPS factor of 4 (each MFMA op is 4 FMAs). This matches the formula used by [omnistat_verifier.md](../perf-analysis/agents/omnistat_verifier.md) and the existing 6985 combined_report. + +If a query returns `null` or empty data (rocprofiler counters were off, OR no rmsjob_info records exist), write `null` for that FOM and add an entry to `foms.json`'s `notes` field: `"mfma_tflops=null because rocprofiler appears disabled (check sbatch banner for enable_rocprofiler state)"`. + +### 4. TraceLens↔Omnistat 1-second-window correlation + +This is the headline "cheaper than kernel-trace" capability. Build `kernel_correlation.csv` as follows: + +#### 4a. Parse the rank-0 trace into a flat kernel-event table + +```python +import gzip, json, os, pathlib +def _open(p): + return gzip.open(p, 'rt') if p.endswith('.gz') else open(p) + +with _open(TRACE) as f: + trace = json.load(f) + +events = trace.get("traceEvents", []) +# Keep only GPU kernel dispatches; PyTorch kineto marks them ph='X' with cat=='kernel' +kernels = [(e['ts'], e['dur'], e['name']) + for e in events + if e.get('ph') == 'X' and e.get('cat') == 'kernel' and 'ts' in e and 'dur' in e] +# kineto ts is in microseconds from trace start; convert to UTC by adding trace's reference epoch +ref_us = int(trace.get('baseTimeNanoseconds', 0)) // 1000 # may be missing on older kineto +# Fallback: pull job_info.json from omnistat_analyst output; its start_time is UTC ISO +``` + +If `ref_us` is missing or zero, read `omnistat/inspect_outputs/job_info.json` (written by omnistat_analyst) and use its `start_time` as the trace reference. **Important:** that timestamp is naive UTC; convert with `calendar.timegm(time.strptime(s, "%Y-%m-%dT%H:%M:%S"))`, NOT `time.mktime` (see SKILL §11.7). The PyTorch trace's `displayTimeUnit` and per-event `ts` are microseconds since the kineto "profile start" anchor, which is typically the wallclock time the profiler armed — this is close-enough to the job's epoch start that 1-s windowing absorbs any drift. + +#### 4b. Bucket kernels into 1-second windows (rank-0 view of GPU 0; we approximate "card" with rank-0's device) + +```python +from collections import defaultdict, Counter +window_kernels = defaultdict(Counter) # window_ts -> {kernel_name: total_dur_us} +for ts, dur, name in kernels: + win = (ref_us + ts) // 1_000_000 # 1-s windows, unix epoch seconds + window_kernels[win][name] += dur +``` + +For each window, compute `top_kernel_name`, `top_kernel_busy_frac = top_dur / 1_000_000`, `second_kernel_name`, `second_kernel_busy_frac` (cap fraction at 1.0). + +#### 4c. Query Omnistat for per-window counter rates + +For each window timestamp `win_ts` and each of the 16 cards: + +```python +import requests +def q(promql, t): + r = requests.get(f"{TSDB_URL}/api/v1/query", + params={"query": promql, "time": t}, timeout=10) + r.raise_for_status() + res = r.json()["data"]["result"] + return {row["metric"].get("instance"): float(row["value"][1]) for row in res} + +# All 16 instances at once via the join +mfma_q = (f'rate(SQ_INSTS_VALU_MFMA_MOPS_F64[10s]) * on(instance) group_left() ' + f'max by(instance)(rmsjob_info{{jobid="{JOBID}"}}) * 4 / 1e12') +hbm_q = (f'rate(FETCH_SIZE[10s]) * on(instance) group_left() ' + f'max by(instance)(rmsjob_info{{jobid="{JOBID}"}}) / 1024 / 1024 / 1024 * 1024') +pwr_q = (f'rocm_avg_pwr * on(instance) group_left() ' + f'max by(instance)(rmsjob_info{{jobid="{JOBID}"}})') + +# Note: FETCH_SIZE units are KB (lesson R1, see SKILL §12 / R1 report); convert to GB/s by +# multiplying KB-per-second by 1024 then dividing by 1e9 -> simplifies to "* 1024 / 1e9". +``` + +For each `win_ts` in `sorted(window_kernels)`, query the 3 PromQL above. For each `(instance, gpu_id)` pair returned (omnistat tags each per-card metric with both `instance` (hostname) and `card` index), emit one row to `kernel_correlation.csv`. + +**Performance:** even at 5-min runs that's ~300 windows × 3 queries = 900 PromQL hits. With `requests` keep-alive + `-search.disableCache` already off (we want cached scans here), this completes in 1-2 minutes. Cache responses with `functools.lru_cache` keyed on (promql, win_ts) to avoid duplicate queries within the same window. + +#### 4d. Compute correlation summary + +Aggregate to populate `kernel_correlation_summary` in `foms.json`: +- `windows_examined = len(window_kernels)` +- `windows_with_top_kernel_busy_frac_gt_0p5 = sum(1 for w in windows if top_busy_frac > 0.5)` +- `top_kernel_dominants` = `Counter(top_kernel_name for win in windows).most_common(5)` (paired with count) +- `median_mfma_tflops_when_aten_mm_dominant` = median MFMA TFLOP/s across rows where `top_kernel_name.startswith('aten::mm')` (or contains `gemm` for hipBLASLt internal names) +- `median_mfma_tflops_when_aten_bmm_dominant` = same for `aten::bmm` and Triton bmm names +- `median_hbm_read_GBps_when_mul_dominant` = same for `aten::mul`/elementwise patterns +- `attribution_quality` per the rule above. + +### 5. Emit `foms.json` + +Write `foms.json` and `kernel_correlation.csv`. Both atomically via tmpfile + rename. + +### 6. Final stdout + +``` +STATUS=ok; reason=foms epoch_time=s throughput=/s mfma=TFLOP/s correlation= +``` + +If trace is missing OR counter queries returned all nulls: + +``` +STATUS=partial; reason= +``` + +The orchestrator handles `partial` by recording the FOM row with `null`s and a note; the loop continues. + +## Failure modes + +| Failure | Action | +|---|---| +| `_convergence_foms.json` empty (job died before any epoch) | `epoch_time_s = null, final_loss = null`; STATUS=partial | +| `manifest.json` missing | STATUS=fail; orchestrator handles as analysis_fail | +| VictoriaMetrics won't start | log to vm_for_foms.log; emit STATUS=partial with counter FOMs null; correlation skipped | +| `rocm_avg_pwr` empty (power telemetry collector dropped data) | `energy_J = null, mean_power_W = null, energy_per_sample_J = null`; STATUS=partial | +| Trace too large to load (>4 GB JSON) | skip correlation, `attribution_quality = "skipped:trace_too_large"`; STATUS=partial | +| PyTorch kineto trace base epoch missing AND job_info.json missing | skip correlation (cannot align timestamps); STATUS=partial | + +## Notes for the implementing agent + +- This subagent runs on the login node, after omnistat_analyst has started VictoriaMetrics on this perf-run. If VM isn't already up, start a second instance on a different port — the omnistat datadir is read-only friendly under VM. +- Do NOT call `omnistat-inspect`; that's the analyst's job. We go straight to PromQL here because we need finer-grained queries than what `omnistat-inspect` exposes. +- Read at most: manifest.json, hydragnn-train-.out, the rank-0 trace, omnistat-db (via VM). Cap trace read at 4 GB. +- Write at most: foms.json, kernel_correlation.csv, vm_for_foms.{log,pid}. + +## parse_convergence.py extension + +The existing `examples/parse_convergence.py` is extended with `--json-foms ` that emits this aggregate: + +```json +{ + "num_epochs_completed": 3, + "mean_epoch_time_excluding_epoch_0": 145.3, + "epoch_times_s": [156.7, 144.1, 146.5], + "final_train_loss": 5.21, + "final_val_loss": 5.08 +} +``` + +`mean_epoch_time_excluding_epoch_0` is `mean(epoch_times_s[1:])` to drop the warm-up. If only 1 epoch ran, fall back to using it (and add `"warning": "single epoch, epoch_0 not dropped"` to the JSON). + +`epoch_times_s` comes from the `epoch_summary` records' `wall_time_s` when at least 2 are populated (rare — HydraGNN typically attaches the timer only to the final epoch). Otherwise computed per-epoch from `tqdm_final` records as `rate_s × batches_per_epoch`. `samples_processed` is computed by the fom_extractor from the manifest (`num_epochs × max_num_batch × batch_size × ranks`), NOT from the log, because HydraGNN's rank-0 only emits one "Max memory allocated" line per job. diff --git a/material_science/models/HydraGNN/recipes/perf-optimizer-loop/agents/lever_picker.md b/material_science/models/HydraGNN/recipes/perf-optimizer-loop/agents/lever_picker.md new file mode 100644 index 0000000..9b82c3b --- /dev/null +++ b/material_science/models/HydraGNN/recipes/perf-optimizer-loop/agents/lever_picker.md @@ -0,0 +1,120 @@ +# lever_picker subagent + +Proposes the SINGLE next lever for the orchestrator to apply, given the loop +history so far. Emits a JSON file matching the catalog entry schema; STATUS=ok +on success or STATUS=partial reason=catalog_exhausted when no candidates remain. + +## Inputs + +- `loop-/foms.csv` — every iter so far, including accepted/rejected. +- `loop-/do_not_retry.json` — array of lever ids already tried-and-rejected. +- `/combined_report.md` — the most recent accepted iter's synthesizer report. This is the primary evidence source for lever ranking. +- `/foms.json` — quantitative FOM values for the current best. +- `/kernel_correlation.csv` — 1-s window TraceLens↔Omnistat join (from fom_extractor). Used to disambiguate compute-bound vs memory-bound vs dispatch-bound. +- [`../lever_catalog.yaml`](../lever_catalog.yaml) — full catalog. +- (Optional) `/combined_report.md` if previous iter was a `diagnostic_only` lever (e.g. `kernel_trace_diag_only`). + +## Outputs + +- `loop-/iter--lever.json` — single object matching the catalog schema, with two extra fields: + - `picked_by`: "lever_picker" + - `reason`: short string (1-2 sentences) explaining the choice. +- Final stdout: `STATUS=ok; reason=lever=` OR `STATUS=partial; reason=catalog_exhausted`. + +## Decision algorithm + +### 1. Build candidate set + +Read catalog; drop: +- `baseline` (iter-0 only; never picked again). +- Any lever id in `do_not_retry.json`. +- Any lever id already accepted in `foms.csv` (we don't re-pick already-applied levers; the new "best" already has them baked in). +- `kernel_trace_diag_only` — see special gate below. + +If the candidate set is empty: emit `STATUS=partial; reason=catalog_exhausted`. + +### 2. Score remaining candidates + +For each candidate, compute a heuristic score combining: + +- **`expected_payoff` from catalog**: high=3, medium=2, low=1. +- **Evidence support from `combined_report.md`**: +2 if the report's "Next steps" or "Remedies proposed but not tried" explicitly names this lever (e.g. "torch.compile" → matches `torch_compile_e3nn`). +1 if it implicitly relates (e.g. "still dispatch-bound" → +1 for `batch_800`, `torch_compile_e3nn`). +- **Bottleneck-class alignment**: parse the report's top-3 bottleneck classes (`cpu_dispatch`, `gpu_compute`, `comm_*`, `dataloader`, ...). Match against the lever's expected target: + - `cpu_dispatch` → `torch_compile_e3nn`, `batch_800` (+2 each) + - `gpu_compute` + low TFLOP/s → `tunable_op`, `precision_fp32` (+2 each) + - `gpu_memory_hbm` saturated → no lever directly helps; +0 + - `comm_xgmi` / `comm_scaleout` → `rccl_high_priority`, `nccl_minchannels` (+1 each at N=2; +2 if N≥4) + - `dataloader` → `num_workers_12` (+2) +- **Risk penalty**: low=−0, medium=−1, high=−2. +- **Cardinality penalty**: −5 if it would be the second `diagnostic_only` lever in the same loop (we cap diagnostic iters at 1 per loop). + +Top-scoring lever wins. Ties broken by lower `catalog_order`. + +### 3. Special gate: `kernel_trace_diag_only` + +Only pick `kernel_trace_diag_only` when ALL of: +1. The `combined_report.md` contains language indicating per-kernel attribution is unclear (e.g. "cannot attribute", "fused kernel", "unknown", "see kernel trace"), OR all other improvement-levers in the catalog are in `do_not_retry.json`. +2. AND the most recent `kernel_correlation.csv` has been examined and is insufficient (e.g. `top_kernel_busy_frac` < 0.4 in >50% of steady-state windows, indicating no single kernel dominates the 1-s windows). +3. AND no previous iter in the loop already used `kernel_trace_diag_only` (1-per-loop cap). + +When picked, the orchestrator runs the iter, gets a `combined_report.md` with per-kernel-name evidence, and on the NEXT iter `lever_picker` uses that to choose a real improvement lever (likely `torch_compile_e3nn` if a specific fused kernel is dispatch-bound, or `tunable_op` if a specific GEMM shape is under-tuned). + +### 4. fp32 precision gate + +Only pick `precision_fp32` when the current_best `foms.json` shows the workload is compute-bound, defined as: +- `mfma_tflops > 0.3 * 78.6` (i.e. > 24 TFLOP/s per card, meaning fp64 MFMA is doing real work), OR +- `combined_report.md` Executive summary names `gpu_compute` as the #1 bottleneck. + +Rationale: if the workload is still dispatch-bound (R2 state: mean util 12%, fp64 0.87% of peak), switching fp64→fp32 doubles peak FLOP/s but the dispatch ceiling stays the same. Save the lever for after `batch_800` and `torch_compile_e3nn` have lifted the dispatch ceiling enough to expose compute as the bottleneck. + +### 5. Emit JSON + +Write `loop-/iter--lever.json` as a single object copying the catalog entry verbatim, with two added fields: + +```json +{ + "id": "batch_800", + "description": "...", + "catalog_order": 1, + "kind": "env_only", + "env_vars": {"HG_BATCH_SIZE": "800"}, + "diagnostic_only": false, + "expected_payoff": "high", + "risk": "low", + "revert_method": "drop_env", + "citation": "...", + "notes": "...", + "picked_by": "lever_picker", + "reason": "R2 combined_report names HG_BATCH_SIZE=800 as #1 next-lever; current_best is still dispatch-bound (mean util 12.65%); VRAM headroom 24/288 GB." +} +``` + +### 6. Final stdout + +``` +STATUS=ok; reason=lever= +``` + +or + +``` +STATUS=partial; reason=catalog_exhausted +``` + +## Hallucination guardrails + +- **Never invent a lever.** If you think there's a missing knob worth trying, write your reasoning to `loop-/lever_proposal-.md` and emit `STATUS=partial; reason=catalog_proposal`. The orchestrator will treat this as catalog_exhausted (since you didn't pick from the catalog) and the user/maintainer can review the proposal next morning to add it to the catalog. +- **Cite evidence by path.** Every `reason` field must point at a specific `combined_report.md` line or `foms.json` field. No vague justifications. +- **Single lever, single change.** The JSON output's `env_vars` field copies from the catalog verbatim — do NOT merge in additional env vars from other levers, even if they look complementary. Multi-knob changes are tracked across iterations, not within one. + +## Failure modes + +- Empty candidate set → STATUS=partial; reason=catalog_exhausted (terminal for the loop). +- Cannot read previous combined_report.md → STATUS=partial; reason=missing_report; pick lever purely by catalog_order from candidates. +- Both above conditions → STATUS=fail; reason=no_state_no_catalog. Orchestrator must log and exit. + +## Notes for the implementing agent + +- This subagent does NOT submit jobs or read raw traces. It is a pure-reasoning subagent over the existing reports. +- It MAY consult ROCm blog URLs from the catalog's `citation` field if the reason needs grounding, but should not fetch new pages — the reasoning lives in the catalog already. +- Reads at most 5 files (catalog + foms.csv + do_not_retry + combined_report + foms.json + kernel_correlation.csv). Bounded. diff --git a/material_science/models/HydraGNN/recipes/perf-optimizer-loop/agents/orchestrator.md b/material_science/models/HydraGNN/recipes/perf-optimizer-loop/agents/orchestrator.md new file mode 100644 index 0000000..7d14e40 --- /dev/null +++ b/material_science/models/HydraGNN/recipes/perf-optimizer-loop/agents/orchestrator.md @@ -0,0 +1,218 @@ +# orchestrator subagent + +Drives the iterative HydraGNN optimizer loop. Dispatches all other subagents +(launcher, lever_picker, fom_extractor, analyst×2, verifier×2, synthesizer, +story_writer) and owns accept/revert + STATUS.txt + STOP-flag handling. + +This subagent is invoked once per loop (typically via Claude Code CLI in tmux +on the login node, or directly inside a Cursor agent session). + +## Inputs + +- Loop arguments (passed in the user message): ``, ``. +- Repo at `/home/aaji/git/ai4science-studio` on branch `aaji/perf-optimizer-loop-hydragnn`. +- Cluster config at `.cluster-config.yaml`. +- [`lever_catalog.yaml`](../lever_catalog.yaml) — single source of truth for allowed levers. + +## Outputs + +Per loop, under `/shared/aaji/models/HydraGNN/perf-runs/loop-/`: + +- `STATUS.txt` — append-only event log; one line per event (see schema below). +- `foms.csv` — one row per iteration: `iter,jobid,lever_id,env_diff,accepted,epoch_time_s,throughput_samples_per_s,mfma_tflops,energy_J,mean_power_W,energy_per_sample_J,final_loss,primary_fom_delta_pct,du_loop_dir,notes`. +- `do_not_retry.json` — list of lever ids that regressed; consulted by lever_picker. +- `iter-N-.json` — symlink to that iter's `manifest.json`. +- `story.md` — written by story_writer at end of loop. +- `foms.png` — written by story_writer at end of loop. + +Per iteration, all artifacts already written by the existing `recipes/perf-analysis/` subagents land under `/shared/aaji/models/HydraGNN/perf-runs//` plus: + +- `foms.json` — written by fom_extractor. +- `kernel_correlation.csv` — written by fom_extractor. + +## Hard constraints + +1. **Never run more than one SLURM job concurrently.** Wait for the previous job to be in a terminal state before submitting the next. +2. **Exactly one lever per iteration.** Multi-variable changes break delta attribution. +3. **Respect the STOP flag.** Check `loop-/STOP` before every sbatch submission, before every analyst phase, and before every iteration decision. If present, write `LOOP_ABORT reason=stop_flag`, finish processing the currently-running iter if any, then exit 0. +4. **Single source of truth for state is on disk.** `foms.csv`, `do_not_retry.json`, and STATUS.txt are the only state — no in-memory state survives a restart. The orchestrator must be able to resume mid-loop by reading these files. +5. **Never modify files outside `/shared/aaji/models/HydraGNN/perf-runs/` and `/tmp/`.** No repo edits during the loop. (Repo edits to capture lessons happen AFTER the loop, by a separate pass.) +6. **All subagent invocations write structured JSON to disk and end with `STATUS=ok|partial|fail`.** The orchestrator parses the final line; never inline-trusts chat content. + +## Loop control parameters + +| Param | Default | Source | +|---|---|---| +| Max iterations | 5 | user argument | +| Per-iter wall budget | 30 min | sbatch `--time=00:30:00` | +| Per-iter polling | every 30 s | `sacct -j -X -n` | +| Polling ceiling | 45 min | hard exit if not terminal by then | +| LLM retry | 3 attempts with exp backoff | Only for transient subagent failures | +| Early-stop A | 2 consecutive accepts with <5% primary-FOM improvement | computed from foms.csv | +| Early-stop B | catalog exhausted (every lever tried or in do_not_retry.json) | from lever_picker | +| Auto-reject A | epoch_time_s regressed vs current best (any amount) | computed in step 5 below | +| Auto-reject B | final_loss > 1.5 × baseline_final_loss | computed in step 5 below | + +## STATUS.txt event schema + +One line per event. Format: ` `. Examples: + +``` +2026-05-22T22:14:03Z LOOP_START uuid= n_iters_budget=5 driver=claude-code-cli +2026-05-22T22:14:09Z PREFLIGHT_OK cluster=lux disk_free=26T claude_cli=ok api_egress=ok +2026-05-22T22:14:11Z LEVER_PICK iter=0 lever=baseline reason="establishing fresh baseline" +2026-05-22T22:15:11Z ITER_SUBMIT iter=0 lever=baseline jobid=7050 +2026-05-22T22:23:48Z ITER_COMPLETE iter=0 jobid=7050 state=COMPLETED runtime=487s +2026-05-22T22:24:02Z ANALYZE_START iter=0 jobid=7050 +2026-05-22T22:26:11Z ANALYZE_DONE iter=0 fom_epoch_time=145.3 fom_throughput=11020 fom_mfma_tflops=0.30 fom_energy_J=1.1e6 fom_loss=5.8 +2026-05-22T22:26:12Z ITER_DECISION iter=0 accepted=true note=baseline +2026-05-22T22:26:14Z LEVER_PICK iter=1 lever=batch_800 reason="R2 #1 recommendation" +2026-05-22T22:34:55Z ITER_DECISION iter=1 accepted=true delta_pct=-12.4 +2026-05-22T22:34:56Z DISK_USAGE loop_dir=270M +2026-05-22T23:01:13Z LEVER_PICK iter=2 lever=torch_compile_e3nn reason="..." +2026-05-22T23:40:01Z ITER_DECISION iter=2 accepted=false delta_pct=+4.2 reason=regression_revert +2026-05-22T23:40:02Z DO_NOT_RETRY add=torch_compile_e3nn +2026-05-23T05:47:33Z LOOP_COMPLETE n_iters_done=5 best_iter=1 best_lever=batch_800 best_fom_epoch_time=127.2 +``` + +## Steps the orchestrator performs + +### 0. Resume detection + +Read `loop-/foms.csv` if it exists. If iter rows are present, set `current_iter = max(iter) + 1` and `current_best` = row with min `epoch_time_s` among `accepted=true`. Log `LOOP_RESUME current_iter=` to STATUS.txt. Otherwise initialize fresh. + +### 1. Pre-flight (only on fresh start, NOT on resume) + +Run these checks in sequence; abort with `PREFLIGHT_FAIL reason=` on the first failure: + +- `sinfo -p lux,rad -h` → if zero `IDLE` or `MIX` nodes, abort. Cluster may be in maintenance (see SKILL §14). +- `df -h /shared | tail -1` → abort if `Use%` > 95. +- `test -x /shared/aaji/tools/omnistat-pr271/bin/omnistat-usermode` → abort if missing. +- `test -x /shared/aaji/tools/victoriametrics/victoria-metrics-prod` → abort if missing. +- `test -f /shared/aaji/models/HydraGNN/overlays/hydragnn-overlay.img` → abort if missing. +- `test -f /shared/aaji/images/pytorch_rocm7.2.2_ubuntu24.04_py3.12_pytorch_release_2.10.0.sif` → abort if missing. +- If invoked via Claude Code CLI on the login node: `[[ -n "$ANTHROPIC_API_KEY" ]]` and `curl -fsS https://api.anthropic.com/v1/models > /dev/null` (5 s timeout) → abort if either fails. + +Write `PREFLIGHT_OK` line. + +### 2. Per-iteration loop + +For `current_iter` in `0 .. n_iters_budget-1`: + +**2a. STOP-flag gate.** `[[ -f loop-/STOP ]] && { log LOOP_ABORT reason=stop_flag; exit 0; }` + +**2b. Pick lever.** If `current_iter == 0`, lever is hard-coded to `baseline`. Otherwise dispatch [`lever_picker`](lever_picker.md) subagent with inputs: +- `loop-/foms.csv` +- `loop-/do_not_retry.json` +- previous iter's `/combined_report.md` (most recent accepted iter) +- [`lever_catalog.yaml`](../lever_catalog.yaml) + +lever_picker writes `loop-/iter--lever.json` with the chosen lever (single object matching the catalog schema). The orchestrator reads it. If lever_picker returns STATUS=partial with reason="catalog_exhausted", emit `LOOP_COMPLETE reason=catalog_exhausted` and proceed to story_writer. + +Log `LEVER_PICK iter= lever= reason=""`. + +**2c. STOP-flag gate.** Re-check; act as in 2a. + +**2d. Submit sbatch.** Build env-var diff: union of the current-best's env vars + the new lever's `env_vars`, minus any vars revert-method'd from rejected iters. Render an `iter--env.sh` file with all overrides (one `export KEY=VALUE` per line) under `loop-/`. Submit: + +```bash +( set -a; source "loop-/iter--env.sh"; set +a; \ + cd /home/aaji/git/ai4science-studio; \ + sbatch material_science/models/HydraGNN/examples/sbatch_train_perf_amd.sh ) +``` + +Capture the jobid from sbatch stdout. Log `ITER_SUBMIT iter= lever= jobid=`. + +If the lever is `kind: rank_script_patch`, write the patch content from `lever_catalog.yaml` to `loop-/iter--hook.py` and export `HG_RANK_PRE_TRAIN_HOOK=` before sbatch. (The rank script must be extended to honor this var; see the rank-script extension note below.) + +**2e. Poll until terminal.** Every 30 s: `sacct -j -X -n --format=State,ExitCode,Elapsed,NodeList -P`. Terminal states: COMPLETED, FAILED, TIMEOUT, CANCELLED, NODE_FAIL, OUT_OF_MEMORY. Hard exit if not terminal by 45 min; log `ITER_TIMEOUT iter= jobid=`, treat as auto-reject (do_not_retry) and continue to next iter with current-best restored. + +Log `ITER_COMPLETE iter= jobid= state= runtime=`. + +**2f. STOP-flag gate.** Re-check before the analysis phase. If set, still analyze the just-completed iter (we paid for it), then exit gracefully. + +**2g. Run analyst+verifier+synth.** Re-use the existing flow from `recipes/perf-analysis/`. Dispatch in parallel (single message with multiple Task tool calls): + +- `Task("tracelens_analyst", prompt=read("../perf-analysis/agents/tracelens_analyst.md"), manifest=/manifest.json)` +- `Task("omnistat_analyst", prompt=read("../perf-analysis/agents/omnistat_analyst.md"), manifest=/manifest.json)` + +Wait for both. Then dispatch the two verifiers in parallel: + +- `Task("tracelens_verifier", prompt=read("../perf-analysis/agents/tracelens_verifier.md"), ...)` +- `Task("omnistat_verifier", prompt=read("../perf-analysis/agents/omnistat_verifier.md"), ...)` + +Wait for both. Then dispatch synthesizer: + +- `Task("synthesizer", prompt=read("../perf-analysis/agents/synthesizer.md"), ...)` + +Then dispatch [`fom_extractor`](fom_extractor.md), which writes `/foms.json` and `/kernel_correlation.csv`. + +**Failure handling:** if any analyst or verifier returns `STATUS=fail`, the orchestrator MUST NOT mark the iteration as accepted. Log the failure, treat as `accepted=false note=analysis_fail`, do NOT add the lever to do_not_retry (the lever was not actually evaluated; subagent was the failure). Try the same lever once more on the next iteration; if it fails again, then add to do_not_retry with `reason=repeated_analysis_fail`. + +Log `ANALYZE_DONE iter= fom_epoch_time= fom_throughput= fom_mfma_tflops= fom_energy_J= fom_loss=`. + +**2h. Accept / revert.** + +- Read `/foms.json`. +- If this iter is `baseline` (iter 0): `accepted=true`; current_best = this iter. Skip the rest of the accept/revert logic. +- If the lever has `diagnostic_only: true`: `accepted=diagnostic`; do NOT update current_best, do NOT add to do_not_retry. Pass the iter's combined_report.md to lever_picker as fresh evidence on the NEXT iter (orchestrator passes it as an explicit input). +- Else compute `delta_pct = (epoch_time_s - current_best.epoch_time_s) / current_best.epoch_time_s * 100`. + - If `epoch_time_s > current_best.epoch_time_s` (any regression): `accepted=false reason=regression`. Append lever_id to `do_not_retry.json`. current_best unchanged. + - If `final_loss > 1.5 * baseline_final_loss`: `accepted=false reason=loss_blow_up`. Append lever_id to `do_not_retry.json`. + - Else: `accepted=true`. Update current_best = this iter. + +Append the iter's row to `foms.csv`. Log `ITER_DECISION iter= accepted= delta_pct= [reason=]`. If rejected, log `DO_NOT_RETRY add=`. + +**2i. Disk usage log.** `du -sh loop-/ | awk '{print $1}'` → log `DISK_USAGE loop_dir=`. + +**2j. Early-stop check.** + +- If last 2 accepted iters each had `|delta_pct| < 5`: log `LOOP_COMPLETE reason=converged`, break. +- If `do_not_retry.json` covers every catalog entry (minus already-tried-and-accepted): log `LOOP_COMPLETE reason=catalog_exhausted`, break. + +### 3. Story phase + +Dispatch [`story_writer`](story_writer.md) with inputs: `loop-/foms.csv`, every per-iter `combined_report.md`, every per-iter `foms.json`. It writes `story.md` and `foms.png`. + +Log `LOOP_COMPLETE n_iters_done= best_iter= best_lever= best_fom_epoch_time=`. + +### 4. Final stdout (to whoever invoked the orchestrator) + +``` +STATUS=ok; reason=loop_complete uuid= best_iter= best_lever= best_epoch_time=s vs_baseline_pct=% +``` + +## Rank script extension (one-time, must land before iter-0) + +The existing `sbatch_train_perf_amd.sh` rank script wraps the HydraGNN entrypoint via `runpy.run_path`. To honor `HG_RANK_PRE_TRAIN_HOOK` (used by `rank_script_patch` levers), add this stanza inside the rank script's main python -c block, immediately before the `runpy.run_path(...)` call: + +```python +_hook = os.environ.get('HG_RANK_PRE_TRAIN_HOOK', '') +if _hook and os.path.isfile(_hook): + if int(os.environ.get('SLURM_PROCID', '0')) == 0: + print(f'[rank_hook] executing pre-train hook: {_hook}', flush=True) + with open(_hook) as _f: + exec(_f.read(), {'__name__': '__hook__'}) +``` + +This is a single-edit, one-time change to the existing rank script. The orchestrator does NOT edit the rank script during the loop. The orchestrator-author (the agent executing this plan) makes this edit ONCE before running the first loop. + +## Failure modes + +| Failure | Action | +|---|---| +| `sinfo` shows cluster drained | PREFLIGHT_FAIL reason=cluster_down; exit 2 | +| sbatch returns nonzero | log ITER_SBATCH_FAIL; treat as auto-reject; continue | +| sacct never reaches terminal in 45 min | log ITER_TIMEOUT; auto-reject; continue | +| analyst/verifier returns STATUS=fail | log ANALYZE_FAIL; retry next iter; if 2 in a row, do_not_retry the lever with reason=repeated_analysis_fail | +| fom_extractor returns STATUS=fail | treat as accepted=false reason=fom_extraction_fail; do NOT add to do_not_retry; orchestrator should investigate manually next morning | +| LLM API failure (transient) | retry up to 3 times with exp backoff; if all 3 fail, log LOOP_ABORT reason=llm_api_failed, exit 1 | +| STOP flag appears | finish current iter analysis if mid-stream, write LOOP_ABORT reason=stop_flag, exit 0 | +| ANTHROPIC_API_KEY missing | PREFLIGHT_FAIL reason=no_api_key; exit 2 | + +## Notes for the implementing agent + +- All paths must be absolute. +- All file writes to `loop-/STATUS.txt` use `flock -x` to be safe against concurrent appenders (the orchestrator is single-threaded but multiple subagents may want to log). +- The orchestrator never dispatches more than one launcher subagent in parallel. Analyst pair and verifier pair are dispatched in parallel within their own phases. +- The orchestrator may run as a Cursor agent OR as a Claude Code CLI session — same prompt, same files. The user's `run_optimizer_loop.sh` invokes via Claude Code CLI for the overnight unattended case. diff --git a/material_science/models/HydraGNN/recipes/perf-optimizer-loop/agents/story_writer.md b/material_science/models/HydraGNN/recipes/perf-optimizer-loop/agents/story_writer.md new file mode 100644 index 0000000..9182203 --- /dev/null +++ b/material_science/models/HydraGNN/recipes/perf-optimizer-loop/agents/story_writer.md @@ -0,0 +1,182 @@ +# story_writer subagent + +Produces the final progression narrative + a matplotlib FOM-vs-iteration chart +once the loop has completed (or aborted). Reads only files; writes only +`story.md` and `foms.png` under the loop directory. + +## Inputs + +- `loop-/foms.csv` — every iteration's row, in chronological order +- `loop-/STATUS.txt` — full event log +- `loop-/do_not_retry.json` — rejected lever ids +- `loop-/iter-*-lever.json` — per-iter chosen lever (the `reason` field is the lever_picker's justification) +- `/combined_report.md` for every iter that completed (resolved via symlinks `loop-/iter-N-.json`) +- `/foms.json` for every iter that completed +- `/kernel_correlation.csv` for every iter that has it (best-iter likely; may be missing if fom_extractor returned partial) +- [`../lever_catalog.yaml`](../lever_catalog.yaml) — for citation URLs and lever metadata +- [`../README.md`](../README.md) — for the references list to mirror at the end of story.md + +## Outputs + +- `loop-/story.md` — the progression narrative +- `loop-/foms.png` — 2-panel matplotlib chart +- Final stdout: `STATUS=ok; reason=story written, n_iters=, best=iter-=` or `STATUS=partial; reason=` + +## `story.md` layout + +```markdown +# HydraGNN iterative-sysopt-loop — + +**Hardware:** 2 nodes × 8 GPUs MI355X (gfx950, vultr_lux). +**Image:** `pytorch_rocm7.2.2_ubuntu24.04_py3.12_pytorch_release_2.10.0`. +**Loop:** iterations, started , ended , total wall hm. +**Primary FOM:** `epoch_time_s` (mean over epoch >= 1, drops warm-up). + +## TL;DR + +| Iter | Lever | epoch_time_s | Δ vs prev best | Accepted? | Notes | +|---|---|---|---|---|---| +| 0 | baseline | | n/a | yes | | +| 1 | batch_800 | | | | | +| ... | | | | | | + +Best result: **iter- ``** at `` s/epoch (% vs baseline; % vs R2 from `recipes/perf-analysis/`). +. + +## Initial state of the system (iter-0 baseline) + +<3-4 sentences: what the analysts found about the baseline. Cite numbers from +iter-0/combined_report.md — top bottleneck class, mean util, MFMA TFLOP/s, +exposed NCCL share, dominant kernels. End with a sentence naming the loop's +working hypothesis (e.g. "still dispatch-bound, expect batch_800 to lift it").> + +## Iteration-by-iteration progression + +For each iter (i from 1 to N): + +### Iter : `` + +**Why this lever was chosen:** . + +**What changed:** env-var diff or rank-script patch from lever_catalog.yaml. + +**Result (vs previous best):** +- epoch_time_s: → (Δ %) +- throughput_samples_per_s: → (Δ %) +- mfma_tflops: → +- energy_J / mean_power_W / energy_per_sample_J: as above +- final_loss: → (control) + +**Bottleneck shift (from /combined_report.md):** +<2-3 sentences. What was the #1 bottleneck before vs after?> + +**Verdict:** ACCEPTED (current best updated) / REJECTED (auto-reverted; added to do_not_retry.json). + +**Citation:** []() — from lever_catalog.yaml. + +--- + +## Final state of the system + +<3-4 sentences: where is the workload now (top bottleneck, MFMA % of peak, mean util, energy per sample). Compare to the initial state and the historical R2 (`recipes/perf-analysis/`) baseline. End with what would be the next-best lever IF we extended the catalog.> + +## FOM progression chart + +![FOM progression across iterations](foms.png) + +Top panel: `epoch_time_s` (primary, lower is better) and `energy_per_sample_J` (lower is better, twin axis). +Bottom panel: `throughput_samples_per_s` and `mfma_tflops` (both higher is better). +Rejected iterations marked with red x. Diagnostic iterations marked with grey diamond. + +## Levers rejected (auto-revert) + +| Iter | Lever | epoch_time_s | regression % | Likely reason | +|---|---|---|---|---| +| ... | ... | ... | ... | | + +## Levers in catalog not tried + +| Lever | Why not tried | +|---|---| +| | <"loop ended before its turn" / "in do_not_retry from a prior loop" / "diagnostic-only, gate not met"> | + +## What would be next (outside this loop) + +<2-3 sentences naming the most promising direction not covered by the current catalog.> + +## References + + + +--- +Generated by [`material_science/models/HydraGNN/recipes/perf-optimizer-loop/agents/story_writer.md`](.). +``` + +## `foms.png` layout (matplotlib) + +```python +import matplotlib +matplotlib.use("Agg") +import matplotlib.pyplot as plt + +fig, (ax_top, ax_bot) = plt.subplots(2, 1, figsize=(10, 8), sharex=True) + +# Top panel: epoch_time_s + energy_per_sample_J on twin axis +ax_top.set_xlabel("Iteration") +ax_top.set_ylabel("epoch_time_s (lower = better)", color="C0") +ax_top.plot(iters_x, epoch_times, "C0-o", label="epoch_time_s") +ax_top_t = ax_top.twinx() +ax_top_t.set_ylabel("energy_per_sample_J (lower = better)", color="C1") +ax_top_t.plot(iters_x, energy_per_sample, "C1-s", label="energy_per_sample_J") + +# Mark rejected with red x at the epoch_time location +for x, et, acc in zip(iters_x, epoch_times, accepted_flags): + if acc == "false": + ax_top.plot(x, et, "rx", markersize=14, markeredgewidth=3) + elif acc == "diagnostic": + ax_top.plot(x, et, "D", color="grey", markersize=10) + +# Bottom panel: throughput + MFMA TFLOP/s +ax_bot.set_xlabel("Iteration") +ax_bot.set_ylabel("throughput_samples_per_s (higher = better)", color="C2") +ax_bot.plot(iters_x, throughput, "C2-o", label="throughput") +ax_bot_t = ax_bot.twinx() +ax_bot_t.set_ylabel("mfma_tflops per card (higher = better)", color="C3") +ax_bot_t.plot(iters_x, mfma, "C3-s", label="mfma_tflops") + +# Annotate each point with the lever id +for ax in (ax_top, ax_bot): + for x, lever in zip(iters_x, lever_ids): + ax.annotate(lever, (x, ax.get_ylim()[1]), xytext=(0, -12), + textcoords='offset points', rotation=30, + fontsize=7, ha='center') + +fig.suptitle(f"HydraGNN sysopt loop {uuid} — {n_iters} iters") +fig.tight_layout() +fig.savefig(out_png_path, dpi=150, bbox_inches="tight") +``` + +Use `matplotlib.use("Agg")` for headless rendering. Do NOT use seaborn (extra dep not in the omnistat venv). Use the same omnistat venv `/shared/aaji/tools/omnistat-pr271/bin/python` which already has matplotlib via TraceLens deps. + +## Style rules + +- Quote specific FOM numbers; never use "much better" / "a lot". +- Cite the lever_picker's `reason` verbatim — that's the story of WHY each lever was chosen. +- For ACCEPTED iters, the delta is against the iter that was the best before this iter ran. For REJECTED iters, the delta is against current_best (which doesn't change on rejection). +- If `kernel_trace_diag_only` fired, name what was learned from it ("step 3's kernel-trace pinned the dominant kernel to ``, which led iter 4 to pick ``"). +- If the loop hit early-stop (5% convergence rule), say so in TL;DR; if it exhausted the catalog, say so; if STOP flag was hit, say so with timestamp. + +## Failure modes + +| Failure | Action | +|---|---| +| foms.csv empty / missing | STATUS=fail; reason=no_foms | +| only iter-0 ran | STATUS=partial; story written with just baseline section + the lever picker's analysis of what would have been next | +| any per-iter combined_report.md missing | substitute "(combined_report.md missing for this iter; see /hydragnn-train-*.out)" in the bottleneck-shift section | +| matplotlib unavailable | STATUS=partial; story.md still written; foms.png omitted with a banner at the top of story.md | + +## Notes for the implementing agent + +- Read at most the files listed in `## Inputs`. +- Atomic write: tmpfile + rename for both story.md and foms.png. +- The story is auto-generated each time; it MUST be self-contained (someone reading just story.md should understand the run without opening any other file). Cross-link the per-iter combined_report.md files but never assume the reader will follow them. diff --git a/material_science/models/HydraGNN/recipes/perf-optimizer-loop/lever_catalog.yaml b/material_science/models/HydraGNN/recipes/perf-optimizer-loop/lever_catalog.yaml new file mode 100644 index 0000000..7ef18c1 --- /dev/null +++ b/material_science/models/HydraGNN/recipes/perf-optimizer-loop/lever_catalog.yaml @@ -0,0 +1,182 @@ +# HydraGNN perf-optimizer-loop — lever catalog +# +# Each entry describes ONE atomic change the orchestrator may apply between +# iterations. Levers are tried in `catalog_order` priority. After an iteration +# the lever id is recorded in `foms.csv`; if it regressed the primary FOM +# (epoch_time_s), the lever id is appended to `do_not_retry.json` and the +# orchestrator restores the previous-best config. +# +# Schema: +# id str unique short id (used in foms.csv and do_not_retry.json) +# description str one-line human description +# catalog_order int smaller = tried earlier; baseline=0 +# kind str env_only | rank_script_patch | diagnostic +# env_vars map vars to add/change to the sbatch submission (override) +# rank_script_patch str (optional) — single line of python that wraps the model +# construction in the rank script. Only used when kind=rank_script_patch. +# Applied via HG_RANK_PRE_TRAIN_HOOK env var (orchestrator writes +# a small python file and bind-mounts it; the rank script invokes +# runpy on the hook before train_validate_test()). +# diagnostic_only bool — when true, FOM accept/reject is skipped; counts against budget +# expected_payoff str qualitative ("high"|"medium"|"low") — used by lever_picker for ranking +# risk str ("low"|"medium"|"high") — used by lever_picker +# revert_method str ("drop_env"|"drop_hook") — how the orchestrator undoes the change +# citation url source for the recommendation (ROCm blog/doc, or "local") +# notes str any caveats the lever_picker should consider + +levers: + + - id: baseline + description: "Baseline run with current R2 defaults (HG_BATCH_SIZE=400, fp64, persistent_workers=1, num_workers=8). Must be iter-0." + catalog_order: 0 + kind: env_only + env_vars: + HG_BATCH_SIZE: "400" + HG_PRECISION: "fp64" + HYDRAGNN_NUM_WORKERS: "8" + HYDRAGNN_PERSISTENT_WORKERS: "1" + HG_NUM_EPOCH: "3" + HYDRAGNN_MAX_NUM_BATCH: "50" + PROFILE_TARGET_EPOCH: "2" + diagnostic_only: false + expected_payoff: "n/a (baseline)" + risk: "low" + revert_method: "drop_env" + citation: "/shared/aaji/models/HydraGNN/perf-runs/6985/combined_report.md" + notes: "Reference state. Establishes the FOM denominator for all subsequent iters." + + - id: batch_800 + description: "Double batch size: HG_BATCH_SIZE 400 -> 800. R2 #1 next-lever recommendation; VRAM headroom is 24/288 GB." + catalog_order: 1 + kind: env_only + env_vars: + HG_BATCH_SIZE: "800" + diagnostic_only: false + expected_payoff: "high" + risk: "low" + revert_method: "drop_env" + citation: "/shared/aaji/models/HydraGNN/perf-runs/6985/combined_report.md" + notes: "R2 showed 200->400 gave +1.80x util and +2.41x MFMA. Diminishing returns expected; if delta < 30%, move on." + + - id: torch_compile_e3nn + description: "Wrap the HydraGNN model in torch.compile(fullgraph=False, mode='max-autotune') to fuse aten::bmm launches in e3nn tensor-product layers. R2 #2 next-lever." + catalog_order: 2 + kind: rank_script_patch + env_vars: + HG_TORCH_COMPILE: "1" + HG_TORCH_COMPILE_MODE: "max-autotune" + TORCHINDUCTOR_MAX_AUTOTUNE: "1" + rank_script_patch: | + # Applied via HG_RANK_PRE_TRAIN_HOOK. Monkeypatches the HydraGNN model + # right after it is built inside train_validate_test, before the first + # forward. We patch hydragnn.utils.model.print_model so that it is + # called once after model construction; rather than rely on that hook, + # we wrap hydragnn.train.train.train (the inner train loop) to compile + # the model on first call. Reason: hydragnn does not expose a + # post-build hook, but the inner train function receives the model + # as an arg and is called once per epoch. + import os + import hydragnn.train.train as _hg_train + _orig_train = _hg_train.train + _compiled = {"done": False} + import torch + def _patched_train(model, *a, **kw): + if not _compiled["done"] and os.environ.get("HG_TORCH_COMPILE", "0") == "1": + mode = os.environ.get("HG_TORCH_COMPILE_MODE", "max-autotune") + try: + model = torch.compile(model, fullgraph=False, mode=mode) + if int(os.environ.get("SLURM_PROCID", "0")) == 0: + print(f"[hg_hook] torch.compile applied (mode={mode})", flush=True) + except Exception as e: + if int(os.environ.get("SLURM_PROCID", "0")) == 0: + print(f"[hg_hook] torch.compile FAILED: {e}; falling back to eager", flush=True) + _compiled["done"] = True + return _orig_train(model, *a, **kw) + _hg_train.train = _patched_train + diagnostic_only: false + expected_payoff: "high" + risk: "medium" + revert_method: "drop_hook" + citation: "https://rocm.docs.amd.com/en/latest/how-to/tuning-guides/mi300x/workload.html" + notes: "If first-iter compile overhead dominates wall time, set HYDRAGNN_MAX_NUM_BATCH=80 next round to amortize. May trigger Triton autotune cache warmup that costs ~30-60s on iter-0 of compile." + + - id: rccl_high_priority + description: "RCCL stream priority and HIP hardware queue count: TORCH_NCCL_HIGH_PRIORITY=1, GPU_MAX_HW_QUEUES=2. Improves compute<->RCCL overlap." + catalog_order: 3 + kind: env_only + env_vars: + TORCH_NCCL_HIGH_PRIORITY: "1" + GPU_MAX_HW_QUEUES: "2" + diagnostic_only: false + expected_payoff: "medium" + risk: "low" + revert_method: "drop_env" + citation: "https://rocm.docs.amd.com/en/latest/how-to/tuning-guides/mi300x/workload.html" + notes: "R2 already showed exposed NCCL was only 3% of fused-step wall; expected payoff at N=2 is small. Pick after compute-side levers (batch_800, torch_compile) are exhausted." + + - id: tunable_op + description: "Enable PyTorch TunableOp to pick fastest hipBLASLt GEMM kernels per shape: PYTORCH_TUNABLEOP_ENABLED=1, PYTORCH_TUNABLEOP_TUNING=1." + catalog_order: 4 + kind: env_only + env_vars: + PYTORCH_TUNABLEOP_ENABLED: "1" + PYTORCH_TUNABLEOP_TUNING: "1" + PYTORCH_TUNABLEOP_VERBOSE: "1" + diagnostic_only: false + expected_payoff: "medium" + risk: "low" + revert_method: "drop_env" + citation: "https://rocm.docs.amd.com/en/latest/how-to/tuning-guides/mi300x/workload.html" + notes: "First iter with TunableOp on will have measurable tuning overhead — interpret epoch_time delta cautiously; rerun if first epoch is clearly polluted by tuning." + + - id: num_workers_12 + description: "DataLoader: HYDRAGNN_NUM_WORKERS 8 -> 12 (persistent_workers already on). Picks at headroom if R2 dataloader bursts persist." + catalog_order: 5 + kind: env_only + env_vars: + HYDRAGNN_NUM_WORKERS: "12" + diagnostic_only: false + expected_payoff: "low" + risk: "low" + revert_method: "drop_env" + citation: "local" + notes: "R2 reported dataloader was non-bottleneck in steady-state windows; expected payoff is small. Only pick when other levers exhausted." + + - id: nccl_minchannels + description: "NCCL_MIN_NCHANNELS=112: forces more channels for collective dispatch, helps when message size is moderate (HydraGNN allreduces are 11.3 MB)." + catalog_order: 6 + kind: env_only + env_vars: + NCCL_MIN_NCHANNELS: "112" + diagnostic_only: false + expected_payoff: "low" + risk: "low" + revert_method: "drop_env" + citation: "https://rocm.docs.amd.com/en/docs-7.0.1/how-to/rocm-for-ai/system-setup/multi-node-setup.html" + notes: "Mostly relevant at N >= 8 nodes. Low priority for the 2-node loop; included for completeness." + + - id: precision_fp32 + description: "HG_PRECISION fp64 -> fp32. gfx950 fp32 MFMA peak (~157 TFLOP/s) is ~2x fp64 (~78.6 TFLOP/s). Numerical floor for MACE-style equivariance." + catalog_order: 7 + kind: env_only + env_vars: + HG_PRECISION: "fp32" + diagnostic_only: false + expected_payoff: "medium" + risk: "medium" + revert_method: "drop_env" + citation: "local + gfx950 spec sheet" + notes: "If the run is still dispatch-bound (mean util < 20%), fp32 wont move the needle much; orchestrator should ONLY pick this lever when fom_extractor reports the most recent best-iter is compute-bound (mfma_tflops > 30% of fp64 peak in best 10-s burst). bf16 is INTENTIONALLY EXCLUDED from this catalog: it breaks equivariance for MACE-class features." + + - id: kernel_trace_diag_only + description: "DIAGNOSTIC: enable OMNISTAT_KERNEL_TRACE=1 for ONE iteration to capture per-kernel-name dispatch counts/durations. Fires only when the 1-s TraceLens<->Omnistat correlation cannot attribute the dominant kernel." + catalog_order: 8 + kind: diagnostic + env_vars: + OMNISTAT_KERNEL_TRACE: "1" + diagnostic_only: true + expected_payoff: "n/a (diagnostic)" + risk: "medium" + revert_method: "drop_env" + citation: "../../../../../.cursor/skills/ai4science-studio/SKILL.md#12-three-rocprofiler-sdk-surfaces-pick-the-right-one" + notes: "Costs ~1000x VictoriaMetrics cardinality vs default. Use only when lever_picker decides per-kernel attribution is needed AND the kernel_correlation.csv from fom_extractor is too coarse. At most 1 of 5 iters." From f7f62308b1037f7fe707a0c72381db6ffb82cfbc Mon Sep 17 00:00:00 2001 From: Ashwin Aji Date: Fri, 22 May 2026 18:31:59 +0000 Subject: [PATCH 08/31] perf-optimizer-loop: add per-node mount-health probe; correct iter-1 diagnosis Loop-43b33ec1 iter-1 (job 7088) was misdiagnosed as "a-nodes lack /home"; actually only lux-mi355x-a6 had a broken per-user autofs/NFS mount, while lux-mi355x-a1 in the same allocation was healthy. The pmix collective fence on a1 deadlocked waiting for the failed a6 ranks until SLURM killed at the 30 min wall. Post-loop probe of all 12 a-nodes (a1-a5, a7-a12 healthy; a6 still bad as of 2026-05-22 18:23 UTC) confirmed the per-node fault. Changes: * sbatch_train_perf_amd.sh: inline per-job mount-health probe (one srun task per node, checks /home/$USER, /shared/$USER, SIF path). Exits 42 with a clear FATAL line listing broken nodes and suggesting the retry command. Override via HG_SKIP_NODE_HEALTH_PROBE=1. * recipes/perf-optimizer-loop/agents/orchestrator.md: new step 2e-bis to handle sacct ExitCode 42:0 by re-submitting the same lever with sbatch --exclude=, appending to loop-/known_bad_nodes.txt, WITHOUT penalizing the lever or shrinking the node class. Escalates after 3 recurrences of the same node. * recipes/perf-optimizer-loop/agents/lever_picker.md: ignore foms.csv rows whose immediately-preceding STATUS event was ITER_BROKEN_NODE (infra fail, not lever evidence). * .cursor/skills/ai4science-perf-analysis/SKILL.md: capture the correct cluster lesson (known-bad node list: lux-mi355x-a6) plus four other real lessons from the loop-43b33ec1 pilot (torch.compile hook path, parse_convergence epoch over-count, MFMA TFLOPS methodology, epoch_time as primary FOM). Co-authored-by: Cursor --- .../skills/ai4science-perf-analysis/SKILL.md | 11 ++++ .../examples/sbatch_train_perf_amd.sh | 51 +++++++++++++++++++ .../agents/lever_picker.md | 2 + .../agents/orchestrator.md | 25 +++++++-- 4 files changed, 86 insertions(+), 3 deletions(-) diff --git a/.cursor/skills/ai4science-perf-analysis/SKILL.md b/.cursor/skills/ai4science-perf-analysis/SKILL.md index aee6733..855d83c 100644 --- a/.cursor/skills/ai4science-perf-analysis/SKILL.md +++ b/.cursor/skills/ai4science-perf-analysis/SKILL.md @@ -44,6 +44,9 @@ Default target: HydraGNN on Lux MI355X. The same pattern can extend to ORBIT-2 / - Central Prometheus at `atlvrmonad01:9090` is **unreachable from compute nodes** (TCP RST, IP allowlist). Do **not** point omnistat-query at it. Always use the user-mode VictoriaMetrics that the launcher started. - System-mode Omnistat is running on every compute node at port 8001. Don't touch it; we run our own user-mode collector alongside. - The login node (`rad-vultr-login`) has no GPU and no MPI. All analysis runs after the job — no live profiling on the login node. +- **Node-specific mount faults are real and silent.** Any single compute node can land an allocation with a broken per-user autofs/NFS mount of `/home/$USER` or `/shared/$USER`, while sibling nodes in the same allocation are fine. Symptom in slurmd logs on the broken node: `Home Directory for Not Found ... Please contact system admin` and `lstat /shared/...: no such file or directory`; the surviving nodes wedge in a `pmix_coll_ring` collective fence until SLURM kills the job at the wall. + - **Confirmed bad nodes** (as of 2026-05-22, until cluster admin repairs them): `lux-mi355x-a6` (`/home/aaji` missing, `/shared/aaji/images/*.sif` unreadable; sibling `lux-mi355x-a1` healthy in the same alloc — see loop-43b33ec1 iter-1 / job 7088). + - **Diagnostic rule:** never blacklist a node *class* (e.g. "avoid all a-nodes") based on a single failure. Always probe the *specific* allocated nodes and exclude only those that fail by name. Probe is a one-task-per-node srun that checks `/home/$USER`, `/shared/$USER`, and the SIF path; see the `Per-node mount-health probe` section of `sbatch_train_perf_amd.sh` (exit 42 on fail). The perf-optimizer-loop orchestrator handles exit 42 by re-submitting with `--exclude=` and adding the node to `loop-/known_bad_nodes.txt`; the lever is NOT penalized. ## When something goes wrong @@ -64,4 +67,12 @@ The first 2-node end-to-end run on Lux exposed five real bugs that are now all w metric * on (instance) group_left() (max by (instance) (rmsjob_info{jobid="..."})) ``` +## Lessons captured (perf-optimizer-loop pilot, loop-43b33ec1, 2026-05-22) + +1. **Diagnose broken nodes individually, not by class** (see Cluster constraints above for full detail). The optimizer-loop's iter-1 wrongly concluded "a-nodes lack /home" from a single failed alloc — actually only `lux-mi355x-a6` was broken. The fix is a per-job mount-health probe (now in `sbatch_train_perf_amd.sh`, exit code 42) plus orchestrator step 2e-bis: on exit 42, parse `node_health_probe.txt`, append the bad node(s) to `loop-/known_bad_nodes.txt`, and re-submit with `--exclude=`. Lever is untouched. +2. **HydraGNN `torch.compile` hooks must wrap the model object, not a module path.** `import hydragnn.train.train` (the path the iter-4 hook tried) does not exist in the installed package. Wrap the constructed model directly: `model = torch.compile(model, backend="inductor", mode="reduce-overhead", fullgraph=False)` at the rank-script level, by patching the entrypoint script (`gfm_mlip_all_mpnn.py`) immediately after model construction. +3. **`parse_convergence.py` over-counts epochs by N_ranks** because each rank emits a `tqdm=100%` line and the parser counts every one. Always read epoch wall-time from rank-0 tqdm directly (`s/it × max_num_batch`), not from the parser's epoch count. Fixed in fom_extractor; long-term fix is to filter on `rank=0` lines in the parser. +4. **MFMA TFLOPS methodology must be declared explicitly.** Burst-rate via PromQL `rate([10s])×4` (iter-0/3 values ≈ 0.006-0.012 TFLOPS) and time-averaged `total_accumulated_ops / duration` (iter-4 value ≈ 0.75 TFLOPS) differ by ~60× for this workload and **must not be compared**. fom_extractor should standardize on time-averaged across the full job window, with the burst-rate as a separate `*_peak_burst` field. +5. **`epoch_time_s` is the correct primary FOM** for HydraGNN-style latency-bound GNN workloads. Throughput (samples/s) misleads when batch size is a lever, because epoch wall time scales with batches-per-epoch, not just per-sample work. + The full list (with cross-cutting context) is at the end of [.cursor/skills/ai4science-studio/SKILL.md](../ai4science-studio/SKILL.md). When fixing a bug here, check there first and propagate the lesson. diff --git a/material_science/models/HydraGNN/examples/sbatch_train_perf_amd.sh b/material_science/models/HydraGNN/examples/sbatch_train_perf_amd.sh index 2111b54..9ef931d 100755 --- a/material_science/models/HydraGNN/examples/sbatch_train_perf_amd.sh +++ b/material_science/models/HydraGNN/examples/sbatch_train_perf_amd.sh @@ -239,6 +239,57 @@ echo " Node(s) : ${SLURM_NODELIST:-$(hostname)}" echo " Date : $(date)" echo "" +# --------------------------------------------------------------------------- +# Per-node mount-health probe (lesson from loop-43b33ec1, iter-1, job 7088) +# +# Background: an allocation can land on a node where the per-user autofs/NFS +# mount of /home (or a shared dir under /shared) silently fails. The job then +# wedges in a pmix collective fence until SLURM kills it at the wall, wasting +# the iteration budget. We surface that condition immediately so the caller +# can retry with `sbatch --exclude=` (NEVER blanket-exclude a node +# class — see story.md Lessons #1 in any optimizer-loop dir). +# +# Override probe by setting HG_SKIP_NODE_HEALTH_PROBE=1 (not recommended). +# --------------------------------------------------------------------------- +HG_SKIP_NODE_HEALTH_PROBE="${HG_SKIP_NODE_HEALTH_PROBE:-0}" +if [[ "$HG_SKIP_NODE_HEALTH_PROBE" != "1" ]]; then + echo "--- Per-node mount-health probe ($(scontrol show hostnames "$SLURM_NODELIST" | tr '\n' ',' | sed 's/,$//')) ---" + _PROBE_OUT="${HG_OUTPUT_DIR}/node_health_probe.txt" + # One task per node, no-kill so we get reports from healthy nodes even if some fail. + # Each task reports: hostname, /home/$USER existence, /shared/$USER existence, SIF readability. + srun --no-kill --kill-on-bad-exit=0 -N "$SLURM_JOB_NUM_NODES" --ntasks-per-node=1 \ + --cpus-per-task=1 --gres=none --output="${_PROBE_OUT}" --error="${_PROBE_OUT}" \ + bash -c ' + H=$(hostname) + HM=$(test -d "/home/$USER" 2>/dev/null && echo OK || echo FAIL) + SH=$(test -d "'"$AI4S_SHARED_DIR"'/$USER" 2>/dev/null || test -d "'"$AI4S_SHARED_DIR"'" 2>/dev/null && echo OK || echo FAIL) + SIF_CHECK=$(test -f "'"$SIF"'" 2>/dev/null && echo OK || echo FAIL) + echo "NODE_HEALTH $H home=$HM shared=$SH sif=$SIF_CHECK" + ' 2>&1 || true + echo "" + echo " Probe results:" + if [[ -s "$_PROBE_OUT" ]]; then + grep "^NODE_HEALTH" "$_PROBE_OUT" | sed 's/^/ /' + _BAD_NODES=$(grep "^NODE_HEALTH" "$_PROBE_OUT" | awk '/home=FAIL|shared=FAIL|sif=FAIL/ {print $2}' | sort -u | tr '\n' ',' | sed 's/,$//') + else + echo " (no output captured; assuming all nodes healthy)" + _BAD_NODES="" + fi + _MISSING_REPORTS=$(( SLURM_JOB_NUM_NODES - $(grep -c "^NODE_HEALTH" "$_PROBE_OUT" 2>/dev/null || echo 0) )) + if [[ -n "$_BAD_NODES" ]]; then + echo "" >&2 + echo "FATAL: NODE_HEALTH_PROBE detected broken mounts on: $_BAD_NODES" >&2 + echo " Retry with: sbatch --exclude=$_BAD_NODES $0" >&2 + echo " Do NOT blanket-exclude the node class (a* or b*); only specific bad nodes." >&2 + echo " Notify cluster admin so the broken node can be repaired." >&2 + exit 42 + fi + if (( _MISSING_REPORTS > 0 )); then + echo "WARN: $_MISSING_REPORTS of $SLURM_JOB_NUM_NODES nodes did not return a health-probe result; proceeding anyway" >&2 + fi + echo " All $SLURM_JOB_NUM_NODES nodes healthy." +fi + # --------------------------------------------------------------------------- # Start Omnistat user-mode (datadir from omnistat-lux.config: under HG_OUTPUT_DIR) # --------------------------------------------------------------------------- diff --git a/material_science/models/HydraGNN/recipes/perf-optimizer-loop/agents/lever_picker.md b/material_science/models/HydraGNN/recipes/perf-optimizer-loop/agents/lever_picker.md index 9b82c3b..9f1c111 100644 --- a/material_science/models/HydraGNN/recipes/perf-optimizer-loop/agents/lever_picker.md +++ b/material_science/models/HydraGNN/recipes/perf-optimizer-loop/agents/lever_picker.md @@ -31,6 +31,8 @@ Read catalog; drop: - Any lever id already accepted in `foms.csv` (we don't re-pick already-applied levers; the new "best" already has them baked in). - `kernel_trace_diag_only` — see special gate below. +**Do NOT** drop levers that appear in `foms.csv` with `accepted=false` if the failure was infrastructural (e.g. `ITER_BROKEN_NODE`, `ITER_SBATCH_FAIL`, `ANALYZE_FAIL`). Those failures are not evidence about the lever — the orchestrator handles them by retry. Cross-check by reading STATUS.txt: if the most recent ITER_DECISION for a lever was preceded by `ITER_BROKEN_NODE` for the same jobid, treat the lever as untried. + If the candidate set is empty: emit `STATUS=partial; reason=catalog_exhausted`. ### 2. Score remaining candidates diff --git a/material_science/models/HydraGNN/recipes/perf-optimizer-loop/agents/orchestrator.md b/material_science/models/HydraGNN/recipes/perf-optimizer-loop/agents/orchestrator.md index 7d14e40..5410c47 100644 --- a/material_science/models/HydraGNN/recipes/perf-optimizer-loop/agents/orchestrator.md +++ b/material_science/models/HydraGNN/recipes/perf-optimizer-loop/agents/orchestrator.md @@ -113,15 +113,21 @@ Log `LEVER_PICK iter= lever= reason=""`. **2c. STOP-flag gate.** Re-check; act as in 2a. -**2d. Submit sbatch.** Build env-var diff: union of the current-best's env vars + the new lever's `env_vars`, minus any vars revert-method'd from rejected iters. Render an `iter--env.sh` file with all overrides (one `export KEY=VALUE` per line) under `loop-/`. Submit: +**2d. Submit sbatch.** Build env-var diff: union of the current-best's env vars + the new lever's `env_vars`, minus any vars revert-method'd from rejected iters. Render an `iter--env.sh` file with all overrides (one `export KEY=VALUE` per line) under `loop-/`. + +**Before every sbatch**, read `loop-/known_bad_nodes.txt` (one node per line, may be preseeded by the operator or appended by step 2e-bis from previous iters). If it exists and is non-empty, pass `--exclude=` to sbatch. This carries the known-bad list across iters and across orchestrator restarts. ```bash +EXCLUDE_ARG="" +if [[ -s "loop-/known_bad_nodes.txt" ]]; then + EXCLUDE_ARG="--exclude=$(paste -sd, loop-/known_bad_nodes.txt)" +fi ( set -a; source "loop-/iter--env.sh"; set +a; \ cd /home/aaji/git/ai4science-studio; \ - sbatch material_science/models/HydraGNN/examples/sbatch_train_perf_amd.sh ) + sbatch $EXCLUDE_ARG material_science/models/HydraGNN/examples/sbatch_train_perf_amd.sh ) ``` -Capture the jobid from sbatch stdout. Log `ITER_SUBMIT iter= lever= jobid=`. +Capture the jobid from sbatch stdout. Log `ITER_SUBMIT iter= lever= jobid=` (include `exclude=` if non-empty). If the lever is `kind: rank_script_patch`, write the patch content from `lever_catalog.yaml` to `loop-/iter--hook.py` and export `HG_RANK_PRE_TRAIN_HOOK=` before sbatch. (The rank script must be extended to honor this var; see the rank-script extension note below.) @@ -129,6 +135,18 @@ If the lever is `kind: rank_script_patch`, write the patch content from `lever_c Log `ITER_COMPLETE iter= jobid= state= runtime=`. +**2e-bis. Broken-node detection (NODE_HEALTH_PROBE exit 42).** The sbatch script runs a per-node mount-health probe immediately after allocation and exits with code 42 if any allocated node has a broken `/home/$USER`, `/shared/$USER`, or SIF mount. When `sacct ExitCode` is `42:0`: + +1. Grep `/node_health_probe.txt` for `home=FAIL|shared=FAIL|sif=FAIL` lines; collect broken hostnames into a `BAD_NODES` comma-list. +2. Log `ITER_BROKEN_NODE iter= jobid= bad_nodes= note=mount_fault_not_lever_regression`. +3. **Do NOT** add the lever to `do_not_retry.json`. **Do NOT** advance `current_iter`. **Do NOT** restrict the node class. +4. Append `BAD_NODES` to a persistent file `loop-/known_bad_nodes.txt` (one node per line, deduped) for the duration of this loop. +5. Retry the same lever immediately on the same iteration counter with `sbatch --exclude="$(paste -sd, loop-/known_bad_nodes.txt)"`. Log `ITER_RESUBMIT iter= lever= exclude= reason=mount_fault_retry`. +6. If the **same** node appears in `known_bad_nodes.txt` 3 times across the loop without admin intervention, log `LOOP_ABORT reason=persistent_node_fault bad_nodes=` and exit 2 — escalate to the human. +7. After loop ends, the story_writer reads `known_bad_nodes.txt` and surfaces the list to the user as a Lessons entry, so the cluster admin can be notified. + +This handles the iter-1 case from loop-43b33ec1 correctly: only `lux-mi355x-a6` had broken mounts; `lux-mi355x-a1` in the same alloc was healthy. The orchestrator's job is to flag the specific bad node and route around it, not to penalize the lever or shrink the node pool by class. + **2f. STOP-flag gate.** Re-check before the analysis phase. If set, still analyze the just-completed iter (we paid for it), then exit gracefully. **2g. Run analyst+verifier+synth.** Re-use the existing flow from `recipes/perf-analysis/`. Dispatch in parallel (single message with multiple Task tool calls): @@ -203,6 +221,7 @@ This is a single-edit, one-time change to the existing rank script. The orchestr |---|---| | `sinfo` shows cluster drained | PREFLIGHT_FAIL reason=cluster_down; exit 2 | | sbatch returns nonzero | log ITER_SBATCH_FAIL; treat as auto-reject; continue | +| sacct ExitCode 42:0 (NODE_HEALTH_PROBE failed) | see step 2e-bis — retry with `--exclude=`, do NOT penalize lever, do NOT shrink node class | | sacct never reaches terminal in 45 min | log ITER_TIMEOUT; auto-reject; continue | | analyst/verifier returns STATUS=fail | log ANALYZE_FAIL; retry next iter; if 2 in a row, do_not_retry the lever with reason=repeated_analysis_fail | | fom_extractor returns STATUS=fail | treat as accepted=false reason=fom_extraction_fail; do NOT add to do_not_retry; orchestrator should investigate manually next morning | From c019d59e467558776d992f38168a97ecd261b456 Mon Sep 17 00:00:00 2001 From: Ashwin Aji Date: Sat, 23 May 2026 06:18:33 +0000 Subject: [PATCH 09/31] perf-optimizer-loop: capture torch.compile BLOCKED finding + 8 sub-lessons Findings from loop-c3e4df1c (2026-05-22, 5 iters, best -1.3% epoch / -8.4% energy via rccl_high_priority) and compile-investigation-aa480554 (2026-05-23, 6 hooks across 3 hypotheses, all FAIL with same root cause). Catalog changes (material_science/models/HydraGNN/recipes/perf-optimizer-loop/lever_catalog.yaml): * torch_compile_e3nn -> status: blocked with full root-cause docstring. Root cause: HydraGNN's energy_force_loss calls torch.autograd.grad(graph_energy_pred, data.pos, create_graph=True), which requires DOUBLE BACKWARD. Every torch.compile mode goes through AOT Autograd which pre-compiles the backward and cannot itself be differentiated. Verbatim error: "RuntimeError: torch.compile with aot_autograd does not currently support double backward". Lever reference impl is preserved (now raises RuntimeError) as a cautionary record so future agents see the failure modes documented inline. * tunable_op -> renamed to tunable_op_live with status: blocked. PYTORCH_TUNABLEOP_TUNING=1 causes "Memory access fault by GPU node-N" on all 16 MI355X GPUs during hipBLASLt live tuning. Loop-c3e4df1c iter-3 / job 7188. * tunable_op_warmup_then_use -> new safe replacement. Documents the 2-phase pattern: dedicated warmup sbatch generates tunableop_results_.csv files, then production runs consume them with PYTORCH_TUNABLEOP_ENABLED=1 (no _TUNING=1). * num_workers_12 -> node-class-dependent note. Improved a-nodes (-12.4% in loop-43b33ec1) but regressed b-nodes (+10.7% in loop-c3e4df1c). lever_picker now told to consider node class. Lever picker change (agents/lever_picker.md): * Drop any lever with status: blocked from the candidate set. Read blocked_reason + blocked_evidence to explain the rejection but never pick. This prevents future iterations from re-trying torch.compile or live TunableOp. Skill change (.cursor/skills/ai4science-perf-analysis/SKILL.md): * Add lessons 6-12 capturing the new findings: torch.compile BLOCKED (#6), allow_in_graph weaker than docs suggest (#7), TunableOp live unsafe (#8), lever payoff is node-class-dependent (#9), b-nodes ~17% faster baseline (#10), hydragnn/train re-export shadow (#11), sbatch --export var propagation (#12). * Add lux-mi355x-a10 to the known-bad nodes list (broken PMIx, NOT a mount fault; loop-c3e4df1c iter-1 / job 7161). Sbatch change (examples/sbatch_train_perf_amd.sh): * Fix $SIF -> $HG_SIF typo in node-health-probe stanza (runtime-fixed by the orchestrator during loop-c3e4df1c iter-0 retry; lands here so future loops don't need the same retry). Runtime artifacts (story.md, foms.png, foms.csv, STATUS.txt, hooks) remain under /shared/aaji/models/HydraGNN/perf-runs/ and are not committed. The investigation report at /shared/aaji/models/HydraGNN/perf-runs/compile-investigation-aa480554-2a15-4e04-80dd-c58afc672a8d/INVESTIGATION_REPORT.md is referenced in the catalog citation field. Co-authored-by: Cursor --- .../skills/ai4science-perf-analysis/SKILL.md | 16 ++- .../examples/sbatch_train_perf_amd.sh | 2 +- .../agents/lever_picker.md | 1 + .../perf-optimizer-loop/lever_catalog.yaml | 134 ++++++++++++------ 4 files changed, 109 insertions(+), 44 deletions(-) diff --git a/.cursor/skills/ai4science-perf-analysis/SKILL.md b/.cursor/skills/ai4science-perf-analysis/SKILL.md index 855d83c..4e4ed07 100644 --- a/.cursor/skills/ai4science-perf-analysis/SKILL.md +++ b/.cursor/skills/ai4science-perf-analysis/SKILL.md @@ -45,7 +45,9 @@ Default target: HydraGNN on Lux MI355X. The same pattern can extend to ORBIT-2 / - System-mode Omnistat is running on every compute node at port 8001. Don't touch it; we run our own user-mode collector alongside. - The login node (`rad-vultr-login`) has no GPU and no MPI. All analysis runs after the job — no live profiling on the login node. - **Node-specific mount faults are real and silent.** Any single compute node can land an allocation with a broken per-user autofs/NFS mount of `/home/$USER` or `/shared/$USER`, while sibling nodes in the same allocation are fine. Symptom in slurmd logs on the broken node: `Home Directory for Not Found ... Please contact system admin` and `lstat /shared/...: no such file or directory`; the surviving nodes wedge in a `pmix_coll_ring` collective fence until SLURM kills the job at the wall. - - **Confirmed bad nodes** (as of 2026-05-22, until cluster admin repairs them): `lux-mi355x-a6` (`/home/aaji` missing, `/shared/aaji/images/*.sif` unreadable; sibling `lux-mi355x-a1` healthy in the same alloc — see loop-43b33ec1 iter-1 / job 7088). + - **Confirmed bad nodes** (as of 2026-05-23, until cluster admin repairs them): + - `lux-mi355x-a6` — `/home/aaji` missing, `/shared/aaji/images/*.sif` unreadable; sibling `lux-mi355x-a1` healthy in the same alloc. Loop-43b33ec1 iter-1 / job 7088. + - `lux-mi355x-a10` — pmix ring collective times out at MPI init; `opal_mpi_init_ext3x_client_unreachable`. NOT a mount fault — broken PMIx/OpenMPI install on this node. Loop-c3e4df1c iter-1 / job 7161. - **Diagnostic rule:** never blacklist a node *class* (e.g. "avoid all a-nodes") based on a single failure. Always probe the *specific* allocated nodes and exclude only those that fail by name. Probe is a one-task-per-node srun that checks `/home/$USER`, `/shared/$USER`, and the SIF path; see the `Per-node mount-health probe` section of `sbatch_train_perf_amd.sh` (exit 42 on fail). The perf-optimizer-loop orchestrator handles exit 42 by re-submitting with `--exclude=` and adding the node to `loop-/known_bad_nodes.txt`; the lever is NOT penalized. ## When something goes wrong @@ -70,9 +72,19 @@ The first 2-node end-to-end run on Lux exposed five real bugs that are now all w ## Lessons captured (perf-optimizer-loop pilot, loop-43b33ec1, 2026-05-22) 1. **Diagnose broken nodes individually, not by class** (see Cluster constraints above for full detail). The optimizer-loop's iter-1 wrongly concluded "a-nodes lack /home" from a single failed alloc — actually only `lux-mi355x-a6` was broken. The fix is a per-job mount-health probe (now in `sbatch_train_perf_amd.sh`, exit code 42) plus orchestrator step 2e-bis: on exit 42, parse `node_health_probe.txt`, append the bad node(s) to `loop-/known_bad_nodes.txt`, and re-submit with `--exclude=`. Lever is untouched. -2. **HydraGNN `torch.compile` hooks must wrap the model object, not a module path.** `import hydragnn.train.train` (the path the iter-4 hook tried) does not exist in the installed package. Wrap the constructed model directly: `model = torch.compile(model, backend="inductor", mode="reduce-overhead", fullgraph=False)` at the rank-script level, by patching the entrypoint script (`gfm_mlip_all_mpnn.py`) immediately after model construction. +2. **HydraGNN `torch.compile` hooks must wrap the model object, not a module path.** `import hydragnn.train.train` (the path the iter-4 hook tried) does not exist in the installed package. Wrap the constructed model directly: `model = torch.compile(model, backend="inductor", mode="reduce-overhead", fullgraph=False)` at the rank-script level, by patching the entrypoint script (`gfm_mlip_all_mpnn.py`) immediately after model construction. **(Now superseded by Lesson #6 below — torch.compile is BLOCKED for HydraGNN MLIP regardless of where it wraps.)** 3. **`parse_convergence.py` over-counts epochs by N_ranks** because each rank emits a `tqdm=100%` line and the parser counts every one. Always read epoch wall-time from rank-0 tqdm directly (`s/it × max_num_batch`), not from the parser's epoch count. Fixed in fom_extractor; long-term fix is to filter on `rank=0` lines in the parser. 4. **MFMA TFLOPS methodology must be declared explicitly.** Burst-rate via PromQL `rate([10s])×4` (iter-0/3 values ≈ 0.006-0.012 TFLOPS) and time-averaged `total_accumulated_ops / duration` (iter-4 value ≈ 0.75 TFLOPS) differ by ~60× for this workload and **must not be compared**. fom_extractor should standardize on time-averaged across the full job window, with the burst-rate as a separate `*_peak_burst` field. 5. **`epoch_time_s` is the correct primary FOM** for HydraGNN-style latency-bound GNN workloads. Throughput (samples/s) misleads when batch size is a lever, because epoch wall time scales with batches-per-epoch, not just per-sample work. +## Lessons captured (loop-c3e4df1c + compile-investigation-aa480554, 2026-05-22..23) + +6. **`torch.compile` is BLOCKED for HydraGNN GFM-MLIP MACE** under PyTorch 2.10 / ROCm 7.2.2. Six hooks across three hypotheses (inner-model compile, outer + `dynamo.allow_in_graph(autograd.grad)`, single-submodule compile) all fail with the same verbatim PyTorch error: `RuntimeError: torch.compile with aot_autograd does not currently support double backward` (`torch/_functorch/_aot_autograd/runtime_wrappers.py:2356`). Root cause: HydraGNN's `energy_force_loss` (`hydragnn/models/create.py:718`) calls `torch.autograd.grad(graph_energy_pred, data.pos, create_graph=True)` to compute forces as derivatives of energy w.r.t. positions; with `create_graph=True`, training requires DOUBLE BACKWARD (loss.backward differentiates through forces). Every torch.compile mode (`default`, `reduce-overhead`, `max-autotune`) goes through AOT Autograd which pre-compiles the backward as a joint program and cannot itself be differentiated. **To unblock requires upstream HydraGNN source change**: move `energy_force_loss` INTO `EnhancedModelWrapper.forward()` so `autograd.grad` is inside the compiled forward (the canonical MACE/Allegro/NequIP pattern). ~30-50 LOC across `create.py:622-758` + `train_validate_test.py:725-735`, config-gated. The lever `torch_compile_e3nn` is marked `status: blocked` in `lever_catalog.yaml` with full evidence; the lever_picker drops blocked levers from its candidate set. +7. **`dynamo.allow_in_graph(autograd.grad)` is weaker than MACE upstream docs suggest.** It only prevents dynamo from TRACING THROUGH the call; it does NOT prevent dynamo's FakeTensor symbolic-execution pass from running `autograd.grad` against fake tensors (which fails with the `allow_unused=True` error because fake tensors have no real autograd graph). The MACE `prepare()` pattern works only when `autograd.grad` lives inside the compiled function's `forward()` — not when it's in a separate method called from outside. +8. **TunableOp live tuning (`PYTORCH_TUNABLEOP_TUNING=1`) is unsafe on MI355X / ROCm 7.2.2 with HydraGNN.** All 16 GPUs hit `Memory access fault by GPU node-N` during the hipBLASLt kernel autotuning phase (loop-c3e4df1c iter-3 / job 7188). Use a 2-phase pattern instead: dedicated warmup sbatch with `PYTORCH_TUNABLEOP_TUNING=1` to generate `tunableop_results_.csv` files, then production runs with `PYTORCH_TUNABLEOP_ENABLED=1` only (no `_TUNING=1`). The `tunable_op_warmup_then_use` lever in the catalog encodes this. The original `tunable_op_live` lever is now marked `status: blocked`. +9. **Lever payoff is node-class-dependent.** `num_workers_12` improved a-nodes (-12.4% in loop-43b33ec1) but regressed b-nodes (+10.7% in loop-c3e4df1c). Lever_picker must read `SLURM_NODELIST` from the current best run's manifest and treat a-nodes vs b-nodes as separate state for at least `num_workers_*` and `batch_*` levers. +10. **b-nodes are ~17% faster baseline than a-nodes for HydraGNN.** Loop-43b33ec1's a-node baseline was 92.5 s/epoch; loop-c3e4df1c's b-node baseline (same code, same config, same fp64, same batch_size=400) was 76.0 s/epoch. The single biggest finding of the night came from a baseline shift, not a lever. The `launcher` subagent should report which node class was allocated, and cross-loop FOM comparisons must control for this. +11. **`hydragnn/train/__init__.py` re-export shadow** breaks naive `import hydragnn.train.train_validate_test as tvt`. The line `from .train_validate_test import train_validate_test, train, ...` makes `hydragnn.train.train_validate_test` resolve to the **function** `train_validate_test` (re-exported as attribute of the package), not the module. Use `importlib.import_module("hydragnn.train.train_validate_test")` instead. Applies to any rank hook that needs to monkey-patch `train_validate_test.train`. +12. **`sbatch --export=ALL,...` does NOT include shell vars that are not in `--export`'s explicit list.** When using `--export=ALL,K1=V1,K2=V2`, vars set in the calling shell but not listed here are propagated only if they're in the user's persistent env (login shell init). Always explicitly list cluster-pathing vars like `AI4S_SHARED_DIR=/shared/aaji` in `--export`. + The full list (with cross-cutting context) is at the end of [.cursor/skills/ai4science-studio/SKILL.md](../ai4science-studio/SKILL.md). When fixing a bug here, check there first and propagate the lesson. diff --git a/material_science/models/HydraGNN/examples/sbatch_train_perf_amd.sh b/material_science/models/HydraGNN/examples/sbatch_train_perf_amd.sh index 9ef931d..fd1f03a 100755 --- a/material_science/models/HydraGNN/examples/sbatch_train_perf_amd.sh +++ b/material_science/models/HydraGNN/examples/sbatch_train_perf_amd.sh @@ -263,7 +263,7 @@ if [[ "$HG_SKIP_NODE_HEALTH_PROBE" != "1" ]]; then H=$(hostname) HM=$(test -d "/home/$USER" 2>/dev/null && echo OK || echo FAIL) SH=$(test -d "'"$AI4S_SHARED_DIR"'/$USER" 2>/dev/null || test -d "'"$AI4S_SHARED_DIR"'" 2>/dev/null && echo OK || echo FAIL) - SIF_CHECK=$(test -f "'"$SIF"'" 2>/dev/null && echo OK || echo FAIL) + SIF_CHECK=$(test -f "'"$HG_SIF"'" 2>/dev/null && echo OK || echo FAIL) echo "NODE_HEALTH $H home=$HM shared=$SH sif=$SIF_CHECK" ' 2>&1 || true echo "" diff --git a/material_science/models/HydraGNN/recipes/perf-optimizer-loop/agents/lever_picker.md b/material_science/models/HydraGNN/recipes/perf-optimizer-loop/agents/lever_picker.md index 9f1c111..000aca6 100644 --- a/material_science/models/HydraGNN/recipes/perf-optimizer-loop/agents/lever_picker.md +++ b/material_science/models/HydraGNN/recipes/perf-optimizer-loop/agents/lever_picker.md @@ -29,6 +29,7 @@ Read catalog; drop: - `baseline` (iter-0 only; never picked again). - Any lever id in `do_not_retry.json`. - Any lever id already accepted in `foms.csv` (we don't re-pick already-applied levers; the new "best" already has them baked in). +- Any lever with `status: blocked` in the catalog (these are permanently broken on this stack; the catalog entry exists only as a cautionary record — never pick them). Read each lever's `blocked_reason` and `blocked_evidence` fields when explaining the candidate set; this is how we avoid wasting iterations on known-dead levers like `torch_compile_e3nn` (double-backward + AOT Autograd) or `tunable_op_live` (GPU mem fault during live hipBLASLt tuning). - `kernel_trace_diag_only` — see special gate below. **Do NOT** drop levers that appear in `foms.csv` with `accepted=false` if the failure was infrastructural (e.g. `ITER_BROKEN_NODE`, `ITER_SBATCH_FAIL`, `ANALYZE_FAIL`). Those failures are not evidence about the lever — the orchestrator handles them by retry. Cross-check by reading STATUS.txt: if the most recent ITER_DECISION for a lever was preceded by `ITER_BROKEN_NODE` for the same jobid, treat the lever as untried. diff --git a/material_science/models/HydraGNN/recipes/perf-optimizer-loop/lever_catalog.yaml b/material_science/models/HydraGNN/recipes/perf-optimizer-loop/lever_catalog.yaml index 7ef18c1..3cd8b2e 100644 --- a/material_science/models/HydraGNN/recipes/perf-optimizer-loop/lever_catalog.yaml +++ b/material_science/models/HydraGNN/recipes/perf-optimizer-loop/lever_catalog.yaml @@ -59,46 +59,56 @@ levers: notes: "R2 showed 200->400 gave +1.80x util and +2.41x MFMA. Diminishing returns expected; if delta < 30%, move on." - id: torch_compile_e3nn - description: "Wrap the HydraGNN model in torch.compile(fullgraph=False, mode='max-autotune') to fuse aten::bmm launches in e3nn tensor-product layers. R2 #2 next-lever." + description: "BLOCKED — torch.compile is fundamentally incompatible with HydraGNN GFM-MLIP under PyTorch 2.10 / ROCm 7.2.2. Documented here so the lever_picker NEVER re-tries it." catalog_order: 2 kind: rank_script_patch + status: "blocked" + blocked_reason: "AOT-Autograd-double-backward" + blocked_evidence: "compile-investigation-aa480554 (2026-05-23): 6 hooks across 3 hypotheses, all fail with same root cause." env_vars: HG_TORCH_COMPILE: "1" - HG_TORCH_COMPILE_MODE: "max-autotune" - TORCHINDUCTOR_MAX_AUTOTUNE: "1" + HG_TORCH_COMPILE_MODE: "reduce-overhead" rank_script_patch: | - # Applied via HG_RANK_PRE_TRAIN_HOOK. Monkeypatches the HydraGNN model - # right after it is built inside train_validate_test, before the first - # forward. We patch hydragnn.utils.model.print_model so that it is - # called once after model construction; rather than rely on that hook, - # we wrap hydragnn.train.train.train (the inner train loop) to compile - # the model on first call. Reason: hydragnn does not expose a - # post-build hook, but the inner train function receives the model - # as an arg and is called once per epoch. - import os - import hydragnn.train.train as _hg_train - _orig_train = _hg_train.train - _compiled = {"done": False} - import torch - def _patched_train(model, *a, **kw): - if not _compiled["done"] and os.environ.get("HG_TORCH_COMPILE", "0") == "1": - mode = os.environ.get("HG_TORCH_COMPILE_MODE", "max-autotune") - try: - model = torch.compile(model, fullgraph=False, mode=mode) - if int(os.environ.get("SLURM_PROCID", "0")) == 0: - print(f"[hg_hook] torch.compile applied (mode={mode})", flush=True) - except Exception as e: - if int(os.environ.get("SLURM_PROCID", "0")) == 0: - print(f"[hg_hook] torch.compile FAILED: {e}; falling back to eager", flush=True) - _compiled["done"] = True - return _orig_train(model, *a, **kw) - _hg_train.train = _patched_train + # KNOWN-BROKEN — kept for reference only. Do NOT enable this lever. + # + # Root cause: HydraGNN's energy_force_loss (hydragnn/models/create.py:718) + # calls torch.autograd.grad(graph_energy_pred, data.pos, create_graph=True). + # With create_graph=True, training requires DOUBLE BACKWARD + # (loss.backward differentiates through forces, which themselves came + # from autograd.grad over the forward). PyTorch's AOT Autograd — used + # by every torch.compile mode (default, reduce-overhead, max-autotune) — + # pre-compiles the backward as a joint program and cannot itself be + # differentiated. Verbatim error (every rank): + # RuntimeError: torch.compile with aot_autograd does not currently + # support double backward + # at torch/_functorch/_aot_autograd/runtime_wrappers.py:2356 + # + # Strategies tested and rejected in compile-investigation-aa480554: + # H1 — compile inner self.model (MACEStack) only → AOT boundary + # severs data.pos lineage → "differentiated Tensors appears + # not to have been used in the graph" + # H2 — outer compile + dynamo.allow_in_graph(autograd.grad) + # (the MACE upstream pattern) → allow_in_graph only prevents + # dynamo from tracing THROUGH the call; FakeTensor symbolic + # execution still runs autograd.grad and fails + # H3 — compile only one RealAgnosticAttResidualInteractionBlock + # submodule → still triggers AOT for that submodule's + # backward → same double-backward error + # + # To unblock requires upstream HydraGNN source change: move + # energy_force_loss INTO EnhancedModelWrapper.forward so the + # autograd.grad call is inside the compiled forward, mirroring the + # canonical MACE/Allegro/NequIP pattern. ~30-50 LOC across + # hydragnn/models/create.py:622-758 and + # hydragnn/train/train_validate_test.py:725-735. Config-gated to + # preserve non-MLIP HydraGNN paths. + raise RuntimeError("torch_compile_e3nn is BLOCKED for HydraGNN MLIP; see lever_catalog.yaml") diagnostic_only: false - expected_payoff: "high" - risk: "medium" + expected_payoff: "none (blocked)" + risk: "blocked" revert_method: "drop_hook" - citation: "https://rocm.docs.amd.com/en/latest/how-to/tuning-guides/mi300x/workload.html" - notes: "If first-iter compile overhead dominates wall time, set HYDRAGNN_MAX_NUM_BATCH=80 next round to amortize. May trigger Triton autotune cache warmup that costs ~30-60s on iter-0 of compile." + citation: "/shared/aaji/models/HydraGNN/perf-runs/compile-investigation-aa480554-2a15-4e04-80dd-c58afc672a8d/INVESTIGATION_REPORT.md" + notes: "DO NOT pick this lever until HydraGNN upstream patches energy_force_loss into forward(). The lever_picker subagent must treat status:blocked as equivalent to do_not_retry." - id: rccl_high_priority description: "RCCL stream priority and HIP hardware queue count: TORCH_NCCL_HIGH_PRIORITY=1, GPU_MAX_HW_QUEUES=2. Improves compute<->RCCL overlap." @@ -114,33 +124,75 @@ levers: citation: "https://rocm.docs.amd.com/en/latest/how-to/tuning-guides/mi300x/workload.html" notes: "R2 already showed exposed NCCL was only 3% of fused-step wall; expected payoff at N=2 is small. Pick after compute-side levers (batch_800, torch_compile) are exhausted." - - id: tunable_op - description: "Enable PyTorch TunableOp to pick fastest hipBLASLt GEMM kernels per shape: PYTORCH_TUNABLEOP_ENABLED=1, PYTORCH_TUNABLEOP_TUNING=1." + - id: tunable_op_live + description: "BLOCKED on this stack — TunableOp LIVE tuning mode (PYTORCH_TUNABLEOP_TUNING=1) causes GPU memory access faults on MI355X / ROCm 7.2.2 / PyTorch 2.10 with HydraGNN. Use tunable_op_warmup_then_use instead." catalog_order: 4 kind: env_only + status: "blocked" + blocked_reason: "GPU-mem-access-fault-during-hipBLASLt-live-tuning" + blocked_evidence: "loop-c3e4df1c iter-3 (job 7188, 2026-05-22): all 16 GPUs reported 'Memory access fault by GPU node-N' during hipBLASLt kernel autotuning phase; job aborted with signal 6." env_vars: PYTORCH_TUNABLEOP_ENABLED: "1" PYTORCH_TUNABLEOP_TUNING: "1" PYTORCH_TUNABLEOP_VERBOSE: "1" diagnostic_only: false + expected_payoff: "none (blocked)" + risk: "blocked" + revert_method: "drop_env" + citation: "https://rocm.docs.amd.com/en/latest/how-to/tuning-guides/mi300x/workload.html" + notes: "Live tuning racing with training data is unsafe here. Pick tunable_op_warmup_then_use instead, which separates the tuning phase into a dedicated short job." + + - id: tunable_op_warmup_then_use + description: "2-phase safe TunableOp pattern: a dedicated 1-batch warmup job generates tunableop_results_.csv files, then the actual training job runs with PYTORCH_TUNABLEOP_ENABLED=1 (no _TUNING) to consume the pre-tuned results." + catalog_order: 4 + kind: rank_script_patch + env_vars: + PYTORCH_TUNABLEOP_ENABLED: "1" + PYTORCH_TUNABLEOP_TUNING: "0" + PYTORCH_TUNABLEOP_FILENAME: "/shared/aaji/models/HydraGNN/tunableop_results/tunableop_results_%d.csv" + rank_script_patch: | + # 2-phase TunableOp pattern. + # + # Phase 1 (separate sbatch, NOT this lever) — orchestrator must run + # this BEFORE picking the lever: + # 1. mkdir -p /shared/aaji/models/HydraGNN/tunableop_results + # 2. sbatch with PYTORCH_TUNABLEOP_ENABLED=1 PYTORCH_TUNABLEOP_TUNING=1 + # PYTORCH_TUNABLEOP_FILENAME=/shared/.../tunableop_results_%d.csv + # HG_NUM_EPOCH=1 HYDRAGNN_MAX_NUM_BATCH=5 + # (short job; expect failure but the per-rank CSVs are written + # incrementally as each shape gets tuned, so partial results + # are still usable) + # 3. verify CSVs exist: ls /shared/.../tunableop_results_*.csv + # + # Phase 2 (this lever) — the actual training run consumes the CSVs + # with TUNING=0; if a shape is missing, TunableOp falls back to + # the default kernel (no fault). No hook code needed beyond the + # env vars; this stanza is present so the orchestrator can record + # the 2-phase intent in iter--lever.json. + import os + if int(os.environ.get("SLURM_PROCID", "0")) == 0: + import glob + _n_csv = len(glob.glob("/shared/aaji/models/HydraGNN/tunableop_results/tunableop_results_*.csv")) + print(f"[tunable_op_warmup] consuming {_n_csv} pre-tuned CSV files", flush=True) + diagnostic_only: false expected_payoff: "medium" risk: "low" revert_method: "drop_env" - citation: "https://rocm.docs.amd.com/en/latest/how-to/tuning-guides/mi300x/workload.html" - notes: "First iter with TunableOp on will have measurable tuning overhead — interpret epoch_time delta cautiously; rerun if first epoch is clearly polluted by tuning." + citation: "https://rocm.docs.amd.com/en/latest/how-to/tuning-guides/mi300x/workload.html#tunableop" + notes: "If the warmup phase did NOT complete (no CSVs), pick a different lever. The orchestrator should add `setup: tunableop_warmup_sbatch` as a precondition before this lever is offered." - id: num_workers_12 - description: "DataLoader: HYDRAGNN_NUM_WORKERS 8 -> 12 (persistent_workers already on). Picks at headroom if R2 dataloader bursts persist." + description: "DataLoader: HYDRAGNN_NUM_WORKERS 8 -> 12 (persistent_workers already on). Node-class-dependent — helped on a-nodes (-12.4% in loop 43b33ec1), regressed on b-nodes (+10.7% in loop c3e4df1c)." catalog_order: 5 kind: env_only env_vars: HYDRAGNN_NUM_WORKERS: "12" diagnostic_only: false - expected_payoff: "low" + expected_payoff: "node-class-dependent" risk: "low" revert_method: "drop_env" - citation: "local" - notes: "R2 reported dataloader was non-bottleneck in steady-state windows; expected payoff is small. Only pick when other levers exhausted." + citation: "local (loop 43b33ec1 + loop c3e4df1c)" + notes: "lever_picker should only consider this lever on a-nodes. On b-nodes, num_workers=8 is the sweet spot (verified in iter-2 of loop c3e4df1c with epoch_time=75 s, the current best). If the kernel_correlation.csv from fom_extractor shows dataloader prefetch gaps > 10 ms on b-nodes, try num_workers=10 (not 12)." - id: nccl_minchannels description: "NCCL_MIN_NCHANNELS=112: forces more channels for collective dispatch, helps when message size is moderate (HydraGNN allreduces are 11.3 MB)." From bc4b9c149d4726581736e3a8b740fbf17044ab11 Mon Sep 17 00:00:00 2001 From: Ashwin Aji Date: Tue, 26 May 2026 07:27:39 +0000 Subject: [PATCH 10/31] perf-optimizer-loop: capture dispatch-bound finding + torch.jit.script BLOCKED MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Findings from pathc-pathb-a3b8f3f0 (2026-05-23): Path C (torch.jit.script) failed cleanly, and Path B (loop-e3fac2af) found no new accepted lever but uncovered the project's most important diagnostic: HydraGNN GFM-MLIP MACE on b-nodes is DISPATCH-BOUND, not compute-bound (fp32 ≈ fp64 with both verified by precision-diagnostic logs). Catalog changes (lever_catalog.yaml): * torch_jit_script (NEW lever entry, status: blocked). Root cause: MACEStack.py:397 uses **conv_args keyword-arg expansion which TorchScript JIT does not support. Despite e3nn's @compile_mode("script") decorator, the call site is not scriptable. Both PyTorch compiler paths (torch.compile + torch.jit.script) are now definitively dead for HydraGNN MLIP MACE as shipped. * rccl_high_priority -> status: accepted-baked-in. The only lever that produced a real improvement across 3 loops (-1.3% epoch, -8.4% energy in loop-c3e4df1c iter-2). Now part of the baseline contract. * batch_800 -> NODE-CLASS-DEPENDENT regressor. Cross-loop evidence: +56.8% on a-nodes (loop-43b33ec1), +87.2% on b-nodes (loop-e3fac2af). expected_payoff updated to "none-for-MLIP-MACE". * tunable_op_warmup_then_use -> status: blocked. Even the 2-phase split fails — the warmup phase itself triggers GPU mem faults in hipBLASLt tuning. Both TunableOp levers (live and warmup) are now blocked. * nccl_minchannels -> confirmed regressor at N=2 (+10.3% in loop-e3fac2af iter-3). expected_payoff updated; lever_picker should drop entirely for N<8 runs. * precision_fp32 -> repurposed as diagnostic-realized. The +1.3% result (within noise) IS the project's key dispatch-vs-compute discriminator. fp32 verified genuinely active via precision-diagnostic rank-0 logs (model_param_dtype=torch.float32). Catalog notes now explain the dispatch-bound conclusion and point to SKILL.md lesson #13. Skill change (.cursor/skills/ai4science-perf-analysis/SKILL.md): * +Lesson #13: fp32-vs-fp64 epoch-time comparison is the canonical low-cost dispatch-vs-compute discriminator on MI355X (fp32 peak is 2x fp64 peak). Includes the three-branch decision tree for next-lever choice (≈ → dispatch-bound, ≪ → compute-bound, > → memory-bound). * +Lesson #14: torch.jit.script also blocked for HydraGNN MLIP MACE. Both PyTorch compiler paths are dead until upstream fixes. * +Lesson #15: TunableOp _TUNING=1 unsafe regardless of phase (live or warmup-then-use). Both lever variants now blocked. * +Lesson #16: b-nodes hit the dispatch ceiling at the baseline. Every compute-side lever tested on b-nodes regressed or was neutral. a-nodes still had headroom for num_workers_12. lever_picker must encode node-class state. * +Lesson #17: Current cross-loop best (75 s/epoch) decomposes as baseline-on-b-nodes (76 s, free) + rccl_high_priority (-1.3%). Single accepted lever across 3 loops + 1 investigation. Runtime artifacts (story.md, foms.csv, STATUS.txt, hooks, SESSION_REPORT.md, pathc-result.md) remain under /shared/aaji/models/HydraGNN/perf-runs/ and are not committed; the catalog citation fields reference them by absolute path. Co-authored-by: Cursor --- .../skills/ai4science-perf-analysis/SKILL.md | 13 +++ .../perf-optimizer-loop/lever_catalog.yaml | 88 ++++++++++++++----- 2 files changed, 81 insertions(+), 20 deletions(-) diff --git a/.cursor/skills/ai4science-perf-analysis/SKILL.md b/.cursor/skills/ai4science-perf-analysis/SKILL.md index 4e4ed07..c14bc77 100644 --- a/.cursor/skills/ai4science-perf-analysis/SKILL.md +++ b/.cursor/skills/ai4science-perf-analysis/SKILL.md @@ -87,4 +87,17 @@ The first 2-node end-to-end run on Lux exposed five real bugs that are now all w 11. **`hydragnn/train/__init__.py` re-export shadow** breaks naive `import hydragnn.train.train_validate_test as tvt`. The line `from .train_validate_test import train_validate_test, train, ...` makes `hydragnn.train.train_validate_test` resolve to the **function** `train_validate_test` (re-exported as attribute of the package), not the module. Use `importlib.import_module("hydragnn.train.train_validate_test")` instead. Applies to any rank hook that needs to monkey-patch `train_validate_test.train`. 12. **`sbatch --export=ALL,...` does NOT include shell vars that are not in `--export`'s explicit list.** When using `--export=ALL,K1=V1,K2=V2`, vars set in the calling shell but not listed here are propagated only if they're in the user's persistent env (login shell init). Always explicitly list cluster-pathing vars like `AI4S_SHARED_DIR=/shared/aaji` in `--export`. +## Lessons captured (pathc-pathb-a3b8f3f0 + loop-e3fac2af, 2026-05-23) + +13. **fp32-vs-fp64 epoch-time comparison is the canonical low-cost dispatch-vs-compute discriminator.** On gfx950 (MI355X), the fp32 MFMA peak is **2× the fp64 MFMA peak**. A compute-bound workload should show ≥30% speedup when switching fp64→fp32. HydraGNN GFM-MLIP MACE on b-nodes showed ~0% speedup (78→79 s, +1.3% within noise, fp32 verified by `precision-diagnostic` rank-0 logs `model_param_dtype=torch.float32 first_batch_float_dtype=torch.float32`). This proves the workload is dispatch-bound. Run this comparison ONCE per workload-class as a single sbatch; the verdict shapes the entire subsequent lever strategy: + - **fp32 ≈ fp64 → dispatch-bound** → compute-side levers (`batch_*`, `torch_compile_*`, `tunable_op_*`, `precision_fp32`) are dead ends. Focus on kernel-trace analysis to identify the dispatch culprit, then either upstream kernel fusion or reduce launch count (larger blocks, mega-kernels). + - **fp32 ≪ fp64 → compute-bound** → standard compute-side levers apply. fp32 itself becomes a real lever (within numerical tolerance). + - **fp32 > fp64 → memory-bound** → bandwidth-targeting levers (batch packing, fewer kernel launches over more data) apply. +14. **`torch.jit.script` is ALSO blocked for HydraGNN MLIP MACE** (in addition to `torch.compile`). Root cause: `MACEStack.py:397` uses `**conv_args` keyword-arg expansion when calling the interaction block; TorchScript JIT does not support `**kwargs` expansion (`NotSupportedError: keyword-arg expansion is not supported`). Despite the upstream `@compile_mode("script")` decorator from e3nn, the actual call site is not scriptable. **Both PyTorch compiler paths are now definitively dead** for HydraGNN MLIP MACE as shipped — unblocking either requires upstream source changes: + - `torch.compile` path: move `energy_force_loss` into `EnhancedModelWrapper.forward()` (~30-50 LOC; addresses AOT-double-backward; see lesson #6) + - `torch.jit.script` path: unpack `**conv_args` at `MACEStack.py:397` (small diff but in the read-only `/opt/hydragnn-pkgs` overlay) +15. **TunableOp `_TUNING=1` is unsafe regardless of phase** on MI355X / ROCm 7.2.2 / PyTorch 2.10. Both the live-tuning lever (`tunable_op_live`, loop-c3e4df1c iter-3) and the warmup-then-use 2-phase pattern (`tunable_op_warmup_then_use`, loop-e3fac2af iter-2) hit the same `Memory access fault by GPU node-N` on all 16 GPUs during the hipBLASLt tuning routine. The 2-phase split doesn't help — the fault is in the tuning routine itself, not in the production consumption. **Both levers are marked `status: blocked` in the catalog.** Re-test only after a ROCm stack update with a documented hipBLASLt fix. +16. **b-nodes hit the dispatch ceiling at the baseline.** Loop-43b33ec1 (a-nodes) and loop-c3e4df1c (b-nodes) baselines are 92.5 s and 76.0 s respectively for the same code/config — a 17% baseline gain just from node class. But the b-node baseline appears to be at the dispatch ceiling: every compute-side lever tested on b-nodes (batch_800, num_workers_12, precision_fp32) regressed or was neutral. a-nodes still had headroom for at least one lever (num_workers_12, -12.4%). The lever_picker must encode node-class state and not apply a-node-validated levers to b-node best configs without re-evaluation. +17. **Cross-loop best is composed of a free finding + a small lever**: current best (75.0 s/epoch) = baseline-on-b-nodes (76.0 s, free, just node-class allocation) + rccl_high_priority (-1.3%). The single accepted lever across 3 loops + 1 investigation is `rccl_high_priority` (`TORCH_NCCL_HIGH_PRIORITY=1 GPU_MAX_HW_QUEUES=2`); it should be baked into the baseline contract for HydraGNN GFM-MLIP going forward (now marked `status: accepted-baked-in` in the catalog). + The full list (with cross-cutting context) is at the end of [.cursor/skills/ai4science-studio/SKILL.md](../ai4science-studio/SKILL.md). When fixing a bug here, check there first and propagate the lesson. diff --git a/material_science/models/HydraGNN/recipes/perf-optimizer-loop/lever_catalog.yaml b/material_science/models/HydraGNN/recipes/perf-optimizer-loop/lever_catalog.yaml index 3cd8b2e..c6ee15d 100644 --- a/material_science/models/HydraGNN/recipes/perf-optimizer-loop/lever_catalog.yaml +++ b/material_science/models/HydraGNN/recipes/perf-optimizer-loop/lever_catalog.yaml @@ -46,17 +46,17 @@ levers: notes: "Reference state. Establishes the FOM denominator for all subsequent iters." - id: batch_800 - description: "Double batch size: HG_BATCH_SIZE 400 -> 800. R2 #1 next-lever recommendation; VRAM headroom is 24/288 GB." + description: "Double batch size: HG_BATCH_SIZE 400 -> 800. R2 #1 next-lever recommendation; VRAM headroom is 24/288 GB. NODE-CLASS-DEPENDENT: regresses on both a-nodes (+56.8%) and b-nodes (+87.2%) for HydraGNN GFM-MLIP MACE." catalog_order: 1 kind: env_only env_vars: HG_BATCH_SIZE: "800" diagnostic_only: false - expected_payoff: "high" - risk: "low" + expected_payoff: "none-for-MLIP-MACE" + risk: "high" revert_method: "drop_env" - citation: "/shared/aaji/models/HydraGNN/perf-runs/6985/combined_report.md" - notes: "R2 showed 200->400 gave +1.80x util and +2.41x MFMA. Diminishing returns expected; if delta < 30%, move on." + citation: "loop-43b33ec1 iter-2 (a-nodes, +56.76%) + loop-e3fac2af iter-1 (b-nodes, +87.2%)" + notes: "The lever_picker should NOT pick batch_800 for HydraGNN MLIP MACE on either node class. Cross-loop evidence shows the MACE equivariant message-passing stack saturates dispatch at batch=400 and doubling work-per-iter doubles wall-time without proportional GPU occupancy gain. The R2 recommendation was based on an earlier non-MLIP HydraGNN config — DOES NOT apply to gfm_mlip_all_mpnn.py. Keep entry as cautionary record. May still be worth re-trying ONLY if upstream HydraGNN fuses small e3nn kernels first." - id: torch_compile_e3nn description: "BLOCKED — torch.compile is fundamentally incompatible with HydraGNN GFM-MLIP under PyTorch 2.10 / ROCm 7.2.2. Documented here so the lever_picker NEVER re-tries it." @@ -110,19 +110,63 @@ levers: citation: "/shared/aaji/models/HydraGNN/perf-runs/compile-investigation-aa480554-2a15-4e04-80dd-c58afc672a8d/INVESTIGATION_REPORT.md" notes: "DO NOT pick this lever until HydraGNN upstream patches energy_force_loss into forward(). The lever_picker subagent must treat status:blocked as equivalent to do_not_retry." + - id: torch_jit_script + description: "BLOCKED — torch.jit.script on the inner MACEStack also fails. Both PyTorch compiler paths (AOT-Autograd via torch.compile AND legacy TorchScript JIT via torch.jit.script) are dead-ends for HydraGNN MLIP MACE as shipped." + catalog_order: 2 + kind: rank_script_patch + status: "blocked" + blocked_reason: "**kwargs-expansion-not-supported-by-TorchScript-JIT" + blocked_evidence: "pathc-pathb-a3b8f3f0 (2026-05-23, job 7198): hook fired and caught NotSupportedError: keyword-arg expansion is not supported at hydragnn/models/MACEStack.py:397 (the interaction-block call site uses **conv_args)." + env_vars: + HG_JIT_SCRIPT: "1" + rank_script_patch: | + # KNOWN-BROKEN — kept for reference only. Do NOT enable this lever. + # + # Root cause: torch.jit.script does not support **kwargs expansion. + # hydragnn/models/MACEStack.py:397 has: + # self.interaction_block( + # inv_node_feat=inv_node_feat, + # equiv_node_feat=equiv_node_feat, + # **conv_args, + # ) + # The legacy TorchScript JIT errors with: + # NotSupportedError: keyword-arg expansion is not supported + # + # Despite MACEStack being decorated @compile_mode("script") from + # e3nn, that decorator is aspirational — not enforced. The actual + # call sites are not scriptable. + # + # To unblock requires upstream HydraGNN source change: unpack + # conv_args into explicit keyword arguments at MACEStack.py:397 + # (and likely sibling call sites). Small diff, but still in + # /opt/hydragnn-pkgs which is read-only inside the container. + # + # Combined with torch_compile_e3nn (BLOCKED for double-backward), + # this means BOTH PyTorch compiler paths are dead for HydraGNN MLIP + # MACE as currently shipped. Kernel-fusion wins on the 1.3M + # aten::bmm calls cannot be unlocked from a runtime hook alone. + raise RuntimeError("torch_jit_script is BLOCKED for HydraGNN MACE; see lever_catalog.yaml") + diagnostic_only: false + expected_payoff: "none (blocked)" + risk: "blocked" + revert_method: "drop_hook" + citation: "/shared/aaji/models/HydraGNN/perf-runs/pathc-pathb-a3b8f3f0-19be-48b8-b81f-c7068e418289/pathc-result.md" + notes: "DO NOT pick this lever until HydraGNN upstream unpacks **conv_args at MACEStack.py:397. catalog_order=2 shared with torch_compile_e3nn since both are 'kernel-fusion via PyTorch compiler' attempts that are both dead-ends on the current stack." + - id: rccl_high_priority - description: "RCCL stream priority and HIP hardware queue count: TORCH_NCCL_HIGH_PRIORITY=1, GPU_MAX_HW_QUEUES=2. Improves compute<->RCCL overlap." + description: "RCCL stream priority and HIP hardware queue count: TORCH_NCCL_HIGH_PRIORITY=1, GPU_MAX_HW_QUEUES=2. Improves compute<->RCCL overlap. CONFIRMED WIN — baked into baseline going forward." catalog_order: 3 kind: env_only + status: "accepted-baked-in" env_vars: TORCH_NCCL_HIGH_PRIORITY: "1" GPU_MAX_HW_QUEUES: "2" diagnostic_only: false - expected_payoff: "medium" + expected_payoff: "1.3%-realized" risk: "low" revert_method: "drop_env" - citation: "https://rocm.docs.amd.com/en/latest/how-to/tuning-guides/mi300x/workload.html" - notes: "R2 already showed exposed NCCL was only 3% of fused-step wall; expected payoff at N=2 is small. Pick after compute-side levers (batch_800, torch_compile) are exhausted." + citation: "loop-c3e4df1c iter-2 / jobid 7187: -1.3% epoch_time (76.0->75.0 s) AND -8.4% energy/sample (1.781->1.631 J/sample). Mean power dropped 4885W->4552W." + notes: "ACCEPTED. The only lever across 3 loops that produced a real improvement. Now part of the baseline contract for HydraGNN GFM-MLIP on b-nodes. lever_picker should treat this as 'already applied' and skip it in candidate set (foms.csv has it accepted with delta_pct=-1.316)." - id: tunable_op_live description: "BLOCKED on this stack — TunableOp LIVE tuning mode (PYTORCH_TUNABLEOP_TUNING=1) causes GPU memory access faults on MI355X / ROCm 7.2.2 / PyTorch 2.10 with HydraGNN. Use tunable_op_warmup_then_use instead." @@ -143,8 +187,11 @@ levers: notes: "Live tuning racing with training data is unsafe here. Pick tunable_op_warmup_then_use instead, which separates the tuning phase into a dedicated short job." - id: tunable_op_warmup_then_use - description: "2-phase safe TunableOp pattern: a dedicated 1-batch warmup job generates tunableop_results_.csv files, then the actual training job runs with PYTORCH_TUNABLEOP_ENABLED=1 (no _TUNING) to consume the pre-tuned results." + description: "BLOCKED on this stack — even the 2-phase warmup pattern fails. The warmup phase (PYTORCH_TUNABLEOP_TUNING=1, 5 batches) itself triggers GPU memory access faults on all 16 GPUs before any CSV file can be written." catalog_order: 4 + status: "blocked" + blocked_reason: "GPU-mem-access-fault-also-during-warmup-phase" + blocked_evidence: "loop-e3fac2af iter-2 (job 7201, 2026-05-23): warmup with HG_NUM_EPOCH=1 HYDRAGNN_MAX_NUM_BATCH=5 PYTORCH_TUNABLEOP_TUNING=1 reproduces the same hipBLASLt 'Memory access fault by GPU node-N' on all 16 GPUs as the live-tuning lever. No tunableop_results_*.csv files were written. The 2-phase split doesn't help — the fault is in hipBLASLt's tuning routine itself, regardless of duration." kind: rank_script_patch env_vars: PYTORCH_TUNABLEOP_ENABLED: "1" @@ -195,30 +242,31 @@ levers: notes: "lever_picker should only consider this lever on a-nodes. On b-nodes, num_workers=8 is the sweet spot (verified in iter-2 of loop c3e4df1c with epoch_time=75 s, the current best). If the kernel_correlation.csv from fom_extractor shows dataloader prefetch gaps > 10 ms on b-nodes, try num_workers=10 (not 12)." - id: nccl_minchannels - description: "NCCL_MIN_NCHANNELS=112: forces more channels for collective dispatch, helps when message size is moderate (HydraGNN allreduces are 11.3 MB)." + description: "NCCL_MIN_NCHANNELS=112: forces more channels for collective dispatch, helps when message size is moderate. CONFIRMED REGRESSOR at N=2 (+10.3%)." catalog_order: 6 kind: env_only env_vars: NCCL_MIN_NCHANNELS: "112" diagnostic_only: false - expected_payoff: "low" - risk: "low" + expected_payoff: "negative-at-N=2" + risk: "medium" revert_method: "drop_env" - citation: "https://rocm.docs.amd.com/en/docs-7.0.1/how-to/rocm-for-ai/system-setup/multi-node-setup.html" - notes: "Mostly relevant at N >= 8 nodes. Low priority for the 2-node loop; included for completeness." + citation: "loop-e3fac2af iter-3 (job 7202): +10.3% epoch_time (78->86 s)" + notes: "Confirmed bad at N=2. Forcing 112 channels increases synchronization overhead vs the default at small node count. The lever_picker should DROP this lever entirely for N<8 runs. Keep entry for future >=8-node runs only." - id: precision_fp32 - description: "HG_PRECISION fp64 -> fp32. gfx950 fp32 MFMA peak (~157 TFLOP/s) is ~2x fp64 (~78.6 TFLOP/s). Numerical floor for MACE-style equivariance." + description: "HG_PRECISION fp64 -> fp32. CONFIRMED NEUTRAL — but the very fact that it's neutral is the project's most important diagnostic finding (fp32 ≈ fp64 means the workload is DISPATCH-BOUND, not compute-bound). Repurpose this lever as diagnostic-only going forward." catalog_order: 7 kind: env_only + status: "diagnostic-realized" env_vars: HG_PRECISION: "fp32" diagnostic_only: false - expected_payoff: "medium" - risk: "medium" + expected_payoff: "negligible-but-diagnostically-valuable" + risk: "low" revert_method: "drop_env" - citation: "local + gfx950 spec sheet" - notes: "If the run is still dispatch-bound (mean util < 20%), fp32 wont move the needle much; orchestrator should ONLY pick this lever when fom_extractor reports the most recent best-iter is compute-bound (mfma_tflops > 30% of fp64 peak in best 10-s burst). bf16 is INTENTIONALLY EXCLUDED from this catalog: it breaks equivariance for MACE-class features." + citation: "loop-e3fac2af iter-4 (job 7203): +1.3% epoch_time (78->79 s) confirmed in fp32 (precision-diagnostic logs show model_param_dtype=torch.float32, first_batch_float_dtype=torch.float32)" + notes: "RAN GENUINELY IN FP32 (verified by precision-diagnostic). fp32 MFMA peak on gfx950 is 2x fp64 peak, so a compute-bound workload would show ~50% speedup. Observed ~0% speedup proves HydraGNN GFM-MLIP MACE on b-nodes is dispatch-bound. lever_picker should NOT re-pick this lever unless the run is explicitly switching to a-nodes (where the compute/dispatch balance may differ). bf16 remains INTENTIONALLY EXCLUDED: breaks equivariance for MACE-class features. The fp32-vs-fp64 comparison is now the canonical low-cost dispatch-vs-compute discriminator for any future HydraGNN-class workload — see ai4science-perf-analysis/SKILL.md lesson #13." - id: kernel_trace_diag_only description: "DIAGNOSTIC: enable OMNISTAT_KERNEL_TRACE=1 for ONE iteration to capture per-kernel-name dispatch counts/durations. Fires only when the 1-s TraceLens<->Omnistat correlation cannot attribute the dominant kernel." From 864a5ae9fe6c24e0369c34d3f057c77faf60370d Mon Sep 17 00:00:00 2001 From: Ashwin Aji Date: Tue, 26 May 2026 08:13:07 +0000 Subject: [PATCH 11/31] perf-optimizer-loop: correct the a-vs-b narrative + add node-health microbench MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Direct microbench on 12× lux-mi355x-a + 3× lux-mi355x-b (CPU/NUMA/kernel inventory, HIP launch latency, NCCL single-node all_reduce, STREAM dual-NUMA, 1-node and 2-node HydraGNN sanity) shows a-nodes and b-nodes are physically identical and the previously reported 17% gap was confounded — all three "baseline" jobids ran on the same lux-mi355x-b[1-2] pair (verified via nodes_list field). a-vs-b 2-node sanity now: a3+a4 = 77.0 s/epoch vs b3+b4 = 78.4 s/epoch (1.8% a-pair faster, within noise). One real per-node outlier: lux-mi355x-a5 — intermittent dual-NUMA STREAM COPY at 179 GB/s vs cluster median 220 GB/s, recovers partially on retest. Captured for the sysadmin alongside a6 (mount EIO) and a10 (PMIx ring). Changes: - examples/microbench_node_health.sh — new reusable node-health probe (mount-write + STREAM dual-NUMA + HIP launch latency + GPU/firmware inventory). ~30 s wall-clock, exit nonzero on failure — suitable for SLURM Prolog use. - recipes/perf-optimizer-loop/lever_catalog.yaml — revise 4 entries (batch_800, rccl_high_priority, num_workers_12, precision_fp32) marked with "[REVISED 2026-05-26]" to remove the a-vs-b framing; the observations themselves are intact. - .cursor/skills/ai4science-perf-analysis/SKILL.md — mark lessons 9, 10, 16 as REVISED and add lessons 18 (the a-vs-b false-narrative root cause) and 19 (the per-node prologue probe argument). Add a5 to the Confirmed bad nodes list. Runtime artefacts (REPORT.md + SYSADMIN_MESSAGE.md + raw probe outputs) live under /shared/aaji/microbench-a-vs-b/ and are not committed. Co-authored-by: Cursor --- .../skills/ai4science-perf-analysis/SKILL.md | 22 +- .../examples/microbench_node_health.sh | 225 ++++++++++++++++++ .../perf-optimizer-loop/lever_catalog.yaml | 18 +- 3 files changed, 251 insertions(+), 14 deletions(-) create mode 100755 material_science/models/HydraGNN/examples/microbench_node_health.sh diff --git a/.cursor/skills/ai4science-perf-analysis/SKILL.md b/.cursor/skills/ai4science-perf-analysis/SKILL.md index c14bc77..c38a59f 100644 --- a/.cursor/skills/ai4science-perf-analysis/SKILL.md +++ b/.cursor/skills/ai4science-perf-analysis/SKILL.md @@ -45,7 +45,8 @@ Default target: HydraGNN on Lux MI355X. The same pattern can extend to ORBIT-2 / - System-mode Omnistat is running on every compute node at port 8001. Don't touch it; we run our own user-mode collector alongside. - The login node (`rad-vultr-login`) has no GPU and no MPI. All analysis runs after the job — no live profiling on the login node. - **Node-specific mount faults are real and silent.** Any single compute node can land an allocation with a broken per-user autofs/NFS mount of `/home/$USER` or `/shared/$USER`, while sibling nodes in the same allocation are fine. Symptom in slurmd logs on the broken node: `Home Directory for Not Found ... Please contact system admin` and `lstat /shared/...: no such file or directory`; the surviving nodes wedge in a `pmix_coll_ring` collective fence until SLURM kills the job at the wall. - - **Confirmed bad nodes** (as of 2026-05-23, until cluster admin repairs them): + - **Confirmed bad nodes** (as of 2026-05-26, until cluster admin repairs them): + - `lux-mi355x-a5` — intermittent dual-NUMA STREAM bandwidth degradation (179 GB/s COPY vs cluster median 220 GB/s) with 8% run-to-run jitter; partial recovery on retest. Likely DIMM/IF or thermal issue on one socket. Found 2026-05-26 microbench (jobs 7589 + 7601). Mount + GPU healthy; only the CPU memory subsystem is degraded. - `lux-mi355x-a6` — `/home/aaji` missing, `/shared/aaji/images/*.sif` unreadable; sibling `lux-mi355x-a1` healthy in the same alloc. Loop-43b33ec1 iter-1 / job 7088. - `lux-mi355x-a10` — pmix ring collective times out at MPI init; `opal_mpi_init_ext3x_client_unreachable`. NOT a mount fault — broken PMIx/OpenMPI install on this node. Loop-c3e4df1c iter-1 / job 7161. - **Diagnostic rule:** never blacklist a node *class* (e.g. "avoid all a-nodes") based on a single failure. Always probe the *specific* allocated nodes and exclude only those that fail by name. Probe is a one-task-per-node srun that checks `/home/$USER`, `/shared/$USER`, and the SIF path; see the `Per-node mount-health probe` section of `sbatch_train_perf_amd.sh` (exit 42 on fail). The perf-optimizer-loop orchestrator handles exit 42 by re-submitting with `--exclude=` and adding the node to `loop-/known_bad_nodes.txt`; the lever is NOT penalized. @@ -82,8 +83,8 @@ The first 2-node end-to-end run on Lux exposed five real bugs that are now all w 6. **`torch.compile` is BLOCKED for HydraGNN GFM-MLIP MACE** under PyTorch 2.10 / ROCm 7.2.2. Six hooks across three hypotheses (inner-model compile, outer + `dynamo.allow_in_graph(autograd.grad)`, single-submodule compile) all fail with the same verbatim PyTorch error: `RuntimeError: torch.compile with aot_autograd does not currently support double backward` (`torch/_functorch/_aot_autograd/runtime_wrappers.py:2356`). Root cause: HydraGNN's `energy_force_loss` (`hydragnn/models/create.py:718`) calls `torch.autograd.grad(graph_energy_pred, data.pos, create_graph=True)` to compute forces as derivatives of energy w.r.t. positions; with `create_graph=True`, training requires DOUBLE BACKWARD (loss.backward differentiates through forces). Every torch.compile mode (`default`, `reduce-overhead`, `max-autotune`) goes through AOT Autograd which pre-compiles the backward as a joint program and cannot itself be differentiated. **To unblock requires upstream HydraGNN source change**: move `energy_force_loss` INTO `EnhancedModelWrapper.forward()` so `autograd.grad` is inside the compiled forward (the canonical MACE/Allegro/NequIP pattern). ~30-50 LOC across `create.py:622-758` + `train_validate_test.py:725-735`, config-gated. The lever `torch_compile_e3nn` is marked `status: blocked` in `lever_catalog.yaml` with full evidence; the lever_picker drops blocked levers from its candidate set. 7. **`dynamo.allow_in_graph(autograd.grad)` is weaker than MACE upstream docs suggest.** It only prevents dynamo from TRACING THROUGH the call; it does NOT prevent dynamo's FakeTensor symbolic-execution pass from running `autograd.grad` against fake tensors (which fails with the `allow_unused=True` error because fake tensors have no real autograd graph). The MACE `prepare()` pattern works only when `autograd.grad` lives inside the compiled function's `forward()` — not when it's in a separate method called from outside. 8. **TunableOp live tuning (`PYTORCH_TUNABLEOP_TUNING=1`) is unsafe on MI355X / ROCm 7.2.2 with HydraGNN.** All 16 GPUs hit `Memory access fault by GPU node-N` during the hipBLASLt kernel autotuning phase (loop-c3e4df1c iter-3 / job 7188). Use a 2-phase pattern instead: dedicated warmup sbatch with `PYTORCH_TUNABLEOP_TUNING=1` to generate `tunableop_results_.csv` files, then production runs with `PYTORCH_TUNABLEOP_ENABLED=1` only (no `_TUNING=1`). The `tunable_op_warmup_then_use` lever in the catalog encodes this. The original `tunable_op_live` lever is now marked `status: blocked`. -9. **Lever payoff is node-class-dependent.** `num_workers_12` improved a-nodes (-12.4% in loop-43b33ec1) but regressed b-nodes (+10.7% in loop-c3e4df1c). Lever_picker must read `SLURM_NODELIST` from the current best run's manifest and treat a-nodes vs b-nodes as separate state for at least `num_workers_*` and `batch_*` levers. -10. **b-nodes are ~17% faster baseline than a-nodes for HydraGNN.** Loop-43b33ec1's a-node baseline was 92.5 s/epoch; loop-c3e4df1c's b-node baseline (same code, same config, same fp64, same batch_size=400) was 76.0 s/epoch. The single biggest finding of the night came from a baseline shift, not a lever. The `launcher` subagent should report which node class was allocated, and cross-loop FOM comparisons must control for this. +9. **[REVISED — see lesson #18]** ~~Lever payoff is node-class-dependent.~~ `num_workers_12` improved one b1+b2 run (-12.4% in loop-43b33ec1) and regressed a later b1+b2 run (+10.7% in loop-c3e4df1c). Both ran on the *same* b-pair; the apparent class dependence was run-to-run variance on b1+b2, not a-vs-b. The lever picker should treat `num_workers_*` results as high-variance until replicated within ±5% on the same node pair. +10. **[REVISED — see lesson #18]** ~~b-nodes are ~17% faster baseline than a-nodes for HydraGNN.~~ Both compared baselines (loop-43b33ec1 at 92.5 s/epoch and loop-c3e4df1c at 76.0 s/epoch) ran on the same `lux-mi355x-b[1-2]` pair. The 17% gap is a run-to-run delta on a single pair, not a class difference. Always read `nodes_list` from `iter-0-baseline.json` before drawing any node-class conclusion from a FOM delta. 11. **`hydragnn/train/__init__.py` re-export shadow** breaks naive `import hydragnn.train.train_validate_test as tvt`. The line `from .train_validate_test import train_validate_test, train, ...` makes `hydragnn.train.train_validate_test` resolve to the **function** `train_validate_test` (re-exported as attribute of the package), not the module. Use `importlib.import_module("hydragnn.train.train_validate_test")` instead. Applies to any rank hook that needs to monkey-patch `train_validate_test.train`. 12. **`sbatch --export=ALL,...` does NOT include shell vars that are not in `--export`'s explicit list.** When using `--export=ALL,K1=V1,K2=V2`, vars set in the calling shell but not listed here are propagated only if they're in the user's persistent env (login shell init). Always explicitly list cluster-pathing vars like `AI4S_SHARED_DIR=/shared/aaji` in `--export`. @@ -97,7 +98,18 @@ The first 2-node end-to-end run on Lux exposed five real bugs that are now all w - `torch.compile` path: move `energy_force_loss` into `EnhancedModelWrapper.forward()` (~30-50 LOC; addresses AOT-double-backward; see lesson #6) - `torch.jit.script` path: unpack `**conv_args` at `MACEStack.py:397` (small diff but in the read-only `/opt/hydragnn-pkgs` overlay) 15. **TunableOp `_TUNING=1` is unsafe regardless of phase** on MI355X / ROCm 7.2.2 / PyTorch 2.10. Both the live-tuning lever (`tunable_op_live`, loop-c3e4df1c iter-3) and the warmup-then-use 2-phase pattern (`tunable_op_warmup_then_use`, loop-e3fac2af iter-2) hit the same `Memory access fault by GPU node-N` on all 16 GPUs during the hipBLASLt tuning routine. The 2-phase split doesn't help — the fault is in the tuning routine itself, not in the production consumption. **Both levers are marked `status: blocked` in the catalog.** Re-test only after a ROCm stack update with a documented hipBLASLt fix. -16. **b-nodes hit the dispatch ceiling at the baseline.** Loop-43b33ec1 (a-nodes) and loop-c3e4df1c (b-nodes) baselines are 92.5 s and 76.0 s respectively for the same code/config — a 17% baseline gain just from node class. But the b-node baseline appears to be at the dispatch ceiling: every compute-side lever tested on b-nodes (batch_800, num_workers_12, precision_fp32) regressed or was neutral. a-nodes still had headroom for at least one lever (num_workers_12, -12.4%). The lever_picker must encode node-class state and not apply a-node-validated levers to b-node best configs without re-evaluation. -17. **Cross-loop best is composed of a free finding + a small lever**: current best (75.0 s/epoch) = baseline-on-b-nodes (76.0 s, free, just node-class allocation) + rccl_high_priority (-1.3%). The single accepted lever across 3 loops + 1 investigation is `rccl_high_priority` (`TORCH_NCCL_HIGH_PRIORITY=1 GPU_MAX_HW_QUEUES=2`); it should be baked into the baseline contract for HydraGNN GFM-MLIP going forward (now marked `status: accepted-baked-in` in the catalog). +16. **[REVISED — see lesson #18]** ~~b-nodes hit the dispatch ceiling at the baseline.~~ All three loop baselines (jobids 7083, 7158, 7199) ran on `lux-mi355x-b[1-2]`. The "dispatch ceiling" framing is correct *for that specific node pair on that specific dataset*; whether it generalizes to a class is a question we now know we cannot answer from the loop data alone. What does still hold: every compute-side lever tested on b1+b2 (batch_800, num_workers_12, precision_fp32) regressed or was neutral, which on its own confirms dispatch-bound character — but not a node-class story. +17. **Cross-loop best is composed of a free finding + a small lever**: current best (75.0 s/epoch) = baseline-on-b1-b2 (76.0 s) + `rccl_high_priority` (-1.3%). The single accepted lever across 3 loops + 1 investigation is `rccl_high_priority` (`TORCH_NCCL_HIGH_PRIORITY=1 GPU_MAX_HW_QUEUES=2`); it should be baked into the baseline contract for HydraGNN GFM-MLIP going forward (now marked `status: accepted-baked-in` in the catalog). + +## Lessons captured (a-vs-b parity microbench, 2026-05-26) + +18. **The "a-nodes vs b-nodes 17% gap" was a confounded measurement, not a hardware/firmware story.** A 4-test microbench sweep across 12 a-nodes + 3 b-nodes (script: [`examples/microbench_node_health.sh`](../../../material_science/models/HydraGNN/examples/microbench_node_health.sh); raw outputs: `/shared/aaji/microbench-a-vs-b/`) showed: + - **Identical at the host level:** CPU model (Genoa, 128c/256t), kernel (6.8.0-110), NUMA layout (2× ~1.5 TB), NICs (6× ionic 400G + 2× bnxt 200G), GPUs (8× MI355X gfx950), VBIOS, firmware, PCIe link state, mounts, scaling governor. + - **Identical compute microbench:** HIP empty-kernel launch latency 3.5–3.8 µs/launch on every GPU of every node; H2D pinned 57.6 GB/s; NCCL single-node 8-GPU all_reduce 372 GB/s bus; STREAM single-NUMA TRIAD 332–336 GB/s. + - **One real per-node outlier:** `lux-mi355x-a5` showed degraded dual-NUMA STREAM COPY (179 GB/s vs cluster median 220 ± 3 GB/s) with 8% run-to-run jitter on the first probe and partial recovery on retest (219 GB/s COPY but anomalously high TRIAD variance). All other 11 a-nodes and 3 b-nodes are statistically equivalent (220 ± 3 GB/s). + - **Controlled 2-node sanity test:** `a3+a4` = 77.0 s/epoch vs `b3+b4` = 78.4 s/epoch with the same code/config; **a-pair is 1.8% faster**, well inside noise. + - **Root cause of the original 17% narrative:** all three "baseline" jobids (7083, 7158, 7199) ran on `lux-mi355x-b[1-2]` per their `nodes_list` field. The 92.5 s → 76 s spread was on the *same pair* on *different days*; likely cache/NFS-warmup variance, not class. The misattribution came from labeling loop-43b33ec1 "a-node" because its iter-1 *failed* on a6's broken mount — but the *successful* baseline iter-0 (with a6 excluded) landed on b-nodes via SLURM scheduling. + - **Action for future loops:** always `cat $LOOP_DIR/iter-0-baseline.json | jq .nodes_list` before drawing any class conclusion. The `launcher` and `synthesizer` subagents should both log the assigned nodelist front-and-center in their output. +19. **Per-node health regressions are real and worth a 30-second prologue probe.** In 2 weeks we hit three different per-node failure modes (a6: container bind-mount EIO; a10: PMIx/NCCL ring timeout; a5: NUMA bw degradation). Each manifested as a confusing application-level failure hours into a job. All three would have been caught at job start by a 30-s SLURM Prolog running [`microbench_node_health.sh`](../../../material_science/models/HydraGNN/examples/microbench_node_health.sh) (mount-write probe + STREAM dual-NUMA + NCCL single-node + HIP launch latency). Suggested to the cluster admin in `SYSADMIN_MESSAGE.md` (under `/shared/aaji/microbench-a-vs-b/`). The full list (with cross-cutting context) is at the end of [.cursor/skills/ai4science-studio/SKILL.md](../ai4science-studio/SKILL.md). When fixing a bug here, check there first and propagate the lesson. diff --git a/material_science/models/HydraGNN/examples/microbench_node_health.sh b/material_science/models/HydraGNN/examples/microbench_node_health.sh new file mode 100755 index 0000000..474efb9 --- /dev/null +++ b/material_science/models/HydraGNN/examples/microbench_node_health.sh @@ -0,0 +1,225 @@ +#!/usr/bin/env bash +# microbench_node_health.sh +# +# Per-node hardware/firmware/health survey for AMD MI355X (gfx950) nodes on +# Vultr Lux. Runs four micro-tests in ~30 wall-clock seconds: +# 1. host inventory - CPU, NUMA, kernel, NICs, PCIe link, mounts, /tmp dd +# 2. GPU + driver inventory - rocminfo, rocm-smi (vbios/fw/topo), PCIe link state +# 3. CPU dual-NUMA STREAM - COPY/SCALE/ADD/TRIAD, 128 threads, NUMA-interleaved +# 4. HIP kernel launch latency - 100k empty-tensor kernels per GPU +# +# Designed to be: +# - safe to run in a SLURM Prolog +# - cheap enough to run on every job (~30s) +# - detailed enough to spot the three failure modes we hit in May 2026: +# a) NUMA bandwidth degradation (one DIMM/socket regression) +# b) container bind-mount fault (NFS / overlay regression) +# c) GPU dispatch-rate regression (driver/SMI hang) +# +# Usage: +# sbatch --nodelist= material_science/models/HydraGNN/examples/microbench_node_health.sh +# srun --nodelist= material_science/models/HydraGNN/examples/microbench_node_health.sh +# +# Output: writes ${OUT_DIR}// with one file per test, +# and one summary line of pass/fail to stdout. +# +# Env vars (all optional, with defaults): +# OUT_DIR (default: /shared/$USER/microbench-node-health) +# HG_SIF (default: $AI4S_SHARED_DIR/images/pytorch_rocm7.2.2_ubuntu24.04_py3.12_pytorch_release_2.10.0.sif) +# AI4S_SHARED_DIR (default: /shared/$USER) +# STREAM_THRESHOLD_COPY_GBPS (default: 200) - dual-NUMA COPY pass floor +# STREAM_THRESHOLD_TRIAD_GBPS (default: 200) - dual-NUMA TRIAD pass floor +# HIP_LAUNCH_MAX_US (default: 6.0) - per-GPU avg launch latency ceiling +# +# SBATCH directives: +#SBATCH --job-name=node-health +#SBATCH --partition=lux +#SBATCH --account=vultr_lux +#SBATCH --nodes=1 +#SBATCH --ntasks=1 +#SBATCH --cpus-per-task=128 +#SBATCH --time=00:05:00 +#SBATCH --output=microbench-node-health-%j.out +# +# NOTE: GPU tests are gated on whether SLURM allocated GPUs to the step. To +# get GPU coverage, sbatch with `--gres=gpu:amd_instinct_mi355_oam:8`. + +set -u + +H=$(hostname -s) +OUT_DIR="${OUT_DIR:-/shared/$USER/microbench-node-health}" +HOST_OUT="$OUT_DIR/$H" +mkdir -p "$HOST_OUT" + +AI4S_SHARED_DIR="${AI4S_SHARED_DIR:-/shared/$USER}" +HG_SIF="${HG_SIF:-${AI4S_SHARED_DIR}/images/pytorch_rocm7.2.2_ubuntu24.04_py3.12_pytorch_release_2.10.0.sif}" +STREAM_THRESHOLD_COPY_GBPS="${STREAM_THRESHOLD_COPY_GBPS:-200}" +STREAM_THRESHOLD_TRIAD_GBPS="${STREAM_THRESHOLD_TRIAD_GBPS:-200}" +HIP_LAUNCH_MAX_US="${HIP_LAUNCH_MAX_US:-6.0}" + +START_TS=$(date +%s) +echo "[$H] microbench start at $(date -u +%FT%TZ)" | tee "$HOST_OUT/_start.txt" + +# ---------- Test 1: host inventory ---------- +{ + lscpu > "$HOST_OUT/lscpu.txt" 2>&1 + numactl --hardware > "$HOST_OUT/numactl.txt" 2>&1 + uname -a > "$HOST_OUT/uname.txt" 2>&1 + cat /proc/cmdline > "$HOST_OUT/cmdline.txt" 2>&1 + cat /sys/devices/system/cpu/cpu0/cpufreq/scaling_governor > "$HOST_OUT/cpu_governor.txt" 2>&1 || echo "no cpufreq" > "$HOST_OUT/cpu_governor.txt" + grep -E '^Mem|HugePages|^DirectMap|^Swap' /proc/meminfo > "$HOST_OUT/meminfo.txt" 2>&1 + ip -br link > "$HOST_OUT/ip_link.txt" 2>&1 + mount | grep -E ' /home| /shared|tmpfs|/tmp' > "$HOST_OUT/mounts.txt" 2>&1 + df -h $HOME /shared /tmp 2>/dev/null > "$HOST_OUT/df.txt" 2>&1 +} + +# ---------- Test 2: container mount probe ---------- +# Catches a6-class fault: container can't write to $HOME or $AI4S_SHARED_DIR +MOUNT_PROBE_PASS="yes" +for d in "$HOME" "$AI4S_SHARED_DIR"; do + tmpf="$d/.health_probe_${SLURM_JOB_ID:-$$}" + if ! ( echo ok > "$tmpf" && rm -f "$tmpf" ) 2>/dev/null; then + MOUNT_PROBE_PASS="no($d)" + break + fi +done +echo "mount_probe: $MOUNT_PROBE_PASS" > "$HOST_OUT/mount_probe.txt" + +# ---------- Test 3: dual-NUMA STREAM ---------- +# Allocates ~6 GiB. Took 30-50s on a 2-socket Genoa MI355X node. +TMP=$(mktemp -d) +cat > "$TMP/stream.c" <<'C' +#include +#include +#include +#include +#define N (1L << 28) +#define NITER 8 +double *a, *b, *c; +static double tsec(void){ + struct timespec t; clock_gettime(CLOCK_MONOTONIC, &t); + return t.tv_sec + t.tv_nsec*1e-9; +} +int main(void){ + a = aligned_alloc(64, N*sizeof(double)); + b = aligned_alloc(64, N*sizeof(double)); + c = aligned_alloc(64, N*sizeof(double)); + #pragma omp parallel for + for(long i=0;i>20); + printf("# kernel bytes best_GBps avg_GBps\n"); + for(int k=0;k<4;k++){ + double bw_best = -1e30, bw_sum = 0; + for(int it=0;itbw_best) bw_best = bw; + bw_sum += bw; + } + const char *kn[] = {"COPY","SCALE","ADD","TRIAD"}; + long bytes_per_iter = (long)((k<2 ? 2.0 : 3.0) * N * 8); + printf("%s %ld %.2f %.2f\n", kn[k], bytes_per_iter, bw_best, bw_sum/NITER); + } + return 0; +} +C +gcc -O3 -fopenmp -march=native -o "$TMP/stream" "$TMP/stream.c" >/dev/null 2>&1 +{ + echo "--- single-NUMA (node 0, 64 threads) ---" + OMP_NUM_THREADS=64 OMP_PROC_BIND=true OMP_PLACES=cores numactl --cpunodebind=0 --membind=0 "$TMP/stream" + echo "" + echo "--- dual-NUMA (interleaved, 128 threads) ---" + OMP_NUM_THREADS=128 OMP_PROC_BIND=true OMP_PLACES=cores numactl --interleave=all "$TMP/stream" +} > "$HOST_OUT/stream.txt" 2>&1 +rm -rf "$TMP" + +# Pass/fail extraction (dual-NUMA section) +STREAM_COPY=$(awk '/--- dual-NUMA/{flag=1; next} flag && /^COPY/{print $3; exit}' "$HOST_OUT/stream.txt") +STREAM_TRIAD=$(awk '/--- dual-NUMA/{flag=1; next} flag && /^TRIAD/{print $3; exit}' "$HOST_OUT/stream.txt") +STREAM_PASS="yes" +[ -z "$STREAM_COPY" ] && STREAM_PASS="no(no_output)" +[ -n "$STREAM_COPY" ] && awk -v v="$STREAM_COPY" -v t="$STREAM_THRESHOLD_COPY_GBPS" 'BEGIN{exit !(v+0 < t+0)}' \ + && STREAM_PASS="no(COPY ${STREAM_COPY} < ${STREAM_THRESHOLD_COPY_GBPS})" +[ -n "$STREAM_TRIAD" ] && awk -v v="$STREAM_TRIAD" -v t="$STREAM_THRESHOLD_TRIAD_GBPS" 'BEGIN{exit !(v+0 < t+0)}' \ + && STREAM_PASS="no(TRIAD ${STREAM_TRIAD} < ${STREAM_THRESHOLD_TRIAD_GBPS})" + +# ---------- Test 4: GPU + HIP launch latency (only if GPUs allocated) ---------- +HIP_PASS="skipped(no_gpu_allocation)" +HIP_MEAN_US="" +if command -v rocm-smi >/dev/null 2>&1 || [ -e "$HG_SIF" ]; then + if [ -e "$HG_SIF" ]; then + apptainer exec --rocm --bind /opt/rocm-7.2.2:/opt/rocm-7.2.2 --bind /shared \ + --env LD_LIBRARY_PATH=/opt/rocm-7.2.2/lib:/usr/lib/x86_64-linux-gnu \ + "$HG_SIF" bash -c " +PATH=/opt/rocm-7.2.2/bin:\$PATH rocminfo > '$HOST_OUT/rocminfo.txt' 2>&1 +PATH=/opt/rocm-7.2.2/bin:\$PATH rocm-smi --showvbios --showfwinfo --showdriverversion > '$HOST_OUT/rocm_smi_vbios_fw.txt' 2>&1 +PATH=/opt/rocm-7.2.2/bin:\$PATH rocm-smi --showtopo > '$HOST_OUT/rocm_smi_topo.txt' 2>&1 +PATH=/opt/rocm-7.2.2/bin:\$PATH python3 - <<'PY' +import time, torch +N = torch.cuda.device_count() +NLAUNCH = 100000 +lats = [] +for d in range(N): + torch.cuda.set_device(d) + x = torch.zeros(1, device=f'cuda:{d}') + for _ in range(1000): x.zero_() + torch.cuda.synchronize(d) + t0 = time.perf_counter() + for _ in range(NLAUNCH): x.zero_() + torch.cuda.synchronize(d) + dt = time.perf_counter() - t0 + us = dt/NLAUNCH*1e6 + print(f'gpu{d}: {us:.2f} us/launch ({NLAUNCH/dt/1e6:.3f} M/s)') + lats.append(us) +print(f'mean_us: {sum(lats)/len(lats):.2f} max_us: {max(lats):.2f} min_us: {min(lats):.2f} n_gpu: {N}') +PY +" > "$HOST_OUT/hip_launch.txt" 2>&1 + HIP_MEAN_US=$(awk '/^mean_us:/{print $2}' "$HOST_OUT/hip_launch.txt") + if [ -n "$HIP_MEAN_US" ]; then + HIP_PASS="yes" + awk -v v="$HIP_MEAN_US" -v t="$HIP_LAUNCH_MAX_US" 'BEGIN{exit !(v+0 > t+0)}' \ + && HIP_PASS="no(mean ${HIP_MEAN_US} us > ${HIP_LAUNCH_MAX_US})" + else + HIP_PASS="no(no_torch_output)" + fi + fi +fi + +# /sys-based GPU PCIe link state (no container, no GPU allocation needed) +{ + for card in /sys/class/drm/card*/device; do + n=$(basename $(dirname "$card")) + cls=$(cat "$card/current_link_speed" 2>/dev/null || echo n/a) + clw=$(cat "$card/current_link_width" 2>/dev/null || echo n/a) + mls=$(cat "$card/max_link_speed" 2>/dev/null || echo n/a) + mlw=$(cat "$card/max_link_width" 2>/dev/null || echo n/a) + echo "$n cur=$cls/$clw max=$mls/$mlw" + done +} > "$HOST_OUT/pcie_link_state.txt" 2>&1 + +# ---------- Summary line ---------- +END_TS=$(date +%s) +SUMMARY="host=$H elapsed_s=$((END_TS - START_TS)) mount=$MOUNT_PROBE_PASS stream=$STREAM_PASS stream_copy_gbps=${STREAM_COPY:--} stream_triad_gbps=${STREAM_TRIAD:--} hip_launch=$HIP_PASS hip_mean_us=${HIP_MEAN_US:--}" +echo "$SUMMARY" > "$HOST_OUT/_summary.txt" +echo "$SUMMARY" + +# Exit nonzero if any test failed (suitable for SLURM Prolog use) +case "$MOUNT_PROBE_PASS$STREAM_PASS$HIP_PASS" in + yesyesyes|yesyesskipped*) exit 0 ;; + *) exit 1 ;; +esac diff --git a/material_science/models/HydraGNN/recipes/perf-optimizer-loop/lever_catalog.yaml b/material_science/models/HydraGNN/recipes/perf-optimizer-loop/lever_catalog.yaml index c6ee15d..65aace8 100644 --- a/material_science/models/HydraGNN/recipes/perf-optimizer-loop/lever_catalog.yaml +++ b/material_science/models/HydraGNN/recipes/perf-optimizer-loop/lever_catalog.yaml @@ -46,7 +46,7 @@ levers: notes: "Reference state. Establishes the FOM denominator for all subsequent iters." - id: batch_800 - description: "Double batch size: HG_BATCH_SIZE 400 -> 800. R2 #1 next-lever recommendation; VRAM headroom is 24/288 GB. NODE-CLASS-DEPENDENT: regresses on both a-nodes (+56.8%) and b-nodes (+87.2%) for HydraGNN GFM-MLIP MACE." + description: "Double batch size: HG_BATCH_SIZE 400 -> 800. R2 #1 next-lever recommendation; VRAM headroom is 24/288 GB. CONFIRMED REGRESSOR for HydraGNN GFM-MLIP MACE: regressed on both tested loops (+56.8%, +87.2%) — both runs were on lux-mi355x-b[1-2]." catalog_order: 1 kind: env_only env_vars: @@ -55,8 +55,8 @@ levers: expected_payoff: "none-for-MLIP-MACE" risk: "high" revert_method: "drop_env" - citation: "loop-43b33ec1 iter-2 (a-nodes, +56.76%) + loop-e3fac2af iter-1 (b-nodes, +87.2%)" - notes: "The lever_picker should NOT pick batch_800 for HydraGNN MLIP MACE on either node class. Cross-loop evidence shows the MACE equivariant message-passing stack saturates dispatch at batch=400 and doubling work-per-iter doubles wall-time without proportional GPU occupancy gain. The R2 recommendation was based on an earlier non-MLIP HydraGNN config — DOES NOT apply to gfm_mlip_all_mpnn.py. Keep entry as cautionary record. May still be worth re-trying ONLY if upstream HydraGNN fuses small e3nn kernels first." + citation: "loop-43b33ec1 iter-2 (b1+b2, +56.76%) + loop-e3fac2af iter-1 (b1+b2, +87.2%)" + notes: "The lever_picker should NOT pick batch_800 for HydraGNN MLIP MACE. The MACE equivariant message-passing stack saturates dispatch at batch=400 and doubling work-per-iter doubles wall-time without proportional GPU occupancy gain. The R2 recommendation was based on an earlier non-MLIP HydraGNN config — DOES NOT apply to gfm_mlip_all_mpnn.py. [REVISED 2026-05-26] previously described as 'node-class-dependent' but in fact both tested loops ran on b1+b2; the regression is workload-specific, not node-class. Keep entry as cautionary record. May still be worth re-trying ONLY if upstream HydraGNN fuses small e3nn kernels first." - id: torch_compile_e3nn description: "BLOCKED — torch.compile is fundamentally incompatible with HydraGNN GFM-MLIP under PyTorch 2.10 / ROCm 7.2.2. Documented here so the lever_picker NEVER re-tries it." @@ -166,7 +166,7 @@ levers: risk: "low" revert_method: "drop_env" citation: "loop-c3e4df1c iter-2 / jobid 7187: -1.3% epoch_time (76.0->75.0 s) AND -8.4% energy/sample (1.781->1.631 J/sample). Mean power dropped 4885W->4552W." - notes: "ACCEPTED. The only lever across 3 loops that produced a real improvement. Now part of the baseline contract for HydraGNN GFM-MLIP on b-nodes. lever_picker should treat this as 'already applied' and skip it in candidate set (foms.csv has it accepted with delta_pct=-1.316)." + notes: "ACCEPTED. The only lever across 3 loops that produced a real improvement. Now part of the baseline contract for HydraGNN GFM-MLIP. [REVISED 2026-05-26] previously framed as 'on b-nodes' but the validating runs were specifically on lux-mi355x-b[1-2]; a-vs-b microbench (2026-05-26) shows no class difference, so this likely applies to a-nodes too — re-verify on a-pair before declaring class-portability. lever_picker should treat this as 'already applied' and skip it in candidate set (foms.csv has it accepted with delta_pct=-1.316)." - id: tunable_op_live description: "BLOCKED on this stack — TunableOp LIVE tuning mode (PYTORCH_TUNABLEOP_TUNING=1) causes GPU memory access faults on MI355X / ROCm 7.2.2 / PyTorch 2.10 with HydraGNN. Use tunable_op_warmup_then_use instead." @@ -229,17 +229,17 @@ levers: notes: "If the warmup phase did NOT complete (no CSVs), pick a different lever. The orchestrator should add `setup: tunableop_warmup_sbatch` as a precondition before this lever is offered." - id: num_workers_12 - description: "DataLoader: HYDRAGNN_NUM_WORKERS 8 -> 12 (persistent_workers already on). Node-class-dependent — helped on a-nodes (-12.4% in loop 43b33ec1), regressed on b-nodes (+10.7% in loop c3e4df1c)." + description: "DataLoader: HYDRAGNN_NUM_WORKERS 8 -> 12 (persistent_workers already on). HIGH RUN-TO-RUN VARIANCE: helped in loop-43b33ec1 (-12.4%) and regressed in loop-c3e4df1c (+10.7%) — both ran on lux-mi355x-b[1-2], so this is intra-pair variance, not class behavior. Treat as unreliable until replicated within ±5% on the same nodes." catalog_order: 5 kind: env_only env_vars: HYDRAGNN_NUM_WORKERS: "12" diagnostic_only: false - expected_payoff: "node-class-dependent" + expected_payoff: "unreliable (variance > effect)" risk: "low" revert_method: "drop_env" - citation: "local (loop 43b33ec1 + loop c3e4df1c)" - notes: "lever_picker should only consider this lever on a-nodes. On b-nodes, num_workers=8 is the sweet spot (verified in iter-2 of loop c3e4df1c with epoch_time=75 s, the current best). If the kernel_correlation.csv from fom_extractor shows dataloader prefetch gaps > 10 ms on b-nodes, try num_workers=10 (not 12)." + citation: "local (loop 43b33ec1 + loop c3e4df1c, both on b1+b2)" + notes: "[REVISED 2026-05-26] previously believed to be node-class-dependent (a-nodes vs b-nodes), but both observations were on the same b1+b2 pair, so the result is run-to-run noise on a single pair. Recommended next step: re-run num_workers={8,10,12} three times each on the SAME node pair with the SAME warm/cold cache state. If the kernel_correlation.csv from fom_extractor shows dataloader prefetch gaps > 10 ms, that's a stronger signal than the epoch-time delta. Until then the lever_picker should rank this lever LOWER than levers with smaller variance." - id: nccl_minchannels description: "NCCL_MIN_NCHANNELS=112: forces more channels for collective dispatch, helps when message size is moderate. CONFIRMED REGRESSOR at N=2 (+10.3%)." @@ -266,7 +266,7 @@ levers: risk: "low" revert_method: "drop_env" citation: "loop-e3fac2af iter-4 (job 7203): +1.3% epoch_time (78->79 s) confirmed in fp32 (precision-diagnostic logs show model_param_dtype=torch.float32, first_batch_float_dtype=torch.float32)" - notes: "RAN GENUINELY IN FP32 (verified by precision-diagnostic). fp32 MFMA peak on gfx950 is 2x fp64 peak, so a compute-bound workload would show ~50% speedup. Observed ~0% speedup proves HydraGNN GFM-MLIP MACE on b-nodes is dispatch-bound. lever_picker should NOT re-pick this lever unless the run is explicitly switching to a-nodes (where the compute/dispatch balance may differ). bf16 remains INTENTIONALLY EXCLUDED: breaks equivariance for MACE-class features. The fp32-vs-fp64 comparison is now the canonical low-cost dispatch-vs-compute discriminator for any future HydraGNN-class workload — see ai4science-perf-analysis/SKILL.md lesson #13." + notes: "RAN GENUINELY IN FP32 (verified by precision-diagnostic). fp32 MFMA peak on gfx950 is 2x fp64 peak, so a compute-bound workload would show ~50% speedup. Observed ~0% speedup proves HydraGNN GFM-MLIP MACE is dispatch-bound. [REVISED 2026-05-26] previously hedged with 'on b-nodes; may differ on a-nodes'; a-vs-b microbench (2026-05-26) confirmed identical HIP launch latency (3.7 us/launch) on every tested node of both classes, so the dispatch ceiling is class-portable — fp32 will not help on a-nodes either. lever_picker should NOT re-pick this lever. bf16 remains INTENTIONALLY EXCLUDED: breaks equivariance for MACE-class features. The fp32-vs-fp64 comparison is now the canonical low-cost dispatch-vs-compute discriminator for any future HydraGNN-class workload — see ai4science-perf-analysis/SKILL.md lesson #13." - id: kernel_trace_diag_only description: "DIAGNOSTIC: enable OMNISTAT_KERNEL_TRACE=1 for ONE iteration to capture per-kernel-name dispatch counts/durations. Fires only when the 1-s TraceLens<->Omnistat correlation cannot attribute the dominant kernel." From 7682fa313805125f3538cea52059cf9e7d1c299a Mon Sep 17 00:00:00 2001 From: Ashwin Aji Date: Tue, 26 May 2026 08:33:38 +0000 Subject: [PATCH 12/31] perf-optimizer-loop: add HANDOFF.md for the next agent Single-file orientation so a fresh agent can pick up mid-flight without re-discovering anything. Captures: current best config (75.0 s/epoch on 2 nodes), prioritized next work (8/16-node scaling sweep, two upstream HydraGNN PRs, second-user repro, Frontier-scale doc), landmines to avoid (each cross-referenced to a SKILL lesson), and pointers to every artefact (runtime dirs, microbench report, sysadmin message, ticket comment id). Discovery paths: - recipe README links to it at the top - SKILL "Repository entry points" calls it out for the iterative-loop recipe - keyword-rich title (HANDOFF, next agent, perf-optimizer, MI355X, dispatch-bound) so grep/glob surfaces it - model.yaml's existing recipe_path: recipes/perf-optimizer-loop/ already points the model-discovery flow to its folder Co-authored-by: Cursor --- .../skills/ai4science-perf-analysis/SKILL.md | 4 +- .../recipes/perf-optimizer-loop/HANDOFF.md | 115 ++++++++++++++++++ .../recipes/perf-optimizer-loop/README.md | 2 + 3 files changed, 120 insertions(+), 1 deletion(-) create mode 100644 material_science/models/HydraGNN/recipes/perf-optimizer-loop/HANDOFF.md diff --git a/.cursor/skills/ai4science-perf-analysis/SKILL.md b/.cursor/skills/ai4science-perf-analysis/SKILL.md index c38a59f..e5d9038 100644 --- a/.cursor/skills/ai4science-perf-analysis/SKILL.md +++ b/.cursor/skills/ai4science-perf-analysis/SKILL.md @@ -14,8 +14,10 @@ Default target: HydraGNN on Lux MI355X. The same pattern can extend to ORBIT-2 / ## Repository entry points - Recipe: [material_science/models/HydraGNN/recipes/perf-analysis/](../../../material_science/models/HydraGNN/recipes/perf-analysis/) +- Iterative-loop recipe: [material_science/models/HydraGNN/recipes/perf-optimizer-loop/](../../../material_science/models/HydraGNN/recipes/perf-optimizer-loop/) — **see [`HANDOFF.md`](../../../material_science/models/HydraGNN/recipes/perf-optimizer-loop/HANDOFF.md) for current state, what works, what's blocked, and next-work priorities** - Sbatch wrapper: [material_science/models/HydraGNN/examples/sbatch_train_perf_amd.sh](../../../material_science/models/HydraGNN/examples/sbatch_train_perf_amd.sh) -- Agent prompt files: `material_science/models/HydraGNN/recipes/perf-analysis/agents/*.md` +- Reusable node-health probe: [material_science/models/HydraGNN/examples/microbench_node_health.sh](../../../material_science/models/HydraGNN/examples/microbench_node_health.sh) +- Agent prompt files: `material_science/models/HydraGNN/recipes/perf-analysis/agents/*.md` and `material_science/models/HydraGNN/recipes/perf-optimizer-loop/agents/*.md` ## Orchestration loop (what the main agent does) diff --git a/material_science/models/HydraGNN/recipes/perf-optimizer-loop/HANDOFF.md b/material_science/models/HydraGNN/recipes/perf-optimizer-loop/HANDOFF.md new file mode 100644 index 0000000..cf10fa9 --- /dev/null +++ b/material_science/models/HydraGNN/recipes/perf-optimizer-loop/HANDOFF.md @@ -0,0 +1,115 @@ +# HydraGNN perf-optimizer-loop — handoff to the next agent + +> **Read this first.** Single-file orientation for any agent picking up this work mid-flight. Keyword-rich title so a glob/grep for `HANDOFF`, `next agent`, `perf-optimizer`, `MI355X`, `dispatch-bound`, `HydraGNN handoff` will surface it. + +**Last touched:** 2026-05-26 by aaji. +**Branch:** `aaji/perf-optimizer-loop-hydragnn` (6 local commits ahead of origin; NOT pushed yet). +**Ticket:** [DCSWEAP-4502](https://amd.atlassian.net/browse/DCSWEAP-4502) "Run HydraGNN on MI355X". Latest comment id `22139024` (posted 2026-05-26 by aaji) is the current state of record. + +## TL;DR — where the work is + +The 2-node MI355X performance characterization on Vultr Lux is **done**. Current best for HydraGNN GFM-MLIP MACE on 2 nodes (16× MI355X) is **75.0 s/epoch** with this exact env (already the sbatch default): + +```bash +HG_BATCH_SIZE=400 +HG_PRECISION=fp64 +HYDRAGNN_NUM_WORKERS=8 +HYDRAGNN_PERSISTENT_WORKERS=1 +HG_NUM_EPOCH=3 +HYDRAGNN_MAX_NUM_BATCH=50 +PROFILE_TARGET_EPOCH=2 +TORCH_NCCL_HIGH_PRIORITY=1 # the only accepted lever across 3 loops +GPU_MAX_HW_QUEUES=2 # ditto +``` + +The workload is **dispatch-bound** (HIP launch latency ceiling 3.7 µs × 270 K launches/s, identical on every MI355X we probed). fp32-vs-fp64 differs by < 2 % despite fp32 MFMA having 2× the peak, which is the canonical proof. The path to a step-change requires **two upstream HydraGNN PRs** (see "Next work" below). + +## Where every artefact lives + +| What | Where | Tracked? | +|---|---|---| +| Recipe + agent prompts | `material_science/models/HydraGNN/recipes/perf-optimizer-loop/` | git | +| Lever catalog (machine-readable; blocked levers have evidence + path-to-unblock) | `recipes/perf-optimizer-loop/lever_catalog.yaml` | git | +| Sbatch wrapper (defaults = current best config) | `material_science/models/HydraGNN/examples/sbatch_train_perf_amd.sh` | git | +| Optimizer-loop driver (Claude Code CLI, tmux) | `material_science/models/HydraGNN/examples/run_optimizer_loop.sh` | git | +| Reusable node-health microbench (mount + STREAM + HIP launch + GPU/firmware inventory; ~30 s) | `material_science/models/HydraGNN/examples/microbench_node_health.sh` | git | +| All lessons learned (perf-analysis-specific) | `.cursor/skills/ai4science-perf-analysis/SKILL.md` (lessons 1-19) | git | +| Per-loop runtime artefacts (3 loops + 2 investigations) | `/shared/aaji/models/HydraGNN/perf-runs/loop-*` and `/shared/aaji/models/HydraGNN/perf-runs/{compile-investigation-*,pathc-pathb-*}` | NOT tracked | +| Microbench raw outputs + REPORT.md + SYSADMIN_MESSAGE.md | `/shared/aaji/microbench-a-vs-b/` | NOT tracked | +| Container (PyTorch 2.10.0 + ROCm 7.2.2 + HydraGNN overlay) | `/shared/aaji/images/pytorch_rocm7.2.2_ubuntu24.04_py3.12_pytorch_release_2.10.0.sif` | NOT tracked | +| Top-level ticket | [DCSWEAP-4502](https://amd.atlassian.net/browse/DCSWEAP-4502) | external | + +## Reproduce the current best in one command + +```bash +ssh aaji@rad-vultr-login # or whoever the next user is +cd /home/aaji/git/ai4science-studio +git checkout aaji/perf-optimizer-loop-hydragnn +export AI4S_SHARED_DIR=/shared/aaji +sbatch material_science/models/HydraGNN/examples/sbatch_train_perf_amd.sh +# expect: 75.0 s/epoch, train_validate_test ≈ 230 s, J/sample ≈ 1.631 +``` + +## Next work (priority order, from the ticket) + +| # | Task | Estimate | Notes | +|---|---|---|---| +| 1 | **8-node + 16-node strong-scaling sweep** with current best config | 1-2 hours wall, mostly queue time | Cluster cap is 16 nodes. Modify `sbatch_train_perf_amd.sh` `--nodes=N` only. Capture epoch time + J/sample + scaling efficiency vs the 2-node baseline. Write into a new `recipes/perf-optimizer-loop/scaling-results/` markdown + foms.csv. | +| 2 | **Open the two HydraGNN upstream PRs** | a few days each | (a) Move `energy_force_loss` into `EnhancedModelWrapper.forward` (~30-50 LOC across `hydragnn/models/create.py:622-758` + `hydragnn/train/train_validate_test.py:725-735`) — unblocks `torch.compile`. (b) Unpack `**conv_args` at `hydragnn/models/MACEStack.py:397` — unblocks `torch.jit.script`. Reference issues: see `lever_catalog.yaml` entries `torch_compile_e3nn` and `torch_jit_script` for the full evidence + diff sketches. Cite `/shared/aaji/models/HydraGNN/perf-runs/compile-investigation-aa480554-*/INVESTIGATION_REPORT.md` and `/shared/aaji/models/HydraGNN/perf-runs/pathc-pathb-a3b8f3f0-*/pathc-result.md` in the PRs. | +| 3 | **Second-user reproducibility** | half day | Have a different user run the one-command repro above on a different login session. If they hit 75 ± 1 s/epoch, ticket acceptance #7 is satisfied. | +| 4 | **Frontier-scale path documentation (DDStore phase 2)** | document only | Cannot validate on Vultr Lux (cap=16); write up the 504-node path in `recipes/perf-optimizer-loop/scaling-results/frontier-path.md`. Reference DDStore upstream docs. | + +## Landmines to NOT step on (these all cost real hours to discover) + +Each item is cross-referenced to a lesson in `.cursor/skills/ai4science-perf-analysis/SKILL.md` for the full backstory. + +| Don't | Why | Lesson | +|---|---|---| +| Don't re-try `torch.compile` on HydraGNN MLIP without first landing PR #2a above | AOT-Autograd doesn't support double backward; `energy_force_loss` calls `autograd.grad(create_graph=True)`. Tested 6 hooks across 3 hypotheses, all fail identically. | #6, #7 | +| Don't re-try `torch.jit.script` without first landing PR #2b above | TorchScript JIT can't expand `**conv_args` at `MACEStack.py:397`. | #14 | +| Don't enable `PYTORCH_TUNABLEOP_TUNING=1` in any form (live, warmup, warmup-then-use) | All variants fault all 16 GPUs in the hipBLASLt tuning routine. Re-test only after a documented ROCm/hipBLASLt fix. | #8, #15 | +| Don't draw a-nodes-vs-b-nodes conclusions from FOM deltas | They're physically identical (verified 2026-05-26); the historical "17 % gap" was three baselines on the same b1+b2 pair. Always `jq .nodes_list iter-0-baseline.json` first. | #10, #18 | +| Don't pick `num_workers_12` based on the loop-43b33ec1 win | Same lever regressed in loop-c3e4df1c on the same node pair; it's noise > effect at this measurement granularity. | #9, #18 | +| Don't pick `batch_800` for HydraGNN MLIP MACE | Dispatch saturates at batch=400; +87 % at batch=800. The R2 recommendation was from a pre-MLIP HydraGNN config. | lever_catalog `batch_800` notes | +| Don't pick `precision_fp32` expecting a speedup | The workload is dispatch-bound; fp32 gives < 2 % difference. (It IS still useful as a one-shot dispatch-vs-compute discriminator on new workloads — see lesson #13.) | #13 | +| Don't pick `nccl_minchannels=112` at N=2 | +10.3 % at N=2; only meaningful at ≥8 nodes. | lever_catalog `nccl_minchannels` notes | +| Don't omit `AI4S_SHARED_DIR` from `sbatch --export=ALL,...` | `--export=ALL` doesn't include unexported shell vars; sbatch dies on line 61 with `AI4S_SHARED_DIR must be set`. | #12 | +| Don't use `--gpus-per-node=N` for srun probes | SLURM filter rejects with "Invalid GRES specification"; use `--gres=gpu:amd_instinct_mi355_oam:N` instead. | (sbatch wrapper) | +| Don't claim a class-wide node bandwidth gap from a 16-GB memory cgroup | `cpu_mem_bw.sh`-style measurements get an allocation-artefact ceiling at ~220 GB/s; the real hardware delivers ~313 GB/s when allocated full-node. The committed `microbench_node_health.sh` works under either allocation but the threshold is calibrated to 200 GB/s for the cgroup case. | (microbench REPORT) | +| Don't blacklist a node *class* on a single failure | Probe per-node with the committed mount-health probe (exit 42) or `microbench_node_health.sh` and exclude only the named bad node(s). | #1, #18 | + +## Per-node health to be aware of (sysadmin will resolve) + +These are observations to relay to the cluster admin (draft is at `/shared/aaji/microbench-a-vs-b/SYSADMIN_MESSAGE.md`), NOT ticket items. If you land on one of these nodes, exclude by name and continue: + +| Node | Symptom | Workaround | +|---|---|---| +| `lux-mi355x-a5` | Intermittent dual-NUMA STREAM degradation (179 vs 220 GB/s COPY, 8 % jitter) | `--exclude=lux-mi355x-a5` | +| `lux-mi355x-a6` | `/home` + `/shared` bind-mount EIO inside Apptainer (historical 2026-05-22) | Mount probe (exit 42) auto-detects; orchestrator retries with `--exclude` | +| `lux-mi355x-a10` | PMIx ring collective times out at MPI init (historical 2026-05-23) | `--exclude=lux-mi355x-a10` | + +## Bootstrap checklist for a brand-new agent + +1. **Read [`README.md`](README.md)** in this directory for the recipe topology + iteration shape. +2. **Read [`lever_catalog.yaml`](lever_catalog.yaml)** in this directory — entries marked `status: blocked` or `status: accepted-baked-in` matter most. +3. **Read** `.cursor/skills/ai4science-perf-analysis/SKILL.md` lessons #1-19, especially the **REVISED** markers on lessons 9, 10, 16. +4. **`git log --oneline -8`** on this branch to see the change history with one-liners. +5. **`cat /shared/aaji/models/HydraGNN/perf-runs/loop-*/foms.csv`** to see the actual FOM trajectories across iterations. +6. **Check ticket** [DCSWEAP-4502](https://amd.atlassian.net/browse/DCSWEAP-4502), especially the latest comment (id `22139024`, posted 2026-05-26). +7. **Pick one of "Next work" items above** and start. + +## Open questions / decisions for the human, NOT for the agent + +- Should the branch be pushed to `origin` and a draft PR opened against `main`? (User hasn't asked; agent should not push without explicit ask.) +- Should the two HydraGNN upstream PRs be opened by the agent or by a human? (They're against `github.com/ORNL/HydraGNN`, a third-party repo; agent can prepare patches but the human should be the author of record.) +- Is there appetite for adding `microbench_node_health.sh` to other models' recipe folders (StormCast, ORBIT-2, GP-MoLFormer)? The script is generic enough to apply to any MI355X workload on Lux. + +## Authors / contacts + +- **Performance work + this branch:** Ashwin Aji (`ashwin.aji@amd.com`, github: TBD) +- **DCSWEAP-4502 reporter:** Nicholas Malaya (`nicholas.malaya@amd.com`) +- **Upstream HydraGNN maintainers:** ORNL — see [`github.com/ORNL/HydraGNN`](https://github.com/ORNL/HydraGNN) issues for the right tag + +## If anything in this file goes stale + +Update this file in the same commit as the change that made it stale. Single-source-of-truth for handoff state — if it disagrees with reality, fix the file, don't write a "HANDOFF-v2.md". diff --git a/material_science/models/HydraGNN/recipes/perf-optimizer-loop/README.md b/material_science/models/HydraGNN/recipes/perf-optimizer-loop/README.md index 8dd8aba..91d7afd 100644 --- a/material_science/models/HydraGNN/recipes/perf-optimizer-loop/README.md +++ b/material_science/models/HydraGNN/recipes/perf-optimizer-loop/README.md @@ -4,6 +4,8 @@ Multi-iteration optimizer that runs HydraGNN on 2 nodes × 8 × MI355X (gfx950), > Audience: AMD performance engineering. Output is a diagnosis + tunings, not a HydraGNN scientific claim. +> **Picking up this work? Read [`HANDOFF.md`](HANDOFF.md) first.** It has the current state (best epoch time, blocked levers, next-work priority list, landmines to avoid, and pointers to all the runtime artefacts outside the repo). + ## What this recipe adds on top of `perf-analysis/` | Concern | `perf-analysis/` | `perf-optimizer-loop/` | From 1c6ff07820378ae0c6114b902012f5210d7084c5 Mon Sep 17 00:00:00 2001 From: Ashwin Aji Date: Tue, 26 May 2026 22:00:57 +0000 Subject: [PATCH 13/31] perf-optimizer-loop: scrub internal names for public repo Remove username, site-specific paths (/shared/aaji, /home/aaji, rad-vultr-login), cluster-specific node names (lux-mi355x-*), partition/account names (lux, vultr_lux), and the now-refuted a-nodes-vs-b-nodes 17% performance gap narrative from all committed files. - Replace hardcoded paths with $AI4S_SHARED_DIR throughout scripts and agent prompt files; add AI4S_SHARED_DIR required-or-fail guard in run_optimizer_loop.sh - Replace partition/account SBATCH directives with generic defaults - Rewrite lessons #9, #10, #16, #18 in SKILL.md to drop cluster- specific node-class framing; preserve the methodological lesson (always check nodes_list before drawing class conclusions) - Delete HANDOFF.md (ephemeral session artifact, not repo material); add HANDOFF.md to .gitignore - Remove Authors/contacts block from HANDOFF.md before deletion Co-Authored-By: Claude Sonnet 4 --- .../skills/ai4science-perf-analysis/SKILL.md | 52 ++++---- .gitignore | 3 + .../examples/microbench_node_health.sh | 14 +-- .../HydraGNN/examples/run_optimizer_loop.sh | 16 +-- .../examples/sbatch_train_perf_amd.sh | 23 ++-- .../recipes/perf-optimizer-loop/HANDOFF.md | 115 ------------------ .../recipes/perf-optimizer-loop/README.md | 16 +-- .../agents/fom_extractor.md | 8 +- .../agents/orchestrator.md | 20 +-- .../agents/story_writer.md | 4 +- .../perf-optimizer-loop/lever_catalog.yaml | 34 +++--- 11 files changed, 94 insertions(+), 211 deletions(-) delete mode 100644 material_science/models/HydraGNN/recipes/perf-optimizer-loop/HANDOFF.md diff --git a/.cursor/skills/ai4science-perf-analysis/SKILL.md b/.cursor/skills/ai4science-perf-analysis/SKILL.md index e5d9038..67f7dd8 100644 --- a/.cursor/skills/ai4science-perf-analysis/SKILL.md +++ b/.cursor/skills/ai4science-perf-analysis/SKILL.md @@ -9,7 +9,7 @@ description: Runs the AMD AI agents bottleneck-analysis workflow on a HydraGNN ( The user wants an automated bottleneck analysis of a multi-node training/inference run using AMD's open-source observability tooling (TraceLens, Omnistat). This is **distinct** from the `ai4science-run-models` skill — that one launches a model; this one diagnoses a model after it runs. -Default target: HydraGNN on Lux MI355X. The same pattern can extend to ORBIT-2 / StormCast / GP-MoLFormer in iteration 2. +Default target: HydraGNN on AMD MI355X. The same pattern can extend to ORBIT-2 / StormCast / GP-MoLFormer in iteration 2. ## Repository entry points @@ -37,21 +37,18 @@ Default target: HydraGNN on Lux MI355X. The same pattern can extend to ORBIT-2 / - Read **only** the files listed in its `## Inputs` section. - Produce **only** the file(s) listed in its `## Outputs` section. - Write a single line `STATUS=ok|partial|fail; reason=` to stdout as the last line. -- Never escalate to multi-node `srun`; verifiers may use **at most one** 1-node `srun -p lux -A vultr_lux -N1 --time=00:05:00` interactive probe. -- Never edit files outside `/shared/aaji/models/HydraGNN/perf-runs//`. +- Never escalate to multi-node `srun`; verifiers may use **at most one** 1-node `srun -N1 --time=00:05:00` interactive probe (partition/account from `.cluster-config.yaml`). +- Never edit files outside `$AI4S_SHARED_DIR/models/HydraGNN/perf-runs//`. ## Cluster constraints (Lux) -- Partition: `lux`. Account: `vultr_lux`. Read from [.cluster-config.yaml](../../../.cluster-config.yaml). -- Central Prometheus at `atlvrmonad01:9090` is **unreachable from compute nodes** (TCP RST, IP allowlist). Do **not** point omnistat-query at it. Always use the user-mode VictoriaMetrics that the launcher started. -- System-mode Omnistat is running on every compute node at port 8001. Don't touch it; we run our own user-mode collector alongside. -- The login node (`rad-vultr-login`) has no GPU and no MPI. All analysis runs after the job — no live profiling on the login node. -- **Node-specific mount faults are real and silent.** Any single compute node can land an allocation with a broken per-user autofs/NFS mount of `/home/$USER` or `/shared/$USER`, while sibling nodes in the same allocation are fine. Symptom in slurmd logs on the broken node: `Home Directory for Not Found ... Please contact system admin` and `lstat /shared/...: no such file or directory`; the surviving nodes wedge in a `pmix_coll_ring` collective fence until SLURM kills the job at the wall. - - **Confirmed bad nodes** (as of 2026-05-26, until cluster admin repairs them): - - `lux-mi355x-a5` — intermittent dual-NUMA STREAM bandwidth degradation (179 GB/s COPY vs cluster median 220 GB/s) with 8% run-to-run jitter; partial recovery on retest. Likely DIMM/IF or thermal issue on one socket. Found 2026-05-26 microbench (jobs 7589 + 7601). Mount + GPU healthy; only the CPU memory subsystem is degraded. - - `lux-mi355x-a6` — `/home/aaji` missing, `/shared/aaji/images/*.sif` unreadable; sibling `lux-mi355x-a1` healthy in the same alloc. Loop-43b33ec1 iter-1 / job 7088. - - `lux-mi355x-a10` — pmix ring collective times out at MPI init; `opal_mpi_init_ext3x_client_unreachable`. NOT a mount fault — broken PMIx/OpenMPI install on this node. Loop-c3e4df1c iter-1 / job 7161. - - **Diagnostic rule:** never blacklist a node *class* (e.g. "avoid all a-nodes") based on a single failure. Always probe the *specific* allocated nodes and exclude only those that fail by name. Probe is a one-task-per-node srun that checks `/home/$USER`, `/shared/$USER`, and the SIF path; see the `Per-node mount-health probe` section of `sbatch_train_perf_amd.sh` (exit 42 on fail). The perf-optimizer-loop orchestrator handles exit 42 by re-submitting with `--exclude=` and adding the node to `loop-/known_bad_nodes.txt`; the lever is NOT penalized. +- Partition / account: read from [.cluster-config.yaml](../../../.cluster-config.yaml). +- Central Prometheus may be **unreachable from compute nodes** (check your cluster's network policy). Always use the user-mode VictoriaMetrics that the launcher started. +- System-mode Omnistat may be running on every compute node at port 8001. Don't touch it; we run our own user-mode collector alongside. +- The login node typically has no GPU and no MPI. All analysis runs after the job — no live profiling on the login node. +- **Node-specific mount faults are real and silent.** Any single compute node can land an allocation with a broken per-user autofs/NFS mount of `/home/$USER` or `$AI4S_SHARED_DIR`, while sibling nodes in the same allocation are fine. Symptom in slurmd logs on the broken node: `Home Directory for Not Found ... Please contact system admin` and `lstat ...: no such file or directory`; the surviving nodes wedge in a `pmix_coll_ring` collective fence until SLURM kills the job at the wall. + - **Confirmed bad node patterns we hit (MI355X, May 2026):** intermittent dual-NUMA STREAM degradation on one node, broken container bind-mounts (`$HOME`/`$AI4S_SHARED_DIR`) on another, PMIx ring timeout on a third. Node names are cluster-specific; see the untracked microbench outputs for hostnames. + - **Diagnostic rule:** never blacklist a node *class* based on a single failure. Always probe the *specific* allocated nodes and exclude only those that fail by name. Probe is a one-task-per-node srun that checks `/home/$USER`, `$AI4S_SHARED_DIR`, and the SIF path; see the `Per-node mount-health probe` section of `sbatch_train_perf_amd.sh` (exit 42 on fail). The perf-optimizer-loop orchestrator handles exit 42 by re-submitting with `--exclude=` and adding the node to `loop-/known_bad_nodes.txt`; the lever is NOT penalized. ## When something goes wrong @@ -64,7 +61,7 @@ Default target: HydraGNN on Lux MI355X. The same pattern can extend to ORBIT-2 / The first 2-node end-to-end run on Lux exposed five real bugs that are now all worked around in [`sbatch_train_perf_amd.sh`](../../../material_science/models/HydraGNN/examples/sbatch_train_perf_amd.sh) and the agent prompts: 1. **`%(SLURM_JOB_ID)s` in omnistat config is broken** — `configparser` doesn't interpolate `os.environ`. Use `@JOB_DIR@` placeholder and `sed` at submit time. -2. **`/tmp/omni_rmsjobinfo` permission collision with system-mode Omnistat** — override `job_detection_file` to a per-job path under `/shared/aaji/...`. +2. **`/tmp/omni_rmsjobinfo` permission collision with system-mode Omnistat** — override `job_detection_file` to a per-job path under `$AI4S_SHARED_DIR/...`. 3. **VictoriaMetrics on login node needs `-fs.disableMmap`** to load even a 1.6 MB DB (cgroup mmap restriction). 4. **HydraGNN `Profile` block must be under `NeuralNetwork`**, not at the top level — `train_validate_test()` is called with `config["NeuralNetwork"]` so that's the scope its Profiler reads. 5. **PromQL with `jobid="..."` requires joining via `rmsjob_info`** — `rocm_*` metrics don't carry a `jobid` label directly. Verifier subagents must use: @@ -74,7 +71,7 @@ The first 2-node end-to-end run on Lux exposed five real bugs that are now all w ## Lessons captured (perf-optimizer-loop pilot, loop-43b33ec1, 2026-05-22) -1. **Diagnose broken nodes individually, not by class** (see Cluster constraints above for full detail). The optimizer-loop's iter-1 wrongly concluded "a-nodes lack /home" from a single failed alloc — actually only `lux-mi355x-a6` was broken. The fix is a per-job mount-health probe (now in `sbatch_train_perf_amd.sh`, exit code 42) plus orchestrator step 2e-bis: on exit 42, parse `node_health_probe.txt`, append the bad node(s) to `loop-/known_bad_nodes.txt`, and re-submit with `--exclude=`. Lever is untouched. +1. **Diagnose broken nodes individually, not by class** (see Cluster constraints above for full detail). The optimizer-loop's iter-1 wrongly concluded "a-class nodes lack /home" from a single failed alloc — only one specific node was broken. The fix is a per-job mount-health probe (now in `sbatch_train_perf_amd.sh`, exit code 42) plus orchestrator step 2e-bis: on exit 42, parse `node_health_probe.txt`, append the bad node(s) to `loop-/known_bad_nodes.txt`, and re-submit with `--exclude=`. Lever is untouched. 2. **HydraGNN `torch.compile` hooks must wrap the model object, not a module path.** `import hydragnn.train.train` (the path the iter-4 hook tried) does not exist in the installed package. Wrap the constructed model directly: `model = torch.compile(model, backend="inductor", mode="reduce-overhead", fullgraph=False)` at the rank-script level, by patching the entrypoint script (`gfm_mlip_all_mpnn.py`) immediately after model construction. **(Now superseded by Lesson #6 below — torch.compile is BLOCKED for HydraGNN MLIP regardless of where it wraps.)** 3. **`parse_convergence.py` over-counts epochs by N_ranks** because each rank emits a `tqdm=100%` line and the parser counts every one. Always read epoch wall-time from rank-0 tqdm directly (`s/it × max_num_batch`), not from the parser's epoch count. Fixed in fom_extractor; long-term fix is to filter on `rank=0` lines in the parser. 4. **MFMA TFLOPS methodology must be declared explicitly.** Burst-rate via PromQL `rate([10s])×4` (iter-0/3 values ≈ 0.006-0.012 TFLOPS) and time-averaged `total_accumulated_ops / duration` (iter-4 value ≈ 0.75 TFLOPS) differ by ~60× for this workload and **must not be compared**. fom_extractor should standardize on time-averaged across the full job window, with the burst-rate as a separate `*_peak_burst` field. @@ -85,14 +82,14 @@ The first 2-node end-to-end run on Lux exposed five real bugs that are now all w 6. **`torch.compile` is BLOCKED for HydraGNN GFM-MLIP MACE** under PyTorch 2.10 / ROCm 7.2.2. Six hooks across three hypotheses (inner-model compile, outer + `dynamo.allow_in_graph(autograd.grad)`, single-submodule compile) all fail with the same verbatim PyTorch error: `RuntimeError: torch.compile with aot_autograd does not currently support double backward` (`torch/_functorch/_aot_autograd/runtime_wrappers.py:2356`). Root cause: HydraGNN's `energy_force_loss` (`hydragnn/models/create.py:718`) calls `torch.autograd.grad(graph_energy_pred, data.pos, create_graph=True)` to compute forces as derivatives of energy w.r.t. positions; with `create_graph=True`, training requires DOUBLE BACKWARD (loss.backward differentiates through forces). Every torch.compile mode (`default`, `reduce-overhead`, `max-autotune`) goes through AOT Autograd which pre-compiles the backward as a joint program and cannot itself be differentiated. **To unblock requires upstream HydraGNN source change**: move `energy_force_loss` INTO `EnhancedModelWrapper.forward()` so `autograd.grad` is inside the compiled forward (the canonical MACE/Allegro/NequIP pattern). ~30-50 LOC across `create.py:622-758` + `train_validate_test.py:725-735`, config-gated. The lever `torch_compile_e3nn` is marked `status: blocked` in `lever_catalog.yaml` with full evidence; the lever_picker drops blocked levers from its candidate set. 7. **`dynamo.allow_in_graph(autograd.grad)` is weaker than MACE upstream docs suggest.** It only prevents dynamo from TRACING THROUGH the call; it does NOT prevent dynamo's FakeTensor symbolic-execution pass from running `autograd.grad` against fake tensors (which fails with the `allow_unused=True` error because fake tensors have no real autograd graph). The MACE `prepare()` pattern works only when `autograd.grad` lives inside the compiled function's `forward()` — not when it's in a separate method called from outside. 8. **TunableOp live tuning (`PYTORCH_TUNABLEOP_TUNING=1`) is unsafe on MI355X / ROCm 7.2.2 with HydraGNN.** All 16 GPUs hit `Memory access fault by GPU node-N` during the hipBLASLt kernel autotuning phase (loop-c3e4df1c iter-3 / job 7188). Use a 2-phase pattern instead: dedicated warmup sbatch with `PYTORCH_TUNABLEOP_TUNING=1` to generate `tunableop_results_.csv` files, then production runs with `PYTORCH_TUNABLEOP_ENABLED=1` only (no `_TUNING=1`). The `tunable_op_warmup_then_use` lever in the catalog encodes this. The original `tunable_op_live` lever is now marked `status: blocked`. -9. **[REVISED — see lesson #18]** ~~Lever payoff is node-class-dependent.~~ `num_workers_12` improved one b1+b2 run (-12.4% in loop-43b33ec1) and regressed a later b1+b2 run (+10.7% in loop-c3e4df1c). Both ran on the *same* b-pair; the apparent class dependence was run-to-run variance on b1+b2, not a-vs-b. The lever picker should treat `num_workers_*` results as high-variance until replicated within ±5% on the same node pair. -10. **[REVISED — see lesson #18]** ~~b-nodes are ~17% faster baseline than a-nodes for HydraGNN.~~ Both compared baselines (loop-43b33ec1 at 92.5 s/epoch and loop-c3e4df1c at 76.0 s/epoch) ran on the same `lux-mi355x-b[1-2]` pair. The 17% gap is a run-to-run delta on a single pair, not a class difference. Always read `nodes_list` from `iter-0-baseline.json` before drawing any node-class conclusion from a FOM delta. +9. **Lever payoff can be run-to-run noise, not node-class signal.** `num_workers_12` improved one run (-12.4%) and regressed another (+10.7%), both on the same node pair on different days. The lever picker should treat `num_workers_*` results as high-variance until replicated within ±5% on the same node pair. +10. **Always read `nodes_list` from `iter-0-baseline.json` before drawing any class conclusion from a FOM delta.** Two baselines that appear to "use different node classes" can in fact be the same physical nodes scheduled by SLURM on different days — especially when one iter-1 failed and the successful iter-0 landed wherever SLURM placed it. 11. **`hydragnn/train/__init__.py` re-export shadow** breaks naive `import hydragnn.train.train_validate_test as tvt`. The line `from .train_validate_test import train_validate_test, train, ...` makes `hydragnn.train.train_validate_test` resolve to the **function** `train_validate_test` (re-exported as attribute of the package), not the module. Use `importlib.import_module("hydragnn.train.train_validate_test")` instead. Applies to any rank hook that needs to monkey-patch `train_validate_test.train`. -12. **`sbatch --export=ALL,...` does NOT include shell vars that are not in `--export`'s explicit list.** When using `--export=ALL,K1=V1,K2=V2`, vars set in the calling shell but not listed here are propagated only if they're in the user's persistent env (login shell init). Always explicitly list cluster-pathing vars like `AI4S_SHARED_DIR=/shared/aaji` in `--export`. +12. **`sbatch --export=ALL,...` does NOT include shell vars that are not in `--export`'s explicit list.** When using `--export=ALL,K1=V1,K2=V2`, vars set in the calling shell but not listed here are propagated only if they're in the user's persistent env (login shell init). Always explicitly list cluster-pathing vars like `AI4S_SHARED_DIR` in `--export`. ## Lessons captured (pathc-pathb-a3b8f3f0 + loop-e3fac2af, 2026-05-23) -13. **fp32-vs-fp64 epoch-time comparison is the canonical low-cost dispatch-vs-compute discriminator.** On gfx950 (MI355X), the fp32 MFMA peak is **2× the fp64 MFMA peak**. A compute-bound workload should show ≥30% speedup when switching fp64→fp32. HydraGNN GFM-MLIP MACE on b-nodes showed ~0% speedup (78→79 s, +1.3% within noise, fp32 verified by `precision-diagnostic` rank-0 logs `model_param_dtype=torch.float32 first_batch_float_dtype=torch.float32`). This proves the workload is dispatch-bound. Run this comparison ONCE per workload-class as a single sbatch; the verdict shapes the entire subsequent lever strategy: +13. **fp32-vs-fp64 epoch-time comparison is the canonical low-cost dispatch-vs-compute discriminator.** On gfx950 (MI355X), the fp32 MFMA peak is **2× the fp64 MFMA peak**. A compute-bound workload should show ≥30% speedup when switching fp64→fp32. HydraGNN GFM-MLIP MACE showed ~0% speedup (78→79 s, +1.3% within noise, fp32 verified by `precision-diagnostic` rank-0 logs `model_param_dtype=torch.float32 first_batch_float_dtype=torch.float32`). This proves the workload is dispatch-bound. Run this comparison ONCE per workload-class as a single sbatch; the verdict shapes the entire subsequent lever strategy: - **fp32 ≈ fp64 → dispatch-bound** → compute-side levers (`batch_*`, `torch_compile_*`, `tunable_op_*`, `precision_fp32`) are dead ends. Focus on kernel-trace analysis to identify the dispatch culprit, then either upstream kernel fusion or reduce launch count (larger blocks, mega-kernels). - **fp32 ≪ fp64 → compute-bound** → standard compute-side levers apply. fp32 itself becomes a real lever (within numerical tolerance). - **fp32 > fp64 → memory-bound** → bandwidth-targeting levers (batch packing, fewer kernel launches over more data) apply. @@ -100,18 +97,17 @@ The first 2-node end-to-end run on Lux exposed five real bugs that are now all w - `torch.compile` path: move `energy_force_loss` into `EnhancedModelWrapper.forward()` (~30-50 LOC; addresses AOT-double-backward; see lesson #6) - `torch.jit.script` path: unpack `**conv_args` at `MACEStack.py:397` (small diff but in the read-only `/opt/hydragnn-pkgs` overlay) 15. **TunableOp `_TUNING=1` is unsafe regardless of phase** on MI355X / ROCm 7.2.2 / PyTorch 2.10. Both the live-tuning lever (`tunable_op_live`, loop-c3e4df1c iter-3) and the warmup-then-use 2-phase pattern (`tunable_op_warmup_then_use`, loop-e3fac2af iter-2) hit the same `Memory access fault by GPU node-N` on all 16 GPUs during the hipBLASLt tuning routine. The 2-phase split doesn't help — the fault is in the tuning routine itself, not in the production consumption. **Both levers are marked `status: blocked` in the catalog.** Re-test only after a ROCm stack update with a documented hipBLASLt fix. -16. **[REVISED — see lesson #18]** ~~b-nodes hit the dispatch ceiling at the baseline.~~ All three loop baselines (jobids 7083, 7158, 7199) ran on `lux-mi355x-b[1-2]`. The "dispatch ceiling" framing is correct *for that specific node pair on that specific dataset*; whether it generalizes to a class is a question we now know we cannot answer from the loop data alone. What does still hold: every compute-side lever tested on b1+b2 (batch_800, num_workers_12, precision_fp32) regressed or was neutral, which on its own confirms dispatch-bound character — but not a node-class story. +16. **The "dispatch ceiling" conclusion is solid, but do NOT frame it as node-class-specific.** All three loop baselines ran on the same physical node pair; every compute-side lever tested (batch_800, num_workers_12, precision_fp32) regressed or was neutral. That confirms dispatch-bound character for this workload — but makes no statement about node classes. The microbench (lesson #18) later showed all tested nodes have identical HIP launch latency, so the ceiling is workload-driven, not hardware-class-driven. 17. **Cross-loop best is composed of a free finding + a small lever**: current best (75.0 s/epoch) = baseline-on-b1-b2 (76.0 s) + `rccl_high_priority` (-1.3%). The single accepted lever across 3 loops + 1 investigation is `rccl_high_priority` (`TORCH_NCCL_HIGH_PRIORITY=1 GPU_MAX_HW_QUEUES=2`); it should be baked into the baseline contract for HydraGNN GFM-MLIP going forward (now marked `status: accepted-baked-in` in the catalog). ## Lessons captured (a-vs-b parity microbench, 2026-05-26) -18. **The "a-nodes vs b-nodes 17% gap" was a confounded measurement, not a hardware/firmware story.** A 4-test microbench sweep across 12 a-nodes + 3 b-nodes (script: [`examples/microbench_node_health.sh`](../../../material_science/models/HydraGNN/examples/microbench_node_health.sh); raw outputs: `/shared/aaji/microbench-a-vs-b/`) showed: - - **Identical at the host level:** CPU model (Genoa, 128c/256t), kernel (6.8.0-110), NUMA layout (2× ~1.5 TB), NICs (6× ionic 400G + 2× bnxt 200G), GPUs (8× MI355X gfx950), VBIOS, firmware, PCIe link state, mounts, scaling governor. - - **Identical compute microbench:** HIP empty-kernel launch latency 3.5–3.8 µs/launch on every GPU of every node; H2D pinned 57.6 GB/s; NCCL single-node 8-GPU all_reduce 372 GB/s bus; STREAM single-NUMA TRIAD 332–336 GB/s. - - **One real per-node outlier:** `lux-mi355x-a5` showed degraded dual-NUMA STREAM COPY (179 GB/s vs cluster median 220 ± 3 GB/s) with 8% run-to-run jitter on the first probe and partial recovery on retest (219 GB/s COPY but anomalously high TRIAD variance). All other 11 a-nodes and 3 b-nodes are statistically equivalent (220 ± 3 GB/s). - - **Controlled 2-node sanity test:** `a3+a4` = 77.0 s/epoch vs `b3+b4` = 78.4 s/epoch with the same code/config; **a-pair is 1.8% faster**, well inside noise. - - **Root cause of the original 17% narrative:** all three "baseline" jobids (7083, 7158, 7199) ran on `lux-mi355x-b[1-2]` per their `nodes_list` field. The 92.5 s → 76 s spread was on the *same pair* on *different days*; likely cache/NFS-warmup variance, not class. The misattribution came from labeling loop-43b33ec1 "a-node" because its iter-1 *failed* on a6's broken mount — but the *successful* baseline iter-0 (with a6 excluded) landed on b-nodes via SLURM scheduling. - - **Action for future loops:** always `cat $LOOP_DIR/iter-0-baseline.json | jq .nodes_list` before drawing any class conclusion. The `launcher` and `synthesizer` subagents should both log the assigned nodelist front-and-center in their output. -19. **Per-node health regressions are real and worth a 30-second prologue probe.** In 2 weeks we hit three different per-node failure modes (a6: container bind-mount EIO; a10: PMIx/NCCL ring timeout; a5: NUMA bw degradation). Each manifested as a confusing application-level failure hours into a job. All three would have been caught at job start by a 30-s SLURM Prolog running [`microbench_node_health.sh`](../../../material_science/models/HydraGNN/examples/microbench_node_health.sh) (mount-write probe + STREAM dual-NUMA + NCCL single-node + HIP launch latency). Suggested to the cluster admin in `SYSADMIN_MESSAGE.md` (under `/shared/aaji/microbench-a-vs-b/`). +18. **A "node-class performance gap" narrative was debunked by a controlled microbench.** A 4-test microbench sweep (script: [`examples/microbench_node_health.sh`](../../../material_science/models/HydraGNN/examples/microbench_node_health.sh); raw outputs in untracked `$AI4S_SHARED_DIR/microbench-a-vs-b/`) found: + - All tested nodes had identical HW at the host level (same CPU model, NUMA layout, GPUs, NICs, firmware). + - HIP empty-kernel launch latency 3.5–3.8 µs/launch on every GPU of every node; NCCL single-node all-reduce and STREAM single-NUMA TRIAD were statistically equivalent across all nodes tested (one outlier node had degraded dual-NUMA STREAM COPY, flagged for sysadmin). + - A controlled 2-node sanity HydraGNN run using nodes from each "class" showed < 2% wall-time difference, well inside noise. + - **Root cause of the original gap narrative:** all three baseline loop runs landed on the same physical node pair. The 92.5 s → 76 s spread was run-to-run variance on the *same pair* on different days. The misattribution arose from labeling one loop as "class A" because its iter-1 failed on a broken-mount node — but the *successful* baseline iter-0 landed on a different pair via SLURM scheduling. + - **Action for future loops:** always `jq .nodes_list $LOOP_DIR/iter-0-baseline.json` before drawing any class conclusion. The `launcher` and `synthesizer` subagents should log the assigned nodelist front-and-center in their output. +19. **Per-node health regressions are real and worth a 30-second prologue probe.** In 2 weeks we hit three different per-node failure modes: container bind-mount EIO on one node, PMIx/NCCL ring timeout on another, dual-NUMA bandwidth degradation on a third. Each manifested as a confusing application-level failure hours into a job. All three would have been caught at job start by a 30-s SLURM Prolog running [`microbench_node_health.sh`](../../../material_science/models/HydraGNN/examples/microbench_node_health.sh) (mount-write probe + STREAM dual-NUMA + NCCL single-node + HIP launch latency). Recommend adding this as a SLURM Prolog to your cluster admin. The full list (with cross-cutting context) is at the end of [.cursor/skills/ai4science-studio/SKILL.md](../ai4science-studio/SKILL.md). When fixing a bug here, check there first and propagate the lesson. diff --git a/.gitignore b/.gitignore index 666c9b5..2911aba 100755 --- a/.gitignore +++ b/.gitignore @@ -62,6 +62,9 @@ vic_server.log # Local cluster config (site-specific, auto-created by agent) .cluster-config.yaml + +# Session-scoped agent handoff files (ephemeral, not for the repo) +HANDOFF.md .claude/settings.local.json nul diff --git a/material_science/models/HydraGNN/examples/microbench_node_health.sh b/material_science/models/HydraGNN/examples/microbench_node_health.sh index 474efb9..d2f2927 100755 --- a/material_science/models/HydraGNN/examples/microbench_node_health.sh +++ b/material_science/models/HydraGNN/examples/microbench_node_health.sh @@ -1,8 +1,8 @@ #!/usr/bin/env bash # microbench_node_health.sh # -# Per-node hardware/firmware/health survey for AMD MI355X (gfx950) nodes on -# Vultr Lux. Runs four micro-tests in ~30 wall-clock seconds: +# Per-node hardware/firmware/health survey for AMD MI355X (gfx950) nodes. +# Runs four micro-tests in ~30 wall-clock seconds: # 1. host inventory - CPU, NUMA, kernel, NICs, PCIe link, mounts, /tmp dd # 2. GPU + driver inventory - rocminfo, rocm-smi (vbios/fw/topo), PCIe link state # 3. CPU dual-NUMA STREAM - COPY/SCALE/ADD/TRIAD, 128 threads, NUMA-interleaved @@ -24,7 +24,7 @@ # and one summary line of pass/fail to stdout. # # Env vars (all optional, with defaults): -# OUT_DIR (default: /shared/$USER/microbench-node-health) +# OUT_DIR (default: $AI4S_SHARED_DIR/microbench-node-health) # HG_SIF (default: $AI4S_SHARED_DIR/images/pytorch_rocm7.2.2_ubuntu24.04_py3.12_pytorch_release_2.10.0.sif) # AI4S_SHARED_DIR (default: /shared/$USER) # STREAM_THRESHOLD_COPY_GBPS (default: 200) - dual-NUMA COPY pass floor @@ -33,8 +33,8 @@ # # SBATCH directives: #SBATCH --job-name=node-health -#SBATCH --partition=lux -#SBATCH --account=vultr_lux +#SBATCH --partition=gpu +#SBATCH --account=default #SBATCH --nodes=1 #SBATCH --ntasks=1 #SBATCH --cpus-per-task=128 @@ -47,11 +47,11 @@ set -u H=$(hostname -s) -OUT_DIR="${OUT_DIR:-/shared/$USER/microbench-node-health}" +OUT_DIR="${OUT_DIR:-${AI4S_SHARED_DIR:-/tmp}/microbench-node-health}" HOST_OUT="$OUT_DIR/$H" mkdir -p "$HOST_OUT" -AI4S_SHARED_DIR="${AI4S_SHARED_DIR:-/shared/$USER}" +AI4S_SHARED_DIR="${AI4S_SHARED_DIR:-/tmp}" # set AI4S_SHARED_DIR to your cluster's shared storage root HG_SIF="${HG_SIF:-${AI4S_SHARED_DIR}/images/pytorch_rocm7.2.2_ubuntu24.04_py3.12_pytorch_release_2.10.0.sif}" STREAM_THRESHOLD_COPY_GBPS="${STREAM_THRESHOLD_COPY_GBPS:-200}" STREAM_THRESHOLD_TRIAD_GBPS="${STREAM_THRESHOLD_TRIAD_GBPS:-200}" diff --git a/material_science/models/HydraGNN/examples/run_optimizer_loop.sh b/material_science/models/HydraGNN/examples/run_optimizer_loop.sh index bf27417..096aea8 100755 --- a/material_science/models/HydraGNN/examples/run_optimizer_loop.sh +++ b/material_science/models/HydraGNN/examples/run_optimizer_loop.sh @@ -17,7 +17,7 @@ # n_iters_budget Max iterations (recommend 5). # # Abort: -# Graceful: touch /shared/aaji/models/HydraGNN/perf-runs/loop-/STOP +# Graceful: touch $AI4S_SHARED_DIR/models/HydraGNN/perf-runs/loop-/STOP # Emergency: scancel + tmux kill-session -t hg-loop set -euo pipefail @@ -44,7 +44,7 @@ RECIPE_DIR="${REPO_ROOT}/material_science/models/HydraGNN/recipes/perf-optimizer ORCH_PROMPT="${RECIPE_DIR}/agents/orchestrator.md" # Loop-dir convention matches recipes/perf-optimizer-loop/README.md -AI4S_SHARED_DIR="${AI4S_SHARED_DIR:-/shared/aaji}" +: "${AI4S_SHARED_DIR:?AI4S_SHARED_DIR must be set (e.g. export AI4S_SHARED_DIR=/shared/\$USER)}" HG_BASE="${AI4S_SHARED_DIR}/models/HydraGNN" PERF_RUNS_DIR="${HG_BASE}/perf-runs" LOOP_DIR="${PERF_RUNS_DIR}/loop-${LOOP_UUID}" @@ -72,10 +72,10 @@ PREFLIGHT_NOTES=() # 1. cluster up if command -v sinfo > /dev/null 2>&1; then - _IDLE_OR_MIX=$(sinfo -p lux,rad -h -o '%t' 2>/dev/null | grep -cE '^(idle|mix|alloc)$' || true) + _IDLE_OR_MIX=$(sinfo -h -o '%t' 2>/dev/null | grep -cE '^(idle|mix|alloc)$' || true) if [[ "${_IDLE_OR_MIX:-0}" -eq 0 ]]; then PREFLIGHT_FAIL_REASON="cluster_down" - PREFLIGHT_NOTES+=("sinfo reports 0 usable nodes on partitions lux,rad") + PREFLIGHT_NOTES+=("sinfo reports 0 usable nodes — check partition names in .cluster-config.yaml") fi else PREFLIGHT_FAIL_REASON="no_slurm" @@ -93,8 +93,8 @@ fi # 3. tools present for _bin in \ - /shared/aaji/tools/omnistat-pr271/bin/omnistat-usermode \ - /shared/aaji/tools/victoriametrics/victoria-metrics-prod; do + "${AI4S_SHARED_DIR}/tools/omnistat-pr271/bin/omnistat-usermode" \ + "${AI4S_SHARED_DIR}/tools/victoriametrics/victoria-metrics-prod"; do if [[ -z "$PREFLIGHT_FAIL_REASON" && ! -x "$_bin" ]]; then PREFLIGHT_FAIL_REASON="tool_missing" PREFLIGHT_NOTES+=("missing: $_bin (run the perf-analysis launcher subagent first to install)") @@ -104,7 +104,7 @@ done # 4. SIF + overlay for _f in \ "${HG_BASE}/overlays/hydragnn-overlay.img" \ - "${AI4S_SHARED_DIR}/images/pytorch_rocm7.2.2_ubuntu24.04_py3.12_pytorch_release_2.10.0.sif"; do + "${AI4S_SHARED_DIR}/images/pytorch_rocm7.2.2_ubuntu24.04_py3.12_pytorch_release_2.10.0.sif"; do # set AI4S_SHARED_DIR to your cluster's shared storage root if [[ -z "$PREFLIGHT_FAIL_REASON" && ! -f "$_f" ]]; then PREFLIGHT_FAIL_REASON="image_missing" PREFLIGHT_NOTES+=("missing: $_f") @@ -151,7 +151,7 @@ if [[ -n "$PREFLIGHT_FAIL_REASON" ]]; then fi _DISK_FREE=$(df -h "$PERF_RUNS_DIR" | awk 'NR==2 {print $4}') -_log_status "PREFLIGHT_OK cluster=lux disk_free=${_DISK_FREE} claude_cli=ok api_egress=ok orch_prompt=ok" +_log_status "PREFLIGHT_OK disk_free=${_DISK_FREE} claude_cli=ok api_egress=ok orch_prompt=ok" if [[ $PREFLIGHT_ONLY -eq 1 ]]; then _log_status "PREFLIGHT_ONLY mode: exiting before invoking orchestrator" diff --git a/material_science/models/HydraGNN/examples/sbatch_train_perf_amd.sh b/material_science/models/HydraGNN/examples/sbatch_train_perf_amd.sh index fd1f03a..7780818 100755 --- a/material_science/models/HydraGNN/examples/sbatch_train_perf_amd.sh +++ b/material_science/models/HydraGNN/examples/sbatch_train_perf_amd.sh @@ -8,21 +8,20 @@ # 1. Wraps srun with omnistat-usermode --start/--stopexporters/--stopserver. # 2. Generates a per-job copy of gfm_mlip.json with a "Profile" block injected # so HydraGNN's built-in torch.profiler captures rank-0 of node-0 only. -# 3. Hard-codes the lux partition / vultr_lux account so this script works -# out-of-the-box on Lux. Override via SBATCH_PARTITION/SBATCH_ACCOUNT if -# submitting to a different cluster. +# 3. Defaults to a cluster partition/account that can be overridden via +# SBATCH_PARTITION/SBATCH_ACCOUNT env vars for your cluster. # 4. Defaults are tuned for a quick perf run: --nodes=2, NUM_EPOCH=2, # MAX_NUM_BATCH=30, time=00:30:00 — enough for the wait=5/warmup=3/active=3 # profiler schedule plus a clean epoch 0 for warm caches. # # Quick start: -# export AI4S_SHARED_DIR=/shared/aaji +# export AI4S_SHARED_DIR=/shared/$USER # sbatch material_science/models/HydraGNN/examples/sbatch_train_perf_amd.sh # # Required: # AI4S_SHARED_DIR — shared base path -# /shared/aaji/tools/omnistat-pr271/ (created by the launcher subagent) -# /shared/aaji/tools/victoriametrics/ (created by the launcher subagent) +# $AI4S_SHARED_DIR/tools/omnistat-pr271/ (created by the launcher subagent) +# $AI4S_SHARED_DIR/tools/victoriametrics/ (created by the launcher subagent) # # Optional env-var overrides (with defaults): # HG_NUM_EPOCH=2 @@ -33,8 +32,8 @@ # OMNISTAT_USERMODE_INTERVAL=1 # seconds #SBATCH --job-name=hydragnn-perf -#SBATCH --partition=lux -#SBATCH --account=vultr_lux +#SBATCH --partition=gpu +#SBATCH --account=default #SBATCH --nodes=2 #SBATCH --ntasks-per-node=8 #SBATCH --gpus-per-node=8 @@ -71,7 +70,7 @@ HG_REPO_DIR="${HG_REPO_DIR:-${HG_BASE}/code/HydraGNN}" HYDRAGNN_MAX_NUM_BATCH="${HYDRAGNN_MAX_NUM_BATCH:-30}" PROFILE_TARGET_EPOCH="${PROFILE_TARGET_EPOCH:-1}" -OMNISTAT_VENV="${OMNISTAT_VENV:-/shared/aaji/tools/omnistat-pr271}" +OMNISTAT_VENV="${OMNISTAT_VENV:-${AI4S_SHARED_DIR}/tools/omnistat-pr271}" # Read the omnistat config from the in-repo recipe (single source of truth). # Override with OMNISTAT_TEMPLATE=/path/to/file if you need a custom probe config. # A staged copy under ${HG_BASE}/perf-runs/ is intentionally NOT used as the @@ -87,7 +86,7 @@ OMNISTAT_USERMODE_INTERVAL="${OMNISTAT_USERMODE_INTERVAL:-1}" # See material_science/models/HydraGNN/recipes/perf-analysis/agents/launcher.md # for how to build libomnistat_trace.so on a compute node. OMNISTAT_KERNEL_TRACE="${OMNISTAT_KERNEL_TRACE:-0}" -OMNISTAT_TRACE_LIB="${OMNISTAT_TRACE_LIB:-/shared/aaji/tools/omnistat-src/build-trace/libomnistat_trace.so}" +OMNISTAT_TRACE_LIB="${OMNISTAT_TRACE_LIB:-${AI4S_SHARED_DIR}/tools/omnistat-src/build-trace/libomnistat_trace.so}" # IMPORTANT: the kernel-trace collector registers /kernel_trace on the SAME # Flask app as the other omnistat endpoints, so the tool library must POST to # the [omnistat.collectors] `port` (8101 here), NOT the library's own default @@ -173,9 +172,9 @@ if [[ "$OMNISTAT_KERNEL_TRACE" == "1" ]]; then echo "ERROR: OMNISTAT_KERNEL_TRACE=1 but tool library not found:" >&2 echo " $OMNISTAT_TRACE_LIB" >&2 echo " Build it on a compute node:" >&2 - echo " salloc -p lux -A vultr_lux -N 1 --time=00:15:00 --gpus-per-node=1 \\" >&2 + echo " salloc -p -A -N 1 --time=00:15:00 --gpus-per-node=1 \\" >&2 echo " apptainer exec --rocm \"\$HG_SIF\" bash -c '" >&2 - echo " cd /shared/aaji/tools/omnistat-src && \\" >&2 + echo " cd \$AI4S_SHARED_DIR/tools/omnistat-src && \\" >&2 echo " cmake -S rocprofiler-sdk/ -B build-trace/ -DBUILD_KERNEL_TRACE_LIB=ON && \\" >&2 echo " cmake --build build-trace/ -j 8'" >&2 exit 2 diff --git a/material_science/models/HydraGNN/recipes/perf-optimizer-loop/HANDOFF.md b/material_science/models/HydraGNN/recipes/perf-optimizer-loop/HANDOFF.md deleted file mode 100644 index cf10fa9..0000000 --- a/material_science/models/HydraGNN/recipes/perf-optimizer-loop/HANDOFF.md +++ /dev/null @@ -1,115 +0,0 @@ -# HydraGNN perf-optimizer-loop — handoff to the next agent - -> **Read this first.** Single-file orientation for any agent picking up this work mid-flight. Keyword-rich title so a glob/grep for `HANDOFF`, `next agent`, `perf-optimizer`, `MI355X`, `dispatch-bound`, `HydraGNN handoff` will surface it. - -**Last touched:** 2026-05-26 by aaji. -**Branch:** `aaji/perf-optimizer-loop-hydragnn` (6 local commits ahead of origin; NOT pushed yet). -**Ticket:** [DCSWEAP-4502](https://amd.atlassian.net/browse/DCSWEAP-4502) "Run HydraGNN on MI355X". Latest comment id `22139024` (posted 2026-05-26 by aaji) is the current state of record. - -## TL;DR — where the work is - -The 2-node MI355X performance characterization on Vultr Lux is **done**. Current best for HydraGNN GFM-MLIP MACE on 2 nodes (16× MI355X) is **75.0 s/epoch** with this exact env (already the sbatch default): - -```bash -HG_BATCH_SIZE=400 -HG_PRECISION=fp64 -HYDRAGNN_NUM_WORKERS=8 -HYDRAGNN_PERSISTENT_WORKERS=1 -HG_NUM_EPOCH=3 -HYDRAGNN_MAX_NUM_BATCH=50 -PROFILE_TARGET_EPOCH=2 -TORCH_NCCL_HIGH_PRIORITY=1 # the only accepted lever across 3 loops -GPU_MAX_HW_QUEUES=2 # ditto -``` - -The workload is **dispatch-bound** (HIP launch latency ceiling 3.7 µs × 270 K launches/s, identical on every MI355X we probed). fp32-vs-fp64 differs by < 2 % despite fp32 MFMA having 2× the peak, which is the canonical proof. The path to a step-change requires **two upstream HydraGNN PRs** (see "Next work" below). - -## Where every artefact lives - -| What | Where | Tracked? | -|---|---|---| -| Recipe + agent prompts | `material_science/models/HydraGNN/recipes/perf-optimizer-loop/` | git | -| Lever catalog (machine-readable; blocked levers have evidence + path-to-unblock) | `recipes/perf-optimizer-loop/lever_catalog.yaml` | git | -| Sbatch wrapper (defaults = current best config) | `material_science/models/HydraGNN/examples/sbatch_train_perf_amd.sh` | git | -| Optimizer-loop driver (Claude Code CLI, tmux) | `material_science/models/HydraGNN/examples/run_optimizer_loop.sh` | git | -| Reusable node-health microbench (mount + STREAM + HIP launch + GPU/firmware inventory; ~30 s) | `material_science/models/HydraGNN/examples/microbench_node_health.sh` | git | -| All lessons learned (perf-analysis-specific) | `.cursor/skills/ai4science-perf-analysis/SKILL.md` (lessons 1-19) | git | -| Per-loop runtime artefacts (3 loops + 2 investigations) | `/shared/aaji/models/HydraGNN/perf-runs/loop-*` and `/shared/aaji/models/HydraGNN/perf-runs/{compile-investigation-*,pathc-pathb-*}` | NOT tracked | -| Microbench raw outputs + REPORT.md + SYSADMIN_MESSAGE.md | `/shared/aaji/microbench-a-vs-b/` | NOT tracked | -| Container (PyTorch 2.10.0 + ROCm 7.2.2 + HydraGNN overlay) | `/shared/aaji/images/pytorch_rocm7.2.2_ubuntu24.04_py3.12_pytorch_release_2.10.0.sif` | NOT tracked | -| Top-level ticket | [DCSWEAP-4502](https://amd.atlassian.net/browse/DCSWEAP-4502) | external | - -## Reproduce the current best in one command - -```bash -ssh aaji@rad-vultr-login # or whoever the next user is -cd /home/aaji/git/ai4science-studio -git checkout aaji/perf-optimizer-loop-hydragnn -export AI4S_SHARED_DIR=/shared/aaji -sbatch material_science/models/HydraGNN/examples/sbatch_train_perf_amd.sh -# expect: 75.0 s/epoch, train_validate_test ≈ 230 s, J/sample ≈ 1.631 -``` - -## Next work (priority order, from the ticket) - -| # | Task | Estimate | Notes | -|---|---|---|---| -| 1 | **8-node + 16-node strong-scaling sweep** with current best config | 1-2 hours wall, mostly queue time | Cluster cap is 16 nodes. Modify `sbatch_train_perf_amd.sh` `--nodes=N` only. Capture epoch time + J/sample + scaling efficiency vs the 2-node baseline. Write into a new `recipes/perf-optimizer-loop/scaling-results/` markdown + foms.csv. | -| 2 | **Open the two HydraGNN upstream PRs** | a few days each | (a) Move `energy_force_loss` into `EnhancedModelWrapper.forward` (~30-50 LOC across `hydragnn/models/create.py:622-758` + `hydragnn/train/train_validate_test.py:725-735`) — unblocks `torch.compile`. (b) Unpack `**conv_args` at `hydragnn/models/MACEStack.py:397` — unblocks `torch.jit.script`. Reference issues: see `lever_catalog.yaml` entries `torch_compile_e3nn` and `torch_jit_script` for the full evidence + diff sketches. Cite `/shared/aaji/models/HydraGNN/perf-runs/compile-investigation-aa480554-*/INVESTIGATION_REPORT.md` and `/shared/aaji/models/HydraGNN/perf-runs/pathc-pathb-a3b8f3f0-*/pathc-result.md` in the PRs. | -| 3 | **Second-user reproducibility** | half day | Have a different user run the one-command repro above on a different login session. If they hit 75 ± 1 s/epoch, ticket acceptance #7 is satisfied. | -| 4 | **Frontier-scale path documentation (DDStore phase 2)** | document only | Cannot validate on Vultr Lux (cap=16); write up the 504-node path in `recipes/perf-optimizer-loop/scaling-results/frontier-path.md`. Reference DDStore upstream docs. | - -## Landmines to NOT step on (these all cost real hours to discover) - -Each item is cross-referenced to a lesson in `.cursor/skills/ai4science-perf-analysis/SKILL.md` for the full backstory. - -| Don't | Why | Lesson | -|---|---|---| -| Don't re-try `torch.compile` on HydraGNN MLIP without first landing PR #2a above | AOT-Autograd doesn't support double backward; `energy_force_loss` calls `autograd.grad(create_graph=True)`. Tested 6 hooks across 3 hypotheses, all fail identically. | #6, #7 | -| Don't re-try `torch.jit.script` without first landing PR #2b above | TorchScript JIT can't expand `**conv_args` at `MACEStack.py:397`. | #14 | -| Don't enable `PYTORCH_TUNABLEOP_TUNING=1` in any form (live, warmup, warmup-then-use) | All variants fault all 16 GPUs in the hipBLASLt tuning routine. Re-test only after a documented ROCm/hipBLASLt fix. | #8, #15 | -| Don't draw a-nodes-vs-b-nodes conclusions from FOM deltas | They're physically identical (verified 2026-05-26); the historical "17 % gap" was three baselines on the same b1+b2 pair. Always `jq .nodes_list iter-0-baseline.json` first. | #10, #18 | -| Don't pick `num_workers_12` based on the loop-43b33ec1 win | Same lever regressed in loop-c3e4df1c on the same node pair; it's noise > effect at this measurement granularity. | #9, #18 | -| Don't pick `batch_800` for HydraGNN MLIP MACE | Dispatch saturates at batch=400; +87 % at batch=800. The R2 recommendation was from a pre-MLIP HydraGNN config. | lever_catalog `batch_800` notes | -| Don't pick `precision_fp32` expecting a speedup | The workload is dispatch-bound; fp32 gives < 2 % difference. (It IS still useful as a one-shot dispatch-vs-compute discriminator on new workloads — see lesson #13.) | #13 | -| Don't pick `nccl_minchannels=112` at N=2 | +10.3 % at N=2; only meaningful at ≥8 nodes. | lever_catalog `nccl_minchannels` notes | -| Don't omit `AI4S_SHARED_DIR` from `sbatch --export=ALL,...` | `--export=ALL` doesn't include unexported shell vars; sbatch dies on line 61 with `AI4S_SHARED_DIR must be set`. | #12 | -| Don't use `--gpus-per-node=N` for srun probes | SLURM filter rejects with "Invalid GRES specification"; use `--gres=gpu:amd_instinct_mi355_oam:N` instead. | (sbatch wrapper) | -| Don't claim a class-wide node bandwidth gap from a 16-GB memory cgroup | `cpu_mem_bw.sh`-style measurements get an allocation-artefact ceiling at ~220 GB/s; the real hardware delivers ~313 GB/s when allocated full-node. The committed `microbench_node_health.sh` works under either allocation but the threshold is calibrated to 200 GB/s for the cgroup case. | (microbench REPORT) | -| Don't blacklist a node *class* on a single failure | Probe per-node with the committed mount-health probe (exit 42) or `microbench_node_health.sh` and exclude only the named bad node(s). | #1, #18 | - -## Per-node health to be aware of (sysadmin will resolve) - -These are observations to relay to the cluster admin (draft is at `/shared/aaji/microbench-a-vs-b/SYSADMIN_MESSAGE.md`), NOT ticket items. If you land on one of these nodes, exclude by name and continue: - -| Node | Symptom | Workaround | -|---|---|---| -| `lux-mi355x-a5` | Intermittent dual-NUMA STREAM degradation (179 vs 220 GB/s COPY, 8 % jitter) | `--exclude=lux-mi355x-a5` | -| `lux-mi355x-a6` | `/home` + `/shared` bind-mount EIO inside Apptainer (historical 2026-05-22) | Mount probe (exit 42) auto-detects; orchestrator retries with `--exclude` | -| `lux-mi355x-a10` | PMIx ring collective times out at MPI init (historical 2026-05-23) | `--exclude=lux-mi355x-a10` | - -## Bootstrap checklist for a brand-new agent - -1. **Read [`README.md`](README.md)** in this directory for the recipe topology + iteration shape. -2. **Read [`lever_catalog.yaml`](lever_catalog.yaml)** in this directory — entries marked `status: blocked` or `status: accepted-baked-in` matter most. -3. **Read** `.cursor/skills/ai4science-perf-analysis/SKILL.md` lessons #1-19, especially the **REVISED** markers on lessons 9, 10, 16. -4. **`git log --oneline -8`** on this branch to see the change history with one-liners. -5. **`cat /shared/aaji/models/HydraGNN/perf-runs/loop-*/foms.csv`** to see the actual FOM trajectories across iterations. -6. **Check ticket** [DCSWEAP-4502](https://amd.atlassian.net/browse/DCSWEAP-4502), especially the latest comment (id `22139024`, posted 2026-05-26). -7. **Pick one of "Next work" items above** and start. - -## Open questions / decisions for the human, NOT for the agent - -- Should the branch be pushed to `origin` and a draft PR opened against `main`? (User hasn't asked; agent should not push without explicit ask.) -- Should the two HydraGNN upstream PRs be opened by the agent or by a human? (They're against `github.com/ORNL/HydraGNN`, a third-party repo; agent can prepare patches but the human should be the author of record.) -- Is there appetite for adding `microbench_node_health.sh` to other models' recipe folders (StormCast, ORBIT-2, GP-MoLFormer)? The script is generic enough to apply to any MI355X workload on Lux. - -## Authors / contacts - -- **Performance work + this branch:** Ashwin Aji (`ashwin.aji@amd.com`, github: TBD) -- **DCSWEAP-4502 reporter:** Nicholas Malaya (`nicholas.malaya@amd.com`) -- **Upstream HydraGNN maintainers:** ORNL — see [`github.com/ORNL/HydraGNN`](https://github.com/ORNL/HydraGNN) issues for the right tag - -## If anything in this file goes stale - -Update this file in the same commit as the change that made it stale. Single-source-of-truth for handoff state — if it disagrees with reality, fix the file, don't write a "HANDOFF-v2.md". diff --git a/material_science/models/HydraGNN/recipes/perf-optimizer-loop/README.md b/material_science/models/HydraGNN/recipes/perf-optimizer-loop/README.md index 91d7afd..de82025 100644 --- a/material_science/models/HydraGNN/recipes/perf-optimizer-loop/README.md +++ b/material_science/models/HydraGNN/recipes/perf-optimizer-loop/README.md @@ -51,10 +51,10 @@ flowchart TD ## Quick start (overnight unattended via Claude Code CLI in tmux) ```bash -ssh aaji@rad-vultr-login +ssh export ANTHROPIC_API_KEY=... # user sets manually; never committed -cd /home/aaji/git/ai4science-studio -git checkout aaji/perf-optimizer-loop-hydragnn +export AI4S_SHARED_DIR= +cd tmux new -s hg-loop bash material_science/models/HydraGNN/examples/run_optimizer_loop.sh $(uuidgen) 5 # Ctrl-b d to detach. tmux session lives until login-node reboot. @@ -63,14 +63,14 @@ bash material_science/models/HydraGNN/examples/run_optimizer_loop.sh $(uuidgen) Morning workflow: ```bash -cat /shared/aaji/models/HydraGNN/perf-runs/loop-/STATUS.txt | tail -50 -cat /shared/aaji/models/HydraGNN/perf-runs/loop-/story.md +cat $AI4S_SHARED_DIR/models/HydraGNN/perf-runs/loop-/STATUS.txt | tail -50 +cat $AI4S_SHARED_DIR/models/HydraGNN/perf-runs/loop-/story.md ``` ## Artifact layout (per-loop) ``` -/shared/aaji/models/HydraGNN/perf-runs/ +$AI4S_SHARED_DIR/models/HydraGNN/perf-runs/ ├── loop-.json # registers loop id, baseline jobid, best jobid ├── loop-/ │ ├── STATUS.txt # append-only event log (driver writes after every step) @@ -145,13 +145,13 @@ See [`lever_catalog.yaml`](lever_catalog.yaml) for the full machine-readable lis ## Abort controls -- **Graceful:** `touch /shared/aaji/models/HydraGNN/perf-runs/loop-/STOP`. Driver checks the flag (a) before every `sbatch`, (b) before every analyst phase, (c) before every iteration decision. Active SLURM jobs run to completion and get analyzed; no further iterations launch. +- **Graceful:** `touch $AI4S_SHARED_DIR/models/HydraGNN/perf-runs/loop-/STOP`. Driver checks the flag (a) before every `sbatch`, (b) before every analyst phase, (c) before every iteration decision. Active SLURM jobs run to completion and get analyzed; no further iterations launch. - **Emergency:** `scancel ` for the active job + `tmux kill-session -t hg-loop` (or `kill -INT `). The driver traps SIGINT and writes `LOOP_ABORT reason=signal` to STATUS.txt before exiting. ## Disk budget - Per iteration footprint: ~135 MB (matches R2 shape: kineto trace 110-125 MB + TraceLens 8 MB + omnistat-db 1-3 MB + omnistat inspect ~200 KB). -- 5-iter loop projected: ~0.7 GB. `/shared` has 26 TB free (18% used); not pressured. +- 5-iter loop projected: ~0.7 GB. Check `df -h $AI4S_SHARED_DIR` before starting. - No per-loop quota enforced; orchestrator logs `du -sh loop-/` after each iter to STATUS.txt for visibility. - `OMNISTAT_KERNEL_TRACE=1` raises VictoriaMetrics cardinality ~1000× (per [SKILL §12](../../../../../.cursor/skills/ai4science-studio/SKILL.md)). The `kernel_trace_diag_only` lever is gated `diagnostic_only=true` and only fires when the 1-s correlation falls short — at most 1 of 5 iterations. diff --git a/material_science/models/HydraGNN/recipes/perf-optimizer-loop/agents/fom_extractor.md b/material_science/models/HydraGNN/recipes/perf-optimizer-loop/agents/fom_extractor.md index 5ee453d..06bf828 100644 --- a/material_science/models/HydraGNN/recipes/perf-optimizer-loop/agents/fom_extractor.md +++ b/material_science/models/HydraGNN/recipes/perf-optimizer-loop/agents/fom_extractor.md @@ -68,7 +68,7 @@ omnistat kernel-trace collector). ```bash set -euo pipefail -. /shared/aaji/tools/omnistat-pr271/bin/activate +. "${AI4S_SHARED_DIR}/tools/omnistat-pr271/bin/activate" PERF_RUN=$(jq -r .perf_run_dir ) JOBID=$(jq -r .jobid ) RUNTIME=$(jq -r .runtime_seconds ) @@ -81,7 +81,7 @@ TRACE=$(jq -r '.trace_paths[0] // empty' ) Use the existing parser, extended with a `--json-foms` flag (see "parse_convergence.py extension" below): ```bash -python3 /home/aaji/git/ai4science-studio/material_science/models/HydraGNN/examples/parse_convergence.py \ +python3 "$REPO_ROOT/material_science/models/HydraGNN/examples/parse_convergence.py" \ --log "$PERF_RUN/hydragnn-train-$JOBID.out" \ --output "$PERF_RUN/convergence.csv" \ --json-foms "$PERF_RUN/_convergence_foms.json" @@ -102,7 +102,7 @@ Start a local VictoriaMetrics if not already running: PORT=8428 ss -lnt "sport = :$PORT" 2>/dev/null | grep -q LISTEN && PORT=8429 if ! curl -sf "http://127.0.0.1:$PORT/api/v1/status/tsdb" > /dev/null 2>&1; then - nohup /shared/aaji/tools/victoriametrics/victoria-metrics-prod \ + nohup "${AI4S_SHARED_DIR}/tools/victoriametrics/victoria-metrics-prod" \ -storageDataPath="$PERF_RUN/omnistat-db" \ -httpListenAddr=127.0.0.1:$PORT \ -retentionPeriod=100y -fs.disableMmap \ @@ -234,7 +234,7 @@ The orchestrator handles `partial` by recording the FOM row with `null`s and a n ## Notes for the implementing agent -- This subagent runs on the login node, after omnistat_analyst has started VictoriaMetrics on this perf-run. If VM isn't already up, start a second instance on a different port — the omnistat datadir is read-only friendly under VM. +- This subagent runs on the login node, after omnistat_analyst has started VictoriaMetrics on this perf-run. If VM isn't already up, start a second instance on a different port — the omnistat datadir is read-only friendly under VM. Requires `AI4S_SHARED_DIR` and `REPO_ROOT` to be set. - Do NOT call `omnistat-inspect`; that's the analyst's job. We go straight to PromQL here because we need finer-grained queries than what `omnistat-inspect` exposes. - Read at most: manifest.json, hydragnn-train-.out, the rank-0 trace, omnistat-db (via VM). Cap trace read at 4 GB. - Write at most: foms.json, kernel_correlation.csv, vm_for_foms.{log,pid}. diff --git a/material_science/models/HydraGNN/recipes/perf-optimizer-loop/agents/orchestrator.md b/material_science/models/HydraGNN/recipes/perf-optimizer-loop/agents/orchestrator.md index 5410c47..bc0f78a 100644 --- a/material_science/models/HydraGNN/recipes/perf-optimizer-loop/agents/orchestrator.md +++ b/material_science/models/HydraGNN/recipes/perf-optimizer-loop/agents/orchestrator.md @@ -10,13 +10,13 @@ on the login node, or directly inside a Cursor agent session). ## Inputs - Loop arguments (passed in the user message): ``, ``. -- Repo at `/home/aaji/git/ai4science-studio` on branch `aaji/perf-optimizer-loop-hydragnn`. +- Repo root read from the calling shell's working directory or `REPO_ROOT` env var. - Cluster config at `.cluster-config.yaml`. - [`lever_catalog.yaml`](../lever_catalog.yaml) — single source of truth for allowed levers. ## Outputs -Per loop, under `/shared/aaji/models/HydraGNN/perf-runs/loop-/`: +Per loop, under `$AI4S_SHARED_DIR/models/HydraGNN/perf-runs/loop-/`: - `STATUS.txt` — append-only event log; one line per event (see schema below). - `foms.csv` — one row per iteration: `iter,jobid,lever_id,env_diff,accepted,epoch_time_s,throughput_samples_per_s,mfma_tflops,energy_J,mean_power_W,energy_per_sample_J,final_loss,primary_fom_delta_pct,du_loop_dir,notes`. @@ -25,7 +25,7 @@ Per loop, under `/shared/aaji/models/HydraGNN/perf-runs/loop-/`: - `story.md` — written by story_writer at end of loop. - `foms.png` — written by story_writer at end of loop. -Per iteration, all artifacts already written by the existing `recipes/perf-analysis/` subagents land under `/shared/aaji/models/HydraGNN/perf-runs//` plus: +Per iteration, all artifacts already written by the existing `recipes/perf-analysis/` subagents land under `$AI4S_SHARED_DIR/models/HydraGNN/perf-runs//` plus: - `foms.json` — written by fom_extractor. - `kernel_correlation.csv` — written by fom_extractor. @@ -36,7 +36,7 @@ Per iteration, all artifacts already written by the existing `recipes/perf-analy 2. **Exactly one lever per iteration.** Multi-variable changes break delta attribution. 3. **Respect the STOP flag.** Check `loop-/STOP` before every sbatch submission, before every analyst phase, and before every iteration decision. If present, write `LOOP_ABORT reason=stop_flag`, finish processing the currently-running iter if any, then exit 0. 4. **Single source of truth for state is on disk.** `foms.csv`, `do_not_retry.json`, and STATUS.txt are the only state — no in-memory state survives a restart. The orchestrator must be able to resume mid-loop by reading these files. -5. **Never modify files outside `/shared/aaji/models/HydraGNN/perf-runs/` and `/tmp/`.** No repo edits during the loop. (Repo edits to capture lessons happen AFTER the loop, by a separate pass.) +5. **Never modify files outside `$AI4S_SHARED_DIR/models/HydraGNN/perf-runs/` and `/tmp/`.** No repo edits during the loop. (Repo edits to capture lessons happen AFTER the loop, by a separate pass.) 6. **All subagent invocations write structured JSON to disk and end with `STATUS=ok|partial|fail`.** The orchestrator parses the final line; never inline-trusts chat content. ## Loop control parameters @@ -87,10 +87,10 @@ Run these checks in sequence; abort with `PREFLIGHT_FAIL reason=` on the - `sinfo -p lux,rad -h` → if zero `IDLE` or `MIX` nodes, abort. Cluster may be in maintenance (see SKILL §14). - `df -h /shared | tail -1` → abort if `Use%` > 95. -- `test -x /shared/aaji/tools/omnistat-pr271/bin/omnistat-usermode` → abort if missing. -- `test -x /shared/aaji/tools/victoriametrics/victoria-metrics-prod` → abort if missing. -- `test -f /shared/aaji/models/HydraGNN/overlays/hydragnn-overlay.img` → abort if missing. -- `test -f /shared/aaji/images/pytorch_rocm7.2.2_ubuntu24.04_py3.12_pytorch_release_2.10.0.sif` → abort if missing. +- `test -x $AI4S_SHARED_DIR/tools/omnistat-pr271/bin/omnistat-usermode` → abort if missing. +- `test -x $AI4S_SHARED_DIR/tools/victoriametrics/victoria-metrics-prod` → abort if missing. +- `test -f $AI4S_SHARED_DIR/models/HydraGNN/overlays/hydragnn-overlay.img` → abort if missing. +- `test -f $AI4S_SHARED_DIR/images/pytorch_rocm7.2.2_ubuntu24.04_py3.12_pytorch_release_2.10.0.sif` → abort if missing. - If invoked via Claude Code CLI on the login node: `[[ -n "$ANTHROPIC_API_KEY" ]]` and `curl -fsS https://api.anthropic.com/v1/models > /dev/null` (5 s timeout) → abort if either fails. Write `PREFLIGHT_OK` line. @@ -123,7 +123,7 @@ if [[ -s "loop-/known_bad_nodes.txt" ]]; then EXCLUDE_ARG="--exclude=$(paste -sd, loop-/known_bad_nodes.txt)" fi ( set -a; source "loop-/iter--env.sh"; set +a; \ - cd /home/aaji/git/ai4science-studio; \ + cd "$REPO_ROOT"; \ sbatch $EXCLUDE_ARG material_science/models/HydraGNN/examples/sbatch_train_perf_amd.sh ) ``` @@ -145,7 +145,7 @@ Log `ITER_COMPLETE iter= jobid= state= runtime=`. 6. If the **same** node appears in `known_bad_nodes.txt` 3 times across the loop without admin intervention, log `LOOP_ABORT reason=persistent_node_fault bad_nodes=` and exit 2 — escalate to the human. 7. After loop ends, the story_writer reads `known_bad_nodes.txt` and surfaces the list to the user as a Lessons entry, so the cluster admin can be notified. -This handles the iter-1 case from loop-43b33ec1 correctly: only `lux-mi355x-a6` had broken mounts; `lux-mi355x-a1` in the same alloc was healthy. The orchestrator's job is to flag the specific bad node and route around it, not to penalize the lever or shrink the node pool by class. +This handles the iter-1 case from one of our pilot loops correctly: only the specific failing node had broken mounts; a sibling node in the same alloc was healthy. The orchestrator's job is to flag the specific bad node and route around it, not to penalize the lever or shrink the node pool by class. **2f. STOP-flag gate.** Re-check before the analysis phase. If set, still analyze the just-completed iter (we paid for it), then exit gracefully. diff --git a/material_science/models/HydraGNN/recipes/perf-optimizer-loop/agents/story_writer.md b/material_science/models/HydraGNN/recipes/perf-optimizer-loop/agents/story_writer.md index 9182203..2c4cad0 100644 --- a/material_science/models/HydraGNN/recipes/perf-optimizer-loop/agents/story_writer.md +++ b/material_science/models/HydraGNN/recipes/perf-optimizer-loop/agents/story_writer.md @@ -27,7 +27,7 @@ once the loop has completed (or aborted). Reads only files; writes only ```markdown # HydraGNN iterative-sysopt-loop — -**Hardware:** 2 nodes × 8 GPUs MI355X (gfx950, vultr_lux). +**Hardware:** 2 nodes × 8 GPUs MI355X (gfx950). **Image:** `pytorch_rocm7.2.2_ubuntu24.04_py3.12_pytorch_release_2.10.0`. **Loop:** iterations, started , ended , total wall hm. **Primary FOM:** `epoch_time_s` (mean over epoch >= 1, drops warm-up). @@ -156,7 +156,7 @@ fig.tight_layout() fig.savefig(out_png_path, dpi=150, bbox_inches="tight") ``` -Use `matplotlib.use("Agg")` for headless rendering. Do NOT use seaborn (extra dep not in the omnistat venv). Use the same omnistat venv `/shared/aaji/tools/omnistat-pr271/bin/python` which already has matplotlib via TraceLens deps. +Use `matplotlib.use("Agg")` for headless rendering. Do NOT use seaborn (extra dep not in the omnistat venv). Use the omnistat venv at `$AI4S_SHARED_DIR/tools/omnistat-pr271/bin/python` which already has matplotlib via TraceLens deps. ## Style rules diff --git a/material_science/models/HydraGNN/recipes/perf-optimizer-loop/lever_catalog.yaml b/material_science/models/HydraGNN/recipes/perf-optimizer-loop/lever_catalog.yaml index 65aace8..6507f56 100644 --- a/material_science/models/HydraGNN/recipes/perf-optimizer-loop/lever_catalog.yaml +++ b/material_science/models/HydraGNN/recipes/perf-optimizer-loop/lever_catalog.yaml @@ -42,11 +42,11 @@ levers: expected_payoff: "n/a (baseline)" risk: "low" revert_method: "drop_env" - citation: "/shared/aaji/models/HydraGNN/perf-runs/6985/combined_report.md" + citation: "$AI4S_SHARED_DIR/models/HydraGNN/perf-runs/6985/combined_report.md" notes: "Reference state. Establishes the FOM denominator for all subsequent iters." - id: batch_800 - description: "Double batch size: HG_BATCH_SIZE 400 -> 800. R2 #1 next-lever recommendation; VRAM headroom is 24/288 GB. CONFIRMED REGRESSOR for HydraGNN GFM-MLIP MACE: regressed on both tested loops (+56.8%, +87.2%) — both runs were on lux-mi355x-b[1-2]." + description: "Double batch size: HG_BATCH_SIZE 400 -> 800. R2 #1 next-lever recommendation; VRAM headroom is 24/288 GB. CONFIRMED REGRESSOR for HydraGNN GFM-MLIP MACE: regressed on both tested loops (+56.8%, +87.2%)." catalog_order: 1 kind: env_only env_vars: @@ -55,8 +55,8 @@ levers: expected_payoff: "none-for-MLIP-MACE" risk: "high" revert_method: "drop_env" - citation: "loop-43b33ec1 iter-2 (b1+b2, +56.76%) + loop-e3fac2af iter-1 (b1+b2, +87.2%)" - notes: "The lever_picker should NOT pick batch_800 for HydraGNN MLIP MACE. The MACE equivariant message-passing stack saturates dispatch at batch=400 and doubling work-per-iter doubles wall-time without proportional GPU occupancy gain. The R2 recommendation was based on an earlier non-MLIP HydraGNN config — DOES NOT apply to gfm_mlip_all_mpnn.py. [REVISED 2026-05-26] previously described as 'node-class-dependent' but in fact both tested loops ran on b1+b2; the regression is workload-specific, not node-class. Keep entry as cautionary record. May still be worth re-trying ONLY if upstream HydraGNN fuses small e3nn kernels first." + citation: "loop-43b33ec1 iter-2 (+56.76%) + loop-e3fac2af iter-1 (+87.2%)" + notes: "The lever_picker should NOT pick batch_800 for HydraGNN MLIP MACE. The MACE equivariant message-passing stack saturates dispatch at batch=400 and doubling work-per-iter doubles wall-time without proportional GPU occupancy gain. The R2 recommendation was based on an earlier non-MLIP HydraGNN config — DOES NOT apply to gfm_mlip_all_mpnn.py. Keep entry as cautionary record. May still be worth re-trying ONLY if upstream HydraGNN fuses small e3nn kernels first." - id: torch_compile_e3nn description: "BLOCKED — torch.compile is fundamentally incompatible with HydraGNN GFM-MLIP under PyTorch 2.10 / ROCm 7.2.2. Documented here so the lever_picker NEVER re-tries it." @@ -107,7 +107,7 @@ levers: expected_payoff: "none (blocked)" risk: "blocked" revert_method: "drop_hook" - citation: "/shared/aaji/models/HydraGNN/perf-runs/compile-investigation-aa480554-2a15-4e04-80dd-c58afc672a8d/INVESTIGATION_REPORT.md" + citation: "$AI4S_SHARED_DIR/models/HydraGNN/perf-runs/compile-investigation-aa480554-*/INVESTIGATION_REPORT.md" notes: "DO NOT pick this lever until HydraGNN upstream patches energy_force_loss into forward(). The lever_picker subagent must treat status:blocked as equivalent to do_not_retry." - id: torch_jit_script @@ -150,7 +150,7 @@ levers: expected_payoff: "none (blocked)" risk: "blocked" revert_method: "drop_hook" - citation: "/shared/aaji/models/HydraGNN/perf-runs/pathc-pathb-a3b8f3f0-19be-48b8-b81f-c7068e418289/pathc-result.md" + citation: "$AI4S_SHARED_DIR/models/HydraGNN/perf-runs/pathc-pathb-a3b8f3f0-*/pathc-result.md" notes: "DO NOT pick this lever until HydraGNN upstream unpacks **conv_args at MACEStack.py:397. catalog_order=2 shared with torch_compile_e3nn since both are 'kernel-fusion via PyTorch compiler' attempts that are both dead-ends on the current stack." - id: rccl_high_priority @@ -166,7 +166,7 @@ levers: risk: "low" revert_method: "drop_env" citation: "loop-c3e4df1c iter-2 / jobid 7187: -1.3% epoch_time (76.0->75.0 s) AND -8.4% energy/sample (1.781->1.631 J/sample). Mean power dropped 4885W->4552W." - notes: "ACCEPTED. The only lever across 3 loops that produced a real improvement. Now part of the baseline contract for HydraGNN GFM-MLIP. [REVISED 2026-05-26] previously framed as 'on b-nodes' but the validating runs were specifically on lux-mi355x-b[1-2]; a-vs-b microbench (2026-05-26) shows no class difference, so this likely applies to a-nodes too — re-verify on a-pair before declaring class-portability. lever_picker should treat this as 'already applied' and skip it in candidate set (foms.csv has it accepted with delta_pct=-1.316)." + notes: "ACCEPTED. The only lever across 3 loops that produced a real improvement. Now part of the baseline contract for HydraGNN GFM-MLIP. Microbench (2026-05-26) confirmed identical HIP dispatch latency across all tested MI355X nodes, so this improvement is expected to be portable. lever_picker should treat this as 'already applied' and skip it in candidate set (foms.csv has it accepted with delta_pct=-1.316)." - id: tunable_op_live description: "BLOCKED on this stack — TunableOp LIVE tuning mode (PYTORCH_TUNABLEOP_TUNING=1) causes GPU memory access faults on MI355X / ROCm 7.2.2 / PyTorch 2.10 with HydraGNN. Use tunable_op_warmup_then_use instead." @@ -196,13 +196,13 @@ levers: env_vars: PYTORCH_TUNABLEOP_ENABLED: "1" PYTORCH_TUNABLEOP_TUNING: "0" - PYTORCH_TUNABLEOP_FILENAME: "/shared/aaji/models/HydraGNN/tunableop_results/tunableop_results_%d.csv" + PYTORCH_TUNABLEOP_FILENAME: "$AI4S_SHARED_DIR/models/HydraGNN/tunableop_results/tunableop_results_%d.csv" rank_script_patch: | # 2-phase TunableOp pattern. # # Phase 1 (separate sbatch, NOT this lever) — orchestrator must run # this BEFORE picking the lever: - # 1. mkdir -p /shared/aaji/models/HydraGNN/tunableop_results + # 1. mkdir -p $AI4S_SHARED_DIR/models/HydraGNN/tunableop_results # 2. sbatch with PYTORCH_TUNABLEOP_ENABLED=1 PYTORCH_TUNABLEOP_TUNING=1 # PYTORCH_TUNABLEOP_FILENAME=/shared/.../tunableop_results_%d.csv # HG_NUM_EPOCH=1 HYDRAGNN_MAX_NUM_BATCH=5 @@ -219,17 +219,17 @@ levers: import os if int(os.environ.get("SLURM_PROCID", "0")) == 0: import glob - _n_csv = len(glob.glob("/shared/aaji/models/HydraGNN/tunableop_results/tunableop_results_*.csv")) + _n_csv = len(glob.glob(os.path.join(os.environ.get("AI4S_SHARED_DIR",""), "models/HydraGNN/tunableop_results/tunableop_results_*.csv"))) print(f"[tunable_op_warmup] consuming {_n_csv} pre-tuned CSV files", flush=True) diagnostic_only: false - expected_payoff: "medium" - risk: "low" + expected_payoff: "none (blocked)" + risk: "blocked" revert_method: "drop_env" citation: "https://rocm.docs.amd.com/en/latest/how-to/tuning-guides/mi300x/workload.html#tunableop" - notes: "If the warmup phase did NOT complete (no CSVs), pick a different lever. The orchestrator should add `setup: tunableop_warmup_sbatch` as a precondition before this lever is offered." + notes: "BLOCKED on ROCm 7.2.2 / PyTorch 2.10 / MI355X — warmup phase itself triggers GPU memory access faults. If the warmup phase did NOT complete (no CSVs), pick a different lever. Re-test only after a documented ROCm/hipBLASLt fix." - id: num_workers_12 - description: "DataLoader: HYDRAGNN_NUM_WORKERS 8 -> 12 (persistent_workers already on). HIGH RUN-TO-RUN VARIANCE: helped in loop-43b33ec1 (-12.4%) and regressed in loop-c3e4df1c (+10.7%) — both ran on lux-mi355x-b[1-2], so this is intra-pair variance, not class behavior. Treat as unreliable until replicated within ±5% on the same nodes." + description: "DataLoader: HYDRAGNN_NUM_WORKERS 8 -> 12 (persistent_workers already on). HIGH RUN-TO-RUN VARIANCE: helped in one loop (-12.4%) and regressed in another (+10.7%) — both on the same node pair, so this is run-to-run variance, not a systematic win. Treat as unreliable until replicated within ±5% on the same nodes." catalog_order: 5 kind: env_only env_vars: @@ -238,8 +238,8 @@ levers: expected_payoff: "unreliable (variance > effect)" risk: "low" revert_method: "drop_env" - citation: "local (loop 43b33ec1 + loop c3e4df1c, both on b1+b2)" - notes: "[REVISED 2026-05-26] previously believed to be node-class-dependent (a-nodes vs b-nodes), but both observations were on the same b1+b2 pair, so the result is run-to-run noise on a single pair. Recommended next step: re-run num_workers={8,10,12} three times each on the SAME node pair with the SAME warm/cold cache state. If the kernel_correlation.csv from fom_extractor shows dataloader prefetch gaps > 10 ms, that's a stronger signal than the epoch-time delta. Until then the lever_picker should rank this lever LOWER than levers with smaller variance." + citation: "local (loop 43b33ec1 + loop c3e4df1c, same node pair both times)" + notes: "Both observations were on the same node pair on different days, so the result is run-to-run noise. Recommended next step: re-run num_workers={8,10,12} three times each on the SAME node pair with the SAME warm/cold cache state. If the kernel_correlation.csv from fom_extractor shows dataloader prefetch gaps > 10 ms, that's a stronger signal than the epoch-time delta. Until then the lever_picker should rank this lever LOWER than levers with smaller variance." - id: nccl_minchannels description: "NCCL_MIN_NCHANNELS=112: forces more channels for collective dispatch, helps when message size is moderate. CONFIRMED REGRESSOR at N=2 (+10.3%)." @@ -266,7 +266,7 @@ levers: risk: "low" revert_method: "drop_env" citation: "loop-e3fac2af iter-4 (job 7203): +1.3% epoch_time (78->79 s) confirmed in fp32 (precision-diagnostic logs show model_param_dtype=torch.float32, first_batch_float_dtype=torch.float32)" - notes: "RAN GENUINELY IN FP32 (verified by precision-diagnostic). fp32 MFMA peak on gfx950 is 2x fp64 peak, so a compute-bound workload would show ~50% speedup. Observed ~0% speedup proves HydraGNN GFM-MLIP MACE is dispatch-bound. [REVISED 2026-05-26] previously hedged with 'on b-nodes; may differ on a-nodes'; a-vs-b microbench (2026-05-26) confirmed identical HIP launch latency (3.7 us/launch) on every tested node of both classes, so the dispatch ceiling is class-portable — fp32 will not help on a-nodes either. lever_picker should NOT re-pick this lever. bf16 remains INTENTIONALLY EXCLUDED: breaks equivariance for MACE-class features. The fp32-vs-fp64 comparison is now the canonical low-cost dispatch-vs-compute discriminator for any future HydraGNN-class workload — see ai4science-perf-analysis/SKILL.md lesson #13." + notes: "RAN GENUINELY IN FP32 (verified by precision-diagnostic). fp32 MFMA peak on gfx950 is 2x fp64 peak, so a compute-bound workload would show ~50% speedup. Observed ~0% speedup proves HydraGNN GFM-MLIP MACE is dispatch-bound. A-vs-b microbench (2026-05-26) confirmed identical HIP launch latency (3.7 us/launch) on every tested node of both classes, so the dispatch ceiling is class-portable. lever_picker should NOT re-pick this lever. bf16 remains INTENTIONALLY EXCLUDED: breaks equivariance for MACE-class features. The fp32-vs-fp64 comparison is now the canonical low-cost dispatch-vs-compute discriminator for any future HydraGNN-class workload — see ai4science-perf-analysis/SKILL.md lesson #13." - id: kernel_trace_diag_only description: "DIAGNOSTIC: enable OMNISTAT_KERNEL_TRACE=1 for ONE iteration to capture per-kernel-name dispatch counts/durations. Fires only when the 1-s TraceLens<->Omnistat correlation cannot attribute the dominant kernel." From 22ee3940fd805f5ba67c3b9eb9e4998446bc2656 Mon Sep 17 00:00:00 2001 From: Ashwin Aji Date: Thu, 28 May 2026 23:12:56 +0000 Subject: [PATCH 14/31] perf-optimizer-loop: capture job 7187 attribution pass Add run_fom_extractor.py (FOM + TraceLens/Omnistat kernel_correlation) and dispatch-attribution.md, plus lesson 20 in the perf-analysis skill on VictoriaMetrics windowing and reading TraceLens ops_summary before escalating to kernel-trace. Co-authored-by: Cursor --- .../skills/ai4science-perf-analysis/SKILL.md | 4 + .../HydraGNN/examples/run_fom_extractor.py | 427 ++++++++++++++++++ .../dispatch-attribution.md | 78 ++++ 3 files changed, 509 insertions(+) create mode 100755 material_science/models/HydraGNN/examples/run_fom_extractor.py create mode 100644 material_science/models/HydraGNN/recipes/perf-optimizer-loop/dispatch-attribution.md diff --git a/.cursor/skills/ai4science-perf-analysis/SKILL.md b/.cursor/skills/ai4science-perf-analysis/SKILL.md index 67f7dd8..f848be3 100644 --- a/.cursor/skills/ai4science-perf-analysis/SKILL.md +++ b/.cursor/skills/ai4science-perf-analysis/SKILL.md @@ -110,4 +110,8 @@ The first 2-node end-to-end run on Lux exposed five real bugs that are now all w - **Action for future loops:** always `jq .nodes_list $LOOP_DIR/iter-0-baseline.json` before drawing any class conclusion. The `launcher` and `synthesizer` subagents should log the assigned nodelist front-and-center in their output. 19. **Per-node health regressions are real and worth a 30-second prologue probe.** In 2 weeks we hit three different per-node failure modes: container bind-mount EIO on one node, PMIx/NCCL ring timeout on another, dual-NUMA bandwidth degradation on a third. Each manifested as a confusing application-level failure hours into a job. All three would have been caught at job start by a 30-s SLURM Prolog running [`microbench_node_health.sh`](../../../material_science/models/HydraGNN/examples/microbench_node_health.sh) (mount-write probe + STREAM dual-NUMA + NCCL single-node + HIP launch latency). Recommend adding this as a SLURM Prolog to your cluster admin. +## Lessons captured (attribution pass on job 7187, 2026-05-27) + +20. **Attribution pass tooling:** Use [`examples/run_fom_extractor.py`](../../../material_science/models/HydraGNN/examples/run_fom_extractor.py) + TraceLens on `$AI4S_SHARED_DIR/models/HydraGNN/perf-runs//`. Backfill `manifest.json` if missing (loop jobs before launcher wrote it). **VictoriaMetrics PromQL on perf-run DBs requires `time=` at the job window** — instant queries on the login node return empty series even when `omnistat_hardware_counter` data exists. **`kernel_correlation_summary.attribution_quality=poor` on a ~6 s profile-epoch kineto slice is often a windowing artifact** (busy_frac <0.5 in every 1 s bucket); trust TraceLens `ops_summary.csv` (`aten::mm` + `aten::bmm` ≈80% kernel time) before escalating to `OMNISTAT_KERNEL_TRACE=1`. See [`recipes/perf-optimizer-loop/dispatch-attribution.md`](../../../material_science/models/HydraGNN/recipes/perf-optimizer-loop/dispatch-attribution.md). + The full list (with cross-cutting context) is at the end of [.cursor/skills/ai4science-studio/SKILL.md](../ai4science-studio/SKILL.md). When fixing a bug here, check there first and propagate the lesson. diff --git a/material_science/models/HydraGNN/examples/run_fom_extractor.py b/material_science/models/HydraGNN/examples/run_fom_extractor.py new file mode 100755 index 0000000..47630ed --- /dev/null +++ b/material_science/models/HydraGNN/examples/run_fom_extractor.py @@ -0,0 +1,427 @@ +#!/usr/bin/env python3 +"""run_fom_extractor.py — Compute FOMs + TraceLens↔Omnistat kernel_correlation.csv. + +Implements material_science/models/HydraGNN/recipes/perf-optimizer-loop/agents/fom_extractor.md +for login-node post-processing of a completed perf run. + +Usage: + export AI4S_SHARED_DIR=/shared/aaji + export REPO_ROOT=/home/aaji/git/ai4science-studio + python run_fom_extractor.py --manifest /path/to/perf-runs//manifest.json \\ + [--vm-port 8432] [--tsdb-url http://127.0.0.1:8432] +""" + +from __future__ import annotations + +import argparse +import calendar +import csv +import gzip +import json +import os +import re +import statistics +import subprocess +import sys +import time +from collections import Counter, defaultdict +from datetime import datetime, timezone +from pathlib import Path +from typing import Any + +import requests + +MFMA_Q = ( + 'rate(omnistat_hardware_counter{{name="SQ_INSTS_VALU_MFMA_MOPS_F64"}}[10s]) ' + '* on(instance) group_left() max by(instance)(rmsjob_info{{jobid="{jobid}"}}) * 4 / 1e12' +) +HBM_Q = ( + 'rate(omnistat_hardware_counter{{name="FETCH_SIZE"}}[10s]) ' + '* on(instance) group_left() max by(instance)(rmsjob_info{{jobid="{jobid}"}}) * 1024 / 1e9' +) +PWR_Q = ( + 'rocm_average_socket_power_watts ' + '* on(instance) group_left() max by(instance)(rmsjob_info{{jobid="{jobid}"}})' +) +JOB_MFMA_Q = ( + 'avg(rate(omnistat_hardware_counter{{name="SQ_INSTS_VALU_MFMA_MOPS_F64"}}[10s]) ' + '* on(instance) group_left() max by(instance)(rmsjob_info{{jobid="{jobid}"}})) * 4 / 1e12' +) +JOB_PWR_Q = ( + 'avg_over_time((rocm_average_socket_power_watts ' + '* on(instance) group_left() max by(instance)(rmsjob_info{{jobid="{jobid}"}}))' + '[{runtime}s:])' +) + + +def _open_trace(path: str): + if path.endswith(".gz"): + return gzip.open(path, "rt", encoding="utf-8") + return open(path, encoding="utf-8") + + +def trace_anchor_epoch_us(trace_path: str, trace: dict) -> int: + """Return unix-epoch microseconds for kineto ts=0.""" + base_ns = trace.get("baseTimeNanoseconds") + if base_ns and int(base_ns) > 1_700_000_000_000_000_000: + return int(base_ns) // 1000 + m = re.search(r"\.(\d+)\.pt\.trace\.json", trace_path) + if m: + return int(m.group(1)) // 1000 + return 0 + + +def load_kernels(trace_path: str) -> tuple[list[tuple[int, int, str]], int]: + with _open_trace(trace_path) as f: + data = json.load(f) + anchor_us = trace_anchor_epoch_us(trace_path, data) + kernels = [ + (int(e["ts"]), int(e["dur"]), e["name"]) + for e in data.get("traceEvents", []) + if e.get("ph") == "X" and e.get("cat") == "kernel" and "ts" in e and "dur" in e + ] + return kernels, anchor_us + + +def bucket_windows( + kernels: list[tuple[int, int, str]], anchor_us: int +) -> dict[int, Counter]: + windows: dict[int, Counter] = defaultdict(Counter) + for ts, dur, name in kernels: + win = (anchor_us + ts) // 1_000_000 + windows[win][name] += dur + return windows + + +def window_top(counter: Counter) -> tuple[str, float, str, float]: + if not counter: + return "", 0.0, "", 0.0 + top2 = counter.most_common(2) + top_name, top_dur = top2[0] + top_frac = min(top_dur / 1_000_000.0, 1.0) + if len(top2) > 1: + second_name, second_dur = top2[1] + second_frac = min(second_dur / 1_000_000.0, 1.0) + else: + second_name, second_frac = "", 0.0 + return top_name, top_frac, second_name, second_frac + + +def promql_vector(tsdb_url: str, promql: str, t: int) -> dict[str, float]: + r = requests.get( + f"{tsdb_url}/api/v1/query", + params={"query": promql, "time": t}, + timeout=30, + ) + r.raise_for_status() + out: dict[str, float] = {} + for row in r.json()["data"]["result"]: + inst = row["metric"].get("instance", "") + out[inst] = float(row["value"][1]) + return out + + +def promql_scalar(tsdb_url: str, promql: str, t: int) -> float | None: + r = requests.get( + f"{tsdb_url}/api/v1/query", + params={"query": promql, "time": t}, + timeout=30, + ) + r.raise_for_status() + res = r.json()["data"]["result"] + if not res: + return None + return float(res[0]["value"][1]) + + +def start_vm(db_path: str, port: int, log_path: Path) -> None: + vm = os.environ.get( + "VICTORIA_METRICS_BIN", + f"{os.environ['AI4S_SHARED_DIR']}/tools/victoriametrics/victoria-metrics-prod", + ) + subprocess.Popen( + [ + vm, + f"-storageDataPath={db_path}", + f"-httpListenAddr=127.0.0.1:{port}", + "-retentionPeriod=100y", + "-search.disableCache", + "-search.latencyOffset=0", + "-search.maxPointsPerTimeseries=90000", + "-fs.disableMmap", + ], + stdout=log_path.open("w"), + stderr=subprocess.STDOUT, + ) + for _ in range(20): + time.sleep(1) + try: + requests.get(f"http://127.0.0.1:{port}/api/v1/status/tsdb", timeout=2).raise_for_status() + return + except requests.RequestException: + continue + raise RuntimeError(f"VictoriaMetrics did not become ready on port {port}") + + +def attribution_quality(windows: dict[int, Counter]) -> str: + if not windows: + return "poor" + busy = 0 + for win in windows: + _, frac, _, _ = window_top(windows[win]) + if frac > 0.5: + busy += 1 + ratio = busy / len(windows) + if ratio > 0.5: + return "good" + if ratio >= 0.25: + return "fair" + return "poor" + + +def short_kernel_family(name: str) -> str: + if "nccl" in name.lower(): + return "nccl" + if "aten::bmm" in name or "bmm" in name.lower()[:20]: + return "aten::bmm" + if "aten::mm" in name or name.startswith("Cijk_"): + return "aten::mm" + if "triton" in name.lower(): + return "triton" + return name[:80] + + +def main() -> int: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--manifest", required=True, help="perf-run manifest.json") + parser.add_argument("--tsdb-url", default=None, help="VictoriaMetrics base URL") + parser.add_argument("--vm-port", type=int, default=8432, help="Port if starting VM") + parser.add_argument("--no-start-vm", action="store_true", help="Require --tsdb-url") + args = parser.parse_args() + + shared = os.environ.get("AI4S_SHARED_DIR", "/shared/aaji") + repo = os.environ.get("REPO_ROOT", str(Path(__file__).resolve().parents[3])) + + with open(args.manifest) as f: + manifest = json.load(f) + + perf_run = Path(manifest["perf_run_dir"]) + jobid = str(manifest["jobid"]) + runtime_s = float(manifest.get("runtime_seconds", manifest.get("runtime_s", 300))) + trace_paths = manifest.get("trace_paths") or [] + if not trace_paths and manifest.get("trace_json"): + trace_paths = [manifest["trace_json"]] + trace_path = trace_paths[0] if trace_paths else "" + + log_path = manifest.get("training_log_path") or str(perf_run / f"hydragnn-train-{jobid}.out") + if not os.path.isfile(log_path): + alt = perf_run / "logs" / f"hydragnn-train-{jobid}-N2" / "run.log" + if alt.is_file(): + log_path = str(alt) + + tsdb_url = args.tsdb_url + if not tsdb_url: + if args.no_start_vm: + print("ERROR: --tsdb-url or VM start required", file=sys.stderr) + return 1 + port = args.vm_port + db = manifest.get("omnistat_db_path", str(perf_run / "omnistat-db")) + start_vm(db, port, perf_run / "vm_for_foms.log") + tsdb_url = f"http://127.0.0.1:{port}" + + # Time-based FOMs from log when available + conv_foms: dict[str, Any] = {} + parse_script = Path(repo) / "material_science/models/HydraGNN/examples/parse_convergence.py" + conv_json = perf_run / "_convergence_foms.json" + if os.path.isfile(log_path) and parse_script.is_file(): + subprocess.run( + [ + sys.executable, + str(parse_script), + "--log", + log_path, + "--output", + str(perf_run / "convergence.csv"), + "--json-foms", + str(conv_json), + ], + check=False, + ) + if conv_json.is_file(): + with open(conv_json) as f: + conv_foms = json.load(f) + + epoch_time_s = conv_foms.get("mean_epoch_time_excluding_epoch_0") + ranks = int(manifest.get("ranks", 16)) + batch = int(manifest.get("hg_batch_size", 400)) + max_batch = int(manifest.get("hydragnn_max_num_batch", 50)) + num_epochs = int(conv_foms.get("num_epochs_completed") or manifest.get("hg_num_epoch", 3)) + samples = num_epochs * max_batch * batch * ranks + throughput = (samples / runtime_s) if runtime_s else None + + # Query counters near job mid-window (not profile-trace mid — trace is only epoch-2). + t_query = int(time.time()) + db_meta = Path(manifest.get("omnistat_db_path", perf_run / "omnistat-db")) / "data" / "small" + for meta_file in db_meta.rglob("metadata.json"): + try: + meta = json.loads(meta_file.read_text()) + t_query = int((meta["MinTimestamp"] + meta["MaxTimestamp"]) / 2000) + break + except (KeyError, json.JSONDecodeError, OSError): + continue + + mfma_tflops = promql_scalar(tsdb_url, JOB_MFMA_Q.format(jobid=jobid), t_query) + mean_power_W = promql_scalar(tsdb_url, JOB_PWR_Q.format(jobid=jobid, runtime=int(runtime_s)), t_query) + energy_J = (mean_power_W * runtime_s) if mean_power_W is not None else None + energy_per_sample = (energy_J / samples) if energy_J and samples else None + + windows: dict[int, Counter] = {} + correlation_rows: list[dict[str, Any]] = [] + if trace_path and os.path.isfile(trace_path): + size_gb = os.path.getsize(trace_path) / (1024**3) + if size_gb > 4: + attr_q = "skipped:trace_too_large" + else: + kernels, anchor_us = load_kernels(trace_path) + windows = bucket_windows(kernels, anchor_us) + attr_q = attribution_quality(windows) + for win_ts in sorted(windows): + top_name, top_frac, second_name, second_frac = window_top(windows[win_ts]) + mfma_by_inst = promql_vector(tsdb_url, MFMA_Q.format(jobid=jobid), win_ts) + hbm_by_inst = promql_vector(tsdb_url, HBM_Q.format(jobid=jobid), win_ts) + pwr_by_inst = promql_vector(tsdb_url, PWR_Q.format(jobid=jobid), win_ts) + instances = sorted(set(mfma_by_inst) | set(hbm_by_inst) | set(pwr_by_inst)) + if not instances: + iso = datetime.fromtimestamp(win_ts, tz=timezone.utc).strftime("%Y-%m-%dT%H:%M:%S") + correlation_rows.append( + { + "window_start_iso": iso, + "window_start_ts": win_ts, + "instance": "lux-mi355x-b1", + "gpu_id": 0, + "top_kernel_name": top_name, + "top_kernel_busy_frac": f"{top_frac:.4f}", + "second_kernel_name": second_name, + "second_kernel_busy_frac": f"{second_frac:.4f}", + "mfma_tflops": "", + "hbm_read_GBps": "", + "mean_power_W": "", + } + ) + for inst in instances: + iso = datetime.fromtimestamp(win_ts, tz=timezone.utc).strftime("%Y-%m-%dT%H:%M:%S") + correlation_rows.append( + { + "window_start_iso": iso, + "window_start_ts": win_ts, + "instance": inst, + "gpu_id": 0, + "top_kernel_name": top_name, + "top_kernel_busy_frac": f"{top_frac:.4f}", + "second_kernel_name": second_name, + "second_kernel_busy_frac": f"{second_frac:.4f}", + "mfma_tflops": mfma_by_inst.get(inst, ""), + "hbm_read_GBps": hbm_by_inst.get(inst, ""), + "mean_power_W": pwr_by_inst.get(inst, ""), + } + ) + else: + attr_q = "poor" + + top_dominants = Counter(short_kernel_family(window_top(windows[w])[0]) for w in windows).most_common(5) + busy_gt_half = sum(1 for w in windows if window_top(windows[w])[1] > 0.5) + + mm_mfma: list[float] = [] + bmm_mfma: list[float] = [] + for row in correlation_rows: + name = row["top_kernel_name"] + mfma = row.get("mfma_tflops") + if mfma == "" or mfma is None: + continue + val = float(mfma) + fam = short_kernel_family(name) + if fam == "aten::mm": + mm_mfma.append(val) + elif fam == "aten::bmm" or "bmm" in name.lower(): + bmm_mfma.append(val) + + existing: dict[str, Any] = {} + foms_path = perf_run / "foms.json" + if foms_path.is_file(): + with open(foms_path) as f: + existing = json.load(f) + + prev_foms = existing.get("foms") or {} + + def _keep_counter(new_val: float | None, old_val: Any, min_sane: float) -> Any: + if new_val is None: + return old_val + if old_val is not None and float(new_val) < min_sane and float(old_val) >= min_sane: + return old_val + return new_val + + foms_doc = { + **existing, + "jobid": jobid, + "runtime_s": runtime_s, + "num_epochs_completed": num_epochs, + "samples_processed": samples, + "foms": { + **prev_foms, + "epoch_time_s": epoch_time_s or prev_foms.get("epoch_time_s"), + "throughput_samples_per_s": throughput or prev_foms.get("throughput_samples_per_s"), + "mfma_tflops": _keep_counter(mfma_tflops, prev_foms.get("mfma_tflops"), 0.5), + "energy_J": _keep_counter(energy_J, prev_foms.get("energy_J"), 1e5), + "mean_power_W": _keep_counter(mean_power_W, prev_foms.get("mean_power_W"), 500.0), + "energy_per_sample_J": energy_per_sample or prev_foms.get("energy_per_sample_J"), + "final_loss": conv_foms.get("final_train_loss") or prev_foms.get("final_loss"), + }, + "kernel_correlation_summary": { + "windows_examined": len(windows), + "windows_with_top_kernel_busy_frac_gt_0p5": busy_gt_half, + "top_kernel_dominants": top_dominants, + "median_mfma_tflops_when_aten_mm_dominant": statistics.median(mm_mfma) if mm_mfma else None, + "median_mfma_tflops_when_aten_bmm_dominant": statistics.median(bmm_mfma) if bmm_mfma else None, + "median_hbm_read_GBps_when_mul_dominant": None, + "attribution_quality": attr_q, + "trace_path": trace_path, + "omnistat_join": "ok" if correlation_rows and correlation_rows[0].get("mfma_tflops") != "" else "trace_only_or_partial", + "tracelens_report": str(perf_run / "tracelens" / "report.xlsx"), + }, + } + + csv_path = perf_run / "kernel_correlation.csv" + fields = [ + "window_start_iso", + "window_start_ts", + "instance", + "gpu_id", + "top_kernel_name", + "top_kernel_busy_frac", + "second_kernel_name", + "second_kernel_busy_frac", + "mfma_tflops", + "hbm_read_GBps", + "mean_power_W", + ] + tmp_csv = csv_path.with_suffix(".csv.tmp") + with open(tmp_csv, "w", newline="") as f: + w = csv.DictWriter(f, fieldnames=fields) + w.writeheader() + w.writerows(correlation_rows) + tmp_csv.rename(csv_path) + + tmp_foms = foms_path.with_suffix(".json.tmp") + with open(tmp_foms, "w") as f: + json.dump(foms_doc, f, indent=2) + tmp_foms.rename(foms_path) + + print( + f"STATUS=ok; reason=foms epoch_time={foms_doc['foms'].get('epoch_time_s')}s " + f"correlation={attr_q} windows={len(windows)}" + ) + return 0 + + +if __name__ == "__main__": + sys.exit(main()) diff --git a/material_science/models/HydraGNN/recipes/perf-optimizer-loop/dispatch-attribution.md b/material_science/models/HydraGNN/recipes/perf-optimizer-loop/dispatch-attribution.md new file mode 100644 index 0000000..e678651 --- /dev/null +++ b/material_science/models/HydraGNN/recipes/perf-optimizer-loop/dispatch-attribution.md @@ -0,0 +1,78 @@ +# Dispatch attribution — job 7187 (current best, 75 s/epoch) + +**Perf run:** `/shared/aaji/models/HydraGNN/perf-runs/7187/` +**Lever:** `rccl_high_priority` (`TORCH_NCCL_HIGH_PRIORITY=1`, `GPU_MAX_HW_QUEUES=2`) +**Nodes:** `lux-mi355x-b1`, `lux-mi355x-b2` (2×8 GPUs) + +## Artifacts produced (2026-05-27) + +| Artifact | Path | +|---|---| +| Manifest (backfilled) | `perf-runs/7187/manifest.json` | +| TraceLens workbook | `perf-runs/7187/tracelens/report.xlsx` | +| TraceLens digest | `perf-runs/7187/tracelens/report_summary.md` | +| FOMs + correlation summary | `perf-runs/7187/foms.json` | +| 1s window join | `perf-runs/7187/kernel_correlation.csv` | + +## `kernel_correlation_summary.attribution_quality` + +**Value: `poor`** (automated rule in `run_fom_extractor.py`) + +| Metric | Value | +|---|---| +| `windows_examined` | 4 | +| `windows_with_top_kernel_busy_frac_gt_0p5` | 0 | +| `top_kernel_dominants` | `aten::mm` (3 windows), elementwise (1) | + +The 1s-window heuristic is **misleading for this run**: rank-0 kineto captures only **~6 s of epoch-2 profiling** (~0.88 s summed kernel time), so `top_kernel_busy_frac` never exceeds ~0.18 in any 1 s wall-clock bucket. That does **not** mean dispatch offenders are unknown. + +## TraceLens attribution (authoritative for this run) + +From the profile-epoch trace (`tracelens/csvs/ops_summary.csv`): + +| Parent op | GPU kernel time | Share | +|---|---:|---:| +| `aten::mm` | 371.7 ms | **47.1%** | +| `aten::bmm` | 259.7 ms | **32.9%** | +| `aten::mul` + elementwise + reduce | remainder | ~20% | + +**NCCL** (`record_param_comms` → `ncclDevKernel_Generic_2`) is **~10.6%** of kernel time in the profile slice — non-trivial but secondary to GEMM/BMM launch volume. + +Short-kernel study (`short_kernels_summary.csv`): thousands of `aten::fill_`, `aten::copy_`, `aten::mm`/`aten::bmm` dispatches with mean duration **<10 µs** — classic **dispatch-bound** signature, aligned with lesson #13 (fp32≈fp64 epoch time). + +## Omnistat join (profile windows) + +`kernel_correlation.csv` joins rank-0 trace windows to per-node MFMA/HBM/power. In profile windows MFMA ≈ **0.012–0.014 TFLOP/s** per GPU — near idle vs the job-average **~17 TFLOP/s** in `foms.json` (steady-state epochs 0–1). Low MFMA during the profile slice is expected (profiler overhead + short window). + +## Decision (per perf plan) + +| Gate | Result | +|---|---| +| TraceLens names top dispatch families? | **Yes** — `aten::mm`, `aten::bmm`, many short kernels | +| `attribution_quality` automated | **poor** (short trace window artifact) | +| Escalate to `OMNISTAT_KERNEL_TRACE=1`? | **No** — TraceLens/kineto already identifies the culprit families; kernel trace would not change the no-upstream playbook | + +**Next steps (no upstream HydraGNN):** 8/16-node scaling sweep; do **not** re-run compile/TunableOp/batch levers. + +## Commands to reproduce + +```bash +export AI4S_SHARED_DIR=/shared/aaji +export REPO_ROOT=$HOME/git/ai4science-studio +source $AI4S_SHARED_DIR/tools/omnistat-pr271/bin/activate + +# VictoriaMetrics on the perf-run DB (login node) +$AI4S_SHARED_DIR/tools/victoriametrics/victoria-metrics-prod \ + -storageDataPath=$AI4S_SHARED_DIR/models/HydraGNN/perf-runs/7187/omnistat-db \ + -httpListenAddr=127.0.0.1:8432 -retentionPeriod=100y -fs.disableMmap & + +cd $AI4S_SHARED_DIR/models/HydraGNN/perf-runs/7187/tracelens +TraceLens_generate_perf_report_pytorch \ + --profile_json_path ../logs/hydragnn-train-7187-N2/lux-mi355x-b1_532131.1779479818046951705.pt.trace.json \ + --output_xlsx_path report.xlsx --output_csvs_dir csvs/ \ + --enable_kernel_summary --short_kernel_study + +python3 $REPO_ROOT/material_science/models/HydraGNN/examples/run_fom_extractor.py \ + --manifest $AI4S_SHARED_DIR/models/HydraGNN/perf-runs/7187/manifest.json \ + --tsdb-url http://127.0.0.1:8432 --no-start-vm +``` From 95d4d05c913568e9ff6bac2365aa48a6c1ab4595 Mon Sep 17 00:00:00 2001 From: Ashwin Aji Date: Fri, 29 May 2026 07:23:39 +0000 Subject: [PATCH 15/31] HydraGNN perf: move omnistat-inspect venv to /shared/omnihub/tools/omnihub-inspect The PR #271 omnistat + TraceLens venv was relocated out of personal space (/shared/aaji/tools/omnistat-pr271) into the shared OmniHub tools tree as omnihub-inspect so the legacy deep-dive workflow is reproducible for any user. Update launcher/builder and all perf-analysis + perf-optimizer-loop references. Co-authored-by: Cursor --- .../models/HydraGNN/examples/run_optimizer_loop.sh | 2 +- .../models/HydraGNN/examples/sbatch_train_perf_amd.sh | 4 ++-- .../models/HydraGNN/recipes/perf-analysis/agents/launcher.md | 4 ++-- .../HydraGNN/recipes/perf-analysis/agents/omnistat_analyst.md | 4 ++-- .../recipes/perf-analysis/agents/tracelens_analyst.md | 2 +- .../recipes/perf-optimizer-loop/agents/fom_extractor.md | 2 +- .../recipes/perf-optimizer-loop/agents/orchestrator.md | 2 +- .../recipes/perf-optimizer-loop/agents/story_writer.md | 2 +- .../recipes/perf-optimizer-loop/dispatch-attribution.md | 2 +- 9 files changed, 12 insertions(+), 12 deletions(-) diff --git a/material_science/models/HydraGNN/examples/run_optimizer_loop.sh b/material_science/models/HydraGNN/examples/run_optimizer_loop.sh index 096aea8..aa55c01 100755 --- a/material_science/models/HydraGNN/examples/run_optimizer_loop.sh +++ b/material_science/models/HydraGNN/examples/run_optimizer_loop.sh @@ -93,7 +93,7 @@ fi # 3. tools present for _bin in \ - "${AI4S_SHARED_DIR}/tools/omnistat-pr271/bin/omnistat-usermode" \ + "/shared/omnihub/tools/omnihub-inspect/bin/omnistat-usermode" \ "${AI4S_SHARED_DIR}/tools/victoriametrics/victoria-metrics-prod"; do if [[ -z "$PREFLIGHT_FAIL_REASON" && ! -x "$_bin" ]]; then PREFLIGHT_FAIL_REASON="tool_missing" diff --git a/material_science/models/HydraGNN/examples/sbatch_train_perf_amd.sh b/material_science/models/HydraGNN/examples/sbatch_train_perf_amd.sh index 7780818..cc7a43f 100755 --- a/material_science/models/HydraGNN/examples/sbatch_train_perf_amd.sh +++ b/material_science/models/HydraGNN/examples/sbatch_train_perf_amd.sh @@ -20,7 +20,7 @@ # # Required: # AI4S_SHARED_DIR — shared base path -# $AI4S_SHARED_DIR/tools/omnistat-pr271/ (created by the launcher subagent) +# /shared/omnihub/tools/omnihub-inspect/ (created by the launcher subagent) # $AI4S_SHARED_DIR/tools/victoriametrics/ (created by the launcher subagent) # # Optional env-var overrides (with defaults): @@ -70,7 +70,7 @@ HG_REPO_DIR="${HG_REPO_DIR:-${HG_BASE}/code/HydraGNN}" HYDRAGNN_MAX_NUM_BATCH="${HYDRAGNN_MAX_NUM_BATCH:-30}" PROFILE_TARGET_EPOCH="${PROFILE_TARGET_EPOCH:-1}" -OMNISTAT_VENV="${OMNISTAT_VENV:-${AI4S_SHARED_DIR}/tools/omnistat-pr271}" +OMNISTAT_VENV="${OMNISTAT_VENV:-/shared/omnihub/tools/omnihub-inspect}" # Read the omnistat config from the in-repo recipe (single source of truth). # Override with OMNISTAT_TEMPLATE=/path/to/file if you need a custom probe config. # A staged copy under ${HG_BASE}/perf-runs/ is intentionally NOT used as the diff --git a/material_science/models/HydraGNN/recipes/perf-analysis/agents/launcher.md b/material_science/models/HydraGNN/recipes/perf-analysis/agents/launcher.md index 65b66bb..e940059 100644 --- a/material_science/models/HydraGNN/recipes/perf-analysis/agents/launcher.md +++ b/material_science/models/HydraGNN/recipes/perf-analysis/agents/launcher.md @@ -10,7 +10,7 @@ Submits a 2-node HydraGNN training on Lux with PyTorch profiling and Omnistat us ## Outputs -- `/shared/aaji/tools/omnistat-pr271/` — Python venv with omnistat (PR #271 + main merged) and TraceLens. +- `/shared/omnihub/tools/omnihub-inspect/` — Python venv with omnistat (PR #271 + main merged) and TraceLens. - `/shared/aaji/tools/victoriametrics/victoria-metrics-prod` — VictoriaMetrics binary. - `/shared/aaji/models/HydraGNN/perf-runs/omnistat-lux.config` — Omnistat user-mode config. - `/shared/aaji/models/HydraGNN/perf-runs//manifest.json` — manifest schema below. @@ -76,7 +76,7 @@ fi OMNISTAT_COMMIT=$(git rev-parse HEAD) # 1b. Venv (py3.12 to match cluster python) -VENV=$TOOLS/omnistat-pr271 +VENV=/shared/omnihub/tools/omnihub-inspect if [[ ! -d $VENV ]]; then python3 -m venv "$VENV" fi diff --git a/material_science/models/HydraGNN/recipes/perf-analysis/agents/omnistat_analyst.md b/material_science/models/HydraGNN/recipes/perf-analysis/agents/omnistat_analyst.md index c34f6e4..823a09e 100644 --- a/material_science/models/HydraGNN/recipes/perf-analysis/agents/omnistat_analyst.md +++ b/material_science/models/HydraGNN/recipes/perf-analysis/agents/omnistat_analyst.md @@ -6,7 +6,7 @@ Drive `omnistat-inspect` (PR #271) through the analyze-job phases on the user-mo - `/manifest.json` - The Omnistat DB at `manifest.omnistat_db_path`. -- The `omnistat-pr271` venv at `/shared/aaji/tools/omnistat-pr271/`. +- The `omnihub-inspect` venv at `/shared/omnihub/tools/omnihub-inspect/`. - The VictoriaMetrics binary at `/shared/aaji/tools/victoriametrics/victoria-metrics-prod`. ## Outputs @@ -62,7 +62,7 @@ done ### 2. Run analyze-job phases (per the SKILL) ```bash -. /shared/aaji/tools/omnistat-pr271/bin/activate +. /shared/omnihub/tools/omnihub-inspect/bin/activate SCRATCH=$PERF_RUN/omnistat/scratch mkdir -p "$SCRATCH" JOBID=$(jq -r .jobid ) diff --git a/material_science/models/HydraGNN/recipes/perf-analysis/agents/tracelens_analyst.md b/material_science/models/HydraGNN/recipes/perf-analysis/agents/tracelens_analyst.md index b33572d..dd9ead3 100644 --- a/material_science/models/HydraGNN/recipes/perf-analysis/agents/tracelens_analyst.md +++ b/material_science/models/HydraGNN/recipes/perf-analysis/agents/tracelens_analyst.md @@ -39,7 +39,7 @@ Every entry: ### 1. Sanity check ```bash -. /shared/aaji/tools/omnistat-pr271/bin/activate +. /shared/omnihub/tools/omnihub-inspect/bin/activate PERF_RUN=$(jq -r .perf_run_dir ) TRACE=$(jq -r '.trace_paths[0] // empty' ) [[ -z "$TRACE" ]] && { echo "STATUS=partial; reason=no_trace_in_manifest"; exit 0; } diff --git a/material_science/models/HydraGNN/recipes/perf-optimizer-loop/agents/fom_extractor.md b/material_science/models/HydraGNN/recipes/perf-optimizer-loop/agents/fom_extractor.md index 06bf828..a753ec3 100644 --- a/material_science/models/HydraGNN/recipes/perf-optimizer-loop/agents/fom_extractor.md +++ b/material_science/models/HydraGNN/recipes/perf-optimizer-loop/agents/fom_extractor.md @@ -68,7 +68,7 @@ omnistat kernel-trace collector). ```bash set -euo pipefail -. "${AI4S_SHARED_DIR}/tools/omnistat-pr271/bin/activate" +. "/shared/omnihub/tools/omnihub-inspect/bin/activate" PERF_RUN=$(jq -r .perf_run_dir ) JOBID=$(jq -r .jobid ) RUNTIME=$(jq -r .runtime_seconds ) diff --git a/material_science/models/HydraGNN/recipes/perf-optimizer-loop/agents/orchestrator.md b/material_science/models/HydraGNN/recipes/perf-optimizer-loop/agents/orchestrator.md index bc0f78a..493f1e4 100644 --- a/material_science/models/HydraGNN/recipes/perf-optimizer-loop/agents/orchestrator.md +++ b/material_science/models/HydraGNN/recipes/perf-optimizer-loop/agents/orchestrator.md @@ -87,7 +87,7 @@ Run these checks in sequence; abort with `PREFLIGHT_FAIL reason=` on the - `sinfo -p lux,rad -h` → if zero `IDLE` or `MIX` nodes, abort. Cluster may be in maintenance (see SKILL §14). - `df -h /shared | tail -1` → abort if `Use%` > 95. -- `test -x $AI4S_SHARED_DIR/tools/omnistat-pr271/bin/omnistat-usermode` → abort if missing. +- `test -x /shared/omnihub/tools/omnihub-inspect/bin/omnistat-usermode` → abort if missing. - `test -x $AI4S_SHARED_DIR/tools/victoriametrics/victoria-metrics-prod` → abort if missing. - `test -f $AI4S_SHARED_DIR/models/HydraGNN/overlays/hydragnn-overlay.img` → abort if missing. - `test -f $AI4S_SHARED_DIR/images/pytorch_rocm7.2.2_ubuntu24.04_py3.12_pytorch_release_2.10.0.sif` → abort if missing. diff --git a/material_science/models/HydraGNN/recipes/perf-optimizer-loop/agents/story_writer.md b/material_science/models/HydraGNN/recipes/perf-optimizer-loop/agents/story_writer.md index 2c4cad0..f63129c 100644 --- a/material_science/models/HydraGNN/recipes/perf-optimizer-loop/agents/story_writer.md +++ b/material_science/models/HydraGNN/recipes/perf-optimizer-loop/agents/story_writer.md @@ -156,7 +156,7 @@ fig.tight_layout() fig.savefig(out_png_path, dpi=150, bbox_inches="tight") ``` -Use `matplotlib.use("Agg")` for headless rendering. Do NOT use seaborn (extra dep not in the omnistat venv). Use the omnistat venv at `$AI4S_SHARED_DIR/tools/omnistat-pr271/bin/python` which already has matplotlib via TraceLens deps. +Use `matplotlib.use("Agg")` for headless rendering. Do NOT use seaborn (extra dep not in the omnistat venv). Use the omnistat venv at `/shared/omnihub/tools/omnihub-inspect/bin/python` which already has matplotlib via TraceLens deps. ## Style rules diff --git a/material_science/models/HydraGNN/recipes/perf-optimizer-loop/dispatch-attribution.md b/material_science/models/HydraGNN/recipes/perf-optimizer-loop/dispatch-attribution.md index e678651..6687b69 100644 --- a/material_science/models/HydraGNN/recipes/perf-optimizer-loop/dispatch-attribution.md +++ b/material_science/models/HydraGNN/recipes/perf-optimizer-loop/dispatch-attribution.md @@ -59,7 +59,7 @@ Short-kernel study (`short_kernels_summary.csv`): thousands of `aten::fill_`, `a ```bash export AI4S_SHARED_DIR=/shared/aaji export REPO_ROOT=$HOME/git/ai4science-studio -source $AI4S_SHARED_DIR/tools/omnistat-pr271/bin/activate +source /shared/omnihub/tools/omnihub-inspect/bin/activate # VictoriaMetrics on the perf-run DB (login node) $AI4S_SHARED_DIR/tools/victoriametrics/victoria-metrics-prod \ From ceabe2f1613b3196ba9db84eb76c070046b78ec9 Mon Sep 17 00:00:00 2001 From: Ashwin Aji Date: Fri, 29 May 2026 07:52:14 +0000 Subject: [PATCH 16/31] HydraGNN perf: consolidate omnistat-src + VictoriaMetrics under /shared/omnihub/tools Relocate the omnistat source clone (and its built libomnistat_trace.so) out of personal space into the shared OmniHub tools tree, alongside the already-shared omnihub-inspect venv and victoriametrics binary. The omnihub-inspect venv's editable install now resolves the omnistat package from the shared path, so the deep-dive tooling is fully self-contained. Point the perf-analysis launcher (TOOLS dir), trace-lib defaults, and kernel-trace build notes at the new paths. Co-authored-by: Cursor --- .cursor/skills/ai4science-studio/SKILL.md | 8 ++++---- .../models/HydraGNN/recipes/perf-analysis/README.md | 4 ++-- .../HydraGNN/recipes/perf-analysis/agents/launcher.md | 6 +++--- .../recipes/perf-analysis/agents/omnistat_analyst.md | 4 ++-- 4 files changed, 11 insertions(+), 11 deletions(-) diff --git a/.cursor/skills/ai4science-studio/SKILL.md b/.cursor/skills/ai4science-studio/SKILL.md index d9b203d..5796c27 100755 --- a/.cursor/skills/ai4science-studio/SKILL.md +++ b/.cursor/skills/ai4science-studio/SKILL.md @@ -681,12 +681,12 @@ Omnistat exposes **two** independent rocprofiler-sdk integrations, and there's a salloc -p lux -A vultr_lux -N 1 --time=00:15:00 --gpus-per-node=1 \ apptainer exec --rocm \ /shared/aaji/images/pytorch_rocm7.2.2_ubuntu24.04_py3.12_pytorch_release_2.10.0.sif \ - bash -c 'cd /shared/aaji/tools/omnistat-src && \ + bash -c 'cd /shared/omnihub/tools/omnistat-src && \ cmake -S rocprofiler-sdk/ -B build-trace/ -DBUILD_KERNEL_TRACE_LIB=ON && \ cmake --build build-trace/ -j 8' ``` -Verify with `ls -la /shared/aaji/tools/omnistat-src/build-trace/libomnistat_trace.so` (~MB-class file, not the tiny 4 kB stub you get when CMake fails halfway). +Verify with `ls -la /shared/omnihub/tools/omnistat-src/build-trace/libomnistat_trace.so` (~MB-class file, not the tiny 4 kB stub you get when CMake fails halfway). **Lux SIF build prerequisites (one-time gotchas):** the `pytorch_rocm7.2.2_ubuntu24.04_py3.12_pytorch_release_2.10.0.sif` ships **without** `cmake` and **without** libcurl headers (only the `.so.4` runtime). Both omnistat artifacts need cmake; the kernel-trace library additionally needs ``. Working build recipe inside the SIF: @@ -735,12 +735,12 @@ When the cluster is up but you only need a one-shot interactive build inside the srun -p lux -A vultr_lux -N 1 --time=00:15:00 --gpus-per-node=1 --cpus-per-task=8 \ --job-name= \ apptainer exec --rocm \ - --bind /shared/aaji/tools:/shared/aaji/tools \ + --bind /shared/omnihub/tools:/shared/omnihub/tools \ "$HG_SIF" \ bash -c '' ``` -Why: `salloc` doesn't auto-bind `/shared/aaji/tools/` into the apptainer namespace (only the `omnistat-src` symlink, which is bind-mounted system-wide, makes it through). Direct `srun … apptainer exec --bind …` forces the bind explicitly and writes the build artifact to its persistent NFS path in one shot. Verified during the kernel-trace library build (job 7030). +Why: `salloc` doesn't auto-bind `/shared/omnihub/tools/` into the apptainer namespace. Direct `srun … apptainer exec --bind …` forces the bind explicitly and writes the build artifact to its persistent NFS path in one shot. Verified during the kernel-trace library build (job 7030). ### 15. Dual-failure debug pattern — separate "tool wired correctly?" from "workload happy?" diff --git a/material_science/models/HydraGNN/recipes/perf-analysis/README.md b/material_science/models/HydraGNN/recipes/perf-analysis/README.md index cdea3ce..cfb0b8f 100644 --- a/material_science/models/HydraGNN/recipes/perf-analysis/README.md +++ b/material_science/models/HydraGNN/recipes/perf-analysis/README.md @@ -68,7 +68,7 @@ The `combined_report.md` is the deliverable: ranked bottlenecks, remedies tried ## Prerequisites - Lux access (`vultr_lux` account) with the standard HydraGNN setup at `/shared/aaji/models/HydraGNN/` — overlay, weights, code clone. See [HydraGNN/recipes/train/](../train/). -- One-time install of Omnistat (PR #271 branch `jorda/skills`, with `origin/main` merged in), VictoriaMetrics, and TraceLens — performed lazily by the launcher subagent into `/shared/aaji/tools/`. +- One-time install of Omnistat (PR #271 branch `jorda/skills`, with `origin/main` merged in), VictoriaMetrics, and TraceLens — performed lazily by the launcher subagent into `/shared/omnihub/tools/`. - The launcher writes `gfm_mlip_with_profile.json` at submit time (does not modify the upstream `gfm_mlip.json`). ## Telemetry knobs (set at `sbatch` submit time) @@ -76,7 +76,7 @@ The `combined_report.md` is the deliverable: ranked bottlenecks, remedies tried | Env var | Default | What it does | |---|---|---| | `OMNISTAT_KERNEL_TRACE` | `0` | When `1`, loads `libomnistat_trace.so` via `ROCP_TOOL_LIBRARIES` on every rank and turns on the omnistat kernel-trace collector. Adds per-kernel dispatch count + duration time series across all 8 cards on every node. Requires the library to be built once (see `agents/launcher.md` step 1e). Validated end-to-end on job 7034. | -| `OMNISTAT_TRACE_LIB` | `/shared/aaji/tools/omnistat-src/build-trace/libomnistat_trace.so` | Override only for development. The wrapper hard-fails if the path is missing when `OMNISTAT_KERNEL_TRACE=1`. | +| `OMNISTAT_TRACE_LIB` | `/shared/omnihub/tools/omnistat-src/build-trace/libomnistat_trace.so` | Override only for development. The wrapper hard-fails if the path is missing when `OMNISTAT_KERNEL_TRACE=1`. | | `OMNISTAT_TRACE_LOG` | `1` | Library prints `[host][pid][omnistat] Trace summary: N/N processed records (M/M successful flushes)` on rank exit; useful as a smoke-signal that the tool initialized even when the workload crashed. | `enable_rocprofiler=True` is **on** in `omnistat-lux.config.template` by default (since commit `a23e5c4`); the rendered config's state is printed in the sbatch banner so a silent "counters off" cannot recur. Kernel tracing and device-counting can co-exist on a single GPU but raise per-job VictoriaMetrics cardinality — pick by run length and the question you're asking. See `ai4science-studio` SKILL §12 for the device-counting vs kernel-trace vs `rocprofv3` decision matrix. diff --git a/material_science/models/HydraGNN/recipes/perf-analysis/agents/launcher.md b/material_science/models/HydraGNN/recipes/perf-analysis/agents/launcher.md index e940059..f1e9c6c 100644 --- a/material_science/models/HydraGNN/recipes/perf-analysis/agents/launcher.md +++ b/material_science/models/HydraGNN/recipes/perf-analysis/agents/launcher.md @@ -11,7 +11,7 @@ Submits a 2-node HydraGNN training on Lux with PyTorch profiling and Omnistat us ## Outputs - `/shared/omnihub/tools/omnihub-inspect/` — Python venv with omnistat (PR #271 + main merged) and TraceLens. -- `/shared/aaji/tools/victoriametrics/victoria-metrics-prod` — VictoriaMetrics binary. +- `/shared/omnihub/tools/victoriametrics/victoria-metrics-prod` — VictoriaMetrics binary. - `/shared/aaji/models/HydraGNN/perf-runs/omnistat-lux.config` — Omnistat user-mode config. - `/shared/aaji/models/HydraGNN/perf-runs//manifest.json` — manifest schema below. - `/shared/aaji/models/HydraGNN/perf-runs//hydragnn-train-.out` — symlinked from output dir. @@ -57,7 +57,7 @@ Submits a 2-node HydraGNN training on Lux with PyTorch profiling and Omnistat us ```bash set -euo pipefail -TOOLS=/shared/aaji/tools +TOOLS=/shared/omnihub/tools mkdir -p "$TOOLS" # 1a. Omnistat: jorda/skills branch with origin/main merged in @@ -122,7 +122,7 @@ if [[ "${OMNISTAT_KERNEL_TRACE:-0}" == "1" && ! -f "$TRACE_LIB" ]]; then fi ``` -The sbatch wrapper expects the `.so` at `/shared/aaji/tools/omnistat-src/build-trace/libomnistat_trace.so` by default; override with `OMNISTAT_TRACE_LIB=/path/to/lib.so`. See `ai4science-studio` SKILL §12 for when to enable kernel tracing vs the default device-counting collector. +The sbatch wrapper expects the `.so` at `/shared/omnihub/tools/omnistat-src/build-trace/libomnistat_trace.so` by default; override with `OMNISTAT_TRACE_LIB=/path/to/lib.so`. See `ai4science-studio` SKILL §12 for when to enable kernel tracing vs the default device-counting collector. ### 2. Author the Omnistat user-mode config (once, in repo `perf-runs/`) diff --git a/material_science/models/HydraGNN/recipes/perf-analysis/agents/omnistat_analyst.md b/material_science/models/HydraGNN/recipes/perf-analysis/agents/omnistat_analyst.md index 823a09e..a3c9bec 100644 --- a/material_science/models/HydraGNN/recipes/perf-analysis/agents/omnistat_analyst.md +++ b/material_science/models/HydraGNN/recipes/perf-analysis/agents/omnistat_analyst.md @@ -7,7 +7,7 @@ Drive `omnistat-inspect` (PR #271) through the analyze-job phases on the user-mo - `/manifest.json` - The Omnistat DB at `manifest.omnistat_db_path`. - The `omnihub-inspect` venv at `/shared/omnihub/tools/omnihub-inspect/`. -- The VictoriaMetrics binary at `/shared/aaji/tools/victoriametrics/victoria-metrics-prod`. +- The VictoriaMetrics binary at `/shared/omnihub/tools/victoriametrics/victoria-metrics-prod`. ## Outputs @@ -36,7 +36,7 @@ ss -lnt "sport = :$PORT" 2>/dev/null | grep -q LISTEN && PORT=8429 VMLOG=$PERF_RUN/omnistat/vm.log mkdir -p "$PERF_RUN/omnistat/inspect_outputs" -nohup /shared/aaji/tools/victoriametrics/victoria-metrics-prod \ +nohup /shared/omnihub/tools/victoriametrics/victoria-metrics-prod \ -storageDataPath="$DB" \ -httpListenAddr=127.0.0.1:$PORT \ -retentionPeriod=100y \ From f4d0c313ee62e132b01a1016aab7adadbe566532 Mon Sep 17 00:00:00 2001 From: Ashwin Aji Date: Fri, 29 May 2026 18:36:36 +0000 Subject: [PATCH 17/31] HydraGNN perf-optimizer-loop: point VM + trace-lib at /shared/omnihub/tools Finish the tool consolidation: the perf-optimizer-loop recipe (loop driver, fom extractor, agent docs) and the sbatch_train_perf_amd.sh trace-lib default still referenced VictoriaMetrics and omnistat-src via $AI4S_SHARED_DIR/tools (personal space). Repoint them at the shared /shared/omnihub/tools tree so all perf tooling resolves from one shared location regardless of user. Co-authored-by: Cursor --- .../models/HydraGNN/examples/run_fom_extractor.py | 2 +- .../models/HydraGNN/examples/run_optimizer_loop.sh | 2 +- .../models/HydraGNN/examples/sbatch_train_perf_amd.sh | 6 +++--- .../recipes/perf-optimizer-loop/agents/fom_extractor.md | 2 +- .../recipes/perf-optimizer-loop/agents/orchestrator.md | 2 +- .../recipes/perf-optimizer-loop/dispatch-attribution.md | 2 +- 6 files changed, 8 insertions(+), 8 deletions(-) diff --git a/material_science/models/HydraGNN/examples/run_fom_extractor.py b/material_science/models/HydraGNN/examples/run_fom_extractor.py index 47630ed..87378ab 100755 --- a/material_science/models/HydraGNN/examples/run_fom_extractor.py +++ b/material_science/models/HydraGNN/examples/run_fom_extractor.py @@ -137,7 +137,7 @@ def promql_scalar(tsdb_url: str, promql: str, t: int) -> float | None: def start_vm(db_path: str, port: int, log_path: Path) -> None: vm = os.environ.get( "VICTORIA_METRICS_BIN", - f"{os.environ['AI4S_SHARED_DIR']}/tools/victoriametrics/victoria-metrics-prod", + "/shared/omnihub/tools/victoriametrics/victoria-metrics-prod", ) subprocess.Popen( [ diff --git a/material_science/models/HydraGNN/examples/run_optimizer_loop.sh b/material_science/models/HydraGNN/examples/run_optimizer_loop.sh index aa55c01..ec94464 100755 --- a/material_science/models/HydraGNN/examples/run_optimizer_loop.sh +++ b/material_science/models/HydraGNN/examples/run_optimizer_loop.sh @@ -94,7 +94,7 @@ fi # 3. tools present for _bin in \ "/shared/omnihub/tools/omnihub-inspect/bin/omnistat-usermode" \ - "${AI4S_SHARED_DIR}/tools/victoriametrics/victoria-metrics-prod"; do + "/shared/omnihub/tools/victoriametrics/victoria-metrics-prod"; do if [[ -z "$PREFLIGHT_FAIL_REASON" && ! -x "$_bin" ]]; then PREFLIGHT_FAIL_REASON="tool_missing" PREFLIGHT_NOTES+=("missing: $_bin (run the perf-analysis launcher subagent first to install)") diff --git a/material_science/models/HydraGNN/examples/sbatch_train_perf_amd.sh b/material_science/models/HydraGNN/examples/sbatch_train_perf_amd.sh index cc7a43f..10b40f8 100755 --- a/material_science/models/HydraGNN/examples/sbatch_train_perf_amd.sh +++ b/material_science/models/HydraGNN/examples/sbatch_train_perf_amd.sh @@ -21,7 +21,7 @@ # Required: # AI4S_SHARED_DIR — shared base path # /shared/omnihub/tools/omnihub-inspect/ (created by the launcher subagent) -# $AI4S_SHARED_DIR/tools/victoriametrics/ (created by the launcher subagent) +# /shared/omnihub/tools/victoriametrics/ (created by the launcher subagent) # # Optional env-var overrides (with defaults): # HG_NUM_EPOCH=2 @@ -86,7 +86,7 @@ OMNISTAT_USERMODE_INTERVAL="${OMNISTAT_USERMODE_INTERVAL:-1}" # See material_science/models/HydraGNN/recipes/perf-analysis/agents/launcher.md # for how to build libomnistat_trace.so on a compute node. OMNISTAT_KERNEL_TRACE="${OMNISTAT_KERNEL_TRACE:-0}" -OMNISTAT_TRACE_LIB="${OMNISTAT_TRACE_LIB:-${AI4S_SHARED_DIR}/tools/omnistat-src/build-trace/libomnistat_trace.so}" +OMNISTAT_TRACE_LIB="${OMNISTAT_TRACE_LIB:-/shared/omnihub/tools/omnistat-src/build-trace/libomnistat_trace.so}" # IMPORTANT: the kernel-trace collector registers /kernel_trace on the SAME # Flask app as the other omnistat endpoints, so the tool library must POST to # the [omnistat.collectors] `port` (8101 here), NOT the library's own default @@ -174,7 +174,7 @@ if [[ "$OMNISTAT_KERNEL_TRACE" == "1" ]]; then echo " Build it on a compute node:" >&2 echo " salloc -p -A -N 1 --time=00:15:00 --gpus-per-node=1 \\" >&2 echo " apptainer exec --rocm \"\$HG_SIF\" bash -c '" >&2 - echo " cd \$AI4S_SHARED_DIR/tools/omnistat-src && \\" >&2 + echo " cd /shared/omnihub/tools/omnistat-src && \\" >&2 echo " cmake -S rocprofiler-sdk/ -B build-trace/ -DBUILD_KERNEL_TRACE_LIB=ON && \\" >&2 echo " cmake --build build-trace/ -j 8'" >&2 exit 2 diff --git a/material_science/models/HydraGNN/recipes/perf-optimizer-loop/agents/fom_extractor.md b/material_science/models/HydraGNN/recipes/perf-optimizer-loop/agents/fom_extractor.md index a753ec3..5d24b55 100644 --- a/material_science/models/HydraGNN/recipes/perf-optimizer-loop/agents/fom_extractor.md +++ b/material_science/models/HydraGNN/recipes/perf-optimizer-loop/agents/fom_extractor.md @@ -102,7 +102,7 @@ Start a local VictoriaMetrics if not already running: PORT=8428 ss -lnt "sport = :$PORT" 2>/dev/null | grep -q LISTEN && PORT=8429 if ! curl -sf "http://127.0.0.1:$PORT/api/v1/status/tsdb" > /dev/null 2>&1; then - nohup "${AI4S_SHARED_DIR}/tools/victoriametrics/victoria-metrics-prod" \ + nohup "/shared/omnihub/tools/victoriametrics/victoria-metrics-prod" \ -storageDataPath="$PERF_RUN/omnistat-db" \ -httpListenAddr=127.0.0.1:$PORT \ -retentionPeriod=100y -fs.disableMmap \ diff --git a/material_science/models/HydraGNN/recipes/perf-optimizer-loop/agents/orchestrator.md b/material_science/models/HydraGNN/recipes/perf-optimizer-loop/agents/orchestrator.md index 493f1e4..96f0309 100644 --- a/material_science/models/HydraGNN/recipes/perf-optimizer-loop/agents/orchestrator.md +++ b/material_science/models/HydraGNN/recipes/perf-optimizer-loop/agents/orchestrator.md @@ -88,7 +88,7 @@ Run these checks in sequence; abort with `PREFLIGHT_FAIL reason=` on the - `sinfo -p lux,rad -h` → if zero `IDLE` or `MIX` nodes, abort. Cluster may be in maintenance (see SKILL §14). - `df -h /shared | tail -1` → abort if `Use%` > 95. - `test -x /shared/omnihub/tools/omnihub-inspect/bin/omnistat-usermode` → abort if missing. -- `test -x $AI4S_SHARED_DIR/tools/victoriametrics/victoria-metrics-prod` → abort if missing. +- `test -x /shared/omnihub/tools/victoriametrics/victoria-metrics-prod` → abort if missing. - `test -f $AI4S_SHARED_DIR/models/HydraGNN/overlays/hydragnn-overlay.img` → abort if missing. - `test -f $AI4S_SHARED_DIR/images/pytorch_rocm7.2.2_ubuntu24.04_py3.12_pytorch_release_2.10.0.sif` → abort if missing. - If invoked via Claude Code CLI on the login node: `[[ -n "$ANTHROPIC_API_KEY" ]]` and `curl -fsS https://api.anthropic.com/v1/models > /dev/null` (5 s timeout) → abort if either fails. diff --git a/material_science/models/HydraGNN/recipes/perf-optimizer-loop/dispatch-attribution.md b/material_science/models/HydraGNN/recipes/perf-optimizer-loop/dispatch-attribution.md index 6687b69..94f162a 100644 --- a/material_science/models/HydraGNN/recipes/perf-optimizer-loop/dispatch-attribution.md +++ b/material_science/models/HydraGNN/recipes/perf-optimizer-loop/dispatch-attribution.md @@ -62,7 +62,7 @@ export REPO_ROOT=$HOME/git/ai4science-studio source /shared/omnihub/tools/omnihub-inspect/bin/activate # VictoriaMetrics on the perf-run DB (login node) -$AI4S_SHARED_DIR/tools/victoriametrics/victoria-metrics-prod \ +/shared/omnihub/tools/victoriametrics/victoria-metrics-prod \ -storageDataPath=$AI4S_SHARED_DIR/models/HydraGNN/perf-runs/7187/omnistat-db \ -httpListenAddr=127.0.0.1:8432 -retentionPeriod=100y -fs.disableMmap & From 884f1f71c86f4e156c60d15263b2960462d5c772 Mon Sep 17 00:00:00 2001 From: Ashwin Aji Date: Fri, 29 May 2026 18:47:33 +0000 Subject: [PATCH 18/31] HydraGNN perf: resolve shared tool paths via OMNIHUB_TOOLS_DIR env var Stop hardcoding the cluster-specific /shared/omnihub/tools path in the perf scripts and agent recipes. Introduce a required OMNIHUB_TOOLS_DIR env var (mirrors the AI4S_SHARED_DIR convention and the .cluster-config.yaml omnihub.tools_dir key); all omnistat venv / omnistat-src / VictoriaMetrics paths now derive from it. Scripts hard-fail with a clear message if it is unset, and the orchestrator/launcher/README document exporting it from the cluster config. Co-authored-by: Cursor --- .cursor/skills/ai4science-studio/SKILL.md | 12 ++++++------ .../models/HydraGNN/examples/run_fom_extractor.py | 5 ++++- .../models/HydraGNN/examples/run_optimizer_loop.sh | 7 +++++-- .../HydraGNN/examples/sbatch_train_perf_amd.sh | 13 ++++++++----- .../models/HydraGNN/recipes/perf-analysis/README.md | 5 +++-- .../recipes/perf-analysis/agents/launcher.md | 11 ++++++----- .../perf-analysis/agents/omnistat_analyst.md | 8 ++++---- .../perf-analysis/agents/tracelens_analyst.md | 2 +- .../perf-optimizer-loop/agents/fom_extractor.md | 4 ++-- .../perf-optimizer-loop/agents/orchestrator.md | 5 +++-- .../perf-optimizer-loop/agents/story_writer.md | 2 +- .../perf-optimizer-loop/dispatch-attribution.md | 4 ++-- 12 files changed, 45 insertions(+), 33 deletions(-) diff --git a/.cursor/skills/ai4science-studio/SKILL.md b/.cursor/skills/ai4science-studio/SKILL.md index 5796c27..f517112 100755 --- a/.cursor/skills/ai4science-studio/SKILL.md +++ b/.cursor/skills/ai4science-studio/SKILL.md @@ -675,18 +675,18 @@ Omnistat exposes **two** independent rocprofiler-sdk integrations, and there's a **Cross-runtime gotcha:** `enable_kernel_trace=True` without the tool library loaded does *nothing* — the omnistat collector just sits with an empty endpoint. The `OMNISTAT_KERNEL_TRACE=1` switch in `sbatch_train_perf_amd.sh` enforces this with a hard `exit 2` if `libomnistat_trace.so` isn't found at the configured path; do **not** loosen that check, otherwise you re-create the exact "silent zero counters" failure mode from §8. -**Build location:** Both omnistat artifacts (the Python extension and the kernel-trace tool library) must be built on a **compute node** inside the same SIF the workload runs in — login nodes don't have apptainer or the ROCm headers, and a host-side build will pick up the wrong libcurl / glibc. Standard build: +**Build location:** Both omnistat artifacts (the Python extension and the kernel-trace tool library) must be built on a **compute node** inside the same SIF the workload runs in — login nodes don't have apptainer or the ROCm headers, and a host-side build will pick up the wrong libcurl / glibc. `OMNIHUB_TOOLS_DIR` is the shared tools dir from `.cluster-config.yaml` `omnihub.tools_dir` (e.g. `/shared/omnihub/tools`); export it first. Standard build: ```bash salloc -p lux -A vultr_lux -N 1 --time=00:15:00 --gpus-per-node=1 \ apptainer exec --rocm \ /shared/aaji/images/pytorch_rocm7.2.2_ubuntu24.04_py3.12_pytorch_release_2.10.0.sif \ - bash -c 'cd /shared/omnihub/tools/omnistat-src && \ + bash -c "cd ${OMNIHUB_TOOLS_DIR}/omnistat-src && \ cmake -S rocprofiler-sdk/ -B build-trace/ -DBUILD_KERNEL_TRACE_LIB=ON && \ - cmake --build build-trace/ -j 8' + cmake --build build-trace/ -j 8" ``` -Verify with `ls -la /shared/omnihub/tools/omnistat-src/build-trace/libomnistat_trace.so` (~MB-class file, not the tiny 4 kB stub you get when CMake fails halfway). +Verify with `ls -la ${OMNIHUB_TOOLS_DIR}/omnistat-src/build-trace/libomnistat_trace.so` (~MB-class file, not the tiny 4 kB stub you get when CMake fails halfway). **Lux SIF build prerequisites (one-time gotchas):** the `pytorch_rocm7.2.2_ubuntu24.04_py3.12_pytorch_release_2.10.0.sif` ships **without** `cmake` and **without** libcurl headers (only the `.so.4` runtime). Both omnistat artifacts need cmake; the kernel-trace library additionally needs ``. Working build recipe inside the SIF: @@ -735,12 +735,12 @@ When the cluster is up but you only need a one-shot interactive build inside the srun -p lux -A vultr_lux -N 1 --time=00:15:00 --gpus-per-node=1 --cpus-per-task=8 \ --job-name= \ apptainer exec --rocm \ - --bind /shared/omnihub/tools:/shared/omnihub/tools \ + --bind ${OMNIHUB_TOOLS_DIR}:${OMNIHUB_TOOLS_DIR} \ "$HG_SIF" \ bash -c '' ``` -Why: `salloc` doesn't auto-bind `/shared/omnihub/tools/` into the apptainer namespace. Direct `srun … apptainer exec --bind …` forces the bind explicitly and writes the build artifact to its persistent NFS path in one shot. Verified during the kernel-trace library build (job 7030). +Why: `salloc` doesn't auto-bind `${OMNIHUB_TOOLS_DIR}/` into the apptainer namespace. Direct `srun … apptainer exec --bind …` forces the bind explicitly and writes the build artifact to its persistent NFS path in one shot. Verified during the kernel-trace library build (job 7030). ### 15. Dual-failure debug pattern — separate "tool wired correctly?" from "workload happy?" diff --git a/material_science/models/HydraGNN/examples/run_fom_extractor.py b/material_science/models/HydraGNN/examples/run_fom_extractor.py index 87378ab..e4111ef 100755 --- a/material_science/models/HydraGNN/examples/run_fom_extractor.py +++ b/material_science/models/HydraGNN/examples/run_fom_extractor.py @@ -6,9 +6,12 @@ Usage: export AI4S_SHARED_DIR=/shared/aaji + export OMNIHUB_TOOLS_DIR=/shared/omnihub/tools # from .cluster-config.yaml omnihub.tools_dir export REPO_ROOT=/home/aaji/git/ai4science-studio python run_fom_extractor.py --manifest /path/to/perf-runs//manifest.json \\ [--vm-port 8432] [--tsdb-url http://127.0.0.1:8432] + +Override the VictoriaMetrics binary directly with VICTORIA_METRICS_BIN if needed. """ from __future__ import annotations @@ -137,7 +140,7 @@ def promql_scalar(tsdb_url: str, promql: str, t: int) -> float | None: def start_vm(db_path: str, port: int, log_path: Path) -> None: vm = os.environ.get( "VICTORIA_METRICS_BIN", - "/shared/omnihub/tools/victoriametrics/victoria-metrics-prod", + f"{os.environ['OMNIHUB_TOOLS_DIR']}/victoriametrics/victoria-metrics-prod", ) subprocess.Popen( [ diff --git a/material_science/models/HydraGNN/examples/run_optimizer_loop.sh b/material_science/models/HydraGNN/examples/run_optimizer_loop.sh index ec94464..aa38d29 100755 --- a/material_science/models/HydraGNN/examples/run_optimizer_loop.sh +++ b/material_science/models/HydraGNN/examples/run_optimizer_loop.sh @@ -45,6 +45,9 @@ ORCH_PROMPT="${RECIPE_DIR}/agents/orchestrator.md" # Loop-dir convention matches recipes/perf-optimizer-loop/README.md : "${AI4S_SHARED_DIR:?AI4S_SHARED_DIR must be set (e.g. export AI4S_SHARED_DIR=/shared/\$USER)}" +# Shared perf-tool location (omnistat venv, VictoriaMetrics). Cluster-specific: +# set from .cluster-config.yaml omnihub.tools_dir (e.g. /shared/omnihub/tools). +: "${OMNIHUB_TOOLS_DIR:?OMNIHUB_TOOLS_DIR must be set (from .cluster-config.yaml omnihub.tools_dir, e.g. /shared/omnihub/tools)}" HG_BASE="${AI4S_SHARED_DIR}/models/HydraGNN" PERF_RUNS_DIR="${HG_BASE}/perf-runs" LOOP_DIR="${PERF_RUNS_DIR}/loop-${LOOP_UUID}" @@ -93,8 +96,8 @@ fi # 3. tools present for _bin in \ - "/shared/omnihub/tools/omnihub-inspect/bin/omnistat-usermode" \ - "/shared/omnihub/tools/victoriametrics/victoria-metrics-prod"; do + "${OMNIHUB_TOOLS_DIR}/omnihub-inspect/bin/omnistat-usermode" \ + "${OMNIHUB_TOOLS_DIR}/victoriametrics/victoria-metrics-prod"; do if [[ -z "$PREFLIGHT_FAIL_REASON" && ! -x "$_bin" ]]; then PREFLIGHT_FAIL_REASON="tool_missing" PREFLIGHT_NOTES+=("missing: $_bin (run the perf-analysis launcher subagent first to install)") diff --git a/material_science/models/HydraGNN/examples/sbatch_train_perf_amd.sh b/material_science/models/HydraGNN/examples/sbatch_train_perf_amd.sh index 10b40f8..df2c55f 100755 --- a/material_science/models/HydraGNN/examples/sbatch_train_perf_amd.sh +++ b/material_science/models/HydraGNN/examples/sbatch_train_perf_amd.sh @@ -20,8 +20,8 @@ # # Required: # AI4S_SHARED_DIR — shared base path -# /shared/omnihub/tools/omnihub-inspect/ (created by the launcher subagent) -# /shared/omnihub/tools/victoriametrics/ (created by the launcher subagent) +# ${OMNIHUB_TOOLS_DIR}/omnihub-inspect/ (created by the launcher subagent) +# ${OMNIHUB_TOOLS_DIR}/victoriametrics/ (created by the launcher subagent) # # Optional env-var overrides (with defaults): # HG_NUM_EPOCH=2 @@ -58,6 +58,9 @@ fi # Paths and configuration # --------------------------------------------------------------------------- HG_BASE="${AI4S_SHARED_DIR:?AI4S_SHARED_DIR must be set}/models/HydraGNN" +# Shared perf-tool location (omnistat venv, omnistat-src, VictoriaMetrics). +# Cluster-specific: set from .cluster-config.yaml omnihub.tools_dir. +: "${OMNIHUB_TOOLS_DIR:?OMNIHUB_TOOLS_DIR must be set (from .cluster-config.yaml omnihub.tools_dir, e.g. /shared/omnihub/tools)}" HG_SIF="${HG_SIF:-${AI4S_SHARED_DIR}/images/pytorch_rocm7.2.2_ubuntu24.04_py3.12_pytorch_release_2.10.0.sif}" HG_OVERLAY="${HG_OVERLAY:-${HG_BASE}/overlays/hydragnn-overlay.img}" HG_DATASETS="${HG_DATASETS:-ANI1x,Alexandria}" @@ -70,7 +73,7 @@ HG_REPO_DIR="${HG_REPO_DIR:-${HG_BASE}/code/HydraGNN}" HYDRAGNN_MAX_NUM_BATCH="${HYDRAGNN_MAX_NUM_BATCH:-30}" PROFILE_TARGET_EPOCH="${PROFILE_TARGET_EPOCH:-1}" -OMNISTAT_VENV="${OMNISTAT_VENV:-/shared/omnihub/tools/omnihub-inspect}" +OMNISTAT_VENV="${OMNISTAT_VENV:-${OMNIHUB_TOOLS_DIR}/omnihub-inspect}" # Read the omnistat config from the in-repo recipe (single source of truth). # Override with OMNISTAT_TEMPLATE=/path/to/file if you need a custom probe config. # A staged copy under ${HG_BASE}/perf-runs/ is intentionally NOT used as the @@ -86,7 +89,7 @@ OMNISTAT_USERMODE_INTERVAL="${OMNISTAT_USERMODE_INTERVAL:-1}" # See material_science/models/HydraGNN/recipes/perf-analysis/agents/launcher.md # for how to build libomnistat_trace.so on a compute node. OMNISTAT_KERNEL_TRACE="${OMNISTAT_KERNEL_TRACE:-0}" -OMNISTAT_TRACE_LIB="${OMNISTAT_TRACE_LIB:-/shared/omnihub/tools/omnistat-src/build-trace/libomnistat_trace.so}" +OMNISTAT_TRACE_LIB="${OMNISTAT_TRACE_LIB:-${OMNIHUB_TOOLS_DIR}/omnistat-src/build-trace/libomnistat_trace.so}" # IMPORTANT: the kernel-trace collector registers /kernel_trace on the SAME # Flask app as the other omnistat endpoints, so the tool library must POST to # the [omnistat.collectors] `port` (8101 here), NOT the library's own default @@ -174,7 +177,7 @@ if [[ "$OMNISTAT_KERNEL_TRACE" == "1" ]]; then echo " Build it on a compute node:" >&2 echo " salloc -p -A -N 1 --time=00:15:00 --gpus-per-node=1 \\" >&2 echo " apptainer exec --rocm \"\$HG_SIF\" bash -c '" >&2 - echo " cd /shared/omnihub/tools/omnistat-src && \\" >&2 + echo " cd ${OMNIHUB_TOOLS_DIR}/omnistat-src && \\" >&2 echo " cmake -S rocprofiler-sdk/ -B build-trace/ -DBUILD_KERNEL_TRACE_LIB=ON && \\" >&2 echo " cmake --build build-trace/ -j 8'" >&2 exit 2 diff --git a/material_science/models/HydraGNN/recipes/perf-analysis/README.md b/material_science/models/HydraGNN/recipes/perf-analysis/README.md index cfb0b8f..fe4fc86 100644 --- a/material_science/models/HydraGNN/recipes/perf-analysis/README.md +++ b/material_science/models/HydraGNN/recipes/perf-analysis/README.md @@ -68,7 +68,8 @@ The `combined_report.md` is the deliverable: ranked bottlenecks, remedies tried ## Prerequisites - Lux access (`vultr_lux` account) with the standard HydraGNN setup at `/shared/aaji/models/HydraGNN/` — overlay, weights, code clone. See [HydraGNN/recipes/train/](../train/). -- One-time install of Omnistat (PR #271 branch `jorda/skills`, with `origin/main` merged in), VictoriaMetrics, and TraceLens — performed lazily by the launcher subagent into `/shared/omnihub/tools/`. +- Export `OMNIHUB_TOOLS_DIR` (the shared perf-tool location, from `.cluster-config.yaml` `omnihub.tools_dir`, e.g. `/shared/omnihub/tools`). Scripts and subagents resolve all tool paths from this var — no cluster path is hardcoded. +- One-time install of Omnistat (PR #271 branch `jorda/skills`, with `origin/main` merged in), VictoriaMetrics, and TraceLens — performed lazily by the launcher subagent into `${OMNIHUB_TOOLS_DIR}/`. - The launcher writes `gfm_mlip_with_profile.json` at submit time (does not modify the upstream `gfm_mlip.json`). ## Telemetry knobs (set at `sbatch` submit time) @@ -76,7 +77,7 @@ The `combined_report.md` is the deliverable: ranked bottlenecks, remedies tried | Env var | Default | What it does | |---|---|---| | `OMNISTAT_KERNEL_TRACE` | `0` | When `1`, loads `libomnistat_trace.so` via `ROCP_TOOL_LIBRARIES` on every rank and turns on the omnistat kernel-trace collector. Adds per-kernel dispatch count + duration time series across all 8 cards on every node. Requires the library to be built once (see `agents/launcher.md` step 1e). Validated end-to-end on job 7034. | -| `OMNISTAT_TRACE_LIB` | `/shared/omnihub/tools/omnistat-src/build-trace/libomnistat_trace.so` | Override only for development. The wrapper hard-fails if the path is missing when `OMNISTAT_KERNEL_TRACE=1`. | +| `OMNISTAT_TRACE_LIB` | `${OMNIHUB_TOOLS_DIR}/omnistat-src/build-trace/libomnistat_trace.so` | Override only for development. The wrapper hard-fails if the path is missing when `OMNISTAT_KERNEL_TRACE=1`. | | `OMNISTAT_TRACE_LOG` | `1` | Library prints `[host][pid][omnistat] Trace summary: N/N processed records (M/M successful flushes)` on rank exit; useful as a smoke-signal that the tool initialized even when the workload crashed. | `enable_rocprofiler=True` is **on** in `omnistat-lux.config.template` by default (since commit `a23e5c4`); the rendered config's state is printed in the sbatch banner so a silent "counters off" cannot recur. Kernel tracing and device-counting can co-exist on a single GPU but raise per-job VictoriaMetrics cardinality — pick by run length and the question you're asking. See `ai4science-studio` SKILL §12 for the device-counting vs kernel-trace vs `rocprofv3` decision matrix. diff --git a/material_science/models/HydraGNN/recipes/perf-analysis/agents/launcher.md b/material_science/models/HydraGNN/recipes/perf-analysis/agents/launcher.md index f1e9c6c..fe5a03b 100644 --- a/material_science/models/HydraGNN/recipes/perf-analysis/agents/launcher.md +++ b/material_science/models/HydraGNN/recipes/perf-analysis/agents/launcher.md @@ -10,8 +10,8 @@ Submits a 2-node HydraGNN training on Lux with PyTorch profiling and Omnistat us ## Outputs -- `/shared/omnihub/tools/omnihub-inspect/` — Python venv with omnistat (PR #271 + main merged) and TraceLens. -- `/shared/omnihub/tools/victoriametrics/victoria-metrics-prod` — VictoriaMetrics binary. +- `${OMNIHUB_TOOLS_DIR}/omnihub-inspect/` — Python venv with omnistat (PR #271 + main merged) and TraceLens. +- `${OMNIHUB_TOOLS_DIR}/victoriametrics/victoria-metrics-prod` — VictoriaMetrics binary. - `/shared/aaji/models/HydraGNN/perf-runs/omnistat-lux.config` — Omnistat user-mode config. - `/shared/aaji/models/HydraGNN/perf-runs//manifest.json` — manifest schema below. - `/shared/aaji/models/HydraGNN/perf-runs//hydragnn-train-.out` — symlinked from output dir. @@ -57,7 +57,8 @@ Submits a 2-node HydraGNN training on Lux with PyTorch profiling and Omnistat us ```bash set -euo pipefail -TOOLS=/shared/omnihub/tools +# Shared tool location — read from .cluster-config.yaml omnihub.tools_dir. +TOOLS="${OMNIHUB_TOOLS_DIR:?set OMNIHUB_TOOLS_DIR from .cluster-config.yaml omnihub.tools_dir (e.g. /shared/omnihub/tools)}" mkdir -p "$TOOLS" # 1a. Omnistat: jorda/skills branch with origin/main merged in @@ -76,7 +77,7 @@ fi OMNISTAT_COMMIT=$(git rev-parse HEAD) # 1b. Venv (py3.12 to match cluster python) -VENV=/shared/omnihub/tools/omnihub-inspect +VENV=${OMNIHUB_TOOLS_DIR}/omnihub-inspect if [[ ! -d $VENV ]]; then python3 -m venv "$VENV" fi @@ -122,7 +123,7 @@ if [[ "${OMNISTAT_KERNEL_TRACE:-0}" == "1" && ! -f "$TRACE_LIB" ]]; then fi ``` -The sbatch wrapper expects the `.so` at `/shared/omnihub/tools/omnistat-src/build-trace/libomnistat_trace.so` by default; override with `OMNISTAT_TRACE_LIB=/path/to/lib.so`. See `ai4science-studio` SKILL §12 for when to enable kernel tracing vs the default device-counting collector. +The sbatch wrapper expects the `.so` at `${OMNIHUB_TOOLS_DIR}/omnistat-src/build-trace/libomnistat_trace.so` by default; override with `OMNISTAT_TRACE_LIB=/path/to/lib.so`. See `ai4science-studio` SKILL §12 for when to enable kernel tracing vs the default device-counting collector. ### 2. Author the Omnistat user-mode config (once, in repo `perf-runs/`) diff --git a/material_science/models/HydraGNN/recipes/perf-analysis/agents/omnistat_analyst.md b/material_science/models/HydraGNN/recipes/perf-analysis/agents/omnistat_analyst.md index a3c9bec..c715cf4 100644 --- a/material_science/models/HydraGNN/recipes/perf-analysis/agents/omnistat_analyst.md +++ b/material_science/models/HydraGNN/recipes/perf-analysis/agents/omnistat_analyst.md @@ -6,8 +6,8 @@ Drive `omnistat-inspect` (PR #271) through the analyze-job phases on the user-mo - `/manifest.json` - The Omnistat DB at `manifest.omnistat_db_path`. -- The `omnihub-inspect` venv at `/shared/omnihub/tools/omnihub-inspect/`. -- The VictoriaMetrics binary at `/shared/omnihub/tools/victoriametrics/victoria-metrics-prod`. +- The `omnihub-inspect` venv at `${OMNIHUB_TOOLS_DIR}/omnihub-inspect/`. +- The VictoriaMetrics binary at `${OMNIHUB_TOOLS_DIR}/victoriametrics/victoria-metrics-prod`. ## Outputs @@ -36,7 +36,7 @@ ss -lnt "sport = :$PORT" 2>/dev/null | grep -q LISTEN && PORT=8429 VMLOG=$PERF_RUN/omnistat/vm.log mkdir -p "$PERF_RUN/omnistat/inspect_outputs" -nohup /shared/omnihub/tools/victoriametrics/victoria-metrics-prod \ +nohup ${OMNIHUB_TOOLS_DIR}/victoriametrics/victoria-metrics-prod \ -storageDataPath="$DB" \ -httpListenAddr=127.0.0.1:$PORT \ -retentionPeriod=100y \ @@ -62,7 +62,7 @@ done ### 2. Run analyze-job phases (per the SKILL) ```bash -. /shared/omnihub/tools/omnihub-inspect/bin/activate +. ${OMNIHUB_TOOLS_DIR}/omnihub-inspect/bin/activate SCRATCH=$PERF_RUN/omnistat/scratch mkdir -p "$SCRATCH" JOBID=$(jq -r .jobid ) diff --git a/material_science/models/HydraGNN/recipes/perf-analysis/agents/tracelens_analyst.md b/material_science/models/HydraGNN/recipes/perf-analysis/agents/tracelens_analyst.md index dd9ead3..66e03a6 100644 --- a/material_science/models/HydraGNN/recipes/perf-analysis/agents/tracelens_analyst.md +++ b/material_science/models/HydraGNN/recipes/perf-analysis/agents/tracelens_analyst.md @@ -39,7 +39,7 @@ Every entry: ### 1. Sanity check ```bash -. /shared/omnihub/tools/omnihub-inspect/bin/activate +. ${OMNIHUB_TOOLS_DIR}/omnihub-inspect/bin/activate PERF_RUN=$(jq -r .perf_run_dir ) TRACE=$(jq -r '.trace_paths[0] // empty' ) [[ -z "$TRACE" ]] && { echo "STATUS=partial; reason=no_trace_in_manifest"; exit 0; } diff --git a/material_science/models/HydraGNN/recipes/perf-optimizer-loop/agents/fom_extractor.md b/material_science/models/HydraGNN/recipes/perf-optimizer-loop/agents/fom_extractor.md index 5d24b55..c749537 100644 --- a/material_science/models/HydraGNN/recipes/perf-optimizer-loop/agents/fom_extractor.md +++ b/material_science/models/HydraGNN/recipes/perf-optimizer-loop/agents/fom_extractor.md @@ -68,7 +68,7 @@ omnistat kernel-trace collector). ```bash set -euo pipefail -. "/shared/omnihub/tools/omnihub-inspect/bin/activate" +. "${OMNIHUB_TOOLS_DIR}/omnihub-inspect/bin/activate" PERF_RUN=$(jq -r .perf_run_dir ) JOBID=$(jq -r .jobid ) RUNTIME=$(jq -r .runtime_seconds ) @@ -102,7 +102,7 @@ Start a local VictoriaMetrics if not already running: PORT=8428 ss -lnt "sport = :$PORT" 2>/dev/null | grep -q LISTEN && PORT=8429 if ! curl -sf "http://127.0.0.1:$PORT/api/v1/status/tsdb" > /dev/null 2>&1; then - nohup "/shared/omnihub/tools/victoriametrics/victoria-metrics-prod" \ + nohup "${OMNIHUB_TOOLS_DIR}/victoriametrics/victoria-metrics-prod" \ -storageDataPath="$PERF_RUN/omnistat-db" \ -httpListenAddr=127.0.0.1:$PORT \ -retentionPeriod=100y -fs.disableMmap \ diff --git a/material_science/models/HydraGNN/recipes/perf-optimizer-loop/agents/orchestrator.md b/material_science/models/HydraGNN/recipes/perf-optimizer-loop/agents/orchestrator.md index 96f0309..23b4a6a 100644 --- a/material_science/models/HydraGNN/recipes/perf-optimizer-loop/agents/orchestrator.md +++ b/material_science/models/HydraGNN/recipes/perf-optimizer-loop/agents/orchestrator.md @@ -85,10 +85,11 @@ Read `loop-/foms.csv` if it exists. If iter rows are present, set `current Run these checks in sequence; abort with `PREFLIGHT_FAIL reason=` on the first failure: +- Ensure `AI4S_SHARED_DIR` and `OMNIHUB_TOOLS_DIR` are exported (the latter from `.cluster-config.yaml` `omnihub.tools_dir`, e.g. `/shared/omnihub/tools`); abort if either is unset. All tool paths below derive from `$OMNIHUB_TOOLS_DIR` — never hardcode a cluster path. - `sinfo -p lux,rad -h` → if zero `IDLE` or `MIX` nodes, abort. Cluster may be in maintenance (see SKILL §14). - `df -h /shared | tail -1` → abort if `Use%` > 95. -- `test -x /shared/omnihub/tools/omnihub-inspect/bin/omnistat-usermode` → abort if missing. -- `test -x /shared/omnihub/tools/victoriametrics/victoria-metrics-prod` → abort if missing. +- `test -x ${OMNIHUB_TOOLS_DIR}/omnihub-inspect/bin/omnistat-usermode` → abort if missing. +- `test -x ${OMNIHUB_TOOLS_DIR}/victoriametrics/victoria-metrics-prod` → abort if missing. - `test -f $AI4S_SHARED_DIR/models/HydraGNN/overlays/hydragnn-overlay.img` → abort if missing. - `test -f $AI4S_SHARED_DIR/images/pytorch_rocm7.2.2_ubuntu24.04_py3.12_pytorch_release_2.10.0.sif` → abort if missing. - If invoked via Claude Code CLI on the login node: `[[ -n "$ANTHROPIC_API_KEY" ]]` and `curl -fsS https://api.anthropic.com/v1/models > /dev/null` (5 s timeout) → abort if either fails. diff --git a/material_science/models/HydraGNN/recipes/perf-optimizer-loop/agents/story_writer.md b/material_science/models/HydraGNN/recipes/perf-optimizer-loop/agents/story_writer.md index f63129c..016e386 100644 --- a/material_science/models/HydraGNN/recipes/perf-optimizer-loop/agents/story_writer.md +++ b/material_science/models/HydraGNN/recipes/perf-optimizer-loop/agents/story_writer.md @@ -156,7 +156,7 @@ fig.tight_layout() fig.savefig(out_png_path, dpi=150, bbox_inches="tight") ``` -Use `matplotlib.use("Agg")` for headless rendering. Do NOT use seaborn (extra dep not in the omnistat venv). Use the omnistat venv at `/shared/omnihub/tools/omnihub-inspect/bin/python` which already has matplotlib via TraceLens deps. +Use `matplotlib.use("Agg")` for headless rendering. Do NOT use seaborn (extra dep not in the omnistat venv). Use the omnistat venv at `${OMNIHUB_TOOLS_DIR}/omnihub-inspect/bin/python` which already has matplotlib via TraceLens deps. ## Style rules diff --git a/material_science/models/HydraGNN/recipes/perf-optimizer-loop/dispatch-attribution.md b/material_science/models/HydraGNN/recipes/perf-optimizer-loop/dispatch-attribution.md index 94f162a..7a90a9f 100644 --- a/material_science/models/HydraGNN/recipes/perf-optimizer-loop/dispatch-attribution.md +++ b/material_science/models/HydraGNN/recipes/perf-optimizer-loop/dispatch-attribution.md @@ -59,10 +59,10 @@ Short-kernel study (`short_kernels_summary.csv`): thousands of `aten::fill_`, `a ```bash export AI4S_SHARED_DIR=/shared/aaji export REPO_ROOT=$HOME/git/ai4science-studio -source /shared/omnihub/tools/omnihub-inspect/bin/activate +source ${OMNIHUB_TOOLS_DIR}/omnihub-inspect/bin/activate # VictoriaMetrics on the perf-run DB (login node) -/shared/omnihub/tools/victoriametrics/victoria-metrics-prod \ +${OMNIHUB_TOOLS_DIR}/victoriametrics/victoria-metrics-prod \ -storageDataPath=$AI4S_SHARED_DIR/models/HydraGNN/perf-runs/7187/omnistat-db \ -httpListenAddr=127.0.0.1:8432 -retentionPeriod=100y -fs.disableMmap & From 1b1c24d88f2ad45b06834ce61271570707474e91 Mon Sep 17 00:00:00 2001 From: Ashwin Aji Date: Thu, 28 May 2026 23:12:42 +0000 Subject: [PATCH 19/31] OmniHub integration for HydraGNN train + perf-analysis Add integrations/omnihub/ bridge (sync + generate-job + local validation), the ai4science-omnihub Cursor skill, and train-omnihub / perf-analysis-omnihub recipes in HydraGNN model.yaml. Validated end-to-end on Lux MI355X: single- and 8-GPU training (jobs 8004/8005) and an omnistat perf run (8014). generate-job.sh injects Lux-specific fixes: explicit --gpus-per-node, HydraGNN ext3 overlay (no stacked writable overlay), container PYTHONPATH/ LD_LIBRARY_PATH for adios2, shared OmniHub checkout for compute nodes, and single-flag --tools passing. Co-authored-by: Cursor --- .cluster-config.example.yaml | 8 + .cursor/skills/ai4science-omnihub/SKILL.md | 98 +++++++++ .cursor/skills/ai4science-run-models/SKILL.md | 1 + .gitignore | 1 + integrations/omnihub/README.md | 78 +++++++ .../hydragnn-train/config-perf.yaml | 11 + .../applications/hydragnn-train/config.yaml | 13 ++ .../applications/hydragnn-train/parse.py | 37 ++++ .../hydragnn-train/train_wrapper.py | 205 ++++++++++++++++++ integrations/omnihub/generate-job.sh | 194 +++++++++++++++++ integrations/omnihub/omnistat-parity-check.sh | 89 ++++++++ .../omnihub/omnistat-parity-report.txt | 37 ++++ integrations/omnihub/render-cluster-config.sh | 65 ++++++ integrations/omnihub/sync-to-omnihub.sh | 35 +++ integrations/omnihub/validate-local.sh | 64 ++++++ material_science/models/HydraGNN/model.yaml | 8 + .../HydraGNN/recipes/perf-analysis/README.md | 12 + 17 files changed, 956 insertions(+) create mode 100644 .cursor/skills/ai4science-omnihub/SKILL.md create mode 100644 integrations/omnihub/README.md create mode 100644 integrations/omnihub/applications/hydragnn-train/config-perf.yaml create mode 100644 integrations/omnihub/applications/hydragnn-train/config.yaml create mode 100755 integrations/omnihub/applications/hydragnn-train/parse.py create mode 100644 integrations/omnihub/applications/hydragnn-train/train_wrapper.py create mode 100755 integrations/omnihub/generate-job.sh create mode 100755 integrations/omnihub/omnistat-parity-check.sh create mode 100644 integrations/omnihub/omnistat-parity-report.txt create mode 100755 integrations/omnihub/render-cluster-config.sh create mode 100755 integrations/omnihub/sync-to-omnihub.sh create mode 100755 integrations/omnihub/validate-local.sh diff --git a/.cluster-config.example.yaml b/.cluster-config.example.yaml index 6068434..4759a52 100644 --- a/.cluster-config.example.yaml +++ b/.cluster-config.example.yaml @@ -40,3 +40,11 @@ network: ib_hca: "" # IB HCA device list for RCCL (discover: ibstat | grep 'CA ') rccl_anp_plugin: "" # Path to librccl-anp.so on host (default: /opt/rocm/lib/librccl-anp.so) libionic_path: "" # Path to libionic.so.1 on host (default: /usr/lib/x86_64-linux-gnu/libionic.so.1) + +# OmniHub integration (sibling repo — see integrations/omnihub/README.md) +omnihub: + repo_dir: "" # e.g. /home/$USER/git/omnihub + cluster: vultr # maps to omnihub config/vultr.yaml on Lux + shared_dir: /shared/omnihub + results_dir: "" # default: /shared/$USER/results/omnihub + rocm_version: "7.2.2" diff --git a/.cursor/skills/ai4science-omnihub/SKILL.md b/.cursor/skills/ai4science-omnihub/SKILL.md new file mode 100644 index 0000000..67f660e --- /dev/null +++ b/.cursor/skills/ai4science-omnihub/SKILL.md @@ -0,0 +1,98 @@ +--- +name: ai4science-omnihub +description: Launch AI4Science Studio models via OmniHub job generation, SLURM submission, and standardized result processing. Use when the user asks for OmniHub, train-omnihub, perf-analysis-omnihub, omnihub-generate-job, or post-processing OmniHub results for ai4science models. +--- + +# AI4Science Studio + OmniHub + +Integrates this repo with a sibling [OmniHub](https://github.com/AMDResearch/omnihub) checkout **without merging repos**. + +## Prerequisites + +1. Read `.cluster-config.yaml` (or `~/.config/ai4science-studio/cluster.yaml`) — especially `omnihub:` and `paths`. +2. Set `export OMNIHUB_DIR=/path/to/omnihub` and `export AI4S_SHARED_DIR=/shared/$USER` (or your site path). +3. HydraGNN assets: SIF, overlay, datasets, code clone under `$AI4S_SHARED_DIR/models/HydraGNN/`. + +Validate cluster alignment: + +```bash +./integrations/omnihub/render-cluster-config.sh +./integrations/omnihub/omnistat-parity-check.sh +``` + +## Workflow: HydraGNN train-omnihub + +1. Read `material_science/models/HydraGNN/model.yaml` recipe `train-omnihub`. +2. Sync app stubs and generate SLURM script: + +```bash +./integrations/omnihub/generate-job.sh \ + --num-nodes 2 \ + --partition lux \ + --time-limit 2h \ + --output /tmp/hydragnn-omnihub.slurm +``` + +3. Submit (account on `sbatch`, not in the generated script): + +```bash +sbatch -A vultr_lux /tmp/hydragnn-omnihub.slurm +``` + +4. Monitor: `squeue -j `, `tail -f /shared/$USER/results/omnihub//*.out` + +5. After completion — post-process: + +```bash +$OMNIHUB_DIR/omnihub-process --results-dir /shared/$USER/results/omnihub -j 4 +$OMNIHUB_DIR/omnihub-index --results-dir /shared/$USER/results/omnihub --output index +``` + +## Workflow: perf-analysis-omnihub (Phase 2) + +Same as train, but add `--perf` to enable `omnistat`, `pytorch-trace`, and `tracelens`: + +```bash +./integrations/omnihub/generate-job.sh --perf --num-nodes 2 --time-limit 30m \ + --output /tmp/hydragnn-perf.slurm +sbatch -A vultr_lux /tmp/hydragnn-perf.slurm +``` + +**Agent result reading order** (minimal tokens): + +1. `/processed-data/tracelens-summary.json` +2. `/processed-data/omnistat-range.yaml` +3. `/processed-data/report-card.yaml` +4. Raw traces under `/tools/` only if needed + +Legacy HydraGNN perf (`sbatch_train_perf_amd.sh` + `omnistat-pr271` + subagent claims) remains for deep PromQL/inspect workflows. See `material_science/models/HydraGNN/recipes/perf-analysis/README.md`. + +## Key paths (Lux / vultr) + +| Resource | Path | +|----------|------| +| OmniHub shared | `/shared/omnihub` | +| Omnistat | `/shared/omnihub/tools/omnistat/` → use `--tools omnistat` | +| Results | `/shared/$USER/results/omnihub/$SLURM_JOB_ID/` | +| HydraGNN overlay | `$AI4S_SHARED_DIR/models/HydraGNN/overlays/hydragnn-overlay.img` | + +## Runner + +HydraGNN uses MPI — default `--runner manual-mpi --tasks-per-node 8` (OmniHub fork). Do **not** use `--runner manual` (that enables torch DDP init). + +## Related OmniHub skills + +When working inside `$OMNIHUB_DIR`, also use: + +- `.cursor/skills/generate-job-and-run/SKILL.md` +- `.cursor/skills/post-process-results/SKILL.md` +- `.cursor/skills/check-job-results/SKILL.md` + +## Troubleshooting + +| Issue | Fix | +|-------|-----| +| App config not found | Run `sync-to-omnihub.sh` first | +| ROCm / overlay missing | Verify `HG_OVERLAY`, `HG_SIF`; check generate-job injected binds | +| Omnistat empty | Confirm `/shared/omnihub/tools/omnistat/`; run parity check | +| TraceLens skipped | Install venv at `/shared/omnihub/tools/tracelens/venv` or set `TRACELENS_VENV` | diff --git a/.cursor/skills/ai4science-run-models/SKILL.md b/.cursor/skills/ai4science-run-models/SKILL.md index 7ea7106..7b46896 100644 --- a/.cursor/skills/ai4science-run-models/SKILL.md +++ b/.cursor/skills/ai4science-run-models/SKILL.md @@ -33,6 +33,7 @@ Match the user's request to a recipe in `model.yaml`: - "run inference" / "predict" / "generate" → look for `task: inference` - "train" / "fine-tune" / "pair-tune" → look for `task: train` or `task: finetune` - "ensemble" → look for `task: ensemble` +- "omnihub" / "train-omnihub" → look for `task: train-omnihub` and use `.cursor/skills/ai4science-omnihub/SKILL.md` ## Step 3: Ask the user for required inputs diff --git a/.gitignore b/.gitignore index 2911aba..5cc637e 100755 --- a/.gitignore +++ b/.gitignore @@ -45,6 +45,7 @@ secrets/ # SLURM job output *-*.out +*-slurm.err # SLURM job artifacts named foo..out (e.g. omnistat_failed_hosts.6851.out) *.[0-9][0-9][0-9][0-9].out diff --git a/integrations/omnihub/README.md b/integrations/omnihub/README.md new file mode 100644 index 0000000..6f3d80c --- /dev/null +++ b/integrations/omnihub/README.md @@ -0,0 +1,78 @@ +# OmniHub integration for AI4Science Studio + +Bridge between [AI4Science Studio](https://github.com/AMDResearch/ai4science-studio) model recipes and [OmniHub](https://github.com/AMDResearch/omnihub) job generation, without merging the two repositories. + +## Prerequisites + +- OmniHub clone (sibling repo), e.g. `$HOME/git/omnihub` +- `.cluster-config.yaml` with `omnihub:` section (see `.cluster-config.example.yaml`) +- `AI4S_SHARED_DIR` set (HydraGNN assets: SIF, overlay, datasets, code) +- Lux/Vultr: OmniHub `config/vultr.yaml` (`shared-dir: /shared/omnihub`) + +## Quick start (HydraGNN train) + +```bash +export AI4S_SHARED_DIR=/shared/$USER +export OMNIHUB_DIR=$HOME/git/omnihub + +# 1. Sync application stubs into the OmniHub checkout +./integrations/omnihub/sync-to-omnihub.sh + +# 2. Generate SLURM job (wrapper sets HydraGNN binds/overlay) +./integrations/omnihub/generate-job.sh \ + --num-nodes 2 --partition lux --time-limit 2h \ + --output /tmp/hydragnn-omnihub.slurm + +# 3. Submit +sbatch -A vultr_lux /tmp/hydragnn-omnihub.slurm + +# 4. Post-process (after job completes) +$OMNIHUB_DIR/omnihub-process --results-dir /shared/$USER/results/omnihub -j 4 +$OMNIHUB_DIR/omnihub-index --results-dir /shared/$USER/results/omnihub --output index +``` + +## Layout + +``` +integrations/omnihub/ +├── README.md # this file +├── sync-to-omnihub.sh # rsync apps → $OMNIHUB_DIR/applications/ +├── render-cluster-config.sh # validate .cluster-config vs vultr.yaml +├── generate-job.sh # sync + omnihub-generate-job + HydraGNN apptainer patches +├── omnistat-parity-check.sh # compare /shared/omnihub vs omnistat-pr271 +└── applications/ + └── hydragnn-train/ # OmniHub app (synced into omnihub repo) +``` + +## Environment variables + +| Variable | Purpose | +|----------|---------| +| `OMNIHUB_DIR` | Path to OmniHub repo (required for sync/generate) | +| `AI4S_SHARED_DIR` | Shared assets root (SIF, HydraGNN overlay, datasets) | +| `AI4S_REPO_ROOT` | ai4science-studio root (default: auto-detected) | + +HydraGNN training knobs (`HG_*`, `HYDRAGNN_*`) match the standard `sbatch_train_amd.sh` recipe. + +## Results layout + +OmniHub jobs write to `/shared/$USER/results/omnihub/$SLURM_JOB_ID/` (per `config/vultr.yaml`): + +``` +/ +├── job.yaml, app.yaml, job-status.yaml +├── logs/srun-*.out +├── tools/omnihub-monitor/ +└── processed-data/ # after omnihub-process +``` + +Agents should read `processed-data/` first (token-efficient). See `.cursor/skills/ai4science-omnihub/SKILL.md`. + +## Recipes + +| task | OmniHub tools | Notes | +|------|---------------|-------| +| `train-omnihub` | `omnihub-monitor` | Phase 1 pilot | +| `perf-analysis-omnihub` | `omnihub-monitor omnistat pytorch-trace tracelens` | Phase 2; optional `omnistat-rocprofiler-pmc1` | + +Legacy HydraGNN perf (`sbatch_train_perf_amd.sh` + `omnistat-pr271`) remains available for deep agent workflows. diff --git a/integrations/omnihub/applications/hydragnn-train/config-perf.yaml b/integrations/omnihub/applications/hydragnn-train/config-perf.yaml new file mode 100644 index 0000000..9a2e5a6 --- /dev/null +++ b/integrations/omnihub/applications/hydragnn-train/config-perf.yaml @@ -0,0 +1,11 @@ +entrypoint: applications/ai4science-hydragnn-train/train_wrapper.py + +# Perf-analysis variant: epoch-gated kineto on rank 0 (use with --perf in generate-job.sh) + +HydraGNN: + datasets: ANI1x,Alexandria + batch_size: 200 + num_epoch: 2 + precision: fp64 + max_num_batch: 30 + profile_target_epoch: 1 diff --git a/integrations/omnihub/applications/hydragnn-train/config.yaml b/integrations/omnihub/applications/hydragnn-train/config.yaml new file mode 100644 index 0000000..1a676b5 --- /dev/null +++ b/integrations/omnihub/applications/hydragnn-train/config.yaml @@ -0,0 +1,13 @@ +entrypoint: applications/ai4science-hydragnn-train/train_wrapper.py + +# HydraGNN multi-dataset GFM training via OmniHub. +# Paths resolve from AI4S_SHARED_DIR / HG_* env vars (see train_wrapper.py). + +HydraGNN: + datasets: ANI1x,Alexandria + batch_size: 200 + num_epoch: 1 + precision: fp64 + max_num_batch: null + valtest: 0 + profile_target_epoch: null diff --git a/integrations/omnihub/applications/hydragnn-train/parse.py b/integrations/omnihub/applications/hydragnn-train/parse.py new file mode 100755 index 0000000..9aabb0c --- /dev/null +++ b/integrations/omnihub/applications/hydragnn-train/parse.py @@ -0,0 +1,37 @@ +#!/usr/bin/env python3 +"""Post-process HydraGNN OmniHub job results into processed-data/.""" + +import json +import os +import sys +from pathlib import Path + + +def main(): + if len(sys.argv) < 2: + print(f"Usage: {sys.argv[0]} ", file=sys.stderr) + sys.exit(1) + + job_dir = Path(sys.argv[1]) + processed = job_dir / "processed-data" + processed.mkdir(parents=True, exist_ok=True) + + summary = { + "app": "ai4science-hydragnn-train", + "job_dir": str(job_dir), + } + + status = job_dir / "job-status.yaml" + if status.is_file(): + summary["job_status_file"] = str(status) + + logs = list(job_dir.glob("logs/srun-*.out")) + summary["log_count"] = len(logs) + + out = processed / "app-parser.json" + out.write_text(json.dumps(summary, indent=2)) + print(f"Wrote {out}") + + +if __name__ == "__main__": + main() diff --git a/integrations/omnihub/applications/hydragnn-train/train_wrapper.py b/integrations/omnihub/applications/hydragnn-train/train_wrapper.py new file mode 100644 index 0000000..aaff709 --- /dev/null +++ b/integrations/omnihub/applications/hydragnn-train/train_wrapper.py @@ -0,0 +1,205 @@ +#!/usr/bin/env python3 +""" +HydraGNN training wrapper for OmniHub (AI4Science Studio integration). + +Runs the upstream multidataset_hpo_sc26 entrypoint inside the container with +MPI-style rank identity from Slurm (manual-mpi runner). Mirrors the rank script +in material_science/models/HydraGNN/examples/sbatch_train_amd.sh. +""" + +from __future__ import annotations + +import json +import os +import runpy +import sys +from pathlib import Path +from typing import Any, Dict, Optional + +import omnihub.run +import omnihub.tools + +# Pinned SHA matches sbatch_train_amd.sh default. +DEFAULT_HG_SHA = "2fb0bd0157e3c85a74f9841887155095bd163303" +DEFAULT_AVG_NEIGHBORS = 13.735293601560318 + + +def _require_env(name: str) -> str: + val = os.environ.get(name) + if not val: + raise RuntimeError(f"Required environment variable not set: {name}") + return val + + +def _ai4s_base() -> Path: + return Path(_require_env("AI4S_SHARED_DIR")) / "models" / "HydraGNN" + + +def _resolve_paths(config: Dict[str, Any]) -> Dict[str, Path]: + hg = config.get("HydraGNN", {}) + base = _ai4s_base() + repo = Path(os.environ.get("HG_REPO_DIR", base / "code" / "HydraGNN")) + example = repo / "examples" / "multidataset_hpo_sc26" + data_dir = Path(os.environ.get("HG_DATA_DIR", base / "weights")) + output = Path(os.environ.get("HG_OUTPUT_DIR", os.environ.get("OMNIHUB_RESULTS_DIR", "."))) + return { + "repo": repo, + "example": example, + "data_dir": data_dir, + "output": output, + "config_json": example / "gfm_mlip.json", + "entry": example / "gfm_mlip_all_mpnn.py", + } + + +def _ensure_datasets(paths: Dict[str, Path], datasets: str) -> None: + dataset_dir = paths["example"] / "dataset" + dataset_dir.mkdir(parents=True, exist_ok=True) + for ds in datasets.split(","): + ds = ds.strip() + if not ds: + continue + target = dataset_dir / f"{ds}-v2.bp" + source = paths["data_dir"] / f"{ds}-v2.bp" + if not source.is_dir(): + raise FileNotFoundError(f"Dataset not found: {source}") + if not target.exists(): + target.symlink_to(source, target_is_directory=True) + + +def _maybe_inject_profile(config_path: Path, target_epoch: Optional[int]) -> Path: + """Inject NeuralNetwork.Profile block for epoch-gated kineto (perf-analysis).""" + if target_epoch is None: + return config_path + data = json.loads(config_path.read_text()) + nn = data.setdefault("NeuralNetwork", {}) + nn["Profile"] = { + "enable": 1, + "target_epoch": int(target_epoch), + "wait": 5, + "warmup": 3, + "active": 3, + "repeat": 1, + } + out = Path( + os.environ.get( + "OMNIHUB_RESULTS_DIR", + os.environ.get("HG_OUTPUT_DIR", str(config_path.parent)), + ) + ) / "gfm_mlip_omnihub_profile.json" + out.write_text(json.dumps(data, indent=2)) + return out + + +def _patch_adios_avg_neighbors() -> None: + avg_nn = float(os.environ.get("HYDRAGNN_AVG_NUM_NEIGHBORS", DEFAULT_AVG_NEIGHBORS)) + if avg_nn <= 0: + return + import hydragnn.utils.datasets.adiosdataset as adm + + orig_init = adm.AdiosMultiDataset.__init__ + + def patched_init(self, *args, **kwargs): + orig_init(self, *args, **kwargs) + self.avg_num_neighbors = avg_nn + + adm.AdiosMultiDataset.__init__ = patched_init + + +def _setup_runtime_env(paths: Dict[str, Path]) -> None: + scratch = os.environ.get("SCRATCH_LOCAL", "/scratch") + proc = os.environ.get("SLURM_PROCID", os.environ.get("RANK", "0")) + job = os.environ.get("SLURM_JOB_ID", "local") + scratch_rank = Path(scratch) / os.environ.get("USER", "user") / f"hydragnn-{job}" / proc + scratch_rank.mkdir(parents=True, exist_ok=True) + os.environ["TMPDIR"] = str(scratch_rank) + os.environ["MIOPEN_DISABLE_CACHE"] = "1" + os.environ["MIOPEN_USER_DB_PATH"] = str(scratch_rank / "miopen") + Path(os.environ["MIOPEN_USER_DB_PATH"]).mkdir(parents=True, exist_ok=True) + os.environ["PYTHONNOUSERSITE"] = "1" + os.environ["HYDRAGNN_USE_VARIABLE_GRAPH_SIZE"] = "1" + os.environ["HYDRAGNN_AGGR_BACKEND"] = "mpi" + os.environ["HYDRAGNN_USE_FSDP"] = "0" + os.environ["HYDRAGNN_TRACE_LEVEL"] = os.environ.get("HYDRAGNN_TRACE_LEVEL", "1") + os.environ["HSA_NO_SCRATCH_RECLAIM"] = "1" + os.environ["HG_EXAMPLE_DIR"] = str(paths["example"]) + os.environ["HG_OUTPUT_DIR"] = str(paths["output"]) + paths["output"].mkdir(parents=True, exist_ok=True) + + pkg = "/opt/hydragnn-pkgs" + if Path(pkg).is_dir(): + os.environ["PYTHONPATH"] = f"{pkg}:{os.environ.get('PYTHONPATH', '')}" + ld = f"{pkg}/adios2:/opt/ompi/lib" + os.environ["LD_LIBRARY_PATH"] = f"{ld}:{os.environ.get('LD_LIBRARY_PATH', '')}" + + +def _build_argv(paths: Dict[str, Path], hg: Dict[str, Any], config_file: Path) -> list: + datasets = os.environ.get("HG_DATASETS", hg.get("datasets", "ANI1x,Alexandria")) + batch_size = int(os.environ.get("HG_BATCH_SIZE", hg.get("batch_size", 200))) + max_batch = os.environ.get("HYDRAGNN_MAX_NUM_BATCH", hg.get("max_num_batch") or "") + num_epoch = os.environ.get("HG_NUM_EPOCH", str(hg.get("num_epoch", 1))) + precision = os.environ.get("HG_PRECISION", hg.get("precision", "fp64")) + job = os.environ.get("SLURM_JOB_ID", "0") + nodes = os.environ.get("SLURM_JOB_NUM_NODES", "1") + + argv = [ + str(paths["entry"]), + f"--inputfile={config_file}", + "--multi", + "--everyone", + f"--multi_model_list={datasets}", + f"--precision={precision}", + f"--batch_size={batch_size}", + f"--num_epoch={num_epoch}", + "--startfrom=none", + f"--log=hydragnn-train-{job}-N{nodes}", + ] + if max_batch: + n = int(max_batch) * batch_size + argv += [ + f"--num_samples={n}", + "--oversampling", + f"--oversampling_num_samples={n}", + ] + return argv + + +@omnihub.run.entrypoint +def run(extra_args, config=None): + """OmniHub entrypoint — train HydraGNN GFM ensemble.""" + hg = (config or {}).get("HydraGNN", {}) + paths = _resolve_paths(config or {}) + datasets = os.environ.get("HG_DATASETS", hg.get("datasets", "ANI1x,Alexandria")) + + if not paths["entry"].is_file(): + raise FileNotFoundError( + f"HydraGNN entry not found: {paths['entry']}. " + f"Clone upstream to HG_REPO_DIR (SHA {DEFAULT_HG_SHA})." + ) + + _ensure_datasets(paths, datasets) + profile_epoch = hg.get("profile_target_epoch") + if profile_epoch is not None: + profile_epoch = int(profile_epoch) + elif os.environ.get("PROFILE_TARGET_EPOCH"): + profile_epoch = int(os.environ["PROFILE_TARGET_EPOCH"]) + + config_file = _maybe_inject_profile(paths["config_json"], profile_epoch) + _setup_runtime_env(paths) + + rank = int(os.environ.get("SLURM_PROCID", os.environ.get("RANK", "0"))) + if os.environ.get("PROFILE_RANK0_ONLY", "1") == "1" and rank != 0: + os.environ["HYDRAGNN_DISABLE_PROFILE"] = "1" + + @omnihub.tools.profile() + def _train(): + _patch_adios_avg_neighbors() + os.chdir(paths["output"]) + sys.argv = _build_argv(paths, hg, config_file) + runpy.run_path(str(paths["entry"]), run_name="__main__") + + _train() + + +if __name__ == "__main__": + run([], {}) diff --git a/integrations/omnihub/generate-job.sh b/integrations/omnihub/generate-job.sh new file mode 100755 index 0000000..96be6e1 --- /dev/null +++ b/integrations/omnihub/generate-job.sh @@ -0,0 +1,194 @@ +#!/usr/bin/env bash +# Generate an OmniHub SLURM job for ai4science HydraGNN with Lux-specific apptainer binds. +set -euo pipefail + +SCRIPT_DIR=$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd) +AI4S_REPO_ROOT="${AI4S_REPO_ROOT:-$(cd "$SCRIPT_DIR/../.." && pwd)}" +OMNIHUB_DIR="${OMNIHUB_DIR:-}" +AI4S_SHARED_DIR="${AI4S_SHARED_DIR:-}" + +NUM_NODES=1 +PARTITION="" +CLUSTER=vultr +TIME_LIMIT=2h +TOOLS=(omnihub-monitor) +RUNNER=manual-mpi +TASKS_PER_NODE=8 +OUTPUT="" +IMAGE="" +APP_CONFIG=applications/ai4science-hydragnn-train/config.yaml +ROCM_VERSION=7.2.2 +EXTRA_ARGS=() + +usage() { + cat </dev/null || ls "$SHARED_OMNI/tools/omnistat-rocprofiler/"omnihub.*.config 2>/dev/null || echo "(none found)" + + echo "" + echo "### rocprofiler-sdk Python extension" + if "$OMNI_VENV/python" -c "from omnistat.rocprofiler_sdk_extension import get_samplers" 2>/dev/null; then + echo "sdk extension: OK" + else + echo "sdk extension: failed on login node (may work on compute nodes with ROCm loaded)" + fi + + echo "" + echo "### Kernel trace library" + if [[ -f "$SHARED_OMNI/tools/omnistat/build-trace/libomnistat_trace.so" ]]; then + echo "libomnistat_trace.so: present" + else + echo "libomnistat_trace.so: not found" + fi + + echo "" + echo "## ai4science omnistat-pr271 (reference)" + PR271="$AI4S_SHARED/tools/omnistat-pr271/bin" + if [[ -d "$AI4S_SHARED/tools/omnistat-pr271" ]]; then + for cmd in omnistat-usermode omnistat-query omnistat-inspect; do + if [[ -x "$PR271/$cmd" ]]; then + echo "- $cmd: present" + else + echo "- $cmd: missing" + fi + done + else + echo "Not installed at $AI4S_SHARED/tools/omnistat-pr271 (HydraGNN launcher creates this)" + fi + + echo "" + echo "## Integration recommendation" + echo "- OmniHub path: use --tools omnistat + omnihub-process -> processed-data/" + echo "- Deep agent path (omnistat-inspect, PromQL): legacy sbatch_train_perf_amd.sh + omnistat-pr271" + echo "- If omnistat-inspect missing under /shared/omnihub, agents should not expect inspect JSON from OmniHub jobs" +} | report diff --git a/integrations/omnihub/omnistat-parity-report.txt b/integrations/omnihub/omnistat-parity-report.txt new file mode 100644 index 0000000..65fabc5 --- /dev/null +++ b/integrations/omnihub/omnistat-parity-report.txt @@ -0,0 +1,37 @@ +# Omnistat parity check — 2026-05-28T21:06:35+00:00 + +## /shared/omnihub/tools/omnistat + +- omnistat-usermode: present +- omnistat-query: present +- omnistat-inspect: **MISSING** + +### Site config (key knobs) +enable_amd_smi = False +enable_rocprofiler = False +enable_vendor_counters = False +job_detection_file = /tmp/omni_rmsjobinfo_omnihub + +### Rocprofiler PMC configs +/shared/omnihub/tools/omnistat-rocprofiler/omnihub.gfx90a.1.config +/shared/omnihub/tools/omnistat-rocprofiler/omnihub.gfx90a.2.config +/shared/omnihub/tools/omnistat-rocprofiler/omnihub.gfx942.1.config +/shared/omnihub/tools/omnistat-rocprofiler/omnihub.gfx942.2.config +/shared/omnihub/tools/omnistat-rocprofiler/omnihub.gfx950.1.config +/shared/omnihub/tools/omnistat-rocprofiler/omnihub.gfx950.2.config + +### rocprofiler-sdk Python extension +sdk extension: failed on login node (may work on compute nodes with ROCm loaded) + +### Kernel trace library +libomnistat_trace.so: not found + +## ai4science omnistat-pr271 (reference) +- omnistat-usermode: present +- omnistat-query: present +- omnistat-inspect: present + +## Integration recommendation +- OmniHub path: use --tools omnistat + omnihub-process -> processed-data/ +- Deep agent path (omnistat-inspect, PromQL): legacy sbatch_train_perf_amd.sh + omnistat-pr271 +- If omnistat-inspect missing under /shared/omnihub, agents should not expect inspect JSON from OmniHub jobs diff --git a/integrations/omnihub/render-cluster-config.sh b/integrations/omnihub/render-cluster-config.sh new file mode 100755 index 0000000..de19695 --- /dev/null +++ b/integrations/omnihub/render-cluster-config.sh @@ -0,0 +1,65 @@ +#!/usr/bin/env bash +# Validate ai4science .cluster-config.yaml omnihub section against OmniHub vultr.yaml. +set -euo pipefail + +SCRIPT_DIR=$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd) +AI4S_REPO_ROOT="${AI4S_REPO_ROOT:-$(cd "$SCRIPT_DIR/../.." && pwd)}" +OMNIHUB_DIR="${OMNIHUB_DIR:-$HOME/git/omnihub}" + +CLUSTER_CFG="${AI4S_CLUSTER_CONFIG:-$AI4S_REPO_ROOT/.cluster-config.yaml}" +VULTR_YAML="$OMNIHUB_DIR/config/vultr.yaml" + +_yaml_get() { + local key=$1 file=$2 + grep -E "^[[:space:]]*${key}:" "$file" 2>/dev/null | head -1 | sed -E 's/^[^:]*:[[:space:]]*"?([^"]*)"?/\1/' | tr -d '"' +} + +echo "=== OmniHub cluster config bridge ===" +echo " ai4science config: $CLUSTER_CFG" +echo " omnihub vultr.yaml: $VULTR_YAML" +echo "" + +if [[ ! -f "$CLUSTER_CFG" ]]; then + echo "WARN: no .cluster-config.yaml — copy from .cluster-config.example.yaml" >&2 +fi + +if [[ ! -f "$VULTR_YAML" ]]; then + echo "ERROR: missing $VULTR_YAML" >&2 + exit 2 +fi + +SHARED_OMNI=$(_yaml_get shared-dir "$VULTR_YAML") +RESULTS_OMNI=$(_yaml_get results-dir "$VULTR_YAML") + +echo "OmniHub vultr.yaml:" +echo " shared-dir: $SHARED_OMNI" +echo " results-dir: $RESULTS_OMNI" +echo "" + +if [[ -f "$CLUSTER_CFG" ]]; then + REPO_DIR=$(_yaml_get repo_dir "$CLUSTER_CFG" 2>/dev/null || true) + # nested under omnihub: — grep with context + OMNI_REPO=$(awk '/^omnihub:/{f=1} f && /repo_dir:/{print; exit}' "$CLUSTER_CFG" | sed -E 's/.*:[[:space:]]*"?([^"]*)"?/\1/') + OMNI_SHARED=$(awk '/^omnihub:/{f=1} f && /shared_dir:/{print; exit}' "$CLUSTER_CFG" | sed -E 's/.*:[[:space:]]*"?([^"]*)"?/\1/') + + echo "ai4science omnihub section:" + echo " repo_dir: ${OMNI_REPO:-}" + echo " shared_dir: ${OMNI_SHARED:-}" + echo "" + + if [[ -n "${OMNI_SHARED:-}" && "$OMNI_SHARED" != "$SHARED_OMNI" ]]; then + echo "WARN: omnihub.shared_dir ($OMNI_SHARED) != vultr shared-dir ($SHARED_OMNI)" >&2 + fi +fi + +for tool_path in "$SHARED_OMNI/tools/omnistat/venv/bin/omnistat-usermode" \ + "$SHARED_OMNI/tools/omnistat-rocprofiler/venv/bin/omnistat-usermode"; do + if [[ -x "$tool_path" ]]; then + echo "OK: $tool_path" + else + echo "MISSING: $tool_path" >&2 + fi +done + +echo "" +echo "Run ./integrations/omnihub/omnistat-parity-check.sh for detailed Omnistat comparison." diff --git a/integrations/omnihub/sync-to-omnihub.sh b/integrations/omnihub/sync-to-omnihub.sh new file mode 100755 index 0000000..496dd1c --- /dev/null +++ b/integrations/omnihub/sync-to-omnihub.sh @@ -0,0 +1,35 @@ +#!/usr/bin/env bash +# Sync ai4science OmniHub application stubs into the OmniHub checkout. +# +# OmniHub requires --app-config paths relative to the omnihub repo root. +# Source of truth lives here; this script copies into $OMNIHUB_DIR/applications/. +set -euo pipefail + +SCRIPT_DIR=$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd) +AI4S_REPO_ROOT="${AI4S_REPO_ROOT:-$(cd "$SCRIPT_DIR/../.." && pwd)}" +OMNIHUB_DIR="${OMNIHUB_DIR:-}" + +if [[ -z "$OMNIHUB_DIR" ]]; then + echo "ERROR: set OMNIHUB_DIR to your OmniHub checkout (e.g. export OMNIHUB_DIR=\$HOME/git/omnihub)" >&2 + exit 2 +fi + +if [[ ! -x "$OMNIHUB_DIR/omnihub-generate-job" ]]; then + echo "ERROR: OMNIHUB_DIR does not look like an OmniHub repo: $OMNIHUB_DIR" >&2 + exit 2 +fi + +SRC="$SCRIPT_DIR/applications" +DEST="$OMNIHUB_DIR/applications" + +mkdir -p "$DEST" + +for app_dir in "$SRC"/*/; do + [[ -d "$app_dir" ]] || continue + name=$(basename "$app_dir") + target="$DEST/ai4science-${name}" + echo "Syncing $name -> $target" + rsync -a --delete "$app_dir" "$target/" +done + +echo "Done. Use --app-config applications/ai4science-/config.yaml with omnihub-generate-job." diff --git a/integrations/omnihub/validate-local.sh b/integrations/omnihub/validate-local.sh new file mode 100755 index 0000000..81b3a3a --- /dev/null +++ b/integrations/omnihub/validate-local.sh @@ -0,0 +1,64 @@ +#!/usr/bin/env bash +# Local validation (no sbatch) for OmniHub integration — Phase 1 checklist. +set -euo pipefail + +SCRIPT_DIR=$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd) +AI4S_REPO_ROOT="${AI4S_REPO_ROOT:-$(cd "$SCRIPT_DIR/../.." && pwd)}" +OMNIHUB_DIR="${OMNIHUB_DIR:-$HOME/git/omnihub}" +AI4S_SHARED_DIR="${AI4S_SHARED_DIR:-/shared/$USER}" + +export OMNIHUB_DIR AI4S_SHARED_DIR AI4S_REPO_ROOT + +echo "=== 1. Sync applications ===" +"$SCRIPT_DIR/sync-to-omnihub.sh" + +echo "=== 2. Cluster config bridge ===" +"$SCRIPT_DIR/render-cluster-config.sh" + +echo "=== 3. Omnistat parity ===" +"$SCRIPT_DIR/omnistat-parity-check.sh" | tee "$SCRIPT_DIR/omnistat-parity-report.txt" + +echo "=== 4. Generate SLURM script ===" +OUT="${TMPDIR:-/tmp}/hydragnn-omnihub-validate.slurm" +"$SCRIPT_DIR/generate-job.sh" --num-nodes 2 --partition lux --time-limit 2h --output "$OUT" + +grep -q 'manual-mpi' "$OUT" && echo "OK: manual-mpi runner" +grep -q 'srun --mpi=pmix' "$OUT" && echo "OK: srun --mpi=pmix" +grep -q 'ai4science-hydragnn-train' "$OUT" && echo "OK: app config path" +grep -q 'AI4S_SHARED_DIR' "$OUT" && echo "OK: HydraGNN env injection" + +if [[ -f "${AI4S_SHARED_DIR}/models/HydraGNN/overlays/hydragnn-overlay.img" ]]; then + grep -q 'hydragnn-overlay' "$OUT" && echo "OK: HydraGNN overlay in apptainer line" || echo "WARN: overlay file exists but not in SLURM script" +else + echo "SKIP: HG overlay not present at ${AI4S_SHARED_DIR}/models/HydraGNN/overlays/" +fi + +echo "=== 5. Mock omnihub-process ===" +MOCK_JOB=$(mktemp -d) +mkdir -p "$MOCK_JOB/logs" "$MOCK_JOB/tools/omnihub-monitor" +cat > "$MOCK_JOB/job.yaml" < "$MOCK_JOB/app.yaml" < "$MOCK_JOB/job-status.yaml" </dev/null || true +if [[ -f "$MOCK_JOB/processed-data/app-parser.json" ]] || ls "$MOCK_JOB/processed-data/"*.yaml &>/dev/null; then + echo "OK: omnihub-process produced processed-data/" +else + echo "OK: omnihub-process ran (partial mock — some parsers may warn)" +fi +rm -rf "$MOCK_JOB" + +echo "=== 6. Python syntax ===" +python3 -m py_compile "$SCRIPT_DIR/applications/hydragnn-train/train_wrapper.py" + +echo "" +echo "Phase 1 local validation complete." +echo "Submit on cluster: sbatch -A $OUT" +echo "Then compare training logs to sbatch_train_amd.sh baseline." diff --git a/material_science/models/HydraGNN/model.yaml b/material_science/models/HydraGNN/model.yaml index 30e546c..29fd157 100644 --- a/material_science/models/HydraGNN/model.yaml +++ b/material_science/models/HydraGNN/model.yaml @@ -25,6 +25,14 @@ recipes: recipe_path: recipes/train/ script: examples/run_train.sh sbatch_script: examples/sbatch_train_amd.sh + - task: train-omnihub + description: Multi-node GFM training via OmniHub job generator (sync + generate-job.sh) + recipe_path: integrations/omnihub/applications/hydragnn-train/ + script: integrations/omnihub/README.md + - task: perf-analysis-omnihub + description: 2-node perf run via OmniHub (omnistat + pytorch-trace + tracelens); read processed-data/ for agents + recipe_path: integrations/omnihub/applications/hydragnn-train/ + script: integrations/omnihub/README.md - task: perf-analysis description: 2-node bottleneck analysis with TraceLens + Omnistat (user-mode) via paired analyst/verifier subagents recipe_path: recipes/perf-analysis/ diff --git a/material_science/models/HydraGNN/recipes/perf-analysis/README.md b/material_science/models/HydraGNN/recipes/perf-analysis/README.md index fe4fc86..c9892fd 100644 --- a/material_science/models/HydraGNN/recipes/perf-analysis/README.md +++ b/material_science/models/HydraGNN/recipes/perf-analysis/README.md @@ -4,6 +4,18 @@ Multi-subagent workflow that launches a 2-node HydraGNN training on Lux (MI355X) > **Audience:** internal AMD performance-engineering. The output is a diagnosis of the run, not a scientific claim about HydraGNN. +## OmniHub alternative (standardized results) + +For token-efficient agent parsing via `processed-data/`, use recipe `perf-analysis-omnihub` in `model.yaml` and `.cursor/skills/ai4science-omnihub/SKILL.md`: + +```bash +./integrations/omnihub/generate-job.sh --perf --num-nodes 2 --output job.slurm +sbatch -A vultr_lux job.slurm +# then omnihub-process + omnihub-index +``` + +This uses `/shared/omnihub/tools/omnistat` and OmniHub's `tracelens` tool. The legacy path below (`sbatch_train_perf_amd.sh` + `omnistat-pr271` + subagent claims) remains for deep `omnistat-inspect` / PromQL workflows. + ## What this recipe does 1. **Launcher** subagent submits a 2-node training (`sbatch_train_perf_amd.sh`) on the `lux` partition with PyTorch profiling armed for one target epoch and Omnistat user-mode collecting on every node. From 4f88ec2d58dd9a2634dab9937deca50f33483680 Mon Sep 17 00:00:00 2001 From: Ashwin Aji Date: Fri, 29 May 2026 06:18:06 +0000 Subject: [PATCH 20/31] omnihub: default to canonical hydragnn-overlay.img MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Smoke-tested the main overlay (job 8027, completed cleanly) — the -pre-mainbump fallback was an unnecessary precaution. Default now matches sbatch_train_amd.sh. Co-authored-by: Cursor --- integrations/omnihub/generate-job.sh | 5 +---- 1 file changed, 1 insertion(+), 4 deletions(-) diff --git a/integrations/omnihub/generate-job.sh b/integrations/omnihub/generate-job.sh index 96be6e1..fadf70b 100755 --- a/integrations/omnihub/generate-job.sh +++ b/integrations/omnihub/generate-job.sh @@ -70,10 +70,7 @@ fi "$SCRIPT_DIR/sync-to-omnihub.sh" HG_BASE="$AI4S_SHARED_DIR/models/HydraGNN" -HG_OVERLAY="${HG_OVERLAY:-$HG_BASE/overlays/hydragnn-overlay-pre-mainbump.img}" -if [[ ! -f "$HG_OVERLAY" ]]; then - HG_OVERLAY="$HG_BASE/overlays/hydragnn-overlay.img" -fi +HG_OVERLAY="${HG_OVERLAY:-$HG_BASE/overlays/hydragnn-overlay.img}" HG_SIF="${HG_SIF:-$AI4S_SHARED_DIR/images/pytorch_rocm7.2.2_ubuntu24.04_py3.12_pytorch_release_2.10.0.sif}" HG_REPO="${HG_REPO_DIR:-$HG_BASE/code/HydraGNN}" HG_DATA="${HG_DATA_DIR:-$HG_BASE/weights}" From 9aa78ea9a9220575d41854eb02861fba266fbe6f Mon Sep 17 00:00:00 2001 From: Ashwin Aji Date: Fri, 29 May 2026 06:56:17 +0000 Subject: [PATCH 21/31] omnihub: drop diagnostic-only scripts from integration Remove one-off cluster diagnostics that don't belong in the recipe: render-cluster-config.sh, omnistat-parity-check.sh, its report, and validate-local.sh. Fold the only durable takeaway (OmniHub omnistat has no omnistat-inspect) into the skill doc instead. Co-authored-by: Cursor --- .cursor/skills/ai4science-omnihub/SKILL.md | 9 +- integrations/omnihub/README.md | 2 - integrations/omnihub/omnistat-parity-check.sh | 89 ------------------- .../omnihub/omnistat-parity-report.txt | 37 -------- integrations/omnihub/render-cluster-config.sh | 65 -------------- integrations/omnihub/validate-local.sh | 64 ------------- 6 files changed, 2 insertions(+), 264 deletions(-) delete mode 100755 integrations/omnihub/omnistat-parity-check.sh delete mode 100644 integrations/omnihub/omnistat-parity-report.txt delete mode 100755 integrations/omnihub/render-cluster-config.sh delete mode 100755 integrations/omnihub/validate-local.sh diff --git a/.cursor/skills/ai4science-omnihub/SKILL.md b/.cursor/skills/ai4science-omnihub/SKILL.md index 67f660e..79d5ff8 100644 --- a/.cursor/skills/ai4science-omnihub/SKILL.md +++ b/.cursor/skills/ai4science-omnihub/SKILL.md @@ -13,12 +13,7 @@ Integrates this repo with a sibling [OmniHub](https://github.com/AMDResearch/omn 2. Set `export OMNIHUB_DIR=/path/to/omnihub` and `export AI4S_SHARED_DIR=/shared/$USER` (or your site path). 3. HydraGNN assets: SIF, overlay, datasets, code clone under `$AI4S_SHARED_DIR/models/HydraGNN/`. -Validate cluster alignment: - -```bash -./integrations/omnihub/render-cluster-config.sh -./integrations/omnihub/omnistat-parity-check.sh -``` +> **Omnistat note:** OmniHub's `/shared/omnihub/tools/omnistat` has no `omnistat-inspect`, so don't expect inspect JSON from OmniHub jobs — read `omnihub-process` output under `processed-data/` instead. For deep `omnistat-inspect`/PromQL workflows use the legacy `omnistat-pr271` path (see perf-analysis recipe). ## Workflow: HydraGNN train-omnihub @@ -94,5 +89,5 @@ When working inside `$OMNIHUB_DIR`, also use: |-------|-----| | App config not found | Run `sync-to-omnihub.sh` first | | ROCm / overlay missing | Verify `HG_OVERLAY`, `HG_SIF`; check generate-job injected binds | -| Omnistat empty | Confirm `/shared/omnihub/tools/omnistat/`; run parity check | +| Omnistat empty | Confirm `/shared/omnihub/tools/omnistat/` exists and VictoriaMetrics binary is present | | TraceLens skipped | Install venv at `/shared/omnihub/tools/tracelens/venv` or set `TRACELENS_VENV` | diff --git a/integrations/omnihub/README.md b/integrations/omnihub/README.md index 6f3d80c..bef792d 100644 --- a/integrations/omnihub/README.md +++ b/integrations/omnihub/README.md @@ -37,9 +37,7 @@ $OMNIHUB_DIR/omnihub-index --results-dir /shared/$USER/results/omnihub --output integrations/omnihub/ ├── README.md # this file ├── sync-to-omnihub.sh # rsync apps → $OMNIHUB_DIR/applications/ -├── render-cluster-config.sh # validate .cluster-config vs vultr.yaml ├── generate-job.sh # sync + omnihub-generate-job + HydraGNN apptainer patches -├── omnistat-parity-check.sh # compare /shared/omnihub vs omnistat-pr271 └── applications/ └── hydragnn-train/ # OmniHub app (synced into omnihub repo) ``` diff --git a/integrations/omnihub/omnistat-parity-check.sh b/integrations/omnihub/omnistat-parity-check.sh deleted file mode 100755 index 49a83d3..0000000 --- a/integrations/omnihub/omnistat-parity-check.sh +++ /dev/null @@ -1,89 +0,0 @@ -#!/usr/bin/env bash -# Compare /shared/omnihub Omnistat installs vs ai4science omnistat-pr271 workflow. -# Run on Lux login node. Writes report to stdout; optional --output file. -set -euo pipefail - -SHARED_OMNI="${OMNIHUB_SHARED_DIR:-/shared/omnihub}" -AI4S_SHARED="${AI4S_SHARED_DIR:-/shared/$USER}" -OUT="${1:-}" - -report() { - if [[ -n "$OUT" ]]; then - tee -a "$OUT" - else - cat - fi -} - -{ - echo "# Omnistat parity check — $(date -Iseconds)" - echo "" - echo "## /shared/omnihub/tools/omnistat" - echo "" - - OMNI_VENV="$SHARED_OMNI/tools/omnistat/venv/bin" - for cmd in omnistat-usermode omnistat-query omnistat-inspect; do - if [[ -x "$OMNI_VENV/$cmd" ]]; then - echo "- $cmd: present" - else - echo "- $cmd: **MISSING**" - fi - done - - echo "" - echo "### Site config (key knobs)" - if [[ -f "$SHARED_OMNI/tools/omnistat/omnistat.config" ]]; then - cfg="$SHARED_OMNI/tools/omnistat/omnistat.config" - elif [[ -f "$SHARED_OMNI/tools/omnistat/omnihub.config" ]]; then - cfg="$SHARED_OMNI/tools/omnistat/omnihub.config" - else - cfg="" - fi - if [[ -n "$cfg" ]]; then - grep -E 'enable_rocprofiler|enable_amd_smi|enable_kernel_trace|job_detection_file|sampling_mode|counters' \ - "$cfg" || echo "(no matching keys)" - else - echo "omnistat.config / omnihub.config not found under $SHARED_OMNI/tools/omnistat/" - fi - - echo "" - echo "### Rocprofiler PMC configs" - ls "$SHARED_OMNI/tools/omnistat-rocprofiler/"*.config 2>/dev/null || ls "$SHARED_OMNI/tools/omnistat-rocprofiler/"omnihub.*.config 2>/dev/null || echo "(none found)" - - echo "" - echo "### rocprofiler-sdk Python extension" - if "$OMNI_VENV/python" -c "from omnistat.rocprofiler_sdk_extension import get_samplers" 2>/dev/null; then - echo "sdk extension: OK" - else - echo "sdk extension: failed on login node (may work on compute nodes with ROCm loaded)" - fi - - echo "" - echo "### Kernel trace library" - if [[ -f "$SHARED_OMNI/tools/omnistat/build-trace/libomnistat_trace.so" ]]; then - echo "libomnistat_trace.so: present" - else - echo "libomnistat_trace.so: not found" - fi - - echo "" - echo "## ai4science omnistat-pr271 (reference)" - PR271="$AI4S_SHARED/tools/omnistat-pr271/bin" - if [[ -d "$AI4S_SHARED/tools/omnistat-pr271" ]]; then - for cmd in omnistat-usermode omnistat-query omnistat-inspect; do - if [[ -x "$PR271/$cmd" ]]; then - echo "- $cmd: present" - else - echo "- $cmd: missing" - fi - done - else - echo "Not installed at $AI4S_SHARED/tools/omnistat-pr271 (HydraGNN launcher creates this)" - fi - - echo "" - echo "## Integration recommendation" - echo "- OmniHub path: use --tools omnistat + omnihub-process -> processed-data/" - echo "- Deep agent path (omnistat-inspect, PromQL): legacy sbatch_train_perf_amd.sh + omnistat-pr271" - echo "- If omnistat-inspect missing under /shared/omnihub, agents should not expect inspect JSON from OmniHub jobs" -} | report diff --git a/integrations/omnihub/omnistat-parity-report.txt b/integrations/omnihub/omnistat-parity-report.txt deleted file mode 100644 index 65fabc5..0000000 --- a/integrations/omnihub/omnistat-parity-report.txt +++ /dev/null @@ -1,37 +0,0 @@ -# Omnistat parity check — 2026-05-28T21:06:35+00:00 - -## /shared/omnihub/tools/omnistat - -- omnistat-usermode: present -- omnistat-query: present -- omnistat-inspect: **MISSING** - -### Site config (key knobs) -enable_amd_smi = False -enable_rocprofiler = False -enable_vendor_counters = False -job_detection_file = /tmp/omni_rmsjobinfo_omnihub - -### Rocprofiler PMC configs -/shared/omnihub/tools/omnistat-rocprofiler/omnihub.gfx90a.1.config -/shared/omnihub/tools/omnistat-rocprofiler/omnihub.gfx90a.2.config -/shared/omnihub/tools/omnistat-rocprofiler/omnihub.gfx942.1.config -/shared/omnihub/tools/omnistat-rocprofiler/omnihub.gfx942.2.config -/shared/omnihub/tools/omnistat-rocprofiler/omnihub.gfx950.1.config -/shared/omnihub/tools/omnistat-rocprofiler/omnihub.gfx950.2.config - -### rocprofiler-sdk Python extension -sdk extension: failed on login node (may work on compute nodes with ROCm loaded) - -### Kernel trace library -libomnistat_trace.so: not found - -## ai4science omnistat-pr271 (reference) -- omnistat-usermode: present -- omnistat-query: present -- omnistat-inspect: present - -## Integration recommendation -- OmniHub path: use --tools omnistat + omnihub-process -> processed-data/ -- Deep agent path (omnistat-inspect, PromQL): legacy sbatch_train_perf_amd.sh + omnistat-pr271 -- If omnistat-inspect missing under /shared/omnihub, agents should not expect inspect JSON from OmniHub jobs diff --git a/integrations/omnihub/render-cluster-config.sh b/integrations/omnihub/render-cluster-config.sh deleted file mode 100755 index de19695..0000000 --- a/integrations/omnihub/render-cluster-config.sh +++ /dev/null @@ -1,65 +0,0 @@ -#!/usr/bin/env bash -# Validate ai4science .cluster-config.yaml omnihub section against OmniHub vultr.yaml. -set -euo pipefail - -SCRIPT_DIR=$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd) -AI4S_REPO_ROOT="${AI4S_REPO_ROOT:-$(cd "$SCRIPT_DIR/../.." && pwd)}" -OMNIHUB_DIR="${OMNIHUB_DIR:-$HOME/git/omnihub}" - -CLUSTER_CFG="${AI4S_CLUSTER_CONFIG:-$AI4S_REPO_ROOT/.cluster-config.yaml}" -VULTR_YAML="$OMNIHUB_DIR/config/vultr.yaml" - -_yaml_get() { - local key=$1 file=$2 - grep -E "^[[:space:]]*${key}:" "$file" 2>/dev/null | head -1 | sed -E 's/^[^:]*:[[:space:]]*"?([^"]*)"?/\1/' | tr -d '"' -} - -echo "=== OmniHub cluster config bridge ===" -echo " ai4science config: $CLUSTER_CFG" -echo " omnihub vultr.yaml: $VULTR_YAML" -echo "" - -if [[ ! -f "$CLUSTER_CFG" ]]; then - echo "WARN: no .cluster-config.yaml — copy from .cluster-config.example.yaml" >&2 -fi - -if [[ ! -f "$VULTR_YAML" ]]; then - echo "ERROR: missing $VULTR_YAML" >&2 - exit 2 -fi - -SHARED_OMNI=$(_yaml_get shared-dir "$VULTR_YAML") -RESULTS_OMNI=$(_yaml_get results-dir "$VULTR_YAML") - -echo "OmniHub vultr.yaml:" -echo " shared-dir: $SHARED_OMNI" -echo " results-dir: $RESULTS_OMNI" -echo "" - -if [[ -f "$CLUSTER_CFG" ]]; then - REPO_DIR=$(_yaml_get repo_dir "$CLUSTER_CFG" 2>/dev/null || true) - # nested under omnihub: — grep with context - OMNI_REPO=$(awk '/^omnihub:/{f=1} f && /repo_dir:/{print; exit}' "$CLUSTER_CFG" | sed -E 's/.*:[[:space:]]*"?([^"]*)"?/\1/') - OMNI_SHARED=$(awk '/^omnihub:/{f=1} f && /shared_dir:/{print; exit}' "$CLUSTER_CFG" | sed -E 's/.*:[[:space:]]*"?([^"]*)"?/\1/') - - echo "ai4science omnihub section:" - echo " repo_dir: ${OMNI_REPO:-}" - echo " shared_dir: ${OMNI_SHARED:-}" - echo "" - - if [[ -n "${OMNI_SHARED:-}" && "$OMNI_SHARED" != "$SHARED_OMNI" ]]; then - echo "WARN: omnihub.shared_dir ($OMNI_SHARED) != vultr shared-dir ($SHARED_OMNI)" >&2 - fi -fi - -for tool_path in "$SHARED_OMNI/tools/omnistat/venv/bin/omnistat-usermode" \ - "$SHARED_OMNI/tools/omnistat-rocprofiler/venv/bin/omnistat-usermode"; do - if [[ -x "$tool_path" ]]; then - echo "OK: $tool_path" - else - echo "MISSING: $tool_path" >&2 - fi -done - -echo "" -echo "Run ./integrations/omnihub/omnistat-parity-check.sh for detailed Omnistat comparison." diff --git a/integrations/omnihub/validate-local.sh b/integrations/omnihub/validate-local.sh deleted file mode 100755 index 81b3a3a..0000000 --- a/integrations/omnihub/validate-local.sh +++ /dev/null @@ -1,64 +0,0 @@ -#!/usr/bin/env bash -# Local validation (no sbatch) for OmniHub integration — Phase 1 checklist. -set -euo pipefail - -SCRIPT_DIR=$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd) -AI4S_REPO_ROOT="${AI4S_REPO_ROOT:-$(cd "$SCRIPT_DIR/../.." && pwd)}" -OMNIHUB_DIR="${OMNIHUB_DIR:-$HOME/git/omnihub}" -AI4S_SHARED_DIR="${AI4S_SHARED_DIR:-/shared/$USER}" - -export OMNIHUB_DIR AI4S_SHARED_DIR AI4S_REPO_ROOT - -echo "=== 1. Sync applications ===" -"$SCRIPT_DIR/sync-to-omnihub.sh" - -echo "=== 2. Cluster config bridge ===" -"$SCRIPT_DIR/render-cluster-config.sh" - -echo "=== 3. Omnistat parity ===" -"$SCRIPT_DIR/omnistat-parity-check.sh" | tee "$SCRIPT_DIR/omnistat-parity-report.txt" - -echo "=== 4. Generate SLURM script ===" -OUT="${TMPDIR:-/tmp}/hydragnn-omnihub-validate.slurm" -"$SCRIPT_DIR/generate-job.sh" --num-nodes 2 --partition lux --time-limit 2h --output "$OUT" - -grep -q 'manual-mpi' "$OUT" && echo "OK: manual-mpi runner" -grep -q 'srun --mpi=pmix' "$OUT" && echo "OK: srun --mpi=pmix" -grep -q 'ai4science-hydragnn-train' "$OUT" && echo "OK: app config path" -grep -q 'AI4S_SHARED_DIR' "$OUT" && echo "OK: HydraGNN env injection" - -if [[ -f "${AI4S_SHARED_DIR}/models/HydraGNN/overlays/hydragnn-overlay.img" ]]; then - grep -q 'hydragnn-overlay' "$OUT" && echo "OK: HydraGNN overlay in apptainer line" || echo "WARN: overlay file exists but not in SLURM script" -else - echo "SKIP: HG overlay not present at ${AI4S_SHARED_DIR}/models/HydraGNN/overlays/" -fi - -echo "=== 5. Mock omnihub-process ===" -MOCK_JOB=$(mktemp -d) -mkdir -p "$MOCK_JOB/logs" "$MOCK_JOB/tools/omnihub-monitor" -cat > "$MOCK_JOB/job.yaml" < "$MOCK_JOB/app.yaml" < "$MOCK_JOB/job-status.yaml" </dev/null || true -if [[ -f "$MOCK_JOB/processed-data/app-parser.json" ]] || ls "$MOCK_JOB/processed-data/"*.yaml &>/dev/null; then - echo "OK: omnihub-process produced processed-data/" -else - echo "OK: omnihub-process ran (partial mock — some parsers may warn)" -fi -rm -rf "$MOCK_JOB" - -echo "=== 6. Python syntax ===" -python3 -m py_compile "$SCRIPT_DIR/applications/hydragnn-train/train_wrapper.py" - -echo "" -echo "Phase 1 local validation complete." -echo "Submit on cluster: sbatch -A $OUT" -echo "Then compare training logs to sbatch_train_amd.sh baseline." From 921ced7dd28b08f708d3babfede84039ed16366c Mon Sep 17 00:00:00 2001 From: Ashwin Aji Date: Fri, 29 May 2026 07:22:01 +0000 Subject: [PATCH 22/31] omnihub: point inspect references to /shared/omnihub/tools/omnihub-inspect The omnistat-inspect / TraceLens venv moved out of personal space (/shared/aaji/tools/omnistat-pr271) into the shared OmniHub tools tree as omnihub-inspect. Update the skill + READMEs and add it to the key-paths table. Co-authored-by: Cursor --- .cursor/skills/ai4science-omnihub/SKILL.md | 6 ++++-- integrations/omnihub/README.md | 2 +- .../models/HydraGNN/recipes/perf-analysis/README.md | 2 +- 3 files changed, 6 insertions(+), 4 deletions(-) diff --git a/.cursor/skills/ai4science-omnihub/SKILL.md b/.cursor/skills/ai4science-omnihub/SKILL.md index 79d5ff8..95132b1 100644 --- a/.cursor/skills/ai4science-omnihub/SKILL.md +++ b/.cursor/skills/ai4science-omnihub/SKILL.md @@ -13,7 +13,7 @@ Integrates this repo with a sibling [OmniHub](https://github.com/AMDResearch/omn 2. Set `export OMNIHUB_DIR=/path/to/omnihub` and `export AI4S_SHARED_DIR=/shared/$USER` (or your site path). 3. HydraGNN assets: SIF, overlay, datasets, code clone under `$AI4S_SHARED_DIR/models/HydraGNN/`. -> **Omnistat note:** OmniHub's `/shared/omnihub/tools/omnistat` has no `omnistat-inspect`, so don't expect inspect JSON from OmniHub jobs — read `omnihub-process` output under `processed-data/` instead. For deep `omnistat-inspect`/PromQL workflows use the legacy `omnistat-pr271` path (see perf-analysis recipe). +> **Omnistat note:** OmniHub's `/shared/omnihub/tools/omnistat` has no `omnistat-inspect`, so don't expect inspect JSON from OmniHub jobs — read `omnihub-process` output under `processed-data/` instead. For deep `omnistat-inspect`/PromQL workflows use the `omnihub-inspect` build at `/shared/omnihub/tools/omnihub-inspect` (see perf-analysis recipe). ## Workflow: HydraGNN train-omnihub @@ -60,7 +60,7 @@ sbatch -A vultr_lux /tmp/hydragnn-perf.slurm 3. `/processed-data/report-card.yaml` 4. Raw traces under `/tools/` only if needed -Legacy HydraGNN perf (`sbatch_train_perf_amd.sh` + `omnistat-pr271` + subagent claims) remains for deep PromQL/inspect workflows. See `material_science/models/HydraGNN/recipes/perf-analysis/README.md`. +Legacy HydraGNN perf (`sbatch_train_perf_amd.sh` + `omnihub-inspect` + subagent claims) remains for deep PromQL/inspect workflows. See `material_science/models/HydraGNN/recipes/perf-analysis/README.md`. ## Key paths (Lux / vultr) @@ -68,6 +68,8 @@ Legacy HydraGNN perf (`sbatch_train_perf_amd.sh` + `omnistat-pr271` + subagent c |----------|------| | OmniHub shared | `/shared/omnihub` | | Omnistat | `/shared/omnihub/tools/omnistat/` → use `--tools omnistat` | +| Omnistat + PMC | `/shared/omnihub/tools/omnistat-rocprofiler/` → use `--tools omnistat-rocprofiler-pmc1` for hardware counters | +| Omnistat inspect | `/shared/omnihub/tools/omnihub-inspect/` → has `omnistat-inspect` + TraceLens CLI (legacy deep-dive venv) | | Results | `/shared/$USER/results/omnihub/$SLURM_JOB_ID/` | | HydraGNN overlay | `$AI4S_SHARED_DIR/models/HydraGNN/overlays/hydragnn-overlay.img` | diff --git a/integrations/omnihub/README.md b/integrations/omnihub/README.md index bef792d..9ca63b7 100644 --- a/integrations/omnihub/README.md +++ b/integrations/omnihub/README.md @@ -73,4 +73,4 @@ Agents should read `processed-data/` first (token-efficient). See `.cursor/skill | `train-omnihub` | `omnihub-monitor` | Phase 1 pilot | | `perf-analysis-omnihub` | `omnihub-monitor omnistat pytorch-trace tracelens` | Phase 2; optional `omnistat-rocprofiler-pmc1` | -Legacy HydraGNN perf (`sbatch_train_perf_amd.sh` + `omnistat-pr271`) remains available for deep agent workflows. +Legacy HydraGNN perf (`sbatch_train_perf_amd.sh` + `omnihub-inspect` at `/shared/omnihub/tools/omnihub-inspect`) remains available for deep agent workflows. diff --git a/material_science/models/HydraGNN/recipes/perf-analysis/README.md b/material_science/models/HydraGNN/recipes/perf-analysis/README.md index c9892fd..4b6247a 100644 --- a/material_science/models/HydraGNN/recipes/perf-analysis/README.md +++ b/material_science/models/HydraGNN/recipes/perf-analysis/README.md @@ -14,7 +14,7 @@ sbatch -A vultr_lux job.slurm # then omnihub-process + omnihub-index ``` -This uses `/shared/omnihub/tools/omnistat` and OmniHub's `tracelens` tool. The legacy path below (`sbatch_train_perf_amd.sh` + `omnistat-pr271` + subagent claims) remains for deep `omnistat-inspect` / PromQL workflows. +This uses `/shared/omnihub/tools/omnistat` and OmniHub's `tracelens` tool. The legacy path below (`sbatch_train_perf_amd.sh` + the `omnihub-inspect` venv at `/shared/omnihub/tools/omnihub-inspect` + subagent claims) remains for deep `omnistat-inspect` / PromQL workflows. ## What this recipe does From a570e667cfcb301d832986f2359108b0a0efa815 Mon Sep 17 00:00:00 2001 From: Ashwin Aji Date: Fri, 29 May 2026 07:34:36 +0000 Subject: [PATCH 23/31] omnihub: make cluster tool paths config-driven via omnihub.tools_dir Add omnihub.tools_dir (+ documented omnistat/omnistat-rocprofiler/omnihub-inspect subdirs) to the cluster-config example so site-specific tool locations live in config, not in recipes. Point the skill's key-paths at the config keys instead of hardcoded literals. Co-authored-by: Cursor --- .cluster-config.example.yaml | 6 ++++++ .cursor/skills/ai4science-omnihub/SKILL.md | 16 +++++++++------- 2 files changed, 15 insertions(+), 7 deletions(-) diff --git a/.cluster-config.example.yaml b/.cluster-config.example.yaml index 4759a52..144558b 100644 --- a/.cluster-config.example.yaml +++ b/.cluster-config.example.yaml @@ -48,3 +48,9 @@ omnihub: shared_dir: /shared/omnihub results_dir: "" # default: /shared/$USER/results/omnihub rocm_version: "7.2.2" + # Shared perf/telemetry tools. Cluster-specific — set per site, do not hardcode + # these into recipes/scripts. Default: /tools. + tools_dir: "" # default: /shared/omnihub/tools + # /omnistat GPU telemetry, no PMC (--tools omnistat) + # /omnistat-rocprofiler adds rocprofiler PMC counters (--tools omnistat-rocprofiler-pmc1) + # /omnihub-inspect omnistat-inspect + TraceLens venv (legacy deep-dive) diff --git a/.cursor/skills/ai4science-omnihub/SKILL.md b/.cursor/skills/ai4science-omnihub/SKILL.md index 95132b1..96428f3 100644 --- a/.cursor/skills/ai4science-omnihub/SKILL.md +++ b/.cursor/skills/ai4science-omnihub/SKILL.md @@ -62,15 +62,17 @@ sbatch -A vultr_lux /tmp/hydragnn-perf.slurm Legacy HydraGNN perf (`sbatch_train_perf_amd.sh` + `omnihub-inspect` + subagent claims) remains for deep PromQL/inspect workflows. See `material_science/models/HydraGNN/recipes/perf-analysis/README.md`. -## Key paths (Lux / vultr) +## Key paths -| Resource | Path | +Cluster-specific paths come from `.cluster-config.yaml` (`omnihub.shared_dir`, `omnihub.tools_dir`, `omnihub.results_dir`). Resolve them from there rather than assuming literals. The values below are the **Lux / vultr** defaults (`shared_dir: /shared/omnihub`, `tools_dir: /shared/omnihub/tools`). + +| Resource | Path (`` = `omnihub.tools_dir`) | |----------|------| -| OmniHub shared | `/shared/omnihub` | -| Omnistat | `/shared/omnihub/tools/omnistat/` → use `--tools omnistat` | -| Omnistat + PMC | `/shared/omnihub/tools/omnistat-rocprofiler/` → use `--tools omnistat-rocprofiler-pmc1` for hardware counters | -| Omnistat inspect | `/shared/omnihub/tools/omnihub-inspect/` → has `omnistat-inspect` + TraceLens CLI (legacy deep-dive venv) | -| Results | `/shared/$USER/results/omnihub/$SLURM_JOB_ID/` | +| OmniHub shared | `omnihub.shared_dir` (Lux: `/shared/omnihub`) | +| Omnistat | `/omnistat/` → use `--tools omnistat` | +| Omnistat + PMC | `/omnistat-rocprofiler/` → use `--tools omnistat-rocprofiler-pmc1` for hardware counters | +| Omnistat inspect | `/omnihub-inspect/` → has `omnistat-inspect` + TraceLens CLI (legacy deep-dive venv) | +| Results | `omnihub.results_dir` (Lux: `/shared/$USER/results/omnihub/$SLURM_JOB_ID/`) | | HydraGNN overlay | `$AI4S_SHARED_DIR/models/HydraGNN/overlays/hydragnn-overlay.img` | ## Runner From 4977c59e514ba5aad788c31395d145085fc754a9 Mon Sep 17 00:00:00 2001 From: Ashwin Aji Date: Fri, 29 May 2026 18:49:12 +0000 Subject: [PATCH 24/31] omnihub: document OMNIHUB_TOOLS_DIR env var mapping for omnihub.tools_dir Note in the cluster-config example that legacy HydraGNN perf recipes consume omnihub.tools_dir via the OMNIHUB_TOOLS_DIR env var. Co-authored-by: Cursor --- .cluster-config.example.yaml | 2 ++ 1 file changed, 2 insertions(+) diff --git a/.cluster-config.example.yaml b/.cluster-config.example.yaml index 144558b..68a342b 100644 --- a/.cluster-config.example.yaml +++ b/.cluster-config.example.yaml @@ -50,6 +50,8 @@ omnihub: rocm_version: "7.2.2" # Shared perf/telemetry tools. Cluster-specific — set per site, do not hardcode # these into recipes/scripts. Default: /tools. + # The legacy HydraGNN perf recipes read this value via the OMNIHUB_TOOLS_DIR + # env var (export OMNIHUB_TOOLS_DIR= before running them). tools_dir: "" # default: /shared/omnihub/tools # /omnistat GPU telemetry, no PMC (--tools omnistat) # /omnistat-rocprofiler adds rocprofiler PMC counters (--tools omnistat-rocprofiler-pmc1) From 51487e5458494bff86a60a1d934adb316bda0232 Mon Sep 17 00:00:00 2001 From: Ashwin Aji Date: Fri, 29 May 2026 23:28:23 +0000 Subject: [PATCH 25/31] HydraGNN perf: finish public-repo scrub + placeholder hygiene Code-review cleanup ahead of the perf-optimization PR. The earlier "scrub internal names" pass missed the perf-analysis recipe, the two SKILL files, and run_fom_extractor.py; complete it here. - Scrub site-specific identifiers from all perf files: username/home paths (/shared/aaji, /home/aaji), cluster/partition/account names (lux, vultr_lux, rad-vultr-login), and node names (lux-mi355x-*). Replace with $AI4S_SHARED_DIR / $REPO_ROOT / / placeholders and generic MI355X (gfx950) wording. - Rename omnistat-lux.config.template -> omnistat.config.template; templatize victoria_binary via @OMNIHUB_TOOLS_DIR@ and substitute it in sbatch_train_perf_amd.sh's render step. system_name = mi355x. - sbatch_train_perf_amd.sh + microbench_node_health.sh: use the repo's YOUR_GPU_PARTITION / YOUR_ACCOUNT placeholder convention (matches sbatch_train_amd.sh) instead of gpu/default; drop hardcoded /shared. - run_fom_extractor.py: drop unused `shared` var and the hardcoded node-name fallback in the empty-instance correlation row. - parse_convergence.py + train/README.md: align stale SHA reference and fix a leftover /shared/aaji path. Co-authored-by: Cursor --- .../skills/ai4science-perf-analysis/SKILL.md | 4 +- .cursor/skills/ai4science-studio/SKILL.md | 32 +++++++------- .../examples/microbench_node_health.sh | 10 ++--- .../HydraGNN/examples/parse_convergence.py | 2 +- .../HydraGNN/examples/run_fom_extractor.py | 7 ++- .../examples/sbatch_train_perf_amd.sh | 17 ++++--- .../HydraGNN/recipes/perf-analysis/README.md | 18 ++++---- .../recipes/perf-analysis/agents/launcher.md | 44 +++++++++---------- .../perf-analysis/agents/omnistat_analyst.md | 2 +- .../perf-analysis/agents/omnistat_verifier.md | 4 +- .../perf-analysis/agents/tracelens_analyst.md | 6 +-- .../agents/tracelens_verifier.md | 6 +-- ...nfig.template => omnistat.config.template} | 18 ++++---- .../agents/orchestrator.md | 6 +-- .../dispatch-attribution.md | 8 ++-- .../models/HydraGNN/recipes/train/README.md | 2 +- 16 files changed, 94 insertions(+), 92 deletions(-) rename material_science/models/HydraGNN/recipes/perf-analysis/{omnistat-lux.config.template => omnistat.config.template} (90%) diff --git a/.cursor/skills/ai4science-perf-analysis/SKILL.md b/.cursor/skills/ai4science-perf-analysis/SKILL.md index f848be3..0c27543 100644 --- a/.cursor/skills/ai4science-perf-analysis/SKILL.md +++ b/.cursor/skills/ai4science-perf-analysis/SKILL.md @@ -40,7 +40,7 @@ Default target: HydraGNN on AMD MI355X. The same pattern can extend to ORBIT-2 / - Never escalate to multi-node `srun`; verifiers may use **at most one** 1-node `srun -N1 --time=00:05:00` interactive probe (partition/account from `.cluster-config.yaml`). - Never edit files outside `$AI4S_SHARED_DIR/models/HydraGNN/perf-runs//`. -## Cluster constraints (Lux) +## Cluster constraints (AMD Instinct MI355X) - Partition / account: read from [.cluster-config.yaml](../../../.cluster-config.yaml). - Central Prometheus may be **unreachable from compute nodes** (check your cluster's network policy). Always use the user-mode VictoriaMetrics that the launcher started. @@ -58,7 +58,7 @@ Default target: HydraGNN on AMD MI355X. The same pattern can extend to ORBIT-2 / ## Lessons captured (smoke-test on JOB 6762) -The first 2-node end-to-end run on Lux exposed five real bugs that are now all worked around in [`sbatch_train_perf_amd.sh`](../../../material_science/models/HydraGNN/examples/sbatch_train_perf_amd.sh) and the agent prompts: +The first 2-node end-to-end run exposed five real bugs that are now all worked around in [`sbatch_train_perf_amd.sh`](../../../material_science/models/HydraGNN/examples/sbatch_train_perf_amd.sh) and the agent prompts: 1. **`%(SLURM_JOB_ID)s` in omnistat config is broken** — `configparser` doesn't interpolate `os.environ`. Use `@JOB_DIR@` placeholder and `sed` at submit time. 2. **`/tmp/omni_rmsjobinfo` permission collision with system-mode Omnistat** — override `job_detection_file` to a per-job path under `$AI4S_SHARED_DIR/...`. diff --git a/.cursor/skills/ai4science-studio/SKILL.md b/.cursor/skills/ai4science-studio/SKILL.md index f517112..daaaf40 100755 --- a/.cursor/skills/ai4science-studio/SKILL.md +++ b/.cursor/skills/ai4science-studio/SKILL.md @@ -399,7 +399,7 @@ setup_ddp (hydragnn/utils/distributed/distributed.py:254) 3. **Print the verify info on rank 0 from the launching shell only**, never from a python that imports mpi4py-bearing modules. -This bit us once on the HydraGNN perf-analysis recipe (see git log on `aaji/perf-agents-hydragnn`). Watch for it whenever you add `python3 -c "..."` scaffolding around an `srun` that uses `--mpi=pmix`. +This bit us once on the HydraGNN perf-analysis recipe. Watch for it whenever you add `python3 -c "..."` scaffolding around an `srun` that uses `--mpi=pmix`. --- @@ -532,7 +532,7 @@ This **does not work** with stock omnistat-usermode because `omnistat.utils.read **Fix:** Use a `@JOB_DIR@`-style placeholder in the template and run `sed` from the sbatch wrapper at submit time: ```bash -sed -e "s|@JOB_DIR@|${HG_OUTPUT_DIR}|g" omnistat-lux.config.template > $HG_OUTPUT_DIR/omnistat.config +sed -e "s|@JOB_DIR@|${HG_OUTPUT_DIR}|g" omnistat.config.template > $HG_OUTPUT_DIR/omnistat.config ``` ### 2. omnistat-usermode `/tmp/omni_rmsjobinfo` permission collision @@ -543,15 +543,15 @@ If the cluster runs system-mode Omnistat as user `omnidc` (or any non-root, non- ```ini [omnistat.collectors.rms] job_detection_mode = file-based -job_detection_file = /shared/aaji/.../omni_rmsjobinfo -step_detection_file = /shared/aaji/.../omni_rmsjobinfo_step +job_detection_file = $AI4S_SHARED_DIR/.../omni_rmsjobinfo +step_detection_file = $AI4S_SHARED_DIR/.../omni_rmsjobinfo_step ``` Path must be on a shared FS reachable by every compute node in the allocation (NFS/Lustre/etc.) because each node's `omnistat-rms-env` runs via `srun -N $NODES --ntasks-per-node=1` and writes to that path. ### 3. VictoriaMetrics needs `-fs.disableMmap` on the login node -Lux login node `vm.max_map_count` is 1048576 (high), but the cgroup memory limit is 8 GB, and even the small (1.6 MB) per-job Omnistat DB cannot mmap successfully — VM panics with: +A login node may have a high `vm.max_map_count` (e.g. 1048576) but a tight cgroup memory limit (e.g. 8 GB), and even the small (1.6 MB) per-job Omnistat DB cannot mmap successfully — VM panics with: ``` FATAL: cannot mmap "/.../index.bin": cannot mmap file with size 4096 bytes; already memory mapped files: 0: no such device ``` @@ -600,7 +600,7 @@ This catches both tool bugs and analyst hallucinations. Refuted claims stay in ` ### 8. Never read configs from a `/shared` "staged copy" — always from the in-repo template -We hit this once: the omnistat config under `/shared/aaji/models/HydraGNN/perf-runs/omnistat-lux.config` was a stale copy of the in-repo `omnistat-lux.config.template`, with `enable_rocprofiler = False` left over from an early v1 decision. The in-repo template had since been updated to `enable_rocprofiler = True` with the `hbm_flops_f64` profile, but the staged copy was never refreshed, so two perf-analysis runs (jobs 6762 and 6840) silently collected zero hardware counters (no HBM bandwidth, no fp64 VALU/MFMA FLOPs). The TraceLens analytical roofline still ran fine, masking the gap until someone asked "where are the HBM counters?" +We hit this once: the omnistat config under `$AI4S_SHARED_DIR/models/HydraGNN/perf-runs/omnistat.config` was a stale copy of the in-repo `omnistat.config.template`, with `enable_rocprofiler = False` left over from an early v1 decision. The in-repo template had since been updated to `enable_rocprofiler = True` with the `hbm_flops_f64` profile, but the staged copy was never refreshed, so two perf-analysis runs silently collected zero hardware counters (no HBM bandwidth, no fp64 VALU/MFMA FLOPs). The TraceLens analytical roofline still ran fine, masking the gap until someone asked "where are the HBM counters?" Two compounding mistakes: 1. The sbatch script defaulted `OMNISTAT_TEMPLATE` to the `/shared` staged path, not the in-repo path. @@ -637,7 +637,7 @@ ROCm rocprofiler-sdk on gfx950 returns ``` create profile failed with error code 38: Request exceeds the capabilities of the hardware to collect ``` -when a single profile asks for more counters than fit in the per-block register file. Per-block limits are not documented uniformly; we determined empirically on `lux-mi355x-b1` (ROCm 7.2.2): +when a single profile asks for more counters than fit in the per-block register file. Per-block limits are not documented uniformly; we determined empirically on MI355X (gfx950, ROCm 7.2.2): - **TCC (L2 cache) block**: **`FETCH_SIZE` AND `WRITE_SIZE` together fail.** Either one alone works (with up to 4 SQ_INSTS_VALU_*_F64 counters in the same profile). - **SQ (scalar/vector unit) block**: ≥4 SQ_INSTS_VALU_*_F64 counters in one profile work. @@ -647,7 +647,7 @@ Always probe interactively first (1 node `salloc`, run `omnistat-monitor --confi ### 11. Omnistat `sampling_mode = periodic` only fires the first set in this build -In `omnistat 1.12.0+git.65ea9ac` (PR #271 + main merged) on `lux-mi355x-b1` with rocprofiler-sdk freshly built, `sampling_mode = periodic` with two counter sets only fires the FIRST set. Every subsequent set returns identically zero across all cards and timestamps. Half the intended counter coverage is silently lost. +In `omnistat 1.12.0+git.65ea9ac` (PR #271 + main merged) on MI355X (gfx950) with rocprofiler-sdk freshly built, `sampling_mode = periodic` with two counter sets only fires the FIRST set. Every subsequent set returns identically zero across all cards and timestamps. Half the intended counter coverage is silently lost. We surfaced this with the `[[FETCH+4×SQ_F64], [WRITE+4×SQ_F64]]` config: FETCH and MFMA_MOPS samples were healthy (~50% non-zero, ~7% non-zero respectively), but WRITE and FMA/ADD/MUL were uniformly zero. Workaround: use `sampling_mode = constant` and a single counter set that fits in the per-block hardware limits — for fp64 + HBM-read on gfx950 that's: ```ini @@ -663,7 +663,7 @@ Omnistat exposes **two** independent rocprofiler-sdk integrations, and there's a | Surface | What you get | Cardinality / cost | Build artifact | Switch in our stack | |---|---|---|---|---| -| **Device counting** (omnistat collector, `enable_rocprofiler=True`) | Whole-card hardware counters: FETCH_SIZE, fp64 VALU/MFMA, etc. Sampled at 1 Hz, scraped by the omnistat exporter. | One time-series per (card, counter). Constant-mode is essentially free. | `omnistat/rocprofiler_sdk_extension.cpython-*.so` (Python binding, built in-place from `rocprofiler-sdk/CMakeLists.txt` with `BUILD_KERNEL_TRACE_LIB=OFF`) | `enable_rocprofiler = True` in `omnistat-lux.config.template` (default). | +| **Device counting** (omnistat collector, `enable_rocprofiler=True`) | Whole-card hardware counters: FETCH_SIZE, fp64 VALU/MFMA, etc. Sampled at 1 Hz, scraped by the omnistat exporter. | One time-series per (card, counter). Constant-mode is essentially free. | `omnistat/rocprofiler_sdk_extension.cpython-*.so` (Python binding, built in-place from `rocprofiler-sdk/CMakeLists.txt` with `BUILD_KERNEL_TRACE_LIB=OFF`) | `enable_rocprofiler = True` in `omnistat.config.template` (default). | | **Kernel tracing** (omnistat collector, `enable_kernel_trace=True`) | Per-kernel-name dispatch count and total duration on every card, every node. Aggregated in 1 s bins, POSTed via HTTP to the omnistat exporter. | One time-series per (card, kernel-name) — can be 1k+ unique kernel names per training step. Light overhead per dispatch (microseconds), but cardinality blows up VictoriaMetrics if left on for hours. | `build-trace/libomnistat_trace.so` (standalone C++ tool library, built from same CMake with `-DBUILD_KERNEL_TRACE_LIB=ON`) | `OMNISTAT_KERNEL_TRACE=1` env-gate in `sbatch_train_perf_amd.sh` — flips `enable_kernel_trace=True` in the rendered config and exports `ROCP_TOOL_LIBRARIES=` into the apptainer exec. | | **`rocprofv3` standalone trace** (no omnistat) | Full kernel + memcpy + HSA API trace on a single rank. JSON / Perfetto / CSV output you read with TraceLens or chrome://tracing. | One trace file per rank, multi-MB per second of training. Heaviest of the three. | None — uses `/opt/rocm/bin/rocprofv3` directly. | Wrap a single rank's command in `rocprofv3 --kernel-trace -o ./rank0_trace -- python ...`. We don't run this in our standard sbatch yet. | @@ -678,9 +678,9 @@ Omnistat exposes **two** independent rocprofiler-sdk integrations, and there's a **Build location:** Both omnistat artifacts (the Python extension and the kernel-trace tool library) must be built on a **compute node** inside the same SIF the workload runs in — login nodes don't have apptainer or the ROCm headers, and a host-side build will pick up the wrong libcurl / glibc. `OMNIHUB_TOOLS_DIR` is the shared tools dir from `.cluster-config.yaml` `omnihub.tools_dir` (e.g. `/shared/omnihub/tools`); export it first. Standard build: ```bash -salloc -p lux -A vultr_lux -N 1 --time=00:15:00 --gpus-per-node=1 \ +salloc -p -A -N 1 --time=00:15:00 --gpus-per-node=1 \ apptainer exec --rocm \ - /shared/aaji/images/pytorch_rocm7.2.2_ubuntu24.04_py3.12_pytorch_release_2.10.0.sif \ + "${AI4S_SHARED_DIR}/images/pytorch_rocm7.2.2_ubuntu24.04_py3.12_pytorch_release_2.10.0.sif" \ bash -c "cd ${OMNIHUB_TOOLS_DIR}/omnistat-src && \ cmake -S rocprofiler-sdk/ -B build-trace/ -DBUILD_KERNEL_TRACE_LIB=ON && \ cmake --build build-trace/ -j 8" @@ -688,13 +688,13 @@ salloc -p lux -A vultr_lux -N 1 --time=00:15:00 --gpus-per-node=1 \ Verify with `ls -la ${OMNIHUB_TOOLS_DIR}/omnistat-src/build-trace/libomnistat_trace.so` (~MB-class file, not the tiny 4 kB stub you get when CMake fails halfway). -**Lux SIF build prerequisites (one-time gotchas):** the `pytorch_rocm7.2.2_ubuntu24.04_py3.12_pytorch_release_2.10.0.sif` ships **without** `cmake` and **without** libcurl headers (only the `.so.4` runtime). Both omnistat artifacts need cmake; the kernel-trace library additionally needs ``. Working build recipe inside the SIF: +**SIF build prerequisites (one-time gotchas):** the `pytorch_rocm7.2.2_ubuntu24.04_py3.12_pytorch_release_2.10.0.sif` ships **without** `cmake` and **without** libcurl headers (only the `.so.4` runtime). Both omnistat artifacts need cmake; the kernel-trace library additionally needs ``. Working build recipe inside the SIF: 1. `python -m venv /tmp/build_venv && /tmp/build_venv/bin/pip install cmake nanobind zstandard` — installs `cmake` 4.x as a wheel, plus `zstandard` for unpacking modern Ubuntu `.deb` files (which use `data.tar.zst`). 2. Fetch `libcurl4-openssl-dev_*.deb` from `archive.ubuntu.com/ubuntu/pool/main/c/curl/`, `ar x` it, decompress `data.tar.zst` with `zstandard.ZstdDecompressor` in Python, untar to a tmp prefix. 3. Pass `-DCURL_INCLUDE_DIR=/usr/include/x86_64-linux-gnu -DCURL_LIBRARY=/usr/lib/x86_64-linux-gnu/libcurl.so.4` to cmake configure. The library finds rocprofiler-sdk via `/opt/rocm/lib/cmake/rocprofiler-sdk/`. -Verified on `lux-mi355x-b1` job 7030: 533 kB `libomnistat_trace.so`, ELF64, all rocprofiler-sdk symbols resolved, env vars `OMNISTAT_TRACE_{MAX_INTERVAL,BUFFER_SIZE,ENDPOINT_PORT,LOG}` baked in. `apt-get download` does **not** work in the SIF — apt sources aren't configured — so don't waste a node on it. +Verified on MI355X (gfx950): 533 kB `libomnistat_trace.so`, ELF64, all rocprofiler-sdk symbols resolved, env vars `OMNISTAT_TRACE_{MAX_INTERVAL,BUFFER_SIZE,ENDPOINT_PORT,LOG}` baked in. `apt-get download` does **not** work in the SIF — apt sources aren't configured — so don't waste a node on it. **Endpoint port alignment:** the kernel-trace **collector** registers `/kernel_trace` on the same Flask app as the rest of the omnistat exporter, so it listens on `[omnistat.collectors] port` (8101 in our config). The kernel-trace **tool library**'s default `OMNISTAT_TRACE_ENDPOINT_PORT` is 8001, which is wrong for our setup — they must match or every dispatch record gets a connection-refused. The sbatch wrapper auto-derives the port from the rendered config; do not hard-code 8001. @@ -725,14 +725,14 @@ srun ... apptainer exec ... "${OPT_ENVS[@]}" "$HG_SIF" bash "$RANK_SCRIPT" Applied to both `sbatch_train_amd.sh` and `sbatch_train_perf_amd.sh` for HydraGNN. **General lesson** for any model recipe: never use `--env KEY="${KEY:-}"` for variables that the workload reads with `os.getenv(K) is not None` — either pass a numeric default (`--env KEY="${KEY:-0}"`) or use the conditional-array pattern above. Audit every `*/models/*/examples/*.sh` against this when introducing new optional env knobs. -### 14. Cluster-state checks before queueing build/probe jobs on Lux +### 14. Cluster-state checks before queueing build/probe jobs -`sinfo -p lux,rad` is the cheapest way to know whether a build alloc will ever start. Lux and rad partitions get drained for slurm maintenance roughly every few weeks — the symptom is `salloc: job NNNN queued and waiting for resources` followed by `(Nodes required for job are DOWN, DRAINED or reserved for jobs in higher priority partitions)`. The drain reasons are visible via `sinfo -R`; treat any `slurm deploy maintenance`-class reason as "wait, don't queue". Cancel pending allocs with `scancel` when you discover the drain — they don't time out on their own. +`sinfo -p ` is the cheapest way to know whether a build alloc will ever start. Partitions get drained for slurm maintenance roughly every few weeks — the symptom is `salloc: job NNNN queued and waiting for resources` followed by `(Nodes required for job are DOWN, DRAINED or reserved for jobs in higher priority partitions)`. The drain reasons are visible via `sinfo -R`; treat any `slurm deploy maintenance`-class reason as "wait, don't queue". Cancel pending allocs with `scancel` when you discover the drain — they don't time out on their own. When the cluster is up but you only need a one-shot interactive build inside the SIF, **prefer `srun --no-shell`-equivalent direct exec over `salloc + apptainer exec`**: ```bash -srun -p lux -A vultr_lux -N 1 --time=00:15:00 --gpus-per-node=1 --cpus-per-task=8 \ +srun -p -A -N 1 --time=00:15:00 --gpus-per-node=1 --cpus-per-task=8 \ --job-name= \ apptainer exec --rocm \ --bind ${OMNIHUB_TOOLS_DIR}:${OMNIHUB_TOOLS_DIR} \ diff --git a/material_science/models/HydraGNN/examples/microbench_node_health.sh b/material_science/models/HydraGNN/examples/microbench_node_health.sh index d2f2927..a242de0 100755 --- a/material_science/models/HydraGNN/examples/microbench_node_health.sh +++ b/material_science/models/HydraGNN/examples/microbench_node_health.sh @@ -33,8 +33,8 @@ # # SBATCH directives: #SBATCH --job-name=node-health -#SBATCH --partition=gpu -#SBATCH --account=default +#SBATCH --partition=YOUR_GPU_PARTITION +#SBATCH --account=YOUR_ACCOUNT #SBATCH --nodes=1 #SBATCH --ntasks=1 #SBATCH --cpus-per-task=128 @@ -69,8 +69,8 @@ echo "[$H] microbench start at $(date -u +%FT%TZ)" | tee "$HOST_OUT/_start.txt" cat /sys/devices/system/cpu/cpu0/cpufreq/scaling_governor > "$HOST_OUT/cpu_governor.txt" 2>&1 || echo "no cpufreq" > "$HOST_OUT/cpu_governor.txt" grep -E '^Mem|HugePages|^DirectMap|^Swap' /proc/meminfo > "$HOST_OUT/meminfo.txt" 2>&1 ip -br link > "$HOST_OUT/ip_link.txt" 2>&1 - mount | grep -E ' /home| /shared|tmpfs|/tmp' > "$HOST_OUT/mounts.txt" 2>&1 - df -h $HOME /shared /tmp 2>/dev/null > "$HOST_OUT/df.txt" 2>&1 + mount | grep -E " /home| ${AI4S_SHARED_DIR}|tmpfs|/tmp" > "$HOST_OUT/mounts.txt" 2>&1 + df -h "$HOME" "$AI4S_SHARED_DIR" /tmp 2>/dev/null > "$HOST_OUT/df.txt" 2>&1 } # ---------- Test 2: container mount probe ---------- @@ -163,7 +163,7 @@ HIP_PASS="skipped(no_gpu_allocation)" HIP_MEAN_US="" if command -v rocm-smi >/dev/null 2>&1 || [ -e "$HG_SIF" ]; then if [ -e "$HG_SIF" ]; then - apptainer exec --rocm --bind /opt/rocm-7.2.2:/opt/rocm-7.2.2 --bind /shared \ + apptainer exec --rocm --bind /opt/rocm-7.2.2:/opt/rocm-7.2.2 --bind "$AI4S_SHARED_DIR" \ --env LD_LIBRARY_PATH=/opt/rocm-7.2.2/lib:/usr/lib/x86_64-linux-gnu \ "$HG_SIF" bash -c " PATH=/opt/rocm-7.2.2/bin:\$PATH rocminfo > '$HOST_OUT/rocminfo.txt' 2>&1 diff --git a/material_science/models/HydraGNN/examples/parse_convergence.py b/material_science/models/HydraGNN/examples/parse_convergence.py index 0aa558a..03d71d1 100755 --- a/material_science/models/HydraGNN/examples/parse_convergence.py +++ b/material_science/models/HydraGNN/examples/parse_convergence.py @@ -26,7 +26,7 @@ def parse_slurm_log(log_path: str) -> list[dict]: """Parse HydraGNN training metrics from SLURM stdout or run.log. - HydraGNN (pinned SHA 6c45f168) emits these patterns: + HydraGNN (pinned SHA 2fb0bd0) emits these patterns: - Epoch summary (requires HYDRAGNN_VALTEST=1): 0: Epoch: 01, Train Loss: 0.12345678, Val Loss: 0.23456789, Test Loss: 0.34567890 - Per-task losses: diff --git a/material_science/models/HydraGNN/examples/run_fom_extractor.py b/material_science/models/HydraGNN/examples/run_fom_extractor.py index e4111ef..6ec0d42 100755 --- a/material_science/models/HydraGNN/examples/run_fom_extractor.py +++ b/material_science/models/HydraGNN/examples/run_fom_extractor.py @@ -5,9 +5,9 @@ for login-node post-processing of a completed perf run. Usage: - export AI4S_SHARED_DIR=/shared/aaji + export AI4S_SHARED_DIR=/shared/$USER export OMNIHUB_TOOLS_DIR=/shared/omnihub/tools # from .cluster-config.yaml omnihub.tools_dir - export REPO_ROOT=/home/aaji/git/ai4science-studio + export REPO_ROOT=/path/to/ai4science-studio python run_fom_extractor.py --manifest /path/to/perf-runs//manifest.json \\ [--vm-port 8432] [--tsdb-url http://127.0.0.1:8432] @@ -202,7 +202,6 @@ def main() -> int: parser.add_argument("--no-start-vm", action="store_true", help="Require --tsdb-url") args = parser.parse_args() - shared = os.environ.get("AI4S_SHARED_DIR", "/shared/aaji") repo = os.environ.get("REPO_ROOT", str(Path(__file__).resolve().parents[3])) with open(args.manifest) as f: @@ -300,7 +299,7 @@ def main() -> int: { "window_start_iso": iso, "window_start_ts": win_ts, - "instance": "lux-mi355x-b1", + "instance": "", "gpu_id": 0, "top_kernel_name": top_name, "top_kernel_busy_frac": f"{top_frac:.4f}", diff --git a/material_science/models/HydraGNN/examples/sbatch_train_perf_amd.sh b/material_science/models/HydraGNN/examples/sbatch_train_perf_amd.sh index df2c55f..313c34f 100755 --- a/material_science/models/HydraGNN/examples/sbatch_train_perf_amd.sh +++ b/material_science/models/HydraGNN/examples/sbatch_train_perf_amd.sh @@ -8,8 +8,9 @@ # 1. Wraps srun with omnistat-usermode --start/--stopexporters/--stopserver. # 2. Generates a per-job copy of gfm_mlip.json with a "Profile" block injected # so HydraGNN's built-in torch.profiler captures rank-0 of node-0 only. -# 3. Defaults to a cluster partition/account that can be overridden via -# SBATCH_PARTITION/SBATCH_ACCOUNT env vars for your cluster. +# 3. Uses placeholder partition/account (#SBATCH --partition=YOUR_GPU_PARTITION, +# --account=YOUR_ACCOUNT); override via SBATCH_PARTITION/SBATCH_ACCOUNT +# env vars or by editing the directives for your cluster. # 4. Defaults are tuned for a quick perf run: --nodes=2, NUM_EPOCH=2, # MAX_NUM_BATCH=30, time=00:30:00 — enough for the wait=5/warmup=3/active=3 # profiler schedule plus a clean epoch 0 for warm caches. @@ -32,8 +33,8 @@ # OMNISTAT_USERMODE_INTERVAL=1 # seconds #SBATCH --job-name=hydragnn-perf -#SBATCH --partition=gpu -#SBATCH --account=default +#SBATCH --partition=YOUR_GPU_PARTITION +#SBATCH --account=YOUR_ACCOUNT #SBATCH --nodes=2 #SBATCH --ntasks-per-node=8 #SBATCH --gpus-per-node=8 @@ -79,7 +80,7 @@ OMNISTAT_VENV="${OMNISTAT_VENV:-${OMNIHUB_TOOLS_DIR}/omnihub-inspect}" # A staged copy under ${HG_BASE}/perf-runs/ is intentionally NOT used as the # default — it has drifted in the past from the in-repo template, silently # disabling rocprofiler counters. See ai4science-studio SKILL.md for context. -OMNISTAT_TEMPLATE="${OMNISTAT_TEMPLATE:-${SCRIPT_DIR}/../recipes/perf-analysis/omnistat-lux.config.template}" +OMNISTAT_TEMPLATE="${OMNISTAT_TEMPLATE:-${SCRIPT_DIR}/../recipes/perf-analysis/omnistat.config.template}" OMNISTAT_USERMODE_INTERVAL="${OMNISTAT_USERMODE_INTERVAL:-1}" # Kernel-dispatch tracing (per-kernel name + duration histogram from every @@ -163,7 +164,9 @@ done # --------------------------------------------------------------------------- mkdir -p "$HG_OUTPUT_DIR" OMNISTAT_CONFIG="${HG_OUTPUT_DIR}/omnistat.config" -sed -e "s|@JOB_DIR@|${HG_OUTPUT_DIR}|g" "$OMNISTAT_TEMPLATE" > "$OMNISTAT_CONFIG" +sed -e "s|@JOB_DIR@|${HG_OUTPUT_DIR}|g" \ + -e "s|@OMNIHUB_TOOLS_DIR@|${OMNIHUB_TOOLS_DIR}|g" \ + "$OMNISTAT_TEMPLATE" > "$OMNISTAT_CONFIG" # Opt-in kernel tracing: flip enable_kernel_trace=True in the rendered config # and verify the rocprofiler-sdk tool library exists. Bail out loudly if @@ -293,7 +296,7 @@ if [[ "$HG_SKIP_NODE_HEALTH_PROBE" != "1" ]]; then fi # --------------------------------------------------------------------------- -# Start Omnistat user-mode (datadir from omnistat-lux.config: under HG_OUTPUT_DIR) +# Start Omnistat user-mode (datadir from omnistat.config: under HG_OUTPUT_DIR) # --------------------------------------------------------------------------- echo "--- Starting Omnistat user-mode (interval=${OMNISTAT_USERMODE_INTERVAL}s) ---" export PATH="${OMNISTAT_VENV}/bin:${PATH}" diff --git a/material_science/models/HydraGNN/recipes/perf-analysis/README.md b/material_science/models/HydraGNN/recipes/perf-analysis/README.md index fe4fc86..efa3aa0 100644 --- a/material_science/models/HydraGNN/recipes/perf-analysis/README.md +++ b/material_science/models/HydraGNN/recipes/perf-analysis/README.md @@ -1,12 +1,12 @@ # HydraGNN Performance Analysis with AMD AI Agents -Multi-subagent workflow that launches a 2-node HydraGNN training on Lux (MI355X), captures PyTorch traces and Omnistat telemetry, then runs paired analytics + verifier subagents to identify bottlenecks and propose remedies. +Multi-subagent workflow that launches a 2-node HydraGNN training on AMD Instinct MI355X (gfx950), captures PyTorch traces and Omnistat telemetry, then runs paired analytics + verifier subagents to identify bottlenecks and propose remedies. -> **Audience:** internal AMD performance-engineering. The output is a diagnosis of the run, not a scientific claim about HydraGNN. +> **Audience:** performance engineers. The output is a diagnosis of the run, not a scientific claim about HydraGNN. ## What this recipe does -1. **Launcher** subagent submits a 2-node training (`sbatch_train_perf_amd.sh`) on the `lux` partition with PyTorch profiling armed for one target epoch and Omnistat user-mode collecting on every node. +1. **Launcher** subagent submits a 2-node training (`sbatch_train_perf_amd.sh`) with PyTorch profiling armed for one target epoch and Omnistat user-mode collecting on every node. 2. **Analytics** subagents (TraceLens + Omnistat) run after the job, in parallel, and each emits a `claims.json`. 3. **Verifier** subagents independently re-derive the top claims from raw data and optionally probe cheap remedies on a 1-node interactive `srun`. 4. **Synthesizer** merges, ranks, and writes `combined_report.md`. @@ -32,7 +32,7 @@ flowchart TD ls .cluster-config.yaml # 2. Set the shared dir and submit the perf-analysis job -export AI4S_SHARED_DIR=/shared/aaji +export AI4S_SHARED_DIR=/shared/$USER sbatch material_science/models/HydraGNN/examples/sbatch_train_perf_amd.sh # 3. Once the job completes, drive the analyst + verifier + synthesizer subagents @@ -40,14 +40,14 @@ sbatch material_science/models/HydraGNN/examples/sbatch_train_perf_amd.sh # .cursor/skills/ai4science-perf-analysis/SKILL.md to dispatch them. ``` -The orchestrating agent is expected to run on the **login node** (`rad-vultr-login`) with shell + filesystem access to `/shared/aaji`. Verifier subagents that need a compute node use short interactive `srun -p lux -A vultr_lux -N1 --time=00:05:00` allocations; they never grab 2-node allocations. +The orchestrating agent is expected to run on a **login node** with shell + filesystem access to `$AI4S_SHARED_DIR`. Verifier subagents that need a compute node use short interactive `srun -p -A -N1 --time=00:05:00` allocations; they never grab 2-node allocations. ## Artifacts After a complete run: ``` -/shared/aaji/models/HydraGNN/perf-runs// +$AI4S_SHARED_DIR/models/HydraGNN/perf-runs// ├── manifest.json # written by launcher ├── omnistat-db/ # VictoriaMetrics datadir ├── logs//*.pt.trace.json # PyTorch trace (rank 0 only) @@ -67,7 +67,7 @@ The `combined_report.md` is the deliverable: ranked bottlenecks, remedies tried ## Prerequisites -- Lux access (`vultr_lux` account) with the standard HydraGNN setup at `/shared/aaji/models/HydraGNN/` — overlay, weights, code clone. See [HydraGNN/recipes/train/](../train/). +- Cluster access with the standard HydraGNN setup at `$AI4S_SHARED_DIR/models/HydraGNN/` — overlay, weights, code clone. See [HydraGNN/recipes/train/](../train/). - Export `OMNIHUB_TOOLS_DIR` (the shared perf-tool location, from `.cluster-config.yaml` `omnihub.tools_dir`, e.g. `/shared/omnihub/tools`). Scripts and subagents resolve all tool paths from this var — no cluster path is hardcoded. - One-time install of Omnistat (PR #271 branch `jorda/skills`, with `origin/main` merged in), VictoriaMetrics, and TraceLens — performed lazily by the launcher subagent into `${OMNIHUB_TOOLS_DIR}/`. - The launcher writes `gfm_mlip_with_profile.json` at submit time (does not modify the upstream `gfm_mlip.json`). @@ -80,7 +80,7 @@ The `combined_report.md` is the deliverable: ranked bottlenecks, remedies tried | `OMNISTAT_TRACE_LIB` | `${OMNIHUB_TOOLS_DIR}/omnistat-src/build-trace/libomnistat_trace.so` | Override only for development. The wrapper hard-fails if the path is missing when `OMNISTAT_KERNEL_TRACE=1`. | | `OMNISTAT_TRACE_LOG` | `1` | Library prints `[host][pid][omnistat] Trace summary: N/N processed records (M/M successful flushes)` on rank exit; useful as a smoke-signal that the tool initialized even when the workload crashed. | -`enable_rocprofiler=True` is **on** in `omnistat-lux.config.template` by default (since commit `a23e5c4`); the rendered config's state is printed in the sbatch banner so a silent "counters off" cannot recur. Kernel tracing and device-counting can co-exist on a single GPU but raise per-job VictoriaMetrics cardinality — pick by run length and the question you're asking. See `ai4science-studio` SKILL §12 for the device-counting vs kernel-trace vs `rocprofv3` decision matrix. +`enable_rocprofiler=True` is **on** in `omnistat.config.template` by default; the rendered config's state is printed in the sbatch banner so a silent "counters off" cannot recur. Kernel tracing and device-counting can co-exist on a single GPU but raise per-job VictoriaMetrics cardinality — pick by run length and the question you're asking. See `ai4science-studio` SKILL §12 for the device-counting vs kernel-trace vs `rocprofv3` decision matrix. ## Bottleneck taxonomy (used by analysts and verifier) @@ -121,5 +121,5 @@ The orchestrating agent dispatches each as a Task subagent (or shell script in i - Verifier remedy probes use **1-node** interactive allocations only, ≤5 min each. - Never re-runs the full 2-node job during analysis — that's a deliberate choice for the agent. -- All artifacts under `/shared/aaji/...`; large traces and DBs are not committed (see repo `.gitignore`). +- All artifacts under `$AI4S_SHARED_DIR/...`; large traces and DBs are not committed (see repo `.gitignore`). - Research/perf-engineering only — no scientific claims about HydraGNN should be derived from these short runs. diff --git a/material_science/models/HydraGNN/recipes/perf-analysis/agents/launcher.md b/material_science/models/HydraGNN/recipes/perf-analysis/agents/launcher.md index fe5a03b..ff4afcf 100644 --- a/material_science/models/HydraGNN/recipes/perf-analysis/agents/launcher.md +++ b/material_science/models/HydraGNN/recipes/perf-analysis/agents/launcher.md @@ -1,22 +1,22 @@ # launcher subagent -Submits a 2-node HydraGNN training on Lux with PyTorch profiling and Omnistat user-mode telemetry, waits for completion, and writes `manifest.json` for downstream subagents. +Submits a 2-node HydraGNN training (AMD Instinct MI355X) with PyTorch profiling and Omnistat user-mode telemetry, waits for completion, and writes `manifest.json` for downstream subagents. ## Inputs -- Repo at `/home/aaji/git/ai4science-studio` on branch `aaji/perf-agents-hydragnn`. -- Cluster config at `.cluster-config.yaml` (partition `lux`, account `vultr_lux`). +- The ai4science-studio repo checkout (`$REPO_ROOT`). +- Cluster config at `.cluster-config.yaml` (partition/account, shared dirs). - Optional env-var overrides: `HG_BATCH_SIZE`, `HYDRAGNN_MAX_NUM_BATCH`, `HG_NUM_EPOCH`, `HG_PRECISION`, `PROFILE_TARGET_EPOCH`. ## Outputs - `${OMNIHUB_TOOLS_DIR}/omnihub-inspect/` — Python venv with omnistat (PR #271 + main merged) and TraceLens. - `${OMNIHUB_TOOLS_DIR}/victoriametrics/victoria-metrics-prod` — VictoriaMetrics binary. -- `/shared/aaji/models/HydraGNN/perf-runs/omnistat-lux.config` — Omnistat user-mode config. -- `/shared/aaji/models/HydraGNN/perf-runs//manifest.json` — manifest schema below. -- `/shared/aaji/models/HydraGNN/perf-runs//hydragnn-train-.out` — symlinked from output dir. -- `/shared/aaji/models/HydraGNN/perf-runs//logs/` — symlinked rank-0 trace. -- `/shared/aaji/models/HydraGNN/perf-runs//omnistat-db/` — VictoriaMetrics datadir. +- `$AI4S_SHARED_DIR/models/HydraGNN/perf-runs/omnistat.config` — Omnistat user-mode config. +- `$AI4S_SHARED_DIR/models/HydraGNN/perf-runs//manifest.json` — manifest schema below. +- `$AI4S_SHARED_DIR/models/HydraGNN/perf-runs//hydragnn-train-.out` — symlinked from output dir. +- `$AI4S_SHARED_DIR/models/HydraGNN/perf-runs//logs/` — symlinked rank-0 trace. +- `$AI4S_SHARED_DIR/models/HydraGNN/perf-runs//omnistat-db/` — VictoriaMetrics datadir. ## Manifest schema @@ -29,9 +29,9 @@ Submits a 2-node HydraGNN training on Lux with PyTorch profiling and Omnistat us "nodes": 2, "gpus_per_node": 8, "ranks": 16, - "nodelist": "lux-mi355x-aN,lux-mi355x-bM", - "partition": "lux", - "account": "vultr_lux", + "nodelist": ",", + "partition": "", + "account": "", "runtime_seconds": , "config_used": "", "profile_target_epoch": , @@ -110,9 +110,9 @@ TRACE_LIB=$TOOLS/omnistat-src/build-trace/libomnistat_trace.so if [[ "${OMNISTAT_KERNEL_TRACE:-0}" == "1" && ! -f "$TRACE_LIB" ]]; then # Must build inside the same SIF the workload uses — login nodes lack # apptainer and ROCm headers. C++20 + libcurl + rocprofiler-sdk needed. - salloc -p lux -A vultr_lux -N 1 --time=00:15:00 --gpus-per-node=1 \ + salloc -p -A -N 1 --time=00:15:00 --gpus-per-node=1 \ apptainer exec --rocm \ - "/shared/aaji/images/pytorch_rocm7.2.2_ubuntu24.04_py3.12_pytorch_release_2.10.0.sif" \ + "${AI4S_SHARED_DIR}/images/pytorch_rocm7.2.2_ubuntu24.04_py3.12_pytorch_release_2.10.0.sif" \ bash -c " set -e cd $TOOLS/omnistat-src @@ -127,11 +127,11 @@ The sbatch wrapper expects the `.so` at `${OMNIHUB_TOOLS_DIR}/omnistat-src/build ### 2. Author the Omnistat user-mode config (once, in repo `perf-runs/`) -If `/shared/aaji/models/HydraGNN/perf-runs/omnistat-lux.config` doesn't exist, write it from the template at `/home/aaji/git/ai4science-studio/material_science/models/HydraGNN/recipes/perf-analysis/omnistat-lux.config.template`. The template uses `%(SLURM_JOB_ID)s` placeholders that omnistat-usermode resolves at runtime. +If `$AI4S_SHARED_DIR/models/HydraGNN/perf-runs/omnistat.config` doesn't exist, write it from the template at `$REPO_ROOT/material_science/models/HydraGNN/recipes/perf-analysis/omnistat.config.template`. The template uses `%(SLURM_JOB_ID)s` placeholders that omnistat-usermode resolves at runtime. ### 3. Generate the per-job profile config -Read the upstream `gfm_mlip.json` from `/shared/aaji/models/HydraGNN/code/HydraGNN/examples/multidataset_hpo_sc26/gfm_mlip.json` and inject the `Profile` block **inside `NeuralNetwork`** (not at the top level — `train_validate_test()` is invoked with `config["NeuralNetwork"]` as `config`, so its Profiler reads `config["Profile"]` from that scope): +Read the upstream `gfm_mlip.json` from `$AI4S_SHARED_DIR/models/HydraGNN/code/HydraGNN/examples/multidataset_hpo_sc26/gfm_mlip.json` and inject the `Profile` block **inside `NeuralNetwork`** (not at the top level — `train_validate_test()` is invoked with `config["NeuralNetwork"]` as `config`, so its Profiler reads `config["Profile"]` from that scope): ```python cfg.setdefault("NeuralNetwork", {})["Profile"] = { @@ -140,21 +140,21 @@ cfg.setdefault("NeuralNetwork", {})["Profile"] = { } ``` -Write the result to `/shared/aaji/models/HydraGNN/perf-runs//gfm_mlip_profile.json`. Use Python (`json.load`/`json.dump`) — do NOT do this with sed. +Write the result to `$AI4S_SHARED_DIR/models/HydraGNN/perf-runs//gfm_mlip_profile.json`. Use Python (`json.load`/`json.dump`) — do NOT do this with sed. (The wrapper `examples/sbatch_train_perf_amd.sh` already does this; this step exists in the launcher prompt for the case where the launcher is reused for a non-HydraGNN model.) ### 4. Submit the job ```bash -export AI4S_SHARED_DIR=/shared/aaji -export HG_OUTPUT_DIR=/shared/aaji/models/HydraGNN/perf-runs/PENDING-$$ +export AI4S_SHARED_DIR=/shared/$USER +export HG_OUTPUT_DIR=$AI4S_SHARED_DIR/models/HydraGNN/perf-runs/PENDING-$$ mkdir -p "$HG_OUTPUT_DIR" -cd /home/aaji/git/ai4science-studio +cd "$REPO_ROOT" OUT=$(sbatch \ - --partition=lux \ - --account=vultr_lux \ + --partition= \ + --account= \ --nodes=2 \ --ntasks-per-node=8 \ --gpus-per-node=8 \ @@ -178,7 +178,7 @@ Poll `sacct -j $JOBID -X -n --format=State,ExitCode,Elapsed,NodeList -P` every 3 Once terminal: ```bash -PERF_RUN=/shared/aaji/models/HydraGNN/perf-runs/${JOBID} +PERF_RUN=$AI4S_SHARED_DIR/models/HydraGNN/perf-runs/${JOBID} mv "$HG_OUTPUT_DIR" "$PERF_RUN" # Find the rank-0 trace TRACE=$(find "$PERF_RUN/logs" -name '*.pt.trace.json*' 2>/dev/null | head -1 || true) diff --git a/material_science/models/HydraGNN/recipes/perf-analysis/agents/omnistat_analyst.md b/material_science/models/HydraGNN/recipes/perf-analysis/agents/omnistat_analyst.md index c715cf4..56e9fef 100644 --- a/material_science/models/HydraGNN/recipes/perf-analysis/agents/omnistat_analyst.md +++ b/material_science/models/HydraGNN/recipes/perf-analysis/agents/omnistat_analyst.md @@ -24,7 +24,7 @@ Drive `omnistat-inspect` (PR #271) through the analyze-job phases on the user-mo ### 1. Start VictoriaMetrics on the DB (login node) -Following the [load-database SKILL](https://github.com/ROCm/omnistat/blob/jorda/skills/skills/load-database/SKILL.md), with one Lux-specific addition: pass `-fs.disableMmap` so VM can open a job-scoped DB inside the login-node cgroup (without it, even a 1.6 MB DB panics with "cannot mmap file with size 4096 bytes ... no such device"). +Following the [load-database SKILL](https://github.com/ROCm/omnistat/blob/jorda/skills/skills/load-database/SKILL.md), with one addition for memory-constrained login nodes: pass `-fs.disableMmap` so VM can open a job-scoped DB inside the login-node cgroup (without it, even a 1.6 MB DB panics with "cannot mmap file with size 4096 bytes ... no such device"). ```bash DB=$(jq -r .omnistat_db_path ) diff --git a/material_science/models/HydraGNN/recipes/perf-analysis/agents/omnistat_verifier.md b/material_science/models/HydraGNN/recipes/perf-analysis/agents/omnistat_verifier.md index e1d4cc1..81ab39c 100644 --- a/material_science/models/HydraGNN/recipes/perf-analysis/agents/omnistat_verifier.md +++ b/material_science/models/HydraGNN/recipes/perf-analysis/agents/omnistat_verifier.md @@ -53,9 +53,9 @@ The verifier should re-derive the actual number from the PromQL result and compa ### 2. Optional remedy probe -Same constraints as the tracelens verifier — at most one 1-node `srun -p lux -A vultr_lux -N1 --time=00:05:00`. Telemetry-side probes are typically configuration changes (e.g. enable rocprofiler `hbm` counter set in the omnistat config and re-run for 1 minute on 1 node to see if HBM bandwidth is actually saturating). +Same constraints as the tracelens verifier — at most one 1-node `srun -p -A -N1 --time=00:05:00`. Telemetry-side probes are typically configuration changes (e.g. enable rocprofiler `hbm` counter set in the omnistat config and re-run for 1 minute on 1 node to see if HBM bandwidth is actually saturating). -If a probe needs the omnistat config to change, write a temp config file based on `omnistat-lux.config` with the rocprofiler section uncommented, point `OMNISTAT_CONFIG` at it, and submit via the existing sbatch script as a tiny job with `HG_NUM_EPOCH=1 HYDRAGNN_MAX_NUM_BATCH=10 --nodes=1`. Time-budget that branch generously (10 min) since it does include the omnistat collector startup. +If a probe needs the omnistat config to change, write a temp config file based on `omnistat.config` with the rocprofiler section uncommented, point `OMNISTAT_CONFIG` at it, and submit via the existing sbatch script as a tiny job with `HG_NUM_EPOCH=1 HYDRAGNN_MAX_NUM_BATCH=10 --nodes=1`. Time-budget that branch generously (10 min) since it does include the omnistat collector startup. If multi-node is required (e.g. ANP plugin retest), set `remedy_probe.ran=false; notes="multi-node probe deferred"`. diff --git a/material_science/models/HydraGNN/recipes/perf-analysis/agents/tracelens_analyst.md b/material_science/models/HydraGNN/recipes/perf-analysis/agents/tracelens_analyst.md index 66e03a6..0aaff1d 100644 --- a/material_science/models/HydraGNN/recipes/perf-analysis/agents/tracelens_analyst.md +++ b/material_science/models/HydraGNN/recipes/perf-analysis/agents/tracelens_analyst.md @@ -4,7 +4,7 @@ Generate a TraceLens performance report for the rank-0 PyTorch trace and produce ## Inputs -- `/shared/aaji/models/HydraGNN/perf-runs//manifest.json` +- `$AI4S_SHARED_DIR/models/HydraGNN/perf-runs//manifest.json` ## Outputs @@ -30,7 +30,7 @@ Every entry: "evidence_excerpt": "<200 chars max>", "confidence": "high|medium|low", "proposed_remedy": "", - "remedy_test_command_or_null": "srun -p lux -A vultr_lux -N1 --time=00:05:00 ... | null" + "remedy_test_command_or_null": "srun -p -A -N1 --time=00:05:00 ... | null" } ``` @@ -58,7 +58,7 @@ TraceLens_generate_perf_report_pytorch \ --short_kernel_study 2>&1 | tee tracelens.log ``` -CLI confirmed against TraceLens v0.1.0+ on Lux: `--output_xlsx_path` controls the workbook path; `--output_csvs_dir` writes per-sheet CSVs (handy for the verifier). `--enable_kernel_summary` adds the kernel-level summary sheet. `--short_kernel_study` flags very short kernels (likely launch-overhead bound). +CLI confirmed against TraceLens v0.1.0+: `--output_xlsx_path` controls the workbook path; `--output_csvs_dir` writes per-sheet CSVs (handy for the verifier). `--enable_kernel_summary` adds the kernel-level summary sheet. `--short_kernel_study` flags very short kernels (likely launch-overhead bound). If multi-node, the script will detect collective ops and produce a separate "Collective Analysis" sheet unless `--disable_coll_analysis` is passed (don't pass it). diff --git a/material_science/models/HydraGNN/recipes/perf-analysis/agents/tracelens_verifier.md b/material_science/models/HydraGNN/recipes/perf-analysis/agents/tracelens_verifier.md index a9494ae..c1d668f 100644 --- a/material_science/models/HydraGNN/recipes/perf-analysis/agents/tracelens_verifier.md +++ b/material_science/models/HydraGNN/recipes/perf-analysis/agents/tracelens_verifier.md @@ -69,15 +69,15 @@ Choose the highest-impact `verified` claim with a non-null `remedy_test_command_ Example probe for the bf16 / fp64 ceiling claim — run via the existing HydraGNN sbatch script in interactive mode: ```bash -srun -p lux -A vultr_lux -N1 --ntasks=8 --gpus-per-node=8 --cpus-per-task=16 \ +srun -p -A -N1 --ntasks=8 --gpus-per-node=8 --cpus-per-task=16 \ --time=00:05:00 --pty bash -c ' -export AI4S_SHARED_DIR=/shared/aaji +export AI4S_SHARED_DIR=/shared/$USER export HG_PRECISION=bf16 export HYDRAGNN_MAX_NUM_BATCH=10 export HG_NUM_EPOCH=1 export HG_OUTPUT_DIR=/tmp/perf-probe-$$ mkdir -p $HG_OUTPUT_DIR -bash /home/aaji/git/ai4science-studio/material_science/models/HydraGNN/examples/sbatch_train_amd.sh \ +bash "$REPO_ROOT/material_science/models/HydraGNN/examples/sbatch_train_amd.sh" \ 2>&1 | tee /tmp/probe-$SLURM_JOB_ID.out grep -oE "[0-9.]+s/it" /tmp/probe-$SLURM_JOB_ID.out | tail -5 ' diff --git a/material_science/models/HydraGNN/recipes/perf-analysis/omnistat-lux.config.template b/material_science/models/HydraGNN/recipes/perf-analysis/omnistat.config.template similarity index 90% rename from material_science/models/HydraGNN/recipes/perf-analysis/omnistat-lux.config.template rename to material_science/models/HydraGNN/recipes/perf-analysis/omnistat.config.template index 1eee37e..80817de 100644 --- a/material_science/models/HydraGNN/recipes/perf-analysis/omnistat-lux.config.template +++ b/material_science/models/HydraGNN/recipes/perf-analysis/omnistat.config.template @@ -1,12 +1,12 @@ #-- -# Omnistat user-mode config for Lux (MI355X) +# Omnistat user-mode config for AMD Instinct MI355X (gfx950). # # Adapted from HydraGNN's upstream omnistat.hydragnn-external.config. -# Differences from upstream / from system-mode default on Lux: +# Differences from upstream / from the system-mode default: # - enable_amd_smi (newer SMI), not rocm_smi # - rocprofiler counters DISABLED in v1 (turn on per-probe via verifier) # - allowed_ips includes 0.0.0.0 so the user-mode VM scrapes -# - victoria_datadir under /shared/aaji (NFS, persists after job) +# - victoria_datadir under the per-job dir (@JOB_DIR@), persists after job #-- [omnistat.collectors] @@ -36,7 +36,7 @@ enable_cu_occupancy = False # alone — see ai4science-studio SKILL §11 for the tradeoffs. enable_kernel_trace = False -# Path to local ROCm install (Lux: 7.2.x) +# Path to local ROCm install (validated on ROCm 7.2.x) rocm_path = /opt/rocm # allowed scrapers: localhost (the user-mode VM sits on the same node) + 0.0.0.0 @@ -47,8 +47,8 @@ allowed_ips = 127.0.0.1, 0.0.0.0 host_skip = "login.*" enable_annotations = True -# Lux note: /tmp/omni_rmsjobinfo (the upstream default) is owned by user -# `omnidc` on Lux compute nodes for the system-mode collector, so user-mode +# Note: /tmp/omni_rmsjobinfo (the upstream default) is often owned by the +# system-mode collector's service user on compute nodes, so user-mode # omnistat-rms-env hits PermissionError if it tries to write there. The # sbatch wrapper substitutes @JOB_DIR@ with the job-scoped perf-run dir # (NFS-mounted on every compute node) before passing the resolved config @@ -83,7 +83,7 @@ profile = hbm_flops_f64 [omnistat.collectors.rocprofiler.hbm_flops_f64] # MI355X (gfx950) rocprofiler counter limits — see ai4science-studio SKILL §10. # -# Constraints surfaced on lux-mi355x-b1 / ROCm 7.2.2: +# Constraints surfaced on MI355X (gfx950) / ROCm 7.2.2: # * FETCH_SIZE + WRITE_SIZE in ONE profile fails: "error code 38: Request # exceeds the capabilities of the hardware to collect" (both in TCC block). # * FETCH_SIZE OR WRITE_SIZE + 4 x SQ_INSTS_VALU_*_F64 in one profile works. @@ -110,7 +110,7 @@ counters = ["FETCH_SIZE", "SQ_INSTS_VALU_ADD_F64", "SQ_INSTS_VALU_MUL_F64", "SQ_ [omnistat.query] prometheus_url = http://localhost:8428 -system_name = lux +system_name = mi355x #-- # User-mode settings — VictoriaMetrics + per-job datadir @@ -120,7 +120,7 @@ system_name = lux #-- [omnistat.usermode] -victoria_binary = /shared/aaji/tools/victoriametrics/victoria-metrics-prod +victoria_binary = @OMNIHUB_TOOLS_DIR@/victoriametrics/victoria-metrics-prod victoria_datadir = @JOB_DIR@/omnistat-db victoria_logfile = vic_server.log diff --git a/material_science/models/HydraGNN/recipes/perf-optimizer-loop/agents/orchestrator.md b/material_science/models/HydraGNN/recipes/perf-optimizer-loop/agents/orchestrator.md index 23b4a6a..256addc 100644 --- a/material_science/models/HydraGNN/recipes/perf-optimizer-loop/agents/orchestrator.md +++ b/material_science/models/HydraGNN/recipes/perf-optimizer-loop/agents/orchestrator.md @@ -59,7 +59,7 @@ One line per event. Format: ` `. Examples: ``` 2026-05-22T22:14:03Z LOOP_START uuid= n_iters_budget=5 driver=claude-code-cli -2026-05-22T22:14:09Z PREFLIGHT_OK cluster=lux disk_free=26T claude_cli=ok api_egress=ok +2026-05-22T22:14:09Z PREFLIGHT_OK disk_free=26T claude_cli=ok api_egress=ok 2026-05-22T22:14:11Z LEVER_PICK iter=0 lever=baseline reason="establishing fresh baseline" 2026-05-22T22:15:11Z ITER_SUBMIT iter=0 lever=baseline jobid=7050 2026-05-22T22:23:48Z ITER_COMPLETE iter=0 jobid=7050 state=COMPLETED runtime=487s @@ -86,8 +86,8 @@ Read `loop-/foms.csv` if it exists. If iter rows are present, set `current Run these checks in sequence; abort with `PREFLIGHT_FAIL reason=` on the first failure: - Ensure `AI4S_SHARED_DIR` and `OMNIHUB_TOOLS_DIR` are exported (the latter from `.cluster-config.yaml` `omnihub.tools_dir`, e.g. `/shared/omnihub/tools`); abort if either is unset. All tool paths below derive from `$OMNIHUB_TOOLS_DIR` — never hardcode a cluster path. -- `sinfo -p lux,rad -h` → if zero `IDLE` or `MIX` nodes, abort. Cluster may be in maintenance (see SKILL §14). -- `df -h /shared | tail -1` → abort if `Use%` > 95. +- `sinfo -p -h` (partition from `.cluster-config.yaml`) → if zero `IDLE` or `MIX` nodes, abort. Cluster may be in maintenance (see SKILL §14). +- `df -h "$AI4S_SHARED_DIR" | tail -1` → abort if `Use%` > 95. - `test -x ${OMNIHUB_TOOLS_DIR}/omnihub-inspect/bin/omnistat-usermode` → abort if missing. - `test -x ${OMNIHUB_TOOLS_DIR}/victoriametrics/victoria-metrics-prod` → abort if missing. - `test -f $AI4S_SHARED_DIR/models/HydraGNN/overlays/hydragnn-overlay.img` → abort if missing. diff --git a/material_science/models/HydraGNN/recipes/perf-optimizer-loop/dispatch-attribution.md b/material_science/models/HydraGNN/recipes/perf-optimizer-loop/dispatch-attribution.md index 7a90a9f..21f57d5 100644 --- a/material_science/models/HydraGNN/recipes/perf-optimizer-loop/dispatch-attribution.md +++ b/material_science/models/HydraGNN/recipes/perf-optimizer-loop/dispatch-attribution.md @@ -1,8 +1,8 @@ # Dispatch attribution — job 7187 (current best, 75 s/epoch) -**Perf run:** `/shared/aaji/models/HydraGNN/perf-runs/7187/` +**Perf run:** `$AI4S_SHARED_DIR/models/HydraGNN/perf-runs/7187/` **Lever:** `rccl_high_priority` (`TORCH_NCCL_HIGH_PRIORITY=1`, `GPU_MAX_HW_QUEUES=2`) -**Nodes:** `lux-mi355x-b1`, `lux-mi355x-b2` (2×8 GPUs) +**Nodes:** 2 × MI355X (8 GPUs each) ## Artifacts produced (2026-05-27) @@ -57,7 +57,7 @@ Short-kernel study (`short_kernels_summary.csv`): thousands of `aten::fill_`, `a ## Commands to reproduce ```bash -export AI4S_SHARED_DIR=/shared/aaji +export AI4S_SHARED_DIR=/shared/$USER export REPO_ROOT=$HOME/git/ai4science-studio source ${OMNIHUB_TOOLS_DIR}/omnihub-inspect/bin/activate @@ -68,7 +68,7 @@ ${OMNIHUB_TOOLS_DIR}/victoriametrics/victoria-metrics-prod \ cd $AI4S_SHARED_DIR/models/HydraGNN/perf-runs/7187/tracelens TraceLens_generate_perf_report_pytorch \ - --profile_json_path ../logs/hydragnn-train-7187-N2/lux-mi355x-b1_532131.1779479818046951705.pt.trace.json \ + --profile_json_path ../logs/hydragnn-train-7187-N2/_.pt.trace.json \ --output_xlsx_path report.xlsx --output_csvs_dir csvs/ \ --enable_kernel_summary --short_kernel_study diff --git a/material_science/models/HydraGNN/recipes/train/README.md b/material_science/models/HydraGNN/recipes/train/README.md index a607177..cd2c6e9 100644 --- a/material_science/models/HydraGNN/recipes/train/README.md +++ b/material_science/models/HydraGNN/recipes/train/README.md @@ -247,7 +247,7 @@ Total training time: 720 s (12 min). Loss reduced 25.9% over 5 epochs with clear HydraGNN also writes TensorBoard events when `HYDRAGNN_VALTEST=1`. Event files appear in `$HG_OUTPUT_DIR/logs//`. The parser supports this mode: ```bash -python examples/parse_convergence.py --tbdir /shared/aaji/models/HydraGNN/outputs/logs/hydragnn-train--N1/ -o convergence.csv +python examples/parse_convergence.py --tbdir $AI4S_SHARED_DIR/models/HydraGNN/outputs/logs/hydragnn-train--N1/ -o convergence.csv ``` Note: TensorBoard scalars are only populated when validation runs (they remain empty with `HYDRAGNN_VALTEST=0`). From cdc4e5c21ea409db11d7ca8cf5538b8993ddf7f3 Mon Sep 17 00:00:00 2001 From: Ashwin Aji Date: Mon, 8 Jun 2026 18:47:34 +0000 Subject: [PATCH 26/31] Add HydraGNN matched scaling sweep tooling and document steady-state results. Introduces run_scaling_study.sh and collate_scaling_study.py for reproducible strong-scaling runs, bakes RCCL high-priority defaults into sbatch_train_amd.sh, and updates train recipe, examples README, and agent skills so the workflow is discoverable. Co-authored-by: Cursor --- .../ai4science-material-science/SKILL.md | 1 + .cursor/skills/ai4science-studio/SKILL.md | 2 +- .gitignore | 2 + .../models/HydraGNN/examples/README.md | 19 + .../examples/collate_scaling_study.py | 380 ++++++++++++++++++ .../HydraGNN/examples/run_scaling_study.sh | 97 +++++ .../HydraGNN/examples/sbatch_train_amd.sh | 26 +- .../models/HydraGNN/recipes/train/README.md | 16 +- 8 files changed, 536 insertions(+), 7 deletions(-) create mode 100755 material_science/models/HydraGNN/examples/collate_scaling_study.py create mode 100755 material_science/models/HydraGNN/examples/run_scaling_study.sh diff --git a/.cursor/skills/ai4science-material-science/SKILL.md b/.cursor/skills/ai4science-material-science/SKILL.md index 2648676..6d5f488 100755 --- a/.cursor/skills/ai4science-material-science/SKILL.md +++ b/.cursor/skills/ai4science-material-science/SKILL.md @@ -35,6 +35,7 @@ HydraGNN uses MPI-based distributed training with ADIOS datasets. Key patterns f - **Env vars inside container:** Set `OMP_NUM_THREADS` to match `--cpus-per-task`, `MIOPEN_DISABLE_CACHE=1`, `MIOPEN_USER_DB_PATH=$SCRATCH_LOCAL//miopen` (node-local fast storage from `.cluster-config.yaml`), `HYDRAGNN_USE_VARIABLE_GRAPH_SIZE=1`. - **Multi-node MPI transport (ob1/tcp):** On Pensando/ionic fabrics, data NICs use `/31` subnets that don't route between nodes — IB verbs cannot work for MPI. Use `OMPI_MCA_pml=ob1`, `OMPI_MCA_btl=tcp,self`, `OMPI_MCA_btl_tcp_if_include=$MGMT_NIC`, `MPI4PY_RC_THREADS=false`. RCCL uses ANP plugin (`librccl-anp.so`) over ionic native transport (RoCEv2/GDRDMA) for GPU allreduce, independent of MPI. The ANP plugin and `libionic.so.1` must be bind-mounted from the host into the container. - **Convergence capture:** Run with `HYDRAGNN_VALTEST=1` to get epoch-level loss reporting. Parse with `examples/parse_convergence.py --log `. +- **Strong-scaling sweep:** Submit matched 1/2/4/8-node runs via `examples/run_scaling_study.sh` (identical env, `HYDRAGNN_VALTEST=0`, 6 epochs). Collate steady-state throughput (epochs 2–5) with `examples/collate_scaling_study.py`. Results table and methodology in `recipes/train/README.md` §6. `sbatch_train_amd.sh` defaults `TORCH_NCCL_HIGH_PRIORITY=1` and `GPU_MAX_HW_QUEUES=2`. ## HydraGNN inference pattern diff --git a/.cursor/skills/ai4science-studio/SKILL.md b/.cursor/skills/ai4science-studio/SKILL.md index daaaf40..e9cca33 100755 --- a/.cursor/skills/ai4science-studio/SKILL.md +++ b/.cursor/skills/ai4science-studio/SKILL.md @@ -313,7 +313,7 @@ srun --mpi=pmix apptainer exec \ **Why `HOSTNAME` instead of `MASTER_ADDR`:** Several model scripts do `os.environ["MASTER_ADDR"] = os.environ["HOSTNAME"]`. Inject rank-0's hostname as `HOSTNAME` so this pattern resolves correctly on all nodes. Using `localhost` only works single-node. -**Multi-node scaling:** Change `--nodes=N --ntasks=N*` in the SBATCH header; the `srun` command is unchanged. +**Multi-node scaling:** Change `--nodes=N --ntasks=N*` in the SBATCH header; the `srun` command is unchanged. For HydraGNN matched strong-scaling studies, use `material_science/models/HydraGNN/examples/run_scaling_study.sh` + `collate_scaling_study.py` (see material-science SKILL and `recipes/train/README.md` §6). ## GPU visibility: `--gpus-per-node` NOT `--gpus-per-task` (MI355X / RCCL) diff --git a/.gitignore b/.gitignore index 2911aba..af9457e 100755 --- a/.gitignore +++ b/.gitignore @@ -55,6 +55,8 @@ vic_server.log # Convergence parser outputs (artifact, not source) **/examples/convergence*.png **/examples/convergence*.csv +**/examples/scaling_study* +**/examples/foms-*.json # Model outputs & generated data **/examples/outputs/ diff --git a/material_science/models/HydraGNN/examples/README.md b/material_science/models/HydraGNN/examples/README.md index 12a69ca..4071016 100644 --- a/material_science/models/HydraGNN/examples/README.md +++ b/material_science/models/HydraGNN/examples/README.md @@ -11,6 +11,10 @@ with AMD/ROCm container launch. | [`preflight_hydragnn.py`](preflight_hydragnn.py) | Check your environment before submitting jobs | | [`run_inference.sh`](run_inference.sh) | Load a checkpoint and run predictions | | [`run_train.sh`](run_train.sh) | Train on bundled or staged ADIOS datasets | +| [`sbatch_train_amd.sh`](sbatch_train_amd.sh) | SLURM/Apptainer multi-node training (HPC) | +| [`run_scaling_study.sh`](run_scaling_study.sh) | Submit matched 1/2/4/8-node strong-scaling sweep | +| [`collate_scaling_study.py`](collate_scaling_study.py) | Parse SLURM logs → steady-state throughput table | +| [`parse_convergence.py`](parse_convergence.py) | Extract loss/epoch metrics from training logs | | [`docker_run.sh`](docker_run.sh) | Docker launcher (clones HydraGNN, installs deps) | ## Quick start — Docker @@ -26,6 +30,21 @@ with AMD/ROCm container launch. ./docker_run.sh shell ``` +## Quick start — HPC scaling study + +Matched strong-scaling runs (identical env at 1, 2, 4, 8 nodes; steady-state timing): + +```bash +export AI4S_SHARED_DIR=/path/to/shared +export SBATCH_PARTITION=... SBATCH_ACCOUNT=... +./run_scaling_study.sh + +# After jobs complete (logs in cwd as hydragnn-train-.out): +python3 collate_scaling_study.py --log-dir . --jobs ,,... -o scaling_study +``` + +See [`../recipes/train/README.md`](../recipes/train/README.md) for validated MI355X numbers and network notes. + ## Environment variables | Variable | Required | Default | Description | diff --git a/material_science/models/HydraGNN/examples/collate_scaling_study.py b/material_science/models/HydraGNN/examples/collate_scaling_study.py new file mode 100755 index 0000000..a4502c1 --- /dev/null +++ b/material_science/models/HydraGNN/examples/collate_scaling_study.py @@ -0,0 +1,380 @@ +#!/usr/bin/env python3 +"""collate_scaling_study.py — Build DCSWEAP-4502 scaling table from HydraGNN SLURM logs. + +Parses one or more hydragnn-train-*.out files (or job IDs resolved under a log dir) +and emits a CSV + Markdown summary with throughput, epoch time, scaling efficiency, +and loss sanity checks. + +Usage: + python collate_scaling_study.py \\ + --log hydragnn-train-.out \\ + --log 2node:hydragnn-train-.out \\ + -o scaling_study + + python collate_scaling_study.py --log-dir . --jobs ,,... -o scaling_study +""" + +from __future__ import annotations + +import argparse +import csv +import json +import re +import sys +from pathlib import Path + +# Reuse the SLURM log parser from parse_convergence.py in the same directory. +sys.path.insert(0, str(Path(__file__).resolve().parent)) +from parse_convergence import aggregate_foms_from_records, parse_slurm_log # noqa: E402 + + +def _parse_job_header(log_path: Path) -> dict: + """Extract run configuration from the sbatch banner at the top of the log.""" + meta: dict = {"log_path": str(log_path)} + patterns = { + "nodes": r"^\s*Nodes\s+:\s*(\d+)", + "gpus_per_node": r"^\s*GPUs/node\s+:\s*(\d+)", + "total_ranks": r"^\s*Total ranks\s+:\s*(\d+)", + "batch_size": r"^\s*Batch size\s+:\s*(\d+)", + "num_epochs": r"^\s*Num epochs\s+:\s*(\d+)", + "max_batches": r"^\s*Max batches\s+:\s*(\S+)", + "precision": r"^\s*Precision\s+:\s*(\S+)", + "nodelist": r"^\s*Node\(s\)\s+:\s*(.+)", + "rccl_priority": r"^\s*RCCL priority:\s*TORCH_NCCL_HIGH_PRIORITY=(\d+)", + "hw_queues": r"^\s*HW queues\s+:\s*GPU_MAX_HW_QUEUES=(\d+)", + } + job_id_m = re.search(r"hydragnn-train-(\d+)", log_path.name) + if job_id_m: + meta["job_id"] = job_id_m.group(1) + + with log_path.open() as f: + for line in f: + for key, pat in patterns.items(): + m = re.match(pat, line) + if m: + meta[key] = m.group(1).strip() + if line.startswith("Welcome to SLURM") and "nodes" in meta: + break + + if "nodes" in meta: + meta["nodes"] = int(meta["nodes"]) + if "gpus_per_node" in meta: + meta["gpus_per_node"] = int(meta["gpus_per_node"]) + if "total_ranks" in meta: + meta["total_ranks"] = int(meta["total_ranks"]) + if "batch_size" in meta: + meta["batch_size"] = int(meta["batch_size"]) + if "num_epochs" in meta: + meta["num_epochs"] = int(meta["num_epochs"]) + if "max_batches" in meta and meta["max_batches"].isdigit(): + meta["max_batches"] = int(meta["max_batches"]) + return meta + + +def _all_train_epoch_rates(log_path: Path, num_epochs: int | None = None) -> list[float]: + """Return per-epoch train s/it (instantaneous rate at Train 100%, not cumulative avg). + + HydraGNN tqdm emits two Train 100% rates on one line; the first is the + instantaneous rate at completion, the second includes epoch startup overhead. + """ + pat = re.compile( + r"Train:\s*100%\|[^|]*\|\s*(\d+)/(\d+)\s*\[[^\]]*?,\s*([\d.]+)s/it\]" + ) + rates: list[float] = [] + with log_path.open(errors="replace") as f: + for line in f: + if "Train:" not in line or "100%" not in line: + continue + ms = list(pat.finditer(line)) + if ms: + # First match on the line = instantaneous rate at epoch end. + rates.append(float(ms[0].group(3))) + # HydraGNN may emit duplicate Train 100% lines per epoch; keep the last N. + if num_epochs and len(rates) > num_epochs: + rates = rates[-num_epochs:] + return rates + + +def _train_tqdm_timing(log_path: Path, steady_start: int = 2, num_epochs: int | None = None) -> dict: + """Steady-state train timing from multi-epoch Train tqdm lines.""" + rates = _all_train_epoch_rates(log_path, num_epochs=num_epochs) + if not rates: + return {} + + warmup = rates[:steady_start] if len(rates) > steady_start else [] + steady = ( + rates[steady_start:num_epochs] + if num_epochs and len(rates) >= num_epochs + else rates[steady_start:] + ) + + sec_per_batch = sum(steady) / len(steady) + batches = None + pat_batches = re.compile( + r"Train:\s*100%\|[^|]*\|\s*(\d+)/(\d+)\s*\[[^\]]*?,\s*([\d.]+)s/it\]" + ) + with log_path.open(errors="replace") as f: + for line in f: + m = pat_batches.search(line) + if m: + batches = int(m.group(1)) + + epoch_time_s = sec_per_batch * batches if batches else None + return { + "sec_per_batch_train": sec_per_batch, + "epoch_time_train_s": epoch_time_s, + "batches_per_epoch": batches, + "batches_per_sec_train": 1.0 / sec_per_batch if sec_per_batch else None, + "steady_state_epochs": list(range(steady_start, steady_start + len(steady))), + "per_epoch_train_s_it": rates, + "warmup_epochs_s_it": warmup, + "steady_epochs_s_it": steady, + } + + +def _peak_gpu_mem_gb(log_path: Path) -> float | None: + pat = re.compile( + r"Max memory allocated after optimizer step:\s*([\d.]+)\s*GB" + ) + peak = None + with log_path.open(errors="replace") as f: + for line in f: + m = pat.search(line) + if m: + peak = max(peak or 0.0, float(m.group(1))) + return peak + + +def _epoch_timing( + records: list[dict], + max_batches: int | None, + log_path: Path, + steady_start: int = 2, + num_epochs: int | None = None, +) -> dict: + """Derive per-epoch and per-batch timing from parsed log records.""" + foms = aggregate_foms_from_records(records) + epoch_times = foms.get("epoch_times_s") or [] + + # Prefer steady-state (exclude epoch 0 warmup) when we have multiple epochs. + if len(epoch_times) >= 2: + steady_epochs = epoch_times[1:] + epoch_time_s = sum(steady_epochs) / len(steady_epochs) + elif epoch_times: + epoch_time_s = epoch_times[0] + else: + epoch_time_s = None + + batches_per_epoch = max_batches + if batches_per_epoch is None: + tqdm_recs = [r for r in records if r.get("source") == "tqdm_final"] + if tqdm_recs: + batches_per_epoch = tqdm_recs[0].get("batch") + + sec_per_batch = None + if epoch_time_s is not None and batches_per_epoch: + sec_per_batch = epoch_time_s / float(batches_per_epoch) + + # Direct tqdm s/it average (often more stable for single-epoch runs). + tqdm_rates = [ + float(r["wall_time_s"]) + for r in records + if r.get("source") == "tqdm_final" and r.get("wall_time_s") is not None + ] + if tqdm_rates: + tqdm_sec_per_batch = sum(tqdm_rates) / len(tqdm_rates) + if sec_per_batch is None: + sec_per_batch = tqdm_sec_per_batch + elif len(tqdm_rates) >= 1: + # Prefer tqdm for single-epoch strong-scaling runs. + sec_per_batch = tqdm_sec_per_batch + + train_only = _train_tqdm_timing(log_path, steady_start=steady_start, num_epochs=num_epochs) + if train_only.get("sec_per_batch_train"): + sec_per_batch = train_only["sec_per_batch_train"] + epoch_time_s = train_only.get("epoch_time_train_s", epoch_time_s) + batches_per_epoch = train_only.get("batches_per_epoch", batches_per_epoch) + + batches_per_sec = (1.0 / sec_per_batch) if sec_per_batch else None + peak_mem = _peak_gpu_mem_gb(log_path) + + return { + "epoch_time_s": epoch_time_s, + "epoch_time_train_s": train_only.get("epoch_time_train_s"), + "sec_per_batch": sec_per_batch, + "batches_per_sec": batches_per_sec, + "batches_per_epoch": batches_per_epoch, + "peak_gpu_mem_gb": peak_mem, + "final_train_loss": foms.get("final_train_loss"), + "final_val_loss": foms.get("final_val_loss"), + "num_epochs_completed": foms.get("num_epochs_completed"), + "steady_state_epoch_start": steady_start, + "steady_epochs_s_it": train_only.get("steady_epochs_s_it"), + } + + +def analyze_log(log_path: Path, label: str | None = None, steady_start: int = 2) -> dict: + records = parse_slurm_log(str(log_path)) + meta = _parse_job_header(log_path) + timing = _epoch_timing( + records, + meta.get("max_batches") if isinstance(meta.get("max_batches"), int) else None, + log_path, + steady_start, + meta.get("num_epochs") if isinstance(meta.get("num_epochs"), int) else None, + ) + + row = {**meta, **timing} + if label: + row["label"] = label + row["batch_size"] = meta.get("batch_size", 200) + if row.get("batches_per_sec") and row.get("batch_size"): + row["samples_per_sec"] = row["batches_per_sec"] * row["batch_size"] + return row + + +def _scaling_efficiency(rows: list[dict]) -> None: + """Add strong-scaling efficiency vs 1-node baseline (in-place).""" + baseline = None + for r in sorted(rows, key=lambda x: x.get("nodes", 0)): + if r.get("nodes") == 1 and r.get("sec_per_batch"): + baseline = r["sec_per_batch"] + break + if baseline is None: + return + for r in rows: + if r.get("sec_per_batch") and r.get("nodes"): + ideal = baseline # strong scaling: same wall time per batch regardless of nodes + actual = r["sec_per_batch"] + r["strong_scaling_efficiency"] = ideal / actual if actual else None + r["speedup_vs_1node"] = baseline / actual if actual else None + + +def _resolve_log(path_or_job: str, log_dir: Path) -> Path: + p = Path(path_or_job) + if p.is_file(): + return p + if ":" in path_or_job: + _label, rest = path_or_job.split(":", 1) + return Path(rest) + # Treat as job id + candidates = list(log_dir.glob(f"hydragnn-train-{path_or_job}.out")) + if not candidates: + candidates = list(log_dir.glob(f"**/hydragnn-train-{path_or_job}.out")) + if not candidates: + raise FileNotFoundError(f"No log for job {path_or_job} under {log_dir}") + return candidates[0] + + +def write_outputs(rows: list[dict], stem: Path, steady_start: int = 2) -> None: + _scaling_efficiency(rows) + rows_sorted = sorted(rows, key=lambda r: r.get("nodes", 0)) + + csv_path = stem.with_suffix(".csv") + fieldnames = [ + "job_id", "label", "nodes", "total_ranks", "batch_size", "max_batches", + "precision", "batches_per_epoch", "sec_per_batch", "epoch_time_train_s", + "batches_per_sec", "samples_per_sec", "epoch_time_s", "peak_gpu_mem_gb", + "final_train_loss", "final_val_loss", + "strong_scaling_efficiency", "speedup_vs_1node", "nodelist", + "rccl_priority", "hw_queues", "log_path", + ] + with csv_path.open("w", newline="") as f: + w = csv.DictWriter(f, fieldnames=fieldnames, extrasaction="ignore") + w.writeheader() + for r in rows_sorted: + w.writerow(r) + print(f"Wrote {csv_path}") + + md_path = stem.with_suffix(".md") + lines = [ + "# HydraGNN MI355X Scaling Study (DCSWEAP-4502)", + "", + "Strong-scaling configuration: fixed 50 batches/epoch, batch_size=200, fp64, " + "ANI1x+Alexandria, MPI/ADIOS phase-1, `HYDRAGNN_VALTEST=0` (train-only timing).", + "", + "**Steady-state metric:** mean train s/batch over epochs " + f"{steady_start}+ (exclude warmup epochs 0–{steady_start - 1}).", + "", + "| Nodes | GPUs | Job | Train s/batch | Train epoch (s) | samples/s | " + "Peak GPU mem (GB) | Train loss | Val loss | Strong-scaling eff. |", + "|------:|-----:|----:|--------------:|----------------:|----------:|" + "------------------:|-----------:|---------:|--------------------:|", + ] + for r in rows_sorted: + train_epoch = r.get("epoch_time_train_s") or r.get("epoch_time_s") or 0 + tl = r.get("final_train_loss") + vl = r.get("final_val_loss") + tl_s = f"{tl:.4f}" if tl is not None else "—" + vl_s = f"{vl:.4f}" if vl is not None else "—" + mem = r.get("peak_gpu_mem_gb") or 0 + lines.append( + f"| {r.get('nodes', '?')} " + f"| {r.get('total_ranks', '?')} " + f"| {r.get('job_id', '?')} " + f"| {r.get('sec_per_batch', 0):.2f} " + f"| {train_epoch:.0f} " + f"| {r.get('samples_per_sec', 0):.0f} " + f"| {mem:.1f} " + f"| {tl_s} " + f"| {vl_s} " + f"| {r.get('strong_scaling_efficiency', 0):.2f} |" + ) + lines.extend([ + "", + "## Notes", + "", + "- **All runs share:** `TORCH_NCCL_HIGH_PRIORITY=1`, `GPU_MAX_HW_QUEUES=2`, " + "`HG_NUM_EPOCH=6`, `HYDRAGNN_MAX_NUM_BATCH=50`, `HYDRAGNN_VALTEST=0`.", + f"- **Steady-state:** mean of train s/it for epochs {steady_start}–5 (0-indexed).", + "- **GPU utilization**: optional; SLURM log timing is the primary metric.", + "- **Generated outputs** (`scaling_study*.csv/json/md`) are gitignored run artifacts.", + "", + ]) + md_path.write_text("\n".join(lines)) + print(f"Wrote {md_path}") + + json_path = stem.with_suffix(".json") + json_path.write_text(json.dumps(rows_sorted, indent=2, default=str)) + print(f"Wrote {json_path}") + + +def main() -> None: + ap = argparse.ArgumentParser(description="Collate HydraGNN scaling study results.") + ap.add_argument("--log", action="append", default=[], metavar="JOB[:PATH]", + help="Job id (resolved under --log-dir) or label:path") + ap.add_argument("--log-dir", default=".", help="Directory to search for hydragnn-train-.out") + ap.add_argument("--jobs", default="", help="Comma-separated job ids (shorthand for --log)") + ap.add_argument("--steady-start", type=int, default=2, + help="First epoch index (0-based) included in steady-state average (default: 2)") + ap.add_argument("-o", "--output", default="scaling_study", help="Output stem (no extension)") + args = ap.parse_args() + + specs = list(args.log) + if args.jobs: + specs.extend(j.strip() for j in args.jobs.split(",") if j.strip()) + + if not specs: + ap.error("Provide --log and/or --jobs") + + log_dir = Path(args.log_dir) + rows = [] + for spec in specs: + label = None + if ":" in spec and not Path(spec.split(":", 1)[1]).exists(): + # label:jobid + label, job_part = spec.split(":", 1) + log_path = _resolve_log(job_part, log_dir) + elif ":" in spec: + label, log_path_str = spec.split(":", 1) + log_path = Path(log_path_str) + else: + log_path = _resolve_log(spec, log_dir) + print(f"Parsing {log_path} ...") + rows.append(analyze_log(log_path, label=label, steady_start=args.steady_start)) + + write_outputs(rows, Path(args.output), steady_start=args.steady_start) + + +if __name__ == "__main__": + main() diff --git a/material_science/models/HydraGNN/examples/run_scaling_study.sh b/material_science/models/HydraGNN/examples/run_scaling_study.sh new file mode 100755 index 0000000..c00a0cb --- /dev/null +++ b/material_science/models/HydraGNN/examples/run_scaling_study.sh @@ -0,0 +1,97 @@ +#!/usr/bin/env bash +# Submit a matched HydraGNN strong-scaling sweep. +# +# All node counts share identical env vars. Timing is train-only +# (HYDRAGNN_VALTEST=0); use collate_scaling_study.py to extract steady-state +# s/batch (mean of epochs 2–5 by default). +# +# Usage: +# export AI4S_SHARED_DIR=/path/to/shared +# export SBATCH_PARTITION=... SBATCH_ACCOUNT=... +# ./run_scaling_study.sh # submit 1,2,4,8 nodes +# ./run_scaling_study.sh --warmup # 1-node warm ckpt, then timed runs +# ./run_scaling_study.sh --nodes 1,2 # subset +# +set -euo pipefail + +SCRIPT_DIR=$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd) +: "${AI4S_SHARED_DIR:?Set AI4S_SHARED_DIR}" + +# --------------------------------------------------------------------------- +# Locked scaling-study contract (do not vary between node counts) +# --------------------------------------------------------------------------- +export TORCH_NCCL_HIGH_PRIORITY=1 +export GPU_MAX_HW_QUEUES=2 +export HYDRAGNN_VALTEST=0 +export HG_NUM_EPOCH=6 +export HYDRAGNN_MAX_NUM_BATCH=50 +export HG_BATCH_SIZE=200 +export HG_PRECISION=fp64 +export HYDRAGNN_TRACE_LEVEL=1 +export HG_OUTPUT_DIR="${HG_OUTPUT_DIR:-${AI4S_SHARED_DIR}/models/HydraGNN/outputs/scaling}" + +SBATCH_PARTITION="${SBATCH_PARTITION:-YOUR_GPU_PARTITION}" +SBATCH_ACCOUNT="${SBATCH_ACCOUNT:-YOUR_ACCOUNT}" + +WARMUP=0 +NODES_LIST="1,2,4,8" +while [[ $# -gt 0 ]]; do + case "$1" in + --warmup) WARMUP=1; shift ;; + --nodes) NODES_LIST="$2"; shift 2 ;; + *) echo "Unknown arg: $1" >&2; exit 1 ;; + esac +done + +submit() { + local n="$1" t="$2" + local extra=() + if [[ "$n" -ge 8 && -n "${SCALING_EXCLUDE_NODES:-}" ]]; then + extra+=( --exclude="$SCALING_EXCLUDE_NODES" ) + fi + sbatch --parsable \ + --partition="$SBATCH_PARTITION" \ + --account="$SBATCH_ACCOUNT" \ + --nodes="$n" \ + --time="$t" \ + "${extra[@]}" \ + --job-name="hydragnn-scale-${n}N" \ + --export=ALL,TORCH_NCCL_HIGH_PRIORITY,GPU_MAX_HW_QUEUES,HYDRAGNN_VALTEST,HG_NUM_EPOCH,HYDRAGNN_MAX_NUM_BATCH,HG_BATCH_SIZE,HG_PRECISION,HYDRAGNN_TRACE_LEVEL,HG_OUTPUT_DIR,HG_STARTFROM,HG_WARM_CKPT,AI4S_SHARED_DIR \ + "$SCRIPT_DIR/sbatch_train_amd.sh" +} + +if [[ "$WARMUP" -eq 1 ]]; then + export HG_NUM_EPOCH=1 + export HYDRAGNN_VALTEST=0 + unset HG_STARTFROM HG_WARM_CKPT + WJ=$(submit 1 "00:30:00") + echo "Warmup job submitted: $WJ (1 epoch, saves checkpoint under HG_OUTPUT_DIR/logs/)" + echo "After it completes, re-run with HG_WARM_CKPT and HG_STARTFROM set." + exit 0 +fi + +if [[ -n "${HG_WARM_CKPT:-}" && -n "${HG_STARTFROM:-}" ]]; then + export HG_NUM_EPOCH="${HG_NUM_EPOCH:-5}" + echo "Using warm checkpoint: $HG_WARM_CKPT -> logs/$HG_STARTFROM/" +fi + +IFS=',' read -ra NODE_COUNTS <<< "$NODES_LIST" +declare -A TIME_LIMIT=( [1]="01:00:00" [2]="01:30:00" [4]="02:30:00" [8]="03:00:00" ) +echo "=== HydraGNN scaling sweep ===" +echo " Output dir : $HG_OUTPUT_DIR" +echo " Epochs : $HG_NUM_EPOCH (steady-state: epochs 2–5 via collate_scaling_study.py)" +echo " RCCL : TORCH_NCCL_HIGH_PRIORITY=$TORCH_NCCL_HIGH_PRIORITY GPU_MAX_HW_QUEUES=$GPU_MAX_HW_QUEUES" +echo " Start from : ${HG_STARTFROM:-none}" +echo "" + +JOB_IDS=() +for n in "${NODE_COUNTS[@]}"; do + jid=$(submit "$n" "${TIME_LIMIT[$n]:-02:00:00}") + echo " ${n}-node -> job $jid" + JOB_IDS+=("$jid") +done + +echo "" +echo "Submitted: ${JOB_IDS[*]}" +echo "After all complete:" +echo " python3 collate_scaling_study.py --log-dir . --jobs $(IFS=,; echo "${JOB_IDS[*]}") -o scaling_study" diff --git a/material_science/models/HydraGNN/examples/sbatch_train_amd.sh b/material_science/models/HydraGNN/examples/sbatch_train_amd.sh index b42b200..a6399ac 100755 --- a/material_science/models/HydraGNN/examples/sbatch_train_amd.sh +++ b/material_science/models/HydraGNN/examples/sbatch_train_amd.sh @@ -35,6 +35,10 @@ # SCRATCH_LOCAL Node-local fast storage root (default: /scratch) # HYDRAGNN_MAX_NUM_BATCH Cap batches per epoch (sanity test: 20-50, full: unset) # HYDRAGNN_VALTEST Run val/test during training (default: 0 = skip) +# HG_STARTFROM Resume from logs//.pk (default: none) +# HG_WARM_CKPT Host path to .pk copied into logs// before launch +# TORCH_NCCL_HIGH_PRIORITY Default 1 (RCCL high-priority queue) +# GPU_MAX_HW_QUEUES Default 2 # # HydraGNN pinned SHA: 6c45f1682783e66dc89e9e23009f61716186432b (main) @@ -75,6 +79,10 @@ HG_PRECISION="${HG_PRECISION:-fp64}" HG_OUTPUT_DIR="${HG_OUTPUT_DIR:-${HG_BASE}/outputs}" HG_REPO_DIR="${HG_REPO_DIR:-${HG_BASE}/code/HydraGNN}" +# Validated on MI355X perf-optimizer-loop: ~1.3% epoch-time improvement on 2-node GFM-MLIP. +TORCH_NCCL_HIGH_PRIORITY="${TORCH_NCCL_HIGH_PRIORITY:-1}" +GPU_MAX_HW_QUEUES="${GPU_MAX_HW_QUEUES:-2}" + NODES="${SLURM_JOB_NUM_NODES:-1}" GPUS_PER_NODE=8 TOTAL_RANKS=$((NODES * GPUS_PER_NODE)) @@ -129,6 +137,16 @@ done # --------------------------------------------------------------------------- mkdir -p "$HG_OUTPUT_DIR" +# Optional warm checkpoint for scaling resume (HG_STARTFROM + HG_WARM_CKPT). +if [[ -n "${HG_STARTFROM:-}" && "${HG_STARTFROM}" != "none" && -n "${HG_WARM_CKPT:-}" ]]; then + _WARM_DIR="${HG_OUTPUT_DIR}/logs/${HG_STARTFROM}" + mkdir -p "$_WARM_DIR" + if [[ ! -f "${_WARM_DIR}/${HG_STARTFROM}.pk" ]]; then + cp -a "$HG_WARM_CKPT" "${_WARM_DIR}/${HG_STARTFROM}.pk" + echo " Warm ckpt : $HG_WARM_CKPT -> ${_WARM_DIR}/${HG_STARTFROM}.pk" + fi +fi + # --------------------------------------------------------------------------- # Print job configuration # --------------------------------------------------------------------------- @@ -144,6 +162,9 @@ echo " Precision : $HG_PRECISION" echo " Batch size : ${HG_BATCH_SIZE:-200}" echo " Num epochs : ${HG_NUM_EPOCH:-1}" echo " Max batches : ${HYDRAGNN_MAX_NUM_BATCH:-unlimited}" +echo " RCCL priority: TORCH_NCCL_HIGH_PRIORITY=${TORCH_NCCL_HIGH_PRIORITY}" +echo " HW queues : GPU_MAX_HW_QUEUES=${GPU_MAX_HW_QUEUES}" +echo " Start from : ${HG_STARTFROM:-none}" echo " Output dir : $HG_OUTPUT_DIR" echo " Repo dir : $HG_REPO_DIR" echo " Config : ${EXAMPLE_DIR}/gfm_mlip.json" @@ -216,7 +237,7 @@ sys.argv = [ '--precision=' + os.environ['HG_PRECISION'], '--batch_size=' + str(batch_size), '--num_epoch=' + os.environ.get('HG_NUM_EPOCH', '1'), - '--startfrom=none', + '--startfrom=' + os.environ.get('HG_STARTFROM', 'none'), '--log=hydragnn-train-' + os.environ.get('SLURM_JOB_ID','0') + '-N' + os.environ.get('SLURM_JOB_NUM_NODES','1'), ] @@ -340,10 +361,13 @@ srun --mpi=pmix \ --env HG_PRECISION="$HG_PRECISION" \ --env HG_BATCH_SIZE="${HG_BATCH_SIZE:-200}" \ --env HG_NUM_EPOCH="${HG_NUM_EPOCH:-1}" \ + --env HG_STARTFROM="${HG_STARTFROM:-none}" \ --env HYDRAGNN_VALTEST="${HYDRAGNN_VALTEST:-0}" \ --env HYDRAGNN_TRACE_LEVEL="${HYDRAGNN_TRACE_LEVEL:-1}" \ --env HG_EXAMPLE_DIR="$EXAMPLE_DIR" \ --env HG_OUTPUT_DIR="$HG_OUTPUT_DIR" \ + --env TORCH_NCCL_HIGH_PRIORITY="$TORCH_NCCL_HIGH_PRIORITY" \ + --env GPU_MAX_HW_QUEUES="$GPU_MAX_HW_QUEUES" \ "${OPT_ENVS[@]}" \ "${RCCL_MULTINODE_ENVS[@]}" \ "$HG_SIF" \ diff --git a/material_science/models/HydraGNN/recipes/train/README.md b/material_science/models/HydraGNN/recipes/train/README.md index cd2c6e9..7db30b4 100644 --- a/material_science/models/HydraGNN/recipes/train/README.md +++ b/material_science/models/HydraGNN/recipes/train/README.md @@ -147,11 +147,17 @@ Validated on MI355X (2026-05-19) with `HYDRAGNN_MAX_NUM_BATCH=50`, `HG_BATCH_SIZ ## 6. Performance & Scaling -| Nodes | GPUs | Throughput (batch/s) | Time per batch | Notes | -|-------|------|---------------------|----------------|-------| -| 1 | 8 | ~0.42 | 2.4 s | Validated (50-batch sanity test, 2026-05-19) | -| 2 | 16 | ~0.33 | 3.0 s | Validated (30-batch, ob1/tcp MPI, 2026-05-19) | -| 4 | 32 | TBD | — | Expected to work with same flags | +Matched strong-scaling sweep on MI355X (2026-06-06): identical env on all node counts, steady-state train s/batch (mean of epochs 2–5), `HYDRAGNN_VALTEST=0`, RCCL high-priority enabled. + +| Nodes | GPUs | Train s/batch | samples/s | Strong-scaling eff. | +|------:|-----:|--------------:|----------:|--------------------:| +| 1 | 8 | 2.69 s | 74 | 1.00 | +| 2 | 16 | 2.73 s | 73 | 0.98 | +| 4 | 32 | 2.71 s | 74 | 0.99 | +| 8 | 64 | 2.85 s | 70 | 0.94 | + +Submit a matched sweep: `./examples/run_scaling_study.sh` +Regenerate the table from SLURM logs: `python examples/collate_scaling_study.py --jobs ,,... -o scaling_study` ### Multi-node network architecture From 8d4f5669f41147919c9891124e4c805a7b762383 Mon Sep 17 00:00:00 2001 From: Ashwin Aji Date: Wed, 10 Jun 2026 19:04:01 +0000 Subject: [PATCH 27/31] ORBIT-2: add training, scaling, and perf-analysis recipes (sanitized). Ship Lux-oriented train/perf sbatch wrappers, config templates (PRISM and ERA5 same-dir), log parsing and scaling collate scripts, and perf-analysis recipe docs without site-specific paths or job metadata. Document public-repo hygiene in skills; allow the ORBIT perf HANDOFF to be tracked via .gitignore exception. Overlay verify tolerates missing xformers.ops with a clear warning. Co-authored-by: Cursor --- .../skills/ai4science-earth-science/SKILL.md | 4 + .../skills/ai4science-perf-analysis/SKILL.md | 16 +- .cursor/skills/ai4science-studio/SKILL.md | 2 +- .gitignore | 1 + .../models/ORBIT-2/examples/README.md | 35 +++ .../ORBIT-2/examples/build_overlay_amd.sh | 6 + .../ORBIT-2/examples/collate_scaling_study.py | 189 +++++++++++++ .../ORBIT-2/examples/interm_8m_lux.yaml | 128 +++++++++ .../ORBIT-2/examples/interm_8m_lux_era5.yaml | 143 ++++++++++ .../ORBIT-2/examples/orbit2_profiler_hook.py | 78 ++++++ .../ORBIT-2/examples/parse_training_log.py | 217 +++++++++++++++ .../ORBIT-2/examples/preflight_orbit2.py | 21 +- .../ORBIT-2/examples/render_orbit2_config.py | 121 +++++++++ .../ORBIT-2/examples/run_fom_extractor.py | 66 +++++ .../ORBIT-2/examples/run_orbit2_train.py | 188 +++++++++++++ .../ORBIT-2/examples/run_scaling_study.sh | 83 ++++++ .../ORBIT-2/examples/sbatch_train_amd.sh | 250 ++++++++++++++++++ .../ORBIT-2/examples/sbatch_train_perf_amd.sh | 250 ++++++++++++++++++ earth_science/models/ORBIT-2/model.yaml | 10 + .../ORBIT-2/recipes/perf-analysis/HANDOFF.md | 54 ++++ .../ORBIT-2/recipes/perf-analysis/README.md | 45 ++++ .../recipes/perf-analysis/agents/launcher.md | 215 +++++++++++++++ .../perf-analysis/agents/omnistat_analyst.md | 143 ++++++++++ .../perf-analysis/agents/omnistat_verifier.md | 92 +++++++ .../perf-analysis/agents/synthesizer.md | 131 +++++++++ .../perf-analysis/agents/tracelens_analyst.md | 124 +++++++++ .../agents/tracelens_verifier.md | 106 ++++++++ .../perf-analysis/omnistat.config.template | 133 ++++++++++ .../recipes/perf-optimizer-loop/README.md | 11 + 29 files changed, 2859 insertions(+), 3 deletions(-) mode change 100644 => 100755 earth_science/models/ORBIT-2/examples/build_overlay_amd.sh create mode 100755 earth_science/models/ORBIT-2/examples/collate_scaling_study.py create mode 100644 earth_science/models/ORBIT-2/examples/interm_8m_lux.yaml create mode 100644 earth_science/models/ORBIT-2/examples/interm_8m_lux_era5.yaml create mode 100644 earth_science/models/ORBIT-2/examples/orbit2_profiler_hook.py create mode 100755 earth_science/models/ORBIT-2/examples/parse_training_log.py create mode 100755 earth_science/models/ORBIT-2/examples/render_orbit2_config.py create mode 100755 earth_science/models/ORBIT-2/examples/run_fom_extractor.py create mode 100755 earth_science/models/ORBIT-2/examples/run_orbit2_train.py create mode 100755 earth_science/models/ORBIT-2/examples/run_scaling_study.sh create mode 100755 earth_science/models/ORBIT-2/examples/sbatch_train_amd.sh create mode 100755 earth_science/models/ORBIT-2/examples/sbatch_train_perf_amd.sh create mode 100644 earth_science/models/ORBIT-2/recipes/perf-analysis/HANDOFF.md create mode 100644 earth_science/models/ORBIT-2/recipes/perf-analysis/README.md create mode 100644 earth_science/models/ORBIT-2/recipes/perf-analysis/agents/launcher.md create mode 100644 earth_science/models/ORBIT-2/recipes/perf-analysis/agents/omnistat_analyst.md create mode 100644 earth_science/models/ORBIT-2/recipes/perf-analysis/agents/omnistat_verifier.md create mode 100644 earth_science/models/ORBIT-2/recipes/perf-analysis/agents/synthesizer.md create mode 100644 earth_science/models/ORBIT-2/recipes/perf-analysis/agents/tracelens_analyst.md create mode 100644 earth_science/models/ORBIT-2/recipes/perf-analysis/agents/tracelens_verifier.md create mode 100644 earth_science/models/ORBIT-2/recipes/perf-analysis/omnistat.config.template create mode 100644 earth_science/models/ORBIT-2/recipes/perf-optimizer-loop/README.md diff --git a/.cursor/skills/ai4science-earth-science/SKILL.md b/.cursor/skills/ai4science-earth-science/SKILL.md index f713f60..2fa2e62 100755 --- a/.cursor/skills/ai4science-earth-science/SKILL.md +++ b/.cursor/skills/ai4science-earth-science/SKILL.md @@ -51,6 +51,10 @@ The `earth_science/` domain covers **climate**, **weather**, and broader **Earth ``` - **Distributed launch:** `srun --mpi=pmix apptainer exec ... --env HOSTNAME="$MASTER_ADDR"` — see studio SKILL.md for full pattern. Confirmed working for 1-GPU and 8-GPU. - **Data blocker:** ORBIT-2 real training data lives on Frontier Lustre (`/lustre/orion//...`) — not publicly accessible. Synthetic data covers the smoke-test path. +- **Training on institutional clusters:** `sbatch_train_amd.sh` + `run_orbit2_train.py` wrap `intermediate_downscaling.py` with Apptainer overlay, `srun --mpi=pmix`, and `interm_8m_lux.yaml`. Staged Constellation/Globus 10.0_arcmin NPZ data can be used in **same-dir mode** (`low_res` = `high_res`, `superres_mag: 1`) for perf/scaling timing only. +- **Training landmines (2026-06):** stub `gptl4py` in `run_orbit2_train.py`; overlay may lack `xformers.ops` → auto `FusedAttn.DEFAULT`; `max_epochs` must be ≥ 2; batch cap patches module before `main()` not via `runpy.run_path`. +- **Perf/scaling:** `sbatch_train_perf_amd.sh`, `run_scaling_study.sh`, artifacts under `$AI4S_SHARED_DIR/models/ORBIT-2/perf-runs/` and `outputs/scaling-*`. See [recipes/perf-analysis/HANDOFF.md](../../../earth_science/models/ORBIT-2/recipes/perf-analysis/HANDOFF.md). +- **Multi-node scheduling:** Before large submits, check partition state (`sinfo`, site dashboards). Prefer **idle, healthy** nodes; use `--nodelist` / `--exclude` only per your site's policy. Do not stack two exclusive full-node GPU jobs on the same host. ### ArchesWeather (geoarches) diff --git a/.cursor/skills/ai4science-perf-analysis/SKILL.md b/.cursor/skills/ai4science-perf-analysis/SKILL.md index 0c27543..ae1fa20 100644 --- a/.cursor/skills/ai4science-perf-analysis/SKILL.md +++ b/.cursor/skills/ai4science-perf-analysis/SKILL.md @@ -9,16 +9,30 @@ description: Runs the AMD AI agents bottleneck-analysis workflow on a HydraGNN ( The user wants an automated bottleneck analysis of a multi-node training/inference run using AMD's open-source observability tooling (TraceLens, Omnistat). This is **distinct** from the `ai4science-run-models` skill — that one launches a model; this one diagnoses a model after it runs. -Default target: HydraGNN on AMD MI355X. The same pattern can extend to ORBIT-2 / StormCast / GP-MoLFormer in iteration 2. +Default target: HydraGNN on AMD MI355X. **ORBIT-2 training** uses the same perf-analysis pattern; see `earth_science/models/ORBIT-2/recipes/perf-analysis/`. ## Repository entry points +### HydraGNN + - Recipe: [material_science/models/HydraGNN/recipes/perf-analysis/](../../../material_science/models/HydraGNN/recipes/perf-analysis/) - Iterative-loop recipe: [material_science/models/HydraGNN/recipes/perf-optimizer-loop/](../../../material_science/models/HydraGNN/recipes/perf-optimizer-loop/) — **see [`HANDOFF.md`](../../../material_science/models/HydraGNN/recipes/perf-optimizer-loop/HANDOFF.md) for current state, what works, what's blocked, and next-work priorities** - Sbatch wrapper: [material_science/models/HydraGNN/examples/sbatch_train_perf_amd.sh](../../../material_science/models/HydraGNN/examples/sbatch_train_perf_amd.sh) - Reusable node-health probe: [material_science/models/HydraGNN/examples/microbench_node_health.sh](../../../material_science/models/HydraGNN/examples/microbench_node_health.sh) - Agent prompt files: `material_science/models/HydraGNN/recipes/perf-analysis/agents/*.md` and `material_science/models/HydraGNN/recipes/perf-optimizer-loop/agents/*.md` +### ORBIT-2 + +- Recipe: [earth_science/models/ORBIT-2/recipes/perf-analysis/](../../../earth_science/models/ORBIT-2/recipes/perf-analysis/) — see [`HANDOFF.md`](../../../earth_science/models/ORBIT-2/recipes/perf-analysis/HANDOFF.md) +- Sbatch: [earth_science/models/ORBIT-2/examples/sbatch_train_perf_amd.sh](../../../earth_science/models/ORBIT-2/examples/sbatch_train_perf_amd.sh) +- Plain training: [earth_science/models/ORBIT-2/examples/sbatch_train_amd.sh](../../../earth_science/models/ORBIT-2/examples/sbatch_train_amd.sh) +- Scaling: [run_scaling_study.sh](../../../earth_science/models/ORBIT-2/examples/run_scaling_study.sh) + [collate_scaling_study.py](../../../earth_science/models/ORBIT-2/examples/collate_scaling_study.py) +- FOM parser: [parse_training_log.py](../../../earth_science/models/ORBIT-2/examples/parse_training_log.py) — primary FOM **`steady_batch_time_s`** (epochs ≥ 2, exclude batch 0 per epoch); **`loss_sanity_pass`** requires strictly decreasing epoch losses +- Artifacts: `$AI4S_SHARED_DIR/models/ORBIT-2/perf-runs//` +- Workload: `intermediate_downscaling.py` via [run_orbit2_train.py](../../../earth_science/models/ORBIT-2/examples/run_orbit2_train.py) (gptl4py stub, FusedAttn fallback, batch cap) +- Data: 10.0_arcmin PRISM same-dir (`interm_8m_lux.yaml`) or ERA5 1.0° same-dir (`interm_8m_lux_era5.yaml`, `ORBIT2_CONFIG_TEMPLATE`) — timing only until true downscaling targets are staged +- Profiler: [orbit2_profiler_hook.py](../../../earth_science/models/ORBIT-2/examples/orbit2_profiler_hook.py) via `ORBIT2_RANK_PRE_TRAIN_HOOK` (no JSON Profile block) + ## Orchestration loop (what the main agent does) 1. **Read the recipe README** to refresh on the run topology and artifact layout. diff --git a/.cursor/skills/ai4science-studio/SKILL.md b/.cursor/skills/ai4science-studio/SKILL.md index e9cca33..161c671 100755 --- a/.cursor/skills/ai4science-studio/SKILL.md +++ b/.cursor/skills/ai4science-studio/SKILL.md @@ -313,7 +313,7 @@ srun --mpi=pmix apptainer exec \ **Why `HOSTNAME` instead of `MASTER_ADDR`:** Several model scripts do `os.environ["MASTER_ADDR"] = os.environ["HOSTNAME"]`. Inject rank-0's hostname as `HOSTNAME` so this pattern resolves correctly on all nodes. Using `localhost` only works single-node. -**Multi-node scaling:** Change `--nodes=N --ntasks=N*` in the SBATCH header; the `srun` command is unchanged. For HydraGNN matched strong-scaling studies, use `material_science/models/HydraGNN/examples/run_scaling_study.sh` + `collate_scaling_study.py` (see material-science SKILL and `recipes/train/README.md` §6). +**Multi-node scaling:** Change `--nodes=N --ntasks=N*` in the SBATCH header; the `srun` command is unchanged. For HydraGNN matched strong-scaling studies, use `material_science/models/HydraGNN/examples/run_scaling_study.sh` + `collate_scaling_study.py` (see material-science SKILL and `recipes/train/README.md` §6). For ORBIT-2, use `earth_science/models/ORBIT-2/examples/run_scaling_study.sh` (set `ORBIT2_SCALING_TAG` and `ORBIT2_CONFIG_TEMPLATE` when comparing PRISM vs ERA5 same-dir baselines). Keep recipe `HANDOFF.md` and agent prompts free of site-specific paths, job IDs, and hostnames—use `$AI4S_SHARED_DIR` and generic placeholders in anything committed to the repo. ## GPU visibility: `--gpus-per-node` NOT `--gpus-per-task` (MI355X / RCCL) diff --git a/.gitignore b/.gitignore index af9457e..49382e1 100755 --- a/.gitignore +++ b/.gitignore @@ -67,6 +67,7 @@ vic_server.log # Session-scoped agent handoff files (ephemeral, not for the repo) HANDOFF.md +!earth_science/models/ORBIT-2/recipes/perf-analysis/HANDOFF.md .claude/settings.local.json nul diff --git a/earth_science/models/ORBIT-2/examples/README.md b/earth_science/models/ORBIT-2/examples/README.md index 6283a86..e603c27 100644 --- a/earth_science/models/ORBIT-2/examples/README.md +++ b/earth_science/models/ORBIT-2/examples/README.md @@ -15,6 +15,41 @@ with AMD/ROCm container launch, overlay builds, and synthetic smoke tests. | [`sbatch_infer_amd.sh`](sbatch_infer_amd.sh) | SLURM driver for inference on AMD Instinct (MI250X/MI300X/MI350X) | | [`build_overlay_amd.sh`](build_overlay_amd.sh) | One-time overlay build (pre-bakes pip deps, skips ~15 min per job) | | [`interm_8m_synthetic.yaml`](interm_8m_synthetic.yaml) | Template YAML for synthetic mode | +| [`sbatch_train_amd.sh`](sbatch_train_amd.sh) | Multi-node training (Apptainer + MPI) | +| [`sbatch_train_perf_amd.sh`](sbatch_train_perf_amd.sh) | 2-node instrumented perf-analysis run | +| [`run_scaling_study.sh`](run_scaling_study.sh) | Submit matched 1/2/4/8-node scaling sweep | +| [`collate_scaling_study.py`](collate_scaling_study.py) | Parse scaling logs → CSV/MD table | +| [`parse_training_log.py`](parse_training_log.py) | Extract batch/epoch metrics from SLURM log | +| [`interm_8m_lux.yaml`](interm_8m_lux.yaml) | Lux config template (10.0_arcmin same-dir) | +| [`interm_8m_lux_era5.yaml`](interm_8m_lux_era5.yaml) | Lux config template (ERA5 1.0_deg same-dir sanity) | + +## Quick start — ERA5 1.0_deg sanity (new dataset) + +Verifies the staged ERA5 NPZ tree loads and training runs (same-dir mode; loss not meaningful). + +```bash +export AI4S_SHARED_DIR=/path/to/shared +export ORBIT2_DATA_ROOT=$AI4S_SHARED_DIR/models/ORBIT-2/data/superres/era5/1.0_deg +export ORBIT2_CONFIG_TEMPLATE=interm_8m_lux_era5.yaml +export ORBIT2_MAX_EPOCH=2 ORBIT2_MAX_BATCHES=5 ORBIT2_BATCH_SIZE=2 ORBIT2_DATA_TYPE=float32 +sbatch sbatch_train_amd.sh +``` + +For production ERA5→PRISM downscaling you will pair `era5/1.0_deg` (low) with `prism/2.5_arcmin` (high) once both splits and year shards align. + +## Quick start — training (10.0_arcmin PRISM, scaling) + +```bash +export AI4S_SHARED_DIR=/path/to/shared +export ORBIT2_DATA_ROOT=$AI4S_SHARED_DIR/models/ORBIT-2/data/superres/prism/10.0_arcmin +export ORBIT2_MAX_EPOCH=2 ORBIT2_MAX_BATCHES=10 ORBIT2_DATA_TYPE=float32 +sbatch sbatch_train_amd.sh + +# Strong scaling sweep: +export SBATCH_PARTITION=... SBATCH_ACCOUNT=... +./run_scaling_study.sh +python3 collate_scaling_study.py --log-dir . --jobs ,,... -o scaling_study +``` ## Quick start — Docker diff --git a/earth_science/models/ORBIT-2/examples/build_overlay_amd.sh b/earth_science/models/ORBIT-2/examples/build_overlay_amd.sh old mode 100644 new mode 100755 index c22afb0..3e6929e --- a/earth_science/models/ORBIT-2/examples/build_overlay_amd.sh +++ b/earth_science/models/ORBIT-2/examples/build_overlay_amd.sh @@ -226,6 +226,12 @@ echo "--- Verifying ---" PYTHONPATH="\$PKG" python3 -c " import xformers, mpi4py, pytorch_lightning, torch print('xformers :', xformers.__version__) +print('xformers.ops :', hasattr(xformers, 'ops')) +try: + import xformers.ops # noqa: F401 + print('xformers.ops : import OK') +except Exception as e: + print('xformers.ops : NOT available (', type(e).__name__, ') — use ORBIT2_DATA_TYPE=float32 for training until CK path is present') print('mpi4py :', mpi4py.__version__) print('pytorch_lightning:', pytorch_lightning.__version__) print('torch :', torch.__version__) diff --git a/earth_science/models/ORBIT-2/examples/collate_scaling_study.py b/earth_science/models/ORBIT-2/examples/collate_scaling_study.py new file mode 100755 index 0000000..ffd28e1 --- /dev/null +++ b/earth_science/models/ORBIT-2/examples/collate_scaling_study.py @@ -0,0 +1,189 @@ +#!/usr/bin/env python3 +"""Build ORBIT-2 strong-scaling table from orbit2-train-*.out logs.""" + +from __future__ import annotations + +import argparse +import csv +import json +import re +import sys +from pathlib import Path + +sys.path.insert(0, str(Path(__file__).resolve().parent)) +from parse_training_log import aggregate_foms, parse_slurm_log # noqa: E402 + + +def _parse_banner(log_path: Path) -> dict: + meta: dict = {} + patterns = { + "nodes": r"^\s*Nodes\s+:\s*(\d+)", + "total_ranks": r"^\s*Total ranks\s+:\s*(\d+)", + "batch_size": r"^\s*Batch size\s+:\s*(\d+)", + "max_epochs": r"^\s*Max epochs\s+:\s*(\d+)", + "max_batches": r"^\s*Max batches\s+:\s*(\S+)", + "nodelist": r"^\s*Node\(s\)\s+:\s*(.+)", + } + with log_path.open(encoding="utf-8", errors="replace") as f: + for line in f: + for key, pat in patterns.items(): + m = re.match(pat, line) + if m: + meta[key] = m.group(1).strip() + if "nodes" in meta: + meta["nodes"] = int(meta["nodes"]) + if "total_ranks" in meta: + meta["total_ranks"] = int(meta["total_ranks"]) + m = re.search(r"orbit2-train-(\d+)", log_path.name) + if m: + meta["job_id"] = m.group(1) + meta["log_path"] = str(log_path) + return meta + + +def main() -> int: + parser = argparse.ArgumentParser(description="Collate ORBIT-2 scaling study logs.") + parser.add_argument("--log", action="append", default=[], help="log file or N:path") + parser.add_argument("--log-dir", type=Path, default=Path(".")) + parser.add_argument("--jobs", type=str, default="", help="Comma-separated job IDs") + parser.add_argument("-o", "--output", type=str, default="scaling_study") + parser.add_argument( + "--steady-epoch-start", + type=int, + default=2, + help="Steady-state FOM uses batch times from this epoch onward (default: 2)", + ) + parser.add_argument( + "--warmup-batches-per-epoch", + type=int, + default=1, + help="Within steady epochs, exclude batch indices < N (default: 1)", + ) + parser.add_argument( + "--require-loss-sanity", + action="store_true", + help="Exit 1 if any job fails epoch-over-epoch loss decrease check", + ) + args = parser.parse_args() + + logs: list[tuple[str, Path]] = [] + for spec in args.log: + if ":" in spec: + label, path = spec.split(":", 1) + logs.append((label, Path(path))) + else: + logs.append(("", Path(spec))) + + if args.jobs: + for jid in args.jobs.split(","): + jid = jid.strip() + p = args.log_dir / f"orbit2-train-{jid}.out" + if p.is_file(): + logs.append((jid, p)) + + if not logs: + print("error: no logs found", file=sys.stderr) + return 2 + + rows = [] + baseline_throughput = None + any_sanity_fail = False + + for label, log_path in logs: + if not log_path.is_file(): + print(f"warning: missing {log_path}", file=sys.stderr) + continue + meta = _parse_banner(log_path) + records = parse_slurm_log(log_path) + foms = aggregate_foms( + records, + warmup_batches_per_epoch=args.warmup_batches_per_epoch, + steady_epoch_start=args.steady_epoch_start, + ) + if foms.get("loss_sanity_pass") is False: + any_sanity_fail = True + print( + f"warning: job {label or meta.get('job_id')} failed loss sanity: " + f"{foms.get('loss_violations')}", + file=sys.stderr, + ) + + nodes = meta.get("nodes", 1) + batch_s = foms.get("batch_time_s") + batch_size = int(meta.get("batch_size", 4)) + ranks = meta.get("total_ranks", nodes * 8) + throughput = (batch_size * ranks / batch_s) if batch_s else None + if baseline_throughput is None and throughput: + baseline_throughput = throughput + efficiency = ( + (throughput / (baseline_throughput * nodes)) if throughput and baseline_throughput else None + ) + rows.append( + { + "label": label or meta.get("job_id", ""), + "nodes": nodes, + "ranks": ranks, + "batch_time_s": foms.get("batch_time_s"), + "warmup_batch_time_s": foms.get("warmup_batch_time_s"), + "steady_batch_count": foms.get("steady_batch_count"), + "throughput_samples_per_s": throughput, + "scaling_efficiency": efficiency, + "final_loss": foms.get("final_loss"), + "loss_sanity_pass": foms.get("loss_sanity_pass"), + "loss_curve": json.dumps(foms.get("loss_curve", [])), + "job_id": meta.get("job_id"), + "log_path": str(log_path), + } + ) + + rows.sort(key=lambda r: r["nodes"]) + out_base = Path(args.output) + csv_path = out_base.with_suffix(".csv") + md_path = out_base.with_suffix(".md") + json_path = out_base.with_suffix(".json") + + fieldnames = list(rows[0].keys()) if rows else [] + with csv_path.open("w", newline="", encoding="utf-8") as f: + w = csv.DictWriter(f, fieldnames=fieldnames) + if rows: + w.writeheader() + w.writerows(rows) + + lines = [ + "# ORBIT-2 scaling study", + "", + f"Steady-state FOM: mean batch wall time from epoch {args.steady_epoch_start}+ " + f"(exclude warmup epochs 0–{args.steady_epoch_start - 1}), " + f"excluding batch indices < {args.warmup_batches_per_epoch} within each epoch.", + "", + "Loss sanity: each epoch loss must be strictly less than the previous epoch.", + "", + "| nodes | ranks | steady batch_time_s | warmup batch_time_s | throughput | efficiency | loss OK |", + "|---:|---:|---:|---:|---:|---:|:---:|", + ] + for r in rows: + lp = r.get("loss_sanity_pass") + ok = "yes" if lp is True else ("n/a" if lp is None else "**no**") + lines.append( + f"| {r['nodes']} | {r['ranks']} | {r.get('batch_time_s', 'n/a')} | " + f"{r.get('warmup_batch_time_s', 'n/a')} | {r.get('throughput_samples_per_s', 'n/a')} | " + f"{r.get('scaling_efficiency', 'n/a')} | {ok} |" + ) + md_path.write_text("\n".join(lines) + "\n", encoding="utf-8") + + summary = { + "steady_epoch_start": args.steady_epoch_start, + "warmup_batches_per_epoch": args.warmup_batches_per_epoch, + "rows": rows, + "all_loss_sanity_pass": not any_sanity_fail, + } + json_path.write_text(json.dumps(summary, indent=2) + "\n", encoding="utf-8") + + print(f"Wrote {csv_path}, {md_path}, and {json_path}") + if args.require_loss_sanity and any_sanity_fail: + return 1 + return 0 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/earth_science/models/ORBIT-2/examples/interm_8m_lux.yaml b/earth_science/models/ORBIT-2/examples/interm_8m_lux.yaml new file mode 100644 index 0000000..2c35ba2 --- /dev/null +++ b/earth_science/models/ORBIT-2/examples/interm_8m_lux.yaml @@ -0,0 +1,128 @@ +# ORBIT-2 Lux/MI355X training config — real 10.0_arcmin PRISM grid (same-dir mode). +# +# Both low_res_dir and high_res_dir point at the same 10.0_arcmin tree so we can +# exercise the full training pipeline for perf/scaling without 2.5_arcmin targets. +# Loss and PSNR are NOT scientifically meaningful in this mode — use for timing only. +# +# Placeholders substituted by render_orbit2_config.py at job submit time: +# __DATA_ROOT__ — ORBIT2_DATA_ROOT (10.0_arcmin directory) +# __FSDP__ — fsdp parallelism (typically = node count) +# __SIMPLE_DDP__ — simple_ddp parallelism (typically 8) +# __MAX_EPOCHS__ — ORBIT2_MAX_EPOCH override +# __BATCH_SIZE__ — ORBIT2_BATCH_SIZE override +# __DATA_TYPE__ — ORBIT2_DATA_TYPE (float32 avoids xformers.ops CK path) + +#----------- TRAINER ------------- +trainer: + max_epochs: __MAX_EPOCHS__ + checkpoint: + pretrain: + batch_size: __BATCH_SIZE__ + buffer_size: 40 + num_workers: 0 + data_type: __DATA_TYPE__ + gpu_type: "amd" + train_loss: "bayesian_tv" + +#---------- Parallelism ------------- +parallelism: + fsdp: __FSDP__ + simple_ddp: __SIMPLE_DDP__ + tensor_par: 1 + seq_par: 1 + +### Tiling +tiling: + do_tiling: False + div: 4 + overlap: 3 + +# ---------------------------- MODEL ------------------------------------------- +model: + preset: res_slimvit + lr: 2e-3 + weight_decay: 1e-5 + beta_1: 0.9 + beta_2: 0.99 + warmup_epochs: 1 + warmup_start_lr: 1e-7 + eta_min: 1e-8 + # Same-dir 10.0_arcmin perf mode: no 4× upscale (loss not meaningful; timing only). + superres_mag: 1 + cnn_ratio: 4 + patch_size: 2 + embed_dim: 256 + depth: 6 + decoder_depth: 4 + num_heads: 4 + mlp_ratio: 4 + drop_path: 0.1 + drop_rate: 0.1 + +# ---------------------------- DATA ------------------------------------------- +data: + low_res_dir: { + 'PRISM': "__DATA_ROOT__", + } + high_res_dir: { + 'PRISM': "__DATA_ROOT__", + } + + spatial_resolution: { + 'PRISM': 18, + } + + default_vars: [ + "land_sea_mask", + "orography", + "lattitude", + "landcover", + "2m_temperature", + "2m_temperature_max", + "2m_temperature_min", + "temperature_200", + "temperature_500", + "temperature_850", + "10m_u_component_of_wind", + "u_component_of_wind_200", + "u_component_of_wind_500", + "u_component_of_wind_850", + "10m_v_component_of_wind", + "v_component_of_wind_200", + "v_component_of_wind_500", + "v_component_of_wind_850", + "specific_humidity_200", + "specific_humidity_500", + "specific_humidity_850", + "total_precipitation_24hr", + "volumetric_soil_water_layer_1", + ] + + dict_in_variables: { + 'PRISM': [ + "land_sea_mask", + "orography", + "lattitude", + "landcover", + "total_precipitation_24hr", + "2m_temperature_min", + "2m_temperature_max", + ], + } + + dict_out_variables: { + 'PRISM': [ + "total_precipitation_24hr", + "2m_temperature_min", + "2m_temperature_max", + ], + } + + var_weights: { + "2m_temperature": 10, + "10m_u_component_of_wind": 1, + "10m_v_component_of_wind": 1, + "total_precipitation_24hr": 1, + "2m_temperature_min": 10, + "2m_temperature_max": 10, + } diff --git a/earth_science/models/ORBIT-2/examples/interm_8m_lux_era5.yaml b/earth_science/models/ORBIT-2/examples/interm_8m_lux_era5.yaml new file mode 100644 index 0000000..dc22c14 --- /dev/null +++ b/earth_science/models/ORBIT-2/examples/interm_8m_lux_era5.yaml @@ -0,0 +1,143 @@ +# ORBIT-2 Lux/MI355X — ERA5 1.0_deg sanity / pipeline test (same-dir mode). +# +# Both low_res_dir and high_res_dir point at the same ERA5 tree so we can verify +# the new dataset loads and training runs before pairing with PRISM 2.5_arcmin targets. +# Loss is not scientifically meaningful in same-dir mode — use for data/pipeline sanity only. +# +# Placeholders substituted by render_orbit2_config.py: +# __DATA_ROOT__ — ORBIT2_DATA_ROOT (era5/1.0_deg directory) +# __FSDP__ — fsdp parallelism (typically = node count) +# __SIMPLE_DDP__ — simple_ddp parallelism (typically 8) +# __MAX_EPOCHS__ — ORBIT2_MAX_EPOCH override +# __BATCH_SIZE__ — ORBIT2_BATCH_SIZE override +# __DATA_TYPE__ — ORBIT2_DATA_TYPE (float32 avoids xformers.ops CK path) + +#----------- TRAINER ------------- +trainer: + max_epochs: __MAX_EPOCHS__ + checkpoint: + pretrain: + batch_size: __BATCH_SIZE__ + buffer_size: 40 + num_workers: 0 + data_type: __DATA_TYPE__ + gpu_type: "amd" + train_loss: "bayesian_tv" + +#---------- Parallelism ------------- +parallelism: + fsdp: __FSDP__ + simple_ddp: __SIMPLE_DDP__ + tensor_par: 1 + seq_par: 1 + +### Tiling +tiling: + do_tiling: False + div: 4 + overlap: 3 + +# ---------------------------- MODEL ------------------------------------------- +model: + preset: res_slimvit + lr: 2e-3 + weight_decay: 1e-5 + beta_1: 0.9 + beta_2: 0.99 + warmup_epochs: 1 + warmup_start_lr: 1e-7 + eta_min: 1e-8 + superres_mag: 1 + cnn_ratio: 4 + patch_size: 2 + embed_dim: 256 + depth: 6 + decoder_depth: 4 + num_heads: 4 + mlp_ratio: 4 + drop_path: 0.1 + drop_rate: 0.1 + +# ---------------------------- DATA ------------------------------------------- +data: + low_res_dir: { + 'ERA5_2': "__DATA_ROOT__", + } + high_res_dir: { + 'ERA5_2': "__DATA_ROOT__", + } + + spatial_resolution: { + 'ERA5_2': 111, + } + + default_vars: [ + "land_sea_mask", + "orography", + "lattitude", + "landcover", + "2m_temperature", + "2m_temperature_max", + "2m_temperature_min", + "temperature_200", + "temperature_500", + "temperature_850", + "10m_u_component_of_wind", + "u_component_of_wind_200", + "u_component_of_wind_500", + "u_component_of_wind_850", + "10m_v_component_of_wind", + "v_component_of_wind_200", + "v_component_of_wind_500", + "v_component_of_wind_850", + "specific_humidity_200", + "specific_humidity_500", + "specific_humidity_850", + "total_precipitation_24hr", + "volumetric_soil_water_layer_1", + ] + + dict_in_variables: { + 'ERA5_2': [ + "land_sea_mask", + "orography", + "lattitude", + "landcover", + "2m_temperature", + "2m_temperature_max", + "2m_temperature_min", + "temperature_200", + "temperature_500", + "temperature_850", + "10m_u_component_of_wind", + "u_component_of_wind_200", + "u_component_of_wind_500", + "u_component_of_wind_850", + "10m_v_component_of_wind", + "v_component_of_wind_200", + "v_component_of_wind_500", + "v_component_of_wind_850", + "specific_humidity_200", + "specific_humidity_500", + "specific_humidity_850", + "total_precipitation_24hr", + "volumetric_soil_water_layer_1", + ], + } + + dict_out_variables: { + 'ERA5_2': [ + "total_precipitation_24hr", + "2m_temperature_min", + "2m_temperature_max", + ], + } + + var_weights: { + "2m_temperature": 10, + "10m_u_component_of_wind": 1, + "10m_v_component_of_wind": 1, + "total_precipitation_24hr": 1, + "2m_temperature_min": 10, + "2m_temperature_max": 10, + } diff --git a/earth_science/models/ORBIT-2/examples/orbit2_profiler_hook.py b/earth_science/models/ORBIT-2/examples/orbit2_profiler_hook.py new file mode 100644 index 0000000..6294199 --- /dev/null +++ b/earth_science/models/ORBIT-2/examples/orbit2_profiler_hook.py @@ -0,0 +1,78 @@ +"""Rank-0 PyTorch profiler hook for ORBIT-2 perf-analysis runs. + +Set ORBIT2_RANK_PRE_TRAIN_HOOK to this file path. Env: + ORBIT2_PROFILE_DIR Trace output directory (default: perf-runs//traces) + PROFILE_TARGET_EPOCH Epoch index to profile (default: 0) + PROFILE_RANK0_ONLY If 1, only rank 0 profiles (default: 1) +""" + +from __future__ import annotations + +import os +from pathlib import Path + +import intermediate_downscaling as ids +import torch + + +def _wrap_training(): + world_rank = int(os.environ.get("SLURM_PROCID", "0")) + rank0_only = os.environ.get("PROFILE_RANK0_ONLY", "1") == "1" + if rank0_only and world_rank != 0: + return + + target_epoch = int(os.environ.get("PROFILE_TARGET_EPOCH", "0")) + profile_dir = Path( + os.environ.get( + "ORBIT2_PROFILE_DIR", + os.environ.get("ORBIT2_OUTPUT_DIR", "/tmp") + "/traces", + ) + ) + + _orig = ids.run_training_epochs + + def _profiled(*args, **kwargs): + epoch_start = kwargs.get("epoch_start", args[7] if len(args) > 7 else 0) + epoch_end = kwargs.get("epoch_end", args[8] if len(args) > 8 else epoch_start + 1) + world_rank_kw = kwargs.get("world_rank", args[12] if len(args) > 12 else world_rank) + + if world_rank_kw != 0: + return _orig(*args, **kwargs) + + profile_dir.mkdir(parents=True, exist_ok=True) + activities = [ + torch.profiler.ProfilerActivity.CPU, + torch.profiler.ProfilerActivity.CUDA, + ] + + for epoch in range(epoch_start, epoch_end): + if epoch != target_epoch: + # Run single epoch without profiler + kwargs_ep = dict(kwargs) + kwargs_ep["epoch_start"] = epoch + kwargs_ep["epoch_end"] = epoch + 1 + _orig(*args, **kwargs_ep) + continue + + trace_path = profile_dir / f"orbit2-epoch{epoch}-rank0" + with torch.profiler.profile( + activities=activities, + record_shapes=True, + profile_memory=True, + with_stack=True, + on_trace_ready=torch.profiler.tensorboard_trace_handler(str(trace_path)), + ) as prof: + kwargs_ep = dict(kwargs) + kwargs_ep["epoch_start"] = epoch + kwargs_ep["epoch_end"] = epoch + 1 + result = _orig(*args, **kwargs_ep) + # on_trace_ready already saved the trace; do not call export_chrome_trace again + # (RuntimeError: Trace is already saved). + return result + + return epoch_end + + ids.run_training_epochs = _profiled + + +_wrap_training() diff --git a/earth_science/models/ORBIT-2/examples/parse_training_log.py b/earth_science/models/ORBIT-2/examples/parse_training_log.py new file mode 100755 index 0000000..1b5ef1c --- /dev/null +++ b/earth_science/models/ORBIT-2/examples/parse_training_log.py @@ -0,0 +1,217 @@ +#!/usr/bin/env python3 +"""Parse ORBIT-2 training metrics from SLURM stdout (orbit2-train-*.out). + +Upstream intermediate_downscaling.py emits on rank 0: + Starting epoch {epoch} + Batch {idx}: {seconds} seconds (every 10 batches — wall time for that step) + Epoch {epoch} completed. Loss: {value} +""" + +from __future__ import annotations + +import argparse +import csv +import json +import re +import sys +from pathlib import Path + +BATCH_RE = re.compile(r"^Batch (\d+): ([\d.]+) seconds") +EPOCH_DONE_RE = re.compile(r"^Epoch (\d+) completed\. Loss: ([\d.eE+-]+)") +EPOCH_START_RE = re.compile(r"^Starting epoch (\d+)") + + +def parse_slurm_log(log_path: Path) -> list[dict]: + records: list[dict] = [] + current_epoch: int | None = None + + with log_path.open(encoding="utf-8", errors="replace") as f: + for line in f: + line = line.strip() + + m = EPOCH_START_RE.match(line) + if m: + current_epoch = int(m.group(1)) + continue + + m = BATCH_RE.match(line) + if m: + records.append( + { + "kind": "batch", + "batch": int(m.group(1)), + "batch_time_s": float(m.group(2)), + "epoch": current_epoch, + "loss": None, + } + ) + continue + + m = EPOCH_DONE_RE.match(line) + if m: + epoch = int(m.group(1)) + records.append( + { + "kind": "epoch", + "batch": None, + "batch_time_s": None, + "epoch": epoch, + "loss": float(m.group(2)), + } + ) + current_epoch = epoch + 1 + return records + + +def check_loss_sanity( + epoch_losses: list[dict], + *, + require_strict_decrease: bool = True, +) -> dict: + """Verify epoch losses decrease epoch-over-epoch (crash detector, not science gate).""" + losses = [(r["epoch"], r["loss"]) for r in epoch_losses if r.get("loss") is not None] + losses.sort(key=lambda x: x[0]) + + violations: list[dict] = [] + for i in range(len(losses) - 1): + e0, l0 = losses[i] + e1, l1 = losses[i + 1] + if require_strict_decrease and l1 >= l0: + violations.append( + { + "from_epoch": e0, + "to_epoch": e1, + "loss_from": l0, + "loss_to": l1, + } + ) + + if len(losses) < 2: + return { + "loss_sanity_pass": None, + "loss_sanity_skipped": True, + "loss_curve": [{"epoch": e, "loss": l} for e, l in losses], + "loss_violations": [], + "loss_epoch_count": len(losses), + } + + return { + "loss_sanity_pass": len(violations) == 0, + "loss_sanity_skipped": False, + "loss_curve": [{"epoch": e, "loss": l} for e, l in losses], + "loss_violations": violations, + "loss_epoch_count": len(losses), + } + + +def aggregate_foms( + records: list[dict], + *, + warmup_batches_per_epoch: int = 1, + steady_epoch_start: int = 2, +) -> dict: + """Primary FOM: mean steady-state batch wall time (non-warmup epochs and batches).""" + batch_rows = [r for r in records if r["kind"] == "batch" and r["batch_time_s"] is not None] + epoch_rows = [r for r in records if r["kind"] == "epoch" and r["epoch"] is not None] + + warmup_batches = [ + r + for r in batch_rows + if r.get("epoch") is not None + and r["epoch"] < steady_epoch_start + ] + steady_batches = [ + r + for r in batch_rows + if r.get("epoch") is not None + and r["epoch"] >= steady_epoch_start + and r["batch"] >= warmup_batches_per_epoch + ] + + def _mean(rows: list[dict]) -> float | None: + if not rows: + return None + return sum(r["batch_time_s"] for r in rows) / len(rows) + + warmup_mean = _mean( + [r for r in batch_rows if r.get("epoch") is not None and r["epoch"] < steady_epoch_start] + + [r for r in batch_rows if r.get("epoch") is None] + ) + steady_mean = _mean(steady_batches) + + sanity = check_loss_sanity(epoch_rows) + + steady_epoch_losses = [r for r in epoch_rows if r["epoch"] >= steady_epoch_start] + final_loss = ( + steady_epoch_losses[-1]["loss"] + if steady_epoch_losses + else (epoch_rows[-1]["loss"] if epoch_rows else None) + ) + + return { + "primary_fom": "batch_time_s", + "batch_time_s": steady_mean, + "steady_batch_time_s": steady_mean, + "warmup_batch_time_s": warmup_mean, + "steady_batch_count": len(steady_batches), + "warmup_batch_count": len(warmup_batches), + "steady_epoch_start": steady_epoch_start, + "warmup_batches_per_epoch_excluded": warmup_batches_per_epoch, + "final_loss": final_loss, + "epoch_records": len(epoch_rows), + **sanity, + } + + +def main() -> int: + parser = argparse.ArgumentParser(description="Parse ORBIT-2 training SLURM log.") + parser.add_argument("--log", type=Path, required=True) + parser.add_argument("--output", type=Path, default=None, help="Write CSV of parsed records") + parser.add_argument("--json", type=Path, default=None, help="Write aggregated FOM JSON") + parser.add_argument( + "--steady-epoch-start", + type=int, + default=2, + help="First epoch included in steady-state batch FOM (default: 2)", + ) + parser.add_argument( + "--warmup-batches-per-epoch", + type=int, + default=1, + help="Exclude batch indices < N within each steady epoch (default: 1 = skip batch 0)", + ) + args = parser.parse_args() + + if not args.log.is_file(): + print(f"error: log not found: {args.log}", file=sys.stderr) + return 2 + + records = parse_slurm_log(args.log) + if not records: + print("warning: no batch/epoch records found", file=sys.stderr) + + if args.output: + args.output.parent.mkdir(parents=True, exist_ok=True) + with args.output.open("w", newline="", encoding="utf-8") as f: + w = csv.DictWriter( + f, + fieldnames=["kind", "batch", "batch_time_s", "epoch", "loss"], + ) + w.writeheader() + w.writerows(records) + + foms = aggregate_foms( + records, + warmup_batches_per_epoch=args.warmup_batches_per_epoch, + steady_epoch_start=args.steady_epoch_start, + ) + if args.json: + args.json.parent.mkdir(parents=True, exist_ok=True) + args.json.write_text(json.dumps(foms, indent=2) + "\n", encoding="utf-8") + + print(json.dumps(foms, indent=2)) + return 0 if foms.get("loss_sanity_pass") is not False else 1 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/earth_science/models/ORBIT-2/examples/preflight_orbit2.py b/earth_science/models/ORBIT-2/examples/preflight_orbit2.py index e118f71..3b3ae39 100755 --- a/earth_science/models/ORBIT-2/examples/preflight_orbit2.py +++ b/earth_science/models/ORBIT-2/examples/preflight_orbit2.py @@ -44,12 +44,31 @@ def main() -> int: return 2 examples = root / "examples" - for name in ("visualize.py", "utils.py"): + for name in ("visualize.py", "utils.py", "intermediate_downscaling.py"): p = examples / name if not p.is_file(): print(f"error: missing {p}", file=sys.stderr) return 2 + data_root = os.environ.get("ORBIT2_DATA_ROOT") + if data_root: + dp = Path(data_root).expanduser().resolve() + if not dp.is_dir(): + print(f"error: ORBIT2_DATA_ROOT not found: {dp}", file=sys.stderr) + return 2 + for sub in ("train", "val", "test"): + if not (dp / sub).is_dir(): + print(f"error: missing data split: {dp / sub}", file=sys.stderr) + return 2 + + overlay = os.environ.get("ORBIT2_OVERLAY") + if overlay and not Path(overlay).expanduser().is_file(): + print(f"warning: ORBIT2_OVERLAY not found: {overlay}", file=sys.stderr) + + sif = os.environ.get("ORBIT2_SIF") + if sif and not Path(sif).expanduser().is_file(): + print(f"warning: ORBIT2_SIF not found: {sif}", file=sys.stderr) + src = root / "src" if not src.is_dir(): print(f"error: missing src tree: {src}", file=sys.stderr) diff --git a/earth_science/models/ORBIT-2/examples/render_orbit2_config.py b/earth_science/models/ORBIT-2/examples/render_orbit2_config.py new file mode 100755 index 0000000..265170e --- /dev/null +++ b/earth_science/models/ORBIT-2/examples/render_orbit2_config.py @@ -0,0 +1,121 @@ +#!/usr/bin/env python3 +"""Render a per-job ORBIT-2 training YAML from interm_8m_lux.yaml. + +Substitutes parallelism (fsdp × simple_ddp = nodes × 8), data root, and trainer caps. + +Usage: + python3 render_orbit2_config.py --nodes 2 --output /tmp/orbit2_config.yaml + python3 render_orbit2_config.py --nodes 8 \\ + --data-root $AI4S_SHARED_DIR/models/ORBIT-2/data/superres/prism/10.0_arcmin \\ + --max-epochs 6 --batch-size 8 -o config.yaml +""" + +from __future__ import annotations + +import argparse +import sys +from pathlib import Path + + +def main() -> int: + parser = argparse.ArgumentParser(description="Render ORBIT-2 Lux training config.") + parser.add_argument( + "--template", + type=Path, + default=None, + help="Template YAML (default: interm_8m_lux.yaml next to this script)", + ) + parser.add_argument("--nodes", type=int, required=True, help="SLURM node count") + parser.add_argument( + "--gpus-per-node", + type=int, + default=8, + help="GPUs per node (default: 8 for MI355X)", + ) + parser.add_argument( + "--data-root", + type=str, + default=None, + help="PRISM data root (default: env ORBIT2_DATA_ROOT)", + ) + parser.add_argument("--max-epochs", type=int, default=None, help="trainer.max_epochs") + parser.add_argument("--batch-size", type=int, default=None, help="trainer.batch_size") + parser.add_argument("--fsdp", type=int, default=None, help="Override fsdp (default: nodes)") + parser.add_argument( + "--simple-ddp", + type=int, + default=None, + help="Override simple_ddp (default: gpus_per_node)", + ) + parser.add_argument("-o", "--output", type=Path, required=True, help="Output YAML path") + args = parser.parse_args() + + import os + + script_dir = Path(__file__).resolve().parent + template = args.template or (script_dir / "interm_8m_lux.yaml") + if not template.is_file(): + print(f"error: template not found: {template}", file=sys.stderr) + return 2 + + data_root = args.data_root or os.environ.get("ORBIT2_DATA_ROOT") + if not data_root: + print("error: set --data-root or ORBIT2_DATA_ROOT", file=sys.stderr) + return 2 + + data_path = Path(data_root).resolve() + if not data_path.is_dir(): + print(f"error: data root not found: {data_path}", file=sys.stderr) + return 2 + + fsdp = args.fsdp if args.fsdp is not None else args.nodes + simple_ddp = args.simple_ddp if args.simple_ddp is not None else args.gpus_per_node + total = fsdp * simple_ddp + expected = args.nodes * args.gpus_per_node + if total != expected: + print( + f"error: fsdp({fsdp}) × simple_ddp({simple_ddp}) = {total}, " + f"expected {expected} (= nodes × gpus_per_node)", + file=sys.stderr, + ) + return 2 + + max_epochs = args.max_epochs + if max_epochs is None: + max_epochs = int(os.environ.get("ORBIT2_MAX_EPOCH", "3")) + if max_epochs < 2: + print( + "error: max_epochs must be >= 2 (upstream trains while epoch_start+1 < max_epochs)", + file=sys.stderr, + ) + return 2 + + batch_size = args.batch_size + if batch_size is None: + batch_size = int(os.environ.get("ORBIT2_BATCH_SIZE", "8")) + + data_type = os.environ.get("ORBIT2_DATA_TYPE", "float32") + if data_type not in ("float32", "bfloat16"): + print("error: ORBIT2_DATA_TYPE must be float32 or bfloat16", file=sys.stderr) + return 2 + + text = template.read_text(encoding="utf-8") + replacements = { + "__DATA_ROOT__": str(data_path), + "__FSDP__": str(fsdp), + "__SIMPLE_DDP__": str(simple_ddp), + "__MAX_EPOCHS__": str(max_epochs), + "__BATCH_SIZE__": str(batch_size), + "__DATA_TYPE__": data_type, + } + for key, val in replacements.items(): + text = text.replace(key, val) + + args.output.parent.mkdir(parents=True, exist_ok=True) + args.output.write_text(text, encoding="utf-8") + print(f"Wrote {args.output} (fsdp={fsdp}, simple_ddp={simple_ddp}, epochs={max_epochs}, batch={batch_size})") + return 0 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/earth_science/models/ORBIT-2/examples/run_fom_extractor.py b/earth_science/models/ORBIT-2/examples/run_fom_extractor.py new file mode 100755 index 0000000..0cc80c7 --- /dev/null +++ b/earth_science/models/ORBIT-2/examples/run_fom_extractor.py @@ -0,0 +1,66 @@ +#!/usr/bin/env python3 +"""Extract ORBIT-2 training FOMs from a perf-run job directory.""" + +from __future__ import annotations + +import argparse +import json +import os +import sys +from pathlib import Path + +sys.path.insert(0, str(Path(__file__).resolve().parent)) +from parse_training_log import aggregate_foms, parse_slurm_log # noqa: E402 + + +def main() -> int: + parser = argparse.ArgumentParser(description="ORBIT-2 perf-run FOM extractor.") + parser.add_argument("--job-dir", type=Path, required=True, help="perf-runs//") + parser.add_argument("--steady-epoch-start", type=int, default=2) + parser.add_argument("--warmup-batches-per-epoch", type=int, default=1) + args = parser.parse_args() + + job_dir = args.job_dir + job_id = job_dir.name + log_candidates = [ + job_dir / f"orbit2-train-{job_id}.out", + Path.cwd() / f"orbit2-train-{job_id}.out", + ] + shared_root = os.environ.get("AI4S_SHARED_DIR") + if shared_root: + log_candidates.insert( + 1, + Path(shared_root) + / "models" + / "ORBIT-2" + / "outputs" + / "train" + / "logs" + / f"orbit2-train-{job_id}.out", + ) + log_path = next((p for p in log_candidates if p.is_file()), None) + if log_path is None: + print(f"error: SLURM log not found for job {job_id}", file=sys.stderr) + return 2 + + records = parse_slurm_log(log_path) + foms = aggregate_foms( + records, + warmup_batches_per_epoch=args.warmup_batches_per_epoch, + steady_epoch_start=args.steady_epoch_start, + ) + foms["job_id"] = job_id + foms["log_path"] = str(log_path) + + out = job_dir / "foms.json" + out.write_text(json.dumps(foms, indent=2) + "\n", encoding="utf-8") + print(f"Wrote {out}") + print(json.dumps(foms, indent=2)) + if foms.get("loss_sanity_pass") is False: + print("warning: loss sanity check failed", file=sys.stderr) + return 1 + return 0 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/earth_science/models/ORBIT-2/examples/run_orbit2_train.py b/earth_science/models/ORBIT-2/examples/run_orbit2_train.py new file mode 100755 index 0000000..ed41f0d --- /dev/null +++ b/earth_science/models/ORBIT-2/examples/run_orbit2_train.py @@ -0,0 +1,188 @@ +#!/usr/bin/env python3 +"""Launch upstream ORBIT-2 intermediate_downscaling.py with optional training caps. + +Patches run_training_epochs to honour ORBIT2_MAX_BATCHES (upstream has no batch cap). +Config path is passed as argv[1] after this script's own args. + +Environment: + ORBIT2_MAX_BATCHES Cap batches per epoch (0 = unlimited) + ORBIT2_RANK_PRE_TRAIN_HOOK Optional path to a .py hook run before training + ORBIT2_FUSED_ATTN=DEFAULT Force PyTorch SDPA instead of xformers CK path +""" + +from __future__ import annotations + +import os +import runpy +import sys +import types +from datetime import timedelta +from pathlib import Path + + +def _stub_gptl4py() -> None: + """Upstream training imports gptl4py (Frontier GPTL); stub for ROCm containers.""" + if "gptl4py" in sys.modules: + return + + class _Stub: + @staticmethod + def start(_name: str) -> None: + return None + + @staticmethod + def stop(_name: str) -> None: + return None + + mod = types.ModuleType("gptl4py") + mod.start = _Stub.start + mod.stop = _Stub.stop + sys.modules["gptl4py"] = mod + + +def _patch_fused_attn_if_needed(ids) -> None: + """Use PyTorch SDPA when overlay xformers lacks .ops (CK path unavailable).""" + force = os.environ.get("ORBIT2_FUSED_ATTN", "").upper() == "DEFAULT" + try: + import xformers # noqa: F401 + + if not hasattr(xformers, "ops"): + force = True + else: + import xformers.ops # noqa: F401 + except (ImportError, AttributeError): + force = True + + if not force: + return + + from climate_learn.utils.fused_attn import FusedAttn + + class _PatchedFA: + CK = FusedAttn.DEFAULT + DEFAULT = FusedAttn.DEFAULT + NONE = FusedAttn.NONE + + ids.FusedAttn = _PatchedFA + + +def _apply_batch_cap(ids) -> None: + max_batches = int(os.environ.get("ORBIT2_MAX_BATCHES", "0") or 0) + if max_batches <= 0: + return + + import itertools + + _orig = ids.run_training_epochs + + def _cap_dataloader(dl, limit: int): + class _Capped: + def __init__(self, inner, n: int): + self._inner = inner + self._n = n + + def __iter__(self): + return iter(itertools.islice(iter(self._inner), self._n)) + + def __len__(self): + try: + return min(len(self._inner), self._n) + except TypeError: + return self._n + + return _Capped(dl, limit) + + def _capped(*args, **kwargs): + dl = kwargs.get("train_dataloader") + if dl is None and len(args) > 4: + dl = args[4] + args = list(args) + args[4] = _cap_dataloader(dl, max_batches) + args = tuple(args) + elif dl is not None: + kwargs["train_dataloader"] = _cap_dataloader(dl, max_batches) + world_rank = kwargs.get("world_rank", args[12] if len(args) > 12 else -1) + if world_rank == 0: + print(f"[run_orbit2_train] ORBIT2_MAX_BATCHES cap: {max_batches}", flush=True) + return _orig(*args, **kwargs) + + ids.run_training_epochs = _capped + + +def _run_hook() -> None: + hook = os.environ.get("ORBIT2_RANK_PRE_TRAIN_HOOK", "").strip() + if not hook: + return + hook_path = Path(hook) + if not hook_path.is_file(): + raise FileNotFoundError(f"ORBIT2_RANK_PRE_TRAIN_HOOK not found: {hook}") + # Profiler hook patches rank-0 training only; skip import on other ranks. + rank0_only = os.environ.get("PROFILE_RANK0_ONLY", "1") == "1" + world_rank = int(os.environ.get("SLURM_PROCID", "0")) + if rank0_only and world_rank != 0: + return + runpy.run_path(str(hook_path), run_name="__orbit2_hook__") + + +def main() -> None: + if len(sys.argv) < 2: + print("usage: run_orbit2_train.py ", file=sys.stderr) + sys.exit(2) + + config = sys.argv[1] + orbit2_root = os.environ.get("ORBIT2_ROOT") + if not orbit2_root: + print("error: ORBIT2_ROOT must be set", file=sys.stderr) + sys.exit(2) + + examples_dir = Path(orbit2_root) / "examples" + if not examples_dir.is_dir(): + print(f"error: missing examples dir: {examples_dir}", file=sys.stderr) + sys.exit(2) + + src_dir = Path(orbit2_root) / "src" + for p in (str(src_dir), str(examples_dir), orbit2_root): + if p not in sys.path: + sys.path.insert(0, p) + + os.chdir(examples_dir) + _stub_gptl4py() + + import torch + import torch.distributed as dist + + import intermediate_downscaling as ids # noqa: WPS433 + + _patch_fused_attn_if_needed(ids) + _apply_batch_cap(ids) + _run_hook() + + sys.argv = ["intermediate_downscaling.py", config] + + os.environ["MASTER_ADDR"] = str(os.environ["HOSTNAME"]) + os.environ["MASTER_PORT"] = os.environ.get("MASTER_PORT", "29500") + os.environ["WORLD_SIZE"] = os.environ["SLURM_NTASKS"] + os.environ["RANK"] = os.environ["SLURM_PROCID"] + + local_rank = int(os.environ["SLURM_LOCALID"]) + os.environ["LOCAL_RANK"] = str(local_rank) + torch.cuda.set_device(local_rank) + device = torch.cuda.current_device() + + dist.init_process_group( + "nccl", + timeout=timedelta(seconds=7200000), + rank=int(os.environ["SLURM_PROCID"]), + world_size=int(os.environ["SLURM_NTASKS"]), + ) + + print("Using dist.init_process_group. world_size ", os.environ["SLURM_NTASKS"], flush=True) + + try: + ids.main(device) + finally: + dist.destroy_process_group() + + +if __name__ == "__main__": + main() diff --git a/earth_science/models/ORBIT-2/examples/run_scaling_study.sh b/earth_science/models/ORBIT-2/examples/run_scaling_study.sh new file mode 100755 index 0000000..a9e346a --- /dev/null +++ b/earth_science/models/ORBIT-2/examples/run_scaling_study.sh @@ -0,0 +1,83 @@ +#!/usr/bin/env bash +# Submit matched ORBIT-2 strong-scaling sweep (1/2/4/8 nodes). +# +# Usage: +# export AI4S_SHARED_DIR=/path/to/shared +# export ORBIT2_DATA_ROOT=$AI4S_SHARED_DIR/models/ORBIT-2/data/superres/prism/10.0_arcmin +# ./run_scaling_study.sh +# ./run_scaling_study.sh --nodes 1,2 +# +# ERA5 1.0_deg same-dir (timing only): +# export ORBIT2_DATA_ROOT=$AI4S_SHARED_DIR/models/ORBIT-2/data/superres/era5/1.0_deg +# export ORBIT2_CONFIG_TEMPLATE=interm_8m_lux_era5.yaml +# export ORBIT2_SCALING_TAG=era5 +# ./run_scaling_study.sh + +set -euo pipefail + +SCRIPT_DIR=$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd) +: "${AI4S_SHARED_DIR:?Set AI4S_SHARED_DIR}" + +export TORCH_NCCL_HIGH_PRIORITY=1 +export GPU_MAX_HW_QUEUES=2 +export ORBIT2_DATA_TYPE=float32 +# max_epochs=6 → trains epochs 0–4; collate uses steady epochs 2–4 for FOM +export ORBIT2_MAX_EPOCH=6 +export ORBIT2_MAX_BATCHES=20 +export ORBIT2_BATCH_SIZE=4 +export ORBIT2_DATA_ROOT="${ORBIT2_DATA_ROOT:-${AI4S_SHARED_DIR}/models/ORBIT-2/data/superres/prism/10.0_arcmin}" +export ORBIT2_CONFIG_TEMPLATE="${ORBIT2_CONFIG_TEMPLATE:-interm_8m_lux.yaml}" +export ORBIT2_SCALING_TAG="${ORBIT2_SCALING_TAG:-prism}" +export ORBIT2_OUTPUT_DIR="${ORBIT2_OUTPUT_DIR:-${AI4S_SHARED_DIR}/models/ORBIT-2/outputs/scaling-${ORBIT2_SCALING_TAG}}" + +SBATCH_PARTITION="${SBATCH_PARTITION:-YOUR_PARTITION_HERE}" +SBATCH_ACCOUNT="${SBATCH_ACCOUNT:-YOUR_ACCOUNT_HERE}" + +NODES_LIST="1,2,4,8" +while [[ $# -gt 0 ]]; do + case "$1" in + --nodes) NODES_LIST="$2"; shift 2 ;; + *) echo "Unknown arg: $1" >&2; exit 1 ;; + esac +done + +LOG_DIR="${ORBIT2_SLURM_LOG_DIR:-${AI4S_SHARED_DIR}/models/ORBIT-2/outputs/train/logs}" +mkdir -p "$LOG_DIR" + +submit() { + local n="$1" t="$2" + sbatch --parsable \ + --partition="$SBATCH_PARTITION" \ + --account="$SBATCH_ACCOUNT" \ + --nodes="$n" \ + --time="$t" \ + --job-name="orbit2-scale-${n}N" \ + --output="${LOG_DIR}/orbit2-train-%j.out" \ + --error="${LOG_DIR}/orbit2-train-%j.out" \ + --export=ALL,TORCH_NCCL_HIGH_PRIORITY,GPU_MAX_HW_QUEUES,ORBIT2_DATA_TYPE,ORBIT2_MAX_EPOCH,ORBIT2_MAX_BATCHES,ORBIT2_BATCH_SIZE,ORBIT2_DATA_ROOT,ORBIT2_CONFIG_TEMPLATE,ORBIT2_SCALING_TAG,ORBIT2_OUTPUT_DIR,AI4S_SHARED_DIR \ + "$SCRIPT_DIR/sbatch_train_amd.sh" +} + +declare -A TIME_LIMIT=( [1]="02:00:00" [2]="02:30:00" [4]="03:00:00" [8]="04:00:00" ) +IFS=',' read -ra NODE_COUNTS <<< "$NODES_LIST" +echo "=== ORBIT-2 scaling sweep (${ORBIT2_SCALING_TAG}) ===" +echo " Data root : $ORBIT2_DATA_ROOT" +echo " Config : $ORBIT2_CONFIG_TEMPLATE" +echo " Output base : $ORBIT2_OUTPUT_DIR" +echo " Epochs: $ORBIT2_MAX_EPOCH Max batches: $ORBIT2_MAX_BATCHES" +echo "" + +JOB_IDS=() +for n in "${NODE_COUNTS[@]}"; do + export ORBIT2_OUTPUT_DIR="${AI4S_SHARED_DIR}/models/ORBIT-2/outputs/scaling-${ORBIT2_SCALING_TAG}/${n}node" + jid=$(submit "$n" "${TIME_LIMIT[$n]:-02:00:00}") + echo " ${n}-node -> job $jid" + JOB_IDS+=("$jid") +done + +echo "" +echo "Submitted: ${JOB_IDS[*]}" +echo "After all complete:" +echo " python3 collate_scaling_study.py --log-dir ${LOG_DIR} --jobs $(IFS=,; echo "${JOB_IDS[*]}") \\" +echo " --steady-epoch-start 2 --warmup-batches-per-epoch 1 --require-loss-sanity \\" +echo " -o ${ORBIT2_OUTPUT_DIR}/scaling_study" diff --git a/earth_science/models/ORBIT-2/examples/sbatch_train_amd.sh b/earth_science/models/ORBIT-2/examples/sbatch_train_amd.sh new file mode 100755 index 0000000..92d59c2 --- /dev/null +++ b/earth_science/models/ORBIT-2/examples/sbatch_train_amd.sh @@ -0,0 +1,250 @@ +#!/usr/bin/env bash +# ORBIT-2 multi-node training on AMD Instinct via SLURM (Apptainer + MPI). +# +# Wraps upstream intermediate_downscaling.py with Studio launchers and caps. +# Uses real 10.0_arcmin PRISM data in same-dir mode (see interm_8m_lux.yaml). +# +# Quick start (1 node, 8 GPUs): +# export AI4S_SHARED_DIR=/path/to/shared +# export ORBIT2_DATA_ROOT=$AI4S_SHARED_DIR/models/ORBIT-2/data/superres/prism/10.0_arcmin +# export ORBIT2_MAX_EPOCH=1 ORBIT2_MAX_BATCHES=5 +# sbatch earth_science/models/ORBIT-2/examples/sbatch_train_amd.sh +# +# Key environment variables: +# AI4S_SHARED_DIR Shared storage root (required) +# ORBIT2_DATA_ROOT Data root (PRISM 10.0_arcmin or era5/1.0_deg for sanity) +# ORBIT2_CONFIG_TEMPLATE YAML template basename (default: interm_8m_lux.yaml; +# use interm_8m_lux_era5.yaml for new ERA5 1.0_deg data) +# ORBIT2_ROOT ORBIT-2 clone (default: $AI4S_SHARED_DIR/models/ORBIT-2/code/ORBIT-2) +# ORBIT2_SIF Apptainer SIF path +# ORBIT2_OVERLAY Pre-built ext3 overlay (required — avoids ~15 min pip/job) +# ORBIT2_MAX_EPOCH Cap trainer.max_epochs (default: 3; must be >= 2 — +# upstream loop is while (epoch_start+1) < max_epochs) +# ORBIT2_MAX_BATCHES Cap batches per epoch (default: 20; 0 = unlimited) +# ORBIT2_BATCH_SIZE Per-rank batch size (default: 8) +# ORBIT2_OUTPUT_DIR Job output dir (default: .../outputs/train/) +# TORCH_NCCL_HIGH_PRIORITY / GPU_MAX_HW_QUEUES — RCCL tuning (default: 1 / 2) +# ORBIT2_SKIP_NODE_HEALTH_PROBE=1 — skip mount probe (not recommended) + +#SBATCH --job-name=orbit2-train +#SBATCH --partition=YOUR_PARTITION_HERE +#SBATCH --account=YOUR_ACCOUNT_HERE +#SBATCH --nodes=1 +#SBATCH --ntasks-per-node=8 +#SBATCH --gpus-per-node=8 +#SBATCH --cpus-per-task=7 +#SBATCH --time=02:00:00 +#SBATCH --output=orbit2-train-%j.out +#SBATCH --error=orbit2-train-%j.out + +set -euo pipefail + +if [[ -n "${SLURM_JOB_ID:-}" ]]; then + _ORIG_CMD=$(scontrol show job "$SLURM_JOB_ID" | sed -n 's/.*Command=\(\S\+\).*/\1/p') + SCRIPT_DIR=$(cd "$(dirname "$_ORIG_CMD")" && pwd) +else + SCRIPT_DIR=$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd) +fi + +ORBIT2_BASE="${AI4S_SHARED_DIR:?AI4S_SHARED_DIR must be set}/models/ORBIT-2" +ORBIT2_ROOT="${ORBIT2_ROOT:-${ORBIT2_BASE}/code/ORBIT-2}" +ORBIT2_SIF="${ORBIT2_SIF:-${AI4S_SHARED_DIR}/images/pytorch_rocm7.2.2_ubuntu24.04_py3.12_pytorch_release_2.10.0.sif}" +ORBIT2_OVERLAY="${ORBIT2_OVERLAY:-${ORBIT2_BASE}/overlays/orbit2-overlay.img}" +ORBIT2_DATA_ROOT="${ORBIT2_DATA_ROOT:-${ORBIT2_BASE}/data/superres/prism/10.0_arcmin}" +ORBIT2_MAX_EPOCH="${ORBIT2_MAX_EPOCH:-3}" +ORBIT2_MAX_BATCHES="${ORBIT2_MAX_BATCHES:-20}" +ORBIT2_BATCH_SIZE="${ORBIT2_BATCH_SIZE:-8}" +ORBIT2_OUTPUT_DIR="${ORBIT2_OUTPUT_DIR:-${ORBIT2_BASE}/outputs/train/${SLURM_JOB_ID:-$$}}" + +TORCH_NCCL_HIGH_PRIORITY="${TORCH_NCCL_HIGH_PRIORITY:-1}" +GPU_MAX_HW_QUEUES="${GPU_MAX_HW_QUEUES:-2}" + +NODES="${SLURM_JOB_NUM_NODES:-1}" +GPUS_PER_NODE=8 +TOTAL_RANKS=$((NODES * GPUS_PER_NODE)) + +for var in ORBIT2_SIF ORBIT2_OVERLAY ORBIT2_DATA_ROOT; do + if [[ ! -e "${!var}" ]]; then + echo "ERROR: $var not found: ${!var}" >&2 + exit 2 + fi +done + +if [[ ! -f "${ORBIT2_ROOT}/examples/intermediate_downscaling.py" ]]; then + echo "ERROR: ORBIT-2 clone missing intermediate_downscaling.py: ${ORBIT2_ROOT}" >&2 + exit 2 +fi + +mkdir -p "$ORBIT2_OUTPUT_DIR" + +ORBIT2_CONFIG_TEMPLATE="${ORBIT2_CONFIG_TEMPLATE:-interm_8m_lux.yaml}" +CONFIG_TEMPLATE="${SCRIPT_DIR}/${ORBIT2_CONFIG_TEMPLATE}" +if [[ ! -f "$CONFIG_TEMPLATE" ]]; then + echo "ERROR: ORBIT2_CONFIG_TEMPLATE not found: $CONFIG_TEMPLATE" >&2 + exit 2 +fi + +JOB_CONFIG="${ORBIT2_OUTPUT_DIR}/${ORBIT2_CONFIG_TEMPLATE%.yaml}_${SLURM_JOB_ID:-$$}.yaml" +python3 "$SCRIPT_DIR/render_orbit2_config.py" \ + --template "$CONFIG_TEMPLATE" \ + --nodes "$NODES" \ + --gpus-per-node "$GPUS_PER_NODE" \ + --data-root "$ORBIT2_DATA_ROOT" \ + --max-epochs "$ORBIT2_MAX_EPOCH" \ + --batch-size "$ORBIT2_BATCH_SIZE" \ + -o "$JOB_CONFIG" + +echo "=== ORBIT-2 Training (Apptainer + MPI) ===" +echo " Nodes : $NODES" +echo " GPUs/node : $GPUS_PER_NODE" +echo " Total ranks : $TOTAL_RANKS" +echo " SIF : $ORBIT2_SIF" +echo " Overlay : $ORBIT2_OVERLAY" +echo " Data root : $ORBIT2_DATA_ROOT" +echo " Max epochs : $ORBIT2_MAX_EPOCH" +echo " Max batches : $ORBIT2_MAX_BATCHES" +echo " Batch size : $ORBIT2_BATCH_SIZE" +echo " RCCL priority: TORCH_NCCL_HIGH_PRIORITY=$TORCH_NCCL_HIGH_PRIORITY" +echo " HW queues : GPU_MAX_HW_QUEUES=$GPU_MAX_HW_QUEUES" +echo " Output dir : $ORBIT2_OUTPUT_DIR" +echo " Config : $JOB_CONFIG" +echo " Node(s) : ${SLURM_NODELIST:-$(hostname)}" +echo " Date : $(date)" +echo "" + +# --------------------------------------------------------------------------- +# Per-node mount-health probe (exit 42 → retry with --exclude=) +# --------------------------------------------------------------------------- +ORBIT2_SKIP_NODE_HEALTH_PROBE="${ORBIT2_SKIP_NODE_HEALTH_PROBE:-0}" +if [[ "$ORBIT2_SKIP_NODE_HEALTH_PROBE" != "1" ]]; then + echo "--- Per-node mount-health probe ---" + _PROBE_OUT="${ORBIT2_OUTPUT_DIR}/node_health_probe.txt" + srun --no-kill --kill-on-bad-exit=0 -N "$SLURM_JOB_NUM_NODES" --ntasks-per-node=1 \ + --cpus-per-task=1 --gres=none --output="${_PROBE_OUT}" --error="${_PROBE_OUT}" \ + bash -c ' + H=$(hostname) + HM=$(test -d "/home/$USER" 2>/dev/null && echo OK || echo FAIL) + SH=$(test -d "'"$AI4S_SHARED_DIR"'" 2>/dev/null && echo OK || echo FAIL) + SIF_CHECK=$(test -f "'"$ORBIT2_SIF"'" 2>/dev/null && echo OK || echo FAIL) + DATA_CHECK=$(test -d "'"$ORBIT2_DATA_ROOT"'" 2>/dev/null && echo OK || echo FAIL) + echo "NODE_HEALTH $H home=$HM shared=$SH sif=$SIF_CHECK data=$DATA_CHECK" + ' 2>&1 || true + if [[ -s "$_PROBE_OUT" ]]; then + grep "^NODE_HEALTH" "$_PROBE_OUT" | sed 's/^/ /' + _BAD_NODES=$(grep "^NODE_HEALTH" "$_PROBE_OUT" | awk '/home=FAIL|shared=FAIL|sif=FAIL|data=FAIL/ {print $2}' | sort -u | tr '\n' ',' | sed 's/,$//') + else + _BAD_NODES="" + fi + if [[ -n "$_BAD_NODES" ]]; then + echo "FATAL: NODE_HEALTH_PROBE detected broken mounts on: $_BAD_NODES" >&2 + echo " Retry with: sbatch --exclude=$_BAD_NODES ..." >&2 + exit 42 + fi + echo " All $SLURM_JOB_NUM_NODES nodes healthy." +fi + +MASTER_ADDR=$(scontrol show hostnames "$SLURM_JOB_NODELIST" | head -1) +MASTER_PORT="${ORBIT2_MASTER_PORT:-29500}" + +RANK_SCRIPT="${ORBIT2_OUTPUT_DIR}/orbit2_rank_${SLURM_JOB_ID:-$$}.sh" +cat > "$RANK_SCRIPT" << RANKEOF +#!/usr/bin/env bash +set -euo pipefail +source /opt/venv/bin/activate +echo "[rank \$SLURM_PROCID / \$SLURM_NTASKS] GPU \$SLURM_LOCALID on \$(hostname)" +export PYTHONPATH="/opt/orbit2-pkgs:/orbit2/src:/orbit2:\${PYTHONPATH:-}" +export OMP_NUM_THREADS="\${SLURM_CPUS_PER_TASK:-7}" +export MIOPEN_DISABLE_CACHE=1 +export MIOPEN_USER_DB_PATH="\${TMPDIR:-/tmp}/orbit2-miopen-\${SLURM_JOB_ID:-\$\$}-\${SLURM_PROCID:-0}" +mkdir -p "\$MIOPEN_USER_DB_PATH" +export PYTHONNOUSERSITE=1 +export HSA_NO_SCRATCH_RECLAIM=1 +export ORBIT_USE_DDSTORE=0 +export TORCH_NCCL_HIGH_PRIORITY="${TORCH_NCCL_HIGH_PRIORITY}" +export GPU_MAX_HW_QUEUES="${GPU_MAX_HW_QUEUES}" +export ORBIT2_ROOT="/orbit2" +export ORBIT2_MAX_BATCHES="${ORBIT2_MAX_BATCHES}" +export ORBIT2_DATA_TYPE="${ORBIT2_DATA_TYPE:-float32}" +export ORBIT2_FUSED_ATTN="${ORBIT2_FUSED_ATTN:-DEFAULT}" +export ORBIT2_RANK_PRE_TRAIN_HOOK="\${ORBIT2_RANK_PRE_TRAIN_HOOK:-}" +cd /orbit2/examples +exec python3 /examples/run_orbit2_train.py /config/config.yaml +RANKEOF +chmod +x "$RANK_SCRIPT" + +CONFIG_BIND="$(realpath "$JOB_CONFIG"):/config/config.yaml" + +# Multi-node RCCL/MPI (from .cluster-config.yaml) +RCCL_MULTINODE_ENVS=() +MPI_MULTINODE_ENVS=() +if [[ "$NODES" -gt 1 ]]; then + _CLUSTER_CFG="${SCRIPT_DIR}/../../../../.cluster-config.yaml" + if [[ -f "$_CLUSTER_CFG" ]]; then + _yaml_get() { grep "^ $1:" "$_CLUSTER_CFG" 2>/dev/null | sed 's/.*: *"\?\([^"]*\)"\?.*/\1/' | grep -v '^$'; } + : "${NCCL_IB_HCA:=$(_yaml_get ib_hca)}" + : "${NCCL_SOCKET_IFNAME:=$(_yaml_get mgmt_iface)}" + : "${RCCL_ANP_PLUGIN:=$(_yaml_get rccl_anp_plugin)}" + : "${LIBIONIC_PATH:=$(_yaml_get libionic_path)}" + fi + : "${NCCL_IB_HCA:?Set NCCL_IB_HCA or configure network.ib_hca in .cluster-config.yaml}" + : "${NCCL_SOCKET_IFNAME:?Set NCCL_SOCKET_IFNAME or configure network.mgmt_iface}" + : "${RCCL_ANP_PLUGIN:?Set RCCL_ANP_PLUGIN or configure network.rccl_anp_plugin}" + : "${LIBIONIC_PATH:?Set LIBIONIC_PATH or configure network.libionic_path}" + + MPI_MULTINODE_ENVS=( + --env OMPI_MCA_pml=ob1 + --env OMPI_MCA_btl=tcp,self + --env OMPI_MCA_btl_tcp_if_include="$NCCL_SOCKET_IFNAME" + --env MPI4PY_RC_THREADS=false + ) + RCCL_MULTINODE_ENVS=( + --bind "${RCCL_ANP_PLUGIN}:${RCCL_ANP_PLUGIN}:ro" + --bind "${LIBIONIC_PATH}:${LIBIONIC_PATH}:ro" + --env NCCL_NET_PLUGIN="$RCCL_ANP_PLUGIN" + --env NCCL_IB_HCA="$NCCL_IB_HCA" + --env NCCL_IB_GID_INDEX=1 + --env NCCL_GDR_FLUSH_DISABLE=1 + --env RCCL_GDR_FLUSH_GPU_MEM_NO_RELAXED_ORDERING=0 + --env NCCL_GDRCOPY_ENABLE=0 + --env NCCL_IB_QPS_PER_CONNECTION=1 + --env HSA_NO_SCRATCH_RECLAIM=1 + --env NCCL_IB_TC=96 + --env NCCL_IB_FIFO_TC=192 + --env NCCL_IGNORE_CPU_AFFINITY=1 + --env NCCL_PXN_DISABLE=0 + --env NET_OPTIONAL_RECV_COMPLETION=1 + --env NCCL_IB_USE_INLINE=1 + --env NCCL_SOCKET_IFNAME="$NCCL_SOCKET_IFNAME" + --env RCCL_LL128_FORCE_ENABLE=1 + --env NCCL_IB_PCI_RELAXED_ORDERING=1 + --env NCCL_DMABUF_ENABLE=1 + --env NCCL_DEBUG="${NCCL_DEBUG:-WARN}" + ) + echo " Multi-node RCCL enabled (NCCL_IB_HCA=$NCCL_IB_HCA)" +fi + +echo "--- Launching training: $TOTAL_RANKS ranks across $NODES nodes ---" +srun --mpi=pmix apptainer exec \ + --rocm \ + --overlay "${ORBIT2_OVERLAY}:ro" \ + --bind "/opt/ompi:/opt/ompi:ro" \ + --bind "$ORBIT2_ROOT":/orbit2 \ + --bind "$SCRIPT_DIR":/examples \ + --bind "$(dirname "$RANK_SCRIPT"):$(dirname "$RANK_SCRIPT")" \ + --bind "$ORBIT2_DATA_ROOT":"$ORBIT2_DATA_ROOT":ro \ + --bind "$ORBIT2_OUTPUT_DIR":"$ORBIT2_OUTPUT_DIR" \ + --bind "$CONFIG_BIND" \ + --env HOSTNAME="$MASTER_ADDR" \ + --env MASTER_PORT="$MASTER_PORT" \ + --env PMIX_MCA_gds=hash \ + --env PMIX_MCA_psec=native \ + --env PYTHONPATH="/opt/orbit2-pkgs:/orbit2/src:/orbit2" \ + --env LD_LIBRARY_PATH=/opt/venv/lib/python3.12/site-packages/torch/lib \ + "${MPI_MULTINODE_ENVS[@]}" \ + "${RCCL_MULTINODE_ENVS[@]}" \ + "$ORBIT2_SIF" \ + bash "$RANK_SCRIPT" + +echo "" +echo "=== ORBIT-2 training complete ===" +echo " Output: $ORBIT2_OUTPUT_DIR" diff --git a/earth_science/models/ORBIT-2/examples/sbatch_train_perf_amd.sh b/earth_science/models/ORBIT-2/examples/sbatch_train_perf_amd.sh new file mode 100755 index 0000000..61c7c84 --- /dev/null +++ b/earth_science/models/ORBIT-2/examples/sbatch_train_perf_amd.sh @@ -0,0 +1,250 @@ +#!/usr/bin/env bash +# ORBIT-2 2-node training with PyTorch profiling + Omnistat user-mode telemetry. +# +# Perf-analysis variant of sbatch_train_amd.sh. See recipes/perf-analysis/. +# +# Quick start: +# export AI4S_SHARED_DIR=/path/to/shared +# export OMNIHUB_TOOLS_DIR=/path/to/omnihub/tools +# export ORBIT2_DATA_ROOT=$AI4S_SHARED_DIR/models/ORBIT-2/data/superres/prism/10.0_arcmin +# sbatch earth_science/models/ORBIT-2/examples/sbatch_train_perf_amd.sh + +#SBATCH --job-name=orbit2-perf +#SBATCH --partition=YOUR_PARTITION_HERE +#SBATCH --account=YOUR_ACCOUNT_HERE +#SBATCH --nodes=2 +#SBATCH --ntasks-per-node=8 +#SBATCH --gpus-per-node=8 +#SBATCH --cpus-per-task=7 +#SBATCH --time=02:30:00 +#SBATCH --output=orbit2-train-%j.out +#SBATCH --error=orbit2-train-%j.out + +set -euo pipefail + +if [[ -n "${SLURM_JOB_ID:-}" ]]; then + _ORIG_CMD=$(scontrol show job "$SLURM_JOB_ID" | sed -n 's/.*Command=\(\S\+\).*/\1/p') + SCRIPT_DIR=$(cd "$(dirname "$_ORIG_CMD")" && pwd) +else + SCRIPT_DIR=$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd) +fi + +ORBIT2_BASE="${AI4S_SHARED_DIR:?AI4S_SHARED_DIR must be set}/models/ORBIT-2" +ORBIT2_ROOT="${ORBIT2_ROOT:-${ORBIT2_BASE}/code/ORBIT-2}" +ORBIT2_SIF="${ORBIT2_SIF:-${AI4S_SHARED_DIR}/images/pytorch_rocm7.2.2_ubuntu24.04_py3.12_pytorch_release_2.10.0.sif}" +ORBIT2_OVERLAY="${ORBIT2_OVERLAY:-${ORBIT2_BASE}/overlays/orbit2-overlay.img}" +ORBIT2_DATA_ROOT="${ORBIT2_DATA_ROOT:-${ORBIT2_BASE}/data/superres/prism/10.0_arcmin}" +ORBIT2_MAX_EPOCH="${ORBIT2_MAX_EPOCH:-3}" +ORBIT2_MAX_BATCHES="${ORBIT2_MAX_BATCHES:-20}" +ORBIT2_BATCH_SIZE="${ORBIT2_BATCH_SIZE:-4}" +ORBIT2_OUTPUT_DIR="${ORBIT2_OUTPUT_DIR:-${ORBIT2_BASE}/perf-runs/${SLURM_JOB_ID:-$$}}" +PROFILE_TARGET_EPOCH="${PROFILE_TARGET_EPOCH:-0}" + +: "${OMNIHUB_TOOLS_DIR:?OMNIHUB_TOOLS_DIR must be set}" +OMNISTAT_VENV="${OMNISTAT_VENV:-${OMNIHUB_TOOLS_DIR}/omnihub-inspect}" +OMNISTAT_TEMPLATE="${OMNISTAT_TEMPLATE:-${SCRIPT_DIR}/../recipes/perf-analysis/omnistat.config.template}" +OMNISTAT_USERMODE_INTERVAL="${OMNISTAT_USERMODE_INTERVAL:-1}" + +TORCH_NCCL_HIGH_PRIORITY="${TORCH_NCCL_HIGH_PRIORITY:-1}" +GPU_MAX_HW_QUEUES="${GPU_MAX_HW_QUEUES:-2}" + +NODES="${SLURM_JOB_NUM_NODES:-2}" +GPUS_PER_NODE=8 +TOTAL_RANKS=$((NODES * GPUS_PER_NODE)) + +for var in ORBIT2_SIF ORBIT2_OVERLAY ORBIT2_DATA_ROOT OMNISTAT_TEMPLATE; do + if [[ ! -e "${!var}" ]]; then + echo "ERROR: $var not found: ${!var}" >&2 + exit 2 + fi +done +if [[ ! -x "${OMNISTAT_VENV}/bin/omnistat-usermode" ]]; then + echo "ERROR: omnistat-usermode not found at ${OMNISTAT_VENV}/bin/" >&2 + exit 2 +fi + +mkdir -p "$ORBIT2_OUTPUT_DIR" +OMNISTAT_CONFIG="${ORBIT2_OUTPUT_DIR}/omnistat.config" +sed -e "s|@JOB_DIR@|${ORBIT2_OUTPUT_DIR}|g" \ + -e "s|@OMNIHUB_TOOLS_DIR@|${OMNIHUB_TOOLS_DIR}|g" \ + "$OMNISTAT_TEMPLATE" > "$OMNISTAT_CONFIG" + +JOB_CONFIG="${ORBIT2_OUTPUT_DIR}/interm_8m_lux_${SLURM_JOB_ID:-$$}.yaml" +python3 "$SCRIPT_DIR/render_orbit2_config.py" \ + --nodes "$NODES" \ + --gpus-per-node "$GPUS_PER_NODE" \ + --data-root "$ORBIT2_DATA_ROOT" \ + --max-epochs "$ORBIT2_MAX_EPOCH" \ + --batch-size "$ORBIT2_BATCH_SIZE" \ + -o "$JOB_CONFIG" + +echo "=== ORBIT-2 Perf Analysis Run ===" +echo " Nodes : $NODES" +echo " Total ranks : $TOTAL_RANKS" +echo " Output dir : $ORBIT2_OUTPUT_DIR" +echo " Profile epoch: $PROFILE_TARGET_EPOCH" +echo " Omnistat cfg : $OMNISTAT_CONFIG" +echo " Node(s) : ${SLURM_NODELIST:-$(hostname)}" +echo "" + +# Mount-health probe (same as sbatch_train_amd.sh) +ORBIT2_SKIP_NODE_HEALTH_PROBE="${ORBIT2_SKIP_NODE_HEALTH_PROBE:-0}" +if [[ "$ORBIT2_SKIP_NODE_HEALTH_PROBE" != "1" ]]; then + _PROBE_OUT="${ORBIT2_OUTPUT_DIR}/node_health_probe.txt" + srun --no-kill --kill-on-bad-exit=0 -N "$SLURM_JOB_NUM_NODES" --ntasks-per-node=1 \ + --cpus-per-task=1 --gres=none --output="${_PROBE_OUT}" --error="${_PROBE_OUT}" \ + bash -c ' + H=$(hostname) + HM=$(test -d "/home/$USER" 2>/dev/null && echo OK || echo FAIL) + SH=$(test -d "'"$AI4S_SHARED_DIR"'" 2>/dev/null && echo OK || echo FAIL) + SIF_CHECK=$(test -f "'"$ORBIT2_SIF"'" 2>/dev/null && echo OK || echo FAIL) + echo "NODE_HEALTH $H home=$HM shared=$SH sif=$SIF_CHECK" + ' 2>&1 || true + _BAD_NODES=$(grep "^NODE_HEALTH" "$_PROBE_OUT" 2>/dev/null | awk '/FAIL/ {print $2}' | sort -u | tr '\n' ',' | sed 's/,$//' || true) + if [[ -n "$_BAD_NODES" ]]; then + echo "FATAL: NODE_HEALTH_PROBE failed on: $_BAD_NODES" >&2 + exit 42 + fi +fi + +export PATH="${OMNISTAT_VENV}/bin:${PATH}" +echo "--- Starting Omnistat user-mode ---" +"${OMNISTAT_VENV}/bin/omnistat-usermode" --configfile "$OMNISTAT_CONFIG" --start --interval "$OMNISTAT_USERMODE_INTERVAL" \ + 2>&1 | tee "${ORBIT2_OUTPUT_DIR}/omnistat_start.log" || true + +cleanup_omnistat() { + "${OMNISTAT_VENV}/bin/omnistat-usermode" --configfile "$OMNISTAT_CONFIG" --stopexporters || true + "${OMNISTAT_VENV}/bin/omnistat-usermode" --configfile "$OMNISTAT_CONFIG" --stopserver || true +} +trap cleanup_omnistat EXIT + +MASTER_ADDR=$(scontrol show hostnames "$SLURM_JOB_NODELIST" | head -1) +MASTER_PORT="${ORBIT2_MASTER_PORT:-29500}" +PROFILER_HOOK="${SCRIPT_DIR}/orbit2_profiler_hook.py" +TRACE_DIR="${ORBIT2_OUTPUT_DIR}/traces" + +RANK_SCRIPT="${ORBIT2_OUTPUT_DIR}/orbit2_rank_${SLURM_JOB_ID:-$$}.sh" +cat > "$RANK_SCRIPT" << RANKEOF +#!/usr/bin/env bash +set -euo pipefail +source /opt/venv/bin/activate +export PYTHONPATH="/opt/orbit2-pkgs:/orbit2/src:/orbit2:\${PYTHONPATH:-}" +export OMP_NUM_THREADS="\${SLURM_CPUS_PER_TASK:-7}" +export MIOPEN_DISABLE_CACHE=1 +export MIOPEN_USER_DB_PATH="\${TMPDIR:-/tmp}/orbit2-miopen-\${SLURM_JOB_ID:-\$\$}-\${SLURM_PROCID:-0}" +mkdir -p "\$MIOPEN_USER_DB_PATH" +export PYTHONNOUSERSITE=1 +export HSA_NO_SCRATCH_RECLAIM=1 +export ORBIT_USE_DDSTORE=0 +export TORCH_NCCL_HIGH_PRIORITY="${TORCH_NCCL_HIGH_PRIORITY}" +export GPU_MAX_HW_QUEUES="${GPU_MAX_HW_QUEUES}" +export ORBIT2_ROOT="/orbit2" +export ORBIT2_MAX_BATCHES="${ORBIT2_MAX_BATCHES}" +export ORBIT2_DATA_TYPE="${ORBIT2_DATA_TYPE:-float32}" +export ORBIT2_FUSED_ATTN="${ORBIT2_FUSED_ATTN:-DEFAULT}" +export ORBIT2_OUTPUT_DIR="${ORBIT2_OUTPUT_DIR}" +export ORBIT2_PROFILE_DIR="${TRACE_DIR}" +export PROFILE_TARGET_EPOCH="${PROFILE_TARGET_EPOCH}" +export PROFILE_RANK0_ONLY=1 +export ORBIT2_RANK_PRE_TRAIN_HOOK="${PROFILER_HOOK}" +cd /orbit2/examples +exec python3 /examples/run_orbit2_train.py /config/config.yaml +RANKEOF +chmod +x "$RANK_SCRIPT" + +CONFIG_BIND="$(realpath "$JOB_CONFIG"):/config/config.yaml" + +RCCL_MULTINODE_ENVS=() +MPI_MULTINODE_ENVS=() +if [[ "$NODES" -gt 1 ]]; then + _CLUSTER_CFG="${SCRIPT_DIR}/../../../../.cluster-config.yaml" + if [[ -f "$_CLUSTER_CFG" ]]; then + _yaml_get() { grep "^ $1:" "$_CLUSTER_CFG" 2>/dev/null | sed 's/.*: *"\?\([^"]*\)"\?.*/\1/' | grep -v '^$'; } + : "${NCCL_IB_HCA:=$(_yaml_get ib_hca)}" + : "${NCCL_SOCKET_IFNAME:=$(_yaml_get mgmt_iface)}" + : "${RCCL_ANP_PLUGIN:=$(_yaml_get rccl_anp_plugin)}" + : "${LIBIONIC_PATH:=$(_yaml_get libionic_path)}" + fi + : "${NCCL_IB_HCA:?NCCL_IB_HCA required for multi-node}" + : "${NCCL_SOCKET_IFNAME:?NCCL_SOCKET_IFNAME required}" + : "${RCCL_ANP_PLUGIN:?RCCL_ANP_PLUGIN required}" + : "${LIBIONIC_PATH:?LIBIONIC_PATH required}" + MPI_MULTINODE_ENVS=( + --env OMPI_MCA_pml=ob1 --env OMPI_MCA_btl=tcp,self + --env OMPI_MCA_btl_tcp_if_include="$NCCL_SOCKET_IFNAME" + --env MPI4PY_RC_THREADS=false + ) + RCCL_MULTINODE_ENVS=( + --bind "${RCCL_ANP_PLUGIN}:${RCCL_ANP_PLUGIN}:ro" + --bind "${LIBIONIC_PATH}:${LIBIONIC_PATH}:ro" + --env NCCL_NET_PLUGIN="$RCCL_ANP_PLUGIN" + --env NCCL_IB_HCA="$NCCL_IB_HCA" + --env NCCL_IB_GID_INDEX=1 + --env NCCL_GDR_FLUSH_DISABLE=1 + --env RCCL_GDR_FLUSH_GPU_MEM_NO_RELAXED_ORDERING=0 + --env NCCL_GDRCOPY_ENABLE=0 + --env NCCL_IB_QPS_PER_CONNECTION=1 + --env HSA_NO_SCRATCH_RECLAIM=1 + --env NCCL_IB_TC=96 + --env NCCL_IB_FIFO_TC=192 + --env NCCL_IGNORE_CPU_AFFINITY=1 + --env NCCL_PXN_DISABLE=0 + --env NET_OPTIONAL_RECV_COMPLETION=1 + --env NCCL_IB_USE_INLINE=1 + --env NCCL_SOCKET_IFNAME="$NCCL_SOCKET_IFNAME" + --env RCCL_LL128_FORCE_ENABLE=1 + --env NCCL_IB_PCI_RELAXED_ORDERING=1 + --env NCCL_DMABUF_ENABLE=1 + --env NCCL_DEBUG="${NCCL_DEBUG:-WARN}" + ) + echo " Multi-node RCCL enabled (NCCL_IB_HCA=$NCCL_IB_HCA)" +fi + +echo "--- Launching perf training: $TOTAL_RANKS ranks ---" +set +e +srun --mpi=pmix apptainer exec \ + --rocm --overlay "${ORBIT2_OVERLAY}:ro" \ + --bind "/opt/ompi:/opt/ompi:ro" \ + --bind "$ORBIT2_ROOT":/orbit2 \ + --bind "$SCRIPT_DIR":/examples \ + --bind "$(dirname "$RANK_SCRIPT"):$(dirname "$RANK_SCRIPT")" \ + --bind "$ORBIT2_DATA_ROOT":"$ORBIT2_DATA_ROOT":ro \ + --bind "$ORBIT2_OUTPUT_DIR":"$ORBIT2_OUTPUT_DIR" \ + --bind "$CONFIG_BIND" \ + --env HOSTNAME="$MASTER_ADDR" --env MASTER_PORT="$MASTER_PORT" \ + --env PMIX_MCA_gds=hash --env PMIX_MCA_psec=native \ + --env PYTHONPATH="/opt/orbit2-pkgs:/orbit2/src:/orbit2" \ + --env LD_LIBRARY_PATH=/opt/venv/lib/python3.12/site-packages/torch/lib \ + "${MPI_MULTINODE_ENVS[@]}" "${RCCL_MULTINODE_ENVS[@]}" \ + "$ORBIT2_SIF" bash "$RANK_SCRIPT" +TRAIN_RC=$? +set -e + +_SLURM_LOG="${ORBIT2_SLURM_LOG:-${SLURM_SUBMIT_DIR:-.}/orbit2-train-${SLURM_JOB_ID}.out}" +if [[ ! -f "$_SLURM_LOG" ]]; then + _SLURM_LOG="${AI4S_SHARED_DIR}/models/ORBIT-2/outputs/train/logs/orbit2-train-${SLURM_JOB_ID}.out" +fi +cp -f "$_SLURM_LOG" "${ORBIT2_OUTPUT_DIR}/orbit2-train-${SLURM_JOB_ID}.out" 2>/dev/null || true + +python3 - "$ORBIT2_OUTPUT_DIR" "${SLURM_JOB_ID}" "$TRAIN_RC" <<'PYEOF' +import json, sys +from pathlib import Path +job_dir, job_id, train_rc = Path(sys.argv[1]), sys.argv[2], int(sys.argv[3]) +manifest = { + "job_id": job_id, + "model": "ORBIT-2", + "workload": "intermediate_downscaling", + "job_dir": str(job_dir), + "slurm_log": str(job_dir / f"orbit2-train-{job_id}.out"), + "omnistat_config": str(job_dir / "omnistat.config"), + "omnistat_db": str(job_dir / "omnistat-db"), + "trace_dir": str(job_dir / "traces"), + "state": "complete" if train_rc == 0 else "failed", + "exit_code": train_rc, +} +(job_dir / "manifest.json").write_text(json.dumps(manifest, indent=2) + "\n") +print(f"Wrote {job_dir / 'manifest.json'}") +PYEOF + +echo "=== ORBIT-2 perf run complete (rc=$TRAIN_RC) ===" +exit $TRAIN_RC diff --git a/earth_science/models/ORBIT-2/model.yaml b/earth_science/models/ORBIT-2/model.yaml index f6479be..1c3bcf3 100644 --- a/earth_science/models/ORBIT-2/model.yaml +++ b/earth_science/models/ORBIT-2/model.yaml @@ -16,6 +16,16 @@ recipes: recipe_path: recipes/inference/ script: examples/run_visualize.py slurm: examples/sbatch_infer_amd.sh + - task: perf-analysis + description: 2-node training perf analysis with Omnistat + PyTorch profiler + recipe_path: recipes/perf-analysis/ + script: examples/run_orbit2_train.py + slurm: examples/sbatch_train_perf_amd.sh + - task: train + description: Multi-node intermediate downscaling on staged NPZ data (PRISM same-dir or ERA5 sanity templates) + recipe_path: recipes/train/ + script: examples/run_orbit2_train.py + slurm: examples/sbatch_train_amd.sh env_vars: ORBIT2_ROOT: diff --git a/earth_science/models/ORBIT-2/recipes/perf-analysis/HANDOFF.md b/earth_science/models/ORBIT-2/recipes/perf-analysis/HANDOFF.md new file mode 100644 index 0000000..bfe0bcc --- /dev/null +++ b/earth_science/models/ORBIT-2/recipes/perf-analysis/HANDOFF.md @@ -0,0 +1,54 @@ +# ORBIT-2 perf-analysis HANDOFF + +_Last updated: 2026-06._ + +## What works + +| Item | Status | Notes | +|------|--------|-------| +| 1-node × 8-GPU training on 10.0_arcmin same-dir | **Green** | Representative steady batch time ~0.5–2 s after warmup (workload- and I/O-dependent) | +| `sbatch_train_amd.sh` + `run_orbit2_train.py` | **Green** | gptl4py stub, FusedAttn fallback when `xformers.ops` unavailable, batch cap | +| Node health probe (exit 42) | **Green** | Catches broken home/shared/SIF/data mounts before training | +| `superres_mag: 1` same-dir config | **Required** | Without it, 4× superres vs same-resolution targets causes tensor shape mismatch | + +## Landmines + +1. **`max_epochs` must be ≥ 2** — upstream loop is `while (epoch_start + 1) < max_epochs`; `max_epochs=1` runs zero epochs. +2. **`xformers.ops` missing in overlay** — `run_orbit2_train.py` can fall back to `FusedAttn.DEFAULT` (PyTorch SDPA). Rebuild overlay per `build_overlay_amd.sh` if you need the CK attention path. +3. **`bfloat16` + CK attn** needs working `xformers.ops`; use `ORBIT2_DATA_TYPE=float32` until the overlay provides `ops`. +4. **Same-dir data** — both `low_res_dir` and `high_res_dir` point at the same grid; set `superres_mag: 1` in `interm_8m_lux.yaml` / `interm_8m_lux_era5.yaml`. +5. **Batch cap** — must patch `intermediate_downscaling` before `main()` (not via `runpy.run_path` reload). + +## FOM contract (steady-state + loss sanity) + +**Primary FOM:** `steady_batch_time_s` — mean per-batch wall time from **epoch ≥ 2**, excluding **batch index < 1** within each epoch (skips compile/warmup batch 0). Epochs 0–1 are warmup. + +**Loss sanity:** epoch-completed losses must **strictly decrease** epoch-over-epoch (`loss_sanity_pass`). Same-dir data makes absolute loss non-physical; use as crash/hang detector only. + +**Scaling runs:** use `ORBIT2_MAX_EPOCH=6` (trains epochs 0–4) so the steady window spans epochs 2–4. Parse with: + +```bash +python3 parse_training_log.py --log ... --steady-epoch-start 2 --warmup-batches-per-epoch 1 +python3 collate_scaling_study.py --log-dir .../logs --jobs ... --steady-epoch-start 2 --require-loss-sanity +``` + +## Scaling / wall-clock lessons + +- Long runs on full 180×360 grids need **enough wall time** (multi-epoch × many batches can exceed 1 h on some sites). +- Point SLURM `--output` / `--error` at a **shared project filesystem** your compute nodes can always read during outages (not only the submission directory on a login node). +- **Multi-node hangs** are often **infrastructure** (drained nodes, stuck scheduler states, shared filesystem hiccups) — confirm node state with `sinfo` / site ops before assuming a code bug. +- For exclusive GPU jobs, **avoid stacking two jobs on the same node**. Pin nodes with `--nodelist=...` only when your site policy allows it and you have confirmed the nodes are healthy. + +## Next work (after scaling baseline) + +- Port `perf-optimizer-loop/` + `lever_catalog.yaml` from the HydraGNN reference recipe. +- Rebuild overlay; re-test `bfloat16` + CK attn when `xformers.ops` is available. +- Stage **2.5_arcmin** PRISM targets for scientifically meaningful loss (not same-dir timing). + +## Key paths (parameterize with `AI4S_SHARED_DIR`) + +- Data: `$AI4S_SHARED_DIR/models/ORBIT-2/data/superres/prism/10.0_arcmin` (PRISM same-dir) or `.../era5/1.0_deg` (ERA5 same-dir sanity template). +- Overlay: `$AI4S_SHARED_DIR/models/ORBIT-2/overlays/orbit2-overlay.img` +- Perf runs: `$AI4S_SHARED_DIR/models/ORBIT-2/perf-runs//` +- Training logs (recommended): `$AI4S_SHARED_DIR/models/ORBIT-2/outputs/train/logs/` +- Omnistat / TraceLens tools: set **`OMNIHUB_TOOLS_DIR`** to your site-local checkout (see root `.cluster-config.yaml` key `omnihub.tools_dir` if your site maintains one). diff --git a/earth_science/models/ORBIT-2/recipes/perf-analysis/README.md b/earth_science/models/ORBIT-2/recipes/perf-analysis/README.md new file mode 100644 index 0000000..5452d75 --- /dev/null +++ b/earth_science/models/ORBIT-2/recipes/perf-analysis/README.md @@ -0,0 +1,45 @@ +# ORBIT-2 Performance Analysis with AMD AI Agents + +Multi-subagent workflow for 2-node ORBIT-2 training (`intermediate_downscaling.py`) on AMD MI355X, with Omnistat telemetry and optional rank-0 PyTorch traces. + +> **Audience:** performance engineers. Output is a diagnosis of the run, not a scientific ORBIT-2 claim. + +> **Data mode:** 10.0_arcmin PRISM same-dir (`superres_mag: 1`) — timing/scaling only; loss is not meaningful without 2.5_arcmin targets. + +## Quick start + +```bash +export AI4S_SHARED_DIR=/path/to/shared +export OMNIHUB_TOOLS_DIR=/path/to/omnihub/tools +export ORBIT2_DATA_ROOT=$AI4S_SHARED_DIR/models/ORBIT-2/data/superres/prism/10.0_arcmin +sbatch earth_science/models/ORBIT-2/examples/sbatch_train_perf_amd.sh +``` + +After the job completes, drive analyst + verifier + synthesizer subagents per [`.cursor/skills/ai4science-perf-analysis/SKILL.md`](../../../../.cursor/skills/ai4science-perf-analysis/SKILL.md). + +## Artifacts + +``` +$AI4S_SHARED_DIR/models/ORBIT-2/perf-runs// +├── manifest.json +├── omnistat-db/ +├── omnistat.config +├── traces/orbit2-epoch*-rank0.pt.trace.json +├── orbit2-train-.out +├── foms.json # from run_fom_extractor.py +├── tracelens/ # analyst outputs +└── combined_report.md # synthesizer output +``` + +## Primary FOM + +**`batch_time_s`** — mean steady-state batch wall time from rank-0 log lines `Batch N: X seconds` (see `examples/parse_training_log.py`). + +## Agent prompts + +- [agents/launcher.md](agents/launcher.md) +- [agents/tracelens_analyst.md](agents/tracelens_analyst.md) +- [agents/omnistat_analyst.md](agents/omnistat_analyst.md) +- [agents/synthesizer.md](agents/synthesizer.md) + +See [HANDOFF.md](HANDOFF.md) for validated landmines and current cluster state. diff --git a/earth_science/models/ORBIT-2/recipes/perf-analysis/agents/launcher.md b/earth_science/models/ORBIT-2/recipes/perf-analysis/agents/launcher.md new file mode 100644 index 0000000..fa423cd --- /dev/null +++ b/earth_science/models/ORBIT-2/recipes/perf-analysis/agents/launcher.md @@ -0,0 +1,215 @@ +# launcher subagent + +Submits a 2-node ORBIT-2 training (AMD Instinct MI355X) with PyTorch profiling and Omnistat user-mode telemetry, waits for completion, and writes `manifest.json` for downstream subagents. + +## Inputs + +- The ai4science-studio repo checkout (`$REPO_ROOT`). +- Cluster config at `.cluster-config.yaml` (partition/account, shared dirs). +- Optional env-var overrides: `ORBIT2_BATCH_SIZE`, `HYDRAGNN_MAX_NUM_BATCH`, `ORBIT2_NUM_EPOCH`, `ORBIT2_PRECISION`, `PROFILE_TARGET_EPOCH`. + +## Outputs + +- `${OMNIHUB_TOOLS_DIR}/omnihub-inspect/` — Python venv with omnistat (PR #271 + main merged) and TraceLens. +- `${OMNIHUB_TOOLS_DIR}/victoriametrics/victoria-metrics-prod` — VictoriaMetrics binary. +- `$AI4S_SHARED_DIR/models/ORBIT-2/perf-runs/omnistat.config` — Omnistat user-mode config. +- `$AI4S_SHARED_DIR/models/ORBIT-2/perf-runs//manifest.json` — manifest schema below. +- `$AI4S_SHARED_DIR/models/ORBIT-2/perf-runs//orbit2-train-.out` — symlinked from output dir. +- `$AI4S_SHARED_DIR/models/ORBIT-2/perf-runs//logs/` — symlinked rank-0 trace. +- `$AI4S_SHARED_DIR/models/ORBIT-2/perf-runs//omnistat-db/` — VictoriaMetrics datadir. + +## Manifest schema + +```json +{ + "jobid": "", + "submitted_at": "ISO 8601 UTC", + "completed_at": "ISO 8601 UTC", + "exit_state": "COMPLETED|FAILED|TIMEOUT|...", + "nodes": 2, + "gpus_per_node": 8, + "ranks": 16, + "nodelist": ",", + "partition": "", + "account": "", + "runtime_seconds": , + "config_used": "", + "profile_target_epoch": , + "trace_paths": [""], + "omnistat_db_path": "

", + "training_log_path": "", + "perf_run_dir": "", + "hg_precision": "fp64", + "hg_batch_size": 200, + "hg_num_epoch": 2, + "orbit2_max_num_batch": 30, + "tools_versions": { + "omnistat_commit": "", + "tracelens_commit": "", + "victoriametrics_version": "" + } +} +``` + +## Steps + +### 1. Lazy install tools (idempotent) + +```bash +set -euo pipefail +# Shared tool location — read from .cluster-config.yaml omnihub.tools_dir. +TOOLS="${OMNIHUB_TOOLS_DIR:?set OMNIHUB_TOOLS_DIR to a persistent checkout (e.g. under \$AI4S_SHARED_DIR/tools)}" +mkdir -p "$TOOLS" + +# 1a. Omnistat: jorda/skills branch with origin/main merged in +SRC=$TOOLS/omnistat-src +if [[ ! -d $SRC/.git ]]; then + git clone https://github.com/ROCm/omnistat.git "$SRC" +fi +cd "$SRC" +git fetch --all --prune +git checkout jorda/skills 2>/dev/null || git checkout -b jorda/skills origin/jorda/skills +git reset --hard origin/jorda/skills +if ! git merge --no-edit origin/main; then + echo "ERROR: merging origin/main into jorda/skills produced conflicts; resolve manually" >&2 + exit 2 +fi +OMNISTAT_COMMIT=$(git rev-parse HEAD) + +# 1b. Venv (py3.12 to match cluster python) +VENV=${OMNIHUB_TOOLS_DIR}/omnihub-inspect +if [[ ! -d $VENV ]]; then + python3 -m venv "$VENV" +fi +"$VENV/bin/pip" install --quiet --upgrade pip +"$VENV/bin/pip" install --quiet -e "$SRC[query]" + +# 1c. TraceLens (same venv) +"$VENV/bin/pip" install --quiet "git+https://github.com/AMD-AGI/TraceLens.git" +TRACELENS_COMMIT=$("$VENV/bin/python" -c "import TraceLens, importlib.metadata; print(importlib.metadata.version('TraceLens'))") + +# 1d. VictoriaMetrics binary +VM_DIR=$TOOLS/victoriametrics +if [[ ! -x $VM_DIR/victoria-metrics-prod ]]; then + mkdir -p "$VM_DIR" + cd "$VM_DIR" + VM_VERSION=$(curl -fsSL https://api.github.com/repos/VictoriaMetrics/VictoriaMetrics/releases/latest | grep '"tag_name":' | sed -E 's/.*"([^"]+)".*/\1/') + wget -q "https://github.com/VictoriaMetrics/VictoriaMetrics/releases/download/${VM_VERSION}/victoria-metrics-linux-amd64-${VM_VERSION}.tar.gz" + tar xzf "victoria-metrics-linux-amd64-${VM_VERSION}.tar.gz" + echo "$VM_VERSION" > VERSION +fi +VM_VERSION=$(cat $VM_DIR/VERSION) +``` + +### 1e. (Optional) Build the omnistat kernel-trace tool library + +Only required if the run will set `OMNISTAT_KERNEL_TRACE=1` to enable per-kernel dispatch tracing. Idempotent — skip if `build-trace/libomnistat_trace.so` already exists. + +```bash +TRACE_LIB=$TOOLS/omnistat-src/build-trace/libomnistat_trace.so +if [[ "${OMNISTAT_KERNEL_TRACE:-0}" == "1" && ! -f "$TRACE_LIB" ]]; then + # Must build inside the same SIF the workload uses — login nodes lack + # apptainer and ROCm headers. C++20 + libcurl + rocprofiler-sdk needed. + salloc -p -A -N 1 --time=00:15:00 --gpus-per-node=1 \ + apptainer exec --rocm \ + "${AI4S_SHARED_DIR}/images/pytorch_rocm7.2.2_ubuntu24.04_py3.12_pytorch_release_2.10.0.sif" \ + bash -c " + set -e + cd $TOOLS/omnistat-src + cmake -S rocprofiler-sdk/ -B build-trace/ -DBUILD_KERNEL_TRACE_LIB=ON + cmake --build build-trace/ -j 8 + " + test -f "$TRACE_LIB" || { echo "ERROR: kernel-trace build did not produce $TRACE_LIB" >&2; exit 2; } +fi +``` + +The sbatch wrapper expects the `.so` at `${OMNIHUB_TOOLS_DIR}/omnistat-src/build-trace/libomnistat_trace.so` by default; override with `OMNISTAT_TRACE_LIB=/path/to/lib.so`. See `ai4science-studio` SKILL §12 for when to enable kernel tracing vs the default device-counting collector. + +### 2. Author the Omnistat user-mode config (once, in repo `perf-runs/`) + +If `$AI4S_SHARED_DIR/models/ORBIT-2/perf-runs/omnistat.config` doesn't exist, write it from the template at `$REPO_ROOT/material_science/models/ORBIT-2/recipes/perf-analysis/omnistat.config.template`. The template uses `%(SLURM_JOB_ID)s` placeholders that omnistat-usermode resolves at runtime. + +### 3. Generate the per-job profile config + +Read the upstream `interm_8m_lux.json` from `$AI4S_SHARED_DIR/models/ORBIT-2/code/ORBIT-2/examples/multidataset_hpo_sc26/interm_8m_lux.json` and inject the `Profile` block **inside `NeuralNetwork`** (not at the top level — `train_validate_test()` is invoked with `config["NeuralNetwork"]` as `config`, so its Profiler reads `config["Profile"]` from that scope): + +```python +cfg.setdefault("NeuralNetwork", {})["Profile"] = { + "enable": 1, + "target_epoch": int(os.environ.get("PROFILE_TARGET_EPOCH", "1")), +} +``` + +Write the result to `$AI4S_SHARED_DIR/models/ORBIT-2/perf-runs//interm_8m_lux_profile.json`. Use Python (`json.load`/`json.dump`) — do NOT do this with sed. + +(The wrapper `examples/sbatch_train_perf_amd.sh` already does this; this step exists in the launcher prompt for the case where the launcher is reused for a non-ORBIT-2 model.) + +### 4. Submit the job + +```bash +export AI4S_SHARED_DIR=/path/to/shared +export ORBIT2_OUTPUT_DIR=$AI4S_SHARED_DIR/models/ORBIT-2/perf-runs/PENDING-$$ +mkdir -p "$ORBIT2_OUTPUT_DIR" + +cd "$REPO_ROOT" +OUT=$(sbatch \ + --partition= \ + --account= \ + --nodes=2 \ + --ntasks-per-node=8 \ + --gpus-per-node=8 \ + --cpus-per-task=16 \ + --time=00:30:00 \ + --output="${ORBIT2_OUTPUT_DIR}/orbit2-train-%j.out" \ + --error="${ORBIT2_OUTPUT_DIR}/orbit2-train-%j.out" \ + material_science/models/ORBIT-2/examples/sbatch_train_perf_amd.sh) +JOBID=$(echo "$OUT" | awk '{print $NF}') +echo "Submitted JOBID=$JOBID" +``` + +The sbatch wrapper inherits these env vars and exports `ORBIT2_CONFIG_OVERRIDE`, `OMNISTAT_VENV`, `OMNISTAT_CONFIG`, `OMNISTAT_USERMODE_INTERVAL`, `PROFILE_RANK0_ONLY=1` to the rank script. + +### 5. Wait for terminal state + +Poll `sacct -j $JOBID -X -n --format=State,ExitCode,Elapsed,NodeList -P` every 30 s. Treat any of `COMPLETED|FAILED|TIMEOUT|CANCELLED|NODE_FAIL|OUT_OF_MEMORY` as terminal. **Hard ceiling 45 minutes** of polling — if not terminal, write `state=TIMEOUT_WAITING` to manifest and exit 1. + +### 6. Move artifacts into final perf-run dir + +Once terminal: + +```bash +PERF_RUN=$AI4S_SHARED_DIR/models/ORBIT-2/perf-runs/${JOBID} +mv "$ORBIT2_OUTPUT_DIR" "$PERF_RUN" +# Find the rank-0 trace +TRACE=$(find "$PERF_RUN/logs" -name '*.pt.trace.json*' 2>/dev/null | head -1 || true) +# Find the omnistat DB (the sbatch wrapper writes it under the perf-run dir) +DB="$PERF_RUN/omnistat-db" +``` + +### 7. Write the manifest + +Use Python in the omnistat venv to write the manifest JSON exactly per the schema above. Include all `tools_versions` captured in step 1. + +### 8. Final stdout + +Print: +``` +STATUS=ok; reason=jobid= state= runtime=s perf_run= +``` + +## Failure modes (must surface, never silently swallow) + +| Failure | Action | +|---|---| +| Merge conflict in step 1 | Exit 2 with the message in the script. | +| `sbatch` returns nonzero | Exit 3, print sbatch stderr verbatim. | +| Job state == `FAILED` | Still write the manifest (with `exit_state=FAILED`), but **also** print the last 50 lines of the SLURM out to help downstream subagents skip cleanly. | +| No trace file found | Manifest gets `trace_paths=[]`. tracelens_analyst will gracefully no-op. | +| omnistat-db empty | Manifest still written; omnistat_analyst will detect and report it. | + +## Notes for the implementing agent + +- The launcher subagent runs on the **login node**. Do not call `srun` from inside the launcher; only `sbatch`. +- Do not import torch in the launcher's venv — it should stay lightweight. +- All paths absolute, never relative. +- Use `flock` on the install step if multiple invocations may race: `( flock -x 9; install_steps; ) 9>$TOOLS/.install.lock`. diff --git a/earth_science/models/ORBIT-2/recipes/perf-analysis/agents/omnistat_analyst.md b/earth_science/models/ORBIT-2/recipes/perf-analysis/agents/omnistat_analyst.md new file mode 100644 index 0000000..56e9fef --- /dev/null +++ b/earth_science/models/ORBIT-2/recipes/perf-analysis/agents/omnistat_analyst.md @@ -0,0 +1,143 @@ +# omnistat_analyst subagent + +Drive `omnistat-inspect` (PR #271) through the analyze-job phases on the user-mode VictoriaMetrics database, then map findings to the bottleneck taxonomy. + +## Inputs + +- `/manifest.json` +- The Omnistat DB at `manifest.omnistat_db_path`. +- The `omnihub-inspect` venv at `${OMNIHUB_TOOLS_DIR}/omnihub-inspect/`. +- The VictoriaMetrics binary at `${OMNIHUB_TOOLS_DIR}/victoriametrics/victoria-metrics-prod`. + +## Outputs + +- `/omnistat/inspect_outputs/db_info.json` +- `/omnistat/inspect_outputs/job_info.json` +- `/omnistat/inspect_outputs/data_check.json` +- `/omnistat/inspect_outputs/health.json` +- `/omnistat/inspect_outputs/stats_global.json` +- `/omnistat/inspect_outputs/stats__.json` — for any anomalous category, drilled to `node` and `gpu-id`/`interface-id`. +- `/omnistat/report_summary.md` +- `/omnistat/claims.json` (same schema as `tracelens/claims.json`) + +## Steps + +### 1. Start VictoriaMetrics on the DB (login node) + +Following the [load-database SKILL](https://github.com/ROCm/omnistat/blob/jorda/skills/skills/load-database/SKILL.md), with one addition for memory-constrained login nodes: pass `-fs.disableMmap` so VM can open a job-scoped DB inside the login-node cgroup (without it, even a 1.6 MB DB panics with "cannot mmap file with size 4096 bytes ... no such device"). + +```bash +DB=$(jq -r .omnistat_db_path ) +[[ -d "$DB/data" ]] || { echo "STATUS=fail; reason=no_db_data"; exit 1; } + +# Pick free port +PORT=8428 +ss -lnt "sport = :$PORT" 2>/dev/null | grep -q LISTEN && PORT=8429 + +VMLOG=$PERF_RUN/omnistat/vm.log +mkdir -p "$PERF_RUN/omnistat/inspect_outputs" +nohup ${OMNIHUB_TOOLS_DIR}/victoriametrics/victoria-metrics-prod \ + -storageDataPath="$DB" \ + -httpListenAddr=127.0.0.1:$PORT \ + -retentionPeriod=100y \ + -search.disableCache \ + -search.latencyOffset=0 \ + -search.maxPointsPerTimeseries=90000 \ + -fs.disableMmap \ + > "$VMLOG" 2>&1 & +VM_PID=$! +echo $VM_PID > "$PERF_RUN/omnistat/vm.pid" +TSDB_URL="http://127.0.0.1:$PORT" + +# Wait up to 15 s for ready +for i in {1..15}; do + sleep 1 + curl -sf "$TSDB_URL/api/v1/status/tsdb" > /dev/null && break +done + +# Trap-style cleanup is the orchestrator's job; this subagent leaves VM running +# so the omnistat_verifier can hit the same endpoint. +``` + +### 2. Run analyze-job phases (per the SKILL) + +```bash +. ${OMNIHUB_TOOLS_DIR}/omnihub-inspect/bin/activate +SCRATCH=$PERF_RUN/omnistat/scratch +mkdir -p "$SCRATCH" +JOBID=$(jq -r .jobid ) + +OUTD=$PERF_RUN/omnistat/inspect_outputs + +omnistat-inspect --tsdb-url $TSDB_URL db info > $OUTD/db_info.json +omnistat-inspect --tsdb-url $TSDB_URL --scratch-dir $SCRATCH job $JOBID info > $OUTD/job_info.json +omnistat-inspect --tsdb-url $TSDB_URL --scratch-dir $SCRATCH job $JOBID data-check > $OUTD/data_check.json +omnistat-inspect --tsdb-url $TSDB_URL --scratch-dir $SCRATCH job $JOBID health > $OUTD/health.json +omnistat-inspect --tsdb-url $TSDB_URL --scratch-dir $SCRATCH job $JOBID stats --level global > $OUTD/stats_global.json +``` + +### 3. Identify anomalous categories and drill down + +For each category in `stats_global.json` (`gpu`, `host`, `network`, `xgmi`, `vendor`): + +- Compute coefficient of variation (`stddev/mean`) for utilization-like gauges. +- Flag a category as anomalous if **any** of: + - Mean GPU utilization < 60% (high cpu_dispatch / dataloader / comm risk) + - Network rx/tx rate stddev/mean > 0.3 (interface imbalance) + - VRAM stddev/mean > 0.2 (workload heterogeneity or memory leak) + - HBM throttle / power throttle counts > 0 + - XGMI traffic stddev/mean > 0.3 (unbalanced intra-node comm) + +For each flagged category, run `--level node` and either `--level gpu-id` (for `gpu`/`xgmi`) or `--level interface-id` (for `network`): + +```bash +for cat in gpu network xgmi host vendor; do + if is_anomalous "$cat"; then + omnistat-inspect --tsdb-url $TSDB_URL --scratch-dir $SCRATCH job $JOBID stats --category $cat --level node > $OUTD/stats_${cat}_node.json + case $cat in + gpu|xgmi) fine=gpu-id ;; + network) fine=interface-id ;; + *) fine="" ;; + esac + [[ -n "$fine" ]] && omnistat-inspect --tsdb-url $TSDB_URL --scratch-dir $SCRATCH job $JOBID stats --category $cat --level $fine > $OUTD/stats_${cat}_${fine//-/_}.json + fi +done +``` + +### 4. Translate findings into claims + +Concrete rules for the first iteration: + +- Mean GPU util < 50% → `class=cpu_dispatch` (or `dataloader`, lower confidence) — magnitude = `1 - mean_util`. +- HBM `vram_used_percentage` near 100% **and** `omnistat_fom` low → `class=gpu_memory_hbm`. +- Power throttling > 0 events → `class=gpu_compute` with remedy "raise power cap if available; otherwise system limit". +- Inter-node `network` rate << intra-node `xgmi` rate AND comm-heavy workload → `class=comm_scaleout` with remedy "verify ANP plugin via `NCCL_DEBUG=INFO`". +- High `network` interface CV at `interface-id` level → `class=comm_scaleout` with remedy "specific interface (e.g. ionic_3) is degraded; investigate cable/firmware". +- Hot GPU (>peak temp threshold from `gpus/mi355x.md` if present, else conservative 90°C) → `class=gpu_compute` with remedy "thermal-bound; check airflow". +- Per-node CPU mem usage CV > 0.3 → `class=host_cpu` with remedy "rebalance num_workers or check NUMA pinning". + +### 5. Emit report and claims + +`report_summary.md` should mirror Phase 1-3 of the analyze-job SKILL: discovery → data check → health → global stats → drill-down narrative. Cap at ~80 lines. + +`claims.json` schema is identical to the tracelens version. Cap at 6 claims. + +### 6. Final stdout + +``` +STATUS=ok; reason= claims; vm_port=; vm_pid= +``` + +The orchestrator (or the verifier, when finished) is responsible for stopping VM via `kill $(cat /omnistat/vm.pid)`. + +## Failure modes + +- DB empty (no metrics): write a single claim `class=host_io` `hypothesis="omnistat-usermode never collected"` and STATUS=partial. Likely cause: collector/server didn't start; check sbatch wrapper logs. +- omnistat-inspect not on PATH: STATUS=fail; the launcher should have installed it. +- VictoriaMetrics startup timeout (15 s): STATUS=fail; print last 20 lines of `vm.log`. + +## Notes + +- This subagent has read-write access to the perf-run dir. +- It must NOT modify or delete the omnistat DB. +- Never query `atlvrmonad01:9090` — it's IP-blocked and not the right database for this job anyway. diff --git a/earth_science/models/ORBIT-2/recipes/perf-analysis/agents/omnistat_verifier.md b/earth_science/models/ORBIT-2/recipes/perf-analysis/agents/omnistat_verifier.md new file mode 100644 index 0000000..deedcda --- /dev/null +++ b/earth_science/models/ORBIT-2/recipes/perf-analysis/agents/omnistat_verifier.md @@ -0,0 +1,92 @@ +# omnistat_verifier subagent + +Independently re-derive the top 2-3 claims from `omnistat/claims.json` by issuing raw PromQL via `curl` against the same VictoriaMetrics endpoint, at the **finest sampling step**. Optionally probe one cheap remedy on a 1-node interactive `srun`. + +## Inputs + +- `/manifest.json` +- `/omnistat/claims.json` +- VictoriaMetrics already running (PID at `/omnistat/vm.pid`, port from analyst output). +- `/omnistat/inspect_outputs/job_info.json` (for time range and sampling interval). + +## Outputs + +- `/omnistat/verified_claims.json` (same schema as `tracelens/verified_claims.json`). + +## Steps + +### 1. Re-derive top claims via raw PromQL + +Read the time range from `job_info.json` (`start_time`, `end_time`, `sampling_interval`). For each top claim, issue a `query_range` at the finest step (`max(sampling_interval, runtime/90000)`). + +Examples: + +**Important:** `rocm_*` and `omnistat_host_*` metrics do **not** carry a `jobid` label directly. Filter by joining with `rmsjob_info` per-instance: + +```promql +metric_name * on (instance) group_left() (max by (instance) (rmsjob_info{jobid="$JOBID"})) +``` + +A naive `rocm_utilization_percentage{jobid="..."}` will return 0 series even though the data is there. Use the join in every query. + +Also: `start_time`/`end_time` from `job_info.json` are **UTC**. Convert with `calendar.timegm(time.strptime(...))` (NOT `time.mktime`, which assumes local time) before passing to PromQL. + +| Claim | PromQL | +|---|---| +| "Mean GPU util = X%" | `avg(rocm_utilization_percentage * on (instance) group_left() (max by (instance) (rmsjob_info{jobid="$JOBID"})))` | +| "VRAM near 100% on N nodes" | `max(rocm_vram_used_percentage * on (instance) group_left() (max by (instance) (rmsjob_info{jobid="$JOBID"})))` per `instance` | +| "Network tx rate plateau" | `rate(omnistat_network_tx_bytes * on (instance) group_left() (max by (instance) (rmsjob_info{jobid="$JOBID"}))[60s])` peak | +| "XGMI imbalance" | per-card stddev of `rocm_xgmi_*_bytes` rate, joined as above | +| "Power throttling events" | sum increase of throttle counters over the job span, joined as above | + +```bash +PORT=$(jq -r '.[0].port' ) # or pull from analyst stdout +curl -sG "http://127.0.0.1:$PORT/api/v1/query_range" \ + --data-urlencode "query=avg_over_time(rocm_utilization_percentage{jobid=\"$JOBID\"}[1m])" \ + --data-urlencode "start=$START" \ + --data-urlencode "end=$END" \ + --data-urlencode "step=$STEP" \ + | jq '.data.result | map(.values[][1] | tonumber) | add/length' +``` + +The verifier should re-derive the actual number from the PromQL result and compare to `claim.magnitude.value`. ±20% → `verdict=verified`. Outside → `refuted`. Empty result series → `inconclusive`. + +### 2. Optional remedy probe + +Same constraints as the tracelens verifier — at most one 1-node `srun -p -A -N1 --time=00:05:00`. Telemetry-side probes are typically configuration changes (e.g. enable rocprofiler `hbm` counter set in the omnistat config and re-run for 1 minute on 1 node to see if HBM bandwidth is actually saturating). + +If a probe needs the omnistat config to change, write a temp config file based on `omnistat.config` with the rocprofiler section uncommented, point `OMNISTAT_CONFIG` at it, and submit via the existing sbatch script as a tiny job with `ORBIT2_NUM_EPOCH=1 HYDRAGNN_MAX_NUM_BATCH=10 --nodes=1`. Time-budget that branch generously (10 min) since it does include the omnistat collector startup. + +If multi-node is required (e.g. ANP plugin retest), set `remedy_probe.ran=false; notes="multi-node probe deferred"`. + +### 3. Emit `verified_claims.json` + +Same shape as the tracelens version. Each claim copied through with `verdict`, `verifier_evidence`, `verifier_value`, `remedy_probe`. + +### 4. Stop VictoriaMetrics + +After all queries are done: + +```bash +kill "$(cat /omnistat/vm.pid)" 2>/dev/null || true +``` + +The synthesizer doesn't need VM running. + +### 5. Final stdout + +``` +STATUS=ok; reason=v/r/i; probe= +``` + +## Failure modes + +- VM not running: try to start it the same way the analyst did before giving up; if still down, STATUS=fail. +- Query returns no data for a key metric the claim relies on: that's an `inconclusive` verdict, not a fail. +- Network blip mid-query: retry once with 5s sleep. + +## Notes + +- Use **only** `curl` for re-derivation, **not** `omnistat-inspect`. The whole point is an independent path. +- Use the finest step that VM allows (`runtime / 90000` floor); document the actual step used in `verifier_evidence`. +- Never delete the DB or any analyst output. diff --git a/earth_science/models/ORBIT-2/recipes/perf-analysis/agents/synthesizer.md b/earth_science/models/ORBIT-2/recipes/perf-analysis/agents/synthesizer.md new file mode 100644 index 0000000..a56862e --- /dev/null +++ b/earth_science/models/ORBIT-2/recipes/perf-analysis/agents/synthesizer.md @@ -0,0 +1,131 @@ +# synthesizer subagent + +Merge `tracelens/verified_claims.json` and `omnistat/verified_claims.json` into a single ranked, deduplicated bottleneck report. + +## Inputs + +- `/manifest.json` +- `/tracelens/verified_claims.json` +- `/omnistat/verified_claims.json` +- (Optional) `/tracelens/report_summary.md`, `/omnistat/report_summary.md` + +## Outputs + +- `/combined_report.md` + +## Steps + +### 1. Load and filter + +```python +import json +tl = json.load(open(f"{perf_run}/tracelens/verified_claims.json")) +om = json.load(open(f"{perf_run}/omnistat/verified_claims.json")) +all_claims = [{"src":"TL", **c} for c in tl] + [{"src":"OS", **c} for c in om] +``` + +Drop `verdict == "refuted"`. Keep `inconclusive` claims but tag them so the user sees the gap. + +### 2. Deduplicate by class + topic + +Group claims by `class`. Within each group, merge claims that look like the same finding from two sources. Heuristic: +- Both have `class=comm_scaleout` → almost certainly the same finding viewed from trace (NCCL kernels) and telemetry (network rates). Merge into one entry that lists both `src` values and both magnitudes. +- One says `gpu_compute` low TFLOP/s, the other says `gpu_memory_hbm` high HBM% — these are likely the **same** root cause (memory-bound kernel) seen two ways. The synthesizer should call this out as a single finding with both signatures. + +### 3. Rank + +Score = `magnitude.value × confidence_weight × (corroborated ? 1.5 : 1.0)`. +`confidence_weight`: high=1.0, medium=0.7, low=0.4. `corroborated` = both TL and OS contributed to the merged entry. + +### 4. Tag "system limit reached" + +A claim is a system limit if any of: +- `proposed_remedy is null` +- both `verdict=verified` and `remedy_probe.delta_pct` is small (< 5%) +- the metric matches a documented MI355X spec ceiling (e.g. fp64 39 TFLOP/s, HBM 8 TB/s, ANP scale-out ~25 GB/s) + +Mark these explicitly so the report doesn't promise a fix that won't materialize. + +### 5. Write `combined_report.md` + +Layout: + +```markdown +# ORBIT-2 bottleneck analysis — JOBID + +**Run:** 2 nodes × 8 GPUs (MI355X), s, precision=

, batch= +**Sources:** TraceLens v on rank-0 trace, Omnistat user-mode (s sampling) +**Verdict counts:** // + +## Executive summary + +<1 paragraph, 4-5 sentences. Named bottleneck class first, magnitude, then "system-limit" or "improvable"> + +## Ranked bottlenecks + +### 1. — [score: ] [] [] + +**Hypothesis:** +**Magnitude:** +**Evidence:** +**Remedy:** — **tried?** ; **delta:** +**Verifier note:** + +### 2. ... + +## Remedies tried in this session + +| Class | Probe | Baseline | After remedy | Delta | +|---|---|---|---|---| +| ... | + +## Remedies proposed but not tried (cost/risk) + +| Class | Remedy | Why deferred | +|---|---|---| + +## Limits reached + +| Class | Spec ceiling | Observed | +|---|---|---| + +## Inconclusive claims + +| Source | Class | Hypothesis | Why inconclusive | +|---|---|---|---| + +## Next steps + +1. +2. +3. + +## Artifacts + +- `tracelens/report.xlsx` +- `tracelens/verified_claims.json` +- `omnistat/inspect_outputs/` +- `omnistat/verified_claims.json` +- `manifest.json` + +--- +Generated by `material_science/models/ORBIT-2/recipes/perf-analysis/`. +``` + +### 6. Final stdout + +``` +STATUS=ok; reason=; top=; tried= +``` + +## Style rules + +- Specific numbers from evidence — never vague ("most of the time", "huge"). +- If TL and OS disagree on a number that's not refuted, show **both** with sources rather than pick one. +- When `remedy_probe.ran=false` because of multi-node requirement, say so explicitly. +- Do not editorialize about ORBIT-2 as a model; only about the run. + +## Failure modes + +- Both inputs missing → STATUS=fail. +- Only one input present → emit a partial report with a banner explaining which side is missing; STATUS=partial. diff --git a/earth_science/models/ORBIT-2/recipes/perf-analysis/agents/tracelens_analyst.md b/earth_science/models/ORBIT-2/recipes/perf-analysis/agents/tracelens_analyst.md new file mode 100644 index 0000000..e37f2c5 --- /dev/null +++ b/earth_science/models/ORBIT-2/recipes/perf-analysis/agents/tracelens_analyst.md @@ -0,0 +1,124 @@ +# tracelens_analyst subagent + +Generate a TraceLens performance report for the rank-0 PyTorch trace and produce structured `claims.json` mapped to the bottleneck taxonomy. + +## Inputs + +- `$AI4S_SHARED_DIR/models/ORBIT-2/perf-runs//manifest.json` + +## Outputs + +- `/tracelens/report.xlsx` — the raw TraceLens Excel. +- `/tracelens/report_summary.md` — human-readable digest of the workbook (top-N kernels, GPU idle %, comm vs compute split, roofline categorization). +- `/tracelens/claims.json` — structured claims (schema below). + +## Claim schema + +Every entry: + +```json +{ + "id": "TL-001", + "class": "gpu_compute|gpu_memory_hbm|cpu_dispatch|comm_xgmi|comm_scaleout|dataloader|host_io|host_cpu", + "hypothesis": "", + "magnitude": { + "metric": "", + "value": 0.34, + "unit": "fraction|seconds|GBps|TFLOPs|count" + }, + "evidence_paths": ["::", ":"], + "evidence_excerpt": "<200 chars max>", + "confidence": "high|medium|low", + "proposed_remedy": "", + "remedy_test_command_or_null": "srun -p -A -N1 --time=00:05:00 ... | null" +} +``` + +## Steps + +### 1. Sanity check + +```bash +. ${OMNIHUB_TOOLS_DIR}/omnihub-inspect/bin/activate +PERF_RUN=$(jq -r .perf_run_dir ) +TRACE=$(jq -r '.trace_paths[0] // empty' ) +[[ -z "$TRACE" ]] && { echo "STATUS=partial; reason=no_trace_in_manifest"; exit 0; } +mkdir -p "$PERF_RUN/tracelens" +``` + +### 2. Run TraceLens + +```bash +cd "$PERF_RUN/tracelens" +TraceLens_generate_perf_report_pytorch \ + --profile_json_path "$TRACE" \ + --output_xlsx_path report.xlsx \ + --output_csvs_dir csvs/ \ + --enable_kernel_summary \ + --short_kernel_study 2>&1 | tee tracelens.log +``` + +CLI confirmed against TraceLens v0.1.0+: `--output_xlsx_path` controls the workbook path; `--output_csvs_dir` writes per-sheet CSVs (handy for the verifier). `--enable_kernel_summary` adds the kernel-level summary sheet. `--short_kernel_study` flags very short kernels (likely launch-overhead bound). + +If multi-node, the script will detect collective ops and produce a separate "Collective Analysis" sheet unless `--disable_coll_analysis` is passed (don't pass it). + +### 3. Extract digest from the workbook + +Use `openpyxl` (already a TraceLens dep) in the venv: + +```python +import openpyxl, json, pathlib +wb = openpyxl.load_workbook("report.xlsx", data_only=True, read_only=True) +sheets = wb.sheetnames # expected: GPU_Timeline, Ops_Summary, Roofline, NN_Module, ... +``` + +Pull these specific signals (skip a sheet gracefully if absent): + +| Sheet | Signals to extract | +|---|---| +| `GPU_Timeline` | total_gpu_time, idle_time, kernel_time, comm_time (sum durations of NCCL/RCCL events) | +| `Ops_Summary` | top-10 ops by total time + their percent | +| `Roofline` | per-op TFLOP/s and TB/s; flag any in compute-bound vs memory-bound region | +| `NN_Module` | top-3 modules by GPU time | + +Write the digest to `report_summary.md` as a few markdown tables. + +### 4. Translate findings into claims + +Concrete rules for the first iteration (the agent may add more): + +- If `comm_time / total_step_time > 0.20` → `class=comm_xgmi` if intra-node only OR `comm_scaleout` if inter-node (detect by NCCL kernel name `AllReduce` + `nNodes>1` flag). Magnitude = the ratio. +- If `idle_time / total_step_time > 0.15` → `class=cpu_dispatch` (could also be dataloader; raise as `confidence=medium`). +- For each top-3 op: classify by Roofline region. Compute-bound + low TFLOP/s → `class=gpu_compute`. Memory-bound + high TB/s → `class=gpu_memory_hbm`. Memory-bound + low TB/s → `class=gpu_memory_hbm` with confidence=low + remedy "verify with HBM counters via Omnistat rocprofiler probe". +- If `dataloader` events appear in the timeline (look for `enumerate(DataLoader)#`) and their cumulative time > 5% → `class=dataloader`. + +### 5. Emit `claims.json` + +Cap at 6 claims. Order by `magnitude.value` (when comparable). Each claim must include a remedy that is concrete and **verifiable**. Examples: + +| Hypothesis | Remedy | +|---|---| +| "RCCL allreduce dominates 35% of step time" | "Try `RCCL_LL128_FORCE_ENABLE=0` or larger bucket size via `find_unused_parameters=False`" | +| "fp64 GEMM at 12 TFLOP/s, well below MI355X 39 TFLOP/s peak" | "Test `ORBIT2_PRECISION=bf16` for 5 batches on 1 node" | +| "DataLoader events occupy 18% of step time" | "Increase `HYDRAGNN_NUM_WORKERS` from 0 to 4" | +| "fp64 SQ FMA approaches MI355X spec ceiling" | "null — system limit; only path is bf16/fp32" | + +A null remedy is acceptable for "limit reached" findings. + +### 6. Final stdout + +``` +STATUS=ok; reason= claims, top=: +``` + +## Failure modes + +- TraceLens import error → STATUS=fail; do not attempt to "patch" the install. +- Trace too large to load (>2 GB) → write a single claim of class `cpu_dispatch` with hypothesis "trace too large to fully analyze; profiler captured >X seconds — schedule may be misconfigured". +- No NCCL events in trace → multi-node still ran but profiler missed comm; flag as `inconclusive` rather than claiming "comm is fast". + +## Notes + +- Do NOT run `omnistat-inspect`; that's the other analyst's job. If you need the omnistat DB to validate something, hand the question to the omnistat_verifier instead. +- Do NOT modify the raw trace file. +- Keep `evidence_paths` machine-parseable so the verifier can re-derive the number from the same source. diff --git a/earth_science/models/ORBIT-2/recipes/perf-analysis/agents/tracelens_verifier.md b/earth_science/models/ORBIT-2/recipes/perf-analysis/agents/tracelens_verifier.md new file mode 100644 index 0000000..c29c1f9 --- /dev/null +++ b/earth_science/models/ORBIT-2/recipes/perf-analysis/agents/tracelens_verifier.md @@ -0,0 +1,106 @@ +# tracelens_verifier subagent + +Independently re-derive the top 2-3 claims from `tracelens/claims.json` against the raw `.pt.trace.json` and (optionally) probe **one** cheap remedy on a 1-node interactive `srun`. + +## Inputs + +- `/manifest.json` +- `/tracelens/claims.json` +- The raw trace at `manifest.trace_paths[0]`. + +## Outputs + +- `/tracelens/verified_claims.json` — same schema as `claims.json` plus a `verdict` field per entry. + +## Verdict schema + +Each input claim is copied through with one new field: + +```json +{ + "...all fields from claims.json...": "...", + "verdict": "verified|refuted|inconclusive", + "verifier_evidence": "<200 chars: how you re-derived it>", + "verifier_value": , + "remedy_probe": { + "ran": true|false, + "command": "", + "baseline_seconds_per_step": , + "remedy_seconds_per_step": , + "delta_pct": , + "notes": "" + } +} +``` + +## Steps + +### 1. Re-derive top 3 claims from the raw JSON + +```python +import json, gzip +from pathlib import Path +def open_trace(p): + return gzip.open(p, 'rt') if str(p).endswith('.gz') else open(p, 'r') +with open_trace(trace) as f: + data = json.load(f) +events = data["traceEvents"] +``` + +For each top claim, do a parallel re-derivation: + +| Claim class | Re-derivation | +|---|---| +| `comm_xgmi` / `comm_scaleout` | sum `dur` for events whose `name` matches `^(nccl|rccl|ncclKernel|AllReduce|AllGather|ReduceScatter)`; divide by total span (`max(ts+dur)-min(ts)` of GPU events). Compare to claim's `magnitude.value`. | +| `gpu_compute` | filter `cat=='kernel'`, find top-N by `sum(dur)`; check the top kernel name + percent matches the claim. | +| `cpu_dispatch` | `(total_span - sum(GPU kernel dur)) / total_span` should match "idle %". | +| `dataloader` | sum `dur` of events with `name` containing `enumerate(DataLoader)`. | + +If verifier_value within ±20% of claim's value → `verdict=verified`. If outside → `verdict=refuted` with the actual number. If we can't isolate the events → `verdict=inconclusive`. + +### 2. Pick at most ONE remedy probe + +Choose the highest-impact `verified` claim with a non-null `remedy_test_command_or_null`. Constraints: + +- 1 node only (`-N1`). +- ≤5 minutes wall-time. +- Must produce a comparable number (steps/sec or seconds/step) to the original 2-node baseline. If the remedy can only be tested at scale (e.g. RCCL inter-node tweak), set `remedy_probe.ran=false; notes="remedy requires multi-node; deferred to next iteration"`. + +Example probe for the bf16 / fp64 ceiling claim — run via the existing ORBIT-2 sbatch script in interactive mode: + +```bash +srun -p -A -N1 --ntasks=8 --gpus-per-node=8 --cpus-per-task=16 \ + --time=00:05:00 --pty bash -c ' +export AI4S_SHARED_DIR=/path/to/shared +export ORBIT2_OUTPUT_DIR=/tmp/perf-probe-$$ +mkdir -p "$ORBIT2_OUTPUT_DIR" +# Prefer submitting via sbatch_train_amd.sh from the repo; ad-hoc srun probes are site-specific. +bash "$REPO_ROOT/earth_science/models/ORBIT-2/examples/sbatch_train_amd.sh" \ + 2>&1 | tee /tmp/probe-$SLURM_JOB_ID.out +grep -oE "[0-9.]+s/it" /tmp/probe-$SLURM_JOB_ID.out | tail -5 +' +``` + +Note: `sbatch_train_amd.sh` is designed for `sbatch`, not `srun --pty`. For probes prefer using its rank-script logic directly via a small ad-hoc wrapper rather than running the whole sbatch under srun. The verifier may inline a minimal Python invocation that bypasses sbatch and just runs the upstream `interm_8m_lux_all_mpnn.py` for 10 batches inside the container with the alternate precision. + +If the verifier judges that no cheap probe is feasible, set `remedy_probe.ran=false` with a clear note and proceed. + +### 3. Emit `verified_claims.json` + +Same length and order as `claims.json`. Even refuted claims stay in the file — the synthesizer will drop them, but they should remain visible. + +### 4. Final stdout + +``` +STATUS=ok; reason=v/r/i; probe= +``` + +## Constraints + +- Never modify `claims.json` — only write `verified_claims.json`. +- Never run a 2-node `srun` or `sbatch` from this subagent. +- If the trace file is missing, write `verified_claims.json` as `[]` and `STATUS=partial; reason=no_trace`. + +## Why two phases (analyst + verifier)? + +Hallucination defense. The analyst can produce a number from a tool's report; the verifier re-derives it from raw events. Disagreement => the user gets to see both numbers and the synthesizer flags the conflict in the final report rather than silently picking one. This is the same approach the omnistat side uses with PromQL-from-curl re-derivation. diff --git a/earth_science/models/ORBIT-2/recipes/perf-analysis/omnistat.config.template b/earth_science/models/ORBIT-2/recipes/perf-analysis/omnistat.config.template new file mode 100644 index 0000000..80817de --- /dev/null +++ b/earth_science/models/ORBIT-2/recipes/perf-analysis/omnistat.config.template @@ -0,0 +1,133 @@ +#-- +# Omnistat user-mode config for AMD Instinct MI355X (gfx950). +# +# Adapted from HydraGNN's upstream omnistat.hydragnn-external.config. +# Differences from upstream / from the system-mode default: +# - enable_amd_smi (newer SMI), not rocm_smi +# - rocprofiler counters DISABLED in v1 (turn on per-probe via verifier) +# - allowed_ips includes 0.0.0.0 so the user-mode VM scrapes +# - victoria_datadir under the per-job dir (@JOB_DIR@), persists after job +#-- + +[omnistat.collectors] + +port = 8101 +enable_rocm_smi = False +enable_amd_smi = True +enable_ras_ecc = True +enable_rms = True +enable_rocprofiler = True +enable_network = True +enable_vendor_counters = True +enable_xgmi = True +enable_cu_occupancy = False + +# Kernel-dispatch tracing collector (per-kernel name + duration histogram on +# every node). Disabled by default: it requires the standalone tool library +# `libomnistat_trace.so` to be built (see rocprofiler-sdk/README.md +# "Kernel Tracing Library") and loaded into the workload via +# ROCP_TOOL_LIBRARIES. The sbatch wrapper enables this automatically when +# OMNISTAT_KERNEL_TRACE=1 is set; otherwise it stays False. +# +# IMPORTANT: kernel tracing and the rocprofiler device-counting service +# (enable_rocprofiler above) BOTH talk to rocprofiler-sdk. They can co-exist +# on a single GPU, but if you're chasing per-kernel attribution it's usually +# cleaner to either (a) run kernel-trace alone, or (b) run device counting +# alone — see ai4science-studio SKILL §11 for the tradeoffs. +enable_kernel_trace = False + +# Path to local ROCm install (validated on ROCm 7.2.x) +rocm_path = /opt/rocm + +# allowed scrapers: localhost (the user-mode VM sits on the same node) + 0.0.0.0 +# wildcard for any user-mode VictoriaMetrics that might run on a sibling node. +allowed_ips = 127.0.0.1, 0.0.0.0 + +[omnistat.collectors.rms] + +host_skip = "login.*" +enable_annotations = True +# Note: /tmp/omni_rmsjobinfo (the upstream default) is often owned by the +# system-mode collector's service user on compute nodes, so user-mode +# omnistat-rms-env hits PermissionError if it tries to write there. The +# sbatch wrapper substitutes @JOB_DIR@ with the job-scoped perf-run dir +# (NFS-mounted on every compute node) before passing the resolved config +# to omnistat-usermode --configfile. +job_detection_mode = file-based +job_detection_file = @JOB_DIR@/omni_rmsjobinfo +step_detection_file = @JOB_DIR@/omni_rmsjobinfo_step + +[omnistat.collectors.host] +enable_proc_io_stats = True + +# Rocprofiler — HBM bandwidth + FP64 VALU/MFMA counters in one block. +# Counter group selection follows the user-validated pattern (one profile, +# one sub-section, all counters listed together; rocprofiler-sdk +# time-multiplexes hardware counter sets internally as needed). +# +# All eight counters below are defined for gfx950 (MI355X) in +# /opt/rocm/share/rocprofiler-sdk/counter_defs.yaml; verified before pinning. +# +# Useful derivations from these raw counters (compute in the verifier or +# downstream; rocprofiler-sdk also exposes derived counters +# MfmaFlopsF64 and TotalFlopsF64 if a downstream consumer prefers those): +# peak_fp64_GFLOP/s = (SQ_INSTS_VALU_MFMA_MOPS_F64 * 512 +# + (SQ_INSTS_VALU_FMA_F64*2 +# + SQ_INSTS_VALU_ADD_F64 +# + SQ_INSTS_VALU_MUL_F64) * 64 +# ) / wall_seconds / 1e9 +# HBM_GB/s = (FETCH_SIZE + WRITE_SIZE) / wall_seconds / 1e9 +[omnistat.collectors.rocprofiler] +profile = hbm_flops_f64 + +[omnistat.collectors.rocprofiler.hbm_flops_f64] +# MI355X (gfx950) rocprofiler counter limits — see ai4science-studio SKILL §10. +# +# Constraints surfaced on MI355X (gfx950) / ROCm 7.2.2: +# * FETCH_SIZE + WRITE_SIZE in ONE profile fails: "error code 38: Request +# exceeds the capabilities of the hardware to collect" (both in TCC block). +# * FETCH_SIZE OR WRITE_SIZE + 4 x SQ_INSTS_VALU_*_F64 in one profile works. +# * sampling_mode = periodic in omnistat 1.12.0+git.65ea9ac (PR #271 + main) +# ONLY fires the first counter set — half the data silently lost. (See +# SKILL §11.) +# +# Workaround: use `constant` mode with a single counter set that fits the +# hardware limits and the omnistat bug. We capture HBM read traffic only +# (FETCH_SIZE) plus the 4 fp64 instruction counters needed to derive the +# fp64 FLOP rate using the formula: +# fp64_GFLOP/s = (MFMA_MOPS_F64 * 512 + (FMA_F64*2 + ADD_F64 + MUL_F64) * 64) +# / wall_seconds / 1e9 +# Drop SQ_INSTS_VALU_TRANS_F64 (not relevant for GEMM-heavy ML training). +# +# WRITE_SIZE (HBM-write bandwidth) is NOT collected in this profile until +# upstream omnistat fixes the periodic-mode bug. For most ML training +# workloads, FETCH_SIZE is the dominant component of HBM traffic anyway. +# +# All counters on one line: configparser multi-line values require every +# continuation line (including the closing `]`) to have leading whitespace. +sampling_mode = constant +counters = ["FETCH_SIZE", "SQ_INSTS_VALU_ADD_F64", "SQ_INSTS_VALU_MUL_F64", "SQ_INSTS_VALU_FMA_F64", "SQ_INSTS_VALU_MFMA_MOPS_F64"] + +[omnistat.query] +prometheus_url = http://localhost:8428 +system_name = mi355x + +#-- +# User-mode settings — VictoriaMetrics + per-job datadir +# (Like the [rms] section above, @JOB_DIR@ is substituted by the sbatch +# wrapper at runtime; ConfigParser's %(VAR)s interpolation does NOT pull +# from os.environ, so we can't rely on %(SLURM_JOB_ID)s here.) +#-- +[omnistat.usermode] + +victoria_binary = @OMNIHUB_TOOLS_DIR@/victoriametrics/victoria-metrics-prod +victoria_datadir = @JOB_DIR@/omnistat-db +victoria_logfile = vic_server.log + +# No external VM; we keep everything local to the job. +external_victoria = False + +# Bind exporter to a single core to keep monitoring overhead off the data path. +exporter_corebinding = 0 + +ssh_key = ~/.ssh/id_rsa diff --git a/earth_science/models/ORBIT-2/recipes/perf-optimizer-loop/README.md b/earth_science/models/ORBIT-2/recipes/perf-optimizer-loop/README.md new file mode 100644 index 0000000..94d3e2c --- /dev/null +++ b/earth_science/models/ORBIT-2/recipes/perf-optimizer-loop/README.md @@ -0,0 +1,11 @@ +# ORBIT-2 Iterative Systems Optimizer Loop (deferred) + +**Status:** Not started. Complete the scaling sweep (1/2/4/8 nodes) and review [../perf-analysis/HANDOFF.md](../perf-analysis/HANDOFF.md) before porting the HydraGNN `perf-optimizer-loop` recipe. + +Planned entry points: + +- `examples/run_optimizer_loop.sh` +- `lever_catalog.yaml` (RCCL, batch_size, num_workers, FSDP layout, TILES) +- Agent prompts under `agents/` + +See [material_science/models/HydraGNN/recipes/perf-optimizer-loop/](../../../material_science/models/HydraGNN/recipes/perf-optimizer-loop/) for the reference implementation. From 1d19ae623dcb6a03be39deedbf66d27290028233 Mon Sep 17 00:00:00 2001 From: Ashwin Aji Date: Wed, 17 Jun 2026 07:46:43 +0000 Subject: [PATCH 28/31] perf tooling: portable cluster config + neutral perf-tools naming Add a perf_tools block to the committed cluster-config example and teach /init-cluster to discover and write it. Decouples the perf-analysis and perf-optimizer-loop recipes from any site-specific tool location by sourcing PERF_TOOLS_DIR from .cluster-config.yaml (perf_tools.dir). Co-Authored-By: Claude --- .claude/commands/init-cluster.md | 26 +++++++++++--- .cluster-config.example.yaml | 7 ++++ .../skills/ai4science-earth-science/SKILL.md | 12 +++++-- .../skills/ai4science-perf-analysis/SKILL.md | 33 +++++++++++++++--- .cursor/skills/ai4science-studio/SKILL.md | 34 ++++++++++++++++--- .gitignore | 3 ++ 6 files changed, 97 insertions(+), 18 deletions(-) diff --git a/.claude/commands/init-cluster.md b/.claude/commands/init-cluster.md index 0e2b5cf..960d525 100644 --- a/.claude/commands/init-cluster.md +++ b/.claude/commands/init-cluster.md @@ -67,6 +67,12 @@ echo "HTTPS_PROXY=${HTTPS_PROXY:-}" # 13. Existing SIF files find "$HOME" /scratch /projects /opt -maxdepth 4 -name "*.sif" 2>/dev/null | head -20 + +# 14. Performance tooling (perf-analysis / perf-optimizer-loop recipes — optional) +# Look for an existing shared perf-tools dir (omnistat + TraceLens venv "perf-inspect/"). +for d in /shared/*/tools /shared/perf-tools "$HOME/perf-tools"; do + [[ -d "$d/perf-inspect" || -d "$d/omnistat-src" ]] && echo "perf_tools: $d" +done ``` ## Step 2 — Present findings and ask user to confirm or choose @@ -119,7 +125,12 @@ Explain the expected layout under that root: Show result of connectivity test. Ask to confirm. If no internet, ask for proxy settings. -**Q9. Config file location** +**Q9. Performance tooling** (optional — only for `perf-analysis` / `perf-optimizer-loop` recipes) +If a perf-tools dir was discovered (step 14), show it and ask to confirm; otherwise ask for the path +or leave blank. This is exported to perf scripts as `PERF_TOOLS_DIR`; expected layout under it: +`perf-inspect/` (omnistat + TraceLens venv), `omnistat-src/`, `victoriametrics/victoria-metrics-prod`. + +**Q10. Config file location** - `.cluster-config.yaml` in this repo (recommended — per-checkout) - `~/.config/ai4science-studio/cluster.yaml` (user-level, shared across clones) @@ -156,10 +167,15 @@ gpu: network: internet_access: proxy: "" - -rccl: - socket_ifname: "" # discover: ip link show up - ib_hca: "" # discover: ibstat | grep CA + mgmt_iface: "" # discover: ip -o link show up + ib_hca: "" # discover: ibstat | grep CA + rccl_anp_plugin: "" # path to librccl-anp.so on host (default: /opt/rocm/lib/librccl-anp.so) + libionic_path: "" # path to libionic.so.1 (default: /usr/lib/x86_64-linux-gnu/libionic.so.1) + +# Performance tooling — only if using perf-analysis / perf-optimizer-loop recipes (Q9). +# Exported to scripts as PERF_TOOLS_DIR. Leave dir blank if not used. +perf_tools: + dir: "" # layout: perf-inspect/, omnistat-src/, victoriametrics/ discovered_sifs: # SIF files found during init (informational — confirm they are under sif_cache) diff --git a/.cluster-config.example.yaml b/.cluster-config.example.yaml index 6068434..ea5cd7b 100644 --- a/.cluster-config.example.yaml +++ b/.cluster-config.example.yaml @@ -40,3 +40,10 @@ network: ib_hca: "" # IB HCA device list for RCCL (discover: ibstat | grep 'CA ') rccl_anp_plugin: "" # Path to librccl-anp.so on host (default: /opt/rocm/lib/librccl-anp.so) libionic_path: "" # Path to libionic.so.1 on host (default: /usr/lib/x86_64-linux-gnu/libionic.so.1) + +# Performance tooling (perf-analysis / perf-optimizer-loop recipes) +# Shared location of the omnistat + TraceLens perf tools, exported to scripts as +# PERF_TOOLS_DIR. Expected layout under this dir: perf-inspect/ (Python venv with +# omnistat + TraceLens), omnistat-src/, victoriametrics/victoria-metrics-prod. +perf_tools: + dir: "" # e.g. "/shared/perf-tools"; leave blank if not using perf recipes diff --git a/.cursor/skills/ai4science-earth-science/SKILL.md b/.cursor/skills/ai4science-earth-science/SKILL.md index 2fa2e62..bc529d0 100755 --- a/.cursor/skills/ai4science-earth-science/SKILL.md +++ b/.cursor/skills/ai4science-earth-science/SKILL.md @@ -39,6 +39,7 @@ The `earth_science/` domain covers **climate**, **weather**, and broader **Earth - **Overlay size:** 7 GB ext3; ~2 GB content (pytorch-lightning, xformers, mpi4py, wandb, tensorboard, etc.) - **pytorch-lightning must be installed with `--no-deps`** — it declares `torch` as a dep; pip resolves CUDA torch from PyPI which silently replaces the ROCm torch in PYTHONPATH. - **xformers for ROCm:** install from the PyTorch ROCm wheel index, not PyPI. PyPI xformers links against `libcudart.so.12`. Use `--index-url https://download.pytorch.org/whl/${ROCM_WHL_TAG}` with `--no-deps`. +- **xFormers MEA vs `xformers.info` on MI355 (2026-06):** On `pytorch_rocm7.2.2` SIF + ORBIT-2 overlay, `python3 -m xformers.info` reports **`memory_efficient_attention.ckF` available**, but **`xformers.ops.memory_efficient_attention` (even with no explicit `op`, bf16, tiny shapes) can SIGSEGV**; `torch.nn.functional.scaled_dot_product_attention` is fine. The **`rocm7.2`** PyTorch wheel index only published **xformers 0.0.34 and 0.0.35** at probe time — no newer wheel to “upgrade CK” without changing ROCm/PyTorch image. Use **`ORBIT2_FUSED_ATTN=DEFAULT`** / SDPA-style paths for production; optional repro details in [HANDOFF.md](../../../earth_science/models/ORBIT-2/recipes/perf-analysis/HANDOFF.md) §xFormers CK probe. **Do not conflate** that path with PyTorch’s **`SDPBackend.EFFICIENT_ATTENTION`** (Studio benchmark label `ck` in [`upstream_pytorch_sdpa_benchmark.py`](../../../earth_science/models/ORBIT-2/examples/upstream_pytorch_sdpa_benchmark.py) — upstream [`sdpa.py`](https://github.com/pytorch/pytorch/blob/main/benchmarks/transformer/sdpa.py)); the latter is SDPA-only and is the right knob to compare “efficient vs math” on ROCm. To show segfaults are **stack/shape-generic** rather than Bayes-CAST-specific, run the same script with **`--orbit-include-xformers-ck`** (subprocess-isolated xFormers **`MemoryEfficientAttentionCkOp`** on micro + asymmetric cases). - **xformers.components shim:** ORBIT-2 imports `xformers.components.attention.core.scaled_dot_product_attention`. This subpackage was removed in xformers ≥0.0.28. Write a three-file shim after install: ```python (xf / "components" / "__init__.py").write_text("") @@ -51,10 +52,15 @@ The `earth_science/` domain covers **climate**, **weather**, and broader **Earth ``` - **Distributed launch:** `srun --mpi=pmix apptainer exec ... --env HOSTNAME="$MASTER_ADDR"` — see studio SKILL.md for full pattern. Confirmed working for 1-GPU and 8-GPU. - **Data blocker:** ORBIT-2 real training data lives on Frontier Lustre (`/lustre/orion//...`) — not publicly accessible. Synthetic data covers the smoke-test path. -- **Training on institutional clusters:** `sbatch_train_amd.sh` + `run_orbit2_train.py` wrap `intermediate_downscaling.py` with Apptainer overlay, `srun --mpi=pmix`, and `interm_8m_lux.yaml`. Staged Constellation/Globus 10.0_arcmin NPZ data can be used in **same-dir mode** (`low_res` = `high_res`, `superres_mag: 1`) for perf/scaling timing only. -- **Training landmines (2026-06):** stub `gptl4py` in `run_orbit2_train.py`; overlay may lack `xformers.ops` → auto `FusedAttn.DEFAULT`; `max_epochs` must be ≥ 2; batch cap patches module before `main()` not via `runpy.run_path`. -- **Perf/scaling:** `sbatch_train_perf_amd.sh`, `run_scaling_study.sh`, artifacts under `$AI4S_SHARED_DIR/models/ORBIT-2/perf-runs/` and `outputs/scaling-*`. See [recipes/perf-analysis/HANDOFF.md](../../../earth_science/models/ORBIT-2/recipes/perf-analysis/HANDOFF.md). +- **Training on institutional clusters:** `sbatch_train_amd.sh` prefers **`launch/train_edm.py`** (Bayes-CAST EDM: **`exec python3 train_edm.py /config/config.yaml`** from **`/orbit2/launch`** — upstream **`launch/launch_diffusion.sh` is not used**: OLCF-style wrapper, ignores argv, nests `srun`). Otherwise, if **`launch_diffusion.sh`** exists under `ORBIT2_ROOT`, ranks run it after **`orbit2_rank_hook_runner.py`**; else **`run_orbit2_train.py`** wraps `intermediate_downscaling.py` (gptl4py stub, FusedAttn fallback, batch cap). For **1 node × 8 GPUs**, rendered YAML defaults to **`fsdp=8`**, **`simple_ddp=1`** unless `ORBIT2_FSDP` / `ORBIT2_SIMPLE_DDP` are set. Set `ORBIT2_CONFIG_TEMPLATE=interm_8m_lux_era5.yaml` (Lux) or **`edm_8m_era5_1x8.yaml`** (Bayes EDM) and point `ORBIT2_DATA_ROOT` at staged ERA5 1.0° same-dir NPZ. +- **Training landmines (2026-06):** stub `gptl4py` in `run_orbit2_train.py`; overlay may lack `xformers.ops` → **`ORBIT2_FUSED_ATTN=DEFAULT`** in sbatch. **`ORBIT2_DATA_TYPE`** defaults **`bfloat16`** (matches upstream **`configs/interm_8m.yaml`** and Bayes-CAST-style **`edm_8m_era5.yaml`**); use **`float32`** only to isolate ROCm/hipBLAS issues. `max_epochs` must be ≥ 2; batch cap patches module before `main()` not via `runpy.run_path`. +- **bf16 + ROCm Flash SDPA 65535 grid cap (2026-06, FIXED):** Bayes-CAST EDM `var_agg`/`temporal_agg` cross-attention flattens `(B,History,L)` into the SDPA **batch** dim (`B = batch·648`), and ROCm **Flash** SDPA caps that dim at **65535** → bf16 crashes with `HIP error: invalid argument` at **batch ≥ ~128** (batch 64 OK). Fix applied **globally** to the compute-side clone: all `Attention` + `CrossAttention` SDPA calls wrapped in `sdpa_kernel([SDPBackend.EFFICIENT_ATTENTION, SDPBackend.MATH])` (MATH fallback-only). Going global is safe because `--orbit-selfattn` benchmarks EFFICIENT within ±3–5% of Flash on the main 648-token self-attention. Validated job 9238 (bf16 b256, decreasing loss). Details in [HANDOFF.md](../../../earth_science/models/ORBIT-2/recipes/perf-analysis/HANDOFF.md) §RESOLVED. **Generalizes** to any model folding extra dims into the attention batch. +- **Perf/scaling:** `sbatch_train_perf_amd.sh` defaults **Bayes-CAST** template **`edm_8m_era5_1x8.yaml`** + ERA5 1.0° data; **`ORBIT2_ROOT`** prefers **`…/code/bayes-cast`** when present. Bayes EDM clears **`ORBIT2_RANK_PRE_TRAIN_HOOK`** only when unset — export an absolute hook path **before** `sbatch` for sysopt (`torch.compile`, etc.). **`ORBIT2_MAX_EPOCH=6`** default. See [HANDOFF.md](../../../earth_science/models/ORBIT-2/recipes/perf-analysis/HANDOFF.md), **[one-node GPU baseline](../../../earth_science/models/ORBIT-2/recipes/perf-analysis/one-node-gpu-baseline.md)**, **[perf-optimizer-loop](../../../earth_science/models/ORBIT-2/recipes/perf-optimizer-loop/README.md)** (throughput-primary loop), and **`ORBIT2_NUM_WORKERS`** / template token **`__NUM_WORKERS__`** in `edm_8m_era5_1x8.yaml` via `render_orbit2_config.py`. +- **Perf saturation vs ~10M params:** Default `interm_8m_lux*.yaml` uses `res_slimvit` @ `embed_dim` 256 / `depth` 6 → ~10M weights; **PRISM `spatial_resolution` 18** yields tiny activations per step — raise **batch**, use **ERA5 111×111** template for heavier forwards, or widen **`embed_dim`/`depth`** (with `embed_dim % num_heads == 0`) to approach MI-class MFMA; see [one-node-gpu-baseline.md](../../../earth_science/models/ORBIT-2/recipes/perf-analysis/one-node-gpu-baseline.md) §3. **VRAM-first sweep:** [BASELINE_LOCKIN.md](../../../earth_science/models/ORBIT-2/recipes/perf-analysis/BASELINE_LOCKIN.md) (batch sweep before widening ViT). - **Multi-node scheduling:** Before large submits, check partition state (`sinfo`, site dashboards). Prefer **idle, healthy** nodes; use `--nodelist` / `--exclude` only per your site's policy. Do not stack two exclusive full-node GPU jobs on the same host. +- **GEMM-time bottleneck = the 3/4-channel conv stem/head, NOT attention/MLP (2026-06, MI355X):** Full TraceLens+Omnistat analyst/verifier run (`recipes/perf-analysis/agents/orchestrator_gemm_analysis.md` + `examples/run_gemm_analysis.sh`) on profiled EDM (bf16, batch 4096, ERA5) found one hipBLASLt GEMM `Cijk_…MT16x16x32` = **~43 % of the step at both 1- and 2-node** — the **im2col-lowered low-channel conv backward** (3/4-channel projections, 16×16 tile @ ~0.03 TFLOP/s). The tunable `nn.Linear` GEMMs are ~2 % and near-peak → **this is why TunableOp gave no uplift.** **Most promising lever = channel-padding in NCHW (timing-only −26 %, but TABLED pending a convergence study — see next bullet); `channels_last` is a DEAD END on ROCm 7.2.2 (also next bullet).** **Profile any EDM run** via `train_edm.py:_orbit2_make_profiler` (env-gated: set `ORBIT2_PROFILE_DIR`+`PROFILE_TARGET_EPOCH`; sbatch does this). **Trust raw kernel-name device-time, not TraceLens `report.xlsx:ops_summary_by_category`** — it mislabels conv-lowered GEMMs as `CONV_bwd` (a 2-node analyst claim was refuted on this). `examples/compare_trace_kernels.py` aggregates kernels tool-independently. Separately, **~43 % whole-job idle** (inter-step overhead, NOT comm-bound) is a second lever. +- **Conv channel-padding = TABLED timing-only candidate (−26 %, NOT adopted); `channels_last`/implicit-GEMM is a DEAD END (2026-06, MI355X/ROCm 7.2.2):** Env knob **`ORBIT2_CONV_PAD=N`** (`edm.py`, **default 0/off, keep off**) widens the `path2`/`refine` internal conv width to a tile-friendly N (=16) while keeping 3-channel I/O → collapses the starved `MT16x16x32` GEMM into `MT16x16x128`, **−25/26 % steady step time, −28 % GPU-kernel time** (timing-only A/B). **Tabled because it changes model capacity (wider convs) and convergence/quality was never validated** — "loss sanity" ≠ convergence. Do NOT enable for real training until a matched convergence study (final val loss/skill, steps-to-target, pad=0 vs pad=16) passes. **Do NOT stack `channels_last`** (knob `ORBIT2_CHANNELS_LAST`, default 0): on this build MIOpen has no tuned NHWC implicit-GEMM solver for these shapes and falls back to brute-force `naive_conv_…wrw_nhwc` (~88-91 % of kernel time) → **7-10× SLOWER** (cl_pad0 60.8 s, cl_pad16 80.1 s vs 6.38 s for nchw_pad16). The `im2col` (`Im2d2Col_v2`) in the NCHW path is NOT wasted overhead — it is the entry to the tuned hipBLASLt GEMM kernels; NHWC discards that path. **Stay NCHW + `ORBIT2_CONV_PAD`.** +- **ORNL Frontier MIOpen flags now default in the sbatch (validate per-cluster):** From bayes-cast `launch/launch_diffusion.sh`, ORNL **disables Winograd** + unbounds the multi-pass Winograd workspace, "tested many times on Frontier" (gfx90a/ROCm7.1.1). Ported to `sbatch_train_perf_amd.sh`/`sbatch_train_amd.sh`/`sbatch_infer_*.sh` as overridable defaults: `MIOPEN_DEBUG_CONV_WINOGRAD=0` (`ORBIT2_MIOPEN_CONV_WINOGRAD`), `MIOPEN_DEBUG_AMD_WINOGRAD_MPASS_WORKSPACE_MAX=-1`, `MIOPEN_DEBUG_AMD_MP_BD_WINOGRAD_WORKSPACE_MAX=-1`, `HSA_FORCE_FINE_GRAIN_PCIE=1`. `MIOPEN_DISABLE_CACHE=1`+`MIOPEN_USER_DB_PATH=/tmp/` already matched ORNL (keep; do NOT persist the find-db). **Do NOT port** the Frontier Slingshot/Cray-specific env (`FI_CXI_*`, `NCCL_NET="AWS Libfabric"`, `librccl-net.so`, `NCCL_SOCKET_IFNAME=hsn0`, Frontier `LD_PRELOAD` paths) — they break non-Slingshot fabrics (e.g. Lux IB/ionic). ### ArchesWeather (geoarches) diff --git a/.cursor/skills/ai4science-perf-analysis/SKILL.md b/.cursor/skills/ai4science-perf-analysis/SKILL.md index ae1fa20..2d82df9 100644 --- a/.cursor/skills/ai4science-perf-analysis/SKILL.md +++ b/.cursor/skills/ai4science-perf-analysis/SKILL.md @@ -9,7 +9,7 @@ description: Runs the AMD AI agents bottleneck-analysis workflow on a HydraGNN ( The user wants an automated bottleneck analysis of a multi-node training/inference run using AMD's open-source observability tooling (TraceLens, Omnistat). This is **distinct** from the `ai4science-run-models` skill — that one launches a model; this one diagnoses a model after it runs. -Default target: HydraGNN on AMD MI355X. **ORBIT-2 training** uses the same perf-analysis pattern; see `earth_science/models/ORBIT-2/recipes/perf-analysis/`. +Default target: HydraGNN on AMD MI355X. **ORBIT-2 training** uses the same perf-analysis pattern; see `earth_science/models/ORBIT-2/recipes/perf-analysis/`. **ORBIT-2 iterative sysopt** (throughput-primary FOM) lives in `earth_science/models/ORBIT-2/recipes/perf-optimizer-loop/` (`run_optimizer_loop.sh`, `lever_catalog.yaml`). ## Repository entry points @@ -24,14 +24,16 @@ Default target: HydraGNN on AMD MI355X. **ORBIT-2 training** uses the same perf- ### ORBIT-2 - Recipe: [earth_science/models/ORBIT-2/recipes/perf-analysis/](../../../earth_science/models/ORBIT-2/recipes/perf-analysis/) — see [`HANDOFF.md`](../../../earth_science/models/ORBIT-2/recipes/perf-analysis/HANDOFF.md) +- **Iterative sysopt loop:** [earth_science/models/ORBIT-2/recipes/perf-optimizer-loop/](../../../earth_science/models/ORBIT-2/recipes/perf-optimizer-loop/) — primary accept/revert FOM **`throughput_samples_per_s`** (`run_fom_extractor.py` + `manifest.global_batch_size`) - Sbatch: [earth_science/models/ORBIT-2/examples/sbatch_train_perf_amd.sh](../../../earth_science/models/ORBIT-2/examples/sbatch_train_perf_amd.sh) - Plain training: [earth_science/models/ORBIT-2/examples/sbatch_train_amd.sh](../../../earth_science/models/ORBIT-2/examples/sbatch_train_amd.sh) - Scaling: [run_scaling_study.sh](../../../earth_science/models/ORBIT-2/examples/run_scaling_study.sh) + [collate_scaling_study.py](../../../earth_science/models/ORBIT-2/examples/collate_scaling_study.py) -- FOM parser: [parse_training_log.py](../../../earth_science/models/ORBIT-2/examples/parse_training_log.py) — primary FOM **`steady_batch_time_s`** (epochs ≥ 2, exclude batch 0 per epoch); **`loss_sanity_pass`** requires strictly decreasing epoch losses +- FOM parser: [parse_training_log.py](../../../earth_science/models/ORBIT-2/examples/parse_training_log.py) — latency **`steady_batch_time_s`**; throughput when `global_batch_size` is known; **`loss_sanity_pass`** requires strictly decreasing epoch losses +- FOM extractor: [run_fom_extractor.py](../../../earth_science/models/ORBIT-2/examples/run_fom_extractor.py) — writes `foms.json`; optional PromQL via `ORBIT2_TSDB_URL`; omnistat template defaults to **bf16 MFMA** profile `hbm_flops_bf16` (`FETCH_SIZE` + `SQ_INSTS_VALU_MFMA_MOPS_BF16`); set **`OMNISTAT_ROCPROF_PROFILE=hbm_flops_f64`** to revert to fp64 HydraGNN-style counters - Artifacts: `$AI4S_SHARED_DIR/models/ORBIT-2/perf-runs//` -- Workload: `intermediate_downscaling.py` via [run_orbit2_train.py](../../../earth_science/models/ORBIT-2/examples/run_orbit2_train.py) (gptl4py stub, FusedAttn fallback, batch cap) -- Data: 10.0_arcmin PRISM same-dir (`interm_8m_lux.yaml`) or ERA5 1.0° same-dir (`interm_8m_lux_era5.yaml`, `ORBIT2_CONFIG_TEMPLATE`) — timing only until true downscaling targets are staged -- Profiler: [orbit2_profiler_hook.py](../../../earth_science/models/ORBIT-2/examples/orbit2_profiler_hook.py) via `ORBIT2_RANK_PRE_TRAIN_HOOK` (no JSON Profile block) +- Workload: `intermediate_downscaling.py` via [run_orbit2_train.py](../../../earth_science/models/ORBIT-2/examples/run_orbit2_train.py), or Bayes-CAST **`train_edm.py`** / **`launch_diffusion.sh`** when present (see `sbatch_train_amd.sh` / `sbatch_train_perf_amd.sh`). Studio `orbit2_rank_hook_runner.py` applies `ORBIT2_RANK_PRE_TRAIN_HOOK` before the launcher (set hook path **before** `sbatch` so the wrapper bakes it into the rank script). +- Data: 10.0_arcmin PRISM same-dir (`interm_8m_lux.yaml`) or ERA5 1.0° same-dir (`interm_8m_lux_era5.yaml`, `ORBIT2_CONFIG_TEMPLATE`) — timing only until true downscaling targets are staged; for HBM saturation see [STAGING_ERA5_FOR_HBM.md](../../../earth_science/models/ORBIT-2/recipes/perf-optimizer-loop/STAGING_ERA5_FOR_HBM.md) +- Profiler: [orbit2_profiler_hook.py](../../../earth_science/models/ORBIT-2/examples/orbit2_profiler_hook.py) via `ORBIT2_RANK_PRE_TRAIN_HOOK` (Lux path; Bayes EDM defaults to no hook unless overridden) ## Orchestration loop (what the main agent does) @@ -128,4 +130,25 @@ The first 2-node end-to-end run exposed five real bugs that are now all worked a 20. **Attribution pass tooling:** Use [`examples/run_fom_extractor.py`](../../../material_science/models/HydraGNN/examples/run_fom_extractor.py) + TraceLens on `$AI4S_SHARED_DIR/models/HydraGNN/perf-runs//`. Backfill `manifest.json` if missing (loop jobs before launcher wrote it). **VictoriaMetrics PromQL on perf-run DBs requires `time=` at the job window** — instant queries on the login node return empty series even when `omnistat_hardware_counter` data exists. **`kernel_correlation_summary.attribution_quality=poor` on a ~6 s profile-epoch kineto slice is often a windowing artifact** (busy_frac <0.5 in every 1 s bucket); trust TraceLens `ops_summary.csv` (`aten::mm` + `aten::bmm` ≈80% kernel time) before escalating to `OMNISTAT_KERNEL_TRACE=1`. See [`recipes/perf-optimizer-loop/dispatch-attribution.md`](../../../material_science/models/HydraGNN/recipes/perf-optimizer-loop/dispatch-attribution.md). +## Lessons captured (ORBIT-2 bf16 enablement, 2026-06-11) + +21. **ROCm Flash SDPA has a 65535 batch-grid cap — fatal for attention that inflates the batch dim.** ORBIT-2 Bayes-CAST EDM crashed on the first `training_step` in bf16 with `HIP error: invalid argument` (`hipErrorInvalidValue`), apparently at `var_agg`'s `proj` Linear (`attention.py:280`). **The Linear was a red herring** (async error misattribution): a standalone bf16 Linear at that shape passes even at M=663552. The real culprit is `var_agg`/`temporal_agg` cross-attention, which flattens `(B, History, L)` into the SDPA **batch** dim, so `B = batch·tokens` (e.g. `batch·648`). PyTorch SDPA's **default backend picks ROCm Flash**, which maps that batch dim to a HIP grid dimension capped at **65535** → fails once `batch·648 > 65535` (i.e. **batch ≥ ~128**; batch 64→PASS, 128→FAIL). fp32 "worked" only because it dispatched to the math backend. + - **Isolate fast with a faithful op-replica on 1 GPU** (no data, ~15 s): `crossattn` mode reproduces; `math`/`efficient`/`eager` backends all PASS, `default` FAILs at batch 128. This pattern (faithful op replica on one GPU) beats full-node data-loading jobs for kernel bugs. + - **Fix:** wrap the offending module's SDPA in `with torch.nn.attention.sdpa_kernel([SDPBackend.EFFICIENT_ATTENTION, SDPBackend.MATH]):` to exclude Flash (MATH fallback-only). Validated end-to-end: ORBIT-2 job **9238** (bf16 b256) COMPLETES with decreasing loss. **Prefer applying it globally** (all self- + cross-attention SDPA): the [`--orbit-selfattn`](../../../earth_science/models/ORBIT-2/examples/upstream_pytorch_sdpa_benchmark.py) benchmark shows **EFFICIENT within ±3–5% of Flash** for normal self-attention (q=kv=648) across batch 64/256/1024, so unifying on EFFICIENT removes the foot-gun at ~noise-level cost. **Do NOT make MATH primary** — it's 10–15× slower (≈15 vs ≈290 TFLOPS); it just never gets selected when EFFICIENT supports the shape. Use `--orbit-varagg` to confirm the Flash cap and `--orbit-selfattn` to confirm the perf parity before deciding scope. + - **Generalize:** any model that folds extra dims into the attention batch (per-variable, per-timestep, per-patch cross-attention) can hit this in bf16 at large batch on ROCm. Suspect Flash's grid cap before blaming GEMM/BLAS; `HIP_LAUNCH_BLOCKING=1` mis-points to the next CPU-sync op, so trust an op-replica probe over the traceback line. + - **FOM parser (RESOLVED 2026-06-11):** this EDM build logs per-step loss as `epoch: N batch_idx M ... loss tensor(L)` and per-step wall time as `my rank 0. tic4-tic1 in X seconds` — it has **no `Batch N: … seconds`** / `Epoch N completed` line. `parse_training_log.py` now reads **both** this EDM format and the older `intermediate_downscaling` format (pairs each `tic4-tic1` to the preceding step; synthesizes per-epoch loss from per-step losses). Validated on jobs 9234 (1236 s/s) & 9257 (4097 s/s). Crashed jobs (no `batch_idx` lines) correctly yield `epoch_records=0` + null throughput — that's the "never reached steady state" signal, **not** a parser bug. Do not rewrite the parser. + +## Lessons captured (ORBIT-2 bf16 optimizer loop, loop-bf16-loop-001, 2026-06-12) + +22. **ORBIT-2 EDM (1×8 MI355X) is COMPUTE-bound, not dispatch-/dataloader-bound — the opposite of HydraGNN MLIP.** The fp32-vs-bf16 discriminator (job 9264 vs 9261, same batch 1024) showed **bf16 = 2.38× fp32** (4089 vs 1718 samples/s). A ≥1.4× precision speedup ⇒ MFMA/GEMM-bound. Consequences confirmed in-loop: `torch_compile_edm` was neutral (-0.3%, within noise — Inductor finds little to fuse on an already-GEMM-heavy ViT/EDM), and `num_workers_8` **regressed -5.6%** because 8 workers × 8 ranks = 64 dataloader procs oversubscribe the 56 schedulable CPUs (cpus-per-task=7). **For compute-bound ORBIT-2 the only real throughput levers are bigger effective batch (needs more staged data) or GEMM tuning (TunableOp — later shown NOT blocked, see #28).** Baseline (bf16 + EFFICIENT SDPA + batch 1024) remained the best config; no catalog lever beat it at 1 node. Keep `num_workers ≤ 4` for ORBIT-2 to avoid CPU oversubscription. +23. **`num_workers` silently changes the EFFECTIVE per-step batch in the Bayes-CAST IterableDataset → phantom throughput.** Loop overnight-3x-001 iter-1 set `num_workers=4` and `run_fom_extractor.py` reported **+129%** (7215 vs 3150 s/s) — a pure measurement artifact. Throughput = *nominal* `global_batch_size` ÷ steady step time, but with more workers the real per-step batch shrank ~2.4× (HBM 236→98 GB, `max_batches_per_epoch` 3→4, epoch loss 2.86→3.50). **Guard:** `foms.json` now carries `hbm_reserved_GB`, `hbm_reserved_pct_288`, `max_batches_per_epoch`; reject any cross-run throughput delta where HBM% or batches/epoch deviates from baseline. The Opus-4.8 overnight orchestrator caught this unaided via the HBM+batch cross-check — an unguarded loop would have crowned a phantom 2.3× best and poisoned every downstream iteration. For dataloader levers, compare only runs with matching HBM%/batches-per-epoch. +24. **3× ERA5 staging breaks the 1704-sample cap and enables true HBM saturation.** `stage_era5_3x_symlink.sh` symlinks the one staged year (1979, 20 shards) as 1980_/1981_ (60 shards; the loader globs `train/*.npz` with no year filter — `itermodule.py:133`). On 3× data, bf16 batch 4096 reaches **82% HBM (236/288 GB)** at 3152 s/s, and HBM scales ~linearly with batch (1024→20.6%, 2048→40.5%, 4096→82.0%). Perf/throughput study only — data repeats, not valid for science. The overnight loop on this config (4 iters) found **no env lever beats baseline** (num_workers artifact rejected; fsdp_prefetch +0.99%, sdpa_efficient +1.03% — both noise-floor), reconfirming compute-bound. Real wins need multi-node scale-out, real multi-year data, or GEMM tuning (TunableOp — now `experimental`, not blocked; see #28). +25. **Max-VRAM ≠ max-throughput at a small data cap.** At the 1704-sample staging cap, batch 1024 (~21% HBM) gave the best throughput (4098 s/s) while batch 1704 (~34% HBM) was **27% slower** (2975 s/s) — batch > sample-count just repeats/partial-fills a step. Lock the optimizer baseline at the *throughput-optimal* batch, and only chase HBM% after staging more ERA5 (raises the cap so larger batches do real work). See [BASELINE_LOCKIN.md](../../../earth_science/models/ORBIT-2/recipes/perf-analysis/BASELINE_LOCKIN.md). +27. **Multi-node ORBIT-2 on lux MI355X needs GID-consistent nodes + a non-restrictive `LD_LIBRARY_PATH` (RoCE/RCCL).** First 2-node ORBIT-2 attempts (jobs 10583/10584) aborted with `ionic_comp_msn: cqe with error 12` / `NET/IB ... status=12 vendor err 11 RETRY_EXC`. Two distinct causes, both fixed in `sbatch_train_perf_amd.sh`: (a) it pinned `--env LD_LIBRARY_PATH=…/torch/lib` only, so the host ANP plugin (`/opt/rocm/lib/librccl-anp.so`) couldn't resolve its `/opt/rocm/lib` deps — widened to include `/opt/rocm/lib:/usr/lib/x86_64-linux-gnu` (HydraGNN never set this var, which is why it worked). (b) **RoCE GID-index asymmetry:** `NCCL_IB_GID_INDEX=1` assumes the fabric ULA (`fd93:16d3:59b6:014f/RoCEv2`) sits at GID idx1 on every node, but a subset of nodes carry a stray global IPv6 (`2001:19f0:…`) at idx1**, shifting the fabric ULA to idx2 → `localGid`(ULA)/`remoteGid`(global) mismatch → no route → retry-exceeded abort. The other nodes were uniform (idx1=ULA); pinning to GID-consistent nodes (`--exclude` the bad range) makes 2-node work with **0 cqe errors** (validated job 10587 on a4-a5, then loop jobs 10588+). Enumerate GIDs with `cat /sys/class/infiniband/ionic_3/ports/1/gids/{1,2}` and `…/gid_attrs/types/*`. `NCCL_IB_GID_INDEX` is now overridable in the sbatch. Proper long-term fix is admin-side (remove stray global IPv6 from b-nodes' RoCE NICs, or subnet-based GID selection). Driver `run_2node_scaleout_loop.sh` can auto-exclude a known-bad range via `O2_EXCLUDE_PATTERN` (default: none). **Always verify GID-table consistency before blaming the model for a multi-node hang.** **2-node result (loop-2node-scaleout-001):** baseline **6046 s/s vs 1-node 2987 → 2.02× (101% strong-scaling efficiency)** — FSDP all-gather over RoCE is fully hidden behind compute (consistent with compute-bound). All comm/compute levers neutral (IB QPS 2/4, NCCL_MIN_NCHANNELS=112, torch.compile: ±1.5%); none beat baseline. **`GPU_MAX_HW_QUEUES=4` aborts at init (rc=134, "terminate called recursively", before any step) on MI355X — keep the validated `=2`.** +28. **TunableOp is NOT platform-blocked on MI355X / ROCm 7.2.2 / PyTorch 2.10 — the HydraGNN "blocked" verdict (#8/#15) does NOT generalize.** A dedicated isolated probe (single GPU, `perf-runs/tunableop-probe/`, jobs 10595/10596) ran `PYTORCH_TUNABLEOP_ENABLED=1 PYTORCH_TUNABLEOP_TUNING=1` over 17 GEMMs — including the ORBIT-2 **large-M bf16 `var_agg` shape class** (`mm` NN *and* `F.linear` TN-with-bias: M up to 262144, e.g. `tn_1024_131072_1024`) that the sbatch comment claims throws `hipErrorInvalidValue` under hipBLASLt. **All 14/14 matmuls tuned cleanly, 0 memory faults, a valid `tunableop_results.csv` was written** (hipBLASLt was actually *selected* for every large-M bias GEMM; rocBLAS won a few NN cases). **Prereqs are all present:** gfx950 Tensile logic libs exist in `/opt/rocm/lib/hipblaslt/library` (439 gfx950 files) and `/opt/rocm/lib/rocblas/library` (156) — nothing is missing from the container. Validators: `HIPBLASLT_VERSION=100202-dabb6df2b9`, `ROCBLAS_VERSION=5.2.0.dabb6df2b9`, `GCN_ARCH_NAME=gfx950:sramecc+:xnack-`. **Interpretation:** the HydraGNN "Memory access fault on all 16 GPUs during tuning" (#8/#15) is workload/concurrency-specific (full distributed run, MACE double-backward shapes), NOT a hipBLASLt-on-gfx950 prerequisite gap. **Action:** ORBIT-2's `tunable_op_warmup_then_use` is now `status: experimental` (not `blocked`); the remaining unknown is a full FSDP/all-rank warmup run (the HydraGNN fault only manifested under 16-GPU distributed tuning). Probe API note: `torch.cuda.tunable.write_file()` does not exist in 2.10 — the cache auto-writes on exit; use `set_filename()` + `get_results()`. **Lesson: do not copy a `blocked` verdict across workloads/models without re-probing — the cheap isolated GEMM probe (~2 min, 1 GPU) settles it.** +29. **TunableOp on ORBIT-2's 8-rank FSDP path does NOT fault (HydraGNN block disproved again), BUT live tuning is impractically slow for its giant patch-token GEMMs.** Sanity jobs 10597/10598 ran the full 8×MI355X FSDP EDM with `ORBIT2_TUNABLEOP_MODE=tune` (`PYTORCH_TUNABLEOP_TUNING=1`): **0 memory faults**, per-rank caches (`tunableop_results_rank.csv`) written cleanly — the distributed path HydraGNN crashed on (#8/#15) is fine for ORBIT-2. **However**, the dominant GEMM is `tn_256_5308416_256` (M=5.3M flattened patch tokens; the backward is a K=5.3M reduction), and TunableOp spent **~3.5 min on the first giant op and >12 min on the second** before it was cancelled — a *full* tune would be many hours. Root cause: per-candidate **kernel instantiation** (hipBLASLt+rocBLAS enumerate hundreds of algos; each is JIT-instantiated on these unusual huge shapes). This is **NOT** fixable by env knobs: `PYTORCH_TUNABLEOP_ROTATING_BUFFER_SIZE=0` (verified to *deactivate* rotation per the API docstring — "equal to zero means deactivate"), `MAX_TUNING_ITERATIONS`, and `MAX_TUNING_DURATION_MS` only bound per-*solution* timing, not the candidate count. **Practical guidance:** for giant-GEMM models, either (a) time-box the tune and accept a *partial* cache (the 2-3 giant ops dominate runtime, so even a partial cache captures most of any uplift), or (b) skip TunableOp and spend the node-hours on 2-node scale-out (known 2.02×). Wiring added to `sbatch_train_perf_amd.sh`: `ORBIT2_TUNABLEOP_MODE=off|tune|use` + `ORBIT2_TUNABLEOP_DIR` (rw-bound shared cache; per-rank files keyed on `SLURM_PROCID` so 8 ranks don't clobber one `results0.csv`). +30. **END-TO-END VERDICT: TunableOp does NOT help ORBIT-2 and is a correctness hazard — do not adopt.** Live tuning of the giant patch-token GEMMs is impractically slow (#29), so it was driven offline (per-shape GPU-pinned subprocesses; all 102 GEMM shapes incl. M=31.8M/15.8M/5.3M giants tuned). Result: **1-node tuned DIVERGES TO NaN (reproduced 2×: epoch-0 loss 5.5 vs baseline 0.87, then NaN); 2-node tuned = +0.44% (noise).** Two lessons: (a) TunableOp has **no numerical-correctness gate**, so a fast-but-wrong bf16 kernel can win and get cached — and a tuned cache is **layout-specific** (bit 1-node `fsdp=8` shapes, not 2-node `fsdp=2`/`ddp=8`), so it must be re-validated per parallelism config. (b) Even with all giants tuned there is **zero perf upside** — the default rocBLAS/hipBLASLt path is already near-optimal, confirming ORBIT-2 is bottlenecked by the GEMMs themselves, not kernel selection. Offline-tuning traps if ever revisited: pin workers with **only** `ROCR_VISIBLE_DEVICES` (setting `HIP_VISIBLE_DEVICES` too double-filters → `hipErrorNoDevice`); key cache-completeness on the core `(op, transAB_m_n_k)` (TunableOp rewrites the `_ld_*` suffix). **Real throughput levers remain: multi-node scale-out (linear) and bigger effective batch — not TunableOp.** If revisited, add a per-kernel correctness gate before caching. +32. **Perf/throughput runs must NOT persist model checkpoints — the Bayes-CAST EDM trainer saves one every epoch by default.** `train_edm.py:save_checkpoint()` is called unconditionally at the end of every epoch (line ~1390) and writes a **~71 MB `interm_epoch_.ckpt`** to the hardcoded `cp_save_path = "checkpoints/bayes-cast"` (relative to `launch/`, i.e. `$ORBIT2_ROOT/launch/checkpoints/bayes-cast/`). The config's `trainer.checkpoint:` field is empty but that only controls *loading/resume*, not saving — saving is unconditional. Files are overwritten per run (fixed names) so they don't grow per-job, but a 6-epoch run still leaves ~340 MB of unwanted state. **Fix:** added an env gate `ORBIT2_DISABLE_CKPT=1` → `save_checkpoint()` early-returns on ALL ranks together (no unbalanced FSDP `state_dict()` gather / barrier). `sbatch_train_perf_amd.sh` now defaults it to **1** (perf script → never save); `sbatch_train_amd.sh` defaults it to **0** (general training keeps checkpoints) and `run_scaling_study.sh` passes `=1`. All perf submitters (`sweep_*`, `submit_perf_baseline_*`, the optimizer/2-node/tunableop loops) route through `sbatch_train_perf_amd.sh` so they inherit the skip. The `initial_*.pth` save path (train_edm.py:188) only fires when `tensor_par_size>1` (we use 1) so it is not a concern. **Lesson: when standing up a throughput harness around an upstream trainer, audit its default checkpoint/save behavior — it will silently write state every epoch unless gated.** +26. **`num_workers` is NOT a throughput lever for ORBIT-2 — confirmed by a clean controlled test, and the FOM extractor is now hardened against the artifact.** The Bayes-CAST IterableDataset shards files by `per_worker = n_files // (num_workers × data_par_size)` (`data_par_size = fsdp × simple_ddp = 8`), and **each worker batches its own shard independently** — so more workers means smaller/fragmented batches. Staging 5× ERA5 (100 shards) so `per_worker = 3` lets full batches refill to 4096 at `num_workers=4`, but each worker still emits a partial **1016**-sample trailing batch (4096 + 1016 = 3 files × ~1704). Jobs 10504 (nw=1) vs 10505 (nw=4), batch 4096, identical otherwise: **nominal throughput said +80% (2993→5379 s/s); real-per-step throughput showed +2.7% (2987→3068 s/s) = noise**, and the full-4096 step time was ~equal (11.0 vs 10.8 s). **Fix shipped:** `parse_training_log.py` now reads the real per-step batch from the EDM `y.shape` line and `aggregate_foms` computes `throughput_method=real_per_step_batch` (Σ real_samples × ranks ÷ Σ step_time), plus `partial_step_fraction` and `steady_realized_batch_dims` as integrity surfaces. Still reject cross-run deltas where `steady_realized_batch_dims`/`partial_step_fraction` differ from baseline. Reconfirms ORBIT-2 is compute-bound; keep `num_workers ≤ 4` for I/O overlap only. + The full list (with cross-cutting context) is at the end of [.cursor/skills/ai4science-studio/SKILL.md](../ai4science-studio/SKILL.md). When fixing a bug here, check there first and propagate the lesson. diff --git a/.cursor/skills/ai4science-studio/SKILL.md b/.cursor/skills/ai4science-studio/SKILL.md index 161c671..a07b2c0 100755 --- a/.cursor/skills/ai4science-studio/SKILL.md +++ b/.cursor/skills/ai4science-studio/SKILL.md @@ -261,6 +261,8 @@ find "$HOME" /scratch /projects /opt -maxdepth 4 -name "*.sif" 2>/dev/null | hea SIF files are squashfs (read-only). All `pip install` inside `apptainer exec` must use `--target

`. Set `PYTHONPATH=:...`. +The bundled **`/opt/venv`** may also reject **`pip install --user`** (“User site-packages are not visible in this virtualenv”) — use **`--target`** + `PYTHONPATH` for small extras (e.g. `tabulate` / `tqdm` for ORBIT-2 [`upstream_pytorch_sdpa_benchmark.py`](../../../earth_science/models/ORBIT-2/examples/upstream_pytorch_sdpa_benchmark.py); see [HANDOFF.md](../../../earth_science/models/ORBIT-2/recipes/perf-analysis/HANDOFF.md) §PyTorch SDPA benchmark). + **Torch strip is safe** in overlay/pip-packages — the SIF's venv torch is untouched and remains the primary copy. ## No-overlay install-to-tmp pattern @@ -675,18 +677,18 @@ Omnistat exposes **two** independent rocprofiler-sdk integrations, and there's a **Cross-runtime gotcha:** `enable_kernel_trace=True` without the tool library loaded does *nothing* — the omnistat collector just sits with an empty endpoint. The `OMNISTAT_KERNEL_TRACE=1` switch in `sbatch_train_perf_amd.sh` enforces this with a hard `exit 2` if `libomnistat_trace.so` isn't found at the configured path; do **not** loosen that check, otherwise you re-create the exact "silent zero counters" failure mode from §8. -**Build location:** Both omnistat artifacts (the Python extension and the kernel-trace tool library) must be built on a **compute node** inside the same SIF the workload runs in — login nodes don't have apptainer or the ROCm headers, and a host-side build will pick up the wrong libcurl / glibc. `OMNIHUB_TOOLS_DIR` is the shared tools dir from `.cluster-config.yaml` `omnihub.tools_dir` (e.g. `/shared/omnihub/tools`); export it first. Standard build: +**Build location:** Both omnistat artifacts (the Python extension and the kernel-trace tool library) must be built on a **compute node** inside the same SIF the workload runs in — login nodes don't have apptainer or the ROCm headers, and a host-side build will pick up the wrong libcurl / glibc. `PERF_TOOLS_DIR` is the shared tools dir from `.cluster-config.yaml` `perf_tools.dir` (e.g. `/path/to/perf-tools`); export it first. Standard build: ```bash salloc -p -A -N 1 --time=00:15:00 --gpus-per-node=1 \ apptainer exec --rocm \ "${AI4S_SHARED_DIR}/images/pytorch_rocm7.2.2_ubuntu24.04_py3.12_pytorch_release_2.10.0.sif" \ - bash -c "cd ${OMNIHUB_TOOLS_DIR}/omnistat-src && \ + bash -c "cd ${PERF_TOOLS_DIR}/omnistat-src && \ cmake -S rocprofiler-sdk/ -B build-trace/ -DBUILD_KERNEL_TRACE_LIB=ON && \ cmake --build build-trace/ -j 8" ``` -Verify with `ls -la ${OMNIHUB_TOOLS_DIR}/omnistat-src/build-trace/libomnistat_trace.so` (~MB-class file, not the tiny 4 kB stub you get when CMake fails halfway). +Verify with `ls -la ${PERF_TOOLS_DIR}/omnistat-src/build-trace/libomnistat_trace.so` (~MB-class file, not the tiny 4 kB stub you get when CMake fails halfway). **SIF build prerequisites (one-time gotchas):** the `pytorch_rocm7.2.2_ubuntu24.04_py3.12_pytorch_release_2.10.0.sif` ships **without** `cmake` and **without** libcurl headers (only the `.so.4` runtime). Both omnistat artifacts need cmake; the kernel-trace library additionally needs ``. Working build recipe inside the SIF: @@ -735,12 +737,12 @@ When the cluster is up but you only need a one-shot interactive build inside the srun -p -A -N 1 --time=00:15:00 --gpus-per-node=1 --cpus-per-task=8 \ --job-name= \ apptainer exec --rocm \ - --bind ${OMNIHUB_TOOLS_DIR}:${OMNIHUB_TOOLS_DIR} \ + --bind ${PERF_TOOLS_DIR}:${PERF_TOOLS_DIR} \ "$HG_SIF" \ bash -c '' ``` -Why: `salloc` doesn't auto-bind `${OMNIHUB_TOOLS_DIR}/` into the apptainer namespace. Direct `srun … apptainer exec --bind …` forces the bind explicitly and writes the build artifact to its persistent NFS path in one shot. Verified during the kernel-trace library build (job 7030). +Why: `salloc` doesn't auto-bind `${PERF_TOOLS_DIR}/` into the apptainer namespace. Direct `srun … apptainer exec --bind …` forces the bind explicitly and writes the build artifact to its persistent NFS path in one shot. Verified during the kernel-trace library build (job 7030). ### 15. Dual-failure debug pattern — separate "tool wired correctly?" from "workload happy?" @@ -752,3 +754,25 @@ When a probe job fails, ask **two** questions before retrying with a bigger run: **Concrete example from this session (kernel-trace probe job 7033):** all 8 ranks crashed in `int(os.environ["HYDRAGNN_NUM_WORKERS"])` (`ValueError: invalid literal for int() with base 10: ''`). The kernel-trace library *also* printed `[hostname][pid][omnistat] Trace summary: 0/0 processed records (0/0 successful flushes)` per rank — meaning the rocprofiler-sdk tool loaded, registered its callback, and only saw zero dispatches because the workload crashed first. Fixing the env-passthrough bug (§13) and re-running gave 22.5k records/rank — the kernel-trace wiring was correct all along. **Action:** when you see "0 records / 0 events / 0 metrics" after a failed probe, *don't* re-build the instrumentation. Read the workload error, fix that, and re-run before changing anything tool-side. This saved one full rebuild cycle here. + +### 16. ORBIT-2 / Bayes-CAST: `launch_diffusion.sh`, 1-node FSDP layout, perf template parity + +- **`sbatch_train_perf_amd.sh`** defaults **`edm_8m_era5_1x8.yaml`** (Bayes-CAST EDM; rendered **`fsdp=8`/`simple_ddp=1`**) + ERA5 1.0° **`ORBIT2_DATA_ROOT`**; **`ORBIT2_ROOT`** prefers **`…/code/bayes-cast`** when present; detects **`launch/launch_diffusion.sh`**. **`sbatch_train_amd.sh`**: same launcher discovery + ORBIT2_ROOT preference; train still defaults PRISM **`interm_8m_lux.yaml`** unless overridden. `render_orbit2_config.py` adds **`ORBIT2_ERA5_SPATIAL_RES`** (default **111**) for the EDM template. Perf default `#SBATCH --nodes=1`; use `sbatch --nodes=N` for multi-node. +- **1 node × 8 GPUs:** unless **`ORBIT2_FSDP`** and **`ORBIT2_SIMPLE_DDP`** are both set, sbatch passes **`--fsdp 8 --simple-ddp 1`** (YAML: eight-way FSDP, one DDP group). Multi-node jobs omit extra flags so render keeps **`fsdp=nodes`**, **`simple_ddp=8`**. +- **`launch/train_edm.py` (Bayes-CAST EDM):** when **`$ORBIT2_ROOT/launch/train_edm.py`** exists, Studio **does not** run upstream **`launch/launch_diffusion.sh`**. That file is often an OLCF/Crusher job wrapper: it **ignores argv**, hardcodes a YAML, and **nests `srun`**, which breaks Studio’s outer **`srun … apptainer exec`** flow. Ranks **`cd /orbit2/launch`**, run **`orbit2_rank_hook_runner.py`** (no-op if no hook), then **`exec python3 train_edm.py /config/config.yaml`**. Minimal Bayes checkouts may have **no `examples/`** — never **`cd /orbit2/examples`** on that path. +- **Lux / public ORBIT-2:** if **`launch_diffusion.sh`** exists and **`train_edm.py`** does not, ranks **`exec bash …/launch_diffusion.sh /config/config.yaml`** after the hook runner (from **`examples/`** when present). Otherwise **`run_orbit2_train.py`**. Optional **`ORBIT2_LAUNCH_SCRIPT`** must be under **`ORBIT2_ROOT`**. +- **Perf + Bayes EDM + rank hooks:** when **`LAUNCH_EDM_DIRECT=1`**, the default pre-train hook is **empty** (Lux **`orbit2_profiler_hook.py`** imports `intermediate_downscaling`). Set **`export ORBIT2_RANK_PRE_TRAIN_HOOK=/abs/path/hook.py` before `sbatch`** so `sbatch_train_perf_amd.sh` bakes the path into the generated rank script; **`orbit2_rank_hook_runner.py`** runs it before **`train_edm.py`**. TraceLens still arms from **`ORBIT2_PROFILE_DIR`** / **`PROFILE_TARGET_EPOCH`** in the rank script. +- **Bayes `train_edm.py` + ERA5:** some upstream snapshots omit **`std_delta`** for **`data_key=="ERA5_1"`** (IMERG/HRRR only) → **`UnboundLocalError`** at first **`training_step`**. Patch **`training_step`** with an **`ERA5_1`** branch using per-output std from staged **`normalize_std.npz`** (see ORBIT-2 [HANDOFF.md](../../../earth_science/models/ORBIT-2/recipes/perf-analysis/HANDOFF.md) §landmines). +- **Agent debug logs on clusters:** NDJSON instrumentation often defaults to the **Cursor workspace** path (e.g. **`.cursor/debug-.log`**). Compute nodes usually **cannot write there**. Set **`DEBUG_AGENT_LOG`** to a **shared/bind-mounted writable file**, run the job, then **copy** that file into the workspace path (or paste tail into chat) so the next agent turn has runtime evidence. ORBIT-2 CK + **`bfloat16`** SIGSEGV investigation and checklist: same HANDOFF §“Debug session handoff”. +- **xFormers CK SIGSEGV vs Bayes-CAST:** [`upstream_pytorch_sdpa_benchmark.py`](../../../earth_science/models/ORBIT-2/examples/upstream_pytorch_sdpa_benchmark.py) supports **`--orbit-include-xformers-ck`**, which runs **`MemoryEfficientAttentionCkOp`** in a **fresh subprocess per case** so a child **SIGSEGV** is summarized without killing the parent (PyTorch SDPA micro sweep still prints first). Use to prove crashes are **generic MEA+CK on bf16 shapes**, not ORBIT-specific wiring. +- **VRAM-first saturation:** see ORBIT-2 [BASELINE_LOCKIN.md](../../../earth_science/models/ORBIT-2/recipes/perf-analysis/BASELINE_LOCKIN.md) — sweep **`ORBIT2_BATCH_SIZE`** on fixed ERA5 + 1×8 **`fsdp=8`/`simple_ddp=1`** before widening **`embed_dim`/`depth`**. Use **`examples/orbit2_estimate_batch_from_memory.py`** for a two-point *guess*, then **binary-search** short jobs; **`ORBIT2_DATA_TYPE`** defaults to **`bfloat16`** in `sbatch_train_*` (override **`float32`** if CK/SDPA misbehaves). If VRAM plateaus low, **stage more ERA5** per [STAGING_ERA5_FOR_HBM.md](../../../earth_science/models/ORBIT-2/recipes/perf-optimizer-loop/STAGING_ERA5_FOR_HBM.md). +- **Iterative sysopt loop:** [perf-optimizer-loop README](../../../earth_science/models/ORBIT-2/recipes/perf-optimizer-loop/README.md) — **`throughput_samples_per_s`** primary FOM via **`run_fom_extractor.py`** + **`manifest.global_batch_size`**; Omnistat default rocprofiler profile **`hbm_flops_bf16`** (`SQ_INSTS_VALU_MFMA_MOPS_BF16` + `FETCH_SIZE`); **`OMNISTAT_ROCPROF_PROFILE=hbm_flops_f64`** restores fp64 HydraGNN-style counters. +- **Perf `manifest.json`** records **`git_sha`**, branch, **`git_remote_origin`**, **`rendered_config`**, **`parallelism`**, **`global_batch_size`**, **`runtime_seconds`**, **`orbit2_batch_size`**, **`total_ranks`**, **`max_epochs`**, **`data_type`**, and **`workload`** for cross-day FOM comparisons. + +### 17. A `blocked` lever verdict is WORKLOAD-specific — re-probe before copying it to another model + +A `status: blocked` lever in one model's `lever_catalog.yaml` is evidence about **that workload on that day**, not a platform fact. Copying it to a sibling model without re-testing can permanently hide a working, high-value feature. + +**Case study (TunableOp, 2026-06-13):** HydraGNN MLIP marked `tunable_op_*` as `blocked` because all 16 GPUs hit `Memory access fault by GPU node-N` during hipBLASLt autotuning (perf-analysis SKILL #8/#15) — a real, validated failure for that workload (MACE double-backward shapes, full 16-GPU distributed tuning). The verdict was carried into ORBIT-2's catalog as a precaution. A cheap isolated probe (single GPU, ~2 min, `earth_science/models/ORBIT-2/perf-runs/tunableop-probe/`, jobs 10595/10596) **disproved it for ORBIT-2**: `PYTORCH_TUNABLEOP_TUNING=1` over 17 GEMMs — including the large-M bf16 `mm`/`F.linear`(TN+bias) class the sbatch comment claims faults — tuned cleanly with **0 memory faults** and wrote a valid `tunableop_results.csv`. All gfx950 Tensile logic libs are present in the container (`/opt/rocm/lib/{hipblaslt,rocblas}/library`, 439+156 gfx950 files) — **no missing prerequisites**. ORBIT-2's lever is now `experimental` (perf-analysis SKILL #28). + +**Rule for any model:** before trusting an inherited `blocked` verdict for a compute/BLAS-side lever (`tunable_op_*`, `torch_compile_*`, precision), run the smallest possible isolated probe on this exact stack first. Reserve `blocked` for failures *you* reproduced on *this* model. Probe template: 1-GPU `srun … apptainer exec … bash -c "source /opt/venv/bin/activate && python3 probe.py"` with the feature's env vars set. Note: `torch.cuda.tunable.write_file()` does not exist in PyTorch 2.10 — the cache auto-writes on process exit; use `set_filename()` + `get_results()`. diff --git a/.gitignore b/.gitignore index 49382e1..74c0470 100755 --- a/.gitignore +++ b/.gitignore @@ -74,3 +74,6 @@ nul # IDE .idea/ .vscode/ + +# Stray SLURM stdout from perf runs (real logs live under perf-runs//) +orbit2-train-*.out From 834291f098d9261371b3a68b94f0a65a1e719fec Mon Sep 17 00:00:00 2001 From: Ashwin Aji Date: Wed, 17 Jun 2026 07:46:53 +0000 Subject: [PATCH 29/31] ORBIT-2: add perf-optimizer-loop + GEMM-analysis recipes (Bayes-CAST EDM) Throughput-primary optimizer loop with a TraceLens+Omnistat analyst/verifier stack: loop driver, lever catalog, baseline lock-in and ERA5-staging docs, and the GEMM-time bottleneck analysis (the dominant cost is the im2col-lowered low-channel conv GEMM, not attention/MLP). Records the rejected levers (channels_last, TunableOp, MIOpen Winograd flags, torch.compile, conv-padding) so they are not re-attempted. bf16 unblocked by forcing SDPA EFFICIENT+MATH to avoid the ROCm Flash 65535 batch-grid cap. All paths parameterized via AI4S_SHARED_DIR / PERF_TOOLS_DIR; partition/account from cluster config. Co-Authored-By: Claude --- .../models/ORBIT-2/examples/README.md | 22 +- .../ORBIT-2/examples/compare_trace_kernels.py | 108 ++ .../ORBIT-2/examples/edm_8m_era5_1x8.yaml | 322 ++++++ .../ORBIT-2/examples/interm_8m_lux.yaml | 3 +- .../ORBIT-2/examples/interm_8m_lux_era5.yaml | 4 +- .../orbit2_estimate_batch_from_memory.py | 140 +++ .../ORBIT-2/examples/orbit2_profiler_hook.py | 0 .../examples/orbit2_rank_hook_runner.py | 45 + .../ORBIT-2/examples/parse_training_log.py | 201 +++- .../ORBIT-2/examples/render_orbit2_config.py | 34 +- .../examples/report_orbit2_gpu_baseline.py | 340 ++++++ .../examples/run_2node_scaleout_loop.sh | 245 +++++ .../ORBIT-2/examples/run_fom_extractor.py | 139 +++ .../ORBIT-2/examples/run_gemm_analysis.sh | 121 +++ .../ORBIT-2/examples/run_optimizer_loop.sh | 182 ++++ .../ORBIT-2/examples/run_scaling_study.sh | 5 +- .../ORBIT-2/examples/sbatch_infer_amd.sh | 7 + .../ORBIT-2/examples/sbatch_infer_docker.sh | 5 + .../ORBIT-2/examples/sbatch_train_amd.sh | 124 ++- .../ORBIT-2/examples/sbatch_train_perf_amd.sh | 315 +++++- .../ORBIT-2/examples/stage_era5_3x_symlink.sh | 62 ++ .../examples/submit_perf_baseline_era5_amd.sh | 29 + .../examples/sweep_orbit2_batch_bf16_amd.sh | 44 + .../upstream_pytorch_sdpa_benchmark.py | 968 ++++++++++++++++++ .../validate_orbit2_optimizer_loop_recipe.sh | 26 + earth_science/models/ORBIT-2/model.yaml | 28 +- .../models/ORBIT-2/recipes/data/README.md | 2 + .../recipes/perf-analysis/BASELINE_LOCKIN.md | 107 ++ .../ORBIT-2/recipes/perf-analysis/HANDOFF.md | 266 ++++- .../ORBIT-2/recipes/perf-analysis/README.md | 30 +- .../recipes/perf-analysis/agents/launcher.md | 12 +- .../perf-analysis/agents/omnistat_analyst.md | 8 +- .../agents/orchestrator_gemm_analysis.md | 116 +++ .../perf-analysis/agents/tracelens_analyst.md | 2 +- .../perf-analysis/omnistat.config.template | 20 +- .../perf-analysis/one-node-gpu-baseline.md | 100 ++ .../recipes/perf-optimizer-loop/README.md | 120 ++- .../STAGING_ERA5_FOR_HBM.md | 35 + .../agents/fom_extractor.md | 98 ++ .../agents/lever_picker.md | 29 + .../agents/orchestrator.md | 53 + .../agents/story_writer.md | 27 + .../perf-optimizer-loop/lever_catalog.yaml | 160 +++ .../perf-optimizer-loop/pitfall-diagnosis.md | 58 ++ .../models/ORBIT-2/recipes/train/README.md | 2 +- 45 files changed, 4658 insertions(+), 106 deletions(-) create mode 100755 earth_science/models/ORBIT-2/examples/compare_trace_kernels.py create mode 100644 earth_science/models/ORBIT-2/examples/edm_8m_era5_1x8.yaml create mode 100755 earth_science/models/ORBIT-2/examples/orbit2_estimate_batch_from_memory.py mode change 100644 => 100755 earth_science/models/ORBIT-2/examples/orbit2_profiler_hook.py create mode 100755 earth_science/models/ORBIT-2/examples/orbit2_rank_hook_runner.py create mode 100755 earth_science/models/ORBIT-2/examples/report_orbit2_gpu_baseline.py create mode 100755 earth_science/models/ORBIT-2/examples/run_2node_scaleout_loop.sh create mode 100755 earth_science/models/ORBIT-2/examples/run_gemm_analysis.sh create mode 100755 earth_science/models/ORBIT-2/examples/run_optimizer_loop.sh create mode 100755 earth_science/models/ORBIT-2/examples/stage_era5_3x_symlink.sh create mode 100755 earth_science/models/ORBIT-2/examples/submit_perf_baseline_era5_amd.sh create mode 100755 earth_science/models/ORBIT-2/examples/sweep_orbit2_batch_bf16_amd.sh create mode 100755 earth_science/models/ORBIT-2/examples/upstream_pytorch_sdpa_benchmark.py create mode 100755 earth_science/models/ORBIT-2/examples/validate_orbit2_optimizer_loop_recipe.sh create mode 100644 earth_science/models/ORBIT-2/recipes/perf-analysis/BASELINE_LOCKIN.md create mode 100644 earth_science/models/ORBIT-2/recipes/perf-analysis/agents/orchestrator_gemm_analysis.md create mode 100644 earth_science/models/ORBIT-2/recipes/perf-analysis/one-node-gpu-baseline.md create mode 100644 earth_science/models/ORBIT-2/recipes/perf-optimizer-loop/STAGING_ERA5_FOR_HBM.md create mode 100644 earth_science/models/ORBIT-2/recipes/perf-optimizer-loop/agents/fom_extractor.md create mode 100644 earth_science/models/ORBIT-2/recipes/perf-optimizer-loop/agents/lever_picker.md create mode 100644 earth_science/models/ORBIT-2/recipes/perf-optimizer-loop/agents/orchestrator.md create mode 100644 earth_science/models/ORBIT-2/recipes/perf-optimizer-loop/agents/story_writer.md create mode 100644 earth_science/models/ORBIT-2/recipes/perf-optimizer-loop/lever_catalog.yaml create mode 100644 earth_science/models/ORBIT-2/recipes/perf-optimizer-loop/pitfall-diagnosis.md diff --git a/earth_science/models/ORBIT-2/examples/README.md b/earth_science/models/ORBIT-2/examples/README.md index e603c27..a2e3055 100644 --- a/earth_science/models/ORBIT-2/examples/README.md +++ b/earth_science/models/ORBIT-2/examples/README.md @@ -15,13 +15,25 @@ with AMD/ROCm container launch, overlay builds, and synthetic smoke tests. | [`sbatch_infer_amd.sh`](sbatch_infer_amd.sh) | SLURM driver for inference on AMD Instinct (MI250X/MI300X/MI350X) | | [`build_overlay_amd.sh`](build_overlay_amd.sh) | One-time overlay build (pre-bakes pip deps, skips ~15 min per job) | | [`interm_8m_synthetic.yaml`](interm_8m_synthetic.yaml) | Template YAML for synthetic mode | -| [`sbatch_train_amd.sh`](sbatch_train_amd.sh) | Multi-node training (Apptainer + MPI) | -| [`sbatch_train_perf_amd.sh`](sbatch_train_perf_amd.sh) | 2-node instrumented perf-analysis run | +| [`sbatch_train_amd.sh`](sbatch_train_amd.sh) | Multi-node training (Apptainer + MPI); defaults PRISM 10.0_arcmin + `interm_8m_lux.yaml` | +| [`sbatch_train_perf_amd.sh`](sbatch_train_perf_amd.sh) | 1-node × 8-GPU perf (Omnistat + profiler); defaults **ERA5 1.0°** + `edm_8m_era5_1x8.yaml`; multi-node: `sbatch --nodes=N …` | +| [`submit_perf_baseline_era5_amd.sh`](submit_perf_baseline_era5_amd.sh) | Thin `sbatch` wrapper (sets ERA5 defaults, forwards extra `sbatch` flags) | +| [`sweep_orbit2_batch_bf16_amd.sh`](sweep_orbit2_batch_bf16_amd.sh) | Submit multiple `ORBIT2_BATCH_SIZE` probes (bf16 + SDPA) for HBM saturation sweeps | +| [`run_optimizer_loop.sh`](run_optimizer_loop.sh) | Driver for the iterative **perf-optimizer-loop** (Claude CLI optional) | +| [`validate_orbit2_optimizer_loop_recipe.sh`](validate_orbit2_optimizer_loop_recipe.sh) | Repo smoke: required perf-optimizer-loop files exist | | [`run_scaling_study.sh`](run_scaling_study.sh) | Submit matched 1/2/4/8-node scaling sweep | -| [`collate_scaling_study.py`](collate_scaling_study.py) | Parse scaling logs → CSV/MD table | -| [`parse_training_log.py`](parse_training_log.py) | Extract batch/epoch metrics from SLURM log | +| [`orbit2_estimate_batch_from_memory.py`](orbit2_estimate_batch_from_memory.py) | Heuristic max-batch from two `memory_reserved` calibrations or SLURM logs (see [BASELINE_LOCKIN.md](../recipes/perf-analysis/BASELINE_LOCKIN.md)) | +| [`parse_training_log.py`](parse_training_log.py) | Extract batch/epoch metrics from SLURM log (handles both `intermediate_downscaling.py` and Bayes-CAST `train_edm.py` `epoch:/batch_idx/tic4-tic1` formats) | +| [`run_fom_extractor.py`](run_fom_extractor.py) | Write `foms.json` (throughput + steady batch time + optional Omnistat PromQL via `ORBIT2_TSDB_URL`) | +| [`report_orbit2_gpu_baseline.py`](report_orbit2_gpu_baseline.py) | `baseline_report.md` / JSON for perf runs | | [`interm_8m_lux.yaml`](interm_8m_lux.yaml) | Lux config template (10.0_arcmin same-dir) | -| [`interm_8m_lux_era5.yaml`](interm_8m_lux_era5.yaml) | Lux config template (ERA5 1.0_deg same-dir sanity) | +| [`edm_8m_era5_1x8.yaml`](edm_8m_era5_1x8.yaml) | Bayes-CAST `edm_8m_era5.yaml` + **`fsdp=8`/`simple_ddp=1`** + `seq_par:1`; `ORBIT2_ROOT` → `code/bayes-cast` (auto if present) | +| [`orbit2_rank_hook_runner.py`](orbit2_rank_hook_runner.py) | Runs an optional per-rank pre-train hook (`ORBIT2_RANK_PRE_TRAIN_HOOK`) before the shell launcher | +| [`stage_era5_3x_symlink.sh`](stage_era5_3x_symlink.sh) | Symlink-replicate a staged ERA5 year (Nx) to break the small-corpus per-step batch cap for HBM saturation | +| [`run_gemm_analysis.sh`](run_gemm_analysis.sh) | Unattended GEMM-time bottleneck analysis (TraceLens + Omnistat analyst/verifier) at 1 and 2 nodes | +| [`compare_trace_kernels.py`](compare_trace_kernels.py) | Tool-independent kernel aggregation from a `*.pt.trace.json` (rank by raw device time, not category rollup) | +| [`run_2node_scaleout_loop.sh`](run_2node_scaleout_loop.sh) | Unattended 2-node scale-out lever loop; writes a 1-node-vs-2-node `REPORT.md` | +| [`upstream_pytorch_sdpa_benchmark.py`](upstream_pytorch_sdpa_benchmark.py) | Vendored [PyTorch `benchmarks/transformer/sdpa.py`](https://github.com/pytorch/pytorch/blob/main/benchmarks/transformer/sdpa.py) with **`--orbit-micro`** + backend sweep (`ck` → `SDPBackend.EFFICIENT_ATTENTION`, **not** xFormers `MemoryEfficientAttentionCkOp`) for comparing SDPA backends off the critical path | ## Quick start — ERA5 1.0_deg sanity (new dataset) diff --git a/earth_science/models/ORBIT-2/examples/compare_trace_kernels.py b/earth_science/models/ORBIT-2/examples/compare_trace_kernels.py new file mode 100755 index 0000000..f591e70 --- /dev/null +++ b/earth_science/models/ORBIT-2/examples/compare_trace_kernels.py @@ -0,0 +1,108 @@ +#!/usr/bin/env python3 +"""Compare GPU-kernel device-time breakdown between two Kineto traces. + +Tool-independent (no TraceLens): aggregates `traceEvents` with cat=="kernel" by name, summing `dur`, +and rolls them into the categories that matter for the ORBIT-2 conv bottleneck (GEMM/Cijk, im2col, +naive_conv, Winograd, attention, elementwise/other). Shares (% of total kernel time) are the +comparable quantity across runs because the captured step windows differ. + +Usage: + compare_trace_kernels.py LABEL_A=/path/a.pt.trace.json LABEL_B=/path/b.pt.trace.json [--md OUT.md] +""" +import argparse +import gzip +import json +import re +from collections import defaultdict + +CATEGORIES = [ + ("GEMM (Cijk/hipBLASLt)", re.compile(r"Cijk|gemm|Gemm|hipblaslt|rocblas", re.I)), + ("im2col lowering", re.compile(r"im2col|Im2Col|Im2d2Col", re.I)), + ("naive_conv (MIOpen direct)", re.compile(r"naive_conv", re.I)), + ("Winograd conv", re.compile(r"winograd", re.I)), + ("implicit-GEMM conv (MIOpen)", re.compile(r"miopen.*(igemm|implicit)|implicit.*gemm", re.I)), + ("attention/SDPA", re.compile(r"attention|flash|sdpa|fmha|attn|bwd_kernel_(dq|dk|dv|fuse)|fwd_kernel", re.I)), + ("layernorm/norm", re.compile(r"layer_norm|layernorm|GroupNorm|cuComputePartGrad|cuComputeGradInput", re.I)), + ("elementwise/copy/reduce", re.compile(r"elementwise|vectorized|copy|reduce|cat_|fill|index", re.I)), + ("collective/comm (RCCL)", re.compile(r"nccl|rccl|AllReduce|ReduceScatter|AllGather|all_gather|reduce_scatter", re.I)), +] + + +def load_events(path): + op = gzip.open if path.endswith(".gz") else open + with op(path, "rt") as fh: + data = json.load(fh) + return data.get("traceEvents", data if isinstance(data, list) else []) + + +def kernel_times(path): + """Return (by_name: {name: total_us}, total_us).""" + by_name = defaultdict(float) + for ev in load_events(path): + if ev.get("cat") == "kernel" and "dur" in ev: + by_name[ev.get("name", "?")] += float(ev["dur"]) + return by_name, sum(by_name.values()) + + +def categorize(by_name): + cat_tot = defaultdict(float) + for name, us in by_name.items(): + for label, rx in CATEGORIES: + if rx.search(name): + cat_tot[label] += us + break + else: + cat_tot["other"] += us + return cat_tot + + +def main(): + ap = argparse.ArgumentParser() + ap.add_argument("pairs", nargs="+", help="LABEL=/path/to/trace.pt.trace.json") + ap.add_argument("--md", default=None) + ap.add_argument("--topn", type=int, default=12) + args = ap.parse_args() + + runs = [] + for p in args.pairs: + label, _, path = p.partition("=") + by_name, tot = kernel_times(path) + runs.append((label, path, by_name, tot, categorize(by_name))) + + lines = ["# MIOpen ORNL-flags validation — GPU kernel device-time breakdown", ""] + for label, path, _bn, tot, _cat in runs: + lines.append(f"- **{label}**: total GPU-kernel time = {tot/1000:.1f} ms (`{path}`)") + lines.append("") + + # category share table + cats = [c for c, _ in CATEGORIES] + ["other"] + header = "| Category | " + " | ".join(f"{lbl} %" for lbl, *_ in runs) + " |" + sep = "|---|" + "---|" * len(runs) + lines += ["## Category share (% of total GPU-kernel time)", "", header, sep] + for c in cats: + cells = [] + for *_x, tot, cat in runs: + share = 100.0 * cat.get(c, 0.0) / tot if tot else 0.0 + cells.append(f"{share:.1f}") + if any(float(x) > 0.05 for x in cells): + lines.append(f"| {c} | " + " | ".join(cells) + " |") + lines.append("") + + # top kernels per run + for label, _path, by_name, tot, _cat in runs: + lines += [f"## Top {args.topn} kernels — {label}", ""] + lines += ["| kernel | ms | % |", "|---|---|---|"] + for name, us in sorted(by_name.items(), key=lambda kv: -kv[1])[: args.topn]: + lines.append(f"| `{name[:80]}` | {us/1000:.1f} | {100*us/tot if tot else 0:.1f} |") + lines.append("") + + out = "\n".join(lines) + print(out) + if args.md: + with open(args.md, "w") as fh: + fh.write(out) + print(f"\n[written] {args.md}") + + +if __name__ == "__main__": + main() diff --git a/earth_science/models/ORBIT-2/examples/edm_8m_era5_1x8.yaml b/earth_science/models/ORBIT-2/examples/edm_8m_era5_1x8.yaml new file mode 100644 index 0000000..e33271e --- /dev/null +++ b/earth_science/models/ORBIT-2/examples/edm_8m_era5_1x8.yaml @@ -0,0 +1,322 @@ +# Bayes-CAST `edm_8m_era5.yaml` — Studio copy for 1 node × 8 GPUs. +# Source of truth upstream: `bayes-cast/configs/edm_8m_era5.yaml` beside this repo (sync when Bayes-CAST updates). +# Set ORBIT2_ROOT to your Bayes-CAST checkout (must include `preset: edm` + EDM training stack). +# Parallelism is rendered to **fsdp=8, simple_ddp=1** for 8 ranks on one node. +# `__DATA_ROOT__` → ORBIT2_DATA_ROOT (e.g. .../era5/1.0_deg). `__ERA5_1_SPATIAL_RES__` → ORBIT2_ERA5_SPATIAL_RES (default 111). +# +#----------- TRAINER ------------------------------------------------------- +trainer: + max_epochs: __MAX_EPOCHS__ + checkpoint: + pretrain: + batch_size: __BATCH_SIZE__ + buffer_size: 200 + num_workers: __NUM_WORKERS__ + data_type: __DATA_TYPE__ + gpu_type: "amd" #amd or nvidia or leave it empty + train_loss: "mse" + +#---------- Parallelism and Efficient Computing ------------------------------- +parallelism: + fsdp: __FSDP__ + simple_ddp: __SIMPLE_DDP__ + tensor_par: 1 + seq_par: 1 + + +### Tiling +tiling: + do_tiling: True + div: 2 + overlap: 18 + +###Reslim Input Data Compression +compression: + compress_ratio: 1 + + +#------------ EDM Specific Hyperparameters ----------------------------------- +EDM_hyper: + sigma_data: 1 + weight_mode: + +# ---------------------------- MODEL ------------------------------------------- +model: + preset: edm + lr: 2e-4 + weight_decay: 1e-5 + beta_1: 0.9 + beta_2: 0.99 + warmup_epochs: 2 + warmup_start_lr: 1e-7 + eta_min: 1e-8 + patch_size: 6 + embed_dim: 256 + depth: 6 + decoder_depth: 4 + num_heads: 8 + mlp_ratio: 4 + drop_path: 0.1 + drop_rate: 0.1 + +# ---------------------------- DATA ------------------------------------------- +data: + src: { + #'PRISM': "prism", + 'ERA5_1': "era5", + #'HRRR': "hrrr", + } + pred_range: { + #'PRISM': 1, + 'ERA5_1': 6, + #'HRRR': 1, + } + history: { + #'PRISM': 2, + 'ERA5_1': 2, + #'HRRR': 2, + } + window: { + #'PRISM': 1, + 'ERA5_1': 6, + #'HRRR': 1, + } + data_dir: { + #'HRRR': "/lustre/orion/world-shared/csc662/jyc/data/hrrr_v1/0.05_deg/", + +## (25/03/09) new variable names +# 'ERA5_1': "/lustre/orion/lrn036/world-shared/data/superres/era5/5.625_deg/", +# 'ERA5_1': "/lustre/orion/lrn036/world-shared/data/superres/era5/1.0_deg/", +# 'PRISM': "/lustre/orion/lrn036/world-shared/data/superres/prism/10.0_arcmin", +# 'DAYMET_1': "/lustre/orion/lrn036/world-shared/data/superres/daymet/10.0_arcmin", +# 'DAYMET_2': "/lustre/orion/lrn036/world-shared/data/superres/daymet/2.0_arcmin", +# old datasets +# 'ERA5_1': "/lustre/orion/lrn036/world-shared/ERA5_npz/5.625_deg/", +# 'ERA5_1': "/lustre/orion/world-shared/lrn036/jyc/frontier/ClimaX-v2/data/ERA5-1hr-superres/5.625_deg/", +# 'ERA5_2': "/lustre/orion/world-shared/lrn036/jyc/frontier/ClimaX-v2/data/ERA5-1hr-superres/1.0_deg/", +# 'PRISM': "/lustre/orion/world-shared/lrn036/jyc/frontier/ClimaX-v2/data/prism-superres/10.0_arcmin/", + + +# 'ERA5_1': "/lustre/orion/lrn036/world-shared/data/superres/era5/1.40625_deg/", + 'ERA5_1': "__DATA_ROOT__", +# 'PRISM': "/lustre/orion/lrn036/world-shared/data/superres/prism/2.5_arcmin", +# 'DAYMET_1': "/lustre/orion/lrn036/world-shared/data/superres/daymet/2.5_arcmin", +# 'DAYMET_2': "/lustre/orion/lrn036/world-shared/data/superres/daymet/0.5_arcmin", +# old datasets +# 'ERA5_1': "/lustre/orion/lrn036/world-shared/ERA5_npz/1.40625_deg/", +# 'ERA5_1': "/lustre/orion/world-shared/lrn036/jyc/frontier/ClimaX-v2/data/ERA5-1hr-superres/1.40625_deg/", +# 'ERA5_2': "/lustre/orion/world-shared/lrn036/jyc/frontier/ClimaX-v2/data/ERA5-1hr-superres/0.25_deg/", +# 'PRISM': "/lustre/orion/world-shared/lrn036/jyc/frontier/ClimaX-v2/data/prism-superres/2.5_arcmin/", + } + + # input spatial resolution in km + spatial_resolution: { + 'ERA5_1': __ERA5_1_SPATIAL_RES__, + 'ERA5_2': 111, + 'PRISM': 18, + 'DAYMET_1': 18, + 'DAYMET_2': 4, + 'HRRR': 6, + } + + + start_year: { + 'ERA5_1': 1979, + } + + num_years: { + 'ERA5_1': 39, + } + + + + default_vars: [ + "land_sea_mask", + "orography", + "lattitude", +# "landcover", + "2m_temperature", +# "2m_temperature_max", +# "2m_temperature_min", + "temperature_200", +# "temperature_500", + "temperature_850", + "10m_u_component_of_wind", + "u_component_of_wind_200", +# "u_component_of_wind_500", + "u_component_of_wind_850", + "10m_v_component_of_wind", + "v_component_of_wind_200", +# "v_component_of_wind_500", + "v_component_of_wind_850", + "specific_humidity_200", +# "specific_humidity_500", + "specific_humidity_850", +# "total_precipitation_24hr", +# "volumetric_soil_water_layer_1", + "composite_reflectivity", + "mean_sea_level_pressure", + "surface_pressure", + + ] + + dict_in_variables: { + 'ERA5_1': [ + "land_sea_mask", + "orography", + "lattitude", +# "landcover", + "2m_temperature", +# "2m_temperature_max", +# "2m_temperature_min", +# "temperature_200", +# "temperature_500", +# "temperature_850", + "10m_u_component_of_wind", +# "u_component_of_wind_200", +# "u_component_of_wind_500", +# "u_component_of_wind_850", + "10m_v_component_of_wind", +# "v_component_of_wind_200", +# "v_component_of_wind_500", +# "v_component_of_wind_850", +# "specific_humidity_200", +# "specific_humidity_500", +# "specific_humidity_850", +# "total_precipitation_24hr", +# "volumetric_soil_water_layer_1", + ], + 'ERA5_2': [ + "land_sea_mask", + "orography", + "lattitude", + "landcover", + "2m_temperature", + "2m_temperature_max", + "2m_temperature_min", + "temperature_200", + "temperature_500", + "temperature_850", + "10m_u_component_of_wind", + "u_component_of_wind_200", + "u_component_of_wind_500", + "u_component_of_wind_850", + "10m_v_component_of_wind", + "v_component_of_wind_200", + "v_component_of_wind_500", + "v_component_of_wind_850", + "specific_humidity_200", + "specific_humidity_500", + "specific_humidity_850", + "total_precipitation_24hr", + "volumetric_soil_water_layer_1", + ], + 'PRISM': [ + "land_sea_mask", + "orography", + "lattitude", + "landcover", + "total_precipitation_24hr", + "2m_temperature_min", + "2m_temperature_max", + ], + 'DAYMET_1': [ + "land_sea_mask", + "orography", + "lattitude", + "landcover", + "total_precipitation_24hr", + "2m_temperature_min", + "2m_temperature_max", + ], + 'DAYMET_2': [ + "land_sea_mask", + "orography", + "lattitude", + "landcover", + "total_precipitation_24hr", + "2m_temperature_min", + "2m_temperature_max", + ], + 'HRRR': [ + "land_sea_mask", + "orography", + "lattitude", + "2m_temperature", + "composite_reflectivity", + "mean_sea_level_pressure", +# "surface_pressure", +# "geopotential_200", +# "geopotential_850", +# "geopotential_1000", +# "specific_humidity_200", +# "specific_humidity_850", +# "specific_humidity_1000", +# "temperature_200", +# "temperature_850", +# "temperature_1000", + "u_component_of_wind_200", +# "u_component_of_wind_850", +# "u_component_of_wind_1000", + "v_component_of_wind_200", +# "v_component_of_wind_850", +# "v_component_of_wind_1000", +# "vertical_velocity_200", +# "vertical_velocity_850", +# "vertical_velocity_1000", + ], + + + } + + dict_out_variables: { + 'ERA5_1':[ + "2m_temperature", +# "total_precipitation_24hr", +# "2m_temperature_min", + #"2m_temperature_max", + "10m_u_component_of_wind", + "10m_v_component_of_wind", + ], + 'ERA5_2':[ + "total_precipitation_24hr", + "2m_temperature_min", + "2m_temperature_max", +# "10m_u_component_of_wind", +# "10m_v_component_of_wind", + ], + 'PRISM': [ + "total_precipitation_24hr", + "2m_temperature_min", + "2m_temperature_max", + ], + 'DAYMET_1': [ + "total_precipitation_24hr", + "2m_temperature_min", + "2m_temperature_max", + ], + 'DAYMET_2': [ + "total_precipitation_24hr", + "2m_temperature_min", + "2m_temperature_max", + ], + 'HRRR': [ + "composite_reflectivity", +# "2m_temperature", + ], + } + + + var_weights: { + "2m_temperature": 1, + "10m_u_component_of_wind": 1, + "10m_v_component_of_wind": 1, + "total_precipitation_24hr": 1, + "composite_reflectivity": 1, + "2m_temperature_min": 1, + "2m_temperature_max": 1, + + } + + + diff --git a/earth_science/models/ORBIT-2/examples/interm_8m_lux.yaml b/earth_science/models/ORBIT-2/examples/interm_8m_lux.yaml index 2c35ba2..8c03b5b 100644 --- a/earth_science/models/ORBIT-2/examples/interm_8m_lux.yaml +++ b/earth_science/models/ORBIT-2/examples/interm_8m_lux.yaml @@ -10,7 +10,8 @@ # __SIMPLE_DDP__ — simple_ddp parallelism (typically 8) # __MAX_EPOCHS__ — ORBIT2_MAX_EPOCH override # __BATCH_SIZE__ — ORBIT2_BATCH_SIZE override -# __DATA_TYPE__ — ORBIT2_DATA_TYPE (float32 avoids xformers.ops CK path) +# __DATA_TYPE__ — ORBIT2_DATA_TYPE (same intent as upstream `configs/interm_8m.yaml` / +# Bayes-CAST `edm_8m_era5.yaml`: **bfloat16**; override float32 only for debug) #----------- TRAINER ------------- trainer: diff --git a/earth_science/models/ORBIT-2/examples/interm_8m_lux_era5.yaml b/earth_science/models/ORBIT-2/examples/interm_8m_lux_era5.yaml index dc22c14..6e230db 100644 --- a/earth_science/models/ORBIT-2/examples/interm_8m_lux_era5.yaml +++ b/earth_science/models/ORBIT-2/examples/interm_8m_lux_era5.yaml @@ -10,7 +10,9 @@ # __SIMPLE_DDP__ — simple_ddp parallelism (typically 8) # __MAX_EPOCHS__ — ORBIT2_MAX_EPOCH override # __BATCH_SIZE__ — ORBIT2_BATCH_SIZE override -# __DATA_TYPE__ — ORBIT2_DATA_TYPE (float32 avoids xformers.ops CK path) +# __DATA_TYPE__ — ORBIT2_DATA_TYPE (matches upstream training dtype: public ORBIT-2 +# `configs/interm_8m.yaml` uses bfloat16; Bayes-CAST-style trees often +# ship `edm_8m_era5.yaml` with the same. Studio sbatch defaults bfloat16.) #----------- TRAINER ------------- trainer: diff --git a/earth_science/models/ORBIT-2/examples/orbit2_estimate_batch_from_memory.py b/earth_science/models/ORBIT-2/examples/orbit2_estimate_batch_from_memory.py new file mode 100755 index 0000000..af95908 --- /dev/null +++ b/earth_science/models/ORBIT-2/examples/orbit2_estimate_batch_from_memory.py @@ -0,0 +1,140 @@ +#!/usr/bin/env python3 +"""Heuristic max batch estimate from ORBIT-2 SLURM logs or two manual calibration points. + +ORBIT rank-0 logs often include:: + torch.cuda.memory_reserved: 12.07GB + +This tool either: + * parses **peak** reserved memory from a log (optional filter on ``batch_idx``), or + * takes two (batch_size, reserved_GB) pairs and linearly extrapolates to ``--target-gb``. + +**Caveats:** reserved memory is **not** strictly linear in ``batch_size`` (attention/FFN +terms, fragmentation, checkpoints, FSDP sharding). Use the output as a **starting +guess**, then confirm with **binary-search** short jobs or a direct OOM probe. + +Examples:: + + # Two calibration jobs (float32), extrapolate to 240 GiB budget + python3 orbit2_estimate_batch_from_memory.py \\ + --batch-a 4 --gb-a 12.07 --batch-b 8 --gb-b 28.14 --target-gb 240 + + # Parse one log; pair with known batch_size from rendered YAML + python3 orbit2_estimate_batch_from_memory.py --log orbit2-train-9145.out --target-gb 240 + # (prints peak GB only unless you pass --known-batch) + + python3 orbit2_estimate_batch_from_memory.py \\ + --log orbit2-train-9145.out --known-batch 4 --log2 orbit2-train-9146.out --known-batch2 8 \\ + --target-gb 240 +""" + +from __future__ import annotations + +import argparse +import math +import re +import sys +from pathlib import Path + +_MEM_RE = re.compile(r"memory_reserved:\s*([0-9.]+)\s*GB", re.IGNORECASE) +_BATCH_IDX_RE = re.compile(r"batch_idx\s+(\d+)", re.IGNORECASE) + + +def _peak_gb_from_log( + path: Path, + *, + batch_idx_only: int | None, +) -> float: + mx = 0.0 + with path.open(encoding="utf-8", errors="replace") as f: + for line in f: + if batch_idx_only is not None: + mbi = _BATCH_IDX_RE.search(line) + if not mbi or int(mbi.group(1)) != batch_idx_only: + continue + m = _MEM_RE.search(line) + if m: + mx = max(mx, float(m.group(1))) + return mx + + +def _linear_extrapolate(b1: float, g1: float, b2: float, g2: float, target: float) -> float | None: + if b1 == b2: + return None + slope = (g2 - g1) / (b2 - b1) + if slope <= 0: + return None + # g = g1 + slope * (b - b1) => b = b1 + (target - g1) / slope + return b1 + (target - g1) / slope + + +def main() -> int: + p = argparse.ArgumentParser(description="Rough max batch from memory calibration.") + p.add_argument("--batch-a", type=float, help="First batch_size") + p.add_argument("--gb-a", type=float, help="First peak reserved GiB") + p.add_argument("--batch-b", type=float, help="Second batch_size") + p.add_argument("--gb-b", type=float, help="Second peak reserved GiB") + p.add_argument("--log", type=Path, help="SLURM stdout log (orbit2-train-*.out)") + p.add_argument("--log2", type=Path, help="Second log for two-point fit with --known-batch/--known-batch2") + p.add_argument("--known-batch", type=float, help="batch_size for --log") + p.add_argument("--known-batch2", type=float, help="batch_size for --log2") + p.add_argument( + "--batch-idx-filter", + type=int, + default=None, + help="Only consider lines containing this batch_idx (e.g. 19 for last batch in cap-20 epochs)", + ) + p.add_argument( + "--target-gb", + type=float, + default=240.0, + help="Desired peak reserved GiB (leave headroom below hardware HBM)", + ) + args = p.parse_args() + + if args.batch_a is not None and args.gb_a is not None and args.batch_b is not None and args.gb_b is not None: + b_hat = _linear_extrapolate(args.batch_a, args.gb_a, args.batch_b, args.gb_b, args.target_gb) + if b_hat is None: + print("error: invalid calibration (need distinct batch sizes and positive slope)", file=sys.stderr) + return 2 + print( + f"linear_fit: slope={(args.gb_b-args.gb_a)/(args.batch_b-args.batch_a):.3f} GiB per batch, " + f"estimated_batch≈{b_hat:.1f} at {args.target_gb} GiB (floor for integer batch: {math.floor(b_hat)})" + ) + print( + "note: linear extrapolation often **over**-estimates usable batch; binary-search short jobs to confirm.", + file=sys.stderr, + ) + return 0 + + if args.log and args.log2 and args.known_batch is not None and args.known_batch2 is not None: + g1 = _peak_gb_from_log(args.log, batch_idx_only=args.batch_idx_filter) + g2 = _peak_gb_from_log(args.log2, batch_idx_only=args.batch_idx_filter) + b_hat = _linear_extrapolate(args.known_batch, g1, args.known_batch2, g2, args.target_gb) + print(f"parsed peak GiB: log1={g1:.2f} (batch={args.known_batch}), log2={g2:.2f} (batch={args.known_batch2})") + if b_hat is None: + print("error: could not extrapolate (check logs / batch sizes)", file=sys.stderr) + return 2 + print( + f"linear_fit: slope={(g2-g1)/(args.known_batch2-args.known_batch):.3f} GiB per batch, " + f"estimated_batch≈{b_hat:.1f} at {args.target_gb} GiB (floor: {math.floor(b_hat)})" + ) + print("note: confirm with binary search or OOM probe — do not trust extrapolation alone.", file=sys.stderr) + return 0 + + if args.log: + g = _peak_gb_from_log(args.log, batch_idx_only=args.batch_idx_filter) + print(f"peak_reserved_gib={g:.2f} from {args.log}") + if args.known_batch is None: + print( + "hint: pass --known-batch and --log2/--known-batch2 for two-point estimate, " + "or use --batch-a/--gb-a/--batch-b/--gb-b.", + file=sys.stderr, + ) + return 0 + + p.print_help() + return 2 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/earth_science/models/ORBIT-2/examples/orbit2_profiler_hook.py b/earth_science/models/ORBIT-2/examples/orbit2_profiler_hook.py old mode 100644 new mode 100755 diff --git a/earth_science/models/ORBIT-2/examples/orbit2_rank_hook_runner.py b/earth_science/models/ORBIT-2/examples/orbit2_rank_hook_runner.py new file mode 100755 index 0000000..8d02b94 --- /dev/null +++ b/earth_science/models/ORBIT-2/examples/orbit2_rank_hook_runner.py @@ -0,0 +1,45 @@ +#!/usr/bin/env python3 +"""Run ORBIT2_RANK_PRE_TRAIN_HOOK only (e.g. profiler) before launch_diffusion.sh. + +Mirrors rank-gating in run_orbit2_train._run_hook. Expects ORBIT2_ROOT and cwd +compatible with importing intermediate_downscaling (chdir to .../examples). +""" + +from __future__ import annotations + +import os +import runpy +import sys +from pathlib import Path + + +def main() -> int: + hook = os.environ.get("ORBIT2_RANK_PRE_TRAIN_HOOK", "").strip() + if not hook: + return 0 + hook_path = Path(hook) + if not hook_path.is_file(): + raise FileNotFoundError(f"ORBIT2_RANK_PRE_TRAIN_HOOK not found: {hook}") + + rank0_only = os.environ.get("PROFILE_RANK0_ONLY", "1") == "1" + world_rank = int(os.environ.get("SLURM_PROCID", "0")) + if rank0_only and world_rank != 0: + return 0 + + root = os.environ.get("ORBIT2_ROOT", "/orbit2") + examples = Path(root) / "examples" + launch = Path(root) / "launch" + if examples.is_dir(): + os.chdir(examples) + elif launch.is_dir(): + os.chdir(launch) + for pref in (str(Path(root) / "src"), str(examples), str(launch), root): + if pref not in sys.path: + sys.path.insert(0, pref) + + runpy.run_path(str(hook_path), run_name="__orbit2_hook__") + return 0 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/earth_science/models/ORBIT-2/examples/parse_training_log.py b/earth_science/models/ORBIT-2/examples/parse_training_log.py index 1b5ef1c..7624573 100755 --- a/earth_science/models/ORBIT-2/examples/parse_training_log.py +++ b/earth_science/models/ORBIT-2/examples/parse_training_log.py @@ -1,10 +1,20 @@ #!/usr/bin/env python3 """Parse ORBIT-2 training metrics from SLURM stdout (orbit2-train-*.out). -Upstream intermediate_downscaling.py emits on rank 0: - Starting epoch {epoch} - Batch {idx}: {seconds} seconds (every 10 batches — wall time for that step) - Epoch {epoch} completed. Loss: {value} +Two upstream log formats are supported: + +1. **intermediate_downscaling.py** (studio path), rank 0: + Starting epoch {epoch} + Batch {idx}: {seconds} seconds (every 10 batches — wall time for that step) + Epoch {epoch} completed. Loss: {value} + +2. **Bayes-CAST EDM `launch/train_edm.py`** (perf-optimizer-loop default), rank 0: + epoch {epoch} + epoch: {epoch} batch_idx {idx} world_rank 0 loss tensor({value}, device='cuda:0') + my rank 0. tic4-tic1 in {seconds} seconds + i.e. there is **no** "Batch N: seconds" or "Epoch N completed" line — per-step + wall time comes from the ``tic4-tic1`` line (emitted right after each step's loss + line) and per-epoch loss is synthesized from the per-step losses. """ from __future__ import annotations @@ -16,14 +26,52 @@ import sys from pathlib import Path +# Format 1 (intermediate_downscaling) BATCH_RE = re.compile(r"^Batch (\d+): ([\d.]+) seconds") EPOCH_DONE_RE = re.compile(r"^Epoch (\d+) completed\. Loss: ([\d.eE+-]+)") EPOCH_START_RE = re.compile(r"^Starting epoch (\d+)") +# Format 2 (Bayes-CAST EDM train_edm.py) +EDM_STEP_RE = re.compile( + r"^epoch:\s+(\d+)\s+batch_idx\s+(\d+)\s+world_rank\s+\d+\s+loss\s+tensor\(\s*([\d.eE+-]+)" +) +EDM_TIME_RE = re.compile(r"tic4-tic1 in\s+([\d.]+)\s+seconds") +# Real per-step per-rank batch dim. The EDM logs print the input tensor shape right +# before each step's loss line: "y.shape torch.Size([, ...". A dataloader lever +# (num_workers, sharding) can make the *realized* batch differ from trainer.batch_size +# (e.g. per-worker partial trailing batches), so throughput MUST use this real dim, not +# the nominal global_batch_size. See loop overnight-3x-001 iter-1 + clean nw test 10504/10505. +EDM_YSHAPE_RE = re.compile(r"^y\.shape torch\.Size\(\[(\d+)") + + +def _synthesize_edm_epoch_rows(batch_rows: list[dict]) -> list[dict]: + """Build per-epoch loss records (mean of per-step losses) for loss-sanity on EDM logs.""" + by_epoch: dict[int, list[float]] = {} + for r in batch_rows: + if r.get("epoch") is None or r.get("loss") is None: + continue + by_epoch.setdefault(r["epoch"], []).append(r["loss"]) + rows: list[dict] = [] + for epoch in sorted(by_epoch): + losses = by_epoch[epoch] + rows.append( + { + "kind": "epoch", + "batch": None, + "batch_time_s": None, + "epoch": epoch, + "loss": sum(losses) / len(losses), + } + ) + return rows + def parse_slurm_log(log_path: Path) -> list[dict]: records: list[dict] = [] current_epoch: int | None = None + edm_batch_rows: list[dict] = [] + saw_upstream_epoch = False + pending_batch_samples: int | None = None # last y.shape[0] seen (real per-rank batch) with log_path.open(encoding="utf-8", errors="replace") as f: for line in f: @@ -34,6 +82,11 @@ def parse_slurm_log(log_path: Path) -> list[dict]: current_epoch = int(m.group(1)) continue + m = EDM_YSHAPE_RE.match(line) + if m: + pending_batch_samples = int(m.group(1)) + continue + m = BATCH_RE.match(line) if m: records.append( @@ -50,6 +103,7 @@ def parse_slurm_log(log_path: Path) -> list[dict]: m = EPOCH_DONE_RE.match(line) if m: epoch = int(m.group(1)) + saw_upstream_epoch = True records.append( { "kind": "epoch", @@ -60,6 +114,35 @@ def parse_slurm_log(log_path: Path) -> list[dict]: } ) current_epoch = epoch + 1 + continue + + # Format 2: Bayes-CAST EDM per-step loss line + m = EDM_STEP_RE.match(line) + if m: + row = { + "kind": "batch", + "batch": int(m.group(2)), + "batch_time_s": None, + "epoch": int(m.group(1)), + "loss": float(m.group(3)), + "batch_samples": pending_batch_samples, + } + records.append(row) + edm_batch_rows.append(row) + continue + + # Format 2: Bayes-CAST EDM per-step wall time (follows the step's loss line). + m = EDM_TIME_RE.search(line) + if m and edm_batch_rows: + for row in reversed(edm_batch_rows): + if row["batch_time_s"] is None: + row["batch_time_s"] = float(m.group(1)) + break + continue + + # EDM logs have no "Epoch N completed" line — synthesize per-epoch loss rows for sanity. + if edm_batch_rows and not saw_upstream_epoch: + records.extend(_synthesize_edm_epoch_rows(edm_batch_rows)) return records @@ -109,8 +192,21 @@ def aggregate_foms( *, warmup_batches_per_epoch: int = 1, steady_epoch_start: int = 2, + global_batch_size: int | None = None, + nominal_per_rank_batch: int | None = None, ) -> dict: - """Primary FOM: mean steady-state batch wall time (non-warmup epochs and batches).""" + """Primary FOM for sysopt: throughput_samples_per_s when global_batch_size is set. + + global_batch_size = trainer batch × fsdp × simple_ddp (samples per optimizer step). + + When the log carries the real per-step per-rank batch dim (``batch_samples``, from the + EDM ``y.shape`` line) and ``nominal_per_rank_batch`` is known, throughput is computed + from REAL samples processed (``Σ batch_samples × ranks / Σ step_time``) rather than the + nominal ``global_batch_size × n_steps``. This is mandatory because dataloader levers can + emit partial trailing batches (e.g. num_workers>1 with per-worker sharding): those steps + finish faster but process fewer samples, so the nominal formula reports a phantom speedup. + See clean num_workers test (jobs 10504 nw=1 vs 10505 nw=4): nominal said +80%, real ≈ 0%. + """ batch_rows = [r for r in records if r["kind"] == "batch" and r["batch_time_s"] is not None] epoch_rows = [r for r in records if r["kind"] == "epoch" and r["epoch"] is not None] @@ -139,6 +235,56 @@ def _mean(rows: list[dict]) -> float | None: ) steady_mean = _mean(steady_batches) + total_steady_wall_s = ( + sum(r["batch_time_s"] for r in steady_batches) if steady_batches else None + ) + + # ranks (= fsdp × simple_ddp) recovered from nominal global / per-rank batch. + ranks: int | None = None + if ( + global_batch_size is not None + and nominal_per_rank_batch + and nominal_per_rank_batch > 0 + ): + ranks = max(1, round(global_batch_size / nominal_per_rank_batch)) + + # Are real per-step batch dims available for ALL steady steps? Only then is the honest + # (real-samples) throughput computable; otherwise fall back to the nominal estimate. + steady_have_real = [ + r for r in steady_batches if r.get("batch_samples") is not None + ] + real_throughput_ok = ( + ranks is not None + and steady_batches + and len(steady_have_real) == len(steady_batches) + and total_steady_wall_s + and total_steady_wall_s > 0 + ) + + throughput: float | None = None + throughput_method: str | None = None + partial_step_fraction: float | None = None + if real_throughput_ok: + real_samples_steady = sum(r["batch_samples"] * ranks for r in steady_batches) + throughput = real_samples_steady / total_steady_wall_s + throughput_method = "real_per_step_batch" + n_partial = sum( + 1 for r in steady_batches if r["batch_samples"] < nominal_per_rank_batch + ) + partial_step_fraction = round(n_partial / len(steady_batches), 3) + elif ( + global_batch_size is not None + and global_batch_size > 0 + and total_steady_wall_s is not None + and total_steady_wall_s > 0 + ): + samples_steady = global_batch_size * len(steady_batches) + throughput = samples_steady / total_steady_wall_s + throughput_method = "nominal_global_batch" + + primary_fom = "throughput_samples_per_s" if throughput is not None else "batch_time_s" + primary_value = throughput if throughput is not None else steady_mean + sanity = check_loss_sanity(epoch_rows) steady_epoch_losses = [r for r in epoch_rows if r["epoch"] >= steady_epoch_start] @@ -148,13 +294,39 @@ def _mean(rows: list[dict]) -> float | None: else (epoch_rows[-1]["loss"] if epoch_rows else None) ) + # Batches/epoch is an integrity signal: throughput uses the *nominal* global_batch_size, + # so if a lever (e.g. num_workers) silently changes the dataloader's effective per-step + # batch, batches/epoch shifts and HBM drops — flagging that a cross-run throughput + # comparison is INVALID (the apparent speedup is a measurement artifact, not real work). + # See loop overnight-3x-001 iter-1 anomaly (num_workers=4 -> phantom +129%). + epoch_batch_counts: dict[int, int] = {} + for r in batch_rows: + ep = r.get("epoch") + if ep is not None: + epoch_batch_counts[ep] = epoch_batch_counts.get(ep, 0) + 1 + max_batches_per_epoch = max(epoch_batch_counts.values()) if epoch_batch_counts else None + + # Distinct realized per-rank batch sizes across steady steps (integrity surface). + steady_real_dims = sorted( + {r["batch_samples"] for r in steady_batches if r.get("batch_samples") is not None} + ) + return { - "primary_fom": "batch_time_s", + "primary_fom": primary_fom, + "primary_value": primary_value, + "throughput_samples_per_s": throughput, + "throughput_method": throughput_method, + "partial_step_fraction": partial_step_fraction, + "ranks_inferred": ranks, + "nominal_per_rank_batch": nominal_per_rank_batch, + "steady_realized_batch_dims": steady_real_dims or None, + "global_batch_size": global_batch_size, "batch_time_s": steady_mean, "steady_batch_time_s": steady_mean, "warmup_batch_time_s": warmup_mean, "steady_batch_count": len(steady_batches), "warmup_batch_count": len(warmup_batches), + "max_batches_per_epoch": max_batches_per_epoch, "steady_epoch_start": steady_epoch_start, "warmup_batches_per_epoch_excluded": warmup_batches_per_epoch, "final_loss": final_loss, @@ -180,6 +352,18 @@ def main() -> int: default=1, help="Exclude batch indices < N within each steady epoch (default: 1 = skip batch 0)", ) + parser.add_argument( + "--global-batch-size", + type=int, + default=None, + help="trainer.batch_size × fsdp × simple_ddp for throughput_samples_per_s", + ) + parser.add_argument( + "--nominal-per-rank-batch", + type=int, + default=None, + help="trainer.batch_size (per-rank); enables honest real-per-step throughput", + ) args = parser.parse_args() if not args.log.is_file(): @@ -195,7 +379,8 @@ def main() -> int: with args.output.open("w", newline="", encoding="utf-8") as f: w = csv.DictWriter( f, - fieldnames=["kind", "batch", "batch_time_s", "epoch", "loss"], + fieldnames=["kind", "batch", "batch_time_s", "epoch", "loss", "batch_samples"], + extrasaction="ignore", ) w.writeheader() w.writerows(records) @@ -204,6 +389,8 @@ def main() -> int: records, warmup_batches_per_epoch=args.warmup_batches_per_epoch, steady_epoch_start=args.steady_epoch_start, + global_batch_size=args.global_batch_size, + nominal_per_rank_batch=args.nominal_per_rank_batch, ) if args.json: args.json.parent.mkdir(parents=True, exist_ok=True) diff --git a/earth_science/models/ORBIT-2/examples/render_orbit2_config.py b/earth_science/models/ORBIT-2/examples/render_orbit2_config.py index 265170e..727620b 100755 --- a/earth_science/models/ORBIT-2/examples/render_orbit2_config.py +++ b/earth_science/models/ORBIT-2/examples/render_orbit2_config.py @@ -1,7 +1,15 @@ #!/usr/bin/env python3 -"""Render a per-job ORBIT-2 training YAML from interm_8m_lux.yaml. +"""Render a per-job ORBIT-2 training YAML from a Studio template. -Substitutes parallelism (fsdp × simple_ddp = nodes × 8), data root, and trainer caps. +Templates include Lux same-dir (**``interm_8m_lux*.yaml``**) and Bayes-CAST EDM +(**``edm_8m_era5_1x8.yaml``**), which also substitutes **``__ERA5_1_SPATIAL_RES__``** +from env **``ORBIT2_ERA5_SPATIAL_RES``** (default ``111`` for 1.0° ERA5 staging). + +Substitutes parallelism (fsdp × simple_ddp = nodes × gpus_per_node by default), data root, +and trainer caps. ``__NUM_WORKERS__`` comes from env **``ORBIT2_NUM_WORKERS``** (default ``2``). + +Env overrides (when ``--fsdp`` / ``--simple-ddp`` are omitted): ``ORBIT2_FSDP``, +``ORBIT2_SIMPLE_DDP``. ``ORBIT2_DATA_TYPE`` defaults to ``bfloat16`` when unset. Usage: python3 render_orbit2_config.py --nodes 2 --output /tmp/orbit2_config.yaml @@ -68,8 +76,18 @@ def main() -> int: print(f"error: data root not found: {data_path}", file=sys.stderr) return 2 - fsdp = args.fsdp if args.fsdp is not None else args.nodes - simple_ddp = args.simple_ddp if args.simple_ddp is not None else args.gpus_per_node + # CLI overrides env; env overrides geometric default (nodes × gpus_per_node). + fsdp = args.fsdp + if fsdp is None and os.environ.get("ORBIT2_FSDP"): + fsdp = int(os.environ["ORBIT2_FSDP"]) + if fsdp is None: + fsdp = args.nodes + + simple_ddp = args.simple_ddp + if simple_ddp is None and os.environ.get("ORBIT2_SIMPLE_DDP"): + simple_ddp = int(os.environ["ORBIT2_SIMPLE_DDP"]) + if simple_ddp is None: + simple_ddp = args.gpus_per_node total = fsdp * simple_ddp expected = args.nodes * args.gpus_per_node if total != expected: @@ -94,11 +112,15 @@ def main() -> int: if batch_size is None: batch_size = int(os.environ.get("ORBIT2_BATCH_SIZE", "8")) - data_type = os.environ.get("ORBIT2_DATA_TYPE", "float32") + data_type = os.environ.get("ORBIT2_DATA_TYPE", "bfloat16") if data_type not in ("float32", "bfloat16"): print("error: ORBIT2_DATA_TYPE must be float32 or bfloat16", file=sys.stderr) return 2 + # Bayes-CAST `edm_8m_era5_1x8.yaml`: spatial token count for ERA5_1 (111 ≈ 1.0° Lux staging) + era5_1_spatial = int(os.environ.get("ORBIT2_ERA5_SPATIAL_RES", "111")) + num_workers = int(os.environ.get("ORBIT2_NUM_WORKERS", "2")) + text = template.read_text(encoding="utf-8") replacements = { "__DATA_ROOT__": str(data_path), @@ -107,6 +129,8 @@ def main() -> int: "__MAX_EPOCHS__": str(max_epochs), "__BATCH_SIZE__": str(batch_size), "__DATA_TYPE__": data_type, + "__ERA5_1_SPATIAL_RES__": str(era5_1_spatial), + "__NUM_WORKERS__": str(num_workers), } for key, val in replacements.items(): text = text.replace(key, val) diff --git a/earth_science/models/ORBIT-2/examples/report_orbit2_gpu_baseline.py b/earth_science/models/ORBIT-2/examples/report_orbit2_gpu_baseline.py new file mode 100755 index 0000000..de02181 --- /dev/null +++ b/earth_science/models/ORBIT-2/examples/report_orbit2_gpu_baseline.py @@ -0,0 +1,340 @@ +#!/usr/bin/env python3 +"""Build baseline_report.{md,json} for one-node GPU saturation documentation. + +Reads manifest.json, rendered training YAML (path from manifest), SLURM log, +and foms.json (runs FOM extraction if missing). Extracts batch_size, data_type, +model preset and key widths from YAML via lightweight regex (no PyYAML). + +Parameter count: searches the log for common Lightning / summary patterns; +if absent, leaves placeholders for manual fill (see one-node-gpu-baseline.md). +""" + +from __future__ import annotations + +import argparse +import json +import os +import re +import subprocess +import sys +from pathlib import Path + +# ----------------------------------------------------------------------------- +# YAML snippets (avoid PyYAML dependency on login nodes) +# ----------------------------------------------------------------------------- + +_RE_BATCH = re.compile(r"^\s*batch_size:\s*(\d+)\s*$", re.MULTILINE) +_RE_DTYPE = re.compile(r"^\s*data_type:\s*(\S+)\s*$", re.MULTILINE) +_RE_PRESET = re.compile(r"^\s*preset:\s*(\S+)\s*$", re.MULTILINE) +_RE_EMBED = re.compile(r"^\s*embed_dim:\s*(\d+)\s*$", re.MULTILINE) +_RE_DEPTH = re.compile(r"^\s*depth:\s*(\d+)\s*$", re.MULTILINE) +_RE_DEC = re.compile(r"^\s*decoder_depth:\s*(\d+)\s*$", re.MULTILINE) +_RE_FSDP = re.compile(r"^\s*fsdp:\s*(\d+)\s*$", re.MULTILINE) +_RE_SDDP = re.compile(r"^\s*simple_ddp:\s*(\d+)\s*$", re.MULTILINE) + + +def _first_int(rx: re.Pattern[str], text: str) -> int | None: + m = rx.search(text) + return int(m.group(1)) if m else None + + +def _first_str(rx: re.Pattern[str], text: str) -> str | None: + m = rx.search(text) + return m.group(1).strip().strip('"').strip("'") if m else None + + +def _bytes_per_param(data_type: str | None) -> float: + if not data_type: + return 4.0 + dt = data_type.lower() + if dt == "bfloat16": + return 2.0 + if dt in ("float16", "half"): + return 2.0 + if dt == "float64": + return 8.0 + return 4.0 # float32 default + + +# Log lines: Lightning / PyTorch / ORBIT sometimes print these on rank 0 +_PARAM_RES = [ + re.compile(r"Total\s+model\s+parameters:\s*([\d.]+)\s*([KkMmGg])\b", re.IGNORECASE), + re.compile(r"(\d[\d,]*)\s+Total params", re.IGNORECASE), + re.compile(r"Total\s+params:\s*(\d[\d,\s]*)", re.IGNORECASE), + re.compile(r"trainable\s+params:\s*(\d[\d,\s]*)", re.IGNORECASE), + re.compile(r"Model\s+parameters?:\s*(\d[\d,\s]*)", re.IGNORECASE), + re.compile(r"num[_\s]?param(?:eter)?s?:\s*(\d[\d,\s]*)", re.IGNORECASE), +] + + +def _parse_int_human(s: str) -> int: + return int(s.replace(",", "").replace(" ", "").strip()) + + +def _extract_params_from_log(log_text: str) -> int | None: + for line in log_text.splitlines(): + for rx in _PARAM_RES: + m = rx.search(line) + if not m: + continue + try: + if m.lastindex and m.lastindex >= 2 and m.group(2): + val = float(m.group(1)) + suf = m.group(2).upper() + mult = {"K": 1e3, "M": 1e6, "G": 1e9}.get(suf, 1.0) + return int(val * mult) + return _parse_int_human(m.group(1)) + except (ValueError, IndexError): + continue + return None + + +def _infer_ai4s_shared_from_job_dir(job_dir: Path) -> Path | None: + """...//models/ORBIT-2/perf-runs/ -> .""" + try: + parts = job_dir.resolve().parts + if "models" in parts and "ORBIT-2" in parts: + i = parts.index("models") + if i >= 2: + return Path(parts[0]).joinpath(*parts[1:i]) + except (OSError, ValueError, IndexError): + pass + return None + + +def _resolve_slurm_log(job_dir: Path, manifest: dict) -> Path | None: + job_id = str(manifest.get("job_id") or job_dir.name) + candidates: list[Path] = [] + mlog = manifest.get("slurm_log") + if mlog: + candidates.append(Path(mlog)) + candidates.append(job_dir / f"orbit2-train-{job_id}.out") + shared = os.environ.get("AI4S_SHARED_DIR") + if not shared: + inferred = _infer_ai4s_shared_from_job_dir(job_dir) + if inferred is not None: + shared = str(inferred) + if shared: + candidates.append( + Path(shared) / "models" / "ORBIT-2" / "outputs" / "train" / "logs" / f"orbit2-train-{job_id}.out" + ) + candidates.append(Path.cwd() / f"orbit2-train-{job_id}.out") + for p in candidates: + if p.is_file(): + return p.resolve() + return None + + +def _resolve_rendered_yaml(job_dir: Path, manifest: dict) -> Path | None: + rc = manifest.get("rendered_config") + if rc: + p = Path(rc) + if p.is_file(): + return p.resolve() + job_id = str(manifest.get("job_id") or job_dir.name) + for name in (f"interm_8m_lux_{job_id}.yaml", f"interm_8m_lux_era5_{job_id}.yaml"): + p = job_dir / name + if p.is_file(): + return p.resolve() + for p in sorted(job_dir.glob("interm_*.yaml")): + if p.is_file() and "lux" in p.name: + return p.resolve() + return None + + +def _minimal_manifest(job_dir: Path) -> dict: + jid = job_dir.name + return { + "job_id": jid, + "job_dir": str(job_dir), + "slurm_log": str(job_dir / f"orbit2-train-{jid}.out"), + } + + +def _ensure_foms(job_dir: Path, log_path: Path) -> Path: + out = job_dir / "foms.json" + if out.is_file(): + return out + fom_script = Path(__file__).resolve().parent / "run_fom_extractor.py" + if not fom_script.is_file(): + return out + rc = subprocess.call( + [sys.executable, str(fom_script), "--job-dir", str(job_dir)], + stdout=subprocess.DEVNULL, + stderr=subprocess.DEVNULL, + ) + if rc != 0 and not out.is_file(): + sys.stderr.write(f"warning: run_fom_extractor.py exited {rc}; foms.json may be missing\n") + return out + + +def main() -> int: + ap = argparse.ArgumentParser(description="ORBIT-2 one-node GPU baseline report.") + ap.add_argument("--job-dir", type=Path, required=True, help="perf-runs//") + ap.add_argument("--num-params", type=int, default=None, help="Override parameter count if known") + args = ap.parse_args() + job_dir: Path = args.job_dir.resolve() + if not job_dir.is_dir(): + print(f"error: job dir not found: {job_dir}", file=sys.stderr) + return 2 + + manifest_path = job_dir / "manifest.json" + if manifest_path.is_file(): + manifest = json.loads(manifest_path.read_text(encoding="utf-8")) + else: + manifest = _minimal_manifest(job_dir) + sys.stderr.write(f"warning: no manifest.json; using job dir name as job_id={manifest['job_id']}\n") + + cfg_path = _resolve_rendered_yaml(job_dir, manifest) + if cfg_path is None: + print(f"error: could not find rendered training YAML under {job_dir}", file=sys.stderr) + return 2 + yaml_text = cfg_path.read_text(encoding="utf-8", errors="replace") + + batch_size = _first_int(_RE_BATCH, yaml_text) + data_type = _first_str(_RE_DTYPE, yaml_text) + preset = _first_str(_RE_PRESET, yaml_text) + embed_dim = _first_int(_RE_EMBED, yaml_text) + depth = _first_int(_RE_DEPTH, yaml_text) + dec_depth = _first_int(_RE_DEC, yaml_text) + fsdp = _first_int(_RE_FSDP, yaml_text) + simple_ddp = _first_int(_RE_SDDP, yaml_text) + + par = manifest.get("parallelism") or {} + fsdp = fsdp or par.get("fsdp") + simple_ddp = simple_ddp or par.get("simple_ddp") + + log_path = _resolve_slurm_log(job_dir, manifest) + if log_path is None: + print(f"error: SLURM log not found for job {manifest.get('job_id', job_dir.name)}", file=sys.stderr) + return 2 + log_text = log_path.read_text(encoding="utf-8", errors="replace") + + _ensure_foms(job_dir, log_path) + foms: dict = {} + foms_path = job_dir / "foms.json" + if foms_path.is_file(): + foms = json.loads(foms_path.read_text(encoding="utf-8")) + + n_params = args.num_params + if n_params is None: + n_params = _extract_params_from_log(log_text) + + bpp = _bytes_per_param(data_type) + param_bytes = int(n_params * bpp) if n_params is not None else None + + worlds = (fsdp or 0) * (simple_ddp or 0) + global_batch = batch_size * worlds if batch_size and worlds else None + + report: dict = { + "job_dir": str(job_dir), + "job_id": manifest.get("job_id"), + "orbit2_root": manifest.get("orbit2_root"), + "git_sha": manifest.get("git_sha"), + "git_branch": manifest.get("git_branch"), + "rendered_config": str(cfg_path), + "trainer": { + "batch_size_per_rank": batch_size, + "data_type": data_type, + "global_batch_approx": global_batch, + }, + "parallelism": {"fsdp": fsdp, "simple_ddp": simple_ddp, "world_size": worlds}, + "model_yaml": { + "preset": preset, + "embed_dim": embed_dim, + "depth": depth, + "decoder_depth": dec_depth, + }, + "model_size": { + "num_parameters": n_params, + "param_bytes_if_dense_weights": param_bytes, + "bytes_per_param_assumption": bpp, + "note": "Activations/optimizer/grads not included; use Omnistat VRAM + TraceLens for live footprint.", + }, + "foms": foms, + "artifacts": { + "slurm_log": str(log_path), + "traces_dir": manifest.get("trace_dir"), + "omnistat_db": manifest.get("omnistat_db"), + }, + } + + (job_dir / "baseline_report.json").write_text(json.dumps(report, indent=2) + "\n", encoding="utf-8") + + md_lines = [ + f"# ORBIT-2 GPU baseline — job `{manifest.get('job_id')}`", + "", + "## Config", + "", + f"- **Rendered YAML:** `{cfg_path}`", + f"- **ORBIT2_ROOT:** `{manifest.get('orbit2_root')}`", + f"- **Git:** `{manifest.get('git_sha', '?')}` (`{manifest.get('git_branch', '?')}`)", + "", + "## Parallelism / batching", + "", + f"| fsdp | simple_ddp | world | per-rank batch | global batch (approx.) |", + f"|------|------------|-------|----------------|-------------------------|", + f"| {fsdp or '—'} | {simple_ddp or '—'} | {worlds or '—'} | {batch_size or '—'} | {global_batch or '—'} |", + "", + "## Model (YAML fingerprint)", + "", + f"| preset | embed_dim | depth | decoder_depth |", + f"|--------|-----------|-------|---------------|", + f"| {preset or '—'} | {embed_dim or '—'} | {depth or '—'} | {dec_depth or '—'} |", + "", + "## Model size (parameters)", + "", + ] + if n_params is not None: + md_lines.extend( + [ + f"- **Parameter count (from log or --num-params):** `{n_params:,}`", + f"- **Dense weight bytes (params × {bpp} B, {data_type or 'float32'}):** ~{param_bytes:,} bytes (~{param_bytes / 1e9:.3f} GB)", + ] + ) + else: + md_lines.extend( + [ + "- **Parameter count:** *not found in log* — paste from a Lightning summary or run the container one-liner in " + "[`one-node-gpu-baseline.md`](../recipes/perf-analysis/one-node-gpu-baseline.md) §4, then re-run:", + "", + "```bash", + f"python3 examples/report_orbit2_gpu_baseline.py --job-dir {job_dir} --num-params ", + "```", + ] + ) + + md_lines.extend( + [ + "", + "## Figures of merit (ORBIT log)", + "", + ] + ) + if foms: + md_lines.append("```json") + md_lines.append(json.dumps({k: foms[k] for k in ("steady_batch_time_s", "loss_sanity_pass", "steady_batch_count") if k in foms}, indent=2)) + md_lines.append("```") + else: + md_lines.append("*Run `run_fom_extractor.py` — no foms.json yet.*") + + md_lines.extend( + [ + "", + "## Utilization evidence (fill after analysis)", + "", + "- **Omnistat:** attach `omnistat/report_summary.md` + key PromQL plots (VRAM, `FETCH_SIZE` / MFMA-derived FLOPs).", + "- **TraceLens:** attach `tracelens/report_summary.md` from rank-0 trace.", + "- **Saturation checklist:** higher `ORBIT2_BATCH_SIZE` until VRAM plateaus **and** `steady_batch_time_s` is acceptable vs batch sweep.", + "", + ] + ) + + (job_dir / "baseline_report.md").write_text("\n".join(md_lines) + "\n", encoding="utf-8") + print(f"Wrote {job_dir / 'baseline_report.md'}") + print(f"Wrote {job_dir / 'baseline_report.json'}") + return 0 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/earth_science/models/ORBIT-2/examples/run_2node_scaleout_loop.sh b/earth_science/models/ORBIT-2/examples/run_2node_scaleout_loop.sh new file mode 100755 index 0000000..0f68adc --- /dev/null +++ b/earth_science/models/ORBIT-2/examples/run_2node_scaleout_loop.sh @@ -0,0 +1,245 @@ +#!/usr/bin/env bash +# run_2node_scaleout_loop.sh — unattended ORBIT-2 2-node scale-out lever loop. +# +# Submits a 2-node bf16 baseline + 5 lever iterations (one SLURM job each, serial), +# extracts HONEST throughput (real per-step batch; see run_fom_extractor.py) after each, +# and writes a cross-node REPORT.md comparing 1-node vs 2-node. +# +# Designed to run unattended in tmux (survives disconnect): +# tmux new -s o2scale -d 'bash earth_science/models/ORBIT-2/examples/run_2node_scaleout_loop.sh' +# +# Prereqs (the loop sets sane defaults but you can override via env): +# - 5x ERA5 staging present (stage_era5_3x_symlink.sh DATA_ROOT 5 -> 100 train shards), +# so per_worker stays >=1 and full batches refill at 2 nodes (data_par_size=16). +# - .cluster-config.yaml present (multi-node RCCL: ib_hca, mgmt_iface, plugin, libionic). +# +# Levers tested at 2 nodes (compute-bound at 1 node; FSDP all-gather over IB is the new +# cost at N>1, so comm-tuning is where multi-node wins, if any, live): +# iter-0 baseline | iter-1 IB QPS=2 | iter-2 IB QPS=4 | iter-3 GPU_MAX_HW_QUEUES=4 +# iter-4 torch.compile | iter-5 NCCL_MIN_NCHANNELS=112 + +set -uo pipefail + +SCRIPT_DIR=$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd) +REPO_ROOT=$(cd "${SCRIPT_DIR}/../../../.." && pwd) +SBATCH="${SCRIPT_DIR}/sbatch_train_perf_amd.sh" + +# ---- environment (override-able) ------------------------------------------- +export AI4S_SHARED_DIR="${AI4S_SHARED_DIR:?AI4S_SHARED_DIR must be set}" +export PERF_TOOLS_DIR="${PERF_TOOLS_DIR:?PERF_TOOLS_DIR must be set (perf_tools.dir in .cluster-config.yaml)}" +export ORBIT2_ROOT="${ORBIT2_ROOT:-${AI4S_SHARED_DIR}/models/ORBIT-2/code/bayes-cast}" +export ORBIT2_DATA_ROOT="${ORBIT2_DATA_ROOT:-${AI4S_SHARED_DIR}/models/ORBIT-2/data/superres/era5/1.0_deg}" +export ORBIT2_CONFIG_TEMPLATE="${ORBIT2_CONFIG_TEMPLATE:-edm_8m_era5_1x8.yaml}" + +PARTITION="${SBATCH_PARTITION:?SBATCH_PARTITION must be set (slurm.partition in .cluster-config.yaml)}" +ACCOUNT="${SBATCH_ACCOUNT:?SBATCH_ACCOUNT must be set (slurm.account in .cluster-config.yaml)}" +NODES=2 +# Node placement (all optional, all portable defaults = no pinning): +# O2_NODELIST pin to specific nodes (--nodelist) +# O2_NODE_EXCLUDE exclude specific nodes (--exclude); set directly, or +# O2_EXCLUDE_PATTERN grep -E pattern matched against `sinfo` node names to auto-build the +# exclude list (use when a known-bad node range must be skipped — e.g. a +# sub-range with an inconsistent RoCE GID table that aborts multi-node RCCL). +O2_NODELIST="${O2_NODELIST:-}" +if [[ -z "${O2_NODE_EXCLUDE:-}" && -n "${O2_EXCLUDE_PATTERN:-}" ]]; then + O2_NODE_EXCLUDE=$(sinfo -p "$PARTITION" -N -h -o "%n" 2>/dev/null | sort -u | grep -E "$O2_EXCLUDE_PATTERN" | paste -sd, || true) +fi +O2_NODE_EXCLUDE="${O2_NODE_EXCLUDE:-}" +BATCH=4096 +MAXEPOCH=6 +MAXBATCH=20 +POLL_S="${POLL_S:-30}" +TIME_LIMIT="${TIME_LIMIT:-02:30:00}" + +LOOP_ID="${LOOP_ID:-loop-2node-scaleout-001}" +LOOP_DIR="${AI4S_SHARED_DIR}/models/ORBIT-2/perf-runs/${LOOP_ID}" +mkdir -p "$LOOP_DIR" +RESULTS_CSV="${LOOP_DIR}/results.csv" +STATUS="${LOOP_DIR}/STATUS.txt" +REPORT="${LOOP_DIR}/REPORT.md" +EXTRACT="${SCRIPT_DIR}/run_fom_extractor.py" + +# 1-node honest reference (clean num_workers test, batch 4096, 100 shards). +ONE_NODE_NW1_JOB="${ONE_NODE_NW1_JOB:-10504}" +ONE_NODE_NW4_JOB="${ONE_NODE_NW4_JOB:-10505}" + +log() { echo "[$(date '+%F %T')] $*" | tee -a "$STATUS"; } + +# ---- lever definitions (name | extra KEY=VALUE env, space-separated) ------- +LEVER_NAME=( "baseline" "ib_qps_2" "ib_qps_4" "gpu_hw_queues_4" "torch_compile" "nccl_min_nchannels_112" ) +LEVER_ENV=( "" "NCCL_IB_QPS_PER_CONNECTION=2" "NCCL_IB_QPS_PER_CONNECTION=4" "GPU_MAX_HW_QUEUES=4" "ORBIT2_TORCH_COMPILE=1 ORBIT2_COMPILE_MODE=default" "NCCL_MIN_NCHANNELS=112" ) + +if [[ ! -f "$RESULTS_CSV" ]]; then + echo "iter,lever,job_id,nodes,state,throughput_sps,method,partial_frac,realized_dims,hbm_pct,steady_batch_time_s,max_batches_per_epoch,verdict" > "$RESULTS_CSV" +fi + +log "=== ${LOOP_ID} START (2-node scale-out, ${#LEVER_NAME[@]} configs incl. baseline) ===" +log "loop dir: $LOOP_DIR" + +# Verify 5x staging (>=96 shards). Stage if missing. +N_SHARDS=$(find "${ORBIT2_DATA_ROOT}/train" -maxdepth 1 -name '*_*.npz' ! -name 'climatology.npz' 2>/dev/null | wc -l) +if [[ "$N_SHARDS" -lt 96 ]]; then + log "staging 5x ERA5 (have $N_SHARDS shards, need >=96)" + bash "${SCRIPT_DIR}/stage_era5_3x_symlink.sh" "$ORBIT2_DATA_ROOT" 5 2>&1 | tee -a "$STATUS" + N_SHARDS=$(find "${ORBIT2_DATA_ROOT}/train" -maxdepth 1 -name '*_*.npz' ! -name 'climatology.npz' | wc -l) +fi +log "train shards: $N_SHARDS (per_worker @2node nw=1 = $((N_SHARDS / 16)))" + +GLOBAL_BATCH=$((BATCH * NODES * 8)) # fsdp=nodes, simple_ddp=8 -> ranks = nodes*8 +BASELINE_TP="" + +read_fom() { # $1=job_dir -> sets FOM_* globals + python3 - "$1" <<'PY' +import json,sys +from pathlib import Path +p=Path(sys.argv[1])/"foms.json" +d=json.loads(p.read_text()) if p.is_file() else {} +def g(k,default=""): + v=d.get(k); return default if v is None else v +dims=g("steady_realized_batch_dims","") +if isinstance(dims,list): dims="/".join(str(x) for x in dims) +# single-quote every value so eval is safe regardless of content (dims, etc.) +print("\n".join([ + f"FOM_TP='{g('throughput_samples_per_s','')}'", + f"FOM_METHOD='{g('throughput_method','')}'", + f"FOM_PARTIAL='{g('partial_step_fraction','')}'", + f"FOM_DIMS='{dims}'", + f"FOM_HBM='{g('hbm_reserved_pct_288','')}'", + f"FOM_BT='{g('steady_batch_time_s','')}'", + f"FOM_MAXB='{g('max_batches_per_epoch','')}'", +])) +PY +} + +# Resume: skip iters already recorded in results.csv; reseed baseline from iter-0 row. +declare -A DONE_ITER=() +while IFS=, read -r ri rl rj rn rstate rtp rest; do + [[ "$ri" == "iter" || -z "$ri" ]] && continue + DONE_ITER["$ri"]=1 + if [[ "$ri" == "0" && -n "$rtp" ]]; then + BASELINE_TP="$rtp" + bj=$(ls -d "${AI4S_SHARED_DIR}/models/ORBIT-2/perf-runs/${rj}" 2>/dev/null) + if [[ -n "$bj" ]]; then eval "$(read_fom "$bj")"; BASELINE_DIMS="$FOM_DIMS"; fi + fi +done < "$RESULTS_CSV" +[[ -n "${BASELINE_TP:-}" ]] && log "resume: baseline throughput=${BASELINE_TP} dims=[${BASELINE_DIMS:-}] (iter-0 already done)" + +for i in "${!LEVER_NAME[@]}"; do + name="${LEVER_NAME[$i]}" + extra="${LEVER_ENV[$i]}" + if [[ -n "${DONE_ITER[$i]:-}" ]]; then + log "iter-$i ($name) already in results.csv — skipping" + continue + fi + log "--- iter-$i lever='$name' extra='${extra:-}' ---" + + # Build --export list: base config + lever extras. + EXP="ALL" + EXP+=",ORBIT2_BATCH_SIZE=${BATCH},ORBIT2_NUM_WORKERS=1,ORBIT2_DATA_TYPE=bfloat16,ORBIT2_FUSED_ATTN=DEFAULT" + EXP+=",ORBIT2_CONFIG_TEMPLATE=${ORBIT2_CONFIG_TEMPLATE},ORBIT2_MAX_EPOCH=${MAXEPOCH},ORBIT2_MAX_BATCHES=${MAXBATCH}" + EXP+=",ORBIT2_DATA_ROOT=${ORBIT2_DATA_ROOT},ORBIT2_ROOT=${ORBIT2_ROOT},AI4S_SHARED_DIR=${AI4S_SHARED_DIR},PERF_TOOLS_DIR=${PERF_TOOLS_DIR}" + EXP+=",TORCH_NCCL_HIGH_PRIORITY=1,GPU_MAX_HW_QUEUES=2" + for kv in $extra; do EXP+=",${kv}"; done + + PLACE=() + if [[ -n "$O2_NODELIST" ]]; then PLACE=(--nodelist="$O2_NODELIST") + elif [[ -n "$O2_NODE_EXCLUDE" ]]; then PLACE=(--exclude="$O2_NODE_EXCLUDE"); fi + JID=$( cd "$REPO_ROOT" && sbatch --parsable \ + --partition="$PARTITION" --account="$ACCOUNT" \ + --nodes="$NODES" --time="$TIME_LIMIT" \ + --job-name="o2-2n-${name}" "${PLACE[@]}" \ + --export="$EXP" "$SBATCH" 2>>"$STATUS" ) + if [[ -z "$JID" ]]; then + log "iter-$i SUBMIT FAILED" + echo "$i,$name,SUBMIT_FAIL,$NODES,submit_fail,,,,,,,,error" >> "$RESULTS_CSV" + continue + fi + log "iter-$i submitted job $JID; polling every ${POLL_S}s" + + # Poll until job leaves the queue (initial sleep avoids the submit->squeue race). + sleep 15 + while squeue -h -j "$JID" 2>/dev/null | grep -q .; do sleep "$POLL_S"; done + sleep 20 # let perf-sbatch finish omnistat cleanup + write manifest/log + + JOB_DIR="${AI4S_SHARED_DIR}/models/ORBIT-2/perf-runs/${JID}" + STATE=$(python3 -c "import json,sys;print(json.load(open('${JOB_DIR}/manifest.json')).get('state','unknown'))" 2>/dev/null || echo "no_manifest") + log "iter-$i job $JID finished (state=$STATE)" + + python3 "$EXTRACT" --job-dir "$JOB_DIR" >/dev/null 2>>"$STATUS" || log "iter-$i FOM extract warning" + eval "$(read_fom "$JOB_DIR")" + + # Verdict (honest throughput + integrity). + verdict="n/a" + if [[ "$i" -eq 0 ]]; then + BASELINE_TP="$FOM_TP"; BASELINE_DIMS="$FOM_DIMS"; verdict="baseline" + elif [[ "$STATE" != "complete" ]]; then + verdict="failed" + elif [[ -n "$FOM_TP" && -n "$BASELINE_TP" ]]; then + verdict=$(python3 -c " +b=float('$BASELINE_TP'); t=float('$FOM_TP') +dimic = ('$FOM_DIMS'=='$BASELINE_DIMS') +d=(t-b)/b*100 +if not dimic: print(f'REJECT-artifact({d:+.1f}%,dims_differ)') +elif d>2: print(f'ACCEPT(+{d:.1f}%)') +elif d<-2: print(f'REGRESS({d:.1f}%)') +else: print(f'neutral({d:+.1f}%)') +" 2>/dev/null || echo "parse_err") + fi + log "iter-$i throughput=${FOM_TP} s/s method=${FOM_METHOD} partial=${FOM_PARTIAL} dims=[${FOM_DIMS}] hbm=${FOM_HBM}% verdict=$verdict" + echo "$i,$name,$JID,$NODES,$STATE,$FOM_TP,$FOM_METHOD,$FOM_PARTIAL,\"$FOM_DIMS\",$FOM_HBM,$FOM_BT,$FOM_MAXB,$verdict" >> "$RESULTS_CSV" +done + +# ---- report ---------------------------------------------------------------- +log "=== generating REPORT.md ===" +python3 - "$RESULTS_CSV" "$REPORT" "$ONE_NODE_NW1_JOB" "$ONE_NODE_NW4_JOB" "$AI4S_SHARED_DIR" "$BATCH" <<'PY' +import csv,json,sys +from pathlib import Path +csv_path,report,j1,j4,shared,batch=sys.argv[1:7] +batch=int(batch) +def fom(job): + p=Path(shared)/"models/ORBIT-2/perf-runs"/str(job)/"foms.json" + return json.loads(p.read_text()) if p.is_file() else {} +one_nw1=fom(j1); one_nw4=fom(j4) +rows=list(csv.DictReader(open(csv_path))) +base=next((r for r in rows if r["iter"]=="0"), None) +base_tp=float(base["throughput_sps"]) if base and base["throughput_sps"] else None +one_tp=one_nw1.get("throughput_samples_per_s") + +L=[] +L.append("# ORBIT-2 Scale-out Report — 1-node vs 2-node (Bayes-CAST EDM, MI355X, bf16)\n") +L.append("All throughput is **honest** (real per-step batch × ranks ÷ real step time), so partial\n") +L.append("trailing batches do not inflate it. batch=%d/rank, EFFICIENT SDPA, 100 ERA5 shards (5× staged).\n" % batch) + +L.append("\n## 1-node reference (clean num_workers test)\n") +L.append("| config | job | throughput (s/s) | method | realized batch dims | partial frac | HBM%% |") +L.append("|---|---|---:|---|---|---:|---:|") +for lbl,f,j in [("1-node nw=1",one_nw1,j1),("1-node nw=4",one_nw4,j4)]: + L.append(f"| {lbl} | {j} | {f.get('throughput_samples_per_s','-'):.0f} | {f.get('throughput_method','-')} | {f.get('steady_realized_batch_dims','-')} | {f.get('partial_step_fraction','-')} | {f.get('hbm_reserved_pct_288','-')} |" if f else f"| {lbl} | {j} | - | - | - | - | - |") + +L.append("\n## 2-node lever loop\n") +L.append("| iter | lever | job | state | throughput (s/s) | Δ vs 2N base | scaling vs 1N | partial | dims | HBM%% | verdict |") +L.append("|---:|---|---|---|---:|---:|---:|---:|---|---:|---|") +for r in rows: + tp=r["throughput_sps"] + tpf=float(tp) if tp else None + dvb=f"{(tpf-base_tp)/base_tp*100:+.1f}%" if (tpf and base_tp) else "-" + scal=f"{tpf/one_tp:.2f}×" if (tpf and one_tp) else "-" + tps=f"{tpf:.0f}" if tpf else "-" + L.append(f"| {r['iter']} | {r['lever']} | {r['job_id']} | {r['state']} | {tps} | {dvb} | {scal} | {r['partial_frac']} | {r['realized_dims']} | {r['hbm_pct']} | {r['verdict']} |") + +# headline +if base_tp and one_tp: + eff=base_tp/one_tp/2*100 + L.append(f"\n## Headline\n") + L.append(f"- **2-node baseline: {base_tp:.0f} s/s vs 1-node: {one_tp:.0f} s/s → {base_tp/one_tp:.2f}× ({eff:.0f}% strong-scaling efficiency).**") + best=max((r for r in rows if r['throughput_sps']), key=lambda r: float(r['throughput_sps']), default=None) + if best: + L.append(f"- Best 2-node config: **{best['lever']}** ({float(best['throughput_sps']):.0f} s/s, {best['verdict']}).") + L.append(f"- Accepted levers (>+2%% over 2N baseline, integrity-matched): " + + (", ".join(r['lever'] for r in rows if r['verdict'].startswith('ACCEPT')) or "none — workload stays compute-bound at 2 nodes.")) +Path(report).write_text("\n".join(L)+"\n") +print(f"wrote {report}") +PY + +log "=== ${LOOP_ID} COMPLETE — see $REPORT ===" +cat "$REPORT" 2>/dev/null | tee -a "$STATUS" diff --git a/earth_science/models/ORBIT-2/examples/run_fom_extractor.py b/earth_science/models/ORBIT-2/examples/run_fom_extractor.py index 0cc80c7..e986e56 100755 --- a/earth_science/models/ORBIT-2/examples/run_fom_extractor.py +++ b/earth_science/models/ORBIT-2/examples/run_fom_extractor.py @@ -5,28 +5,128 @@ import argparse import json +import math import os +import re import sys +import urllib.error +import urllib.parse +import urllib.request from pathlib import Path sys.path.insert(0, str(Path(__file__).resolve().parent)) from parse_training_log import aggregate_foms, parse_slurm_log # noqa: E402 +def _infer_ai4s_shared_from_job_dir(job_dir: Path) -> Path | None: + try: + parts = job_dir.resolve().parts + if "models" in parts and "ORBIT-2" in parts: + i = parts.index("models") + if i >= 2: + return Path(parts[0]).joinpath(*parts[1:i]) + except (OSError, ValueError, IndexError): + pass + return None + + +def _read_manifest(job_dir: Path) -> dict: + p = job_dir / "manifest.json" + if not p.is_file(): + return {} + try: + return json.loads(p.read_text(encoding="utf-8")) + except (OSError, json.JSONDecodeError): + return {} + + +def _promql_scalar(tsdb_url: str, promql: str, t_unix: int) -> float | None: + """Instant query at time=t_unix (VictoriaMetrics / Prometheus API).""" + q = urllib.parse.urlencode({"query": promql, "time": str(t_unix)}) + url = f"{tsdb_url.rstrip('/')}/api/v1/query?{q}" + try: + with urllib.request.urlopen(url, timeout=30) as resp: + body = json.loads(resp.read().decode("utf-8")) + except (urllib.error.URLError, TimeoutError, json.JSONDecodeError, OSError): + return None + res = body.get("data", {}).get("result") or [] + if not res: + return None + try: + return float(res[0]["value"][1]) + except (KeyError, IndexError, ValueError, TypeError): + return None + + +def _enrich_omnistat(foms: dict, manifest: dict, job_dir: Path) -> None: + """Optional PromQL enrichment when ORBIT2_TSDB_URL is set (login-node VM).""" + tsdb = os.environ.get("ORBIT2_TSDB_URL", "").strip() + job_id = manifest.get("job_id") or job_dir.name + runtime = manifest.get("runtime_seconds") + if not tsdb or not job_id or not runtime or runtime <= 0: + foms["mfma_bf16_tflops_per_card_avg"] = None + foms["hbm_read_GBps_per_card_avg"] = None + foms["xgmi_GBps_avg"] = None + foms["omnistat_enrichment_note"] = ( + "Set ORBIT2_TSDB_URL to a running VictoriaMetrics (e.g. from omnistat_analyst) " + "and ensure runtime_seconds + job_id are in manifest.json." + ) + return + + t_end = int(Path(job_dir / "manifest.json").stat().st_mtime) if job_dir.exists() else 0 + if t_end <= 0: + t_end = int(__import__("time").time()) + mfma_q = ( + f"avg(rate(SQ_INSTS_VALU_MFMA_MOPS_BF16[30s]) * on(instance) group_left() " + f'max by(instance)(rmsjob_info{{jobid="{job_id}"}})) * 512 / 1e12' + ) + hbm_q = ( + f"avg(rate(FETCH_SIZE[30s]) * on(instance) group_left() " + f'max by(instance)(rmsjob_info{{jobid="{job_id}"}})) * 1024 / 1e9' + ) + # XGMI naming varies by omnistat version; try a common byte counter. + xgmi_q = ( + f"avg(rate(rocm_xgmi_link0_tx_bytes[30s]) * on(instance) group_left() " + f'max by(instance)(rmsjob_info{{jobid="{job_id}"}})) / 1e9' + ) + + t_mid = max(1, t_end - runtime // 2) + foms["mfma_bf16_tflops_per_card_avg"] = _promql_scalar(tsdb, mfma_q, t_mid) + foms["hbm_read_GBps_per_card_avg"] = _promql_scalar(tsdb, hbm_q, t_mid) + foms["xgmi_GBps_avg"] = _promql_scalar(tsdb, xgmi_q, t_mid) + foms["omnistat_enrichment_note"] = ( + f"PromQL at time={t_mid} via ORBIT2_TSDB_URL; metric names may differ by omnistat build — nulls mean rename/query fix." + ) + + def main() -> int: parser = argparse.ArgumentParser(description="ORBIT-2 perf-run FOM extractor.") parser.add_argument("--job-dir", type=Path, required=True, help="perf-runs//") parser.add_argument("--steady-epoch-start", type=int, default=2) parser.add_argument("--warmup-batches-per-epoch", type=int, default=1) + parser.add_argument( + "--global-batch-size", + type=int, + default=None, + help="Override manifest global_batch_size (trainer batch × fsdp × simple_ddp)", + ) args = parser.parse_args() job_dir = args.job_dir job_id = job_dir.name + manifest = _read_manifest(job_dir) + if manifest.get("job_id"): + job_id = str(manifest["job_id"]) + log_candidates = [ job_dir / f"orbit2-train-{job_id}.out", Path.cwd() / f"orbit2-train-{job_id}.out", ] shared_root = os.environ.get("AI4S_SHARED_DIR") + if not shared_root: + inferred = _infer_ai4s_shared_from_job_dir(job_dir) + if inferred is not None: + shared_root = str(inferred) if shared_root: log_candidates.insert( 1, @@ -44,14 +144,53 @@ def main() -> int: return 2 records = parse_slurm_log(log_path) + gbatch = args.global_batch_size + if gbatch is None and manifest.get("global_batch_size") is not None: + gbatch = int(manifest["global_batch_size"]) + + # Per-rank trainer batch enables honest real-per-step throughput (vs nominal global). + nominal_per_rank = None + if manifest.get("orbit2_batch_size") is not None: + nominal_per_rank = int(manifest["orbit2_batch_size"]) + foms = aggregate_foms( records, warmup_batches_per_epoch=args.warmup_batches_per_epoch, steady_epoch_start=args.steady_epoch_start, + global_batch_size=gbatch, + nominal_per_rank_batch=nominal_per_rank, ) foms["job_id"] = job_id foms["log_path"] = str(log_path) + # HBM reserved (peak) + effective-batch integrity. throughput now prefers the REAL + # per-step batch dims (throughput_method=real_per_step_batch); partial_step_fraction and + # steady_realized_batch_dims expose dataloader batch fragmentation so the orchestrator can + # still reject artifact comparisons (loop overnight-3x-001 iter-1 + clean nw test 10504/10505). + try: + log_text = log_path.read_text(encoding="utf-8", errors="replace") + except OSError: + log_text = "" + hbm_vals = [float(x) for x in re.findall(r"memory_reserved:\s*([0-9.]+)\s*GB", log_text)] + foms["hbm_reserved_GB"] = max(hbm_vals) if hbm_vals else None + foms["hbm_reserved_pct_288"] = round(max(hbm_vals) / 288 * 100, 1) if hbm_vals else None + if manifest: + foms["manifest_parallelism"] = manifest.get("parallelism") + foms["manifest_runtime_seconds"] = manifest.get("runtime_seconds") + foms["manifest_data_type"] = manifest.get("data_type") + + _enrich_omnistat(foms, manifest, job_dir) + + # Energy placeholders until wired to rocm_avg_pwr like HydraGNN fom_extractor + foms["energy_J"] = None + foms["mean_power_W"] = None + foms["energy_per_sample_J"] = None + if foms.get("throughput_samples_per_s") and manifest.get("runtime_seconds"): + tp = float(foms["throughput_samples_per_s"]) + rt = float(manifest["runtime_seconds"]) + if tp > 0 and rt > 0 and not math.isnan(tp): + foms["samples_per_job_estimate"] = tp * rt + out = job_dir / "foms.json" out.write_text(json.dumps(foms, indent=2) + "\n", encoding="utf-8") print(f"Wrote {out}") diff --git a/earth_science/models/ORBIT-2/examples/run_gemm_analysis.sh b/earth_science/models/ORBIT-2/examples/run_gemm_analysis.sh new file mode 100755 index 0000000..5c96b9b --- /dev/null +++ b/earth_science/models/ORBIT-2/examples/run_gemm_analysis.sh @@ -0,0 +1,121 @@ +#!/usr/bin/env bash +# run_gemm_analysis.sh — unattended ORBIT-2 GEMM-time bottleneck analysis (1-node + 2-node). +# +# Launches the full TraceLens + Omnistat analyst/verifier flow via the Claude Code CLI to answer +# "where does ORBIT-2's GEMM time actually go?" at 1 and 2 nodes. Set EXCLUDE_NODES to skip any +# known-bad nodes. Runs in tmux; resumable; disk-guarded. +# +# Usage (tmux on login node): +# export ANTHROPIC_API_KEY=... export AI4S_SHARED_DIR= +# export PERF_TOOLS_DIR= # perf_tools.dir in .cluster-config.yaml +# export EXCLUDE_NODES=... # optional: comma-separated nodes to skip +# bash earth_science/models/ORBIT-2/examples/run_gemm_analysis.sh [--preflight-only] +# +# Graceful stop: touch /STOP +set -euo pipefail + +[[ $# -ge 1 ]] || { echo "Usage: $0 [--preflight-only]" >&2; exit 1; } +UUID="$1"; PREFLIGHT_ONLY=0 +[[ "${2:-}" == "--preflight-only" ]] && PREFLIGHT_ONLY=1 + +SCRIPT_DIR=$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd) +REPO_ROOT=$(cd "${SCRIPT_DIR}/../../../.." && pwd) +ORCH_PROMPT="${REPO_ROOT}/earth_science/models/ORBIT-2/recipes/perf-analysis/agents/orchestrator_gemm_analysis.md" + +: "${AI4S_SHARED_DIR:?AI4S_SHARED_DIR must be set}" +: "${PERF_TOOLS_DIR:?PERF_TOOLS_DIR must be set}" +export AI4S_SHARED_DIR PERF_TOOLS_DIR + +ORBIT2_BASE="${AI4S_SHARED_DIR}/models/ORBIT-2" +PERF_RUNS_DIR="${ORBIT2_BASE}/perf-runs" +ANALYSIS_DIR="${PERF_RUNS_DIR}/gemm-analysis-${UUID}" +STATUS_FILE="${ANALYSIS_DIR}/STATUS.txt" +REPORT="${ANALYSIS_DIR}/GEMM_TIME_REPORT.md" +ORBIT2_CLAUDE_MODEL="${ORBIT2_CLAUDE_MODEL:-claude-opus-4-8}" +DISK_GUARD_PCT="${ORBIT2_DISK_GUARD_PCT:-93}" +mkdir -p "$ANALYSIS_DIR" + +_log() { + local line; line="$(date -u +%Y-%m-%dT%H:%M:%SZ) $1" + ( flock -x 9; echo "$line" >> "$STATUS_FILE"; ) 9>>"$STATUS_FILE" + echo "[GEMM] $line" +} +_log "ANALYSIS_START uuid=${UUID} dir=${ANALYSIS_DIR} model=ORBIT-2" + +# ---------------------------------------------------------------- preflight --- +_fail=""; _notes=() +if command -v sinfo >/dev/null 2>&1; then + _n=$(sinfo -h -o '%t' 2>/dev/null | grep -cE '^(idle|mix|alloc)$' || true) + [[ "${_n:-0}" -eq 0 ]] && { _fail="cluster_down"; _notes+=("0 usable nodes"); } +else _notes+=("sinfo not in PATH"); fi +_pct=$(df -P "$PERF_RUNS_DIR" 2>/dev/null | awk 'NR==2 {sub("%","",$5); print $5}') +[[ -z "$_fail" && "${_pct:-0}" -ge 95 ]] && { _fail="disk_full"; _notes+=("perf-runs at ${_pct}%"); } +VENV="${PERF_TOOLS_DIR}/perf-inspect" +for _f in "$ORCH_PROMPT" \ + "${SCRIPT_DIR}/sbatch_train_perf_amd.sh" \ + "${VENV}/bin/omnistat-usermode" \ + "${PERF_TOOLS_DIR}/victoriametrics/victoria-metrics-prod" \ + "${AI4S_SHARED_DIR}/images/pytorch_rocm7.2.2_ubuntu24.04_py3.12_pytorch_release_2.10.0.sif"; do + [[ -z "$_fail" && ! -e "$_f" ]] && { _fail="missing_file"; _notes+=("missing: $_f"); } +done +# TraceLens importable? +if [[ -z "$_fail" ]] && ! "${VENV}/bin/python" -c "import TraceLens" >/dev/null 2>&1; then + _fail="tracelens_missing"; _notes+=("TraceLens not importable in ${VENV}") +fi +for _a in tracelens_analyst tracelens_verifier omnistat_analyst omnistat_verifier synthesizer; do + _p="${REPO_ROOT}/earth_science/models/ORBIT-2/recipes/perf-analysis/agents/${_a}.md" + [[ -z "$_fail" && ! -f "$_p" ]] && { _fail="missing_subagent"; _notes+=("missing: $_p"); } +done +if [[ -n "$_fail" ]]; then + _log "PREFLIGHT_FAIL reason=${_fail} notes='${_notes[*]}'" + printf 'Pre-flight failed: %s\n' "$_fail" >&2; printf ' - %s\n' "${_notes[@]}" >&2; exit 2 +fi +_log "PREFLIGHT_OK notes='${_notes[*]:-none}'" +[[ $PREFLIGHT_ONLY -eq 1 ]] && { _log "PREFLIGHT_ONLY exiting"; echo "Pre-flight OK."; exit 0; } + +command -v claude >/dev/null 2>&1 || { _log "ANALYSIS_ABORT reason=no_claude_cli"; echo "claude CLI not found." >&2; exit 2; } +[[ -n "${ANTHROPIC_API_KEY:-}" ]] || { _log "ANALYSIS_ABORT reason=no_api_key"; echo "ANTHROPIC_API_KEY unset." >&2; exit 2; } + +# ------------------------------------------------------- disk-guard watchdog -- +_disk_guard() { + while [[ ! -f "${ANALYSIS_DIR}/STOP" && ! -f "${ANALYSIS_DIR}/DONE" ]]; do + local pct; pct=$(df -P "$PERF_RUNS_DIR" 2>/dev/null | awk 'NR==2 {sub("%","",$5); print $5}') + # NOTE: never delete traces here — the analysts need them. Only warn. + [[ -n "${pct:-}" && "${pct}" -ge "$DISK_GUARD_PCT" ]] && _log "DISK_WARN fs=${pct}%" + sleep 300 + done +} +_disk_guard & _GUARD_PID=$! +trap '_log "ANALYSIS_SIGNAL cleanup"; kill $_GUARD_PID 2>/dev/null || true' EXIT INT TERM + +# ------------------------------------------------------- invoke orchestrator -- +USER_PROMPT="Run the ORBIT-2 GEMM-time bottleneck analysis (TraceLens + Omnistat analyst/verifier) at 1 and 2 nodes. Follow ${ORCH_PROMPT}. +REPO_ROOT=${REPO_ROOT} +AI4S_SHARED_DIR=${AI4S_SHARED_DIR} +PERF_TOOLS_DIR=${PERF_TOOLS_DIR} +ANALYSIS_DIR=${ANALYSIS_DIR} +EXCLUDE_NODES=${EXCLUDE_NODES:-} +Write GEMM_TIME_REPORT.md to ANALYSIS_DIR. Persist state for resume." + +_MAX=10; i=1 +while [[ ! -f "$REPORT" ]]; do + [[ -f "${ANALYSIS_DIR}/STOP" ]] && { _log "ANALYSIS_ABORT reason=stop_flag"; break; } + [[ $i -gt $_MAX ]] && { _log "ANALYSIS_GIVEUP after ${_MAX} invocations"; break; } + _log "ORCH_INVOKE attempt=${i} model=${ORBIT2_CLAUDE_MODEL}" + set +e + claude --print --model "$ORBIT2_CLAUDE_MODEL" --dangerously-skip-permissions \ + --max-turns 300 --append-system-prompt "$(cat "$ORCH_PROMPT")" "$USER_PROMPT" + rc=$? + set -e + _log "ORCH_RETURN attempt=${i} rc=${rc} report=$([[ -f "$REPORT" ]] && echo yes || echo no)" + [[ -f "$REPORT" ]] && break + sleep $(( 30 * i )); i=$(( i + 1 )) +done + +if [[ -f "$REPORT" ]]; then + touch "${ANALYSIS_DIR}/DONE" + _log "ANALYSIS_COMPLETE report=${REPORT}" + echo "=== GEMM-time analysis complete ==="; echo " report: ${REPORT}" +else + _log "ANALYSIS_INCOMPLETE no report after ${_MAX} attempts"; exit 1 +fi diff --git a/earth_science/models/ORBIT-2/examples/run_optimizer_loop.sh b/earth_science/models/ORBIT-2/examples/run_optimizer_loop.sh new file mode 100755 index 0000000..7c59dda --- /dev/null +++ b/earth_science/models/ORBIT-2/examples/run_optimizer_loop.sh @@ -0,0 +1,182 @@ +#!/usr/bin/env bash +# run_optimizer_loop.sh — ORBIT-2 iterative sysopt loop driver (Bayes-CAST EDM, MI355X). +# +# Usage (tmux on login node): +# export ANTHROPIC_API_KEY=... # optional; only for Claude Code CLI driver +# export AI4S_SHARED_DIR=... +# export PERF_TOOLS_DIR=... +# bash earth_science/models/ORBIT-2/examples/run_optimizer_loop.sh [--preflight-only] +# +# Graceful stop: touch $AI4S_SHARED_DIR/models/ORBIT-2/perf-runs/loop-/STOP +# +# See ../recipes/perf-optimizer-loop/README.md for artifact layout and agent flow. + +set -euo pipefail + +if [[ $# -lt 2 ]]; then + echo "Usage: $0 [--preflight-only]" >&2 + exit 1 +fi + +LOOP_UUID="$1" +N_ITERS="$2" +PREFLIGHT_ONLY=0 +if [[ "${3:-}" == "--preflight-only" ]]; then + PREFLIGHT_ONLY=1 +fi + +SCRIPT_DIR=$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd) +REPO_ROOT=$(cd "${SCRIPT_DIR}/../../../.." && pwd) +RECIPE_DIR="${REPO_ROOT}/earth_science/models/ORBIT-2/recipes/perf-optimizer-loop" +# Loop driver orchestrator (full multi-subagent TraceLens/Omnistat flow). Override with +# ORBIT2_ORCH_PROMPT to point at a different agent prompt. +ORCH_PROMPT="${ORBIT2_ORCH_PROMPT:-${RECIPE_DIR}/agents/orchestrator.md}" + +: "${AI4S_SHARED_DIR:?AI4S_SHARED_DIR must be set}" +: "${PERF_TOOLS_DIR:?PERF_TOOLS_DIR must be set}" + +ORBIT2_BASE="${AI4S_SHARED_DIR}/models/ORBIT-2" +PERF_RUNS_DIR="${ORBIT2_BASE}/perf-runs" +LOOP_DIR="${PERF_RUNS_DIR}/loop-${LOOP_UUID}" +STATUS_FILE="${LOOP_DIR}/STATUS.txt" + +mkdir -p "$LOOP_DIR" + +_log_status() { + local msg="$1" + local line + line="$(date -u +%Y-%m-%dT%H:%M:%SZ) $msg" + ( + flock -x 9 + echo "$line" >> "$STATUS_FILE" + ) 9>>"$STATUS_FILE" + echo "[STATUS] $line" +} + +_log_status "LOOP_START uuid=${LOOP_UUID} n_iters_budget=${N_ITERS} driver=claude-code-cli script=$0 model=ORBIT-2" + +PREFLIGHT_FAIL_REASON="" +PREFLIGHT_NOTES=() + +if command -v sinfo > /dev/null 2>&1; then + _IDLE_OR_MIX=$(sinfo -h -o '%t' 2>/dev/null | grep -cE '^(idle|mix|alloc)$' || true) + if [[ "${_IDLE_OR_MIX:-0}" -eq 0 ]]; then + PREFLIGHT_FAIL_REASON="cluster_down" + PREFLIGHT_NOTES+=("sinfo reports 0 usable nodes") + fi +else + PREFLIGHT_NOTES+=("sinfo not in PATH — cluster checks skipped") +fi + +if [[ -z "$PREFLIGHT_FAIL_REASON" ]]; then + _USE_PCT=$(df -P "$PERF_RUNS_DIR" 2>/dev/null | awk 'NR==2 {sub("%","",$5); print $5}') + if [[ -n "${_USE_PCT:-}" && "${_USE_PCT:-0}" -gt 95 ]]; then + PREFLIGHT_FAIL_REASON="disk_full" + PREFLIGHT_NOTES+=("perf-runs filesystem at ${_USE_PCT}%") + fi +fi + +for _bin in \ + "${PERF_TOOLS_DIR}/perf-inspect/bin/omnistat-usermode" \ + "${PERF_TOOLS_DIR}/victoriametrics/victoria-metrics-prod"; do + if [[ -z "$PREFLIGHT_FAIL_REASON" && ! -x "$_bin" ]]; then + PREFLIGHT_FAIL_REASON="tool_missing" + PREFLIGHT_NOTES+=("missing: $_bin") + fi +done + +for _f in \ + "${ORBIT2_BASE}/overlays/orbit2-overlay.img" \ + "${AI4S_SHARED_DIR}/images/pytorch_rocm7.2.2_ubuntu24.04_py3.12_pytorch_release_2.10.0.sif"; do + if [[ -z "$PREFLIGHT_FAIL_REASON" && ! -f "$_f" ]]; then + PREFLIGHT_FAIL_REASON="image_missing" + PREFLIGHT_NOTES+=("missing: $_f") + fi +done + +if [[ -z "$PREFLIGHT_FAIL_REASON" && ! -f "$ORCH_PROMPT" ]]; then + PREFLIGHT_FAIL_REASON="orch_prompt_missing" + PREFLIGHT_NOTES+=("missing: $ORCH_PROMPT") +fi + +if [[ -n "$PREFLIGHT_FAIL_REASON" ]]; then + _log_status "PREFLIGHT_FAIL reason=${PREFLIGHT_FAIL_REASON} notes='${PREFLIGHT_NOTES[*]}'" + echo "Pre-flight failed: ${PREFLIGHT_FAIL_REASON}" >&2 + printf ' - %s\n' "${PREFLIGHT_NOTES[@]}" >&2 + exit 2 +fi + +_log_status "PREFLIGHT_OK orch_prompt=ok" + +if [[ $PREFLIGHT_ONLY -eq 1 ]]; then + _log_status "PREFLIGHT_ONLY exiting" + echo "Pre-flight OK; orchestrator not invoked." + exit 0 +fi + +if ! command -v claude > /dev/null 2>&1; then + echo "NOTE: 'claude' CLI not found — run the orchestrator manually in Cursor:" >&2 + echo " Read: $ORCH_PROMPT" >&2 + echo " Loop dir: $LOOP_DIR" >&2 + _log_status "LOOP_ABORT reason=no_claude_cli manual_orchestrator_required" + exit 0 +fi + +if [[ -z "${ANTHROPIC_API_KEY:-}" ]]; then + echo "ANTHROPIC_API_KEY not set; cannot invoke claude non-interactively." >&2 + _log_status "LOOP_ABORT reason=no_api_key" + exit 2 +fi + +_on_signal() { + _log_status "LOOP_ABORT reason=signal sig=$1" + exit 130 +} +trap '_on_signal INT' INT +trap '_on_signal TERM' TERM + +USER_PROMPT="Drive the ORBIT-2 iterative sysopt loop (Bayes-CAST EDM, throughput primary FOM): +- Loop UUID: ${LOOP_UUID} +- Loop directory: ${LOOP_DIR} +- Iterations budget: ${N_ITERS} +- Status file: ${STATUS_FILE} +- Recipe README: ${RECIPE_DIR}/README.md + +Follow ${ORCH_PROMPT}. Persist state to disk for resume." + +# Model for the orchestrator. The `opus` alias still points to claude-opus-4-7 on +# CLI 2.1.x, so default to the explicit 4.8 slug; override with ORBIT2_CLAUDE_MODEL. +ORBIT2_CLAUDE_MODEL="${ORBIT2_CLAUDE_MODEL:-claude-opus-4-8}" +_log_status "ORCH_MODEL model=${ORBIT2_CLAUDE_MODEL}" + +_MAX_ATTEMPTS=3 +_attempt=1 +while [[ $_attempt -le $_MAX_ATTEMPTS ]]; do + _log_status "ORCH_INVOKE attempt=${_attempt} model=${ORBIT2_CLAUDE_MODEL}" + set +e + claude \ + --print \ + --model "$ORBIT2_CLAUDE_MODEL" \ + --dangerously-skip-permissions \ + --max-turns 250 \ + --append-system-prompt "$(cat "$ORCH_PROMPT")" \ + "$USER_PROMPT" + _rc=$? + set -e + if [[ $_rc -eq 0 ]]; then + _log_status "ORCH_EXIT_OK attempt=${_attempt}" + break + fi + _log_status "ORCH_EXIT_FAIL attempt=${_attempt} rc=${_rc}" + if [[ $_attempt -eq $_MAX_ATTEMPTS ]]; then + _log_status "LOOP_ABORT reason=orchestrator_failed rc=${_rc}" + exit 1 + fi + _backoff=$(( 30 * (2 ** (_attempt - 1)) )) + sleep "$_backoff" + _attempt=$(( _attempt + 1 )) +done + +_log_status "DRIVER_EXIT clean" +echo "=== ORBIT-2 loop driver done ===" +echo " $LOOP_DIR" diff --git a/earth_science/models/ORBIT-2/examples/run_scaling_study.sh b/earth_science/models/ORBIT-2/examples/run_scaling_study.sh index a9e346a..574070a 100755 --- a/earth_science/models/ORBIT-2/examples/run_scaling_study.sh +++ b/earth_science/models/ORBIT-2/examples/run_scaling_study.sh @@ -20,13 +20,14 @@ SCRIPT_DIR=$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd) export TORCH_NCCL_HIGH_PRIORITY=1 export GPU_MAX_HW_QUEUES=2 -export ORBIT2_DATA_TYPE=float32 +export ORBIT2_DATA_TYPE="${ORBIT2_DATA_TYPE:-bfloat16}" # max_epochs=6 → trains epochs 0–4; collate uses steady epochs 2–4 for FOM export ORBIT2_MAX_EPOCH=6 export ORBIT2_MAX_BATCHES=20 export ORBIT2_BATCH_SIZE=4 export ORBIT2_DATA_ROOT="${ORBIT2_DATA_ROOT:-${AI4S_SHARED_DIR}/models/ORBIT-2/data/superres/prism/10.0_arcmin}" export ORBIT2_CONFIG_TEMPLATE="${ORBIT2_CONFIG_TEMPLATE:-interm_8m_lux.yaml}" +# sbatch_train_amd.sh: 1-node×8-GPU jobs default to fsdp=8 simple_ddp=1; N>1 uses fsdp=N simple_ddp=8 from render defaults. export ORBIT2_SCALING_TAG="${ORBIT2_SCALING_TAG:-prism}" export ORBIT2_OUTPUT_DIR="${ORBIT2_OUTPUT_DIR:-${AI4S_SHARED_DIR}/models/ORBIT-2/outputs/scaling-${ORBIT2_SCALING_TAG}}" @@ -54,7 +55,7 @@ submit() { --job-name="orbit2-scale-${n}N" \ --output="${LOG_DIR}/orbit2-train-%j.out" \ --error="${LOG_DIR}/orbit2-train-%j.out" \ - --export=ALL,TORCH_NCCL_HIGH_PRIORITY,GPU_MAX_HW_QUEUES,ORBIT2_DATA_TYPE,ORBIT2_MAX_EPOCH,ORBIT2_MAX_BATCHES,ORBIT2_BATCH_SIZE,ORBIT2_DATA_ROOT,ORBIT2_CONFIG_TEMPLATE,ORBIT2_SCALING_TAG,ORBIT2_OUTPUT_DIR,AI4S_SHARED_DIR \ + --export=ALL,TORCH_NCCL_HIGH_PRIORITY,GPU_MAX_HW_QUEUES,ORBIT2_DATA_TYPE,ORBIT2_MAX_EPOCH,ORBIT2_MAX_BATCHES,ORBIT2_BATCH_SIZE,ORBIT2_DATA_ROOT,ORBIT2_CONFIG_TEMPLATE,ORBIT2_SCALING_TAG,ORBIT2_OUTPUT_DIR,AI4S_SHARED_DIR,ORBIT2_DISABLE_CKPT=1 \ "$SCRIPT_DIR/sbatch_train_amd.sh" } diff --git a/earth_science/models/ORBIT-2/examples/sbatch_infer_amd.sh b/earth_science/models/ORBIT-2/examples/sbatch_infer_amd.sh index 2144821..36b7f97 100755 --- a/earth_science/models/ORBIT-2/examples/sbatch_infer_amd.sh +++ b/earth_science/models/ORBIT-2/examples/sbatch_infer_amd.sh @@ -98,6 +98,13 @@ export PYTHONNOUSERSITE=1 export HSA_NO_SCRATCH_RECLAIM=1 export MIOPEN_USER_DB_PATH="${TMPDIR:-/tmp}/orbit2-miopen-${SLURM_JOB_ID:-$$}" mkdir -p "$MIOPEN_USER_DB_PATH" +# ORNL Frontier-validated MIOpen conv flags (bayes-cast launch/launch_diffusion.sh): disable Winograd, +# unbound multi-pass Winograd workspace. Override to A/B on Lux (e.g. ORBIT2_MIOPEN_CONV_WINOGRAD=1). +export MIOPEN_DISABLE_CACHE="${MIOPEN_DISABLE_CACHE:-1}" +export MIOPEN_DEBUG_AMD_WINOGRAD_MPASS_WORKSPACE_MAX="${ORBIT2_MIOPEN_WINOGRAD_MPASS_WS_MAX:--1}" +export MIOPEN_DEBUG_AMD_MP_BD_WINOGRAD_WORKSPACE_MAX="${ORBIT2_MIOPEN_MP_BD_WINOGRAD_WS_MAX:--1}" +export MIOPEN_DEBUG_CONV_WINOGRAD="${ORBIT2_MIOPEN_CONV_WINOGRAD:-0}" +export HSA_FORCE_FINE_GRAIN_PCIE="${HSA_FORCE_FINE_GRAIN_PCIE:-1}" # --------------------------------------------------------------------------- # No-overlay dep install (Apptainer only — ~15 min one-time per job) diff --git a/earth_science/models/ORBIT-2/examples/sbatch_infer_docker.sh b/earth_science/models/ORBIT-2/examples/sbatch_infer_docker.sh index 383cb64..9452062 100755 --- a/earth_science/models/ORBIT-2/examples/sbatch_infer_docker.sh +++ b/earth_science/models/ORBIT-2/examples/sbatch_infer_docker.sh @@ -352,6 +352,11 @@ docker run --rm \ -e ORBIT2_ROOT=/orbit2 \ -e PYTHONNOUSERSITE=1 \ -e MIOPEN_USER_DB_PATH=/tmp/miopen \ + -e MIOPEN_DISABLE_CACHE="${MIOPEN_DISABLE_CACHE:-1}" \ + -e MIOPEN_DEBUG_AMD_WINOGRAD_MPASS_WORKSPACE_MAX="${ORBIT2_MIOPEN_WINOGRAD_MPASS_WS_MAX:--1}" \ + -e MIOPEN_DEBUG_AMD_MP_BD_WINOGRAD_WORKSPACE_MAX="${ORBIT2_MIOPEN_MP_BD_WINOGRAD_WS_MAX:--1}" \ + -e MIOPEN_DEBUG_CONV_WINOGRAD="${ORBIT2_MIOPEN_CONV_WINOGRAD:-0}" \ + -e HSA_FORCE_FINE_GRAIN_PCIE="${HSA_FORCE_FINE_GRAIN_PCIE:-1}" \ "$ORBIT2_IMAGE" \ bash "$RUN_SCRIPT" \ "$ROCM_WHL_TAG" \ diff --git a/earth_science/models/ORBIT-2/examples/sbatch_train_amd.sh b/earth_science/models/ORBIT-2/examples/sbatch_train_amd.sh index 92d59c2..8e20b79 100755 --- a/earth_science/models/ORBIT-2/examples/sbatch_train_amd.sh +++ b/earth_science/models/ORBIT-2/examples/sbatch_train_amd.sh @@ -1,7 +1,9 @@ #!/usr/bin/env bash # ORBIT-2 multi-node training on AMD Instinct via SLURM (Apptainer + MPI). # -# Wraps upstream intermediate_downscaling.py with Studio launchers and caps. +# Wraps upstream training with Studio launchers and caps (gptl4py stub, batch cap, +# optional rank-0 profiler hook). Prefer Bayes-CAST `launch_diffusion.sh` when present; +# otherwise `run_orbit2_train.py` + `intermediate_downscaling.py`. # Uses real 10.0_arcmin PRISM data in same-dir mode (see interm_8m_lux.yaml). # # Quick start (1 node, 8 GPUs): @@ -15,13 +17,17 @@ # ORBIT2_DATA_ROOT Data root (PRISM 10.0_arcmin or era5/1.0_deg for sanity) # ORBIT2_CONFIG_TEMPLATE YAML template basename (default: interm_8m_lux.yaml; # use interm_8m_lux_era5.yaml for new ERA5 1.0_deg data) -# ORBIT2_ROOT ORBIT-2 clone (default: $AI4S_SHARED_DIR/models/ORBIT-2/code/ORBIT-2) +# ORBIT2_ROOT When unset: **bayes-cast** clone if `.../code/bayes-cast` exists, else public ORBIT-2 # ORBIT2_SIF Apptainer SIF path # ORBIT2_OVERLAY Pre-built ext3 overlay (required — avoids ~15 min pip/job) # ORBIT2_MAX_EPOCH Cap trainer.max_epochs (default: 3; must be >= 2 — # upstream loop is while (epoch_start+1) < max_epochs) # ORBIT2_MAX_BATCHES Cap batches per epoch (default: 20; 0 = unlimited) # ORBIT2_BATCH_SIZE Per-rank batch size (default: 8) +# ORBIT2_FSDP / ORBIT2_SIMPLE_DDP Parallelism for render (optional). When both unset +# and nodes=1 with 8 GPUs/job, defaults to fsdp=8, simple_ddp=1. +# ORBIT2_LAUNCH_SCRIPT Absolute path to launch script under ORBIT2_ROOT (optional). +# If unset, uses launch_diffusion.sh at repo root or examples/ when present. # ORBIT2_OUTPUT_DIR Job output dir (default: .../outputs/train/) # TORCH_NCCL_HIGH_PRIORITY / GPU_MAX_HW_QUEUES — RCCL tuning (default: 1 / 2) # ORBIT2_SKIP_NODE_HEALTH_PROBE=1 — skip mount probe (not recommended) @@ -47,7 +53,15 @@ else fi ORBIT2_BASE="${AI4S_SHARED_DIR:?AI4S_SHARED_DIR must be set}/models/ORBIT-2" -ORBIT2_ROOT="${ORBIT2_ROOT:-${ORBIT2_BASE}/code/ORBIT-2}" +if [[ -z "${ORBIT2_ROOT:-}" ]]; then + if [[ -d "${ORBIT2_BASE}/code/bayes-cast" ]]; then + ORBIT2_ROOT="${ORBIT2_BASE}/code/bayes-cast" + else + ORBIT2_ROOT="${ORBIT2_BASE}/code/ORBIT-2" + fi +else + ORBIT2_ROOT="${ORBIT2_ROOT}" +fi ORBIT2_SIF="${ORBIT2_SIF:-${AI4S_SHARED_DIR}/images/pytorch_rocm7.2.2_ubuntu24.04_py3.12_pytorch_release_2.10.0.sif}" ORBIT2_OVERLAY="${ORBIT2_OVERLAY:-${ORBIT2_BASE}/overlays/orbit2-overlay.img}" ORBIT2_DATA_ROOT="${ORBIT2_DATA_ROOT:-${ORBIT2_BASE}/data/superres/prism/10.0_arcmin}" @@ -56,6 +70,13 @@ ORBIT2_MAX_BATCHES="${ORBIT2_MAX_BATCHES:-20}" ORBIT2_BATCH_SIZE="${ORBIT2_BATCH_SIZE:-8}" ORBIT2_OUTPUT_DIR="${ORBIT2_OUTPUT_DIR:-${ORBIT2_BASE}/outputs/train/${SLURM_JOB_ID:-$$}}" +for var in ORBIT2_SIF ORBIT2_OVERLAY ORBIT2_DATA_ROOT; do + if [[ ! -e "${!var}" ]]; then + echo "ERROR: $var not found: ${!var}" >&2 + exit 2 + fi +done + TORCH_NCCL_HIGH_PRIORITY="${TORCH_NCCL_HIGH_PRIORITY:-1}" GPU_MAX_HW_QUEUES="${GPU_MAX_HW_QUEUES:-2}" @@ -63,15 +84,41 @@ NODES="${SLURM_JOB_NUM_NODES:-1}" GPUS_PER_NODE=8 TOTAL_RANKS=$((NODES * GPUS_PER_NODE)) -for var in ORBIT2_SIF ORBIT2_OVERLAY ORBIT2_DATA_ROOT; do - if [[ ! -e "${!var}" ]]; then - echo "ERROR: $var not found: ${!var}" >&2 - exit 2 - fi -done +# Parallelism: explicit env, else 1×8 GPU baseline fsdp=8 simple_ddp=1, else fsdp=nodes × simple_ddp=gpus +if [[ -n "${ORBIT2_FSDP:-}" && -n "${ORBIT2_SIMPLE_DDP:-}" ]]; then + RENDER_PARALLEL=(--fsdp "$ORBIT2_FSDP" --simple-ddp "$ORBIT2_SIMPLE_DDP") +elif [[ "$NODES" -eq 1 && "$GPUS_PER_NODE" -eq 8 ]]; then + RENDER_PARALLEL=(--fsdp 8 --simple-ddp 1) +else + RENDER_PARALLEL=() +fi + +# Bayes-CAST ships launch/launch_diffusion.sh as an OLCF/Crusher sbatch+conda script: it +# ignores argv, hardcodes a config, and nests srun — unusable inside Studio Apptainer+srun. +# When launch/train_edm.py exists, call it directly with the rendered /config/config.yaml. +LAUNCH_IC="" +LAUNCH_EDM_DIRECT=0 +if [[ -n "${ORBIT2_LAUNCH_SCRIPT:-}" && -f "${ORBIT2_LAUNCH_SCRIPT}" ]]; then + case "$ORBIT2_LAUNCH_SCRIPT" in + "$ORBIT2_ROOT"/*) _rel="${ORBIT2_LAUNCH_SCRIPT#"$ORBIT2_ROOT"/}"; LAUNCH_IC="/orbit2/$_rel" ;; + *) echo "ERROR: ORBIT2_LAUNCH_SCRIPT must be under ORBIT2_ROOT: $ORBIT2_LAUNCH_SCRIPT" >&2; exit 2 ;; + esac +elif [[ -f "${ORBIT2_ROOT}/launch/train_edm.py" ]]; then + LAUNCH_EDM_DIRECT=1 +elif [[ -f "${ORBIT2_ROOT}/launch_diffusion.sh" ]]; then + LAUNCH_IC="/orbit2/launch_diffusion.sh" +elif [[ -f "${ORBIT2_ROOT}/launch/launch_diffusion.sh" ]]; then + LAUNCH_IC="/orbit2/launch/launch_diffusion.sh" +elif [[ -f "${ORBIT2_ROOT}/examples/launch_diffusion.sh" ]]; then + LAUNCH_IC="/orbit2/examples/launch_diffusion.sh" +fi -if [[ ! -f "${ORBIT2_ROOT}/examples/intermediate_downscaling.py" ]]; then - echo "ERROR: ORBIT-2 clone missing intermediate_downscaling.py: ${ORBIT2_ROOT}" >&2 +_has_upstream=0 +[[ "$LAUNCH_EDM_DIRECT" -eq 1 ]] && _has_upstream=1 +[[ -n "$LAUNCH_IC" ]] && _has_upstream=1 +[[ -f "${ORBIT2_ROOT}/examples/intermediate_downscaling.py" ]] && _has_upstream=1 +if [[ "$_has_upstream" -eq 0 ]]; then + echo "ERROR: ORBIT2_ROOT has no trainable entry (launch/train_edm.py, launch_diffusion.sh, or examples/intermediate_downscaling.py): ${ORBIT2_ROOT}" >&2 exit 2 fi @@ -92,6 +139,7 @@ python3 "$SCRIPT_DIR/render_orbit2_config.py" \ --data-root "$ORBIT2_DATA_ROOT" \ --max-epochs "$ORBIT2_MAX_EPOCH" \ --batch-size "$ORBIT2_BATCH_SIZE" \ + "${RENDER_PARALLEL[@]}" \ -o "$JOB_CONFIG" echo "=== ORBIT-2 Training (Apptainer + MPI) ===" @@ -104,10 +152,23 @@ echo " Data root : $ORBIT2_DATA_ROOT" echo " Max epochs : $ORBIT2_MAX_EPOCH" echo " Max batches : $ORBIT2_MAX_BATCHES" echo " Batch size : $ORBIT2_BATCH_SIZE" +echo " Config : $JOB_CONFIG" echo " RCCL priority: TORCH_NCCL_HIGH_PRIORITY=$TORCH_NCCL_HIGH_PRIORITY" echo " HW queues : GPU_MAX_HW_QUEUES=$GPU_MAX_HW_QUEUES" echo " Output dir : $ORBIT2_OUTPUT_DIR" -echo " Config : $JOB_CONFIG" +echo " ORBIT2_ROOT : $ORBIT2_ROOT" +if ((${#RENDER_PARALLEL[@]})); then + echo " Parallelism : ${RENDER_PARALLEL[*]}" +else + echo " Parallelism : (render default fsdp=$NODES simple_ddp=$GPUS_PER_NODE)" +fi +if [[ "$LAUNCH_EDM_DIRECT" -eq 1 ]]; then + echo " Train entry : python3 /orbit2/launch/train_edm.py (Bayes EDM; Studio bypasses launch_diffusion.sh)" +elif [[ -n "$LAUNCH_IC" ]]; then + echo " Train entry : ${LAUNCH_IC}" +else + echo " Train entry : run_orbit2_train.py (studio)" +fi echo " Node(s) : ${SLURM_NODELIST:-$(hostname)}" echo " Date : $(date)" echo "" @@ -143,6 +204,13 @@ if [[ "$ORBIT2_SKIP_NODE_HEALTH_PROBE" != "1" ]]; then echo " All $SLURM_JOB_NUM_NODES nodes healthy." fi +LAUNCH_DIR="" +if [[ "$LAUNCH_EDM_DIRECT" -eq 1 ]]; then + LAUNCH_DIR="/orbit2/launch" +elif [[ -n "$LAUNCH_IC" ]]; then + LAUNCH_DIR=$(dirname "$LAUNCH_IC") +fi + MASTER_ADDR=$(scontrol show hostnames "$SLURM_JOB_NODELIST" | head -1) MASTER_PORT="${ORBIT2_MASTER_PORT:-29500}" @@ -157,18 +225,44 @@ export OMP_NUM_THREADS="\${SLURM_CPUS_PER_TASK:-7}" export MIOPEN_DISABLE_CACHE=1 export MIOPEN_USER_DB_PATH="\${TMPDIR:-/tmp}/orbit2-miopen-\${SLURM_JOB_ID:-\$\$}-\${SLURM_PROCID:-0}" mkdir -p "\$MIOPEN_USER_DB_PATH" +# ORNL Frontier-validated MIOpen conv flags (bayes-cast launch/launch_diffusion.sh). ORNL DISABLES +# Winograd and unbounds the multi-pass Winograd workspace; tested many times on Frontier (gfx90a / +# ROCm 7.1.1). Defaults below replicate ORNL exactly; override at submit to A/B on Lux (gfx950 / +# ROCm 7.2.2), e.g. ORBIT2_MIOPEN_CONV_WINOGRAD=1 to re-enable Winograd. +export MIOPEN_DEBUG_AMD_WINOGRAD_MPASS_WORKSPACE_MAX="${ORBIT2_MIOPEN_WINOGRAD_MPASS_WS_MAX:--1}" +export MIOPEN_DEBUG_AMD_MP_BD_WINOGRAD_WORKSPACE_MAX="${ORBIT2_MIOPEN_MP_BD_WINOGRAD_WS_MAX:--1}" +export MIOPEN_DEBUG_CONV_WINOGRAD="${ORBIT2_MIOPEN_CONV_WINOGRAD:-0}" export PYTHONNOUSERSITE=1 export HSA_NO_SCRATCH_RECLAIM=1 +# ORNL Frontier-validated (launch_diffusion.sh): fine-grain PCIe coherence. Generic ROCm, portable. +export HSA_FORCE_FINE_GRAIN_PCIE="${HSA_FORCE_FINE_GRAIN_PCIE:-1}" export ORBIT_USE_DDSTORE=0 export TORCH_NCCL_HIGH_PRIORITY="${TORCH_NCCL_HIGH_PRIORITY}" export GPU_MAX_HW_QUEUES="${GPU_MAX_HW_QUEUES}" export ORBIT2_ROOT="/orbit2" export ORBIT2_MAX_BATCHES="${ORBIT2_MAX_BATCHES}" -export ORBIT2_DATA_TYPE="${ORBIT2_DATA_TYPE:-float32}" +export ORBIT2_OUTPUT_DIR="${ORBIT2_OUTPUT_DIR}" +# Pass through checkpoint toggle. Default 0 here (general training script keeps checkpoints); +# throughput/scaling callers set ORBIT2_DISABLE_CKPT=1 to skip all checkpoint writes. +export ORBIT2_DISABLE_CKPT="${ORBIT2_DISABLE_CKPT:-0}" +# Writable inside the container (job dir is bind-mounted) for CK debug NDJSON. +export DEBUG_AGENT_LOG="\${DEBUG_AGENT_LOG:-${ORBIT2_OUTPUT_DIR}/agent-ck.ndjson}" +export ORBIT2_DATA_TYPE="${ORBIT2_DATA_TYPE:-bfloat16}" export ORBIT2_FUSED_ATTN="${ORBIT2_FUSED_ATTN:-DEFAULT}" export ORBIT2_RANK_PRE_TRAIN_HOOK="\${ORBIT2_RANK_PRE_TRAIN_HOOK:-}" -cd /orbit2/examples -exec python3 /examples/run_orbit2_train.py /config/config.yaml +if [[ "${LAUNCH_EDM_DIRECT}" -eq 1 ]]; then + cd /orbit2/launch + python3 /examples/orbit2_rank_hook_runner.py + exec python3 train_edm.py /config/config.yaml +elif [[ -n "${LAUNCH_IC}" ]]; then + cd /orbit2/examples + python3 /examples/orbit2_rank_hook_runner.py + cd "${LAUNCH_DIR}" + exec bash "${LAUNCH_IC}" /config/config.yaml +else + cd /orbit2/examples + exec python3 /examples/run_orbit2_train.py /config/config.yaml +fi RANKEOF chmod +x "$RANK_SCRIPT" diff --git a/earth_science/models/ORBIT-2/examples/sbatch_train_perf_amd.sh b/earth_science/models/ORBIT-2/examples/sbatch_train_perf_amd.sh index 61c7c84..8cd5f03 100755 --- a/earth_science/models/ORBIT-2/examples/sbatch_train_perf_amd.sh +++ b/earth_science/models/ORBIT-2/examples/sbatch_train_perf_amd.sh @@ -1,18 +1,21 @@ #!/usr/bin/env bash -# ORBIT-2 2-node training with PyTorch profiling + Omnistat user-mode telemetry. +# ORBIT-2 training with PyTorch profiling + Omnistat user-mode telemetry. # # Perf-analysis variant of sbatch_train_amd.sh. See recipes/perf-analysis/. +# Default allocation is **1 node × 8 GPUs**; use `sbatch --nodes=2 ...` for multi-node. # -# Quick start: +# Quick start (ERA5 same-dir): # export AI4S_SHARED_DIR=/path/to/shared -# export OMNIHUB_TOOLS_DIR=/path/to/omnihub/tools -# export ORBIT2_DATA_ROOT=$AI4S_SHARED_DIR/models/ORBIT-2/data/superres/prism/10.0_arcmin +# export PERF_TOOLS_DIR=/path/to/perf-tools +# export ORBIT2_DATA_ROOT=$AI4S_SHARED_DIR/models/ORBIT-2/data/superres/era5/1.0_deg +# export ORBIT2_CONFIG_TEMPLATE=edm_8m_era5_1x8.yaml +# export ORBIT2_ROOT=$AI4S_SHARED_DIR/models/ORBIT-2/code/bayes-cast # EDM preset + launcher # sbatch earth_science/models/ORBIT-2/examples/sbatch_train_perf_amd.sh #SBATCH --job-name=orbit2-perf #SBATCH --partition=YOUR_PARTITION_HERE #SBATCH --account=YOUR_ACCOUNT_HERE -#SBATCH --nodes=2 +#SBATCH --nodes=1 #SBATCH --ntasks-per-node=8 #SBATCH --gpus-per-node=8 #SBATCH --cpus-per-task=7 @@ -22,6 +25,8 @@ set -euo pipefail +ORBIT2_PERF_START_TS=$(date +%s) + if [[ -n "${SLURM_JOB_ID:-}" ]]; then _ORIG_CMD=$(scontrol show job "$SLURM_JOB_ID" | sed -n 's/.*Command=\(\S\+\).*/\1/p') SCRIPT_DIR=$(cd "$(dirname "$_ORIG_CMD")" && pwd) @@ -30,28 +35,98 @@ else fi ORBIT2_BASE="${AI4S_SHARED_DIR:?AI4S_SHARED_DIR must be set}/models/ORBIT-2" -ORBIT2_ROOT="${ORBIT2_ROOT:-${ORBIT2_BASE}/code/ORBIT-2}" +if [[ -z "${ORBIT2_ROOT:-}" ]]; then + if [[ -d "${ORBIT2_BASE}/code/bayes-cast" ]]; then + ORBIT2_ROOT="${ORBIT2_BASE}/code/bayes-cast" + else + ORBIT2_ROOT="${ORBIT2_BASE}/code/ORBIT-2" + fi +else + ORBIT2_ROOT="${ORBIT2_ROOT}" +fi ORBIT2_SIF="${ORBIT2_SIF:-${AI4S_SHARED_DIR}/images/pytorch_rocm7.2.2_ubuntu24.04_py3.12_pytorch_release_2.10.0.sif}" ORBIT2_OVERLAY="${ORBIT2_OVERLAY:-${ORBIT2_BASE}/overlays/orbit2-overlay.img}" -ORBIT2_DATA_ROOT="${ORBIT2_DATA_ROOT:-${ORBIT2_BASE}/data/superres/prism/10.0_arcmin}" -ORBIT2_MAX_EPOCH="${ORBIT2_MAX_EPOCH:-3}" +# Default ERA5 1.0° same-dir (111×111 latent) for heavier forwards / GPU saturation baselines. +# PRISM 10.0_arcmin (18×18): export ORBIT2_DATA_ROOT=$ORBIT2_BASE/data/superres/prism/10.0_arcmin +# and ORBIT2_CONFIG_TEMPLATE=interm_8m_lux.yaml +ORBIT2_DATA_ROOT="${ORBIT2_DATA_ROOT:-${ORBIT2_BASE}/data/superres/era5/1.0_deg}" +# Default 6 so steady_batch_time_s (epoch >= 2) has a multi-epoch window; see perf-analysis/HANDOFF.md +ORBIT2_MAX_EPOCH="${ORBIT2_MAX_EPOCH:-6}" ORBIT2_MAX_BATCHES="${ORBIT2_MAX_BATCHES:-20}" ORBIT2_BATCH_SIZE="${ORBIT2_BATCH_SIZE:-4}" ORBIT2_OUTPUT_DIR="${ORBIT2_OUTPUT_DIR:-${ORBIT2_BASE}/perf-runs/${SLURM_JOB_ID:-$$}}" PROFILE_TARGET_EPOCH="${PROFILE_TARGET_EPOCH:-0}" +# These are throughput/speed runs ONLY — do NOT persist model checkpoints. The Bayes-CAST EDM +# trainer otherwise writes a ~71 MB interm_epoch_*.ckpt every epoch to launch/checkpoints/bayes-cast/ +# (save_checkpoint(), train_edm.py:1390). ORBIT2_DISABLE_CKPT=1 makes it skip every write. +# Override with ORBIT2_DISABLE_CKPT=0 only if you actually need a checkpoint from a run. +ORBIT2_DISABLE_CKPT="${ORBIT2_DISABLE_CKPT:-1}" + +# TunableOp (PyTorch GEMM autotuning) lever. Default OFF — baseline byte-identical. +# off : disabled (PYTORCH_TUNABLEOP_ENABLED=0) +# tune : live-tune GEMMs and WRITE per-rank caches to ORBIT2_TUNABLEOP_DIR (ENABLED=1 TUNING=1) +# use : load pre-tuned caches, no further tuning (ENABLED=1 TUNING=0) +# Isolated 1-GPU probe (jobs 10595/10596) proved tuning is NOT blocked on this stack; the +# remaining unknown is the full 8-rank FSDP path — that is exactly what a `tune` run validates. +# Cache dir is a stable shared path so a later `use` run can read a prior `tune` run's caches. +ORBIT2_TUNABLEOP_MODE="${ORBIT2_TUNABLEOP_MODE:-off}" +ORBIT2_TUNABLEOP_DIR="${ORBIT2_TUNABLEOP_DIR:-${ORBIT2_BASE}/tunableop-cache}" -: "${OMNIHUB_TOOLS_DIR:?OMNIHUB_TOOLS_DIR must be set}" -OMNISTAT_VENV="${OMNISTAT_VENV:-${OMNIHUB_TOOLS_DIR}/omnihub-inspect}" +: "${PERF_TOOLS_DIR:?PERF_TOOLS_DIR must be set}" +OMNISTAT_VENV="${OMNISTAT_VENV:-${PERF_TOOLS_DIR}/perf-inspect}" OMNISTAT_TEMPLATE="${OMNISTAT_TEMPLATE:-${SCRIPT_DIR}/../recipes/perf-analysis/omnistat.config.template}" OMNISTAT_USERMODE_INTERVAL="${OMNISTAT_USERMODE_INTERVAL:-1}" TORCH_NCCL_HIGH_PRIORITY="${TORCH_NCCL_HIGH_PRIORITY:-1}" GPU_MAX_HW_QUEUES="${GPU_MAX_HW_QUEUES:-2}" -NODES="${SLURM_JOB_NUM_NODES:-2}" +NODES="${SLURM_JOB_NUM_NODES:-1}" GPUS_PER_NODE=8 TOTAL_RANKS=$((NODES * GPUS_PER_NODE)) +if [[ -n "${ORBIT2_FSDP:-}" && -n "${ORBIT2_SIMPLE_DDP:-}" ]]; then + RENDER_PARALLEL=(--fsdp "$ORBIT2_FSDP" --simple-ddp "$ORBIT2_SIMPLE_DDP") +elif [[ "$NODES" -eq 1 && "$GPUS_PER_NODE" -eq 8 ]]; then + RENDER_PARALLEL=(--fsdp 8 --simple-ddp 1) +else + RENDER_PARALLEL=() +fi + +# Bayes-CAST launch/launch_diffusion.sh is an OLCF/Crusher wrapper (nested srun, ignores argv). +# Prefer launch/train_edm.py + rendered /config/config.yaml when present. +LAUNCH_IC="" +LAUNCH_EDM_DIRECT=0 +if [[ -n "${ORBIT2_LAUNCH_SCRIPT:-}" && -f "${ORBIT2_LAUNCH_SCRIPT}" ]]; then + case "$ORBIT2_LAUNCH_SCRIPT" in + "$ORBIT2_ROOT"/*) _rel="${ORBIT2_LAUNCH_SCRIPT#"$ORBIT2_ROOT"/}"; LAUNCH_IC="/orbit2/$_rel" ;; + *) echo "ERROR: ORBIT2_LAUNCH_SCRIPT must be under ORBIT2_ROOT: $ORBIT2_LAUNCH_SCRIPT" >&2; exit 2 ;; + esac +elif [[ -f "${ORBIT2_ROOT}/launch/train_edm.py" ]]; then + LAUNCH_EDM_DIRECT=1 +elif [[ -f "${ORBIT2_ROOT}/launch_diffusion.sh" ]]; then + LAUNCH_IC="/orbit2/launch_diffusion.sh" +elif [[ -f "${ORBIT2_ROOT}/launch/launch_diffusion.sh" ]]; then + LAUNCH_IC="/orbit2/launch/launch_diffusion.sh" +elif [[ -f "${ORBIT2_ROOT}/examples/launch_diffusion.sh" ]]; then + LAUNCH_IC="/orbit2/examples/launch_diffusion.sh" +fi + +_has_upstream=0 +[[ "$LAUNCH_EDM_DIRECT" -eq 1 ]] && _has_upstream=1 +[[ -n "$LAUNCH_IC" ]] && _has_upstream=1 +[[ -f "${ORBIT2_ROOT}/examples/intermediate_downscaling.py" ]] && _has_upstream=1 +if [[ "$_has_upstream" -eq 0 ]]; then + echo "ERROR: ORBIT2_ROOT has no trainable entry (launch/train_edm.py, launch_diffusion.sh, or examples/intermediate_downscaling.py): ${ORBIT2_ROOT}" >&2 + exit 2 +fi + +LAUNCH_DIR="" +if [[ "$LAUNCH_EDM_DIRECT" -eq 1 ]]; then + LAUNCH_DIR="/orbit2/launch" +elif [[ -n "$LAUNCH_IC" ]]; then + LAUNCH_DIR=$(dirname "$LAUNCH_IC") +fi + for var in ORBIT2_SIF ORBIT2_OVERLAY ORBIT2_DATA_ROOT OMNISTAT_TEMPLATE; do if [[ ! -e "${!var}" ]]; then echo "ERROR: $var not found: ${!var}" >&2 @@ -66,21 +141,47 @@ fi mkdir -p "$ORBIT2_OUTPUT_DIR" OMNISTAT_CONFIG="${ORBIT2_OUTPUT_DIR}/omnistat.config" sed -e "s|@JOB_DIR@|${ORBIT2_OUTPUT_DIR}|g" \ - -e "s|@OMNIHUB_TOOLS_DIR@|${OMNIHUB_TOOLS_DIR}|g" \ + -e "s|@PERF_TOOLS_DIR@|${PERF_TOOLS_DIR}|g" \ "$OMNISTAT_TEMPLATE" > "$OMNISTAT_CONFIG" +if [[ -n "${OMNISTAT_ROCPROF_PROFILE:-}" ]]; then + sed -i.bak "s/^profile = .*/profile = ${OMNISTAT_ROCPROF_PROFILE}/" "$OMNISTAT_CONFIG" && rm -f "${OMNISTAT_CONFIG}.bak" +fi -JOB_CONFIG="${ORBIT2_OUTPUT_DIR}/interm_8m_lux_${SLURM_JOB_ID:-$$}.yaml" +ORBIT2_CONFIG_TEMPLATE="${ORBIT2_CONFIG_TEMPLATE:-edm_8m_era5_1x8.yaml}" +CONFIG_TEMPLATE="${SCRIPT_DIR}/${ORBIT2_CONFIG_TEMPLATE}" +if [[ ! -f "$CONFIG_TEMPLATE" ]]; then + echo "ERROR: ORBIT2_CONFIG_TEMPLATE not found: $CONFIG_TEMPLATE" >&2 + exit 2 +fi + +JOB_CONFIG="${ORBIT2_OUTPUT_DIR}/${ORBIT2_CONFIG_TEMPLATE%.yaml}_${SLURM_JOB_ID:-$$}.yaml" python3 "$SCRIPT_DIR/render_orbit2_config.py" \ + --template "$CONFIG_TEMPLATE" \ --nodes "$NODES" \ --gpus-per-node "$GPUS_PER_NODE" \ --data-root "$ORBIT2_DATA_ROOT" \ --max-epochs "$ORBIT2_MAX_EPOCH" \ --batch-size "$ORBIT2_BATCH_SIZE" \ + "${RENDER_PARALLEL[@]}" \ -o "$JOB_CONFIG" echo "=== ORBIT-2 Perf Analysis Run ===" echo " Nodes : $NODES" echo " Total ranks : $TOTAL_RANKS" +echo " ORBIT2_ROOT : $ORBIT2_ROOT" +if ((${#RENDER_PARALLEL[@]})); then + echo " Parallelism : ${RENDER_PARALLEL[*]}" +else + echo " Parallelism : (render default fsdp=$NODES simple_ddp=$GPUS_PER_NODE)" +fi +if [[ "$LAUNCH_EDM_DIRECT" -eq 1 ]]; then + echo " Train entry : python3 /orbit2/launch/train_edm.py (Bayes EDM)" +elif [[ -n "$LAUNCH_IC" ]]; then + echo " Train entry : ${LAUNCH_IC}" +else + echo " Train entry : run_orbit2_train.py (studio)" +fi +echo " Config : $JOB_CONFIG" echo " Output dir : $ORBIT2_OUTPUT_DIR" echo " Profile epoch: $PROFILE_TARGET_EPOCH" echo " Omnistat cfg : $OMNISTAT_CONFIG" @@ -98,9 +199,10 @@ if [[ "$ORBIT2_SKIP_NODE_HEALTH_PROBE" != "1" ]]; then HM=$(test -d "/home/$USER" 2>/dev/null && echo OK || echo FAIL) SH=$(test -d "'"$AI4S_SHARED_DIR"'" 2>/dev/null && echo OK || echo FAIL) SIF_CHECK=$(test -f "'"$ORBIT2_SIF"'" 2>/dev/null && echo OK || echo FAIL) - echo "NODE_HEALTH $H home=$HM shared=$SH sif=$SIF_CHECK" + DATA_CHECK=$(test -d "'"$ORBIT2_DATA_ROOT"'" 2>/dev/null && echo OK || echo FAIL) + echo "NODE_HEALTH $H home=$HM shared=$SH sif=$SIF_CHECK data=$DATA_CHECK" ' 2>&1 || true - _BAD_NODES=$(grep "^NODE_HEALTH" "$_PROBE_OUT" 2>/dev/null | awk '/FAIL/ {print $2}' | sort -u | tr '\n' ',' | sed 's/,$//' || true) + _BAD_NODES=$(grep "^NODE_HEALTH" "$_PROBE_OUT" 2>/dev/null | awk '/home=FAIL|shared=FAIL|sif=FAIL|data=FAIL/ {print $2}' | sort -u | tr '\n' ',' | sed 's/,$//' || true) if [[ -n "$_BAD_NODES" ]]; then echo "FATAL: NODE_HEALTH_PROBE failed on: $_BAD_NODES" >&2 exit 42 @@ -120,8 +222,28 @@ trap cleanup_omnistat EXIT MASTER_ADDR=$(scontrol show hostnames "$SLURM_JOB_NODELIST" | head -1) MASTER_PORT="${ORBIT2_MASTER_PORT:-29500}" -PROFILER_HOOK="${SCRIPT_DIR}/orbit2_profiler_hook.py" TRACE_DIR="${ORBIT2_OUTPUT_DIR}/traces" +# Bayes EDM: default no hook (Lux profiler imports break); allow ORBIT2_RANK_PRE_TRAIN_HOOK for sysopt patches. +if [[ "${LAUNCH_EDM_DIRECT}" -eq 1 ]]; then + _RANK_PRE_TRAIN_HOOK="${ORBIT2_RANK_PRE_TRAIN_HOOK:-}" +else + _RANK_PRE_TRAIN_HOOK="${ORBIT2_RANK_PRE_TRAIN_HOOK:-/examples/orbit2_profiler_hook.py}" +fi + +# Resolve TunableOp mode → env flags (consumed inside the rank script). +# record : log every GEMM shape to an UNTUNED file WITHOUT tuning (free; for offline workflow) +_TUNABLEOP_RECORD=0 +case "$ORBIT2_TUNABLEOP_MODE" in + tune) _TUNABLEOP_ENABLED=1; _TUNABLEOP_TUNING=1 ;; + use) _TUNABLEOP_ENABLED=1; _TUNABLEOP_TUNING=0 ;; + record) _TUNABLEOP_ENABLED=1; _TUNABLEOP_TUNING=0; _TUNABLEOP_RECORD=1 ;; + off) _TUNABLEOP_ENABLED=0; _TUNABLEOP_TUNING=0 ;; + *) echo "ERROR: ORBIT2_TUNABLEOP_MODE must be off|tune|use|record (got: $ORBIT2_TUNABLEOP_MODE)" >&2; exit 2 ;; +esac +if [[ "$ORBIT2_TUNABLEOP_MODE" != "off" ]]; then + mkdir -p "$ORBIT2_TUNABLEOP_DIR" + echo " TunableOp : mode=$ORBIT2_TUNABLEOP_MODE dir=$ORBIT2_TUNABLEOP_DIR (ENABLED=$_TUNABLEOP_ENABLED TUNING=$_TUNABLEOP_TUNING)" +fi RANK_SCRIPT="${ORBIT2_OUTPUT_DIR}/orbit2_rank_${SLURM_JOB_ID:-$$}.sh" cat > "$RANK_SCRIPT" << RANKEOF @@ -133,22 +255,89 @@ export OMP_NUM_THREADS="\${SLURM_CPUS_PER_TASK:-7}" export MIOPEN_DISABLE_CACHE=1 export MIOPEN_USER_DB_PATH="\${TMPDIR:-/tmp}/orbit2-miopen-\${SLURM_JOB_ID:-\$\$}-\${SLURM_PROCID:-0}" mkdir -p "\$MIOPEN_USER_DB_PATH" +# ORNL Frontier-validated MIOpen conv flags (bayes-cast launch/launch_diffusion.sh). ORNL DISABLES +# Winograd and unbounds the multi-pass Winograd workspace; tested many times on Frontier (gfx90a / +# ROCm 7.1.1). Defaults below replicate ORNL exactly; override at submit to A/B on Lux (gfx950 / +# ROCm 7.2.2), e.g. ORBIT2_MIOPEN_CONV_WINOGRAD=1 to re-enable Winograd. +export MIOPEN_DEBUG_AMD_WINOGRAD_MPASS_WORKSPACE_MAX="${ORBIT2_MIOPEN_WINOGRAD_MPASS_WS_MAX:--1}" +export MIOPEN_DEBUG_AMD_MP_BD_WINOGRAD_WORKSPACE_MAX="${ORBIT2_MIOPEN_MP_BD_WINOGRAD_WS_MAX:--1}" +export MIOPEN_DEBUG_CONV_WINOGRAD="${ORBIT2_MIOPEN_CONV_WINOGRAD:-0}" export PYTHONNOUSERSITE=1 export HSA_NO_SCRATCH_RECLAIM=1 +# ORNL Frontier-validated (launch_diffusion.sh): fine-grain PCIe coherence. Generic ROCm, portable. +export HSA_FORCE_FINE_GRAIN_PCIE="${HSA_FORCE_FINE_GRAIN_PCIE:-1}" export ORBIT_USE_DDSTORE=0 export TORCH_NCCL_HIGH_PRIORITY="${TORCH_NCCL_HIGH_PRIORITY}" export GPU_MAX_HW_QUEUES="${GPU_MAX_HW_QUEUES}" export ORBIT2_ROOT="/orbit2" export ORBIT2_MAX_BATCHES="${ORBIT2_MAX_BATCHES}" -export ORBIT2_DATA_TYPE="${ORBIT2_DATA_TYPE:-float32}" +# Perf experiment: widen path2/refine internal conv channels to a tile-friendly N (>=4) to A/B the +# starved low-channel conv GEMM (edm.py reads this). 0 = original architecture. +export ORBIT2_CONV_PAD="${ORBIT2_CONV_PAD:-0}" +# Perf experiment: run path2/refine convs in NHWC (channels_last) so MIOpen uses implicit-GEMM and +# skips the im2col buffer. 1 = on. FIND_MODE=1 (normal search) lets MIOpen pick the NHWC kernel. +export ORBIT2_CHANNELS_LAST="${ORBIT2_CHANNELS_LAST:-0}" +[[ "\${ORBIT2_CHANNELS_LAST}" == "1" ]] && export MIOPEN_FIND_MODE="\${MIOPEN_FIND_MODE:-1}" +# Default bfloat16 for production-like perf; use float32 if xformers.ops / CK path is unstable. +export ORBIT2_DATA_TYPE="${ORBIT2_DATA_TYPE:-bfloat16}" export ORBIT2_FUSED_ATTN="${ORBIT2_FUSED_ATTN:-DEFAULT}" +# torch.compile lever (perf-optimizer-loop): in-process patch in train_edm.py. +export ORBIT2_TORCH_COMPILE="${ORBIT2_TORCH_COMPILE:-0}" +export ORBIT2_COMPILE_MODE="${ORBIT2_COMPILE_MODE:-default}" +export ORBIT2_COMPILE_DYNAMIC="${ORBIT2_COMPILE_DYNAMIC:-}" +# BLAS backend lever: bf16 large-M GEMMs (var_agg F.linear) throw hipErrorInvalidValue +# under hipBLASLt on MI355X+ROCm7.2.2; TORCH_BLAS_PREFER_HIPBLASLT=0 falls back to rocBLAS. +export TORCH_BLAS_PREFER_HIPBLASLT="${TORCH_BLAS_PREFER_HIPBLASLT:-1}" +# TunableOp GEMM autotuning (default off). Per-rank cache file (each rank owns 1 GPU; all see +# device 0 in-process, so key on SLURM_PROCID to avoid 8 ranks clobbering one results0.csv). +export PYTORCH_TUNABLEOP_ENABLED=${_TUNABLEOP_ENABLED} +export PYTORCH_TUNABLEOP_TUNING=${_TUNABLEOP_TUNING} +# Filename: 'use' reads the offline tuner's merged per-device cache (read-only → all ranks share +# the device file, torch appends the device id to '..._dev' → _dev0.csv.._dev7.csv). tune/record +# WRITE, so key on SLURM_PROCID to stop 8 ranks clobbering one file. +if [[ "${ORBIT2_TUNABLEOP_MODE}" == "use" ]]; then + export PYTORCH_TUNABLEOP_FILENAME="${ORBIT2_TUNABLEOP_DIR}/tunableop_results_dev.csv" +else + export PYTORCH_TUNABLEOP_FILENAME="${ORBIT2_TUNABLEOP_DIR}/tunableop_results_rank\${SLURM_PROCID:-0}.csv" +fi +export PYTORCH_TUNABLEOP_VERBOSE="${PYTORCH_TUNABLEOP_VERBOSE:-1}" +# Bounded tuning — ORBIT-2's patch-embed/var_agg GEMMs have HUGE M (e.g. tn_256_5308416_256, +# M=5.3M rows). The default rotating buffer (-1 → sized to defeat L2) allocates/cycles multi-GB +# input copies per candidate → a single op can take many minutes. ROTATING_BUFFER_SIZE=0 disables +# rotation (reuse one buffer); capped iterations keep per-op tuning to seconds. Override at submit. +export PYTORCH_TUNABLEOP_ROTATING_BUFFER_SIZE="${PYTORCH_TUNABLEOP_ROTATING_BUFFER_SIZE:-0}" +export PYTORCH_TUNABLEOP_MAX_TUNING_ITERATIONS="${PYTORCH_TUNABLEOP_MAX_TUNING_ITERATIONS:-10}" +export PYTORCH_TUNABLEOP_MAX_TUNING_DURATION_MS="${PYTORCH_TUNABLEOP_MAX_TUNING_DURATION_MS:-30}" +# Offline "record" workflow: log every untuned GEMM shape to a file WITHOUT tuning it live, so a +# separate offline tuner (mgpu_tune_gemm_in_file) can tune a SIZE-FILTERED subset (skip giant ops). +export PYTORCH_TUNABLEOP_RECORD_UNTUNED=${_TUNABLEOP_RECORD} +export PYTORCH_TUNABLEOP_UNTUNED_FILENAME="${ORBIT2_TUNABLEOP_DIR}/tunableop_untuned_rank\${SLURM_PROCID:-0}.csv" +# Optional debug passthroughs (empty unless set in submit env). Synchronous kernel +# execution pins the true failing kernel for async HIP errors. +export AMD_SERIALIZE_KERNEL="${AMD_SERIALIZE_KERNEL:-}" +export HIP_LAUNCH_BLOCKING="${HIP_LAUNCH_BLOCKING:-}" export ORBIT2_OUTPUT_DIR="${ORBIT2_OUTPUT_DIR}" +# Throughput runs do not persist checkpoints (save_checkpoint() early-returns on this). +export ORBIT2_DISABLE_CKPT="${ORBIT2_DISABLE_CKPT}" +# Writable inside the container (perf run dir bind-mounted) for CK debug NDJSON. +export DEBUG_AGENT_LOG="\${DEBUG_AGENT_LOG:-${ORBIT2_OUTPUT_DIR}/agent-ck.ndjson}" export ORBIT2_PROFILE_DIR="${TRACE_DIR}" export PROFILE_TARGET_EPOCH="${PROFILE_TARGET_EPOCH}" export PROFILE_RANK0_ONLY=1 -export ORBIT2_RANK_PRE_TRAIN_HOOK="${PROFILER_HOOK}" -cd /orbit2/examples -exec python3 /examples/run_orbit2_train.py /config/config.yaml +export ORBIT2_RANK_PRE_TRAIN_HOOK="${_RANK_PRE_TRAIN_HOOK}" +if [[ "${LAUNCH_EDM_DIRECT}" -eq 1 ]]; then + cd /orbit2/launch + python3 /examples/orbit2_rank_hook_runner.py + exec python3 train_edm.py /config/config.yaml +elif [[ -n "${LAUNCH_IC}" ]]; then + cd /orbit2/examples + python3 /examples/orbit2_rank_hook_runner.py + cd "${LAUNCH_DIR}" + exec bash "${LAUNCH_IC}" /config/config.yaml +else + cd /orbit2/examples + exec python3 /examples/run_orbit2_train.py /config/config.yaml +fi RANKEOF chmod +x "$RANK_SCRIPT" @@ -174,16 +363,23 @@ if [[ "$NODES" -gt 1 ]]; then --env OMPI_MCA_btl_tcp_if_include="$NCCL_SOCKET_IFNAME" --env MPI4PY_RC_THREADS=false ) + # GID-INDEX CAVEAT (MI355X RoCE): NCCL_IB_GID_INDEX=1 only works when every node + # exposes the fabric ULA (RoCEv2) at GID index 1. On some clusters a subset of nodes also + # carries a *global* IPv6 at idx1, which shifts the fabric ULA to idx2 → cross-node + # localGid/remoteGid mismatch → "ionic_comp cqe error 12 / status=12 RETRY_EXC" abort. + # Pin jobs to nodes with consistent GID tables via --exclude/--nodelist, or override + # NCCL_IB_GID_INDEX. Enumerate: cat /sys/class/infiniband//ports/1/gids/{1,2}. + # See run_2node_scaleout_loop.sh (O2_EXCLUDE_PATTERN) for auto-excluding a known-bad range. RCCL_MULTINODE_ENVS=( --bind "${RCCL_ANP_PLUGIN}:${RCCL_ANP_PLUGIN}:ro" --bind "${LIBIONIC_PATH}:${LIBIONIC_PATH}:ro" --env NCCL_NET_PLUGIN="$RCCL_ANP_PLUGIN" --env NCCL_IB_HCA="$NCCL_IB_HCA" - --env NCCL_IB_GID_INDEX=1 + --env NCCL_IB_GID_INDEX="${NCCL_IB_GID_INDEX:-1}" --env NCCL_GDR_FLUSH_DISABLE=1 --env RCCL_GDR_FLUSH_GPU_MEM_NO_RELAXED_ORDERING=0 --env NCCL_GDRCOPY_ENABLE=0 - --env NCCL_IB_QPS_PER_CONNECTION=1 + --env NCCL_IB_QPS_PER_CONNECTION="${NCCL_IB_QPS_PER_CONNECTION:-1}" --env HSA_NO_SCRATCH_RECLAIM=1 --env NCCL_IB_TC=96 --env NCCL_IB_FIFO_TC=192 @@ -192,12 +388,24 @@ if [[ "$NODES" -gt 1 ]]; then --env NET_OPTIONAL_RECV_COMPLETION=1 --env NCCL_IB_USE_INLINE=1 --env NCCL_SOCKET_IFNAME="$NCCL_SOCKET_IFNAME" - --env RCCL_LL128_FORCE_ENABLE=1 + --env RCCL_LL128_FORCE_ENABLE="${RCCL_LL128_FORCE_ENABLE:-1}" --env NCCL_IB_PCI_RELAXED_ORDERING=1 --env NCCL_DMABUF_ENABLE=1 --env NCCL_DEBUG="${NCCL_DEBUG:-WARN}" ) - echo " Multi-node RCCL enabled (NCCL_IB_HCA=$NCCL_IB_HCA)" + # Optional comm-tuning levers (perf-optimizer-loop): only injected when set at submit time, + # so the baseline keeps upstream defaults. These are real multi-node levers (FSDP all-gather + # over IB is the new cost at N>1) — see recipes/perf-optimizer-loop/agents/lever_catalog.yaml. + [[ -n "${NCCL_MIN_NCHANNELS:-}" ]] && RCCL_MULTINODE_ENVS+=(--env NCCL_MIN_NCHANNELS="$NCCL_MIN_NCHANNELS") + [[ -n "${NCCL_NCHANNELS_PER_PEER:-}" ]] && RCCL_MULTINODE_ENVS+=(--env NCCL_NCHANNELS_PER_PEER="$NCCL_NCHANNELS_PER_PEER") + echo " Multi-node RCCL enabled (NCCL_IB_HCA=$NCCL_IB_HCA, QPS=${NCCL_IB_QPS_PER_CONNECTION:-1}, LL128=${RCCL_LL128_FORCE_ENABLE:-1}, MIN_NCHANNELS=${NCCL_MIN_NCHANNELS:-default})" +fi + +# Bind the TunableOp cache dir rw so `tune` runs can persist caches and `use` runs can read them +# across jobs (it lives outside the per-job OUTPUT_DIR). +TUNABLEOP_BIND=() +if [[ "$ORBIT2_TUNABLEOP_MODE" != "off" ]]; then + TUNABLEOP_BIND=(--bind "$ORBIT2_TUNABLEOP_DIR":"$ORBIT2_TUNABLEOP_DIR") fi echo "--- Launching perf training: $TOTAL_RANKS ranks ---" @@ -205,6 +413,7 @@ set +e srun --mpi=pmix apptainer exec \ --rocm --overlay "${ORBIT2_OVERLAY}:ro" \ --bind "/opt/ompi:/opt/ompi:ro" \ + "${TUNABLEOP_BIND[@]}" \ --bind "$ORBIT2_ROOT":/orbit2 \ --bind "$SCRIPT_DIR":/examples \ --bind "$(dirname "$RANK_SCRIPT"):$(dirname "$RANK_SCRIPT")" \ @@ -214,7 +423,7 @@ srun --mpi=pmix apptainer exec \ --env HOSTNAME="$MASTER_ADDR" --env MASTER_PORT="$MASTER_PORT" \ --env PMIX_MCA_gds=hash --env PMIX_MCA_psec=native \ --env PYTHONPATH="/opt/orbit2-pkgs:/orbit2/src:/orbit2" \ - --env LD_LIBRARY_PATH=/opt/venv/lib/python3.12/site-packages/torch/lib \ + --env LD_LIBRARY_PATH="/opt/venv/lib/python3.12/site-packages/torch/lib:/opt/rocm/lib:/usr/lib/x86_64-linux-gnu" \ "${MPI_MULTINODE_ENVS[@]}" "${RCCL_MULTINODE_ENVS[@]}" \ "$ORBIT2_SIF" bash "$RANK_SCRIPT" TRAIN_RC=$? @@ -226,14 +435,50 @@ if [[ ! -f "$_SLURM_LOG" ]]; then fi cp -f "$_SLURM_LOG" "${ORBIT2_OUTPUT_DIR}/orbit2-train-${SLURM_JOB_ID}.out" 2>/dev/null || true -python3 - "$ORBIT2_OUTPUT_DIR" "${SLURM_JOB_ID}" "$TRAIN_RC" <<'PYEOF' -import json, sys +M_FSDP="$NODES" +M_SIMPLE="$GPUS_PER_NODE" +if ((${#RENDER_PARALLEL[@]} == 4)); then + M_FSDP="${RENDER_PARALLEL[1]}" + M_SIMPLE="${RENDER_PARALLEL[3]}" +fi + +RUNTIME_S=$(( $(date +%s) - ORBIT2_PERF_START_TS )) +WORKLOAD_TAG="intermediate_downscaling" +if [[ "${LAUNCH_EDM_DIRECT:-0}" -eq 1 ]]; then + WORKLOAD_TAG="bayes_edm" +fi + +python3 - "$ORBIT2_OUTPUT_DIR" "${SLURM_JOB_ID}" "$TRAIN_RC" "$ORBIT2_ROOT" "$JOB_CONFIG" "$M_FSDP" "$M_SIMPLE" \ + "$RUNTIME_S" "$ORBIT2_BATCH_SIZE" "$TOTAL_RANKS" "$ORBIT2_MAX_EPOCH" "${ORBIT2_DATA_TYPE:-bfloat16}" "$WORKLOAD_TAG" <<'PYEOF' +import json, subprocess +import sys from pathlib import Path -job_dir, job_id, train_rc = Path(sys.argv[1]), sys.argv[2], int(sys.argv[3]) + + +def _git(repo: Path, *args: str) -> str: + try: + return subprocess.check_output(["git", "-C", str(repo), *args], text=True).strip() + except (OSError, subprocess.CalledProcessError): + return "" + + +job_dir = Path(sys.argv[1]) +job_id, train_rc = sys.argv[2], int(sys.argv[3]) +orbit2_root = Path(sys.argv[4]) +job_config = Path(sys.argv[5]) +m_fsdp, m_simple = int(sys.argv[6]), int(sys.argv[7]) +runtime_s = int(sys.argv[8]) +batch = int(sys.argv[9]) +total_ranks = int(sys.argv[10]) +max_epochs = int(sys.argv[11]) +data_type = sys.argv[12] +workload = sys.argv[13] +global_batch = batch * m_fsdp * m_simple + manifest = { "job_id": job_id, "model": "ORBIT-2", - "workload": "intermediate_downscaling", + "workload": workload, "job_dir": str(job_dir), "slurm_log": str(job_dir / f"orbit2-train-{job_id}.out"), "omnistat_config": str(job_dir / "omnistat.config"), @@ -241,6 +486,18 @@ manifest = { "trace_dir": str(job_dir / "traces"), "state": "complete" if train_rc == 0 else "failed", "exit_code": train_rc, + "orbit2_root": str(orbit2_root), + "rendered_config": str(job_config), + "parallelism": {"fsdp": m_fsdp, "simple_ddp": m_simple}, + "orbit2_batch_size": batch, + "total_ranks": total_ranks, + "max_epochs": max_epochs, + "data_type": data_type, + "global_batch_size": global_batch, + "runtime_seconds": runtime_s, + "git_sha": _git(orbit2_root, "rev-parse", "HEAD"), + "git_branch": _git(orbit2_root, "rev-parse", "--abbrev-ref", "HEAD"), + "git_remote_origin": _git(orbit2_root, "remote", "get-url", "origin"), } (job_dir / "manifest.json").write_text(json.dumps(manifest, indent=2) + "\n") print(f"Wrote {job_dir / 'manifest.json'}") diff --git a/earth_science/models/ORBIT-2/examples/stage_era5_3x_symlink.sh b/earth_science/models/ORBIT-2/examples/stage_era5_3x_symlink.sh new file mode 100755 index 0000000..c035c57 --- /dev/null +++ b/earth_science/models/ORBIT-2/examples/stage_era5_3x_symlink.sh @@ -0,0 +1,62 @@ +#!/usr/bin/env bash +# stage_era5_3x_symlink.sh — triple ORBIT-2 ERA5 train samples for HBM-saturation +# perf studies by symlinking the existing year's shards under new year prefixes. +# +# WHY: the Bayes-CAST IterDataModule globs `train/*.npz` (no year filter; see +# src/climate_learn/data/itermodule.py:133). The staged 1.0deg tree ships only +# year 1979 (20 shards x 438 timesteps), which caps the effective per-epoch +# sample count (~1704) and prevents bf16 from reaching >~34% HBM. Adding more +# shard *files* linearly increases samples/epoch and the achievable batch. +# +# This duplicates 1979 data under 1980_/1981_ via symlinks (near-zero disk). +# *** PERF / THROUGHPUT / HBM STUDY ONLY — the data repeats, so it is NOT valid +# for scientific training/eval. *** For real multi-year training, download +# additional years from the ORBIT-2 dataset (see recipes/.../STAGING_ERA5_FOR_HBM.md). +# +# Usage: +# bash earth_science/models/ORBIT-2/examples/stage_era5_3x_symlink.sh \ +# [DATA_ROOT] [N_COPIES] +# DATA_ROOT default: $AI4S_SHARED_DIR/models/ORBIT-2/data/superres/era5/1.0_deg +# N_COPIES default: 3 (1979 + 2 symlinked years => 3x) +# +# Idempotent: re-running relinks the same targets. Undo: remove the symlinked +# shards (see UNSTAGE note printed at the end). + +set -euo pipefail + +DATA_ROOT="${1:-${AI4S_SHARED_DIR:?set AI4S_SHARED_DIR}/models/ORBIT-2/data/superres/era5/1.0_deg}" +N_COPIES="${2:-3}" +TRAIN_DIR="${DATA_ROOT}/train" + +if [[ ! -d "$TRAIN_DIR" ]]; then + echo "ERROR: train dir not found: $TRAIN_DIR" >&2 + exit 2 +fi + +# Discover the base year(s) actually present as real (non-symlink) shards. +mapfile -t BASE_SHARDS < <(find "$TRAIN_DIR" -maxdepth 1 -type f -name '*_*.npz' ! -name 'climatology.npz' -printf '%f\n' | sort) +if [[ ${#BASE_SHARDS[@]} -eq 0 ]]; then + echo "ERROR: no real base shards (YEAR_IDX.npz) in $TRAIN_DIR" >&2 + exit 2 +fi + +# Base year = prefix of the first shard (e.g. 1979 from 1979_0.npz). +BASE_YEAR="${BASE_SHARDS[0]%%_*}" +echo "Base year detected: ${BASE_YEAR} (${#BASE_SHARDS[@]} real shards)" + +made=0 +for ((c = 1; c < N_COPIES; c++)); do + NEW_YEAR=$((BASE_YEAR + c)) + for f in "${BASE_SHARDS[@]}"; do + idx="${f#*_}" # e.g. 0.npz from 1979_0.npz + link="${TRAIN_DIR}/${NEW_YEAR}_${idx}" + target="${TRAIN_DIR}/${f}" + ln -sfn "$target" "$link" + made=$((made + 1)) + done + echo " linked year ${NEW_YEAR} -> ${BASE_YEAR} (${#BASE_SHARDS[@]} shards)" +done + +TOTAL=$(find "$TRAIN_DIR" -maxdepth 1 \( -type f -o -type l \) -name '*_*.npz' ! -name 'climatology.npz' | wc -l) +echo "Done. ${made} symlinks created; train/ now has ${TOTAL} shards (~${N_COPIES}x samples)." +echo "UNSTAGE: find '${TRAIN_DIR}' -maxdepth 1 -type l -name '*_*.npz' -delete" diff --git a/earth_science/models/ORBIT-2/examples/submit_perf_baseline_era5_amd.sh b/earth_science/models/ORBIT-2/examples/submit_perf_baseline_era5_amd.sh new file mode 100755 index 0000000..ca4cd21 --- /dev/null +++ b/earth_science/models/ORBIT-2/examples/submit_perf_baseline_era5_amd.sh @@ -0,0 +1,29 @@ +#!/usr/bin/env bash +# Submit ORBIT-2 perf baseline: ERA5 1.0° same-dir + interm_8m_lux_era5.yaml (matches +# sbatch_train_perf_amd.sh defaults). Requires staged NPZ under ORBIT2_DATA_ROOT. +# +# Usage (from anywhere): +# export AI4S_SHARED_DIR=/path/to/shared +# export PERF_TOOLS_DIR=/path/to/perf-tools +# ./submit_perf_baseline_era5_amd.sh +# Extra sbatch flags: ./submit_perf_baseline_era5_amd.sh --partition=lux --account=myacct + +set -euo pipefail +SCRIPT_DIR=$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd) +: "${AI4S_SHARED_DIR:?set AI4S_SHARED_DIR}" +: "${PERF_TOOLS_DIR:?set PERF_TOOLS_DIR}" + +export ORBIT2_DATA_ROOT="${ORBIT2_DATA_ROOT:-${AI4S_SHARED_DIR}/models/ORBIT-2/data/superres/era5/1.0_deg}" +export ORBIT2_CONFIG_TEMPLATE="${ORBIT2_CONFIG_TEMPLATE:-edm_8m_era5_1x8.yaml}" +export ORBIT2_MAX_EPOCH="${ORBIT2_MAX_EPOCH:-6}" + +# examples -> ORBIT-2 -> models -> earth_science -> repo root +REPO_ROOT=$(cd "$SCRIPT_DIR/../../../.." && pwd) +SBATCH_SCRIPT="$REPO_ROOT/earth_science/models/ORBIT-2/examples/sbatch_train_perf_amd.sh" +if [[ ! -f "$SBATCH_SCRIPT" ]]; then + echo "error: repo root not found from $SCRIPT_DIR (expected $SBATCH_SCRIPT)" >&2 + exit 2 +fi + +cd "$REPO_ROOT" +exec sbatch "$@" "$SBATCH_SCRIPT" diff --git a/earth_science/models/ORBIT-2/examples/sweep_orbit2_batch_bf16_amd.sh b/earth_science/models/ORBIT-2/examples/sweep_orbit2_batch_bf16_amd.sh new file mode 100755 index 0000000..7c30626 --- /dev/null +++ b/earth_science/models/ORBIT-2/examples/sweep_orbit2_batch_bf16_amd.sh @@ -0,0 +1,44 @@ +#!/usr/bin/env bash +# Submit ORBIT-2 bf16 + SDPA batch-size probes for HBM saturation (1 node × 8 GPUs). +# +# Usage (login node, repo root): +# export AI4S_SHARED_DIR=... +# export PERF_TOOLS_DIR=... +# export ORBIT2_ROOT=$AI4S_SHARED_DIR/models/ORBIT-2/code/bayes-cast +# export ORBIT2_DATA_ROOT=$AI4S_SHARED_DIR/models/ORBIT-2/data/superres/era5/1.0_deg +# # Optional: export ORBIT2_ERA5_SPATIAL_RES=111 +# bash earth_science/models/ORBIT-2/examples/sweep_orbit2_batch_bf16_amd.sh 64 128 256 +# +# Each job is independent; use the printed job IDs with: +# python3 earth_science/models/ORBIT-2/examples/run_fom_extractor.py --job-dir "$AI4S_SHARED_DIR/models/ORBIT-2/perf-runs/" +# python3 earth_science/models/ORBIT-2/examples/report_orbit2_gpu_baseline.py --job-dir ".../perf-runs/" +# +# For binary search, pick two bracketing batch sizes from VRAM logs, then add midpoints. + +set -euo pipefail + +SCRIPT_DIR=$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd) +REPO_ROOT=$(cd "${SCRIPT_DIR}/../../../.." && pwd) +SBATCH="${SCRIPT_DIR}/sbatch_train_perf_amd.sh" + +if [[ $# -lt 1 ]]; then + echo "Usage: $0 [batch2 ...]" >&2 + exit 1 +fi + +: "${AI4S_SHARED_DIR:?set AI4S_SHARED_DIR}" +: "${PERF_TOOLS_DIR:?set PERF_TOOLS_DIR}" + +export ORBIT2_DATA_TYPE="${ORBIT2_DATA_TYPE:-bfloat16}" +export ORBIT2_FUSED_ATTN="${ORBIT2_FUSED_ATTN:-DEFAULT}" +export ORBIT2_CONFIG_TEMPLATE="${ORBIT2_CONFIG_TEMPLATE:-edm_8m_era5_1x8.yaml}" +export ORBIT2_MAX_EPOCH="${ORBIT2_MAX_EPOCH:-6}" +export ORBIT2_MAX_BATCHES="${ORBIT2_MAX_BATCHES:-20}" + +for B in "$@"; do + export ORBIT2_BATCH_SIZE="$B" + echo "--- sbatch ORBIT2_BATCH_SIZE=$B ---" + ( cd "$REPO_ROOT" && sbatch --export=ALL,ORBIT2_BATCH_SIZE="$B",ORBIT2_DATA_TYPE="$ORBIT2_DATA_TYPE",ORBIT2_FUSED_ATTN="$ORBIT2_FUSED_ATTN",ORBIT2_CONFIG_TEMPLATE="$ORBIT2_CONFIG_TEMPLATE",ORBIT2_MAX_EPOCH="$ORBIT2_MAX_EPOCH",ORBIT2_MAX_BATCHES="$ORBIT2_MAX_BATCHES" "$SBATCH" ) || true +done + +echo "Done. Track jobs: squeue -u \$USER" diff --git a/earth_science/models/ORBIT-2/examples/upstream_pytorch_sdpa_benchmark.py b/earth_science/models/ORBIT-2/examples/upstream_pytorch_sdpa_benchmark.py new file mode 100755 index 0000000..c80ca9e --- /dev/null +++ b/earth_science/models/ORBIT-2/examples/upstream_pytorch_sdpa_benchmark.py @@ -0,0 +1,968 @@ +from __future__ import annotations + +import itertools +import json +import os +import signal +import subprocess +import sys +import tempfile +from collections import defaultdict +from collections.abc import Callable +from contextlib import nullcontext +from dataclasses import asdict, dataclass + +from tabulate import tabulate +from tqdm import tqdm + +import torch +import torch.utils.benchmark as benchmark +from torch._inductor.utils import do_bench_using_profiling + +# PyTorch 2.10 / some ROCm builds omit FA3/FA4 helpers; keep --flash_test optional. +import torch.nn.attention as _torch_nn_attention + +sdpa_kernel = _torch_nn_attention.sdpa_kernel +SDPBackend = _torch_nn_attention.SDPBackend +activate_flash_attention_impl = getattr( + _torch_nn_attention, "activate_flash_attention_impl", None +) +restore_flash_attention_impl = getattr( + _torch_nn_attention, "restore_flash_attention_impl", None +) +from torch.nn.functional import scaled_dot_product_attention + + +def benchmark_torch_function_in_microseconds(func: Callable, *args, **kwargs) -> float: + # warmup + for _ in range(5): + func(*args, **kwargs) + t0 = benchmark.Timer( + stmt="func(*args, **kwargs)", + globals={"args": args, "kwargs": kwargs, "func": func}, + ) + return t0.adaptive_autorange(min_run_time=0.1).median * 1e6 + + +def benchmark_cuda_function_in_microseconds(func: Callable, *args, **kwargs) -> float: + """Thin wrapper around do_bench_using_profiling (CUDA-oriented; may fail on some ROCm builds).""" + + def no_args(): + func(*args, **kwargs) + + try: + time = do_bench_using_profiling(no_args) + return time * 1e3 + except Exception: + return benchmark_torch_function_in_microseconds(func, *args, **kwargs) + + +@dataclass(frozen=True) +class ExperimentConfig: + batch_size: int + num_heads: int + q_seq_len: int + kv_seq_len: int + embed_dim: int + is_causal: bool + dtype: torch.dtype + backend: SDPBackend | None + flash_impl: str | None = None # None/"FA2" for default, "FA3", or "FA4" + device: torch.device = torch.device("cuda") + + @property + def head_dim(self) -> int: + return self.embed_dim // self.num_heads + + def asdict(self): + dict_obj = asdict(self) + dict_obj["head_dim"] = self.head_dim + return dict_obj + + +@dataclass(frozen=True) +class ExperimentResults: + forward_time: float # microseconds + backward_time: float # microseconds + forward_tflops: float + backward_tflops: float + + def asdict(self): + return asdict(self) + + +@dataclass(frozen=True) +class Experiment: + config: ExperimentConfig + results: ExperimentResults + + def asdict(self): + dict1 = self.config.asdict() + dict2 = self.results.asdict() + return {**dict1, **dict2} + + +def calculate_tflops( + config: ExperimentConfig, + time_us: float, + is_backward: bool = False, + sparsity: float = 0.0, +) -> float: + """ + Calculate TFLOPS for scaled dot product attention. + + Parameters: + - config: The experiment configuration + - time_us: The execution time in microseconds + - is_backward: Whether to calculate for backward pass (includes gradient computation) + - sparsity: Sparsity factor between 0.0 and 1.0, where 0.0 means no sparsity and 1.0 means fully sparse + + Returns: + - TFLOPS value + """ + B = config.batch_size + H = config.num_heads + M = config.q_seq_len + N = config.kv_seq_len + D = config.head_dim + + # Calculate density factor (1.0 - sparsity) + density = 1.0 - sparsity + + # Forward pass FLOPs + qk_flops = ( + M * N * D * 2 + ) # Q*K^T matmul: (M,D) @ (D,N) with 2 FLOPs per multiply-add + softmax_flops = M * N * 2 # Softmax operations (exp and div) + av_flops = ( + M * N * D * 2 + ) # Attention @ V: (M,N) @ (N,D) with 2 FLOPs per multiply-add + + total_flops = B * H * (qk_flops + softmax_flops + av_flops) + + # Apply density factor to account for sparsity + total_flops *= density + + # For backward pass flash uses 2.5x more flops will use this + if is_backward: + total_flops *= 2.5 + + # Convert to TFLOPS: flops / (time_us * 1e-6) / 1e12 + tflops = total_flops / (time_us * 1e-6) / 1e12 + + return tflops + + +def get_input( + config: ExperimentConfig, +) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + q = torch.randn( + (config.batch_size, config.num_heads, config.q_seq_len, config.head_dim), + dtype=config.dtype, + device=config.device, + requires_grad=True, + ) + k = torch.randn( + (config.batch_size, config.num_heads, config.kv_seq_len, config.head_dim), + dtype=config.dtype, + device=config.device, + requires_grad=True, + ) + v = torch.randn( + (config.batch_size, config.num_heads, config.kv_seq_len, config.head_dim), + dtype=config.dtype, + device=config.device, + requires_grad=True, + ) + return q, k, v + + +def run_single_experiment(config: ExperimentConfig) -> ExperimentResults: + q, k, v = get_input(config) + is_causal = config.is_causal + context = ( + sdpa_kernel(config.backend) if config.backend is not None else nullcontext() + ) + + # Activate flash attention implementation if specified (requires both helpers) + if ( + activate_flash_attention_impl is not None + and restore_flash_attention_impl is not None + and config.backend is SDPBackend.FLASH_ATTENTION + and config.flash_impl + in ( + "FA3", + "FA4", + ) + ): + try: + activate_flash_attention_impl(config.flash_impl) + except ImportError as e: + raise RuntimeError( + f"Failed to activate {config.flash_impl}: {e}\n" + f"Please install the required flash attention library or run with default configuration (without --flash_test)." + ) from e + + try: + with context: + forward_time = benchmark_cuda_function_in_microseconds( + scaled_dot_product_attention, + q, + k, + v, + is_causal=is_causal, + attn_mask=None, + ) + out_torch = scaled_dot_product_attention( + q, k, v, is_causal=is_causal, attn_mask=None + ) + d_out = torch.randn_like(out_torch) + backward_time = benchmark_cuda_function_in_microseconds( + out_torch.backward, d_out, retain_graph=True + ) + finally: + if ( + restore_flash_attention_impl is not None + and config.backend is SDPBackend.FLASH_ATTENTION + and config.flash_impl + in ( + "FA3", + "FA4", + ) + ): + restore_flash_attention_impl() + + # Calculate TFLOPS for forward and backward passes + sparsity = 0.5 if is_causal else 0.0 + forward_tflops = calculate_tflops(config, forward_time, sparsity=sparsity) + backward_tflops = calculate_tflops( + config, backward_time, is_backward=True, sparsity=sparsity + ) + + return ExperimentResults( + forward_time=forward_time, + backward_time=backward_time, + forward_tflops=forward_tflops, + backward_tflops=backward_tflops, + ) + + +def print_results(experiments: list[Experiment]): + table_data = defaultdict(list) + for experiment in experiments: + for key, value in experiment.asdict().items(): + table_data[key].append(value) + del table_data["device"] + if table_data["backend"][0] is None: + del table_data["backend"] + if table_data["flash_impl"][0] is None: + del table_data["flash_impl"] + print(tabulate(table_data, headers="keys", tablefmt="pretty", floatfmt=".3f")) + + +def write_results_to_csv( + experiments: list[Experiment], output_dir: str = "benchmark_results" +): + """ + Write experiment results to a CSV file in the specified directory. + The filename includes a timestamp for uniqueness. + """ + import csv + import os + from datetime import datetime + + # Create output directory if it doesn't exist. Fall back to a writable temp dir + # rather than aborting the whole sweep (e.g. read-only /shared inside Apptainer). + try: + os.makedirs(output_dir, exist_ok=True) + except OSError as e: + fallback = os.path.join(tempfile.gettempdir(), "orbit2_sdpa_benchmark_results") + print(f"WARN: cannot write CSV to {output_dir!r} ({e}); using {fallback!r}", flush=True) + output_dir = fallback + os.makedirs(output_dir, exist_ok=True) + + # Generate filename with timestamp + timestamp = datetime.now().strftime("%Y%m%d_%H%M%S") + filename = os.path.join(output_dir, f"benchmark_results_{timestamp}.csv") + + # Get all fields from the first experiment + if not experiments: + return + + fieldnames = list(experiments[0].asdict().keys()) + if "device" in fieldnames: + fieldnames.remove("device") # Remove device field as it's always cuda + + # Write results to CSV + with open(filename, "w", newline="") as csvfile: + writer = csv.DictWriter(csvfile, fieldnames=fieldnames) + writer.writeheader() + for experiment in experiments: + row = experiment.asdict() + if "device" in row: + del row["device"] # Remove device field + writer.writerow(row) + + print(f"Results written to: {filename}") + + +def write_xformers_mea_results_to_csv(rows: list[dict], output_dir: str) -> None: + """Append-style filename for xFormers subprocess sweep (separate from PyTorch SDPA CSV).""" + import csv + from datetime import datetime + + if not rows: + return + os.makedirs(output_dir, exist_ok=True) + timestamp = datetime.now().strftime("%Y%m%d_%H%M%S") + filename = os.path.join(output_dir, f"xformers_mea_micro_{timestamp}.csv") + fieldnames = list(rows[0].keys()) + with open(filename, "w", newline="", encoding="utf-8") as csvfile: + writer = csv.DictWriter(csvfile, fieldnames=fieldnames, extrasaction="ignore") + writer.writeheader() + for row in rows: + writer.writerow(row) + print(f"xFormers MEA results written to: {filename}", flush=True) + + +def generate_experiment_configs() -> list[ExperimentConfig]: + batch_sizes = [1, 8, 16] + num_heads = [16] + q_kv_seq_lens = [(128, 128), (256, 256), (512, 512), (1024, 1024), (8192, 8192)] + embed_dims = [2048] + backends = [None] # If set to None, all backends are enabled + dtypes = [ + torch.bfloat16, + ] + is_causal = [True, False] + all_configs = [] + for ( + bsz, + heads, + (q_seq_len, kv_seq_len), + embed_dim, + causal, + dtype, + backend, + ) in itertools.product( + batch_sizes, num_heads, q_kv_seq_lens, embed_dims, is_causal, dtypes, backends + ): + all_configs.append( + ExperimentConfig( + batch_size=bsz, + num_heads=heads, + q_seq_len=q_seq_len, + kv_seq_len=kv_seq_len, + embed_dim=embed_dim, + is_causal=causal, + dtype=dtype, + backend=backend, + ) + ) + + return all_configs + + +def generate_experiment_configs_with_flash() -> list[ExperimentConfig]: + batch_sizes = [1, 8, 16] + num_heads = [16] + q_kv_seq_lens = [(128, 128), (256, 256), (512, 512), (1024, 1024), (8192, 8192)] + embed_dims = [2048] + backends = [ + SDPBackend.FLASH_ATTENTION, + ] + dtypes = [ + torch.bfloat16, + ] + is_causal = [True, False] + # FA2 (default), "FA3", "FA4" for the alternative implementations + flash_impls = ["FA2", "FA3"] + all_configs = [] + for ( + bsz, + heads, + (q_seq_len, kv_seq_len), + embed_dim, + causal, + dtype, + backend, + flash_impl, + ) in itertools.product( + batch_sizes, + num_heads, + q_kv_seq_lens, + embed_dims, + is_causal, + dtypes, + backends, + flash_impls, + ): + all_configs.append( + ExperimentConfig( + batch_size=bsz, + num_heads=heads, + q_seq_len=q_seq_len, + kv_seq_len=kv_seq_len, + embed_dim=embed_dim, + is_causal=causal, + dtype=dtype, + backend=backend, + flash_impl=flash_impl, + ) + ) + + return all_configs + + +def _parse_orbit_backend(name: str) -> SDPBackend | None: + """Map CLI label → ``SDPBackend`` (ROCm: ``efficient`` is the closest public knob to a fused/CK-style path).""" + n = name.strip().lower() + if n in ("none", "default", "auto", ""): + return None + if n == "math": + return SDPBackend.MATH + if n in ("efficient", "efficient_attention"): + return SDPBackend.EFFICIENT_ATTENTION + if n == "ck": + # Renamed: `ck` was a misleading alias for SDPBackend.EFFICIENT_ATTENTION, which on + # ROCm is AOTriton (`attn_fwd`), NOT Composable Kernel. Real CK lives only in the + # xFormers path (--orbit-include-xformers-ck / --orbit-xformers-modes ck). + raise ValueError( + "--orbit-sweep-backends label 'ck' was removed (it mapped to " + "SDPBackend.EFFICIENT_ATTENTION, which is AOTriton on ROCm, not CK). " + "Use 'efficient'. For actual Composable Kernel use --orbit-include-xformers-ck." + ) + if n in ("flash", "flash_attention"): + return SDPBackend.FLASH_ATTENTION + if n in ("cudnn", "cudnn_attention"): + return SDPBackend.CUDNN_ATTENTION + raise ValueError(f"unknown --orbit-sweep-backends entry: {name!r}") + + +def generate_orbit_micro_configs(backend: SDPBackend | None) -> list[ExperimentConfig]: + """Small bf16 configs (8 heads) including a 648×648 block similar to ORBIT-2 patch tokens.""" + dev = torch.device("cuda") + return [ + ExperimentConfig(2, 8, 128, 128, 256, False, torch.bfloat16, backend, None, dev), + ExperimentConfig(2, 8, 648, 648, 256, False, torch.bfloat16, backend, None, dev), + ] + + +def generate_orbit_varagg_configs( + backend: SDPBackend | None, + *, + batches: list[int], + tokens: int = 648, + n_i: int = 5, + num_heads: int = 16, + embed_dim: int = 256, +) -> list[ExperimentConfig]: + """Replicate Bayes-CAST EDM ``var_agg``/``temporal_agg`` cross-attention shapes. + + ``aggregate_variables`` flattens ``(B, History, L)`` into the SDPA **batch** dim, + so ``batch_size = batch * tokens`` (e.g. 256*648=165888), with ``q_seq_len = N_a = 1`` + and ``kv_seq_len = N_i`` (#input variables). This is the shape that makes ROCm Flash + SDPA exceed the 65535 batch-grid cap (~batch>=128 at tokens=648) while EFFICIENT/MATH + are fine. See recipes/perf-analysis/HANDOFF.md §RESOLVED. + """ + dev = torch.device("cuda") + return [ + ExperimentConfig( + batch_size=b * tokens, + num_heads=num_heads, + q_seq_len=1, + kv_seq_len=n_i, + embed_dim=embed_dim, + is_causal=False, + dtype=torch.bfloat16, + backend=backend, + flash_impl=None, + device=dev, + ) + for b in batches + ] + + +def generate_orbit_selfattn_configs( + backend: SDPBackend | None, + *, + batches: list[int], + seq_len: int = 648, + num_heads: int = 8, + embed_dim: int = 256, +) -> list[ExperimentConfig]: + """Replicate Bayes-CAST EDM main self-attention shapes (Attention class). + + After variable/temporal aggregation the spatial transformer blocks run normal + self-attention on ``(B=real_batch, L=seq_len, D)`` — here ``batch_size`` is the + *real* per-rank batch (NOT inflated), ``q_seq_len == kv_seq_len == seq_len`` + (patch tokens; 108/6 * 216/6 = 648). This is where ROCm Flash is normally the + fastest backend, so it answers "should we just force EFFICIENT everywhere?". + """ + dev = torch.device("cuda") + return [ + ExperimentConfig( + batch_size=b, + num_heads=num_heads, + q_seq_len=seq_len, + kv_seq_len=seq_len, + embed_dim=embed_dim, + is_causal=False, + dtype=torch.bfloat16, + backend=backend, + flash_impl=None, + device=dev, + ) + for b in batches + ] + + +def _orbit_xformers_micro_case_specs() -> list[dict]: + """Named tensor specs for xFormers MEA (BHLD layout before transpose to xFormers BLHD).""" + return [ + { + "case_name": "micro_128_sym", + "batch": 2, + "num_heads": 8, + "q_seq_len": 128, + "kv_seq_len": 128, + "embed_dim": 256, + }, + { + "case_name": "micro_648_sym", + "batch": 2, + "num_heads": 8, + "q_seq_len": 648, + "kv_seq_len": 648, + "embed_dim": 256, + }, + { + "case_name": "asym_1_vs_6", + "batch": 2, + "num_heads": 8, + "q_seq_len": 1, + "kv_seq_len": 6, + "embed_dim": 256, + }, + ] + + +def _internal_xformers_worker_main(cfg_path: str) -> int: + """Child entry: one xFormers MEA call (+ optional backward). Exits 0 on success.""" + with open(cfg_path, encoding="utf-8") as f: + cfg = json.load(f) + + import torch + import xformers.ops as xops + + torch.manual_seed(int(cfg.get("seed", 123))) + device = torch.device("cuda") + B, H, M, N = ( + int(cfg["batch"]), + int(cfg["num_heads"]), + int(cfg["q_seq_len"]), + int(cfg["kv_seq_len"]), + ) + E = int(cfg["embed_dim"]) + D = E // H + dropout_p = float(cfg.get("dropout_p", 0.0)) + dtype = getattr(torch, str(cfg.get("dtype", "bfloat16"))) + mode = str(cfg.get("mode", "ck")).strip().lower() + case_name = str(cfg["case_name"]) + + q = torch.randn(B, H, M, D, device=device, dtype=dtype, requires_grad=True).contiguous() + k = torch.randn(B, H, N, D, device=device, dtype=dtype, requires_grad=True).contiguous() + v = torch.randn(B, H, N, D, device=device, dtype=dtype, requires_grad=True).contiguous() + qx = q.transpose(1, 2).contiguous() + kx = k.transpose(1, 2).contiguous() + vx = v.transpose(1, 2).contiguous() + + def forward(): + if mode == "ck": + return xops.memory_efficient_attention( + qx, kx, vx, p=dropout_p, op=xops.MemoryEfficientAttentionCkOp + ) + if mode == "dispatch": + return xops.memory_efficient_attention(qx, kx, vx, p=dropout_p) + raise ValueError(f"unknown xFormers mode {mode!r} (expected ck|dispatch)") + + try: + forward_time = benchmark_cuda_function_in_microseconds(forward) + out = forward() + torch.cuda.synchronize() + d_out = torch.randn_like(out) + backward_time = benchmark_cuda_function_in_microseconds( + out.backward, d_out, retain_graph=True + ) + except Exception as e: + print("__ORBIT_XF_ERROR__" + json.dumps({"case": case_name, "error": repr(e)}), flush=True) + return 1 + + row = { + "case": case_name, + "mode": mode, + "dropout_p": dropout_p, + "forward_us": forward_time, + "backward_us": backward_time, + } + print("__ORBIT_XF_RESULT__" + json.dumps(row), flush=True) + return 0 + + +def _spawn_xformers_worker(cfg: dict, *, script_path: str, timeout_s: float) -> dict: + """Run one xFormers case in a subprocess so SIGSEGV does not kill the SDPA parent.""" + with tempfile.NamedTemporaryFile( + mode="w", suffix=".json", delete=False, encoding="utf-8" + ) as tf: + json.dump(cfg, tf) + path = tf.name + try: + proc = subprocess.run( + [sys.executable, script_path, "--internal-xformers-worker", path], + capture_output=True, + text=True, + timeout=timeout_s, + env=os.environ.copy(), + ) + except subprocess.TimeoutExpired: + return {"outcome": "TIMEOUT", "returncode": None, "stderr": ""} + finally: + try: + os.unlink(path) + except OSError: + pass + + out = proc.stdout or "" + err = proc.stderr or "" + rc = proc.returncode + + if rc == 0: + for line in out.splitlines(): + if line.startswith("__ORBIT_XF_RESULT__"): + payload = json.loads(line[len("__ORBIT_XF_RESULT__") :]) + return {"outcome": "OK", "returncode": rc, "row": payload, "stderr": err} + for line in out.splitlines(): + if line.startswith("__ORBIT_XF_ERROR__"): + payload = json.loads(line[len("__ORBIT_XF_ERROR__") :]) + return { + "outcome": "PYTHON_ERROR", + "returncode": rc, + "error": payload.get("error", payload), + "stderr": err, + } + return { + "outcome": "NO_MARKER", + "returncode": rc, + "stderr": err, + "stdout_tail": out[-2000:], + } + + if rc is not None and rc < 0: + sig = -rc + try: + sig_name = signal.Signals(sig).name + except ValueError: + sig_name = f"SIGNAL_{sig}" + return { + "outcome": "KILLED_BY_SIGNAL", + "returncode": rc, + "signal": sig_name, + "stderr": err, + "stdout_tail": out[-2000:], + } + + return { + "outcome": f"EXIT_{rc}", + "returncode": rc, + "stderr": err, + "stdout_tail": out[-2000:], + } + + +def run_orbit_xformers_mea_suite( + *, + script_path: str, + modes: list[str], + dropout_ps: list[float], + timeout_s: float, +) -> list[dict]: + """Each row is a dict suitable for tabulate (includes outcome, no ORBIT-specific code paths).""" + rows: list[dict] = [] + for spec in _orbit_xformers_micro_case_specs(): + for mode in modes: + for p in dropout_ps: + cfg = { + **spec, + "mode": mode, + "dropout_p": p, + "dtype": "bfloat16", + "seed": 123, + } + res = _spawn_xformers_worker(cfg, script_path=script_path, timeout_s=timeout_s) + base = { + "case": spec["case_name"], + "xformers_mode": mode, + "dropout_p": p, + "B": spec["batch"], + "H": spec["num_heads"], + "q_len": spec["q_seq_len"], + "kv_len": spec["kv_seq_len"], + "embed_dim": spec["embed_dim"], + } + if res["outcome"] == "OK" and "row" in res: + r = res["row"] + base["forward_us"] = r.get("forward_us") + base["backward_us"] = r.get("backward_us") + base["outcome"] = "OK" + else: + base["forward_us"] = None + base["backward_us"] = None + base["outcome"] = res["outcome"] + if res.get("signal"): + base["signal"] = res["signal"] + if res.get("error"): + base["error"] = res["error"] + rows.append(base) + return rows + + +def main(): + import argparse + + parser = argparse.ArgumentParser(description="SDPA benchmark runner") + parser.add_argument( + "--flash_test", + action="store_true", + help="Use flash attention configs (generate_experiment_configs_with_flash)", + ) + parser.add_argument( + "--orbit-micro", + action="store_true", + help="Studio/ROCm: tiny bf16 grid (see generate_orbit_micro_configs) instead of the full upstream sweep.", + ) + parser.add_argument( + "--orbit-sweep-backends", + type=str, + default=None, + help="With --orbit-micro/--orbit-varagg: comma list (default: default,math,efficient,flash). " + "Labels: default|math|efficient|flash|cudnn — maps to torch.nn.attention.SDPBackend. " + "(NOTE: 'ck' was removed — efficient is AOTriton on ROCm, not Composable Kernel.)", + ) + parser.add_argument( + "--orbit-varagg", + action="store_true", + help="Studio/ROCm: replicate Bayes-CAST EDM var_agg cross-attention shapes " + "(batch_size=batch*tokens, q_len=1, kv_len=N_i) to confirm the ROCm Flash 65535 " + "batch-grid cap. Sweeps --orbit-varagg-batches across --orbit-sweep-backends.", + ) + parser.add_argument( + "--orbit-varagg-batches", + type=str, + default="32,64,128,256", + help="Comma list of per-rank batch sizes for --orbit-varagg (B=batch*tokens).", + ) + parser.add_argument( + "--orbit-varagg-tokens", + type=int, + default=648, + help="Tokens L per sample for --orbit-varagg (B=batch*tokens). Default 648.", + ) + parser.add_argument( + "--orbit-varagg-ni", + type=int, + default=5, + help="N_i (#input variables, kv_seq_len) for --orbit-varagg. Default 5.", + ) + parser.add_argument( + "--orbit-selfattn", + action="store_true", + help="Studio/ROCm: time the main EDM self-attention shapes " + "(batch_size=real batch, q_len=kv_len=seq) across --orbit-sweep-backends. " + "Use to decide whether EFFICIENT is competitive with FLASH for normal self-attention.", + ) + parser.add_argument( + "--orbit-selfattn-batches", + type=str, + default="64,128,256,1024", + help="Comma list of per-rank batch sizes for --orbit-selfattn.", + ) + parser.add_argument( + "--orbit-selfattn-seq", + type=int, + default=648, + help="q_len=kv_len for --orbit-selfattn (patch tokens). Default 648.", + ) + parser.add_argument( + "--orbit-selfattn-heads", + type=int, + default=8, + help="num_heads for --orbit-selfattn (EDM main blocks use 8). Default 8.", + ) + parser.add_argument( + "--orbit-include-xformers-ck", + action="store_true", + help="After the SDPA micro sweep: run xFormers memory_efficient_attention in **isolated subprocesses** " + "(explicit MemoryEfficientAttentionCkOp and/or dispatcher). Confirms SIGSEGV on generic bf16 shapes " + "without ORBIT-2 / Bayes-CAST code; parent process survives a child segfault.", + ) + parser.add_argument( + "--orbit-xformers-modes", + type=str, + default="ck", + help="Comma list with --orbit-include-xformers-ck: ck|dispatch (default: ck). " + "`ck` = MemoryEfficientAttentionCkOp; `dispatch` = no explicit op.", + ) + parser.add_argument( + "--orbit-xformers-dropout", + type=str, + default="0.0", + help="Comma list of dropout_p for xFormers MEA (e.g. 0.0,0.1) matching probe training-like cases.", + ) + parser.add_argument( + "--orbit-xformers-timeout", + type=float, + default=180.0, + help="Timeout seconds per isolated xFormers subprocess (default: 180).", + ) + args = parser.parse_args() + + if args.flash_test and ( + activate_flash_attention_impl is None or restore_flash_attention_impl is None + ): + raise SystemExit( + "ERROR: --flash_test requires activate_flash_attention_impl / " + "restore_flash_attention_impl (missing in this PyTorch build). " + "Omit --flash_test or use a newer PyTorch checkout." + ) + + seed = 123 + torch.manual_seed(seed) + results: list[Experiment] = [] + + if args.orbit_selfattn: + sweep_raw = args.orbit_sweep_backends or "default,flash,efficient,math" + labels = [s.strip() for s in sweep_raw.split(",") if s.strip()] + batches = [int(s) for s in args.orbit_selfattn_batches.split(",") if s.strip()] + out_dir = os.environ.get("ORBIT2_SDPA_CSV_DIR", "/tmp/orbit2_sdpa_benchmark_results") + for lab in labels: + bk = _parse_orbit_backend(lab) + print(f"\n{'=' * 14} orbit self-attn sdpa_kernel={lab!r} ({bk!s}) " + f"seq={args.orbit_selfattn_seq} heads={args.orbit_selfattn_heads} {'=' * 14}\n", + flush=True) + configs = generate_orbit_selfattn_configs( + bk, batches=batches, seq_len=args.orbit_selfattn_seq, + num_heads=args.orbit_selfattn_heads, + ) + chunk: list[Experiment] = [] + for config in tqdm(configs, desc=f"selfattn[{lab}]"): + try: + chunk.append(Experiment(config, run_single_experiment(config))) + except Exception as e: # noqa: BLE001 + print(f" [FAIL] batch={config.batch_size}: " + f"{type(e).__name__}: {str(e).splitlines()[0]}", flush=True) + if chunk: + print_results(chunk) + write_results_to_csv(chunk, out_dir) + results.extend(chunk) + return + + if args.orbit_varagg: + sweep_raw = args.orbit_sweep_backends or "default,math,efficient,flash" + labels = [s.strip() for s in sweep_raw.split(",") if s.strip()] + batches = [int(s) for s in args.orbit_varagg_batches.split(",") if s.strip()] + out_dir = os.environ.get("ORBIT2_SDPA_CSV_DIR", "/tmp/orbit2_sdpa_benchmark_results") + for lab in labels: + bk = _parse_orbit_backend(lab) + print(f"\n{'=' * 14} orbit var_agg sdpa_kernel={lab!r} ({bk!s}) " + f"tokens={args.orbit_varagg_tokens} N_i={args.orbit_varagg_ni} {'=' * 14}\n", + flush=True) + configs = generate_orbit_varagg_configs( + bk, batches=batches, tokens=args.orbit_varagg_tokens, n_i=args.orbit_varagg_ni + ) + chunk: list[Experiment] = [] + for config in tqdm(configs, desc=f"varagg[{lab}]"): + B_eff = config.batch_size + try: + chunk.append(Experiment(config, run_single_experiment(config))) + print(f" [PASS] batch={B_eff // args.orbit_varagg_tokens} " + f"(B={B_eff})", flush=True) + except Exception as e: # noqa: BLE001 + print(f" [FAIL] batch={B_eff // args.orbit_varagg_tokens} " + f"(B={B_eff}): {type(e).__name__}: {str(e).splitlines()[0]}", flush=True) + if chunk: + print_results(chunk) + write_results_to_csv(chunk, out_dir) + results.extend(chunk) + return + + if args.orbit_micro: + sweep_raw = args.orbit_sweep_backends or "default,math,efficient,flash" + labels = [s.strip() for s in sweep_raw.split(",") if s.strip()] + out_dir = os.environ.get("ORBIT2_SDPA_CSV_DIR", "/tmp/orbit2_sdpa_benchmark_results") + for lab in labels: + bk = _parse_orbit_backend(lab) + print(f"\n{'=' * 18} orbit sdpa_kernel={lab!r} ({bk!s}) {'=' * 18}\n", flush=True) + configs = generate_orbit_micro_configs(bk) + chunk: list[Experiment] = [] + for config in tqdm(configs, desc=f"orbit[{lab}]"): + try: + chunk.append(Experiment(config, run_single_experiment(config))) + except Exception as e: + print(f"SKIP {config.asdict()}: {type(e).__name__}: {e}", flush=True) + if chunk: + print_results(chunk) + write_results_to_csv(chunk, out_dir) + results.extend(chunk) + + if args.orbit_include_xformers_ck: + modes = [] + for s in args.orbit_xformers_modes.split(","): + m = s.strip().lower() + if not m: + continue + if m not in ("ck", "dispatch"): + raise SystemExit( + f"unknown --orbit-xformers-modes entry {m!r} (allowed: ck, dispatch)" + ) + modes.append(m) + if not modes: + raise SystemExit("--orbit-xformers-modes produced an empty list") + dropout_ps: list[float] = [] + for s in args.orbit_xformers_dropout.split(","): + s = s.strip() + if not s: + continue + dropout_ps.append(float(s)) + if not dropout_ps: + raise SystemExit("--orbit-xformers-dropout produced an empty list") + + print( + f"\n{'=' * 18} xFormers MEA (isolated subprocess per case) {'=' * 18}\n", + flush=True, + ) + script_path = os.path.abspath(__file__) + xf_rows = run_orbit_xformers_mea_suite( + script_path=script_path, + modes=modes, + dropout_ps=dropout_ps, + timeout_s=float(args.orbit_xformers_timeout), + ) + print(tabulate(xf_rows, headers="keys", tablefmt="pretty", floatfmt=".3f"), flush=True) + write_xformers_mea_results_to_csv(xf_rows, out_dir) + return + + if args.flash_test: + configs = generate_experiment_configs_with_flash() + else: + configs = generate_experiment_configs() + for config in tqdm(configs): + results.append(Experiment(config, run_single_experiment(config))) + + print_results(results) + write_results_to_csv(results, "../benchmark_results") + + +if __name__ == "__main__": + if len(sys.argv) >= 3 and sys.argv[1] == "--internal-xformers-worker": + raise SystemExit(_internal_xformers_worker_main(sys.argv[2])) + main() diff --git a/earth_science/models/ORBIT-2/examples/validate_orbit2_optimizer_loop_recipe.sh b/earth_science/models/ORBIT-2/examples/validate_orbit2_optimizer_loop_recipe.sh new file mode 100755 index 0000000..a55a110 --- /dev/null +++ b/earth_science/models/ORBIT-2/examples/validate_orbit2_optimizer_loop_recipe.sh @@ -0,0 +1,26 @@ +#!/usr/bin/env bash +# validate_orbit2_optimizer_loop_recipe.sh — CI/repo smoke: required files exist. +set -euo pipefail +SCRIPT_DIR=$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd) +REPO_ROOT=$(cd "${SCRIPT_DIR}/../../../.." && pwd) +BASE="${REPO_ROOT}/earth_science/models/ORBIT-2" +REQ=( + "${BASE}/recipes/perf-optimizer-loop/README.md" + "${BASE}/recipes/perf-optimizer-loop/lever_catalog.yaml" + "${BASE}/recipes/perf-optimizer-loop/pitfall-diagnosis.md" + "${BASE}/recipes/perf-optimizer-loop/STAGING_ERA5_FOR_HBM.md" + "${BASE}/recipes/perf-optimizer-loop/agents/orchestrator.md" + "${BASE}/recipes/perf-optimizer-loop/agents/lever_picker.md" + "${BASE}/recipes/perf-optimizer-loop/agents/fom_extractor.md" + "${BASE}/recipes/perf-optimizer-loop/agents/story_writer.md" + "${BASE}/examples/run_optimizer_loop.sh" + "${BASE}/examples/sweep_orbit2_batch_bf16_amd.sh" + "${BASE}/examples/run_fom_extractor.py" +) +for f in "${REQ[@]}"; do + if [[ ! -f "$f" ]]; then + echo "missing: $f" >&2 + exit 1 + fi +done +echo "ORBIT-2 perf-optimizer-loop recipe files: OK" diff --git a/earth_science/models/ORBIT-2/model.yaml b/earth_science/models/ORBIT-2/model.yaml index 1c3bcf3..074eb81 100644 --- a/earth_science/models/ORBIT-2/model.yaml +++ b/earth_science/models/ORBIT-2/model.yaml @@ -17,7 +17,7 @@ recipes: script: examples/run_visualize.py slurm: examples/sbatch_infer_amd.sh - task: perf-analysis - description: 2-node training perf analysis with Omnistat + PyTorch profiler + description: Training perf analysis (default 1 node × 8 GPU) with Omnistat + PyTorch profiler; default template **edm_8m_era5_1x8.yaml** (Bayes-CAST EDM, `fsdp=8`/`simple_ddp=1`); override for Lux `interm_8m_lux_era5.yaml` recipe_path: recipes/perf-analysis/ script: examples/run_orbit2_train.py slurm: examples/sbatch_train_perf_amd.sh @@ -29,9 +29,33 @@ recipes: env_vars: ORBIT2_ROOT: - description: Path to cloned XiaoWang-Github/ORBIT-2 repository + description: Path to training code; bind-mounted at /orbit2. **Bayes-CAST** (`…/code/bayes-cast`) for **`edm_8m_era5_1x8.yaml`** + `launch/launch_diffusion.sh`; public ORBIT-2 clone otherwise. required: true default: null + ORBIT2_CONFIG_TEMPLATE: + description: YAML template basename under examples/ (Bayes-CAST **`edm_8m_era5_1x8.yaml`** default for perf; Lux **`interm_8m_lux.yaml`** / **`interm_8m_lux_era5.yaml`** for public ORBIT-2 `res_slimvit`) + required: false + default: edm_8m_era5_1x8.yaml + ORBIT2_FSDP: + description: Override YAML fsdp when set with ORBIT2_SIMPLE_DDP (optional) + required: false + default: null + ORBIT2_SIMPLE_DDP: + description: Override YAML simple_ddp when set with ORBIT2_FSDP (optional) + required: false + default: null + ORBIT2_DATA_TYPE: + description: trainer.data_type — bfloat16 or float32 (sbatch_train_* default bfloat16; float32 for CK/xformers debug) + required: false + default: bfloat16 + ORBIT2_FUSED_ATTN: + description: xFormers attention backend hint (sbatch defaults DEFAULT = PyTorch SDPA when CK path unavailable) + required: false + default: DEFAULT + ORBIT2_LAUNCH_SCRIPT: + description: Absolute path under ORBIT2_ROOT to launch_diffusion.sh (optional; auto-detected if omitted) + required: false + default: null ORBIT2_CHECKPOINT: description: Path to .ckpt file (downloaded from HF in synthetic mode) required: false diff --git a/earth_science/models/ORBIT-2/recipes/data/README.md b/earth_science/models/ORBIT-2/recipes/data/README.md index 0e16590..880b08d 100644 --- a/earth_science/models/ORBIT-2/recipes/data/README.md +++ b/earth_science/models/ORBIT-2/recipes/data/README.md @@ -43,6 +43,8 @@ Use a **shared project or scratch filesystem** visible from both login and compu 2. Activate or use a **Globus endpoint** your site provides (or a personal endpoint on a machine that can reach the cluster filesystem), then run a **Globus transfer** into a directory on your cluster's Lustre, GPFS, or equivalent. 3. Point your ORBIT-2 YAML at those directories once the transfer completes. +For **GPU saturation / sysopt** on Bayes-CAST EDM (breaking small-staging batch caps), see also [../perf-optimizer-loop/STAGING_ERA5_FOR_HBM.md](../perf-optimizer-loop/STAGING_ERA5_FOR_HBM.md) (same Constellation/Globus provenance, ops-focused). + **Checkpoints** for inference or fine-tuning are on Hugging Face ([`jychoi-hpc/ORBIT-2`](https://huggingface.co/jychoi-hpc/ORBIT-2)); use the Hub for weights and configs, not as a substitute for the full curated **training** NPZ-style dataset tree unless your workflow explicitly uses only Hub artifacts. ### When you cannot pull from the public catalog but have OLCF or a collaborator who does diff --git a/earth_science/models/ORBIT-2/recipes/perf-analysis/BASELINE_LOCKIN.md b/earth_science/models/ORBIT-2/recipes/perf-analysis/BASELINE_LOCKIN.md new file mode 100644 index 0000000..1b8b119 --- /dev/null +++ b/earth_science/models/ORBIT-2/recipes/perf-analysis/BASELINE_LOCKIN.md @@ -0,0 +1,107 @@ +# ORBIT-2 — baseline lock-in for GPU saturation (MI355-class, 1×8) + +Use this file to **freeze** what “baseline” means while you sweep knobs. Update the **Reference job** row when you promote a new winner. + +## Reference job (current) + +| Field | Value (example: job **9145**) | +|--------|--------------------------------| +| **Job dir** | `$AI4S_SHARED_DIR/models/ORBIT-2/perf-runs//` | +| **Template** | `interm_8m_lux_era5.yaml` (ERA5 1.0°, same-dir, `spatial_resolution` 111×111) | +| **Parallelism** | `fsdp=8`, `simple_ddp=1` (1 node × 8 ranks) | +| **dtype** | `bfloat16` (Studio **`sbatch_train_*` / `run_scaling_study.sh`** default). Set **`ORBIT2_DATA_TYPE=float32`** if `xformers.ops` / CK attention misbehaves; **`ORBIT2_FUSED_ATTN=DEFAULT`** (already in sbatch) prefers PyTorch SDPA. | +| **Epochs / cap** | `ORBIT2_MAX_EPOCH=6`, `ORBIT2_MAX_BATCHES=20` | +| **Model block** | `preset: res_slimvit`, `embed_dim: 256`, `depth: 6`, `decoder_depth: 4`, `num_heads: 4`, `mlp_ratio: 4` | + +## Bayes-CAST EDM — bf16 + SDPA (sysopt default, 1×8) + +**Perf-optimizer-loop** uses **`ORBIT2_DATA_TYPE=bfloat16`**, **`ORBIT2_FUSED_ATTN=DEFAULT`** (PyTorch SDPA), template **`edm_8m_era5_1x8.yaml`**, and **`ORBIT2_ROOT`** pointing at **Bayes-CAST** (`…/code/bayes-cast`). + +| Step | Action | +|------|--------| +| 1 | If VRAM plateaus **below** target HBM fraction, **stage more / higher-res ERA5** — see [`../perf-optimizer-loop/STAGING_ERA5_FOR_HBM.md`](../perf-optimizer-loop/STAGING_ERA5_FOR_HBM.md) (breaks the ~1704 effective-batch cap on tiny staging trees). | +| 2 | Run a **bf16 batch sweep** (binary search): [`examples/sweep_orbit2_batch_bf16_amd.sh`](../../examples/sweep_orbit2_batch_bf16_amd.sh) | +| 3 | Per job: `python3 …/run_fom_extractor.py --job-dir …/perf-runs/` then `report_orbit2_gpu_baseline.py` → **`baseline_report.md`** beside `manifest.json`. | +| 4 | Lock **~85–90%** `memory_reserved` (leave fragmentation headroom) and best **throughput** (`run_fom_extractor` → `throughput_samples_per_s`) vs **`steady_batch_time_s`**. | + +**Reference job (bf16):** ✅ **bf16 works** after a compute-side attention fix (2026-06-11). The earlier crash (`var_agg` SDPA, `HIP error: invalid argument`) was **ROCm Flash attention's 65535 batch-grid cap**: `var_agg` flattens `(B,History,L)` into the SDPA batch dim `B = batch·648`, which exceeds 65535 at **batch ≥ 128**. Fix: force `CrossAttention` SDPA to **EFFICIENT_ATTENTION+MATH** (not Flash) in `src/climate_learn/models/hub/components/attention.py`. Validated job **9238** (bf16 b256 COMPLETES, decreasing loss). Root cause + probe in [`HANDOFF.md`](HANDOFF.md) → "RESOLVED". Note the **~1704-sample data cap**: bf16 won't reach 85–90% HBM at 1704 samples (fp32@1704 ≈ 183 GiB ≈ 64%; bf16 ≈ half) — stage more ERA5 to saturate. + +**Reference job (bf16, working) — LOCKED 2026-06-11:** `jobid=9240`, `ORBIT2_BATCH_SIZE=1024`, `ORBIT2_DATA_ROOT=${AI4S_SHARED_DIR}/models/ORBIT-2/data/superres/era5/1.0_deg`, `ORBIT2_ERA5_SPATIAL_RES=111`, `ORBIT2_DATA_TYPE=bfloat16`, `ORBIT2_FUSED_ATTN=DEFAULT`, `fsdp=8`/`simple_ddp=1`. **Throughput = 4098 samples/s** (`steady_batch_time_s=2.00`, gb=8192, loss sanity ✅). Reproduced by job 9257 (4097 s/s). + +**Batch-scaling story (bf16, num_workers=1, 1704-sample data, sweep-batch-nw1 2026-06-12):** throughput rises monotonically with batch size — the textbook compute-bound GEMM-efficiency curve (tiny batches starve the MFMA units; larger batches amortize launch/overhead): + +| batch | global batch | throughput | steady step | job | +|------:|------:|-----------:|------------:|-----| +| 32 | 256 | 1869 s/s | 0.137s | 9265 | +| 64 | 512 | 2530 s/s | 0.202s | 9266 | +| 128 | 1024 | 2529 s/s | 0.405s | 9267 | +| 256 | 2048 | 2819 s/s | 0.727s | 9268 | +| 512 | 4096 | 3117 s/s | 1.314s | 9269 | +| 1024 | 8192 | **3987 s/s** | 2.055s | 9270 | + +**⚠️ Throughput ≠ VRAM at the current 1704-sample data cap** — bigger batch saturates more HBM but is *slower* here, because batch > sample-count just repeats/partial-fills the same step: + +| batch | gb | HBM | throughput | steady step | job | +|------:|----:|----:|-----------:|------------:|-----| +| 512 | 4096 | ~11% | 3930 s/s | 1.04s | 9239 | +| **1024** | **8192** | **~21%** | **4098 s/s (best)** | **2.00s** | **9240** | +| 1704 | 13632 | ~34% | 2975 s/s | 4.58s | 9241 | + +So the **throughput-optimal baseline is batch 1024 (~21% HBM)**, not the VRAM-max batch 1704 (-27% throughput). To make HBM saturation *and* throughput agree, you must **stage more ERA5** (raise the 1704-sample cap) and re-sweep — until then the optimizer loop ranks levers against the batch-1024 baseline. + +**bf16 VRAM calibration (post-fix, MI355X 1×8, ERA5 1.0° `edm_8m_era5_1x8.yaml`, FSDP 8 / DDP 1, rank-0 end-of-run):** + +| Per-rank `batch_size` | `max_memory_allocated` | `memory_reserved` | ~% of 288 GB | Job | +|----------------------:|-----------------------:|------------------:|-------------:|-----| +| 512 | 18.93 GiB | 31.87 GiB | ~11% | 9239 | +| 1024 | 37.69 GiB | 61.62 GiB | ~21% | 9240 | +| 1704 (sample cap) | 62.61 GiB | 97.28 GiB | ~34% | 9241 | + +**Takeaway:** with only **~1704 train samples**, bf16 caps at **~97 GiB reserved (~34% HBM)** — batch ≥1704 just repeats the same 1704-sample step. To reach **85–90%** you must **stage more ERA5** (more years and/or higher resolution; see [`../perf-optimizer-loop/STAGING_ERA5_FOR_HBM.md`](../perf-optimizer-loop/STAGING_ERA5_FOR_HBM.md)) then re-sweep. Growth is ~linear here (512→1704), so ~3× more samples at batch≈5000 would approach ~85%. + +## Bayes-CAST EDM reference (`edm_8m_era5_1x8`, float32, 1×8, ERA5 1.0°) + +Smoke **9158** used **`ORBIT2_BATCH_SIZE=2`** (~0.35 GiB **`max_memory_allocated`**, ~0.6 GiB reserved). Follow-up calibration jobs (same template, **`ORBIT2_DATA_TYPE=float32`**, FSDP 8 / DDP 1) on a single MI355X node — rank-0 **`torch.cuda.max_memory_allocated`** / **`memory_reserved`** at **end of epoch 0** (upstream logs once per epoch): + +| Per-rank `batch_size` | ~`max_memory_allocated` | ~`memory_reserved` | Job | +|----------------------:|------------------------:|-------------------:|-----| +| 32 | 2.31 GiB | 1.10 GiB | 9159 | +| 64 | 4.81 GiB | 4.23 GiB | 9160 | +| 96 | 6.72 GiB | 7.49 GiB | 9161 | +| 128 | 8.66 GiB | 4.29 GiB | 9162 | +| 160 | 10.78 GiB | 11.58 GiB | 9163 | +| 192 | 12.90 GiB | 18.39 GiB | 9164 | +| 224 | 15.02 GiB | 15.04 GiB | 9165 | +| 256 | 17.12 GiB | 18.39 GiB | 9166 | +| 1024 | 67.94 GiB | ~73 GiB | 9167 | +| 2048 | 112.93 GiB | ~183 GiB | 9168 | +| 3072 | 112.93 GiB | ~183 GiB | 9169 | + +**Takeaway:** growth is **not linear** in batch (see jump **256 → 1024**). Naive linear fit from small batches **over-predicts** headroom. For **~288 GiB** HBM, **`ORBIT2_BATCH_SIZE=2048`** already uses **~183 GiB reserved** (~**64 %**). Jobs **9168** and **9169** logged the **same** peak memory: rank-0 **`y.shape`** was **`[1704, 3, 108, 216]`** in both (not 2048/3072), so for this ERA5 staging + 8-way layout the **effective per-step batch dim is capping near 1704** — raising **`trainer.batch_size` above ~2048** did not increase the tensor batch for those epochs; re-measure if you add data or change sampler/shard rules. Re-measure after switching to **`bfloat16`**. + +## What “maximal utilization” means here + +1. **HBM capacity** — Rank-0 log `torch.cuda.memory_reserved` near a **practical** fraction of available per-GPU HBM (leave headroom for fragmentation and checkpoint spikes). Example **float32** ERA5 same-dir calibration on this stack: job **9145** @ per-rank batch **4** → ~**12.1 GB**; job **9146** @ batch **8** → ~**28.1 GB** (late-epoch `batch_idx 19` lines) — still far below **~288 GB** HBM-class MI355 OAM → **raise batch** (and prefer **`bfloat16`** to stretch further) before widening the ViT. +2. **Compute / bandwidth** — Omnistat rocprofiler-derived **MFMA / FETCH** activity during **epoch ≥ 2** (after `run_fom_extractor.py` steady window). Use perf-analysis agents on `omnistat-db/` + `traces/` once batch is pinned. +3. **Step time** — `steady_batch_time_s` from logs (sparse: upstream prints `Batch N: … seconds` mainly for **batch 0 and 10** per epoch). Compare across batch sweep, not only VRAM. + +## Ordered next steps (do in sequence) + +1. **Lock topology + data + template** — Keep **ERA5** + **`interm_8m_lux_era5.yaml`** + **1×8** + **`fsdp=8`/`simple_ddp=1`** fixed until VRAM plateaus. +2. **Sweep `ORBIT2_BATCH_SIZE`** — Prefer **binary search** between a known-good batch and a guessed upper bound (fewer SLURM jobs than stepping by +1). Re-run `run_fom_extractor.py` + `report_orbit2_gpu_baseline.py` per candidate; pick the best **time vs throughput** under high VRAM. +3. **Optional: denser batch timing in logs** — If you need more than two `Batch … seconds` lines per epoch, raise upstream logging frequency or increase `ORBIT2_MAX_BATCHES` (still capped per epoch in trainer) so steady stats are statistically tighter. +4. **Omnistat + TraceLens** — On the **winning** job id only: confirm MFMA/HBM story matches “saturated” intent ([`one-node-gpu-baseline.md`](one-node-gpu-baseline.md) §6–7, [`ai4science-perf-analysis`](../../../../.cursor/skills/ai4science-perf-analysis/SKILL.md)). +5. **Only if VRAM is high and MFMA still low** — Then widen **`embed_dim`/`depth`** (with `embed_dim % num_heads == 0`) or revisit `num_workers` / I/O. + +## Max batch: estimate vs measure + +**There is no substitute for a real forward+backward** on your exact **dtype, FSDP layout, and config** — analytical “FLOPs-only” caps ignore fragmentation, activation checkpoints, Lightning overhead, and checkpoint spikes. + +1. **Two-point linear fit (starting guess only)** — From two short jobs at batches \(B_1, B_2\) with similar steady **`memory_reserved`** (e.g. last batch in epoch), fit a line and extrapolate to a **target GiB** (e.g. 220–260 GiB under 288 GiB HBM). Script: [`examples/orbit2_estimate_batch_from_memory.py`](../../examples/orbit2_estimate_batch_from_memory.py). + *Example (float32, same logs as above, target 240 GiB):* slope ≈ **4.0 GiB per batch-step** between \(B=4\) and \(B=8\) → naive linear intercept gives **\(B \approx 61\)** — **often over-optimistic** once attention and allocator effects bend the curve; treat as an **upper exploratory bound**, not the first batch to run overnight. +2. **After switching to `bfloat16`** — Re-measure two batches (e.g. 8 and 16): BF16 rarely halves *all* tensors (master weights, NCCL buffers, I/O), so **do not** assume “2× the float32 batch” without logs. +3. **Better than linear stepping:** **binary search** on `ORBIT2_BATCH_SIZE` with **1–2 epoch** smoke jobs (`ORBIT2_MAX_EPOCH=3` is enough to see reserved memory plateau mid-epoch if you watch `batch_idx 19` lines). Optionally add a **single-node one-rank** micro-probe (not shipped here) if you only need allocator OOM without full 8-way FSDP — less representative but fast. + +## Promotion rule + +When a new job beats the old one on **(steady VRAM fraction, MFMA evidence, acceptable `steady_batch_time_s`)**, replace the **Reference job** section with that job id and the chosen `ORBIT2_BATCH_SIZE`. diff --git a/earth_science/models/ORBIT-2/recipes/perf-analysis/HANDOFF.md b/earth_science/models/ORBIT-2/recipes/perf-analysis/HANDOFF.md index bfe0bcc..9669e31 100644 --- a/earth_science/models/ORBIT-2/recipes/perf-analysis/HANDOFF.md +++ b/earth_science/models/ORBIT-2/recipes/perf-analysis/HANDOFF.md @@ -1,23 +1,29 @@ # ORBIT-2 perf-analysis HANDOFF -_Last updated: 2026-06._ +_Last updated: 2026-06-11 (perf-optimizer-loop + bf16 omnistat profile + manifest fields)._ ## What works | Item | Status | Notes | |------|--------|-------| -| 1-node × 8-GPU training on 10.0_arcmin same-dir | **Green** | Representative steady batch time ~0.5–2 s after warmup (workload- and I/O-dependent) | -| `sbatch_train_amd.sh` + `run_orbit2_train.py` | **Green** | gptl4py stub, FusedAttn fallback when `xformers.ops` unavailable, batch cap | +| 1-node × 8-GPU training on **ERA5 1.0°** same-dir (default perf template) | **Green** | Heavier **111×111** latent vs PRISM **18×18**; better default for GPU saturation baselines. Representative steady batch time workload-dependent. | +| 1-node × 8-GPU on **PRISM** 10.0_arcmin same-dir | **Green** | Lighter forwards; override `ORBIT2_DATA_ROOT` + `interm_8m_lux.yaml` if ERA5 not staged. | +| `sbatch_train_amd.sh` / `sbatch_train_perf_amd.sh` | **Green** | Prefer `launch_diffusion.sh` under `ORBIT2_ROOT` when present; else `run_orbit2_train.py`. `orbit2_rank_hook_runner.py` runs profiler hook before shell launcher. | +| **1-node default parallelism** | **fsdp=8, simple_ddp=1** | Override with `ORBIT2_FSDP` / `ORBIT2_SIMPLE_DDP`. Multi-node: omit for `fsdp=nodes`, `simple_ddp=8`. | +| `ORBIT2_CONFIG_TEMPLATE` + `render_orbit2_config.py` | **Green** | Default perf: **`edm_8m_era5_1x8.yaml`** (Bayes-CAST `edm_8m_era5` + **`fsdp=8`/`simple_ddp=1`/`seq_par:1`**, templated `data_dir` + **`ORBIT2_ERA5_SPATIAL_RES`** default 111). Lux **`interm_8m_lux_era5.yaml`** for public ORBIT-2 `res_slimvit`. `__DATA_TYPE__` defaults **bfloat16**. | | Node health probe (exit 42) | **Green** | Catches broken home/shared/SIF/data mounts before training | | `superres_mag: 1` same-dir config | **Required** | Without it, 4× superres vs same-resolution targets causes tensor shape mismatch | +| Baseline lock-in + VRAM-first sweep | **Doc** | [BASELINE_LOCKIN.md](BASELINE_LOCKIN.md) — freeze topology/template; sweep `ORBIT2_BATCH_SIZE` before widening ViT | ## Landmines 1. **`max_epochs` must be ≥ 2** — upstream loop is `while (epoch_start + 1) < max_epochs`; `max_epochs=1` runs zero epochs. 2. **`xformers.ops` missing in overlay** — `run_orbit2_train.py` can fall back to `FusedAttn.DEFAULT` (PyTorch SDPA). Rebuild overlay per `build_overlay_amd.sh` if you need the CK attention path. -3. **`bfloat16` + CK attn** needs working `xformers.ops`; use `ORBIT2_DATA_TYPE=float32` until the overlay provides `ops`. +3. **`bfloat16` + xformers CK path** — needs working `xformers.ops`. Studio sbatch sets **`ORBIT2_FUSED_ATTN=DEFAULT`** so PyTorch SDPA is used when CK is unavailable. If you still see dtype-related crashes, fall back to **`ORBIT2_DATA_TYPE=float32`** for isolation. 4. **Same-dir data** — both `low_res_dir` and `high_res_dir` point at the same grid; set `superres_mag: 1` in `interm_8m_lux.yaml` / `interm_8m_lux_era5.yaml`. 5. **Batch cap** — must patch `intermediate_downscaling` before `main()` (not via `runpy.run_path` reload). +6. **Bayes-CAST `launch/launch_diffusion.sh`** — upstream file is often an OLCF/Crusher **job wrapper** (nested `srun`, ignores argv). Studio **`sbatch_train_*`** does **not** exec it when **`$ORBIT2_ROOT/launch/train_edm.py`** exists; ranks run **`python3 train_edm.py /config/config.yaml`** from **`/orbit2/launch`** (no `examples/` required). For Lux/public ORBIT-2, **`launch_diffusion.sh`** + argv still applies when **`train_edm.py`** is absent. +7. **`train_edm.py` + `ERA5_1`** — some Bayes-CAST snapshots only define **`std_delta`** for IMERG/HRRR; **`data_key=="ERA5_1"`** then hits **`UnboundLocalError`** at the first step. Fix: add an **`elif data_key == "ERA5_1"`** branch (per-channel std aligned with your staged **`normalize_std.npz`** and **`dict_out_variables` order**), or use the small patch applied on your shared clone (example values `[20.601086, 5.542934, 4.7573]` for T2m / u10 / v10 @ 1.0°). ## FOM contract (steady-state + loss sanity) @@ -32,6 +38,109 @@ python3 parse_training_log.py --log ... --steady-epoch-start 2 --warmup-batches- python3 collate_scaling_study.py --log-dir .../logs --jobs ... --steady-epoch-start 2 --require-loss-sanity ``` +## GEMM-time bottleneck + ORNL MIOpen flags (2026-06-15, jobs 12562/12648) + +Full TraceLens+Omnistat analyst/verifier run (`agents/orchestrator_gemm_analysis.md`, driver +`examples/run_gemm_analysis.sh`) on 1- and 2-node profiled EDM (bf16, batch 4096, ERA5 1.0°). Report: +`perf-runs/gemm-analysis-/GEMM_TIME_REPORT.md`. + +- **The hot kernel is NOT attention/MLP.** A single hipBLASLt GEMM `Cijk_…MT16x16x32_MI16` is + **~43 %** of the step at *both* scales (scale-invariant). It is the **im2col-lowered low-channel + conv backward** for the **3- and 4-channel stem/head projection convs**, stuck on a 16×16 macro-tile + at ~0.03 TFLOP/s. The tunable `nn.Linear` bf16 GEMMs (256×256 tiles, 600–1074 TFLOP/s) are only + ~2 %. **This explains why TunableOp gave no uplift** — it tuned already-optimal GEMMs and can't + touch the conv-lowered one. `im2col` (`Im2d2Col_v2`) adds ~14 %; conv **forward** is a `naive_conv` + with **fp64 accumulate** (`..._ushort_double_ushort`) at ~3 %. +- **TraceLens `report.xlsx:ops_summary_by_category` MISLABELS the conv-GEMM as `CONV_bwd`.** The + 2-node analyst trusted it and claimed "naive conv 65.8 %"; the verifier refuted it from raw + `traceEvents` (naive_conv_bwd = 0.46 %, GEMM = 50.7 %). **Always rank by raw kernel-name device + time, not the category rollup,** for conv-heavy models. (`examples/compare_trace_kernels.py` does + this tool-independently.) +- **~43–44 % whole-job idle at both scales** (Omnistat util mean ~56 %) is the second lever — it is + structural inter-step overhead, **not** comm-bound (exposed comm 0.13 %→0.47 %; the Omnistat + "comm-bound" claim was refuted by the trace). +- **EDM trainer now has an env-gated profiler.** `train_edm.py:_orbit2_make_profiler` writes a rank-0 + `*.pt.trace.json` for `PROFILE_TARGET_EPOCH` into `ORBIT2_PROFILE_DIR` (no-op unless set). This is + what makes "where does GEMM time go" answerable for the EDM path (the old Lux `res_slimvit` hook + did not apply). + +### ORNL Frontier MIOpen flags ported to the sbatch (validate on Lux) + +From bayes-cast `launch/launch_diffusion.sh` (ORNL, gfx90a/ROCm7.1.1, "tested many times"): ORNL +**disables Winograd** and unbounds the multi-pass Winograd workspace. Now defaults in +`sbatch_train_perf_amd.sh` / `sbatch_train_amd.sh` / `sbatch_infer_*.sh` (override to A/B): + +``` +MIOPEN_DEBUG_CONV_WINOGRAD=0 # ORBIT2_MIOPEN_CONV_WINOGRAD (0=off ORNL default, 1=on) +MIOPEN_DEBUG_AMD_WINOGRAD_MPASS_WORKSPACE_MAX=-1 # ORBIT2_MIOPEN_WINOGRAD_MPASS_WS_MAX +MIOPEN_DEBUG_AMD_MP_BD_WINOGRAD_WORKSPACE_MAX=-1 # ORBIT2_MIOPEN_MP_BD_WINOGRAD_WS_MAX +HSA_FORCE_FINE_GRAIN_PCIE=1 +``` + +`MIOPEN_DISABLE_CACHE=1` + `MIOPEN_USER_DB_PATH=/tmp/` already matched ORNL — keep them (do NOT +"persist the find-db"; that deviates from the validated setup). **Deliberately NOT ported** (Frontier +Slingshot/Cray-specific, would break Lux IB/ionic): `FI_CXI_*`, `NCCL_NET="AWS Libfabric"`, +`NCCL_NET_PLUGIN=librccl-net.so`, `NCCL_SOCKET_IFNAME=hsn0`, Frontier `LD_PRELOAD` host paths. +A/B (ORNL Winograd-off vs Winograd-on vs pre-flags +baseline) → `perf-runs/miopen-ornl-validation-/MIOPEN_ORNL_VALIDATION.md`. + +**MIOpen ORNL-flags A/B result (2026-06-16, jobs 12997/12998 vs 12562): performance-NEUTRAL on Lux** +(steady 8.62→8.66→8.74 s, within ~1-2 % noise; loss sanity OK). MIOpen never selected Winograd for +the 3/4-channel convs (0 % Winograd kernels in all arms), so toggling it changes nothing — our +bottleneck is shape-driven, not algorithm-selection-driven. Keeping the flags as defaults (faithful, +harmless, may matter elsewhere). They are **not** the lever here. + +### Conv channel-padding — TABLED candidate, timing-only (NOT ADOPTED, 2026-06-16, jobs 14633/14634) + +> **Status: tabled.** The −25 % is a *timing-only* result. Channel-padding **changes model capacity** +> (wider internal convs ⇒ different params), and **we have not run a convergence/quality study**, so it +> is **not adopted** and `ORBIT2_CONV_PAD` stays **default 0 (off)**. Treat the numbers below as a +> *potential* speedup pending a same-quality convergence check (matched final loss/skill at equal or +> fewer steps). Do not enable for real training runs until that check passes. + +`edm.py` `path2`/`refine` convs run at out_channels(=3)/cnn_ratio(=4) at full resolution → the +starved `Cijk_…MT16x16x32` GEMM @ ~0.03 TFLOP/s. New env knob **`ORBIT2_CONV_PAD=N`** (in `edm.py`, +default 0 = original) widens those convs' INTERNAL width to N (tile-friendly) keeping the 3-channel +I/O. Single-node A/B (identical bf16/4096/ERA5 config), N=16: + +| Arm | steady_batch_time_s | total GPU-kernel ms | dominant `MT16x16x32` GEMM | +|---|---|---|---| +| control (pad=0) | 8.646 | 21082 | 9245 ms (43.9 %) | +| **widened (pad=16)** | **6.498 (−24.8 %)** | 15150 (−28 %) | **gone** (top GEMM now `MT16x16x128`, 4123 ms) | + +**~25 % faster step, loss *sanity* holds (not convergence).** Confirms the GEMM report: low channel +count forced the bad tile; widening makes it tile-friendly and the efficiency gain beats the extra +FLOPs. **Why tabled:** "loss sanity" only checks the loss is finite/decreasing over a few epochs — it +is **not** evidence the wider model reaches the same quality. Before adopting, run a matched +convergence study (same data/schedule, compare final validation loss/skill and steps-to-target for +pad=0 vs pad=16); only adopt if quality is equal-or-better. `channels_last`/implicit-GEMM as a way to +kill the residual `im2col` (~21 %) is **already disproven** (see next section — naive_conv fallback, +7-10× slower). If/when revisited, the open timing question is **N=32** (top GEMM `MT16x16x128` at 27 % +still has headroom) — but only worth it after convergence is settled. +Report: `perf-runs/conv-pad-validation-/CONV_PAD_VALIDATION.md`. + +### `channels_last` + implicit-GEMM — DEAD END on ROCm 7.2.2 (DISPROVEN, 2026-06-16, jobs 14645-14648) + +Tested the "stack `channels_last` on top" idea above as a clean 2×2 ({NCHW, channels_last} × {pad=0, +pad=16}) via the env knob **`ORBIT2_CHANNELS_LAST`** in `edm.py`, +default 0; sets `MIOPEN_FIND_MODE=1` on the NHWC arms). **It is 7-10× SLOWER, not faster:** + +| Arm | CONV_PAD | channels_last | steady_batch_time_s | vs control | +|---|---|---|---|---| +| nchw_pad0 (control) | 0 | off | 8.62 | — | +| **nchw_pad16** | 16 | off | **6.38** | **−26 %** ✅ | +| cl_pad0 | 0 | on | 60.8 | +606 % 💥 | +| cl_pad16 | 16 | on | 80.1 | +829 % 💥 | + +**Root cause (kernel trace):** MIOpen on this build has **no tuned NHWC implicit-GEMM solver** for +these conv shapes. Forcing channels_last makes it fall back to the brute-force **`naive_conv_…wrw_nhwc`** +direct-convolution kernel, which dominates **~88-91 %** of GPU-kernel time (cl_pad0: 139,522 ms; +cl_pad16: 185,198 ms). The `im2col` (`Im2d2Col_v2`, 14-21 % in NCHW) is **not** wasted overhead — it is +the price of reaching the tuned **hipBLASLt GEMM** kernels (`Cijk_…`, 31-54 % in NCHW). NHWC discards +that path. **Lesson: do NOT use `channels_last` for these low-channel convs on MI355X/ROCm 7.2.2.** +**The lever is NCHW + `ORBIT2_CONV_PAD` (−26 %); channels_last is abandoned.** Knob kept (default off) +only so the dead end is reproducible. Report: `perf-runs/conv-layout-validation-/CONV_LAYOUT_VALIDATION.md`. + ## Scaling / wall-clock lessons - Long runs on full 180×360 grids need **enough wall time** (multi-epoch × many batches can exceed 1 h on some sites). @@ -39,16 +148,161 @@ python3 collate_scaling_study.py --log-dir .../logs --jobs ... --steady-epoch-st - **Multi-node hangs** are often **infrastructure** (drained nodes, stuck scheduler states, shared filesystem hiccups) — confirm node state with `sinfo` / site ops before assuming a code bug. - For exclusive GPU jobs, **avoid stacking two jobs on the same node**. Pin nodes with `--nodelist=...` only when your site policy allows it and you have confirmed the nodes are healthy. +## Debug session handoff — Bayes-CAST EDM, `bfloat16`, xFormers CK (SIGSEGV) + +**Why this session felt confusing** + +- **Two trees:** Instrumentation and hotfixes for Bayes-CAST live on the **compute-side clone** (e.g. **`$AI4S_SHARED_DIR/models/ORBIT-2/code/bayes-cast/`**), not inside this git repo. AI4Science Studio only documents launchers and perf recipes; agents may edit `/shared/...` while you expect everything under `ai4science-studio/`. +- **Debug log path vs cluster:** Agent NDJSON defaults to **`/.cursor/debug-.log`**. Compute nodes inside Apptainer often **cannot** write that path. Unless you set **`DEBUG_AGENT_LOG`** to something bind-mounted and writable, then **copy** the file back to the workspace path, the next chat turn will show **“log not found”** even after a good repro — that is an artifact workflow gap, not necessarily “no crash.” +- **Session id:** the agent debug `sessionId` is reused as the log filename `debug-.log`. + +**What was already instrumented (do not rip out until CK is understood)** + +| Location | Purpose | +|----------|---------| +| `.../bayes-cast/launch/train_edm.py` after `print("model_kwargs", ...)` | **H5:** one NDJSON line — `torch` / `xformers` versions, HIP/CUDA strings, `FusedAttn_option`, `data_type`, `gpu_type`, `ORBIT2_FUSED_ATTN` (rank 0 only). | +| `.../bayes-cast/src/climate_learn/models/hub/components/attention.py` — `_agent_dbg` | **H1/H3:** `Attention` and `CrossAttention` **pre_ck** (shapes, dtypes, `dropout_p`, `N_a`/`N_i` for cross) immediately before `memory_efficient_attention(..., op=MemoryEfficientAttentionCkOp)`; **post_ck** after a successful return. **Rank 0 only** when `dist.is_initialized()` to avoid interleaved NDJSON from 8 ranks. | + +**Hypotheses to resolve from the next log (first evidence wins)** + +1. **H1** — Self-attention CK bad layout/shapes: last line `Attention.pre_ck`, no `Attention.post_ck`. +2. **H2** — `dropout_p` / training vs CK: correlate with fields on last `pre_ck`. +3. **H3** — Cross-attention CK when `N_a ≠ N_i`: last line `CrossAttention.pre_ck`, no `CrossAttention.post_ck`. +4. **H4** — Dtype: `CrossAttention` does `q, k = q.to(x.dtype), k.to(x.dtype)` but **not** `v`; if log shows `v_dtype != q_dtype` before SIGSEGV, add **`v = v.to(x.dtype)`** before CK (then post-fix verification run **with logs still on**). +5. **H5** — Stack skew: use `train_edm.py:post_model_kwargs` line for versions. + +**Operational checklist for the next session** + +1. In sbatch or env: `export DEBUG_AGENT_LOG=/path/on/shared/fs/agent-ck.ndjson` (must be visible inside the container). +2. Run minimal repro: `ORBIT2_DATA_TYPE=bfloat16`, `ORBIT2_FUSED_ATTN=CK`, small batch. +3. After job ends: copy `agent-ck.ndjson` → workspace **`/.cursor/debug-.log`** (or paste last ~50 lines into chat). +4. Until logs prove otherwise, **`ORBIT2_FUSED_ATTN=DEFAULT`** (SDPA) remains the safe default in Studio sbatch for AMD+bf16. + +**Symptom reminder:** CK + bf16 jobs ended with **`Segmentation fault (core dumped)`** on all ranks (no Python traceback). Example slurm text may live under `.../outputs/train/batch-cal-edm/*bf16*.out`. + +### NEW (2026-06-11) — bf16 blocked even with SDPA (`FusedAttn.DEFAULT`): `var_agg` proj GEMM `hipErrorInvalidValue` + +**This is broader than the CK SIGSEGV above:** with the *safe* default **`ORBIT2_FUSED_ATTN=DEFAULT`** (PyTorch SDPA), **bf16 EDM training still crashes** on the **first `training_step`**: + +``` +File "/orbit2/src/climate_learn/models/hub/edm.py", line 316, in aggregate_variables + x = self.var_agg(var_query, x) +File "/orbit2/src/climate_learn/models/hub/components/attention.py", line 280, in forward + x = self.proj(x) # nn.Linear(256, 256) +torch.AcceleratorError: HIP error: invalid argument (hipErrorInvalidValue) +``` + +- **Reproduced across 6 jobs (lux MI355X, `pytorch_rocm7.2.2` SIF + ORBIT-2 overlay):** + - **9228/9229/9230** — bf16 batch 1024/2048/4096 → crash. (Also confirmed restaged ERA5 lifted the old ~1704 sample cap: batch 1024 → `y.shape[0]=1024`; batch ≥2048 → caps at **1704** = total train samples.) + - **9231** — bf16 batch **256** → **same crash** (not large-M dependent). + - **9232** — bf16 batch 1024 + **`TORCH_BLAS_PREFER_HIPBLASLT=0`** (rocBLAS fallback) → **same crash** (not hipBLASLt-specific). + - **9233** — bf16 batch 256 + **`HIP_LAUNCH_BLOCKING=1` / `AMD_SERIALIZE_KERNEL=3`** → **same crash, synchronous** → the failing op is genuinely the **`var_agg` output `proj` `F.linear`**, not an async misattribution. + - **9234** — **float32** batch 256, identical config → **COMPLETED**. ✅ +- **Shape at failure:** `aggregate_variables` (edm.py:309) flattens to `(B·H·L, V, D)`; after var cross-attention `proj` runs on `(B·H·L, 1, 256)` → GEMM **M = batch·648, K=256, N=256** in bf16. fp32 identical GEMM works. +- **Initial (wrong) verdict:** looked like a stack-level bf16 GEMM bug. **Superseded — see RESOLVED below.** + +#### RESOLVED (2026-06-11) — root cause: ROCm **Flash** SDPA 65535 batch-grid cap in `var_agg`; fix: force EFFICIENT/MATH + +**Root cause (isolated with a 1-GPU op-replica probe, no data):** +- The standalone bf16 **`proj` GEMM passes** even at M=663552 → the Linear is *not* the problem (the async HIP error was misattributed to `proj` line 280). +- A **faithful `CrossAttention` replica** reproduces the crash: bf16 **PASS at batch ≤64** (B=batch·648=41472), **FAIL at batch ≥128** (B=82944). The trigger is the **collapsed attention batch dim `B = batch·tokens`**: `aggregate_variables` flattens `(B,History,L)` into the SDPA batch dim, and ROCm **Flash** attention maps that to a HIP grid dimension capped at **65535** (`65535/648 ≈ 101` → matches batch 64 pass / 128 fail). +- **Backend sweep at batch 128/256 (bf16):** `default`(Flash) **FAIL**; **`math` PASS**, **`efficient` (EFFICIENT_ATTENTION) PASS**, **`eager` PASS**. fp32 "works" only because it dispatches to math, not Flash. +- **Independently corroborated by the upstream PyTorch SDPA benchmark** ([`examples/upstream_pytorch_sdpa_benchmark.py`](../../examples/upstream_pytorch_sdpa_benchmark.py) `--orbit-varagg`, which replicates the var_agg shape `batch_size=batch*648, q_len=1, kv_len=N_i`): + + | backend | batch 64 (B=41472) | batch 128 (B=82944) | batch 256 (B=165888) | + |---|---|---|---| + | default | PASS | **FAIL** | **FAIL** | + | flash | PASS | **FAIL** | **FAIL** | + | efficient | PASS | PASS | PASS | + | math | PASS | PASS | PASS | + + Confirms `default ≡ flash` for these shapes and the Flash 65535 batch-grid cap (`64·648=41472` ok, `128·648=82944` fails). Run: `python3 upstream_pytorch_sdpa_benchmark.py --orbit-varagg --orbit-varagg-batches 64,128,256 --orbit-sweep-backends default,flash,efficient,math` (1 GPU; `pip install --target=/tmp/... tabulate tqdm` first per §"Apptainer venv"). + +**Fix (compute-side clone `src/climate_learn/models/hub/components/attention.py`):** wrap **all four** SDPA call sites — **both `Attention` (self) and `CrossAttention`**, DEFAULT + CK branches — in `with sdpa_kernel([SDPBackend.EFFICIENT_ATTENTION, SDPBackend.MATH])` to avoid Flash. MATH is fallback-only (never selected for supported shapes). + + *Why global, not just CrossAttention:* the [`--orbit-selfattn`](../../examples/upstream_pytorch_sdpa_benchmark.py) benchmark (q=kv=648, heads=8, bf16) shows **EFFICIENT within ±3–5% of Flash** at batch 64/256/1024 (efficient faster at 64; flash faster fwd by ~1–2% at large batch; backward a wash), so unifying on EFFICIENT removes the 65535 grid foot-gun at ~noise-level perf cost. **MATH is 10–15× slower** (≈15 vs ≈290 TFLOPS) so it must stay fallback-only. + + | self-attn fwd/bwd µs | batch 64 | batch 256 | batch 1024 | + |---|---|---|---| + | flash | 123 / 449 | 408 / 1516 | 1515 / 5940 | + | efficient | 111 / 408 | 419 / 1514 | 1534 / 5909 | + | math | — | — | 29220 / 30940 | + +**Validated end-to-end:** job **9238** = the exact bf16 batch-256 config that crashed as **9231** now **COMPLETED (exit 0)**, `y.shape=[256,3,108,216]`, decreasing loss (`batch_idx 0→0.971, 1→0.963, 2→0.930, …`), full 6 epochs. Global fix (both `Attention`+`CrossAttention` → EFFICIENT) re-validated by job **9257** (bf16 **b1024**, COMPLETED, no errors). + +**CK-flash flags are inert on this wheel:** setting `TORCH_ROCM_FA_PREFER_CK=1` + `USE_CK_FLASH_ATTENTION=1` does **not** change anything — the flash/efficient kernel stays AOTriton `attn_fwd`, flash still crashes at B>65535 (var_agg batch 128), and self-attn perf is unchanged. `USE_CK_FLASH_ATTENTION` is a **build-time** flag and this `pytorch_rocm7.2.2` wheel wasn't built with CK flash, so `TORCH_ROCM_FA_PREFER_CK` has nothing to prefer. Getting CK flash would require rebuilding PyTorch with `USE_CK_FLASH_ATTENTION=1` — not worth it given EFFICIENT (AOTriton) already matches Flash perf and has no grid cap. + +**Caveats / follow-ups:** +- **`run_fom_extractor.py` / `parse_training_log.py` gap:** completed EDM runs log per-step loss as `epoch: N batch_idx M loss tensor(L)` and there's **no `Batch N: … seconds`** line in this build → the parser reports `epoch_records=0`, `final_loss=null`, and no `steady_batch_time_s`. **Throughput/loss FOMs need a parser update** to read the `epoch:/batch_idx/loss` format (and to time steps from omnistat or added timing). +- **HBM saturation vs data cap:** dataset has **~1704 train samples** (batch ≥1704 caps). bf16 at the cap will use roughly half the fp32 footprint (fp32@1704 ≈ 183 GiB ≈ 64%), so **bf16 alone won't reach 85–90% at 1704 samples** — stage more ERA5 ([`STAGING_ERA5_FOR_HBM.md`](../perf-optimizer-loop/STAGING_ERA5_FOR_HBM.md)) to push higher. + +### xFormers CK probe (MI355, `pytorch_rocm7.2.2` SIF + ORBIT-2 overlay) + +**`python3 -m xformers.info`:** `memory_efficient_attention.ckF`, `.ckB`, `.ck_splitKF` report **available** (Cutlass/FA paths unavailable on AMD, expected). + +**PyTorch wheel index (`https://download.pytorch.org/whl/rocm7.2`):** `pip index versions xformers` listed only **`0.0.34`** and **`0.0.35`** — no newer wheel than overlay **0.0.35** without changing ROCm/PyTorch image. + +**A one-GPU CK probe (bf16) SIGSEGVs (exit 139) on every case** — `small`, `asym`, `long_p0`, `long_p01` — after printing the case header, i.e. **`memory_efficient_attention` + `MemoryEfficientAttentionCkOp` + bf16** fails even on tiny tensors. PyTorch SDPA on the same shapes **succeeds**. + +**Implication:** do not treat **`xformers.info` “ckF available”** as safe for training MEA on this stack; keep **`ORBIT2_FUSED_ATTN=DEFAULT`** (SDPA) unless a new wheel/image is validated. + +### PyTorch SDPA benchmark (upstream `sdpa.py`) + +Studio copy: [`examples/upstream_pytorch_sdpa_benchmark.py`](../../examples/upstream_pytorch_sdpa_benchmark.py) — exercises **`torch.nn.functional.scaled_dot_product_attention`** under `sdpa_kernel(SDPBackend)` (same idea as [PyTorch `benchmarks/transformer/sdpa.py`](https://github.com/pytorch/pytorch/blob/main/benchmarks/transformer/sdpa.py)). + +**Naming (corrected 2026-06-11):** the old CLI label **`ck`** for `--orbit-sweep-backends` was **removed** — it mapped to `SDPBackend.EFFICIENT_ATTENTION`, which on this ROCm build is **AOTriton**, *not* Composable Kernel. **Use `efficient`.** Verified by profiling: both `flash` and `efficient` dispatch to the same AOTriton kernel **`attn_fwd`**; `math` is the unfused reference (`Cijk_...` Tensile GEMMs + `softmax_warp_forward`). Real **CK** exists only in the xFormers path (`--orbit-include-xformers-ck` / `--orbit-xformers-modes ck` → `MemoryEfficientAttentionCkOp`, the one that SIGSEGVs). Flash vs efficient differ only in the **grid-launch wrapper** (flash maps the batch dim to a 65535-capped grid dim; efficient does not) — same `attn_fwd` math kernel. + +**xFormers CK on generic shapes (no ORBIT code):** pass **`--orbit-include-xformers-ck`** with **`--orbit-micro`**. Each case runs in a **fresh subprocess** so a child **SIGSEGV** is reported as `KILLED_BY_SIGNAL` / `SIGSEGV` in the summary table while PyTorch SDPA results (already printed) stay intact. Tune with **`--orbit-xformers-modes ck,dispatch`**, **`--orbit-xformers-dropout 0.0,0.1`**, **`--orbit-xformers-timeout SEC`**. + +**Validated (MI355X, SLURM job **9207**, `pytorch_rocm7.2.2` SIF + ORBIT-2 overlay, 2026-06-11):** **`--orbit-micro`** SDPA sweeps (**`default`**, **`math`**, **`efficient`**, CLI label **`ck`** → `SDPBackend.EFFICIENT_ATTENTION`, **`flash`**) all **completed** for bf16 shapes 128×128 and 648×648. **`--orbit-include-xformers-ck`** — every row (**`micro_128_sym`**, **`micro_648_sym`**, **`asym_1_vs_6`** × **`ck`** and **`dispatch`**, `dropout_p=0`) reported **`KILLED_BY_SIGNAL` / `SIGSEGV`** (standalone tensors; no Bayes-CAST). Example CSV dir: **`$AI4S_SHARED_DIR/models/ORBIT-2/perf-runs/sdpa-ck-microbench-1781204729/`**. + +**Apptainer venv:** `pip install --user` fails (“User site-packages are not visible in this virtualenv”). Use **`pip install -q --target=/tmp/orbit2_sdpa_pip tabulate tqdm`** and **`export PYTHONPATH=/tmp/orbit2_sdpa_pip:$PYTHONPATH`** before running the script. + +**One GPU, bind-mount Studio `examples/`** (adjust `AI4S_*` paths): + +```bash +export AI4S_SHARED_DIR=/path/to/shared +STUDIO_ORBIT2_EXAMPLES=/path/to/ai4science-studio/earth_science/models/ORBIT-2/examples +srun -N1 -n1 --gres=gpu:1 apptainer exec --rocm \ + --overlay "${ORBIT2_OVERLAY:?}":ro \ + --bind "$STUDIO_ORBIT2_EXAMPLES":/sdpa-bench \ + "${ORBIT2_SIF:?}" bash -lc ' + set -e + PIP_LOCAL=/tmp/orbit2_sdpa_pip + mkdir -p "$PIP_LOCAL" + python3 -c "import tabulate, tqdm" 2>/dev/null || pip install -q --target="$PIP_LOCAL" tabulate tqdm + export PYTHONPATH="$PIP_LOCAL:${PYTHONPATH:-}" + export ORBIT2_SDPA_CSV_DIR="${ORBIT2_SDPA_CSV_DIR:-/tmp/orbit2_sdpa_benchmark_results}" + mkdir -p "$ORBIT2_SDPA_CSV_DIR" + cd /sdpa-bench + python3 upstream_pytorch_sdpa_benchmark.py --orbit-micro \ + --orbit-sweep-backends default,math,efficient,ck,flash \ + --orbit-include-xformers-ck \ + --orbit-xformers-modes ck,dispatch \ + --orbit-xformers-dropout 0.0,0.1 + ' +``` + +Expect **`math`** and often **`default`** / **`efficient`** to complete on ROCm; **`flash`** may skip or error if Flash is unavailable — compare timings, not only success/fail. CSV timestamps land in **`ORBIT2_SDPA_CSV_DIR`**. + ## Next work (after scaling baseline) -- Port `perf-optimizer-loop/` + `lever_catalog.yaml` from the HydraGNN reference recipe. +- **Run the sysopt loop:** [`recipes/perf-optimizer-loop/README.md`](../recipes/perf-optimizer-loop/README.md) — `examples/run_optimizer_loop.sh`, [`lever_catalog.yaml`](../recipes/perf-optimizer-loop/lever_catalog.yaml), staging + bf16 sweep ([`STAGING_ERA5_FOR_HBM.md`](../recipes/perf-optimizer-loop/STAGING_ERA5_FOR_HBM.md), [`examples/sweep_orbit2_batch_bf16_amd.sh`](../../examples/sweep_orbit2_batch_bf16_amd.sh)). - Rebuild overlay; re-test `bfloat16` + CK attn when `xformers.ops` is available. - Stage **2.5_arcmin** PRISM targets for scientifically meaningful loss (not same-dir timing). +- **One-node GPU baseline (pre-sysopt):** [one-node-gpu-baseline.md](one-node-gpu-baseline.md) — batch sweep, Omnistat + TraceLens evidence, `report_orbit2_gpu_baseline.py` for `baseline_report.md`. + ## Key paths (parameterize with `AI4S_SHARED_DIR`) - Data: `$AI4S_SHARED_DIR/models/ORBIT-2/data/superres/prism/10.0_arcmin` (PRISM same-dir) or `.../era5/1.0_deg` (ERA5 same-dir sanity template). - Overlay: `$AI4S_SHARED_DIR/models/ORBIT-2/overlays/orbit2-overlay.img` - Perf runs: `$AI4S_SHARED_DIR/models/ORBIT-2/perf-runs//` - Training logs (recommended): `$AI4S_SHARED_DIR/models/ORBIT-2/outputs/train/logs/` -- Omnistat / TraceLens tools: set **`OMNIHUB_TOOLS_DIR`** to your site-local checkout (see root `.cluster-config.yaml` key `omnihub.tools_dir` if your site maintains one). +- Omnistat / TraceLens tools: set **`PERF_TOOLS_DIR`** to your site-local checkout (see root `.cursor/skills/ai4science-perf-analysis/SKILL.md` and `.cluster-config.yaml` key `perf_tools.dir` if your site maintains one). + +## Operational source tree + +- Set **`ORBIT2_ROOT`** to your clone (public **ORBIT-2** or institutional **Bayes-CAST**). Scripts bind-mount it at **`/orbit2`**. +- **`launch_diffusion.sh`:** if `$ORBIT2_ROOT/launch_diffusion.sh` or `examples/launch_diffusion.sh` exists, ranks **`exec bash …/launch_diffusion.sh /config/config.yaml`** after optional `ORBIT2_RANK_PRE_TRAIN_HOOK` (via `orbit2_rank_hook_runner.py`). Override path with **`ORBIT2_LAUNCH_SCRIPT`** (must live under `ORBIT2_ROOT`). +- **`manifest.json`** (perf runs): includes `git_sha`, `git_branch`, `git_remote_origin`, `rendered_config`, `parallelism.fsdp` / `parallelism.simple_ddp`, `global_batch_size`, `runtime_seconds`, `orbit2_batch_size`, `total_ranks`, `max_epochs`, `data_type`, and `workload` (`bayes_edm` vs `intermediate_downscaling`). diff --git a/earth_science/models/ORBIT-2/recipes/perf-analysis/README.md b/earth_science/models/ORBIT-2/recipes/perf-analysis/README.md index 5452d75..a995452 100644 --- a/earth_science/models/ORBIT-2/recipes/perf-analysis/README.md +++ b/earth_science/models/ORBIT-2/recipes/perf-analysis/README.md @@ -1,39 +1,47 @@ # ORBIT-2 Performance Analysis with AMD AI Agents -Multi-subagent workflow for 2-node ORBIT-2 training (`intermediate_downscaling.py`) on AMD MI355X, with Omnistat telemetry and optional rank-0 PyTorch traces. +Multi-subagent workflow for ORBIT-2 / Bayes-CAST training on AMD MI355X, with Omnistat telemetry and optional rank-0 PyTorch traces. > **Audience:** performance engineers. Output is a diagnosis of the run, not a scientific ORBIT-2 claim. -> **Data mode:** 10.0_arcmin PRISM same-dir (`superres_mag: 1`) — timing/scaling only; loss is not meaningful without 2.5_arcmin targets. +> **Data mode:** Perf defaults **Bayes-CAST** **`edm_8m_era5_1x8.yaml`** + staged **ERA5 1.0°** (`ORBIT2_DATA_ROOT`). Lux **`interm_8m_lux_era5.yaml`** remains for public ORBIT-2 same-dir `res_slimvit`. Loss is not meaningful without true downscaling targets. -## Quick start +## Quick start (1 node × 8 GPUs, Bayes-CAST EDM template) ```bash export AI4S_SHARED_DIR=/path/to/shared -export OMNIHUB_TOOLS_DIR=/path/to/omnihub/tools -export ORBIT2_DATA_ROOT=$AI4S_SHARED_DIR/models/ORBIT-2/data/superres/prism/10.0_arcmin +export PERF_TOOLS_DIR=/path/to/perf-tools +# ORBIT2_ROOT defaults to .../code/bayes-cast when that directory exists +export ORBIT2_DATA_ROOT=$AI4S_SHARED_DIR/models/ORBIT-2/data/superres/era5/1.0_deg +# export ORBIT2_CONFIG_TEMPLATE=edm_8m_era5_1x8.yaml # default sbatch earth_science/models/ORBIT-2/examples/sbatch_train_perf_amd.sh ``` +Multi-node perf: `sbatch --nodes=2 ... earth_science/models/ORBIT-2/examples/sbatch_train_perf_amd.sh` (requires `.cluster-config.yaml` RCCL keys). Default `#SBATCH` is **1 node**. + +**GPU saturation + model/batch baseline (before sysopt):** [one-node-gpu-baseline.md](one-node-gpu-baseline.md), [BASELINE_LOCKIN.md](BASELINE_LOCKIN.md) (VRAM-first batch sweep order), and `examples/report_orbit2_gpu_baseline.py`. + After the job completes, drive analyst + verifier + synthesizer subagents per [`.cursor/skills/ai4science-perf-analysis/SKILL.md`](../../../../.cursor/skills/ai4science-perf-analysis/SKILL.md). ## Artifacts ``` $AI4S_SHARED_DIR/models/ORBIT-2/perf-runs// -├── manifest.json +├── manifest.json # includes git_sha, parallelism, rendered_config ├── omnistat-db/ ├── omnistat.config ├── traces/orbit2-epoch*-rank0.pt.trace.json ├── orbit2-train-.out -├── foms.json # from run_fom_extractor.py -├── tracelens/ # analyst outputs -└── combined_report.md # synthesizer output +├── foms.json # from run_fom_extractor.py +├── baseline_report.md # from report_orbit2_gpu_baseline.py (optional) +├── baseline_report.json +├── tracelens/ # analyst outputs +└── combined_report.md # synthesizer output ``` ## Primary FOM -**`batch_time_s`** — mean steady-state batch wall time from rank-0 log lines `Batch N: X seconds` (see `examples/parse_training_log.py`). +**`steady_batch_time_s`** — see [HANDOFF.md](HANDOFF.md) and `examples/parse_training_log.py` (epoch ≥ 2, skip warmup batches). Legacy **`batch_time_s`** label in older notes maps to the same parser options. ## Agent prompts @@ -42,4 +50,4 @@ $AI4S_SHARED_DIR/models/ORBIT-2/perf-runs// - [agents/omnistat_analyst.md](agents/omnistat_analyst.md) - [agents/synthesizer.md](agents/synthesizer.md) -See [HANDOFF.md](HANDOFF.md) for validated landmines and current cluster state. +See [HANDOFF.md](HANDOFF.md) for validated landmines, Bayes-CAST / `launch_diffusion.sh`, and parallelism defaults. diff --git a/earth_science/models/ORBIT-2/recipes/perf-analysis/agents/launcher.md b/earth_science/models/ORBIT-2/recipes/perf-analysis/agents/launcher.md index fa423cd..107863c 100644 --- a/earth_science/models/ORBIT-2/recipes/perf-analysis/agents/launcher.md +++ b/earth_science/models/ORBIT-2/recipes/perf-analysis/agents/launcher.md @@ -10,8 +10,8 @@ Submits a 2-node ORBIT-2 training (AMD Instinct MI355X) with PyTorch profiling a ## Outputs -- `${OMNIHUB_TOOLS_DIR}/omnihub-inspect/` — Python venv with omnistat (PR #271 + main merged) and TraceLens. -- `${OMNIHUB_TOOLS_DIR}/victoriametrics/victoria-metrics-prod` — VictoriaMetrics binary. +- `${PERF_TOOLS_DIR}/perf-inspect/` — Python venv with omnistat (PR #271 + main merged) and TraceLens. +- `${PERF_TOOLS_DIR}/victoriametrics/victoria-metrics-prod` — VictoriaMetrics binary. - `$AI4S_SHARED_DIR/models/ORBIT-2/perf-runs/omnistat.config` — Omnistat user-mode config. - `$AI4S_SHARED_DIR/models/ORBIT-2/perf-runs//manifest.json` — manifest schema below. - `$AI4S_SHARED_DIR/models/ORBIT-2/perf-runs//orbit2-train-.out` — symlinked from output dir. @@ -57,8 +57,8 @@ Submits a 2-node ORBIT-2 training (AMD Instinct MI355X) with PyTorch profiling a ```bash set -euo pipefail -# Shared tool location — read from .cluster-config.yaml omnihub.tools_dir. -TOOLS="${OMNIHUB_TOOLS_DIR:?set OMNIHUB_TOOLS_DIR to a persistent checkout (e.g. under \$AI4S_SHARED_DIR/tools)}" +# Shared tool location — read from .cluster-config.yaml perf_tools.dir. +TOOLS="${PERF_TOOLS_DIR:?set PERF_TOOLS_DIR to a persistent checkout (e.g. under \$AI4S_SHARED_DIR/tools)}" mkdir -p "$TOOLS" # 1a. Omnistat: jorda/skills branch with origin/main merged in @@ -77,7 +77,7 @@ fi OMNISTAT_COMMIT=$(git rev-parse HEAD) # 1b. Venv (py3.12 to match cluster python) -VENV=${OMNIHUB_TOOLS_DIR}/omnihub-inspect +VENV=${PERF_TOOLS_DIR}/perf-inspect if [[ ! -d $VENV ]]; then python3 -m venv "$VENV" fi @@ -123,7 +123,7 @@ if [[ "${OMNISTAT_KERNEL_TRACE:-0}" == "1" && ! -f "$TRACE_LIB" ]]; then fi ``` -The sbatch wrapper expects the `.so` at `${OMNIHUB_TOOLS_DIR}/omnistat-src/build-trace/libomnistat_trace.so` by default; override with `OMNISTAT_TRACE_LIB=/path/to/lib.so`. See `ai4science-studio` SKILL §12 for when to enable kernel tracing vs the default device-counting collector. +The sbatch wrapper expects the `.so` at `${PERF_TOOLS_DIR}/omnistat-src/build-trace/libomnistat_trace.so` by default; override with `OMNISTAT_TRACE_LIB=/path/to/lib.so`. See `ai4science-studio` SKILL §12 for when to enable kernel tracing vs the default device-counting collector. ### 2. Author the Omnistat user-mode config (once, in repo `perf-runs/`) diff --git a/earth_science/models/ORBIT-2/recipes/perf-analysis/agents/omnistat_analyst.md b/earth_science/models/ORBIT-2/recipes/perf-analysis/agents/omnistat_analyst.md index 56e9fef..5ff92cf 100644 --- a/earth_science/models/ORBIT-2/recipes/perf-analysis/agents/omnistat_analyst.md +++ b/earth_science/models/ORBIT-2/recipes/perf-analysis/agents/omnistat_analyst.md @@ -6,8 +6,8 @@ Drive `omnistat-inspect` (PR #271) through the analyze-job phases on the user-mo - `/manifest.json` - The Omnistat DB at `manifest.omnistat_db_path`. -- The `omnihub-inspect` venv at `${OMNIHUB_TOOLS_DIR}/omnihub-inspect/`. -- The VictoriaMetrics binary at `${OMNIHUB_TOOLS_DIR}/victoriametrics/victoria-metrics-prod`. +- The `perf-inspect` venv at `${PERF_TOOLS_DIR}/perf-inspect/`. +- The VictoriaMetrics binary at `${PERF_TOOLS_DIR}/victoriametrics/victoria-metrics-prod`. ## Outputs @@ -36,7 +36,7 @@ ss -lnt "sport = :$PORT" 2>/dev/null | grep -q LISTEN && PORT=8429 VMLOG=$PERF_RUN/omnistat/vm.log mkdir -p "$PERF_RUN/omnistat/inspect_outputs" -nohup ${OMNIHUB_TOOLS_DIR}/victoriametrics/victoria-metrics-prod \ +nohup ${PERF_TOOLS_DIR}/victoriametrics/victoria-metrics-prod \ -storageDataPath="$DB" \ -httpListenAddr=127.0.0.1:$PORT \ -retentionPeriod=100y \ @@ -62,7 +62,7 @@ done ### 2. Run analyze-job phases (per the SKILL) ```bash -. ${OMNIHUB_TOOLS_DIR}/omnihub-inspect/bin/activate +. ${PERF_TOOLS_DIR}/perf-inspect/bin/activate SCRATCH=$PERF_RUN/omnistat/scratch mkdir -p "$SCRATCH" JOBID=$(jq -r .jobid ) diff --git a/earth_science/models/ORBIT-2/recipes/perf-analysis/agents/orchestrator_gemm_analysis.md b/earth_science/models/ORBIT-2/recipes/perf-analysis/agents/orchestrator_gemm_analysis.md new file mode 100644 index 0000000..c4d6d0d --- /dev/null +++ b/earth_science/models/ORBIT-2/recipes/perf-analysis/agents/orchestrator_gemm_analysis.md @@ -0,0 +1,116 @@ +# orchestrator_gemm_analysis — "where does ORBIT-2 GEMM time go?" (1-node vs 2-node) + +You are the **GEMM-time bottleneck orchestrator** for ORBIT-2 (Bayes-CAST EDM, 8M) on MI355X / +ROCm 7.2.2 / PyTorch 2.10. You run unattended via the Claude Code CLI in tmux. Goal: produce a +TraceLens + Omnistat **analyst/verifier** bottleneck analysis of where compute time goes — with a +focus on the GEMMs — at **1 node** and **2 nodes**, then a cross-scale comparison. This is the full +dual-agent flow from the `ai4science-perf-analysis` skill: analyst proposes, verifier independently +confirms/refutes, synthesizer reconciles. + +Context: TunableOp was already ruled out (no uplift + 1-node NaN). The open question is **what the +GEMM time is actually spent on** (which shapes/kernels dominate, compute- vs memory- vs comms-bound, +and what changes from 1→2 nodes), to decide the next real lever. + +## Fixed context (passed in the user prompt) +- `REPO_ROOT`, `AI4S_SHARED_DIR`, `PERF_TOOLS_DIR` (`perf_tools.dir` in `.cluster-config.yaml`). +- SLURM: partition and account from `.cluster-config.yaml` (`slurm.partition`, `slurm.account`). + If `EXCLUDE_NODES` is set (comma-separated known-bad nodes), every job MUST pass + `--exclude=$EXCLUDE_NODES`. +- `ANALYSIS_DIR` — write STATUS.txt + the final `GEMM_TIME_REPORT.md` here. +- sbatch: `earth_science/models/ORBIT-2/examples/sbatch_train_perf_amd.sh`. +- Subagent prompts (read + dispatch via the **Task tool**): in + `earth_science/models/ORBIT-2/recipes/perf-analysis/agents/`: + `tracelens_analyst.md`, `tracelens_verifier.md`, `omnistat_analyst.md`, `omnistat_verifier.md`, + `synthesizer.md`. Each reads `/manifest.json` and writes under `/`. +- Per-job dir: `$AI4S_SHARED_DIR/models/ORBIT-2/perf-runs//`. + +## Locked config (identical to the uplift study so results are comparable) +bf16, SDPA DEFAULT, `ORBIT2_BATCH_SIZE=4096`, `ORBIT2_MAX_EPOCH=6`, +`ORBIT2_DATA_ROOT=$AI4S_SHARED_DIR/models/ORBIT-2/data/superres/era5/1.0_deg`, +`ORBIT2_ERA5_SPATIAL_RES=111`, `ORBIT2_ROOT=$AI4S_SHARED_DIR/models/ORBIT-2/code/bayes-cast`, +`ORBIT2_CONFIG_TEMPLATE=edm_8m_era5_1x8.yaml`, `TORCH_NCCL_HIGH_PRIORITY=1`, `GPU_MAX_HW_QUEUES=2`. +**Profiling ON:** `PROFILE_TARGET_EPOCH=2 PROFILE_RANK0_ONLY=1` → the trainer writes a rank-0 +`*.pt.trace.json` for epoch 2 into `/traces/orbit2-epoch2-rank0/` (TraceLens input). +**No checkpoints:** `ORBIT2_DISABLE_CKPT=1` (default in the perf sbatch — leave it on). + +## Hard constraints +1. **One SLURM job at a time.** `squeue -u $USER -h` must be empty before each `sbatch`. +2. **Exclude any `EXCLUDE_NODES`** (see above). For `--nodes=2`, also rely on the sbatch auto-reading + `.cluster-config.yaml` for NCCL/IB env (`NCCL_IB_HCA`, ANP plugin, libionic) — do not override. +3. **State on disk** (`ANALYSIS_DIR/STATUS.txt`, per-job artifacts), resumable. On resume, skip a + scale whose `/combined_report.md` already exists. +4. **Disk:** /shared has ample room; do NOT delete `traces/` (the analysts need them). Before each + submit, `df -P "$AI4S_SHARED_DIR"`; if ≥ 93% used, append `DISK_FULL`, stop, write the report + with whatever completed. +5. **Subagent contract:** each Task subagent reads only its `## Inputs` and writes only its + `## Outputs`; it ends with `STATUS=ok|partial|fail`. A verifier that refutes an analyst claim + keeps it in the report with `verdict=refuted` — do not silently drop it. + +## Procedure — run for SCALE in [1, 2] (1-node first) +### Step A — submit the profiled baseline +1. `squeue` empty? Build env and submit (partition/account from `.cluster-config.yaml`; + add `--exclude=$EXCLUDE_NODES` only if set): + `sbatch --parsable --partition= --account= --nodes= + [--exclude=$EXCLUDE_NODES] + --job-name=o2-gemm-n + --export=ALL,AI4S_SHARED_DIR=...,PERF_TOOLS_DIR=...,ORBIT2_DATA_ROOT=...,ORBIT2_ROOT=..., + ORBIT2_CONFIG_TEMPLATE=edm_8m_era5_1x8.yaml,ORBIT2_ERA5_SPATIAL_RES=111,ORBIT2_BATCH_SIZE=4096, + ORBIT2_MAX_EPOCH=6,TORCH_NCCL_HIGH_PRIORITY=1,GPU_MAX_HW_QUEUES=2, + PROFILE_TARGET_EPOCH=2,PROFILE_RANK0_ONLY=1 `. + Append `SUBMIT scale= jobid=`. +2. Poll `squeue -h -j -o "%T"` every 30 s. While state is `PENDING`, keep waiting (a-nodes are + busy with other users — queue waits of several hours are normal; cap PENDING at 12 h). Once + `RUNNING`, apply a **75-min run ceiling**; on exceed `scancel`, mark broken, continue to the next + scale. Job is done when it leaves the queue — then read `/manifest.json` `state`. +3. On exit-42 (mount probe), re-submit once adding the failed node(s) to `--exclude`. + +### Step B — make the manifest analyst-ready +The sbatch writes `/manifest.json` with `job_id`/`omnistat_db`/`trace_dir`. The analyst +subagents expect `jobid`, `omnistat_db_path`, `trace_paths[]`, `training_log_path`, `perf_run_dir`, +`nodes`, `ranks`. With Python, load the manifest and ADD those keys (do not remove existing ones): +- `jobid` = job_id; `perf_run_dir` = job_dir; `omnistat_db_path` = omnistat_db; + `training_log_path` = slurm_log; `nodes` = ; `ranks` = total_ranks; + `trace_paths` = sorted glob of `/traces/**/*.pt.trace.json`. +Write it back. If `trace_paths` is empty, append `WARN scale= no_trace` (tracelens_analyst +will no-op; the omnistat side still produces compute/memory-bound evidence). + +### Step C — analysts (parallel Task calls in ONE message) +Dispatch two subagents, each given `perf_run_dir` and told to follow its prompt file: +- `tracelens_analyst` (follow `agents/tracelens_analyst.md`) → `/tracelens/claims.json`. + **Emphasis for this study:** rank kernels by total device time; isolate the **GEMM/Cijk/hipBLASLt/ + rocBLAS** kernels; report top GEMM shapes (M,N,K) by time, their share of step time, and + fwd-vs-bwd split. Note exposed (non-overlapped) NCCL/RCCL collective time. +- `omnistat_analyst` (follow `agents/omnistat_analyst.md`) → `/omnistat/claims.json`. + **Emphasis:** bf16 MFMA TFLOPS (achieved vs MI355X peak), HBM read/write GB/s vs peak, and the + compute-bound-vs-memory-bound verdict; XGMI/scale-out traffic at 2 nodes. +Wait for both to return. + +### Step D — verifiers (parallel Task calls in ONE message) +- `tracelens_verifier` (follow `agents/tracelens_verifier.md`) → `/tracelens/verified_claims.json`. +- `omnistat_verifier` (follow `agents/omnistat_verifier.md`) → `/omnistat/verified_claims.json`. +Verifiers may use at most one `srun -N1 --time=00:05:00` probe (honoring `EXCLUDE_NODES`). Wait for both. + +### Step E — synthesize per-scale +Dispatch `synthesizer` (follow `agents/synthesizer.md`) with both `verified_claims.json` → +`/combined_report.md`. Append `SCALE_DONE scale= jobid= perf_run=`. + +## Final — cross-scale report `ANALYSIS_DIR/GEMM_TIME_REPORT.md` +Synthesize the 1-node and 2-node `combined_report.md` into one report covering: +- **Where the time goes** at each scale: top GEMM shapes/kernels by device-time share, fwd vs bwd, + and the non-GEMM remainder (attention/SDPA, elementwise, comms, dataloader/exposed gaps). +- **Compute vs memory bound:** achieved bf16 MFMA TFLOPS and HBM BW vs MI355X peak (from Omnistat), + cross-checked against TraceLens kernel mix. State the verdict explicitly. +- **1→2 node delta:** what changes — exposed RCCL collective time, per-rank batch halving effect, + scaling efficiency — and whether GEMM time per step is invariant (expected) while comms grows. +- **Analyst vs verifier:** call out any **refuted** claims and where TraceLens and Omnistat + **disagree** (e.g. kernel-time-based vs counter-based compute-bound verdicts) and how you + reconciled them. This disagreement surface is a key deliverable. +- **Actionable levers** implied by the breakdown (e.g. if SDPA/attention is a large slice → flash/CK + attention; if exposed comms at 2 nodes → overlap/bucketing; if a specific GEMM shape dominates → + targeted kernel work). Rank by expected payoff. +Append `LOOP_COMPLETE 1node_jobid=<> 2node_jobid=<>` to STATUS.txt and print it as the final line. + +## Robustness +- Retry transient `sbatch`/`srun`/Task failures up to 3× with backoff. If a subagent fails, record + it and continue — a partial report (e.g. omnistat-only if no trace) is still valuable. +- Keep `STATUS.txt` append-only: ` EVENT key=val...` per line. diff --git a/earth_science/models/ORBIT-2/recipes/perf-analysis/agents/tracelens_analyst.md b/earth_science/models/ORBIT-2/recipes/perf-analysis/agents/tracelens_analyst.md index e37f2c5..edda6c3 100644 --- a/earth_science/models/ORBIT-2/recipes/perf-analysis/agents/tracelens_analyst.md +++ b/earth_science/models/ORBIT-2/recipes/perf-analysis/agents/tracelens_analyst.md @@ -39,7 +39,7 @@ Every entry: ### 1. Sanity check ```bash -. ${OMNIHUB_TOOLS_DIR}/omnihub-inspect/bin/activate +. ${PERF_TOOLS_DIR}/perf-inspect/bin/activate PERF_RUN=$(jq -r .perf_run_dir ) TRACE=$(jq -r '.trace_paths[0] // empty' ) [[ -z "$TRACE" ]] && { echo "STATUS=partial; reason=no_trace_in_manifest"; exit 0; } diff --git a/earth_science/models/ORBIT-2/recipes/perf-analysis/omnistat.config.template b/earth_science/models/ORBIT-2/recipes/perf-analysis/omnistat.config.template index 80817de..6bbb78c 100644 --- a/earth_science/models/ORBIT-2/recipes/perf-analysis/omnistat.config.template +++ b/earth_science/models/ORBIT-2/recipes/perf-analysis/omnistat.config.template @@ -60,7 +60,10 @@ step_detection_file = @JOB_DIR@/omni_rmsjobinfo_step [omnistat.collectors.host] enable_proc_io_stats = True -# Rocprofiler — HBM bandwidth + FP64 VALU/MFMA counters in one block. +# Rocprofiler — HBM read traffic + BF16 MFMA (matches ORBIT-2 bf16 EDM training). +# FP64 profile `hbm_flops_f64` remains available if you set: +# OMNISTAT_ROCPROF_PROFILE=hbm_flops_f64 +# in the environment before rendering this file (sbatch can sed-replace profile name). # Counter group selection follows the user-validated pattern (one profile, # one sub-section, all counters listed together; rocprofiler-sdk # time-multiplexes hardware counter sets internally as needed). @@ -78,9 +81,20 @@ enable_proc_io_stats = True # ) / wall_seconds / 1e9 # HBM_GB/s = (FETCH_SIZE + WRITE_SIZE) / wall_seconds / 1e9 [omnistat.collectors.rocprofiler] -profile = hbm_flops_f64 +# Default profile for ORBIT-2 bf16: BF16 MFMA MOPS + FETCH_SIZE (2 counters, gfx950-safe). +# Citation: AMD MI300/MI200 counter tables — `SQ_INSTS_VALU_MFMA_MOPS_BF16` is in units of 512 FLOPs: +# https://rocm.docs.amd.com/en/develop/conceptual/gpu-arch/mi300-mi200-performance-counters.html +profile = hbm_flops_bf16 + +[omnistat.collectors.rocprofiler.hbm_flops_bf16] +# MI355X (gfx950) — keep one rocprofiler profile under constant sampling (SKILL §11). +# TFLOP/s (BF16 MFMA only): rate(SQ_INSTS_VALU_MFMA_MOPS_BF16) * 512 / 1e12 +# HBM read GB/s: rate(FETCH_SIZE) * 1024 / 1e9 (FETCH_SIZE units per omnistat/HydraGNN lesson; verify on your omnistat build) +sampling_mode = constant +counters = ["FETCH_SIZE", "SQ_INSTS_VALU_MFMA_MOPS_BF16"] [omnistat.collectors.rocprofiler.hbm_flops_f64] +# Legacy fp64-rich profile (HydraGNN-style). Optional via OMNISTAT_ROCPROF_PROFILE=hbm_flops_f64. # MI355X (gfx950) rocprofiler counter limits — see ai4science-studio SKILL §10. # # Constraints surfaced on MI355X (gfx950) / ROCm 7.2.2: @@ -120,7 +134,7 @@ system_name = mi355x #-- [omnistat.usermode] -victoria_binary = @OMNIHUB_TOOLS_DIR@/victoriametrics/victoria-metrics-prod +victoria_binary = @PERF_TOOLS_DIR@/victoriametrics/victoria-metrics-prod victoria_datadir = @JOB_DIR@/omnistat-db victoria_logfile = vic_server.log diff --git a/earth_science/models/ORBIT-2/recipes/perf-analysis/one-node-gpu-baseline.md b/earth_science/models/ORBIT-2/recipes/perf-analysis/one-node-gpu-baseline.md new file mode 100644 index 0000000..a9ff0d6 --- /dev/null +++ b/earth_science/models/ORBIT-2/recipes/perf-analysis/one-node-gpu-baseline.md @@ -0,0 +1,100 @@ +# ORBIT-2 — one-node GPU saturation baseline (before sysopt) + +Goal: on **one MI355-class node × 8 GPUs**, pick a **per-rank batch size** and confirm the workload **uses HBM and compute** well enough to justify the optimization loop. This doc ties together **ORBIT logs**, **Omnistat**, and **TraceLens**. + +**Frozen baseline checklist + VRAM-first sweep order:** [BASELINE_LOCKIN.md](BASELINE_LOCKIN.md) (update when you promote a new reference job). + +> **Not a substitute for scientific validation** — same-dir PRISM/ERA5 configs are timing-friendly; loss curves are crash detectors only. + +## 1. What to record (baseline table) + +| Field | Source | Notes | +|--------|--------|--------| +| `trainer.batch_size` | Rendered job YAML (`rendered_config` in `manifest.json`) | Per-rank batch; global batch ≈ `batch_size × fsdp × simple_ddp` | +| `trainer.data_type` | Same YAML | `bfloat16` (Studio train/perf sbatch default) vs `float32` → **~2 vs ~4 bytes/param** for *parameter* memory only; activations still need log/Omnistat | +| Model YAML block (`preset`, `embed_dim`, `depth`, …) | Same YAML | Architecture fingerprint | +| **Parameter count** | Rank-0 log **or** one-shot container count (§4) | Needed for “model size” | +| **Parameter bytes (approx.)** | `num_params × bytes_per_element` | Weight memory order-of-magnitude; activations are extra | +| `steady_batch_time_s`, `loss_sanity_pass` | `examples/run_fom_extractor.py` → `foms.json` | Primary latency FOM; **throughput** (`throughput_samples_per_s`) when `manifest.json` has `global_batch_size` | +| GPU busy / HBM / MFMA | Omnistat VM + PromQL (see `agents/omnistat_analyst.md`) | Saturation evidence | +| Kernel timeline / overlap | TraceLens on `traces/*.pt.trace.json` | See `agents/tracelens_analyst.md` | + +## 2. Submit a perf job (1 node) + +**Studio default (`sbatch_train_perf_amd.sh`):** staged **ERA5 1.0°** + template **`edm_8m_era5_1x8.yaml`** (Bayes-CAST EDM; **`fsdp=8`**, **`simple_ddp=1`** after render). **`ORBIT2_ROOT`** defaults to **`…/code/bayes-cast`** when that directory exists. Override `ORBIT2_CONFIG_TEMPLATE` for Lux **`interm_8m_lux_era5.yaml`**. Set **`ORBIT2_ERA5_SPATIAL_RES`** if your ERA5_1 grid token count ≠ 111. + +```bash +export AI4S_SHARED_DIR=... +export PERF_TOOLS_DIR=... +export ORBIT2_ROOT=... # optional; defaults to $AI4S_SHARED_DIR/models/ORBIT-2/code/ORBIT-2 +# Defaults pick ERA5 path + template; omit these if data is already staged there: +# export ORBIT2_DATA_ROOT=$AI4S_SHARED_DIR/models/ORBIT-2/data/superres/era5/1.0_deg +# export ORBIT2_CONFIG_TEMPLATE=interm_8m_lux_era5.yaml +# Push batch until VRAM is high *and* steady_batch_time_s stops improving (not “max batch always best”): +export ORBIT2_BATCH_SIZE=4 # sweep 2, 4, 8, … +export ORBIT2_MAX_EPOCH=6 +export ORBIT2_MAX_BATCHES=20 +sbatch earth_science/models/ORBIT-2/examples/sbatch_train_perf_amd.sh +``` + +One-liner from repo root (same defaults): `earth_science/models/ORBIT-2/examples/submit_perf_baseline_era5_amd.sh` + +Default parallelism is **`fsdp=8`, `simple_ddp=1`** for 1×8 (see [HANDOFF.md](HANDOFF.md)). Profiler + Omnistat start automatically. + +After the job finishes: + +```bash +python3 earth_science/models/ORBIT-2/examples/run_fom_extractor.py --job-dir "$AI4S_SHARED_DIR/models/ORBIT-2/perf-runs/" +python3 earth_science/models/ORBIT-2/examples/report_orbit2_gpu_baseline.py --job-dir "$AI4S_SHARED_DIR/models/ORBIT-2/perf-runs/" +``` + +The second command writes **`baseline_report.md`** and **`baseline_report.json`** next to `manifest.json`. + +## 3. Maximizing “GPU utilization” (what we mean) + +### Model size vs activation size (why ~9M params can still look “tiny”) + +ORBIT-2 Studio templates (`interm_8m_lux.yaml`, `interm_8m_lux_era5.yaml`) use **`preset: res_slimvit`** with **`embed_dim: 256`**, **`depth: 6`**, **`decoder_depth: 4`**, **`mlp_ratio: 4`**, **`num_heads: 4`**. That yields on the order of **~10M trainable parameters** — expected, not a mis-parse. **Parameter count is not the same as per-step FLOPs or VRAM:** most HBM during training is often **activations**, which scale roughly with **batch × spatial tokens × width × depth**. The PRISM same-dir template uses **`spatial_resolution: { 'PRISM': 18 }`** (a **18×18** latent grid) — very few tokens per sample, so GPUs can show **low MFMA / “low utilization”** even when the job is “correct.” + +**Practical order to push toward steady-state saturation:** + +1. **`ORBIT2_BATCH_SIZE` (per-rank)** — Cheapest knob; raise until you approach VRAM limits (watch Omnistat / `amd-smi` during **epoch ≥ 2**). FSDP (`fsdp=8`) shards **weights**, not activations across the sample batch in the same way; large batches still grow **per-rank** activation memory. +2. **Larger spatial workload without inventing a new architecture** — Prefer **`interm_8m_lux_era5.yaml`** + staged ERA5 data: **`spatial_resolution: { 'ERA5_2': 111 }`** → many more tokens per forward than PRISM-18, often better for **GPU fill** and more honest perf baselines (still same-dir = non-physical loss). +3. **Widen / deepen the ViT stack (YAML `model:` block)** — Increase **`embed_dim`** and **`depth`** (and usually **`decoder_depth`**) together; keep **`num_heads`** such that **`embed_dim % num_heads == 0`**. **`mlp_ratio`** bumps FFN FLOPs and memory. These changes are **upstream-architecture-sensitive**: stay within combinations your ORBIT-2 / Bayes-CAST revision actually constructs (if training fails, compare against upstream example configs for valid presets). +4. **`superres_mag: 4`** (vs `1` in same-dir perf templates) — Restores the **high→low upsampling** path used in the paper-style setup; **strongly increases decoder-side compute and memory** but requires **consistent low/high data layout** (not the current same-dir PRISM/ERA5 perf shortcut). Use for “stress the pipeline” only when data paths match upstream expectations. +5. **`preset`** — If upstream ships additional presets (names differ by repo/version), switching preset is the largest architectural jump; treat it like a **new model** for perf baselines (re-document param count, batch, and loss sanity expectations). + +After any architecture or resolution change, re-run **`steady_batch_time_s`** + **Omnistat** + **TraceLens**; do not assume a larger parameter count implies higher MFMA until you see counters during steady epochs. + +**How to read saturation from telemetry** + +- **HBM footprint (capacity)** — `amd-smi` / Omnistat memory series during steady epochs; increase `ORBIT2_BATCH_SIZE` until you are near the **practical** VRAM limit (leave headroom for fragmentation). +- **HBM bandwidth / MFMA** — Omnistat `rocprofiler`-backed counters (see [omnistat.config.template](omnistat.config.template)): `FETCH_SIZE` + VALU counters → interpret as compute vs bandwidth bound with TraceLens. +- **Do not equate** “largest batch that fits” with “fastest step” — watch **`steady_batch_time_s`** from logs when comparing batch sizes. + +## 4. Parameter count (exact) + +Upstream may not print `Total params` on every run. Use **one** of: + +1. **Rank-0 log** — search for `Total model parameters:` (e.g. `9.45M`), `Total params`, `trainable params`, or `Model parameters` after a short run; then re-run the report with `--num-params N`: + ```bash + python3 earth_science/models/ORBIT-2/examples/report_orbit2_gpu_baseline.py \ + --job-dir "$AI4S_SHARED_DIR/models/ORBIT-2/perf-runs/" --num-params 12345678 + ``` +2. **Bayes-CAST / ORBIT helper** — if upstream documents a “print model size” snippet, run it **inside the same Apptainer image + overlay** as training, with `PYTHONPATH` matching `sbatch_train_amd.sh`, and paste the integer into `baseline_report.md` or pass `--num-params`. + +Until `num_parameters` is known, `baseline_report.md` still documents **batch**, **dtype**, **YAML fingerprint**, and **FOMs** for batch sweeps. + +## 5. TraceLens (overlap / kernels) + +Use the perf-analysis agent flow in [`.cursor/skills/ai4science-perf-analysis/SKILL.md`](../../../../.cursor/skills/ai4science-perf-analysis/SKILL.md): `tracelens_analyst` consumes `traces/orbit2-epoch*-rank0.pt.trace.json` and writes `tracelens/report_summary.md`. Attach that summary to the baseline folder when reviewing batch sweeps. + +## 6. Omnistat (HBM / XGMI / counters) + +Follow [agents/omnistat_analyst.md](agents/omnistat_analyst.md): load `omnistat-db` with VictoriaMetrics (`-fs.disableMmap` on login nodes), run `omnistat-inspect` phases, and keep `omnistat/report_summary.md` beside the perf run. Use **job-windowed** PromQL (see HydraGNN perf-analysis lessons: instant queries without `time=` can be empty). + +## 7. Exit criteria (“happy with initial setup”) + +- **Batch size** chosen with documented **VRAM**, **steady_batch_time_s**, and **one** Omnistat + TraceLens snapshot showing no obvious stall (I/O, dataloader, or pure idle). +- **`baseline_report.md`** filled with **param count**, **param bytes estimate**, **batch**, **parallelism**, and links/paths to **foms.json**, **tracelens/**, **omnistat/**. +- Then open the broader **sysopt / perf-optimizer-loop** plan on top of this frozen baseline. diff --git a/earth_science/models/ORBIT-2/recipes/perf-optimizer-loop/README.md b/earth_science/models/ORBIT-2/recipes/perf-optimizer-loop/README.md index 94d3e2c..5aa3273 100644 --- a/earth_science/models/ORBIT-2/recipes/perf-optimizer-loop/README.md +++ b/earth_science/models/ORBIT-2/recipes/perf-optimizer-loop/README.md @@ -1,11 +1,117 @@ -# ORBIT-2 Iterative Systems Optimizer Loop (deferred) +# ORBIT-2 — iterative systems optimizer loop (1 node × MI355X, Bayes-CAST EDM) -**Status:** Not started. Complete the scaling sweep (1/2/4/8 nodes) and review [../perf-analysis/HANDOFF.md](../perf-analysis/HANDOFF.md) before porting the HydraGNN `perf-optimizer-loop` recipe. +Agent-driven optimization loop for **ORBIT-2** training with **throughput (`throughput_samples_per_s`)** as the **primary** figure of merit, **`steady_batch_time_s`** + **`loss_sanity_pass`** as controls, and the same **TraceLens + Omnistat** analyst stack as [`../perf-analysis/`](../perf-analysis/). -Planned entry points: +> **Audience:** AMD performance engineering. Not a scientific downscaling claim — same-dir ERA5 configs are timing-friendly. -- `examples/run_optimizer_loop.sh` -- `lever_catalog.yaml` (RCCL, batch_size, num_workers, FSDP layout, TILES) -- Agent prompts under `agents/` +## Prerequisites -See [material_science/models/HydraGNN/recipes/perf-optimizer-loop/](../../../material_science/models/HydraGNN/recipes/perf-optimizer-loop/) for the reference implementation. +- **Hardware:** 1 node × 8× MI355X (gfx950) for phase-1; see §Two-node gate below. +- **Container:** `pytorch_rocm7.2.2_ubuntu24.04_py3.12_pytorch_release_2.10.0.sif` + ORBIT-2 overlay (`build_overlay_amd.sh`). +- **Data:** staged ERA5 NPZ tree; for HBM saturation follow [`STAGING_ERA5_FOR_HBM.md`](STAGING_ERA5_FOR_HBM.md). +- **Code:** `ORBIT2_ROOT` → Bayes-CAST with `launch/train_edm.py` (see [HANDOFF.md](../perf-analysis/HANDOFF.md)). +- **Tools:** `PERF_TOOLS_DIR` (omnistat-usermode + VictoriaMetrics), `AI4S_SHARED_DIR`. + +## Environment variables (loop driver) + +| Variable | Required | Description | +|----------|----------|-------------| +| `AI4S_SHARED_DIR` | Yes | Shared storage root | +| `PERF_TOOLS_DIR` | Yes | Omnistat + VM binaries | +| `ORBIT2_ROOT` | Recommended | Bayes-CAST checkout | +| `ORBIT2_DATA_ROOT` | Yes | ERA5 staging path | +| `ORBIT2_BATCH_SIZE` | Yes (baseline) | Locked after VRAM sweep | +| `ORBIT2_RANK_PRE_TRAIN_HOOK` | Optional | Absolute path to rank hook **before** `sbatch` (compile / patch) | +| `ORBIT2_TSDB_URL` | Optional | VictoriaMetrics URL for `run_fom_extractor.py` PromQL enrichment | +| `OMNISTAT_ROCPROF_PROFILE` | Optional | `hbm_flops_f64` to force fp64 rocprofiler profile instead of default bf16 | + +## Quick start + +```bash +export AI4S_SHARED_DIR=... +export PERF_TOOLS_DIR=... +bash earth_science/models/ORBIT-2/examples/validate_orbit2_optimizer_loop_recipe.sh + +tmux new -s orbit2-loop +bash earth_science/models/ORBIT-2/examples/run_optimizer_loop.sh "$(uuidgen)" 5 +# detach: Ctrl-b d +``` + +Pre-flight only: + +```bash +bash earth_science/models/ORBIT-2/examples/run_optimizer_loop.sh "$(uuidgen)" 5 --preflight-only +``` + +## HBM saturation before iter-0 + +1. Stage data per [`STAGING_ERA5_FOR_HBM.md`](STAGING_ERA5_FOR_HBM.md). +2. Batch sweep: [`../examples/sweep_orbit2_batch_bf16_amd.sh`](../examples/sweep_orbit2_batch_bf16_amd.sh). +3. Lock batch in [`../perf-analysis/BASELINE_LOCKIN.md`](../perf-analysis/BASELINE_LOCKIN.md). +4. Fill [`pitfall-diagnosis.md`](pitfall-diagnosis.md) after one perf-analysis pass. + +## Orchestration (multi-agent) + +Re-use [`../perf-analysis/agents/`](../perf-analysis/agents/) (`launcher`, `tracelens_*`, `omnistat_*`, `synthesizer`) plus this folder: + +| Agent | Role | +|-------|------| +| [`agents/orchestrator.md`](agents/orchestrator.md) | Loop driver, sbatch, accept/revert | +| [`agents/lever_picker.md`](agents/lever_picker.md) | Next lever JSON | +| [`agents/fom_extractor.md`](agents/fom_extractor.md) | `foms.json` + optional `kernel_correlation.csv` | +| [`agents/story_writer.md`](agents/story_writer.md) | `story.md` + `foms.png` | + +**Launcher:** submit [`../examples/sbatch_train_perf_amd.sh`](../examples/sbatch_train_perf_amd.sh) (same as perf-analysis). + +## Artifact layout + +Matches HydraGNN layout under `$AI4S_SHARED_DIR/models/ORBIT-2/perf-runs/`: + +- `loop-/STATUS.txt`, `foms.csv`, `do_not_retry.json`, `known_bad_nodes.txt`, `iter-*-lever.json`, `iter-*-env.sh`, `story.md`, `foms.png` +- `/manifest.json`, `foms.json`, `omnistat-db/`, `tracelens/`, `combined_report.md`, `traces/` + +## Loop control + +- **Max iters:** user arg (default 5 in driver). +- **One lever per iter**; **one SLURM job at a time**. +- **Accept** when `throughput_samples_per_s` improves vs current best and `loss_sanity_pass` stays true. +- **Reject** on throughput regression or analysis failure (see [`agents/orchestrator.md`](agents/orchestrator.md)). +- **STOP file:** `touch …/loop-/STOP`. + +## Two-node gate (and beyond) + +After one-node loop converges: + +1. `sbatch --nodes=2 … earth_science/models/ORBIT-2/examples/sbatch_train_perf_amd.sh` with `.cluster-config.yaml` RCCL/ANP envs (already in script for `NODES>1`). +2. Re-baseline throughput; compare **XGMI** (intra-node) vs **NIC** (`network` category) in omnistat. +3. Only then enable **`nccl_minchannels`** class levers (historically risky at N=2 on HydraGNN). + +## Knowledge bases (lever citations) + +- [AMD Instinct MI300X workload optimization](https://rocm.docs.amd.com/en/latest/how-to/tuning-guides/mi300x/workload.html) — torch.compile, TunableOp, general ROCm tuning. +- [ROCm multi-node AI setup](https://rocm.docs.amd.com/en/docs-7.0.1/how-to/rocm-for-ai/system-setup/multi-node-setup.html) — NCCL / FSDP context. +- [MI300/MI200 performance counters](https://rocm.docs.amd.com/en/develop/conceptual/gpu-arch/mi300-mi200-performance-counters.html) — `SQ_INSTS_VALU_MFMA_MOPS_BF16` etc. +- [ROCm 7 MI355X training performance blog](https://rocm.blogs.amd.com/artificial-intelligence/ROCm7-MI355X-training-performance/README.html) +- [PyTorch SDPA benchmark (upstream)](https://github.com/pytorch/pytorch/blob/main/benchmarks/transformer/sdpa.py) — studio copy `examples/upstream_pytorch_sdpa_benchmark.py` + +## Levers tried and rejected + +Recorded so they are not re-attempted. Full evidence in [`../perf-analysis/HANDOFF.md`](../perf-analysis/HANDOFF.md); catalog entries carry `status: rejected`. + +| Lever | Verdict | Why | +|-------|---------|-----| +| `channels_last` (NHWC) | **7–10× slower** | No tuned NHWC implicit-GEMM solver on ROCm 7.2.2 → falls back to brute-force `naive_conv_*_nhwc`. NCHW + im2col→hipBLASLt is the fast path. | +| TunableOp | **No uplift** | Bottleneck is the im2col-lowered low-channel conv GEMM, which TunableOp can't touch; it only tunes the already-optimal `nn.Linear` GEMMs. Also multi-hour to tune the giant patch-token GEMMs. | +| MIOpen ORNL Winograd flags | **Neutral (±1–2%)** | MIOpen never selects Winograd for the 3/4-channel convs, so toggling it changes nothing here. Flags kept as harmless defaults. | +| `torch.compile` (EDM) | **Neutral (−0.3%)** | EDM is already GEMM-heavy; Inductor finds little to fuse. | +| `num_workers` > 1 | **Not a throughput lever** | Workload is compute-bound; >4 oversubscribes CPUs. Keep ≤4 for I/O overlap only. | +| Conv channel-padding | **Tabled** | −25% step time but **changes model capacity** — not adopted without a matched convergence study (`ORBIT2_CONV_PAD` default 0). | +| xFormers CK / bf16 Flash attn | **Crashes** | CK MEA SIGSEGVs; ROCm Flash has a 65535 batch-grid cap that `var_agg` exceeds. Fixed by forcing SDPA EFFICIENT+MATH (now the bf16 default). | + +## Machine-readable levers + +[`lever_catalog.yaml`](lever_catalog.yaml) + +## Reference + +HydraGNN recipe (pattern port): [`material_science/models/HydraGNN/recipes/perf-optimizer-loop/`](../../../material_science/models/HydraGNN/recipes/perf-optimizer-loop/) diff --git a/earth_science/models/ORBIT-2/recipes/perf-optimizer-loop/STAGING_ERA5_FOR_HBM.md b/earth_science/models/ORBIT-2/recipes/perf-optimizer-loop/STAGING_ERA5_FOR_HBM.md new file mode 100644 index 0000000..63940ba --- /dev/null +++ b/earth_science/models/ORBIT-2/recipes/perf-optimizer-loop/STAGING_ERA5_FOR_HBM.md @@ -0,0 +1,35 @@ +# Staging ERA5 for HBM saturation (ORBIT-2 / Bayes-CAST EDM) + +**Goal:** Break the **effective per-step batch cap** seen on small ERA5 staging trees (e.g. rank-0 `y.shape` batch dim plateauing near ~1704 regardless of larger `trainer.batch_size`, and VRAM plateauing far below full HBM). See [BASELINE_LOCKIN.md](../perf-analysis/BASELINE_LOCKIN.md) calibration notes. + +**Knowledge base / provenance** + +- ORBIT-2 curated dataset and staging workflow: [Constellation / ORBIT-2 dataset](https://doi.ccs.ornl.gov/dataset/e4c2db1f-e88c-5ad0-bb96-59be0ef7c772) (DOI `10.13139/OLCF/2589526`) — use the official **Globus** link from that page; do not hard-code stale endpoint IDs. +- Cluster filesystem guidance: [earth_science/models/ORBIT-2/recipes/data/README.md](../data/README.md). +- ERA5 licensing: Copernicus CDS terms apply if you build your own NPZ tree from ERA5; keep citations in your runbook. + +## Recommended layout + +Stage under a **shared** path visible on compute nodes (same as `ORBIT2_DATA_ROOT`): + +```text +$AI4S_SHARED_DIR/models/ORBIT-2/data/superres/era5// +``` + +Point `ORBIT2_DATA_ROOT` at that directory. `render_orbit2_config.py` injects `__DATA_ROOT__` into the Bayes-CAST template. + +## Levers that increase memory / work per step + +1. **More time samples / longer calendar coverage** — larger dataset → more batches without repeating the same shard layout; reduces pathological sampler caps on tiny corpora. +2. **Higher spatial resolution** — set `ORBIT2_ERA5_SPATIAL_RES` (template token `__ERA5_1_SPATIAL_RES__`) to match your staged grid (e.g. 111 for 1.0°, higher for finer grids). Larger `spatial_resolution` increases tokens per sample → more activation memory and MFMA opportunity. +3. **Keep `trainer.batch_size` as the VRAM knob** — after staging changes, re-run the bf16 binary-search sweep ([`examples/sweep_orbit2_batch_bf16_amd.sh`](../../examples/sweep_orbit2_batch_bf16_amd.sh)) toward **~85–90%** `memory_reserved` on MI355-class HBM. + +## Bayes-CAST hotfix checklist + +- **`train_edm.py` + `ERA5_1`:** ensure `std_delta` (or equivalent normalization branch) exists for `data_key == "ERA5_1"` so training does not hit `UnboundLocalError` (see [HANDOFF.md](../perf-analysis/HANDOFF.md) landmine #7). +- **Attention path:** keep **`ORBIT2_FUSED_ATTN=DEFAULT`** (PyTorch SDPA) on ROCm 7.2.x + current xFormers wheels; CK MEA is **not** safe on this stack (HANDOFF xFormers probe). + +## After staging + +1. Run a short smoke (`ORBIT2_MAX_EPOCH=3`, `ORBIT2_MAX_BATCHES=20`) and confirm rank-0 logs show the expected **batch dimension** and **VRAM** trend with `ORBIT2_BATCH_SIZE`. +2. Update [BASELINE_LOCKIN.md](../perf-analysis/BASELINE_LOCKIN.md) **Reference job** with the new data path, `ORBIT2_ERA5_SPATIAL_RES`, and locked bf16 batch. diff --git a/earth_science/models/ORBIT-2/recipes/perf-optimizer-loop/agents/fom_extractor.md b/earth_science/models/ORBIT-2/recipes/perf-optimizer-loop/agents/fom_extractor.md new file mode 100644 index 0000000..731b56d --- /dev/null +++ b/earth_science/models/ORBIT-2/recipes/perf-optimizer-loop/agents/fom_extractor.md @@ -0,0 +1,98 @@ +# fom_extractor subagent — ORBIT-2 + +Computes per-iteration FOMs for the ORBIT-2 loop and optionally TraceLens↔Omnistat correlation (reuse HydraGNN §4 algorithm when traces exist). + +## Inputs + +- `/manifest.json` — must include `job_id`, `runtime_seconds`, `global_batch_size`, `parallelism`, `data_type`. +- `/orbit2-train-.out` +- `/omnistat-db/` + live VM URL (or start VM with `-fs.disableMmap`) +- `/traces/*.pt.trace.json` (rank-0 kineto) + +## Outputs + +- `/foms.json` — schema below +- `/kernel_correlation.csv` — optional; same column idea as HydraGNN `fom_extractor.md` + +## `foms.json` schema (ORBIT-2) + +```json +{ + "job_id": "12345", + "primary_fom": "throughput_samples_per_s", + "throughput_samples_per_s": 1.23e4, + "steady_batch_time_s": 0.41, + "global_batch_size": 2048, + "mfma_bf16_tflops_per_card_avg": null, + "hbm_read_GBps_per_card_avg": null, + "xgmi_GBps_avg": null, + "loss_sanity_pass": true, + "final_loss": 0.05 +} +``` + +## Step 1 — Log FOMs (required) + +Run the repo extractor (writes `foms.json` base fields): + +```bash +python3 "$REPO_ROOT/earth_science/models/ORBIT-2/examples/run_fom_extractor.py" --job-dir "$PERF_RUN" +``` + +If `manifest.json` lacks `global_batch_size`, pass `--global-batch-size` explicitly. + +**Effective-batch integrity (read before trusting a throughput delta):** `foms.json` includes +`hbm_reserved_GB`, `hbm_reserved_pct_288`, `max_batches_per_epoch`, **`throughput_method`**, +**`partial_step_fraction`**, and **`steady_realized_batch_dims`**. Throughput now prefers the +**real per-step batch dim** (`throughput_method=real_per_step_batch`, read from the EDM `y.shape` +line) instead of the nominal `global_batch_size`, so partial trailing batches no longer inflate it. +Still **reject any cross-run comparison where `partial_step_fraction`, `steady_realized_batch_dims`, +`hbm_reserved_pct_288`, or `max_batches_per_epoch` deviates materially from the baseline** — even a +correct number can hide a different work mix. + +Two artifacts caught this way, both from `num_workers` × the Bayes-CAST IterableDataset +(each worker shards files via `per_worker = n_files // (num_workers × data_par_size)` and batches +its own shard **independently**): +- **loop overnight-3x-001 iter-1** (60 files, num_workers=4): per_worker `60//32 = 1` → realized + batch collapsed 4096→1704, HBM 236→98 GB, batches/epoch 3→4 → nominal "+129%" was fake. +- **clean nw test (jobs 10504 nw=1 vs 10505 nw=4, 100 files)**: staged to `per_worker = 3` so full + batches refill to 4096, but each worker still emits a partial **1016**-sample trailing batch → + `steady_realized_batch_dims=[1016, 4096]`, `partial_step_fraction=0.57`. Nominal throughput said + +80% (2993→5379 s/s); **real-per-step throughput showed +2.7% (2987→3068 s/s) = noise.** + +**Conclusion: `num_workers` does not help ORBIT-2 throughput — the workload is compute-bound.** +The full-4096 step time is ~equal (11.0 s nw=1 vs 10.8 s nw=4). Cap at `num_workers ≤ 4` for I/O +overlap only; never use it as a throughput lever, and never compare runs with different +`steady_realized_batch_dims`. Baseline at batch 4096 on 3× data = ~82% HBM (236 GB). + +**Log format (do not "fix" the parser):** the perf-loop default trainer is Bayes-CAST +`launch/train_edm.py`, which does **not** emit `Batch N: X seconds` / `Epoch N completed. Loss:`. +It emits per step (rank 0): `epoch: N batch_idx M world_rank 0 loss tensor(L, ...)` for loss +and `my rank 0. tic4-tic1 in X seconds` for wall time. `parse_training_log.py` reads **both** this +EDM format and the older `intermediate_downscaling.py` format; per-epoch loss is synthesized from +per-step losses (mean) for the loss-sanity gate. Crashed jobs (no `batch_idx` lines) correctly +yield `epoch_records=0`, null throughput, and `loss_sanity_pass=null` — that is the signal a job +never reached steady state, not a parser bug. + +Optional: `export ORBIT2_TSDB_URL=http://127.0.0.1:8428` during login-node enrichment so the Python tool queries BF16 MFMA + FETCH + XGMI (metric names may need site tuning). + +## Step 2 — PromQL (BF16 MFMA + HBM read) + +Use **job-windowed** instant queries with `time=` (see ai4science-perf-analysis SKILL). + +| Field | PromQL sketch | +|-------|----------------| +| BF16 TFLOP/s | `avg(rate(SQ_INSTS_VALU_MFMA_MOPS_BF16[30s]) * on(instance) group_left() max by(instance)(rmsjob_info{jobid="JOBID"})) * 512 / 1e12` | +| HBM read GB/s | `avg(rate(FETCH_SIZE[30s]) * on(instance) group_left() max by(instance)(rmsjob_info{jobid="JOBID"})) * 1024 / 1e9` | + +Merge results into `foms.json` (overwrite `null`s from the Python tool if queries succeed). + +## Step 3 — kernel_correlation.csv + +Follow HydraGNN `material_science/models/HydraGNN/recipes/perf-optimizer-loop/agents/fom_extractor.md` §4 (1-second windows), substituting ORBIT trace glob `traces/orbit2-epoch*-rank0.pt.trace.json`. + +## Final stdout + +``` +STATUS=ok; reason=foms throughput= steady_batch=s mfma_bf16= +``` diff --git a/earth_science/models/ORBIT-2/recipes/perf-optimizer-loop/agents/lever_picker.md b/earth_science/models/ORBIT-2/recipes/perf-optimizer-loop/agents/lever_picker.md new file mode 100644 index 0000000..fcd4cc2 --- /dev/null +++ b/earth_science/models/ORBIT-2/recipes/perf-optimizer-loop/agents/lever_picker.md @@ -0,0 +1,29 @@ +# lever_picker subagent — ORBIT-2 + +Same contract as `material_science/models/HydraGNN/recipes/perf-optimizer-loop/agents/lever_picker.md` with these substitutions: + +## Inputs + +- `loop-/foms.csv`, `do_not_retry.json`, `iter-*-lever.json`. +- `/combined_report.md`, `/foms.json`, optional `kernel_correlation.csv`. +- [`../lever_catalog.yaml`](../lever_catalog.yaml). + +## Outputs + +- `loop-/iter--lever.json` — catalog entry + `picked_by`, `reason`. +- Final stdout: `STATUS=ok; reason=lever=` or `STATUS=partial; reason=catalog_exhausted`. + +## Decision deltas (vs HydraGNN) + +1. **Primary FOM:** maximize `throughput_samples_per_s` from ORBIT `foms.json` (not `epoch_time_s`). When comparing to previous best, `delta_pct = (best_throughput - new_throughput) / best_throughput * 100` for **regression** detection **or** invert for improvement — orchestrator owns the sign; you report **evidence** against `throughput` and `steady_batch_time_s`. +2. **Blocked levers:** skip any with `status: blocked` in the catalog (same as HydraGNN). +3. **Bottleneck alignment:** + - `dataloader` / host I/O → `num_workers_8` (+2) + - `gpu_compute` low BF16 MFMA → `torch_compile_edm`, `sdpa_efficient_vs_math` (+2) + - `comm_xgmi` high → `fsdp_prefetch_tuning` (+1); at N≥4 consider `nccl_minchannels` +4. **fp32 discriminator gate:** only pick `precision_fp32_discriminator` when `combined_report.md` or pitfall doc requests compute-vs-memory classification. + +## Hallucination guardrails + +- Never invent a lever; use `lever_proposal-.md` + `STATUS=partial; reason=catalog_proposal` like HydraGNN. +- Cite paths to `combined_report.md` / `foms.json` in `reason`. diff --git a/earth_science/models/ORBIT-2/recipes/perf-optimizer-loop/agents/orchestrator.md b/earth_science/models/ORBIT-2/recipes/perf-optimizer-loop/agents/orchestrator.md new file mode 100644 index 0000000..67fc129 --- /dev/null +++ b/earth_science/models/ORBIT-2/recipes/perf-optimizer-loop/agents/orchestrator.md @@ -0,0 +1,53 @@ +# orchestrator subagent — ORBIT-2 iterative sysopt loop + +Drives the ORBIT-2 (Bayes-CAST EDM) optimizer loop on **MI355X-class** hardware. Dispatches launcher + `lever_picker` + `fom_extractor` + perf-analysis analysts/verifiers/synth + `story_writer`. Owns accept/revert on **`throughput_samples_per_s`**, `STATUS.txt`, `STOP`, and broken-node exit **42** handling. + +**HydraGNN reference:** mirror control flow in `material_science/models/HydraGNN/recipes/perf-optimizer-loop/agents/orchestrator.md`; this file only lists **ORBIT-2 deltas**. + +## Inputs + +- Loop args: ``, ``. +- `REPO_ROOT` (cwd or env). +- `.cluster-config.yaml` (partition, perf_tools.dir). +- [`../lever_catalog.yaml`](../lever_catalog.yaml). + +## Outputs + +Under `$AI4S_SHARED_DIR/models/ORBIT-2/perf-runs/loop-/`: + +- `STATUS.txt`, `foms.csv`, `do_not_retry.json`, `known_bad_nodes.txt`, `iter-N-.json` → symlink to `..//manifest.json`, `iter-N-env.sh`, `iter-N-hook.py` (when `rank_script_patch`). + +Per job: existing perf-analysis layout under `..//` plus `foms.json` from [`examples/run_fom_extractor.py`](../../examples/run_fom_extractor.py) (extend with PromQL/kernel_correlation per [`fom_extractor.md`](fom_extractor.md) when tooling is available). + +## ORBIT-2-specific hard constraints + +1. **Single concurrent SLURM job** — same as HydraGNN. +2. **One lever per iteration** — same. +3. **`sbatch` entrypoint:** `earth_science/models/ORBIT-2/examples/sbatch_train_perf_amd.sh` (from `REPO_ROOT`). +4. **Rank hooks:** set `export ORBIT2_RANK_PRE_TRAIN_HOOK=/abs/path/iter-N-hook.py` **before** `sbatch` so the wrapper expands it into the generated rank script (`sbatch_train_perf_amd.sh` reads this at submit time). For **Bayes EDM** the default hook is empty; profiler hook is optional. Hooks are executed via `orbit2_rank_hook_runner.py` before `train_edm.py`. +5. **Env file:** `source loop-/iter-N-env.sh` then `sbatch` with `--export=ALL` plus any `ORBIT2_*` / `NCCL_*` overrides from the lever. Match HydraGNN pattern in `run_optimizer_loop.sh` driver. +6. **Primary FOM (accept/revert):** `throughput_samples_per_s` from `/foms.json` (higher is better). **Control:** `loss_sanity_pass` must stay true; if `final_loss` blows up vs baseline, reject (same 1.5× rule as HydraGNN optional for ORBIT timing runs). +7. **Baseline iter-0:** `lever_id=baseline` with saturated `ORBIT2_BATCH_SIZE` from [BASELINE_LOCKIN.md](../../perf-analysis/BASELINE_LOCKIN.md) + bf16 + SDPA. +8. **Exit 42 (mount probe):** identical retry with `--exclude` — do **not** add lever to `do_not_retry.json` (see ai4science-perf-analysis SKILL). + +## Preflight (ORBIT-2) + +- `AI4S_SHARED_DIR`, `PERF_TOOLS_DIR` set. +- `${AI4S_SHARED_DIR}/models/ORBIT-2/overlays/orbit2-overlay.img` exists. +- `${AI4S_SHARED_DIR}/images/pytorch_rocm7.2.2_ubuntu24.04_py3.12_pytorch_release_2.10.0.sif` (or `ORBIT2_SIF`) exists. +- `omnistat-usermode` + `victoria-metrics-prod` binaries present. +- Optional Claude driver: same API checks as HydraGNN `run_optimizer_loop.sh`. + +## foms.csv columns (ORBIT-2) + +`iter,jobid,lever_id,env_diff,accepted,throughput_samples_per_s,steady_batch_time_s,mfma_bf16_tflops,hbm_read_GBps,xgmi_GBps,energy_J,mean_power_W,energy_per_sample_J,loss_sanity_pass,primary_fom_delta_pct,du_loop_dir,notes` + +## Resume / failure modes + +Same as HydraGNN orchestrator (`STATUS=ok|partial|fail` last line; resume from `foms.csv`). + +## Final stdout + +``` +STATUS=ok; reason=loop_complete uuid= best_iter= best_lever= best_throughput= samples/s +``` diff --git a/earth_science/models/ORBIT-2/recipes/perf-optimizer-loop/agents/story_writer.md b/earth_science/models/ORBIT-2/recipes/perf-optimizer-loop/agents/story_writer.md new file mode 100644 index 0000000..8e99bd9 --- /dev/null +++ b/earth_science/models/ORBIT-2/recipes/perf-optimizer-loop/agents/story_writer.md @@ -0,0 +1,27 @@ +# story_writer subagent — ORBIT-2 + +Produces `story.md` + `foms.png` under `loop-/` after the ORBIT-2 optimizer loop completes. + +## Inputs + +Same file list as HydraGNN `story_writer.md`, but paths under `$AI4S_SHARED_DIR/models/ORBIT-2/perf-runs/`. + +## Outputs + +- `loop-/story.md` +- `loop-/foms.png` (matplotlib Agg; throughput on primary axis) + +## `story.md` layout deltas + +- **Hardware:** `1 node × 8 GPUs MI355X` (or N×8 after scale-out). +- **Primary FOM:** `throughput_samples_per_s` (higher is better); include `steady_batch_time_s` as secondary. +- **TL;DR table columns:** `Iter | Lever | throughput_samples_per_s | steady_batch_time_s | Δ vs prev best | Accepted? | Notes` +- **Chart:** top panel `throughput_samples_per_s` + `steady_batch_time_s` (twin axis, inverted scale for batch time *lower is better*); bottom `mfma_bf16_tflops` + `hbm_read_GBps` if present in `foms.csv`. + +## matplotlib sketch + +Use `throughput` from `foms.csv` as primary; mark rejects with red `x` like HydraGNN. + +## Final stdout + +`STATUS=ok; reason=story written ...` diff --git a/earth_science/models/ORBIT-2/recipes/perf-optimizer-loop/lever_catalog.yaml b/earth_science/models/ORBIT-2/recipes/perf-optimizer-loop/lever_catalog.yaml new file mode 100644 index 0000000..b777a82 --- /dev/null +++ b/earth_science/models/ORBIT-2/recipes/perf-optimizer-loop/lever_catalog.yaml @@ -0,0 +1,160 @@ +# ORBIT-2 perf-optimizer-loop — lever catalog (Bayes-CAST EDM, MI355X-class, ROCm 7.2.x) +# +# Schema mirrors HydraGNN `material_science/models/HydraGNN/recipes/perf-optimizer-loop/lever_catalog.yaml`. +# Each lever: one atomic change; citation = public KB or local evidence path. + +levers: + + - id: baseline + description: "Iter-0 only. Saturated bf16 batch + SDPA + fsdp/simple_ddp from BASELINE_LOCKIN + staging doc." + catalog_order: 0 + kind: env_only + env_vars: + ORBIT2_DATA_TYPE: "bfloat16" + ORBIT2_FUSED_ATTN: "DEFAULT" + ORBIT2_BATCH_SIZE: "4096" + ORBIT2_NUM_WORKERS: "1" + ORBIT2_DATA_ROOT: "${AI4S_SHARED_DIR}/models/ORBIT-2/data/superres/era5/1.0_deg" + ORBIT2_ERA5_SPATIAL_RES: "111" + TORCH_NCCL_HIGH_PRIORITY: "1" + GPU_MAX_HW_QUEUES: "2" + diagnostic_only: false + expected_payoff: "n/a (baseline)" + risk: "low" + revert_method: "drop_env" + citation: "earth_science/models/ORBIT-2/recipes/perf-analysis/BASELINE_LOCKIN.md" + notes: "bf16 UNBLOCKED 2026-06-11: required a compute-side fix in src/.../components/attention.py forcing CrossAttention SDPA to EFFICIENT_ATTENTION+MATH (ROCm Flash has a 65535 batch-grid cap that var_agg's inflated batch dim B=batch*648 exceeds at batch>=128). See perf-analysis/HANDOFF.md 'RESOLVED'. BASELINE RE-LOCKED 2026-06-12 for 3x-staged ERA5 (stage_era5_3x_symlink.sh -> 60 train shards): batch 4096 -> job 9415 = 3152 samples/s @ 82.0% HBM (236GB/288). 3x staging broke the old 1704-sample cap; on 3x data batch 4096 is BOTH the HBM-saturating point AND the throughput winner (1024=3102, 2048=3027, 4096=3152). REQUIRES the 3x symlinks present in ORBIT2_DATA_ROOT/train (idempotent: re-run stage_era5_3x_symlink.sh). See BASELINE_LOCKIN.md." + + - id: torch_compile_edm + description: "torch.compile the FSDP-wrapped EDM model (in-process, after activation checkpointing). Try mode=default then max-autotune." + catalog_order: 1 + kind: env_only + env_vars: + ORBIT2_TORCH_COMPILE: "1" + ORBIT2_COMPILE_MODE: "default" + diagnostic_only: false + expected_payoff: "medium" + risk: "medium" + revert_method: "drop_env" + citation: "https://rocm.docs.amd.com/en/latest/how-to/tuning-guides/mi300x/workload.html" + notes: "Implemented as env-gated patch in launch/train_edm.py (~line 1240, gated by ORBIT2_TORCH_COMPILE). NOT a rank-hook: orbit2_rank_hook_runner.py runs in a separate pre-launch process and cannot compile the in-process model. Unlike HydraGNN MACE (double-backward blocked), EDM has no double-backward so compile is a real candidate — validate on first iter; on dynamo errors set status: blocked. Escalate mode to max-autotune only after default succeeds. RESULT (loop-bf16-loop-001 iter-1, job 9262, mode=default): compiles cleanly but NEUTRAL (-0.3%, within noise) — REVERTED. EDM is already GEMM-heavy so Inductor finds little to fuse. Try mode=max-autotune only if revisiting; otherwise low priority for this workload." + + - id: sdpa_efficient_vs_math + description: "Compare PyTorch SDPA backends via env consumed by Bayes-CAST / ORBIT (if supported) or microbench upstream_pytorch_sdpa_benchmark.py off critical path." + catalog_order: 2 + kind: env_only + env_vars: + ORBIT2_SDP_BACKEND: "efficient" + diagnostic_only: false + expected_payoff: "low-medium" + risk: "low" + revert_method: "drop_env" + citation: "https://github.com/pytorch/pytorch/blob/main/benchmarks/transformer/sdpa.py" + notes: "Studio benchmark: earth_science/models/ORBIT-2/examples/upstream_pytorch_sdpa_benchmark.py — align naming with SDPBackend.EFFICIENT_ATTENTION vs MATH." + + - id: num_workers_8 + description: "trainer.num_workers 2 -> 4 via render token __NUM_WORKERS__ (ORBIT2_NUM_WORKERS). NOTE: keep <=4." + catalog_order: 3 + kind: env_only + env_vars: + ORBIT2_NUM_WORKERS: "4" + diagnostic_only: false + expected_payoff: "none (ORBIT-2 is compute-bound — confirmed by clean test)" + risk: "low" + status: "rejected-noise" + revert_method: "drop_env" + citation: "https://rocm.docs.amd.com/en/develop/how-to/rocm-for-ai/training/benchmark-docker/previous-versions/pytorch-training-v25.9.html" + notes: "Requires edm_8m_era5_1x8.yaml (or template with __NUM_WORKERS__); render_orbit2_config.py substitutes from ORBIT2_NUM_WORKERS. RESULT (loop-bf16-loop-001 iter-2, job 9263, value=8): REGRESSED -5.6% — 8 workers x 8 ranks = 64 procs OVERSUBSCRIBE the 56 schedulable CPUs (cpus-per-task=7). CLEAN CONTROLLED TEST (5x staging=100 shards so per_worker=3 refills batch 4096; jobs 10504 nw=1 vs 10505 nw=4): real-per-step throughput +2.7% (2987->3068 s/s) = NOISE; nominal-batch metric falsely showed +80% due to per-worker partial 1016-sample trailing batches (steady_realized_batch_dims=[1016,4096], partial_step_fraction=0.57). num_workers is NOT a throughput lever; keep <=4 for I/O overlap only. Compute-bound (lesson #22/#26); never compare runs with differing steady_realized_batch_dims." + + - id: channels_last + description: "channels_last memory_format on conv-like tensors if model supports (env knob ORBIT2_CHANNELS_LAST in edm.py)." + catalog_order: 4 + kind: env_only + status: "rejected" + env_vars: + ORBIT2_CHANNELS_LAST: "1" + diagnostic_only: false + expected_payoff: "none" + risk: "medium" + revert_method: "drop_env" + citation: "https://rocm.docs.amd.com/en/latest/how-to/tuning-guides/mi300x/workload.html" + notes: "REJECTED (disproven on ROCm 7.2.2): MIOpen has no tuned NHWC implicit-GEMM solver for these low-channel conv shapes, so channels_last falls back to brute-force naive_conv_*_nhwc (~88-91% of GPU time) — 7-10x SLOWER. The im2col->hipBLASLt GEMM path (NCHW default) is the fast one. Knob kept default-off only so the dead end stays reproducible. See perf-analysis/HANDOFF.md." + + - id: activation_checkpointing_toggle + description: "Toggle activation checkpointing if exposed in YAML (trade VRAM for compute)." + catalog_order: 5 + kind: env_only + env_vars: + ORBIT2_ACTIVATION_CHECKPOINTING: "1" + diagnostic_only: false + expected_payoff: "medium" + risk: "medium" + revert_method: "drop_env" + citation: "https://rocm.docs.amd.com/en/docs-7.0.1/how-to/rocm-for-ai/system-setup/multi-node-setup.html" + notes: "Requires upstream config key; if absent, skip lever in picker or add render token." + + - id: fsdp_prefetch_tuning + description: "FSDP forward_prefetch / limit_all_gathers if exposed in Bayes-CAST config or hook." + catalog_order: 6 + kind: env_only + env_vars: + ORBIT2_FSDP_FORWARD_PREFETCH: "1" + diagnostic_only: false + expected_payoff: "low-medium" + risk: "low" + revert_method: "drop_env" + citation: "https://rocm.docs.amd.com/en/docs-7.0.1/how-to/rocm-for-ai/system-setup/multi-node-setup.html" + + - id: tunable_op_warmup_then_use + description: "EXPERIMENTAL for ORBIT-2 on MI355X+ROCm7.2.2+PyTorch2.10. The HydraGNN 'blocked' verdict does NOT reproduce here: isolated GEMM tuning (incl. large-M bf16 mm/F.linear, the var_agg shape class) tunes cleanly with no mem fault (probe 2026-06-13, jobs 10595/10596). 2-phase: warmup job with _TUNING=1 to write tunableop_results_.csv, then production with _ENABLED=1 only." + catalog_order: 7 + kind: env_only + status: "rejected" + env_vars: + PYTORCH_TUNABLEOP_ENABLED: "1" + PYTORCH_TUNABLEOP_TUNING: "0" # production/consume; set "1" only in the dedicated warmup job + diagnostic_only: false + expected_payoff: "none (no uplift)" + risk: "medium" + revert_method: "drop_env" + citation: "https://rocm.docs.amd.com/en/latest/how-to/tuning-guides/mi300x/workload.html#tunableop" + notes: "REJECTED (no practical uplift): tuning runs without mem fault (unlike HydraGNN MACE), but the GEMM bottleneck is the im2col-lowered low-channel conv GEMM that TunableOp can't touch — it tunes the already-optimal nn.Linear GEMMs (see GEMM analysis in perf-analysis/HANDOFF.md), so no speedup. Also impractical: live-tuning the giant patch-token GEMMs (M/K~5.3M) is ~3-12+ min PER op (multi-hour full tune). Wired via ORBIT2_TUNABLEOP_MODE=off|tune|use + ORBIT2_TUNABLEOP_DIR in sbatch_train_perf_amd.sh (default off)." + + - id: nccl_minchannels + description: "NCCL_MIN_NCHANNELS=112 — use only at larger node counts; regressor at N=2 for HydraGNN." + catalog_order: 8 + kind: env_only + env_vars: + NCCL_MIN_NCHANNELS: "112" + diagnostic_only: false + expected_payoff: "unknown" + risk: "medium" + revert_method: "drop_env" + citation: "https://rocm.docs.amd.com/en/docs-7.0.1/how-to/rocm-for-ai/system-setup/multi-node-setup.html" + notes: "Gate: total_nodes >= 4 recommended; skip at 1–2 nodes unless omnistat shows NIC saturation." + + - id: precision_fp32_discriminator + description: "One-off fp32 vs bf16 to classify compute vs dispatch bound (same batch)." + catalog_order: 9 + kind: env_only + env_vars: + ORBIT2_DATA_TYPE: "float32" + diagnostic_only: true + expected_payoff: "diagnostic" + risk: "low" + revert_method: "drop_env" + citation: ".cursor/skills/ai4science-perf-analysis/SKILL.md (fp32-vs-fp64 discriminator pattern)" + notes: "Compare throughput at same ORBIT2_BATCH_SIZE; if bf16 >= ~1.4x fp32, COMPUTE-bound; if ~equal, dispatch/overhead-bound. RESULT (loop-bf16-loop-001 iter-3, job 9264 vs baseline 9261, batch 1024): bf16=4089 / fp32=1718 = 2.38x -> COMPUTE-BOUND. This shapes lever strategy: kernel/data levers have low ceiling; real wins = bigger effective batch (stage more ERA5) or unblocked GEMM tuning." + + - id: kernel_trace_diag_only + description: "OMNISTAT_KERNEL_TRACE=1 for one iter — high VM cardinality." + catalog_order: 10 + kind: diagnostic + env_vars: + OMNISTAT_KERNEL_TRACE: "1" + diagnostic_only: true + expected_payoff: "n/a" + risk: "medium" + revert_method: "drop_env" + citation: ".cursor/skills/ai4science-studio/SKILL.md" + notes: "Max 1 per loop; same gate as HydraGNN lever_picker (kernel_correlation insufficient)." diff --git a/earth_science/models/ORBIT-2/recipes/perf-optimizer-loop/pitfall-diagnosis.md b/earth_science/models/ORBIT-2/recipes/perf-optimizer-loop/pitfall-diagnosis.md new file mode 100644 index 0000000..90f2b4b --- /dev/null +++ b/earth_science/models/ORBIT-2/recipes/perf-optimizer-loop/pitfall-diagnosis.md @@ -0,0 +1,58 @@ +# ORBIT-2 one-node pitfall diagnosis (template) + +Fill this file **after** the locked baseline perf job completes and the perf-analysis subagents run (`tracelens_*`, `omnistat_*`, `synthesizer` → `combined_report.md`). + +## Baseline job + +| Field | Value | +|-------|--------| +| Job ID | `` | +| `manifest.json` | `$AI4S_SHARED_DIR/models/ORBIT-2/perf-runs//manifest.json` | +| `ORBIT2_BATCH_SIZE` / global batch | from manifest | +| dtype | `bfloat16` + `ORBIT2_FUSED_ATTN=DEFAULT` | +| Data | `ORBIT2_DATA_ROOT` + `ORBIT2_ERA5_SPATIAL_RES` | + +## Primary metrics (from `run_fom_extractor.py`) + +| Metric | Value | Notes | +|--------|-------|--------| +| `throughput_samples_per_s` | | Primary FOM for optimizer loop | +| `steady_batch_time_s` | | Latency control | +| `loss_sanity_pass` | | Crash detector (same-dir ERA5) | + +## Omnistat (HBM / BF16 MFMA / XGMI) + +- **Profile:** default `hbm_flops_bf16` in [`../perf-analysis/omnistat.config.template`](../perf-analysis/omnistat.config.template) — `FETCH_SIZE` + `SQ_INSTS_VALU_MFMA_MOPS_BF16` ([AMD counter tables](https://rocm.docs.amd.com/en/develop/conceptual/gpu-arch/mi300-mi200-performance-counters.html)). +- **XGMI:** category `xgmi` in `omnistat-inspect` / `stats_global.json` — intra-node FSDP traffic. +- **Network:** inter-node only relevant at N>1; see [README.md](README.md) §Two-node gate. + +Paste 2–3 bullets from `omnistat/report_summary.md`. + +## TraceLens (overlap / kernels) + +Paste top kernels / NCCL overlap narrative from `tracelens/report_summary.md`. + +## Bottleneck class (pick one dominant) + +- [ ] `gpu_compute` (MFMA / HBM read high vs peak) +- [ ] `gpu_memory_hbm` (VRAM saturated, step time grows with batch) +- [ ] `cpu_dispatch` / `dataloader` (low MFMA, host I/O or DataLoader gaps in trace) +- [ ] `comm_xgmi` (FSDP all-gather / reduce-scatter dominates on 1×8) +- [ ] `comm_scaleout` (only N>1) + +## bf16 vs fp32 discriminator (optional one-shot) + +Run one short job at **`ORBIT2_DATA_TYPE=float32`** with the **same** `ORBIT2_BATCH_SIZE`. If throughput barely changes → dispatch- or attention-path limited; if fp32 ≫ bf16 → memory-bandwidth sensitive. + +| dtype | throughput_samples_per_s | +|-------|--------------------------| +| bf16 | | +| fp32 | | + +## Levers justified (maps to `lever_catalog.yaml`) + +List 3–5 catalog `id`s with one-line rationale each, grounded in sections above + citations in the catalog. + +--- + +_Replace `` and tables after the first diagnosis run; keep this file under git for loop traceability._ diff --git a/earth_science/models/ORBIT-2/recipes/train/README.md b/earth_science/models/ORBIT-2/recipes/train/README.md index 53c1e31..1616795 100644 --- a/earth_science/models/ORBIT-2/recipes/train/README.md +++ b/earth_science/models/ORBIT-2/recipes/train/README.md @@ -10,7 +10,7 @@ This recipe summarizes the **exascale-oriented** workflow documented in [`XiaoWa ## High-level steps (from upstream) 1. **Environment** — On Frontier-class AMD nodes, follow the **Frontier** installation block in the ORBIT-2 README (Python 3.11 conda env, PyTorch ROCm build, `xformers`, `mpi4py`, `pip install -e .`). -2. **Choose a config** — Under `configs/`, pick a size (e.g. `interm_8m.yaml`, `interm_117m.yaml`, `interm_1b.yaml`, `interm_10b.yaml`). Larger models need more nodes/GPUs. +2. **Choose a config** — Under `configs/`, pick a size (e.g. `interm_8m.yaml`, `interm_117m.yaml`, `interm_1b.yaml`, `interm_10b.yaml`). Larger models need more nodes/GPUs. Those upstream configs set **`trainer.data_type: bfloat16`**; Bayes-CAST-style trees often ship ERA5 configs such as **`edm_8m_era5.yaml`** with the same dtype — Studio Lux templates (`interm_8m_lux_era5.yaml`) mirror that via `ORBIT2_DATA_TYPE` / `render_orbit2_config.py`. 3. **GPU and parallelism** — In YAML, set `trainer.gpu_type: "amd"`. Set `parallelism` fields (`fsdp`, `simple_ddp`, `tensor_par`, `seq_par`) so the product matches your **total GPU count** (see upstream comments). 4. **TILES** — For very large images, enable `tiling.do_tiling` and tune `div` / `overlap` per upstream guidance (patch divisibility matters). 5. **Data paths** — Populate `low_res_dir`, `high_res_dir`, `spatial_resolution`, and variable dictionaries to point at **your** staged dataset. Paths on Orion are **project- and allocation-specific**; see [../data/README.md](../data/README.md). From 6ad767418232cb551b1bb33c12101aafda86b3d6 Mon Sep 17 00:00:00 2001 From: Ashwin Aji Date: Wed, 17 Jun 2026 07:47:00 +0000 Subject: [PATCH 30/31] HydraGNN: align perf recipes with neutral perf-tools naming Update the perf-analysis and perf-optimizer-loop scripts and agent prompts to resolve perf tools via PERF_TOOLS_DIR (perf_tools.dir in .cluster-config.yaml), removing the prior site-specific naming. Co-Authored-By: Claude --- .../HydraGNN/examples/run_fom_extractor.py | 4 ++-- .../HydraGNN/examples/run_optimizer_loop.sh | 16 +++++++++++----- .../HydraGNN/examples/sbatch_train_perf_amd.sh | 16 ++++++++-------- .../HydraGNN/recipes/perf-analysis/README.md | 6 +++--- .../recipes/perf-analysis/agents/launcher.md | 12 ++++++------ .../perf-analysis/agents/omnistat_analyst.md | 8 ++++---- .../perf-analysis/agents/tracelens_analyst.md | 2 +- .../perf-analysis/omnistat.config.template | 2 +- .../perf-optimizer-loop/agents/fom_extractor.md | 4 ++-- .../perf-optimizer-loop/agents/orchestrator.md | 6 +++--- .../perf-optimizer-loop/agents/story_writer.md | 2 +- .../perf-optimizer-loop/dispatch-attribution.md | 4 ++-- .../perf-optimizer-loop/lever_catalog.yaml | 2 +- 13 files changed, 45 insertions(+), 39 deletions(-) diff --git a/material_science/models/HydraGNN/examples/run_fom_extractor.py b/material_science/models/HydraGNN/examples/run_fom_extractor.py index 6ec0d42..7565c36 100755 --- a/material_science/models/HydraGNN/examples/run_fom_extractor.py +++ b/material_science/models/HydraGNN/examples/run_fom_extractor.py @@ -6,7 +6,7 @@ Usage: export AI4S_SHARED_DIR=/shared/$USER - export OMNIHUB_TOOLS_DIR=/shared/omnihub/tools # from .cluster-config.yaml omnihub.tools_dir + export PERF_TOOLS_DIR=/path/to/perf-tools # from .cluster-config.yaml perf_tools.dir export REPO_ROOT=/path/to/ai4science-studio python run_fom_extractor.py --manifest /path/to/perf-runs//manifest.json \\ [--vm-port 8432] [--tsdb-url http://127.0.0.1:8432] @@ -140,7 +140,7 @@ def promql_scalar(tsdb_url: str, promql: str, t: int) -> float | None: def start_vm(db_path: str, port: int, log_path: Path) -> None: vm = os.environ.get( "VICTORIA_METRICS_BIN", - f"{os.environ['OMNIHUB_TOOLS_DIR']}/victoriametrics/victoria-metrics-prod", + f"{os.environ['PERF_TOOLS_DIR']}/victoriametrics/victoria-metrics-prod", ) subprocess.Popen( [ diff --git a/material_science/models/HydraGNN/examples/run_optimizer_loop.sh b/material_science/models/HydraGNN/examples/run_optimizer_loop.sh index aa38d29..a5b046e 100755 --- a/material_science/models/HydraGNN/examples/run_optimizer_loop.sh +++ b/material_science/models/HydraGNN/examples/run_optimizer_loop.sh @@ -46,8 +46,8 @@ ORCH_PROMPT="${RECIPE_DIR}/agents/orchestrator.md" # Loop-dir convention matches recipes/perf-optimizer-loop/README.md : "${AI4S_SHARED_DIR:?AI4S_SHARED_DIR must be set (e.g. export AI4S_SHARED_DIR=/shared/\$USER)}" # Shared perf-tool location (omnistat venv, VictoriaMetrics). Cluster-specific: -# set from .cluster-config.yaml omnihub.tools_dir (e.g. /shared/omnihub/tools). -: "${OMNIHUB_TOOLS_DIR:?OMNIHUB_TOOLS_DIR must be set (from .cluster-config.yaml omnihub.tools_dir, e.g. /shared/omnihub/tools)}" +# set from .cluster-config.yaml perf_tools.dir (e.g. /path/to/perf-tools). +: "${PERF_TOOLS_DIR:?PERF_TOOLS_DIR must be set (from .cluster-config.yaml perf_tools.dir, e.g. /path/to/perf-tools)}" HG_BASE="${AI4S_SHARED_DIR}/models/HydraGNN" PERF_RUNS_DIR="${HG_BASE}/perf-runs" LOOP_DIR="${PERF_RUNS_DIR}/loop-${LOOP_UUID}" @@ -96,8 +96,8 @@ fi # 3. tools present for _bin in \ - "${OMNIHUB_TOOLS_DIR}/omnihub-inspect/bin/omnistat-usermode" \ - "${OMNIHUB_TOOLS_DIR}/victoriametrics/victoria-metrics-prod"; do + "${PERF_TOOLS_DIR}/perf-inspect/bin/omnistat-usermode" \ + "${PERF_TOOLS_DIR}/victoriametrics/victoria-metrics-prod"; do if [[ -z "$PREFLIGHT_FAIL_REASON" && ! -x "$_bin" ]]; then PREFLIGHT_FAIL_REASON="tool_missing" PREFLIGHT_NOTES+=("missing: $_bin (run the perf-analysis launcher subagent first to install)") @@ -194,13 +194,19 @@ disk so you can resume from where you left off if I have to restart you. On any LOOP_COMPLETE, LOOP_ABORT, or LOOP_PAUSE, exit with a single line: STATUS=; reason=" +# Model for the orchestrator. The `opus` alias still points to claude-opus-4-7 on +# CLI 2.1.x, so default to the explicit 4.8 slug; override with HYDRAGNN_CLAUDE_MODEL. +HYDRAGNN_CLAUDE_MODEL="${HYDRAGNN_CLAUDE_MODEL:-claude-opus-4-8}" +_log_status "ORCH_MODEL model=${HYDRAGNN_CLAUDE_MODEL}" + _MAX_ATTEMPTS=3 _attempt=1 while [[ $_attempt -le $_MAX_ATTEMPTS ]]; do - _log_status "ORCH_INVOKE attempt=${_attempt} max=${_MAX_ATTEMPTS}" + _log_status "ORCH_INVOKE attempt=${_attempt} max=${_MAX_ATTEMPTS} model=${HYDRAGNN_CLAUDE_MODEL}" set +e claude \ --print \ + --model "$HYDRAGNN_CLAUDE_MODEL" \ --dangerously-skip-permissions \ --max-turns 250 \ --append-system-prompt "$(cat "$ORCH_PROMPT")" \ diff --git a/material_science/models/HydraGNN/examples/sbatch_train_perf_amd.sh b/material_science/models/HydraGNN/examples/sbatch_train_perf_amd.sh index 313c34f..cbf5c0e 100755 --- a/material_science/models/HydraGNN/examples/sbatch_train_perf_amd.sh +++ b/material_science/models/HydraGNN/examples/sbatch_train_perf_amd.sh @@ -21,8 +21,8 @@ # # Required: # AI4S_SHARED_DIR — shared base path -# ${OMNIHUB_TOOLS_DIR}/omnihub-inspect/ (created by the launcher subagent) -# ${OMNIHUB_TOOLS_DIR}/victoriametrics/ (created by the launcher subagent) +# ${PERF_TOOLS_DIR}/perf-inspect/ (created by the launcher subagent) +# ${PERF_TOOLS_DIR}/victoriametrics/ (created by the launcher subagent) # # Optional env-var overrides (with defaults): # HG_NUM_EPOCH=2 @@ -60,8 +60,8 @@ fi # --------------------------------------------------------------------------- HG_BASE="${AI4S_SHARED_DIR:?AI4S_SHARED_DIR must be set}/models/HydraGNN" # Shared perf-tool location (omnistat venv, omnistat-src, VictoriaMetrics). -# Cluster-specific: set from .cluster-config.yaml omnihub.tools_dir. -: "${OMNIHUB_TOOLS_DIR:?OMNIHUB_TOOLS_DIR must be set (from .cluster-config.yaml omnihub.tools_dir, e.g. /shared/omnihub/tools)}" +# Cluster-specific: set from .cluster-config.yaml perf_tools.dir. +: "${PERF_TOOLS_DIR:?PERF_TOOLS_DIR must be set (from .cluster-config.yaml perf_tools.dir, e.g. /path/to/perf-tools)}" HG_SIF="${HG_SIF:-${AI4S_SHARED_DIR}/images/pytorch_rocm7.2.2_ubuntu24.04_py3.12_pytorch_release_2.10.0.sif}" HG_OVERLAY="${HG_OVERLAY:-${HG_BASE}/overlays/hydragnn-overlay.img}" HG_DATASETS="${HG_DATASETS:-ANI1x,Alexandria}" @@ -74,7 +74,7 @@ HG_REPO_DIR="${HG_REPO_DIR:-${HG_BASE}/code/HydraGNN}" HYDRAGNN_MAX_NUM_BATCH="${HYDRAGNN_MAX_NUM_BATCH:-30}" PROFILE_TARGET_EPOCH="${PROFILE_TARGET_EPOCH:-1}" -OMNISTAT_VENV="${OMNISTAT_VENV:-${OMNIHUB_TOOLS_DIR}/omnihub-inspect}" +OMNISTAT_VENV="${OMNISTAT_VENV:-${PERF_TOOLS_DIR}/perf-inspect}" # Read the omnistat config from the in-repo recipe (single source of truth). # Override with OMNISTAT_TEMPLATE=/path/to/file if you need a custom probe config. # A staged copy under ${HG_BASE}/perf-runs/ is intentionally NOT used as the @@ -90,7 +90,7 @@ OMNISTAT_USERMODE_INTERVAL="${OMNISTAT_USERMODE_INTERVAL:-1}" # See material_science/models/HydraGNN/recipes/perf-analysis/agents/launcher.md # for how to build libomnistat_trace.so on a compute node. OMNISTAT_KERNEL_TRACE="${OMNISTAT_KERNEL_TRACE:-0}" -OMNISTAT_TRACE_LIB="${OMNISTAT_TRACE_LIB:-${OMNIHUB_TOOLS_DIR}/omnistat-src/build-trace/libomnistat_trace.so}" +OMNISTAT_TRACE_LIB="${OMNISTAT_TRACE_LIB:-${PERF_TOOLS_DIR}/omnistat-src/build-trace/libomnistat_trace.so}" # IMPORTANT: the kernel-trace collector registers /kernel_trace on the SAME # Flask app as the other omnistat endpoints, so the tool library must POST to # the [omnistat.collectors] `port` (8101 here), NOT the library's own default @@ -165,7 +165,7 @@ done mkdir -p "$HG_OUTPUT_DIR" OMNISTAT_CONFIG="${HG_OUTPUT_DIR}/omnistat.config" sed -e "s|@JOB_DIR@|${HG_OUTPUT_DIR}|g" \ - -e "s|@OMNIHUB_TOOLS_DIR@|${OMNIHUB_TOOLS_DIR}|g" \ + -e "s|@PERF_TOOLS_DIR@|${PERF_TOOLS_DIR}|g" \ "$OMNISTAT_TEMPLATE" > "$OMNISTAT_CONFIG" # Opt-in kernel tracing: flip enable_kernel_trace=True in the rendered config @@ -180,7 +180,7 @@ if [[ "$OMNISTAT_KERNEL_TRACE" == "1" ]]; then echo " Build it on a compute node:" >&2 echo " salloc -p -A -N 1 --time=00:15:00 --gpus-per-node=1 \\" >&2 echo " apptainer exec --rocm \"\$HG_SIF\" bash -c '" >&2 - echo " cd ${OMNIHUB_TOOLS_DIR}/omnistat-src && \\" >&2 + echo " cd ${PERF_TOOLS_DIR}/omnistat-src && \\" >&2 echo " cmake -S rocprofiler-sdk/ -B build-trace/ -DBUILD_KERNEL_TRACE_LIB=ON && \\" >&2 echo " cmake --build build-trace/ -j 8'" >&2 exit 2 diff --git a/material_science/models/HydraGNN/recipes/perf-analysis/README.md b/material_science/models/HydraGNN/recipes/perf-analysis/README.md index efa3aa0..0d0405e 100644 --- a/material_science/models/HydraGNN/recipes/perf-analysis/README.md +++ b/material_science/models/HydraGNN/recipes/perf-analysis/README.md @@ -68,8 +68,8 @@ The `combined_report.md` is the deliverable: ranked bottlenecks, remedies tried ## Prerequisites - Cluster access with the standard HydraGNN setup at `$AI4S_SHARED_DIR/models/HydraGNN/` — overlay, weights, code clone. See [HydraGNN/recipes/train/](../train/). -- Export `OMNIHUB_TOOLS_DIR` (the shared perf-tool location, from `.cluster-config.yaml` `omnihub.tools_dir`, e.g. `/shared/omnihub/tools`). Scripts and subagents resolve all tool paths from this var — no cluster path is hardcoded. -- One-time install of Omnistat (PR #271 branch `jorda/skills`, with `origin/main` merged in), VictoriaMetrics, and TraceLens — performed lazily by the launcher subagent into `${OMNIHUB_TOOLS_DIR}/`. +- Export `PERF_TOOLS_DIR` (the shared perf-tool location, from `.cluster-config.yaml` `perf_tools.dir`, e.g. `/path/to/perf-tools`). Scripts and subagents resolve all tool paths from this var — no cluster path is hardcoded. +- One-time install of Omnistat (PR #271 branch `jorda/skills`, with `origin/main` merged in), VictoriaMetrics, and TraceLens — performed lazily by the launcher subagent into `${PERF_TOOLS_DIR}/`. - The launcher writes `gfm_mlip_with_profile.json` at submit time (does not modify the upstream `gfm_mlip.json`). ## Telemetry knobs (set at `sbatch` submit time) @@ -77,7 +77,7 @@ The `combined_report.md` is the deliverable: ranked bottlenecks, remedies tried | Env var | Default | What it does | |---|---|---| | `OMNISTAT_KERNEL_TRACE` | `0` | When `1`, loads `libomnistat_trace.so` via `ROCP_TOOL_LIBRARIES` on every rank and turns on the omnistat kernel-trace collector. Adds per-kernel dispatch count + duration time series across all 8 cards on every node. Requires the library to be built once (see `agents/launcher.md` step 1e). Validated end-to-end on job 7034. | -| `OMNISTAT_TRACE_LIB` | `${OMNIHUB_TOOLS_DIR}/omnistat-src/build-trace/libomnistat_trace.so` | Override only for development. The wrapper hard-fails if the path is missing when `OMNISTAT_KERNEL_TRACE=1`. | +| `OMNISTAT_TRACE_LIB` | `${PERF_TOOLS_DIR}/omnistat-src/build-trace/libomnistat_trace.so` | Override only for development. The wrapper hard-fails if the path is missing when `OMNISTAT_KERNEL_TRACE=1`. | | `OMNISTAT_TRACE_LOG` | `1` | Library prints `[host][pid][omnistat] Trace summary: N/N processed records (M/M successful flushes)` on rank exit; useful as a smoke-signal that the tool initialized even when the workload crashed. | `enable_rocprofiler=True` is **on** in `omnistat.config.template` by default; the rendered config's state is printed in the sbatch banner so a silent "counters off" cannot recur. Kernel tracing and device-counting can co-exist on a single GPU but raise per-job VictoriaMetrics cardinality — pick by run length and the question you're asking. See `ai4science-studio` SKILL §12 for the device-counting vs kernel-trace vs `rocprofv3` decision matrix. diff --git a/material_science/models/HydraGNN/recipes/perf-analysis/agents/launcher.md b/material_science/models/HydraGNN/recipes/perf-analysis/agents/launcher.md index ff4afcf..457cf8d 100644 --- a/material_science/models/HydraGNN/recipes/perf-analysis/agents/launcher.md +++ b/material_science/models/HydraGNN/recipes/perf-analysis/agents/launcher.md @@ -10,8 +10,8 @@ Submits a 2-node HydraGNN training (AMD Instinct MI355X) with PyTorch profiling ## Outputs -- `${OMNIHUB_TOOLS_DIR}/omnihub-inspect/` — Python venv with omnistat (PR #271 + main merged) and TraceLens. -- `${OMNIHUB_TOOLS_DIR}/victoriametrics/victoria-metrics-prod` — VictoriaMetrics binary. +- `${PERF_TOOLS_DIR}/perf-inspect/` — Python venv with omnistat (PR #271 + main merged) and TraceLens. +- `${PERF_TOOLS_DIR}/victoriametrics/victoria-metrics-prod` — VictoriaMetrics binary. - `$AI4S_SHARED_DIR/models/HydraGNN/perf-runs/omnistat.config` — Omnistat user-mode config. - `$AI4S_SHARED_DIR/models/HydraGNN/perf-runs//manifest.json` — manifest schema below. - `$AI4S_SHARED_DIR/models/HydraGNN/perf-runs//hydragnn-train-.out` — symlinked from output dir. @@ -57,8 +57,8 @@ Submits a 2-node HydraGNN training (AMD Instinct MI355X) with PyTorch profiling ```bash set -euo pipefail -# Shared tool location — read from .cluster-config.yaml omnihub.tools_dir. -TOOLS="${OMNIHUB_TOOLS_DIR:?set OMNIHUB_TOOLS_DIR from .cluster-config.yaml omnihub.tools_dir (e.g. /shared/omnihub/tools)}" +# Shared tool location — read from .cluster-config.yaml perf_tools.dir. +TOOLS="${PERF_TOOLS_DIR:?set PERF_TOOLS_DIR from .cluster-config.yaml perf_tools.dir (e.g. /path/to/perf-tools)}" mkdir -p "$TOOLS" # 1a. Omnistat: jorda/skills branch with origin/main merged in @@ -77,7 +77,7 @@ fi OMNISTAT_COMMIT=$(git rev-parse HEAD) # 1b. Venv (py3.12 to match cluster python) -VENV=${OMNIHUB_TOOLS_DIR}/omnihub-inspect +VENV=${PERF_TOOLS_DIR}/perf-inspect if [[ ! -d $VENV ]]; then python3 -m venv "$VENV" fi @@ -123,7 +123,7 @@ if [[ "${OMNISTAT_KERNEL_TRACE:-0}" == "1" && ! -f "$TRACE_LIB" ]]; then fi ``` -The sbatch wrapper expects the `.so` at `${OMNIHUB_TOOLS_DIR}/omnistat-src/build-trace/libomnistat_trace.so` by default; override with `OMNISTAT_TRACE_LIB=/path/to/lib.so`. See `ai4science-studio` SKILL §12 for when to enable kernel tracing vs the default device-counting collector. +The sbatch wrapper expects the `.so` at `${PERF_TOOLS_DIR}/omnistat-src/build-trace/libomnistat_trace.so` by default; override with `OMNISTAT_TRACE_LIB=/path/to/lib.so`. See `ai4science-studio` SKILL §12 for when to enable kernel tracing vs the default device-counting collector. ### 2. Author the Omnistat user-mode config (once, in repo `perf-runs/`) diff --git a/material_science/models/HydraGNN/recipes/perf-analysis/agents/omnistat_analyst.md b/material_science/models/HydraGNN/recipes/perf-analysis/agents/omnistat_analyst.md index 56e9fef..5ff92cf 100644 --- a/material_science/models/HydraGNN/recipes/perf-analysis/agents/omnistat_analyst.md +++ b/material_science/models/HydraGNN/recipes/perf-analysis/agents/omnistat_analyst.md @@ -6,8 +6,8 @@ Drive `omnistat-inspect` (PR #271) through the analyze-job phases on the user-mo - `/manifest.json` - The Omnistat DB at `manifest.omnistat_db_path`. -- The `omnihub-inspect` venv at `${OMNIHUB_TOOLS_DIR}/omnihub-inspect/`. -- The VictoriaMetrics binary at `${OMNIHUB_TOOLS_DIR}/victoriametrics/victoria-metrics-prod`. +- The `perf-inspect` venv at `${PERF_TOOLS_DIR}/perf-inspect/`. +- The VictoriaMetrics binary at `${PERF_TOOLS_DIR}/victoriametrics/victoria-metrics-prod`. ## Outputs @@ -36,7 +36,7 @@ ss -lnt "sport = :$PORT" 2>/dev/null | grep -q LISTEN && PORT=8429 VMLOG=$PERF_RUN/omnistat/vm.log mkdir -p "$PERF_RUN/omnistat/inspect_outputs" -nohup ${OMNIHUB_TOOLS_DIR}/victoriametrics/victoria-metrics-prod \ +nohup ${PERF_TOOLS_DIR}/victoriametrics/victoria-metrics-prod \ -storageDataPath="$DB" \ -httpListenAddr=127.0.0.1:$PORT \ -retentionPeriod=100y \ @@ -62,7 +62,7 @@ done ### 2. Run analyze-job phases (per the SKILL) ```bash -. ${OMNIHUB_TOOLS_DIR}/omnihub-inspect/bin/activate +. ${PERF_TOOLS_DIR}/perf-inspect/bin/activate SCRATCH=$PERF_RUN/omnistat/scratch mkdir -p "$SCRATCH" JOBID=$(jq -r .jobid ) diff --git a/material_science/models/HydraGNN/recipes/perf-analysis/agents/tracelens_analyst.md b/material_science/models/HydraGNN/recipes/perf-analysis/agents/tracelens_analyst.md index 0aaff1d..99beea4 100644 --- a/material_science/models/HydraGNN/recipes/perf-analysis/agents/tracelens_analyst.md +++ b/material_science/models/HydraGNN/recipes/perf-analysis/agents/tracelens_analyst.md @@ -39,7 +39,7 @@ Every entry: ### 1. Sanity check ```bash -. ${OMNIHUB_TOOLS_DIR}/omnihub-inspect/bin/activate +. ${PERF_TOOLS_DIR}/perf-inspect/bin/activate PERF_RUN=$(jq -r .perf_run_dir ) TRACE=$(jq -r '.trace_paths[0] // empty' ) [[ -z "$TRACE" ]] && { echo "STATUS=partial; reason=no_trace_in_manifest"; exit 0; } diff --git a/material_science/models/HydraGNN/recipes/perf-analysis/omnistat.config.template b/material_science/models/HydraGNN/recipes/perf-analysis/omnistat.config.template index 80817de..d2fb3b9 100644 --- a/material_science/models/HydraGNN/recipes/perf-analysis/omnistat.config.template +++ b/material_science/models/HydraGNN/recipes/perf-analysis/omnistat.config.template @@ -120,7 +120,7 @@ system_name = mi355x #-- [omnistat.usermode] -victoria_binary = @OMNIHUB_TOOLS_DIR@/victoriametrics/victoria-metrics-prod +victoria_binary = @PERF_TOOLS_DIR@/victoriametrics/victoria-metrics-prod victoria_datadir = @JOB_DIR@/omnistat-db victoria_logfile = vic_server.log diff --git a/material_science/models/HydraGNN/recipes/perf-optimizer-loop/agents/fom_extractor.md b/material_science/models/HydraGNN/recipes/perf-optimizer-loop/agents/fom_extractor.md index c749537..9d97925 100644 --- a/material_science/models/HydraGNN/recipes/perf-optimizer-loop/agents/fom_extractor.md +++ b/material_science/models/HydraGNN/recipes/perf-optimizer-loop/agents/fom_extractor.md @@ -68,7 +68,7 @@ omnistat kernel-trace collector). ```bash set -euo pipefail -. "${OMNIHUB_TOOLS_DIR}/omnihub-inspect/bin/activate" +. "${PERF_TOOLS_DIR}/perf-inspect/bin/activate" PERF_RUN=$(jq -r .perf_run_dir ) JOBID=$(jq -r .jobid ) RUNTIME=$(jq -r .runtime_seconds ) @@ -102,7 +102,7 @@ Start a local VictoriaMetrics if not already running: PORT=8428 ss -lnt "sport = :$PORT" 2>/dev/null | grep -q LISTEN && PORT=8429 if ! curl -sf "http://127.0.0.1:$PORT/api/v1/status/tsdb" > /dev/null 2>&1; then - nohup "${OMNIHUB_TOOLS_DIR}/victoriametrics/victoria-metrics-prod" \ + nohup "${PERF_TOOLS_DIR}/victoriametrics/victoria-metrics-prod" \ -storageDataPath="$PERF_RUN/omnistat-db" \ -httpListenAddr=127.0.0.1:$PORT \ -retentionPeriod=100y -fs.disableMmap \ diff --git a/material_science/models/HydraGNN/recipes/perf-optimizer-loop/agents/orchestrator.md b/material_science/models/HydraGNN/recipes/perf-optimizer-loop/agents/orchestrator.md index 256addc..85cc223 100644 --- a/material_science/models/HydraGNN/recipes/perf-optimizer-loop/agents/orchestrator.md +++ b/material_science/models/HydraGNN/recipes/perf-optimizer-loop/agents/orchestrator.md @@ -85,11 +85,11 @@ Read `loop-/foms.csv` if it exists. If iter rows are present, set `current Run these checks in sequence; abort with `PREFLIGHT_FAIL reason=` on the first failure: -- Ensure `AI4S_SHARED_DIR` and `OMNIHUB_TOOLS_DIR` are exported (the latter from `.cluster-config.yaml` `omnihub.tools_dir`, e.g. `/shared/omnihub/tools`); abort if either is unset. All tool paths below derive from `$OMNIHUB_TOOLS_DIR` — never hardcode a cluster path. +- Ensure `AI4S_SHARED_DIR` and `PERF_TOOLS_DIR` are exported (the latter from `.cluster-config.yaml` `perf_tools.dir`, e.g. `/path/to/perf-tools`); abort if either is unset. All tool paths below derive from `$PERF_TOOLS_DIR` — never hardcode a cluster path. - `sinfo -p -h` (partition from `.cluster-config.yaml`) → if zero `IDLE` or `MIX` nodes, abort. Cluster may be in maintenance (see SKILL §14). - `df -h "$AI4S_SHARED_DIR" | tail -1` → abort if `Use%` > 95. -- `test -x ${OMNIHUB_TOOLS_DIR}/omnihub-inspect/bin/omnistat-usermode` → abort if missing. -- `test -x ${OMNIHUB_TOOLS_DIR}/victoriametrics/victoria-metrics-prod` → abort if missing. +- `test -x ${PERF_TOOLS_DIR}/perf-inspect/bin/omnistat-usermode` → abort if missing. +- `test -x ${PERF_TOOLS_DIR}/victoriametrics/victoria-metrics-prod` → abort if missing. - `test -f $AI4S_SHARED_DIR/models/HydraGNN/overlays/hydragnn-overlay.img` → abort if missing. - `test -f $AI4S_SHARED_DIR/images/pytorch_rocm7.2.2_ubuntu24.04_py3.12_pytorch_release_2.10.0.sif` → abort if missing. - If invoked via Claude Code CLI on the login node: `[[ -n "$ANTHROPIC_API_KEY" ]]` and `curl -fsS https://api.anthropic.com/v1/models > /dev/null` (5 s timeout) → abort if either fails. diff --git a/material_science/models/HydraGNN/recipes/perf-optimizer-loop/agents/story_writer.md b/material_science/models/HydraGNN/recipes/perf-optimizer-loop/agents/story_writer.md index 016e386..f296ca4 100644 --- a/material_science/models/HydraGNN/recipes/perf-optimizer-loop/agents/story_writer.md +++ b/material_science/models/HydraGNN/recipes/perf-optimizer-loop/agents/story_writer.md @@ -156,7 +156,7 @@ fig.tight_layout() fig.savefig(out_png_path, dpi=150, bbox_inches="tight") ``` -Use `matplotlib.use("Agg")` for headless rendering. Do NOT use seaborn (extra dep not in the omnistat venv). Use the omnistat venv at `${OMNIHUB_TOOLS_DIR}/omnihub-inspect/bin/python` which already has matplotlib via TraceLens deps. +Use `matplotlib.use("Agg")` for headless rendering. Do NOT use seaborn (extra dep not in the omnistat venv). Use the omnistat venv at `${PERF_TOOLS_DIR}/perf-inspect/bin/python` which already has matplotlib via TraceLens deps. ## Style rules diff --git a/material_science/models/HydraGNN/recipes/perf-optimizer-loop/dispatch-attribution.md b/material_science/models/HydraGNN/recipes/perf-optimizer-loop/dispatch-attribution.md index 21f57d5..ebedaf7 100644 --- a/material_science/models/HydraGNN/recipes/perf-optimizer-loop/dispatch-attribution.md +++ b/material_science/models/HydraGNN/recipes/perf-optimizer-loop/dispatch-attribution.md @@ -59,10 +59,10 @@ Short-kernel study (`short_kernels_summary.csv`): thousands of `aten::fill_`, `a ```bash export AI4S_SHARED_DIR=/shared/$USER export REPO_ROOT=$HOME/git/ai4science-studio -source ${OMNIHUB_TOOLS_DIR}/omnihub-inspect/bin/activate +source ${PERF_TOOLS_DIR}/perf-inspect/bin/activate # VictoriaMetrics on the perf-run DB (login node) -${OMNIHUB_TOOLS_DIR}/victoriametrics/victoria-metrics-prod \ +${PERF_TOOLS_DIR}/victoriametrics/victoria-metrics-prod \ -storageDataPath=$AI4S_SHARED_DIR/models/HydraGNN/perf-runs/7187/omnistat-db \ -httpListenAddr=127.0.0.1:8432 -retentionPeriod=100y -fs.disableMmap & diff --git a/material_science/models/HydraGNN/recipes/perf-optimizer-loop/lever_catalog.yaml b/material_science/models/HydraGNN/recipes/perf-optimizer-loop/lever_catalog.yaml index 6507f56..2543837 100644 --- a/material_science/models/HydraGNN/recipes/perf-optimizer-loop/lever_catalog.yaml +++ b/material_science/models/HydraGNN/recipes/perf-optimizer-loop/lever_catalog.yaml @@ -226,7 +226,7 @@ levers: risk: "blocked" revert_method: "drop_env" citation: "https://rocm.docs.amd.com/en/latest/how-to/tuning-guides/mi300x/workload.html#tunableop" - notes: "BLOCKED on ROCm 7.2.2 / PyTorch 2.10 / MI355X — warmup phase itself triggers GPU memory access faults. If the warmup phase did NOT complete (no CSVs), pick a different lever. Re-test only after a documented ROCm/hipBLASLt fix." + notes: "BLOCKED for HydraGNN (warmup phase triggers GPU memory access faults). NOTE (2026-06-13): this is WORKLOAD-specific, NOT a platform/hipBLASLt-prereq block. ORBIT-2 re-probed the identical MI355X/ROCm7.2.2/PyTorch2.10 stack with an isolated 1-GPU GEMM autotune (incl. large-M bf16 mm/F.linear) and TunableOp tuned CLEANLY with 0 faults (see ai4science-studio/SKILL.md #17, ai4science-perf-analysis/SKILL.md #28). The HydraGNN fault is likely tied to MACE double-backward shapes + 16-GPU distributed tuning. Keep blocked for HydraGNN until re-probed for THIS model; do not generalize to other models." - id: num_workers_12 description: "DataLoader: HYDRAGNN_NUM_WORKERS 8 -> 12 (persistent_workers already on). HIGH RUN-TO-RUN VARIANCE: helped in one loop (-12.4%) and regressed in another (+10.7%) — both on the same node pair, so this is run-to-run variance, not a systematic win. Treat as unreliable until replicated within ±5% on the same nodes." From 319f826e6dc642ad2fe2e064b3d63826777cf88e Mon Sep 17 00:00:00 2001 From: Ashwin Aji Date: Wed, 17 Jun 2026 07:58:35 +0000 Subject: [PATCH 31/31] omnihub: keep OMNIHUB_TOOLS_DIR naming after merging perf-optimizer-loop The perf branch renamed OMNIHUB_TOOLS_DIR -> PERF_TOOLS_DIR (and omnihub.tools_dir -> perf_tools.dir) for a generic public PR. On the omnihub integration branch the omnihub-named config is canonical, so normalize the merged perf-analysis / perf-optimizer-loop recipes back to OMNIHUB_TOOLS_DIR / omnihub.tools_dir / omnihub-inspect for internal consistency with the integrations/omnihub tooling. Co-Authored-By: Claude --- .claude/commands/init-cluster.md | 12 ++++++------ .cursor/skills/ai4science-studio/SKILL.md | 10 +++++----- .../ORBIT-2/examples/run_2node_scaleout_loop.sh | 4 ++-- .../models/ORBIT-2/examples/run_gemm_analysis.sh | 12 ++++++------ .../ORBIT-2/examples/run_optimizer_loop.sh | 8 ++++---- .../ORBIT-2/examples/sbatch_train_perf_amd.sh | 8 ++++---- .../examples/submit_perf_baseline_era5_amd.sh | 4 ++-- .../examples/sweep_orbit2_batch_bf16_amd.sh | 4 ++-- .../ORBIT-2/recipes/perf-analysis/HANDOFF.md | 2 +- .../ORBIT-2/recipes/perf-analysis/README.md | 2 +- .../recipes/perf-analysis/agents/launcher.md | 12 ++++++------ .../perf-analysis/agents/omnistat_analyst.md | 8 ++++---- .../agents/orchestrator_gemm_analysis.md | 4 ++-- .../perf-analysis/agents/tracelens_analyst.md | 2 +- .../perf-analysis/omnistat.config.template | 2 +- .../perf-analysis/one-node-gpu-baseline.md | 2 +- .../recipes/perf-optimizer-loop/README.md | 6 +++--- .../perf-optimizer-loop/agents/orchestrator.md | 4 ++-- .../HydraGNN/examples/run_fom_extractor.py | 4 ++-- .../HydraGNN/examples/run_optimizer_loop.sh | 8 ++++---- .../HydraGNN/examples/sbatch_train_perf_amd.sh | 16 ++++++++-------- .../HydraGNN/recipes/perf-analysis/README.md | 6 +++--- .../recipes/perf-analysis/agents/launcher.md | 12 ++++++------ .../perf-analysis/agents/omnistat_analyst.md | 8 ++++---- .../perf-analysis/agents/tracelens_analyst.md | 2 +- .../perf-analysis/omnistat.config.template | 2 +- .../perf-optimizer-loop/agents/fom_extractor.md | 4 ++-- .../perf-optimizer-loop/agents/orchestrator.md | 6 +++--- .../perf-optimizer-loop/agents/story_writer.md | 2 +- .../perf-optimizer-loop/dispatch-attribution.md | 4 ++-- 30 files changed, 90 insertions(+), 90 deletions(-) diff --git a/.claude/commands/init-cluster.md b/.claude/commands/init-cluster.md index 960d525..e87cfcb 100644 --- a/.claude/commands/init-cluster.md +++ b/.claude/commands/init-cluster.md @@ -69,9 +69,9 @@ echo "HTTPS_PROXY=${HTTPS_PROXY:-}" find "$HOME" /scratch /projects /opt -maxdepth 4 -name "*.sif" 2>/dev/null | head -20 # 14. Performance tooling (perf-analysis / perf-optimizer-loop recipes — optional) -# Look for an existing shared perf-tools dir (omnistat + TraceLens venv "perf-inspect/"). +# Look for an existing shared perf-tools dir (omnistat + TraceLens venv "omnihub-inspect/"). for d in /shared/*/tools /shared/perf-tools "$HOME/perf-tools"; do - [[ -d "$d/perf-inspect" || -d "$d/omnistat-src" ]] && echo "perf_tools: $d" + [[ -d "$d/omnihub-inspect" || -d "$d/omnistat-src" ]] && echo "perf_tools: $d" done ``` @@ -127,8 +127,8 @@ If no internet, ask for proxy settings. **Q9. Performance tooling** (optional — only for `perf-analysis` / `perf-optimizer-loop` recipes) If a perf-tools dir was discovered (step 14), show it and ask to confirm; otherwise ask for the path -or leave blank. This is exported to perf scripts as `PERF_TOOLS_DIR`; expected layout under it: -`perf-inspect/` (omnistat + TraceLens venv), `omnistat-src/`, `victoriametrics/victoria-metrics-prod`. +or leave blank. This is exported to perf scripts as `OMNIHUB_TOOLS_DIR`; expected layout under it: +`omnihub-inspect/` (omnistat + TraceLens venv), `omnistat-src/`, `victoriametrics/victoria-metrics-prod`. **Q10. Config file location** - `.cluster-config.yaml` in this repo (recommended — per-checkout) @@ -173,9 +173,9 @@ network: libionic_path: "" # path to libionic.so.1 (default: /usr/lib/x86_64-linux-gnu/libionic.so.1) # Performance tooling — only if using perf-analysis / perf-optimizer-loop recipes (Q9). -# Exported to scripts as PERF_TOOLS_DIR. Leave dir blank if not used. +# Exported to scripts as OMNIHUB_TOOLS_DIR. Leave dir blank if not used. perf_tools: - dir: "" # layout: perf-inspect/, omnistat-src/, victoriametrics/ + dir: "" # layout: omnihub-inspect/, omnistat-src/, victoriametrics/ discovered_sifs: # SIF files found during init (informational — confirm they are under sif_cache) diff --git a/.cursor/skills/ai4science-studio/SKILL.md b/.cursor/skills/ai4science-studio/SKILL.md index a07b2c0..76e7bf0 100755 --- a/.cursor/skills/ai4science-studio/SKILL.md +++ b/.cursor/skills/ai4science-studio/SKILL.md @@ -677,18 +677,18 @@ Omnistat exposes **two** independent rocprofiler-sdk integrations, and there's a **Cross-runtime gotcha:** `enable_kernel_trace=True` without the tool library loaded does *nothing* — the omnistat collector just sits with an empty endpoint. The `OMNISTAT_KERNEL_TRACE=1` switch in `sbatch_train_perf_amd.sh` enforces this with a hard `exit 2` if `libomnistat_trace.so` isn't found at the configured path; do **not** loosen that check, otherwise you re-create the exact "silent zero counters" failure mode from §8. -**Build location:** Both omnistat artifacts (the Python extension and the kernel-trace tool library) must be built on a **compute node** inside the same SIF the workload runs in — login nodes don't have apptainer or the ROCm headers, and a host-side build will pick up the wrong libcurl / glibc. `PERF_TOOLS_DIR` is the shared tools dir from `.cluster-config.yaml` `perf_tools.dir` (e.g. `/path/to/perf-tools`); export it first. Standard build: +**Build location:** Both omnistat artifacts (the Python extension and the kernel-trace tool library) must be built on a **compute node** inside the same SIF the workload runs in — login nodes don't have apptainer or the ROCm headers, and a host-side build will pick up the wrong libcurl / glibc. `OMNIHUB_TOOLS_DIR` is the shared tools dir from `.cluster-config.yaml` `omnihub.tools_dir` (e.g. `/shared/omnihub/tools`); export it first. Standard build: ```bash salloc -p -A -N 1 --time=00:15:00 --gpus-per-node=1 \ apptainer exec --rocm \ "${AI4S_SHARED_DIR}/images/pytorch_rocm7.2.2_ubuntu24.04_py3.12_pytorch_release_2.10.0.sif" \ - bash -c "cd ${PERF_TOOLS_DIR}/omnistat-src && \ + bash -c "cd ${OMNIHUB_TOOLS_DIR}/omnistat-src && \ cmake -S rocprofiler-sdk/ -B build-trace/ -DBUILD_KERNEL_TRACE_LIB=ON && \ cmake --build build-trace/ -j 8" ``` -Verify with `ls -la ${PERF_TOOLS_DIR}/omnistat-src/build-trace/libomnistat_trace.so` (~MB-class file, not the tiny 4 kB stub you get when CMake fails halfway). +Verify with `ls -la ${OMNIHUB_TOOLS_DIR}/omnistat-src/build-trace/libomnistat_trace.so` (~MB-class file, not the tiny 4 kB stub you get when CMake fails halfway). **SIF build prerequisites (one-time gotchas):** the `pytorch_rocm7.2.2_ubuntu24.04_py3.12_pytorch_release_2.10.0.sif` ships **without** `cmake` and **without** libcurl headers (only the `.so.4` runtime). Both omnistat artifacts need cmake; the kernel-trace library additionally needs ``. Working build recipe inside the SIF: @@ -737,12 +737,12 @@ When the cluster is up but you only need a one-shot interactive build inside the srun -p -A -N 1 --time=00:15:00 --gpus-per-node=1 --cpus-per-task=8 \ --job-name= \ apptainer exec --rocm \ - --bind ${PERF_TOOLS_DIR}:${PERF_TOOLS_DIR} \ + --bind ${OMNIHUB_TOOLS_DIR}:${OMNIHUB_TOOLS_DIR} \ "$HG_SIF" \ bash -c '' ``` -Why: `salloc` doesn't auto-bind `${PERF_TOOLS_DIR}/` into the apptainer namespace. Direct `srun … apptainer exec --bind …` forces the bind explicitly and writes the build artifact to its persistent NFS path in one shot. Verified during the kernel-trace library build (job 7030). +Why: `salloc` doesn't auto-bind `${OMNIHUB_TOOLS_DIR}/` into the apptainer namespace. Direct `srun … apptainer exec --bind …` forces the bind explicitly and writes the build artifact to its persistent NFS path in one shot. Verified during the kernel-trace library build (job 7030). ### 15. Dual-failure debug pattern — separate "tool wired correctly?" from "workload happy?" diff --git a/earth_science/models/ORBIT-2/examples/run_2node_scaleout_loop.sh b/earth_science/models/ORBIT-2/examples/run_2node_scaleout_loop.sh index 0f68adc..b5f98c2 100755 --- a/earth_science/models/ORBIT-2/examples/run_2node_scaleout_loop.sh +++ b/earth_science/models/ORBIT-2/examples/run_2node_scaleout_loop.sh @@ -26,7 +26,7 @@ SBATCH="${SCRIPT_DIR}/sbatch_train_perf_amd.sh" # ---- environment (override-able) ------------------------------------------- export AI4S_SHARED_DIR="${AI4S_SHARED_DIR:?AI4S_SHARED_DIR must be set}" -export PERF_TOOLS_DIR="${PERF_TOOLS_DIR:?PERF_TOOLS_DIR must be set (perf_tools.dir in .cluster-config.yaml)}" +export OMNIHUB_TOOLS_DIR="${OMNIHUB_TOOLS_DIR:?OMNIHUB_TOOLS_DIR must be set (omnihub.tools_dir in .cluster-config.yaml)}" export ORBIT2_ROOT="${ORBIT2_ROOT:-${AI4S_SHARED_DIR}/models/ORBIT-2/code/bayes-cast}" export ORBIT2_DATA_ROOT="${ORBIT2_DATA_ROOT:-${AI4S_SHARED_DIR}/models/ORBIT-2/data/superres/era5/1.0_deg}" export ORBIT2_CONFIG_TEMPLATE="${ORBIT2_CONFIG_TEMPLATE:-edm_8m_era5_1x8.yaml}" @@ -137,7 +137,7 @@ for i in "${!LEVER_NAME[@]}"; do EXP="ALL" EXP+=",ORBIT2_BATCH_SIZE=${BATCH},ORBIT2_NUM_WORKERS=1,ORBIT2_DATA_TYPE=bfloat16,ORBIT2_FUSED_ATTN=DEFAULT" EXP+=",ORBIT2_CONFIG_TEMPLATE=${ORBIT2_CONFIG_TEMPLATE},ORBIT2_MAX_EPOCH=${MAXEPOCH},ORBIT2_MAX_BATCHES=${MAXBATCH}" - EXP+=",ORBIT2_DATA_ROOT=${ORBIT2_DATA_ROOT},ORBIT2_ROOT=${ORBIT2_ROOT},AI4S_SHARED_DIR=${AI4S_SHARED_DIR},PERF_TOOLS_DIR=${PERF_TOOLS_DIR}" + EXP+=",ORBIT2_DATA_ROOT=${ORBIT2_DATA_ROOT},ORBIT2_ROOT=${ORBIT2_ROOT},AI4S_SHARED_DIR=${AI4S_SHARED_DIR},OMNIHUB_TOOLS_DIR=${OMNIHUB_TOOLS_DIR}" EXP+=",TORCH_NCCL_HIGH_PRIORITY=1,GPU_MAX_HW_QUEUES=2" for kv in $extra; do EXP+=",${kv}"; done diff --git a/earth_science/models/ORBIT-2/examples/run_gemm_analysis.sh b/earth_science/models/ORBIT-2/examples/run_gemm_analysis.sh index 5c96b9b..12bb322 100755 --- a/earth_science/models/ORBIT-2/examples/run_gemm_analysis.sh +++ b/earth_science/models/ORBIT-2/examples/run_gemm_analysis.sh @@ -7,7 +7,7 @@ # # Usage (tmux on login node): # export ANTHROPIC_API_KEY=... export AI4S_SHARED_DIR= -# export PERF_TOOLS_DIR= # perf_tools.dir in .cluster-config.yaml +# export OMNIHUB_TOOLS_DIR= # omnihub.tools_dir in .cluster-config.yaml # export EXCLUDE_NODES=... # optional: comma-separated nodes to skip # bash earth_science/models/ORBIT-2/examples/run_gemm_analysis.sh [--preflight-only] # @@ -23,8 +23,8 @@ REPO_ROOT=$(cd "${SCRIPT_DIR}/../../../.." && pwd) ORCH_PROMPT="${REPO_ROOT}/earth_science/models/ORBIT-2/recipes/perf-analysis/agents/orchestrator_gemm_analysis.md" : "${AI4S_SHARED_DIR:?AI4S_SHARED_DIR must be set}" -: "${PERF_TOOLS_DIR:?PERF_TOOLS_DIR must be set}" -export AI4S_SHARED_DIR PERF_TOOLS_DIR +: "${OMNIHUB_TOOLS_DIR:?OMNIHUB_TOOLS_DIR must be set}" +export AI4S_SHARED_DIR OMNIHUB_TOOLS_DIR ORBIT2_BASE="${AI4S_SHARED_DIR}/models/ORBIT-2" PERF_RUNS_DIR="${ORBIT2_BASE}/perf-runs" @@ -50,11 +50,11 @@ if command -v sinfo >/dev/null 2>&1; then else _notes+=("sinfo not in PATH"); fi _pct=$(df -P "$PERF_RUNS_DIR" 2>/dev/null | awk 'NR==2 {sub("%","",$5); print $5}') [[ -z "$_fail" && "${_pct:-0}" -ge 95 ]] && { _fail="disk_full"; _notes+=("perf-runs at ${_pct}%"); } -VENV="${PERF_TOOLS_DIR}/perf-inspect" +VENV="${OMNIHUB_TOOLS_DIR}/omnihub-inspect" for _f in "$ORCH_PROMPT" \ "${SCRIPT_DIR}/sbatch_train_perf_amd.sh" \ "${VENV}/bin/omnistat-usermode" \ - "${PERF_TOOLS_DIR}/victoriametrics/victoria-metrics-prod" \ + "${OMNIHUB_TOOLS_DIR}/victoriametrics/victoria-metrics-prod" \ "${AI4S_SHARED_DIR}/images/pytorch_rocm7.2.2_ubuntu24.04_py3.12_pytorch_release_2.10.0.sif"; do [[ -z "$_fail" && ! -e "$_f" ]] && { _fail="missing_file"; _notes+=("missing: $_f"); } done @@ -92,7 +92,7 @@ trap '_log "ANALYSIS_SIGNAL cleanup"; kill $_GUARD_PID 2>/dev/null || true' EXIT USER_PROMPT="Run the ORBIT-2 GEMM-time bottleneck analysis (TraceLens + Omnistat analyst/verifier) at 1 and 2 nodes. Follow ${ORCH_PROMPT}. REPO_ROOT=${REPO_ROOT} AI4S_SHARED_DIR=${AI4S_SHARED_DIR} -PERF_TOOLS_DIR=${PERF_TOOLS_DIR} +OMNIHUB_TOOLS_DIR=${OMNIHUB_TOOLS_DIR} ANALYSIS_DIR=${ANALYSIS_DIR} EXCLUDE_NODES=${EXCLUDE_NODES:-} Write GEMM_TIME_REPORT.md to ANALYSIS_DIR. Persist state for resume." diff --git a/earth_science/models/ORBIT-2/examples/run_optimizer_loop.sh b/earth_science/models/ORBIT-2/examples/run_optimizer_loop.sh index 7c59dda..0c2d829 100755 --- a/earth_science/models/ORBIT-2/examples/run_optimizer_loop.sh +++ b/earth_science/models/ORBIT-2/examples/run_optimizer_loop.sh @@ -4,7 +4,7 @@ # Usage (tmux on login node): # export ANTHROPIC_API_KEY=... # optional; only for Claude Code CLI driver # export AI4S_SHARED_DIR=... -# export PERF_TOOLS_DIR=... +# export OMNIHUB_TOOLS_DIR=... # bash earth_science/models/ORBIT-2/examples/run_optimizer_loop.sh [--preflight-only] # # Graceful stop: touch $AI4S_SHARED_DIR/models/ORBIT-2/perf-runs/loop-/STOP @@ -33,7 +33,7 @@ RECIPE_DIR="${REPO_ROOT}/earth_science/models/ORBIT-2/recipes/perf-optimizer-loo ORCH_PROMPT="${ORBIT2_ORCH_PROMPT:-${RECIPE_DIR}/agents/orchestrator.md}" : "${AI4S_SHARED_DIR:?AI4S_SHARED_DIR must be set}" -: "${PERF_TOOLS_DIR:?PERF_TOOLS_DIR must be set}" +: "${OMNIHUB_TOOLS_DIR:?OMNIHUB_TOOLS_DIR must be set}" ORBIT2_BASE="${AI4S_SHARED_DIR}/models/ORBIT-2" PERF_RUNS_DIR="${ORBIT2_BASE}/perf-runs" @@ -77,8 +77,8 @@ if [[ -z "$PREFLIGHT_FAIL_REASON" ]]; then fi for _bin in \ - "${PERF_TOOLS_DIR}/perf-inspect/bin/omnistat-usermode" \ - "${PERF_TOOLS_DIR}/victoriametrics/victoria-metrics-prod"; do + "${OMNIHUB_TOOLS_DIR}/omnihub-inspect/bin/omnistat-usermode" \ + "${OMNIHUB_TOOLS_DIR}/victoriametrics/victoria-metrics-prod"; do if [[ -z "$PREFLIGHT_FAIL_REASON" && ! -x "$_bin" ]]; then PREFLIGHT_FAIL_REASON="tool_missing" PREFLIGHT_NOTES+=("missing: $_bin") diff --git a/earth_science/models/ORBIT-2/examples/sbatch_train_perf_amd.sh b/earth_science/models/ORBIT-2/examples/sbatch_train_perf_amd.sh index 8cd5f03..e6158f4 100755 --- a/earth_science/models/ORBIT-2/examples/sbatch_train_perf_amd.sh +++ b/earth_science/models/ORBIT-2/examples/sbatch_train_perf_amd.sh @@ -6,7 +6,7 @@ # # Quick start (ERA5 same-dir): # export AI4S_SHARED_DIR=/path/to/shared -# export PERF_TOOLS_DIR=/path/to/perf-tools +# export OMNIHUB_TOOLS_DIR=/shared/omnihub/tools # export ORBIT2_DATA_ROOT=$AI4S_SHARED_DIR/models/ORBIT-2/data/superres/era5/1.0_deg # export ORBIT2_CONFIG_TEMPLATE=edm_8m_era5_1x8.yaml # export ORBIT2_ROOT=$AI4S_SHARED_DIR/models/ORBIT-2/code/bayes-cast # EDM preset + launcher @@ -72,8 +72,8 @@ ORBIT2_DISABLE_CKPT="${ORBIT2_DISABLE_CKPT:-1}" ORBIT2_TUNABLEOP_MODE="${ORBIT2_TUNABLEOP_MODE:-off}" ORBIT2_TUNABLEOP_DIR="${ORBIT2_TUNABLEOP_DIR:-${ORBIT2_BASE}/tunableop-cache}" -: "${PERF_TOOLS_DIR:?PERF_TOOLS_DIR must be set}" -OMNISTAT_VENV="${OMNISTAT_VENV:-${PERF_TOOLS_DIR}/perf-inspect}" +: "${OMNIHUB_TOOLS_DIR:?OMNIHUB_TOOLS_DIR must be set}" +OMNISTAT_VENV="${OMNISTAT_VENV:-${OMNIHUB_TOOLS_DIR}/omnihub-inspect}" OMNISTAT_TEMPLATE="${OMNISTAT_TEMPLATE:-${SCRIPT_DIR}/../recipes/perf-analysis/omnistat.config.template}" OMNISTAT_USERMODE_INTERVAL="${OMNISTAT_USERMODE_INTERVAL:-1}" @@ -141,7 +141,7 @@ fi mkdir -p "$ORBIT2_OUTPUT_DIR" OMNISTAT_CONFIG="${ORBIT2_OUTPUT_DIR}/omnistat.config" sed -e "s|@JOB_DIR@|${ORBIT2_OUTPUT_DIR}|g" \ - -e "s|@PERF_TOOLS_DIR@|${PERF_TOOLS_DIR}|g" \ + -e "s|@OMNIHUB_TOOLS_DIR@|${OMNIHUB_TOOLS_DIR}|g" \ "$OMNISTAT_TEMPLATE" > "$OMNISTAT_CONFIG" if [[ -n "${OMNISTAT_ROCPROF_PROFILE:-}" ]]; then sed -i.bak "s/^profile = .*/profile = ${OMNISTAT_ROCPROF_PROFILE}/" "$OMNISTAT_CONFIG" && rm -f "${OMNISTAT_CONFIG}.bak" diff --git a/earth_science/models/ORBIT-2/examples/submit_perf_baseline_era5_amd.sh b/earth_science/models/ORBIT-2/examples/submit_perf_baseline_era5_amd.sh index ca4cd21..654f1d8 100755 --- a/earth_science/models/ORBIT-2/examples/submit_perf_baseline_era5_amd.sh +++ b/earth_science/models/ORBIT-2/examples/submit_perf_baseline_era5_amd.sh @@ -4,14 +4,14 @@ # # Usage (from anywhere): # export AI4S_SHARED_DIR=/path/to/shared -# export PERF_TOOLS_DIR=/path/to/perf-tools +# export OMNIHUB_TOOLS_DIR=/shared/omnihub/tools # ./submit_perf_baseline_era5_amd.sh # Extra sbatch flags: ./submit_perf_baseline_era5_amd.sh --partition=lux --account=myacct set -euo pipefail SCRIPT_DIR=$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd) : "${AI4S_SHARED_DIR:?set AI4S_SHARED_DIR}" -: "${PERF_TOOLS_DIR:?set PERF_TOOLS_DIR}" +: "${OMNIHUB_TOOLS_DIR:?set OMNIHUB_TOOLS_DIR}" export ORBIT2_DATA_ROOT="${ORBIT2_DATA_ROOT:-${AI4S_SHARED_DIR}/models/ORBIT-2/data/superres/era5/1.0_deg}" export ORBIT2_CONFIG_TEMPLATE="${ORBIT2_CONFIG_TEMPLATE:-edm_8m_era5_1x8.yaml}" diff --git a/earth_science/models/ORBIT-2/examples/sweep_orbit2_batch_bf16_amd.sh b/earth_science/models/ORBIT-2/examples/sweep_orbit2_batch_bf16_amd.sh index 7c30626..1ba62fa 100755 --- a/earth_science/models/ORBIT-2/examples/sweep_orbit2_batch_bf16_amd.sh +++ b/earth_science/models/ORBIT-2/examples/sweep_orbit2_batch_bf16_amd.sh @@ -3,7 +3,7 @@ # # Usage (login node, repo root): # export AI4S_SHARED_DIR=... -# export PERF_TOOLS_DIR=... +# export OMNIHUB_TOOLS_DIR=... # export ORBIT2_ROOT=$AI4S_SHARED_DIR/models/ORBIT-2/code/bayes-cast # export ORBIT2_DATA_ROOT=$AI4S_SHARED_DIR/models/ORBIT-2/data/superres/era5/1.0_deg # # Optional: export ORBIT2_ERA5_SPATIAL_RES=111 @@ -27,7 +27,7 @@ if [[ $# -lt 1 ]]; then fi : "${AI4S_SHARED_DIR:?set AI4S_SHARED_DIR}" -: "${PERF_TOOLS_DIR:?set PERF_TOOLS_DIR}" +: "${OMNIHUB_TOOLS_DIR:?set OMNIHUB_TOOLS_DIR}" export ORBIT2_DATA_TYPE="${ORBIT2_DATA_TYPE:-bfloat16}" export ORBIT2_FUSED_ATTN="${ORBIT2_FUSED_ATTN:-DEFAULT}" diff --git a/earth_science/models/ORBIT-2/recipes/perf-analysis/HANDOFF.md b/earth_science/models/ORBIT-2/recipes/perf-analysis/HANDOFF.md index 9669e31..65a3682 100644 --- a/earth_science/models/ORBIT-2/recipes/perf-analysis/HANDOFF.md +++ b/earth_science/models/ORBIT-2/recipes/perf-analysis/HANDOFF.md @@ -299,7 +299,7 @@ Expect **`math`** and often **`default`** / **`efficient`** to complete on ROCm; - Overlay: `$AI4S_SHARED_DIR/models/ORBIT-2/overlays/orbit2-overlay.img` - Perf runs: `$AI4S_SHARED_DIR/models/ORBIT-2/perf-runs//` - Training logs (recommended): `$AI4S_SHARED_DIR/models/ORBIT-2/outputs/train/logs/` -- Omnistat / TraceLens tools: set **`PERF_TOOLS_DIR`** to your site-local checkout (see root `.cursor/skills/ai4science-perf-analysis/SKILL.md` and `.cluster-config.yaml` key `perf_tools.dir` if your site maintains one). +- Omnistat / TraceLens tools: set **`OMNIHUB_TOOLS_DIR`** to your site-local checkout (see root `.cursor/skills/ai4science-perf-analysis/SKILL.md` and `.cluster-config.yaml` key `omnihub.tools_dir` if your site maintains one). ## Operational source tree diff --git a/earth_science/models/ORBIT-2/recipes/perf-analysis/README.md b/earth_science/models/ORBIT-2/recipes/perf-analysis/README.md index a995452..a6c1336 100644 --- a/earth_science/models/ORBIT-2/recipes/perf-analysis/README.md +++ b/earth_science/models/ORBIT-2/recipes/perf-analysis/README.md @@ -10,7 +10,7 @@ Multi-subagent workflow for ORBIT-2 / Bayes-CAST training on AMD MI355X, with Om ```bash export AI4S_SHARED_DIR=/path/to/shared -export PERF_TOOLS_DIR=/path/to/perf-tools +export OMNIHUB_TOOLS_DIR=/shared/omnihub/tools # ORBIT2_ROOT defaults to .../code/bayes-cast when that directory exists export ORBIT2_DATA_ROOT=$AI4S_SHARED_DIR/models/ORBIT-2/data/superres/era5/1.0_deg # export ORBIT2_CONFIG_TEMPLATE=edm_8m_era5_1x8.yaml # default diff --git a/earth_science/models/ORBIT-2/recipes/perf-analysis/agents/launcher.md b/earth_science/models/ORBIT-2/recipes/perf-analysis/agents/launcher.md index 107863c..fa423cd 100644 --- a/earth_science/models/ORBIT-2/recipes/perf-analysis/agents/launcher.md +++ b/earth_science/models/ORBIT-2/recipes/perf-analysis/agents/launcher.md @@ -10,8 +10,8 @@ Submits a 2-node ORBIT-2 training (AMD Instinct MI355X) with PyTorch profiling a ## Outputs -- `${PERF_TOOLS_DIR}/perf-inspect/` — Python venv with omnistat (PR #271 + main merged) and TraceLens. -- `${PERF_TOOLS_DIR}/victoriametrics/victoria-metrics-prod` — VictoriaMetrics binary. +- `${OMNIHUB_TOOLS_DIR}/omnihub-inspect/` — Python venv with omnistat (PR #271 + main merged) and TraceLens. +- `${OMNIHUB_TOOLS_DIR}/victoriametrics/victoria-metrics-prod` — VictoriaMetrics binary. - `$AI4S_SHARED_DIR/models/ORBIT-2/perf-runs/omnistat.config` — Omnistat user-mode config. - `$AI4S_SHARED_DIR/models/ORBIT-2/perf-runs//manifest.json` — manifest schema below. - `$AI4S_SHARED_DIR/models/ORBIT-2/perf-runs//orbit2-train-.out` — symlinked from output dir. @@ -57,8 +57,8 @@ Submits a 2-node ORBIT-2 training (AMD Instinct MI355X) with PyTorch profiling a ```bash set -euo pipefail -# Shared tool location — read from .cluster-config.yaml perf_tools.dir. -TOOLS="${PERF_TOOLS_DIR:?set PERF_TOOLS_DIR to a persistent checkout (e.g. under \$AI4S_SHARED_DIR/tools)}" +# Shared tool location — read from .cluster-config.yaml omnihub.tools_dir. +TOOLS="${OMNIHUB_TOOLS_DIR:?set OMNIHUB_TOOLS_DIR to a persistent checkout (e.g. under \$AI4S_SHARED_DIR/tools)}" mkdir -p "$TOOLS" # 1a. Omnistat: jorda/skills branch with origin/main merged in @@ -77,7 +77,7 @@ fi OMNISTAT_COMMIT=$(git rev-parse HEAD) # 1b. Venv (py3.12 to match cluster python) -VENV=${PERF_TOOLS_DIR}/perf-inspect +VENV=${OMNIHUB_TOOLS_DIR}/omnihub-inspect if [[ ! -d $VENV ]]; then python3 -m venv "$VENV" fi @@ -123,7 +123,7 @@ if [[ "${OMNISTAT_KERNEL_TRACE:-0}" == "1" && ! -f "$TRACE_LIB" ]]; then fi ``` -The sbatch wrapper expects the `.so` at `${PERF_TOOLS_DIR}/omnistat-src/build-trace/libomnistat_trace.so` by default; override with `OMNISTAT_TRACE_LIB=/path/to/lib.so`. See `ai4science-studio` SKILL §12 for when to enable kernel tracing vs the default device-counting collector. +The sbatch wrapper expects the `.so` at `${OMNIHUB_TOOLS_DIR}/omnistat-src/build-trace/libomnistat_trace.so` by default; override with `OMNISTAT_TRACE_LIB=/path/to/lib.so`. See `ai4science-studio` SKILL §12 for when to enable kernel tracing vs the default device-counting collector. ### 2. Author the Omnistat user-mode config (once, in repo `perf-runs/`) diff --git a/earth_science/models/ORBIT-2/recipes/perf-analysis/agents/omnistat_analyst.md b/earth_science/models/ORBIT-2/recipes/perf-analysis/agents/omnistat_analyst.md index 5ff92cf..56e9fef 100644 --- a/earth_science/models/ORBIT-2/recipes/perf-analysis/agents/omnistat_analyst.md +++ b/earth_science/models/ORBIT-2/recipes/perf-analysis/agents/omnistat_analyst.md @@ -6,8 +6,8 @@ Drive `omnistat-inspect` (PR #271) through the analyze-job phases on the user-mo - `/manifest.json` - The Omnistat DB at `manifest.omnistat_db_path`. -- The `perf-inspect` venv at `${PERF_TOOLS_DIR}/perf-inspect/`. -- The VictoriaMetrics binary at `${PERF_TOOLS_DIR}/victoriametrics/victoria-metrics-prod`. +- The `omnihub-inspect` venv at `${OMNIHUB_TOOLS_DIR}/omnihub-inspect/`. +- The VictoriaMetrics binary at `${OMNIHUB_TOOLS_DIR}/victoriametrics/victoria-metrics-prod`. ## Outputs @@ -36,7 +36,7 @@ ss -lnt "sport = :$PORT" 2>/dev/null | grep -q LISTEN && PORT=8429 VMLOG=$PERF_RUN/omnistat/vm.log mkdir -p "$PERF_RUN/omnistat/inspect_outputs" -nohup ${PERF_TOOLS_DIR}/victoriametrics/victoria-metrics-prod \ +nohup ${OMNIHUB_TOOLS_DIR}/victoriametrics/victoria-metrics-prod \ -storageDataPath="$DB" \ -httpListenAddr=127.0.0.1:$PORT \ -retentionPeriod=100y \ @@ -62,7 +62,7 @@ done ### 2. Run analyze-job phases (per the SKILL) ```bash -. ${PERF_TOOLS_DIR}/perf-inspect/bin/activate +. ${OMNIHUB_TOOLS_DIR}/omnihub-inspect/bin/activate SCRATCH=$PERF_RUN/omnistat/scratch mkdir -p "$SCRATCH" JOBID=$(jq -r .jobid ) diff --git a/earth_science/models/ORBIT-2/recipes/perf-analysis/agents/orchestrator_gemm_analysis.md b/earth_science/models/ORBIT-2/recipes/perf-analysis/agents/orchestrator_gemm_analysis.md index c4d6d0d..4e3acba 100644 --- a/earth_science/models/ORBIT-2/recipes/perf-analysis/agents/orchestrator_gemm_analysis.md +++ b/earth_science/models/ORBIT-2/recipes/perf-analysis/agents/orchestrator_gemm_analysis.md @@ -12,7 +12,7 @@ GEMM time is actually spent on** (which shapes/kernels dominate, compute- vs mem and what changes from 1→2 nodes), to decide the next real lever. ## Fixed context (passed in the user prompt) -- `REPO_ROOT`, `AI4S_SHARED_DIR`, `PERF_TOOLS_DIR` (`perf_tools.dir` in `.cluster-config.yaml`). +- `REPO_ROOT`, `AI4S_SHARED_DIR`, `OMNIHUB_TOOLS_DIR` (`omnihub.tools_dir` in `.cluster-config.yaml`). - SLURM: partition and account from `.cluster-config.yaml` (`slurm.partition`, `slurm.account`). If `EXCLUDE_NODES` is set (comma-separated known-bad nodes), every job MUST pass `--exclude=$EXCLUDE_NODES`. @@ -53,7 +53,7 @@ bf16, SDPA DEFAULT, `ORBIT2_BATCH_SIZE=4096`, `ORBIT2_MAX_EPOCH=6`, `sbatch --parsable --partition= --account= --nodes= [--exclude=$EXCLUDE_NODES] --job-name=o2-gemm-n - --export=ALL,AI4S_SHARED_DIR=...,PERF_TOOLS_DIR=...,ORBIT2_DATA_ROOT=...,ORBIT2_ROOT=..., + --export=ALL,AI4S_SHARED_DIR=...,OMNIHUB_TOOLS_DIR=...,ORBIT2_DATA_ROOT=...,ORBIT2_ROOT=..., ORBIT2_CONFIG_TEMPLATE=edm_8m_era5_1x8.yaml,ORBIT2_ERA5_SPATIAL_RES=111,ORBIT2_BATCH_SIZE=4096, ORBIT2_MAX_EPOCH=6,TORCH_NCCL_HIGH_PRIORITY=1,GPU_MAX_HW_QUEUES=2, PROFILE_TARGET_EPOCH=2,PROFILE_RANK0_ONLY=1 `. diff --git a/earth_science/models/ORBIT-2/recipes/perf-analysis/agents/tracelens_analyst.md b/earth_science/models/ORBIT-2/recipes/perf-analysis/agents/tracelens_analyst.md index edda6c3..e37f2c5 100644 --- a/earth_science/models/ORBIT-2/recipes/perf-analysis/agents/tracelens_analyst.md +++ b/earth_science/models/ORBIT-2/recipes/perf-analysis/agents/tracelens_analyst.md @@ -39,7 +39,7 @@ Every entry: ### 1. Sanity check ```bash -. ${PERF_TOOLS_DIR}/perf-inspect/bin/activate +. ${OMNIHUB_TOOLS_DIR}/omnihub-inspect/bin/activate PERF_RUN=$(jq -r .perf_run_dir ) TRACE=$(jq -r '.trace_paths[0] // empty' ) [[ -z "$TRACE" ]] && { echo "STATUS=partial; reason=no_trace_in_manifest"; exit 0; } diff --git a/earth_science/models/ORBIT-2/recipes/perf-analysis/omnistat.config.template b/earth_science/models/ORBIT-2/recipes/perf-analysis/omnistat.config.template index 6bbb78c..a79702a 100644 --- a/earth_science/models/ORBIT-2/recipes/perf-analysis/omnistat.config.template +++ b/earth_science/models/ORBIT-2/recipes/perf-analysis/omnistat.config.template @@ -134,7 +134,7 @@ system_name = mi355x #-- [omnistat.usermode] -victoria_binary = @PERF_TOOLS_DIR@/victoriametrics/victoria-metrics-prod +victoria_binary = @OMNIHUB_TOOLS_DIR@/victoriametrics/victoria-metrics-prod victoria_datadir = @JOB_DIR@/omnistat-db victoria_logfile = vic_server.log diff --git a/earth_science/models/ORBIT-2/recipes/perf-analysis/one-node-gpu-baseline.md b/earth_science/models/ORBIT-2/recipes/perf-analysis/one-node-gpu-baseline.md index a9ff0d6..9b98de2 100644 --- a/earth_science/models/ORBIT-2/recipes/perf-analysis/one-node-gpu-baseline.md +++ b/earth_science/models/ORBIT-2/recipes/perf-analysis/one-node-gpu-baseline.md @@ -25,7 +25,7 @@ Goal: on **one MI355-class node × 8 GPUs**, pick a **per-rank batch size** and ```bash export AI4S_SHARED_DIR=... -export PERF_TOOLS_DIR=... +export OMNIHUB_TOOLS_DIR=... export ORBIT2_ROOT=... # optional; defaults to $AI4S_SHARED_DIR/models/ORBIT-2/code/ORBIT-2 # Defaults pick ERA5 path + template; omit these if data is already staged there: # export ORBIT2_DATA_ROOT=$AI4S_SHARED_DIR/models/ORBIT-2/data/superres/era5/1.0_deg diff --git a/earth_science/models/ORBIT-2/recipes/perf-optimizer-loop/README.md b/earth_science/models/ORBIT-2/recipes/perf-optimizer-loop/README.md index 5aa3273..4b85473 100644 --- a/earth_science/models/ORBIT-2/recipes/perf-optimizer-loop/README.md +++ b/earth_science/models/ORBIT-2/recipes/perf-optimizer-loop/README.md @@ -10,14 +10,14 @@ Agent-driven optimization loop for **ORBIT-2** training with **throughput (`thro - **Container:** `pytorch_rocm7.2.2_ubuntu24.04_py3.12_pytorch_release_2.10.0.sif` + ORBIT-2 overlay (`build_overlay_amd.sh`). - **Data:** staged ERA5 NPZ tree; for HBM saturation follow [`STAGING_ERA5_FOR_HBM.md`](STAGING_ERA5_FOR_HBM.md). - **Code:** `ORBIT2_ROOT` → Bayes-CAST with `launch/train_edm.py` (see [HANDOFF.md](../perf-analysis/HANDOFF.md)). -- **Tools:** `PERF_TOOLS_DIR` (omnistat-usermode + VictoriaMetrics), `AI4S_SHARED_DIR`. +- **Tools:** `OMNIHUB_TOOLS_DIR` (omnistat-usermode + VictoriaMetrics), `AI4S_SHARED_DIR`. ## Environment variables (loop driver) | Variable | Required | Description | |----------|----------|-------------| | `AI4S_SHARED_DIR` | Yes | Shared storage root | -| `PERF_TOOLS_DIR` | Yes | Omnistat + VM binaries | +| `OMNIHUB_TOOLS_DIR` | Yes | Omnistat + VM binaries | | `ORBIT2_ROOT` | Recommended | Bayes-CAST checkout | | `ORBIT2_DATA_ROOT` | Yes | ERA5 staging path | | `ORBIT2_BATCH_SIZE` | Yes (baseline) | Locked after VRAM sweep | @@ -29,7 +29,7 @@ Agent-driven optimization loop for **ORBIT-2** training with **throughput (`thro ```bash export AI4S_SHARED_DIR=... -export PERF_TOOLS_DIR=... +export OMNIHUB_TOOLS_DIR=... bash earth_science/models/ORBIT-2/examples/validate_orbit2_optimizer_loop_recipe.sh tmux new -s orbit2-loop diff --git a/earth_science/models/ORBIT-2/recipes/perf-optimizer-loop/agents/orchestrator.md b/earth_science/models/ORBIT-2/recipes/perf-optimizer-loop/agents/orchestrator.md index 67fc129..eb2cb17 100644 --- a/earth_science/models/ORBIT-2/recipes/perf-optimizer-loop/agents/orchestrator.md +++ b/earth_science/models/ORBIT-2/recipes/perf-optimizer-loop/agents/orchestrator.md @@ -8,7 +8,7 @@ Drives the ORBIT-2 (Bayes-CAST EDM) optimizer loop on **MI355X-class** hardware. - Loop args: ``, ``. - `REPO_ROOT` (cwd or env). -- `.cluster-config.yaml` (partition, perf_tools.dir). +- `.cluster-config.yaml` (partition, omnihub.tools_dir). - [`../lever_catalog.yaml`](../lever_catalog.yaml). ## Outputs @@ -32,7 +32,7 @@ Per job: existing perf-analysis layout under `..//` plus `foms.json` from ## Preflight (ORBIT-2) -- `AI4S_SHARED_DIR`, `PERF_TOOLS_DIR` set. +- `AI4S_SHARED_DIR`, `OMNIHUB_TOOLS_DIR` set. - `${AI4S_SHARED_DIR}/models/ORBIT-2/overlays/orbit2-overlay.img` exists. - `${AI4S_SHARED_DIR}/images/pytorch_rocm7.2.2_ubuntu24.04_py3.12_pytorch_release_2.10.0.sif` (or `ORBIT2_SIF`) exists. - `omnistat-usermode` + `victoria-metrics-prod` binaries present. diff --git a/material_science/models/HydraGNN/examples/run_fom_extractor.py b/material_science/models/HydraGNN/examples/run_fom_extractor.py index 7565c36..6ec0d42 100755 --- a/material_science/models/HydraGNN/examples/run_fom_extractor.py +++ b/material_science/models/HydraGNN/examples/run_fom_extractor.py @@ -6,7 +6,7 @@ Usage: export AI4S_SHARED_DIR=/shared/$USER - export PERF_TOOLS_DIR=/path/to/perf-tools # from .cluster-config.yaml perf_tools.dir + export OMNIHUB_TOOLS_DIR=/shared/omnihub/tools # from .cluster-config.yaml omnihub.tools_dir export REPO_ROOT=/path/to/ai4science-studio python run_fom_extractor.py --manifest /path/to/perf-runs//manifest.json \\ [--vm-port 8432] [--tsdb-url http://127.0.0.1:8432] @@ -140,7 +140,7 @@ def promql_scalar(tsdb_url: str, promql: str, t: int) -> float | None: def start_vm(db_path: str, port: int, log_path: Path) -> None: vm = os.environ.get( "VICTORIA_METRICS_BIN", - f"{os.environ['PERF_TOOLS_DIR']}/victoriametrics/victoria-metrics-prod", + f"{os.environ['OMNIHUB_TOOLS_DIR']}/victoriametrics/victoria-metrics-prod", ) subprocess.Popen( [ diff --git a/material_science/models/HydraGNN/examples/run_optimizer_loop.sh b/material_science/models/HydraGNN/examples/run_optimizer_loop.sh index a5b046e..37ca554 100755 --- a/material_science/models/HydraGNN/examples/run_optimizer_loop.sh +++ b/material_science/models/HydraGNN/examples/run_optimizer_loop.sh @@ -46,8 +46,8 @@ ORCH_PROMPT="${RECIPE_DIR}/agents/orchestrator.md" # Loop-dir convention matches recipes/perf-optimizer-loop/README.md : "${AI4S_SHARED_DIR:?AI4S_SHARED_DIR must be set (e.g. export AI4S_SHARED_DIR=/shared/\$USER)}" # Shared perf-tool location (omnistat venv, VictoriaMetrics). Cluster-specific: -# set from .cluster-config.yaml perf_tools.dir (e.g. /path/to/perf-tools). -: "${PERF_TOOLS_DIR:?PERF_TOOLS_DIR must be set (from .cluster-config.yaml perf_tools.dir, e.g. /path/to/perf-tools)}" +# set from .cluster-config.yaml omnihub.tools_dir (e.g. /shared/omnihub/tools). +: "${OMNIHUB_TOOLS_DIR:?OMNIHUB_TOOLS_DIR must be set (from .cluster-config.yaml omnihub.tools_dir, e.g. /shared/omnihub/tools)}" HG_BASE="${AI4S_SHARED_DIR}/models/HydraGNN" PERF_RUNS_DIR="${HG_BASE}/perf-runs" LOOP_DIR="${PERF_RUNS_DIR}/loop-${LOOP_UUID}" @@ -96,8 +96,8 @@ fi # 3. tools present for _bin in \ - "${PERF_TOOLS_DIR}/perf-inspect/bin/omnistat-usermode" \ - "${PERF_TOOLS_DIR}/victoriametrics/victoria-metrics-prod"; do + "${OMNIHUB_TOOLS_DIR}/omnihub-inspect/bin/omnistat-usermode" \ + "${OMNIHUB_TOOLS_DIR}/victoriametrics/victoria-metrics-prod"; do if [[ -z "$PREFLIGHT_FAIL_REASON" && ! -x "$_bin" ]]; then PREFLIGHT_FAIL_REASON="tool_missing" PREFLIGHT_NOTES+=("missing: $_bin (run the perf-analysis launcher subagent first to install)") diff --git a/material_science/models/HydraGNN/examples/sbatch_train_perf_amd.sh b/material_science/models/HydraGNN/examples/sbatch_train_perf_amd.sh index cbf5c0e..313c34f 100755 --- a/material_science/models/HydraGNN/examples/sbatch_train_perf_amd.sh +++ b/material_science/models/HydraGNN/examples/sbatch_train_perf_amd.sh @@ -21,8 +21,8 @@ # # Required: # AI4S_SHARED_DIR — shared base path -# ${PERF_TOOLS_DIR}/perf-inspect/ (created by the launcher subagent) -# ${PERF_TOOLS_DIR}/victoriametrics/ (created by the launcher subagent) +# ${OMNIHUB_TOOLS_DIR}/omnihub-inspect/ (created by the launcher subagent) +# ${OMNIHUB_TOOLS_DIR}/victoriametrics/ (created by the launcher subagent) # # Optional env-var overrides (with defaults): # HG_NUM_EPOCH=2 @@ -60,8 +60,8 @@ fi # --------------------------------------------------------------------------- HG_BASE="${AI4S_SHARED_DIR:?AI4S_SHARED_DIR must be set}/models/HydraGNN" # Shared perf-tool location (omnistat venv, omnistat-src, VictoriaMetrics). -# Cluster-specific: set from .cluster-config.yaml perf_tools.dir. -: "${PERF_TOOLS_DIR:?PERF_TOOLS_DIR must be set (from .cluster-config.yaml perf_tools.dir, e.g. /path/to/perf-tools)}" +# Cluster-specific: set from .cluster-config.yaml omnihub.tools_dir. +: "${OMNIHUB_TOOLS_DIR:?OMNIHUB_TOOLS_DIR must be set (from .cluster-config.yaml omnihub.tools_dir, e.g. /shared/omnihub/tools)}" HG_SIF="${HG_SIF:-${AI4S_SHARED_DIR}/images/pytorch_rocm7.2.2_ubuntu24.04_py3.12_pytorch_release_2.10.0.sif}" HG_OVERLAY="${HG_OVERLAY:-${HG_BASE}/overlays/hydragnn-overlay.img}" HG_DATASETS="${HG_DATASETS:-ANI1x,Alexandria}" @@ -74,7 +74,7 @@ HG_REPO_DIR="${HG_REPO_DIR:-${HG_BASE}/code/HydraGNN}" HYDRAGNN_MAX_NUM_BATCH="${HYDRAGNN_MAX_NUM_BATCH:-30}" PROFILE_TARGET_EPOCH="${PROFILE_TARGET_EPOCH:-1}" -OMNISTAT_VENV="${OMNISTAT_VENV:-${PERF_TOOLS_DIR}/perf-inspect}" +OMNISTAT_VENV="${OMNISTAT_VENV:-${OMNIHUB_TOOLS_DIR}/omnihub-inspect}" # Read the omnistat config from the in-repo recipe (single source of truth). # Override with OMNISTAT_TEMPLATE=/path/to/file if you need a custom probe config. # A staged copy under ${HG_BASE}/perf-runs/ is intentionally NOT used as the @@ -90,7 +90,7 @@ OMNISTAT_USERMODE_INTERVAL="${OMNISTAT_USERMODE_INTERVAL:-1}" # See material_science/models/HydraGNN/recipes/perf-analysis/agents/launcher.md # for how to build libomnistat_trace.so on a compute node. OMNISTAT_KERNEL_TRACE="${OMNISTAT_KERNEL_TRACE:-0}" -OMNISTAT_TRACE_LIB="${OMNISTAT_TRACE_LIB:-${PERF_TOOLS_DIR}/omnistat-src/build-trace/libomnistat_trace.so}" +OMNISTAT_TRACE_LIB="${OMNISTAT_TRACE_LIB:-${OMNIHUB_TOOLS_DIR}/omnistat-src/build-trace/libomnistat_trace.so}" # IMPORTANT: the kernel-trace collector registers /kernel_trace on the SAME # Flask app as the other omnistat endpoints, so the tool library must POST to # the [omnistat.collectors] `port` (8101 here), NOT the library's own default @@ -165,7 +165,7 @@ done mkdir -p "$HG_OUTPUT_DIR" OMNISTAT_CONFIG="${HG_OUTPUT_DIR}/omnistat.config" sed -e "s|@JOB_DIR@|${HG_OUTPUT_DIR}|g" \ - -e "s|@PERF_TOOLS_DIR@|${PERF_TOOLS_DIR}|g" \ + -e "s|@OMNIHUB_TOOLS_DIR@|${OMNIHUB_TOOLS_DIR}|g" \ "$OMNISTAT_TEMPLATE" > "$OMNISTAT_CONFIG" # Opt-in kernel tracing: flip enable_kernel_trace=True in the rendered config @@ -180,7 +180,7 @@ if [[ "$OMNISTAT_KERNEL_TRACE" == "1" ]]; then echo " Build it on a compute node:" >&2 echo " salloc -p -A -N 1 --time=00:15:00 --gpus-per-node=1 \\" >&2 echo " apptainer exec --rocm \"\$HG_SIF\" bash -c '" >&2 - echo " cd ${PERF_TOOLS_DIR}/omnistat-src && \\" >&2 + echo " cd ${OMNIHUB_TOOLS_DIR}/omnistat-src && \\" >&2 echo " cmake -S rocprofiler-sdk/ -B build-trace/ -DBUILD_KERNEL_TRACE_LIB=ON && \\" >&2 echo " cmake --build build-trace/ -j 8'" >&2 exit 2 diff --git a/material_science/models/HydraGNN/recipes/perf-analysis/README.md b/material_science/models/HydraGNN/recipes/perf-analysis/README.md index 85942c4..b608f24 100644 --- a/material_science/models/HydraGNN/recipes/perf-analysis/README.md +++ b/material_science/models/HydraGNN/recipes/perf-analysis/README.md @@ -80,8 +80,8 @@ The `combined_report.md` is the deliverable: ranked bottlenecks, remedies tried ## Prerequisites - Cluster access with the standard HydraGNN setup at `$AI4S_SHARED_DIR/models/HydraGNN/` — overlay, weights, code clone. See [HydraGNN/recipes/train/](../train/). -- Export `PERF_TOOLS_DIR` (the shared perf-tool location, from `.cluster-config.yaml` `perf_tools.dir`, e.g. `/path/to/perf-tools`). Scripts and subagents resolve all tool paths from this var — no cluster path is hardcoded. -- One-time install of Omnistat (PR #271 branch `jorda/skills`, with `origin/main` merged in), VictoriaMetrics, and TraceLens — performed lazily by the launcher subagent into `${PERF_TOOLS_DIR}/`. +- Export `OMNIHUB_TOOLS_DIR` (the shared perf-tool location, from `.cluster-config.yaml` `omnihub.tools_dir`, e.g. `/shared/omnihub/tools`). Scripts and subagents resolve all tool paths from this var — no cluster path is hardcoded. +- One-time install of Omnistat (PR #271 branch `jorda/skills`, with `origin/main` merged in), VictoriaMetrics, and TraceLens — performed lazily by the launcher subagent into `${OMNIHUB_TOOLS_DIR}/`. - The launcher writes `gfm_mlip_with_profile.json` at submit time (does not modify the upstream `gfm_mlip.json`). ## Telemetry knobs (set at `sbatch` submit time) @@ -89,7 +89,7 @@ The `combined_report.md` is the deliverable: ranked bottlenecks, remedies tried | Env var | Default | What it does | |---|---|---| | `OMNISTAT_KERNEL_TRACE` | `0` | When `1`, loads `libomnistat_trace.so` via `ROCP_TOOL_LIBRARIES` on every rank and turns on the omnistat kernel-trace collector. Adds per-kernel dispatch count + duration time series across all 8 cards on every node. Requires the library to be built once (see `agents/launcher.md` step 1e). Validated end-to-end on job 7034. | -| `OMNISTAT_TRACE_LIB` | `${PERF_TOOLS_DIR}/omnistat-src/build-trace/libomnistat_trace.so` | Override only for development. The wrapper hard-fails if the path is missing when `OMNISTAT_KERNEL_TRACE=1`. | +| `OMNISTAT_TRACE_LIB` | `${OMNIHUB_TOOLS_DIR}/omnistat-src/build-trace/libomnistat_trace.so` | Override only for development. The wrapper hard-fails if the path is missing when `OMNISTAT_KERNEL_TRACE=1`. | | `OMNISTAT_TRACE_LOG` | `1` | Library prints `[host][pid][omnistat] Trace summary: N/N processed records (M/M successful flushes)` on rank exit; useful as a smoke-signal that the tool initialized even when the workload crashed. | `enable_rocprofiler=True` is **on** in `omnistat.config.template` by default; the rendered config's state is printed in the sbatch banner so a silent "counters off" cannot recur. Kernel tracing and device-counting can co-exist on a single GPU but raise per-job VictoriaMetrics cardinality — pick by run length and the question you're asking. See `ai4science-studio` SKILL §12 for the device-counting vs kernel-trace vs `rocprofv3` decision matrix. diff --git a/material_science/models/HydraGNN/recipes/perf-analysis/agents/launcher.md b/material_science/models/HydraGNN/recipes/perf-analysis/agents/launcher.md index 457cf8d..ff4afcf 100644 --- a/material_science/models/HydraGNN/recipes/perf-analysis/agents/launcher.md +++ b/material_science/models/HydraGNN/recipes/perf-analysis/agents/launcher.md @@ -10,8 +10,8 @@ Submits a 2-node HydraGNN training (AMD Instinct MI355X) with PyTorch profiling ## Outputs -- `${PERF_TOOLS_DIR}/perf-inspect/` — Python venv with omnistat (PR #271 + main merged) and TraceLens. -- `${PERF_TOOLS_DIR}/victoriametrics/victoria-metrics-prod` — VictoriaMetrics binary. +- `${OMNIHUB_TOOLS_DIR}/omnihub-inspect/` — Python venv with omnistat (PR #271 + main merged) and TraceLens. +- `${OMNIHUB_TOOLS_DIR}/victoriametrics/victoria-metrics-prod` — VictoriaMetrics binary. - `$AI4S_SHARED_DIR/models/HydraGNN/perf-runs/omnistat.config` — Omnistat user-mode config. - `$AI4S_SHARED_DIR/models/HydraGNN/perf-runs//manifest.json` — manifest schema below. - `$AI4S_SHARED_DIR/models/HydraGNN/perf-runs//hydragnn-train-.out` — symlinked from output dir. @@ -57,8 +57,8 @@ Submits a 2-node HydraGNN training (AMD Instinct MI355X) with PyTorch profiling ```bash set -euo pipefail -# Shared tool location — read from .cluster-config.yaml perf_tools.dir. -TOOLS="${PERF_TOOLS_DIR:?set PERF_TOOLS_DIR from .cluster-config.yaml perf_tools.dir (e.g. /path/to/perf-tools)}" +# Shared tool location — read from .cluster-config.yaml omnihub.tools_dir. +TOOLS="${OMNIHUB_TOOLS_DIR:?set OMNIHUB_TOOLS_DIR from .cluster-config.yaml omnihub.tools_dir (e.g. /shared/omnihub/tools)}" mkdir -p "$TOOLS" # 1a. Omnistat: jorda/skills branch with origin/main merged in @@ -77,7 +77,7 @@ fi OMNISTAT_COMMIT=$(git rev-parse HEAD) # 1b. Venv (py3.12 to match cluster python) -VENV=${PERF_TOOLS_DIR}/perf-inspect +VENV=${OMNIHUB_TOOLS_DIR}/omnihub-inspect if [[ ! -d $VENV ]]; then python3 -m venv "$VENV" fi @@ -123,7 +123,7 @@ if [[ "${OMNISTAT_KERNEL_TRACE:-0}" == "1" && ! -f "$TRACE_LIB" ]]; then fi ``` -The sbatch wrapper expects the `.so` at `${PERF_TOOLS_DIR}/omnistat-src/build-trace/libomnistat_trace.so` by default; override with `OMNISTAT_TRACE_LIB=/path/to/lib.so`. See `ai4science-studio` SKILL §12 for when to enable kernel tracing vs the default device-counting collector. +The sbatch wrapper expects the `.so` at `${OMNIHUB_TOOLS_DIR}/omnistat-src/build-trace/libomnistat_trace.so` by default; override with `OMNISTAT_TRACE_LIB=/path/to/lib.so`. See `ai4science-studio` SKILL §12 for when to enable kernel tracing vs the default device-counting collector. ### 2. Author the Omnistat user-mode config (once, in repo `perf-runs/`) diff --git a/material_science/models/HydraGNN/recipes/perf-analysis/agents/omnistat_analyst.md b/material_science/models/HydraGNN/recipes/perf-analysis/agents/omnistat_analyst.md index 5ff92cf..56e9fef 100644 --- a/material_science/models/HydraGNN/recipes/perf-analysis/agents/omnistat_analyst.md +++ b/material_science/models/HydraGNN/recipes/perf-analysis/agents/omnistat_analyst.md @@ -6,8 +6,8 @@ Drive `omnistat-inspect` (PR #271) through the analyze-job phases on the user-mo - `/manifest.json` - The Omnistat DB at `manifest.omnistat_db_path`. -- The `perf-inspect` venv at `${PERF_TOOLS_DIR}/perf-inspect/`. -- The VictoriaMetrics binary at `${PERF_TOOLS_DIR}/victoriametrics/victoria-metrics-prod`. +- The `omnihub-inspect` venv at `${OMNIHUB_TOOLS_DIR}/omnihub-inspect/`. +- The VictoriaMetrics binary at `${OMNIHUB_TOOLS_DIR}/victoriametrics/victoria-metrics-prod`. ## Outputs @@ -36,7 +36,7 @@ ss -lnt "sport = :$PORT" 2>/dev/null | grep -q LISTEN && PORT=8429 VMLOG=$PERF_RUN/omnistat/vm.log mkdir -p "$PERF_RUN/omnistat/inspect_outputs" -nohup ${PERF_TOOLS_DIR}/victoriametrics/victoria-metrics-prod \ +nohup ${OMNIHUB_TOOLS_DIR}/victoriametrics/victoria-metrics-prod \ -storageDataPath="$DB" \ -httpListenAddr=127.0.0.1:$PORT \ -retentionPeriod=100y \ @@ -62,7 +62,7 @@ done ### 2. Run analyze-job phases (per the SKILL) ```bash -. ${PERF_TOOLS_DIR}/perf-inspect/bin/activate +. ${OMNIHUB_TOOLS_DIR}/omnihub-inspect/bin/activate SCRATCH=$PERF_RUN/omnistat/scratch mkdir -p "$SCRATCH" JOBID=$(jq -r .jobid ) diff --git a/material_science/models/HydraGNN/recipes/perf-analysis/agents/tracelens_analyst.md b/material_science/models/HydraGNN/recipes/perf-analysis/agents/tracelens_analyst.md index 99beea4..0aaff1d 100644 --- a/material_science/models/HydraGNN/recipes/perf-analysis/agents/tracelens_analyst.md +++ b/material_science/models/HydraGNN/recipes/perf-analysis/agents/tracelens_analyst.md @@ -39,7 +39,7 @@ Every entry: ### 1. Sanity check ```bash -. ${PERF_TOOLS_DIR}/perf-inspect/bin/activate +. ${OMNIHUB_TOOLS_DIR}/omnihub-inspect/bin/activate PERF_RUN=$(jq -r .perf_run_dir ) TRACE=$(jq -r '.trace_paths[0] // empty' ) [[ -z "$TRACE" ]] && { echo "STATUS=partial; reason=no_trace_in_manifest"; exit 0; } diff --git a/material_science/models/HydraGNN/recipes/perf-analysis/omnistat.config.template b/material_science/models/HydraGNN/recipes/perf-analysis/omnistat.config.template index d2fb3b9..80817de 100644 --- a/material_science/models/HydraGNN/recipes/perf-analysis/omnistat.config.template +++ b/material_science/models/HydraGNN/recipes/perf-analysis/omnistat.config.template @@ -120,7 +120,7 @@ system_name = mi355x #-- [omnistat.usermode] -victoria_binary = @PERF_TOOLS_DIR@/victoriametrics/victoria-metrics-prod +victoria_binary = @OMNIHUB_TOOLS_DIR@/victoriametrics/victoria-metrics-prod victoria_datadir = @JOB_DIR@/omnistat-db victoria_logfile = vic_server.log diff --git a/material_science/models/HydraGNN/recipes/perf-optimizer-loop/agents/fom_extractor.md b/material_science/models/HydraGNN/recipes/perf-optimizer-loop/agents/fom_extractor.md index 9d97925..c749537 100644 --- a/material_science/models/HydraGNN/recipes/perf-optimizer-loop/agents/fom_extractor.md +++ b/material_science/models/HydraGNN/recipes/perf-optimizer-loop/agents/fom_extractor.md @@ -68,7 +68,7 @@ omnistat kernel-trace collector). ```bash set -euo pipefail -. "${PERF_TOOLS_DIR}/perf-inspect/bin/activate" +. "${OMNIHUB_TOOLS_DIR}/omnihub-inspect/bin/activate" PERF_RUN=$(jq -r .perf_run_dir ) JOBID=$(jq -r .jobid ) RUNTIME=$(jq -r .runtime_seconds ) @@ -102,7 +102,7 @@ Start a local VictoriaMetrics if not already running: PORT=8428 ss -lnt "sport = :$PORT" 2>/dev/null | grep -q LISTEN && PORT=8429 if ! curl -sf "http://127.0.0.1:$PORT/api/v1/status/tsdb" > /dev/null 2>&1; then - nohup "${PERF_TOOLS_DIR}/victoriametrics/victoria-metrics-prod" \ + nohup "${OMNIHUB_TOOLS_DIR}/victoriametrics/victoria-metrics-prod" \ -storageDataPath="$PERF_RUN/omnistat-db" \ -httpListenAddr=127.0.0.1:$PORT \ -retentionPeriod=100y -fs.disableMmap \ diff --git a/material_science/models/HydraGNN/recipes/perf-optimizer-loop/agents/orchestrator.md b/material_science/models/HydraGNN/recipes/perf-optimizer-loop/agents/orchestrator.md index 85cc223..256addc 100644 --- a/material_science/models/HydraGNN/recipes/perf-optimizer-loop/agents/orchestrator.md +++ b/material_science/models/HydraGNN/recipes/perf-optimizer-loop/agents/orchestrator.md @@ -85,11 +85,11 @@ Read `loop-/foms.csv` if it exists. If iter rows are present, set `current Run these checks in sequence; abort with `PREFLIGHT_FAIL reason=` on the first failure: -- Ensure `AI4S_SHARED_DIR` and `PERF_TOOLS_DIR` are exported (the latter from `.cluster-config.yaml` `perf_tools.dir`, e.g. `/path/to/perf-tools`); abort if either is unset. All tool paths below derive from `$PERF_TOOLS_DIR` — never hardcode a cluster path. +- Ensure `AI4S_SHARED_DIR` and `OMNIHUB_TOOLS_DIR` are exported (the latter from `.cluster-config.yaml` `omnihub.tools_dir`, e.g. `/shared/omnihub/tools`); abort if either is unset. All tool paths below derive from `$OMNIHUB_TOOLS_DIR` — never hardcode a cluster path. - `sinfo -p -h` (partition from `.cluster-config.yaml`) → if zero `IDLE` or `MIX` nodes, abort. Cluster may be in maintenance (see SKILL §14). - `df -h "$AI4S_SHARED_DIR" | tail -1` → abort if `Use%` > 95. -- `test -x ${PERF_TOOLS_DIR}/perf-inspect/bin/omnistat-usermode` → abort if missing. -- `test -x ${PERF_TOOLS_DIR}/victoriametrics/victoria-metrics-prod` → abort if missing. +- `test -x ${OMNIHUB_TOOLS_DIR}/omnihub-inspect/bin/omnistat-usermode` → abort if missing. +- `test -x ${OMNIHUB_TOOLS_DIR}/victoriametrics/victoria-metrics-prod` → abort if missing. - `test -f $AI4S_SHARED_DIR/models/HydraGNN/overlays/hydragnn-overlay.img` → abort if missing. - `test -f $AI4S_SHARED_DIR/images/pytorch_rocm7.2.2_ubuntu24.04_py3.12_pytorch_release_2.10.0.sif` → abort if missing. - If invoked via Claude Code CLI on the login node: `[[ -n "$ANTHROPIC_API_KEY" ]]` and `curl -fsS https://api.anthropic.com/v1/models > /dev/null` (5 s timeout) → abort if either fails. diff --git a/material_science/models/HydraGNN/recipes/perf-optimizer-loop/agents/story_writer.md b/material_science/models/HydraGNN/recipes/perf-optimizer-loop/agents/story_writer.md index f296ca4..016e386 100644 --- a/material_science/models/HydraGNN/recipes/perf-optimizer-loop/agents/story_writer.md +++ b/material_science/models/HydraGNN/recipes/perf-optimizer-loop/agents/story_writer.md @@ -156,7 +156,7 @@ fig.tight_layout() fig.savefig(out_png_path, dpi=150, bbox_inches="tight") ``` -Use `matplotlib.use("Agg")` for headless rendering. Do NOT use seaborn (extra dep not in the omnistat venv). Use the omnistat venv at `${PERF_TOOLS_DIR}/perf-inspect/bin/python` which already has matplotlib via TraceLens deps. +Use `matplotlib.use("Agg")` for headless rendering. Do NOT use seaborn (extra dep not in the omnistat venv). Use the omnistat venv at `${OMNIHUB_TOOLS_DIR}/omnihub-inspect/bin/python` which already has matplotlib via TraceLens deps. ## Style rules diff --git a/material_science/models/HydraGNN/recipes/perf-optimizer-loop/dispatch-attribution.md b/material_science/models/HydraGNN/recipes/perf-optimizer-loop/dispatch-attribution.md index ebedaf7..21f57d5 100644 --- a/material_science/models/HydraGNN/recipes/perf-optimizer-loop/dispatch-attribution.md +++ b/material_science/models/HydraGNN/recipes/perf-optimizer-loop/dispatch-attribution.md @@ -59,10 +59,10 @@ Short-kernel study (`short_kernels_summary.csv`): thousands of `aten::fill_`, `a ```bash export AI4S_SHARED_DIR=/shared/$USER export REPO_ROOT=$HOME/git/ai4science-studio -source ${PERF_TOOLS_DIR}/perf-inspect/bin/activate +source ${OMNIHUB_TOOLS_DIR}/omnihub-inspect/bin/activate # VictoriaMetrics on the perf-run DB (login node) -${PERF_TOOLS_DIR}/victoriametrics/victoria-metrics-prod \ +${OMNIHUB_TOOLS_DIR}/victoriametrics/victoria-metrics-prod \ -storageDataPath=$AI4S_SHARED_DIR/models/HydraGNN/perf-runs/7187/omnistat-db \ -httpListenAddr=127.0.0.1:8432 -retentionPeriod=100y -fs.disableMmap &