diff --git a/.claude/commands/init-cluster.md b/.claude/commands/init-cluster.md index 0e2b5cf..9bb77a9 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: "" + 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) -rccl: - socket_ifname: "" # discover: ip link show up - ib_hca: "" # discover: ibstat | grep CA +# 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) @@ -177,6 +193,7 @@ After writing, confirm to the user: Suggest what the user can do next: - Export `AI4S_SHARED_DIR=` (or add to `~/.bashrc`) so all model scripts resolve paths correctly - Run a model: `/run-stormcast`, `/run-orbit2`, etc. +- If **Q9 (Performance tooling)** was configured: run ORBIT-2 or HydraGNN perf recipes with `/run-perf-orbit2` or `/run-perf-hydragnn` - The run commands will read cluster config automatically for SLURM partition/account defaults $ARGUMENTS diff --git a/.claude/commands/run-perf-hydragnn.md b/.claude/commands/run-perf-hydragnn.md new file mode 100644 index 0000000..c98abdf --- /dev/null +++ b/.claude/commands/run-perf-hydragnn.md @@ -0,0 +1,146 @@ +# Run HydraGNN perf-analysis or perf-optimizer-loop on an AMD cluster + +Guide the user through **performance engineering** for HydraGNN Predictive GFM: one-shot **perf-analysis** (2-node training + TraceLens + Omnistat diagnosis) or iterative **perf-optimizer-loop** (accept/revert on **`epoch_time_s`**). This is **not** ensemble inference — use `/run-hydragnn` for that. + +**Recipe docs:** [material_science/models/HydraGNN/recipes/perf-analysis/README.md](../material_science/models/HydraGNN/recipes/perf-analysis/README.md), [material_science/models/HydraGNN/recipes/perf-optimizer-loop/README.md](../material_science/models/HydraGNN/recipes/perf-optimizer-loop/README.md). + +**Attribution / levers:** [dispatch-attribution.md](../material_science/models/HydraGNN/recipes/perf-optimizer-loop/dispatch-attribution.md), [lever_catalog.yaml](../material_science/models/HydraGNN/recipes/perf-optimizer-loop/lever_catalog.yaml). + +**Orchestration reference:** `.cursor/skills/ai4science-perf-analysis/SKILL.md` (launcher → parallel analysts → parallel verifiers → synthesizer). + +## Step 0 — Cluster config check (required) + +Read `.cluster-config.yaml` (repo root) or `~/.config/ai4science-studio/cluster.yaml`. + +**Required for perf recipes:** + +- Set **`AI4S_SHARED_DIR`** from cluster config (`paths.scratch` or documented scratch root). +- Set **`PERF_TOOLS_DIR`** from `perf_tools.dir`. If missing, stop and send the user to **`/init-cluster`** Q9 (Performance tooling). + +Pre-fill SLURM partition/account from config. + +**Runtime:** Apptainer + SLURM only (`sbatch_train_perf_amd.sh`). + +--- + +## Step 1 — Questionnaire (ask ALL before acting) + +**Q0. Mode** + +- **perf-analysis** — One 2-node perf job + post-hoc multi-subagent analysis (default topology for this recipe). +- **perf-optimizer-loop** — Iterative tuning; primary FOM **`epoch_time_s`** (lower is better); see `recipes/perf-optimizer-loop/README.md` for stop rules and `do_not_retry.json`. + +**Q1. HydraGNN checkout (`HG_REPO_DIR`)** + +Default in sbatch: `$AI4S_SHARED_DIR/models/HydraGNN/code/HydraGNN` (cloned at pinned SHA if missing). + +- **Use default** / **Yes, custom path** / **No** (clone instructions) / **Auto-discover** — `find "$HOME" /scratch /projects /opt -maxdepth 5 -type d -name "HydraGNN" 2>/dev/null | head -20` and confirm `examples/multidataset_hpo_sc26` exists. + +**Q2. SIF (`HG_SIF`)** + +Default: `$AI4S_SHARED_DIR/images/pytorch_rocm7.2.2_ubuntu24.04_py3.12_pytorch_release_2.10.0.sif`. + +- **Yes** / **No** / **Auto-discover** — same SIF discovery pattern as `/run-orbit2`. + +**Q3. Overlay (`HG_OVERLAY`)** + +Default: `$AI4S_SHARED_DIR/models/HydraGNN/overlays/hydragnn-overlay.img`. + +- **Yes** / **No, build one** — `sbatch material_science/models/HydraGNN/examples/build_overlay_amd.sh` (edit `#SBATCH` partition/account first) / **Auto-discover** — `find ... -name "*hydragnn*overlay*"`. + +**Q4. Datasets (`HG_DATASETS`, `HG_DATA_DIR`)** + +Default datasets and paths are set in `sbatch_train_perf_amd.sh` (`ANI1x,Alexandria` under `HG_DATA_DIR`). Confirm `.bp` files exist or user will stage per [recipes/data/README.md](../material_science/models/HydraGNN/recipes/data/README.md). + +**Q5. Precision** + +Perf recipe defaults to **`fp64`** / **`HG_PRECISION=fp64`** (MACE-style equivariance). **`bf16` is not used** for this GFM path — do not suggest bf16 as a lever; the lever catalog documents blocked options. + +**Q6. Blocked or high-risk levers (set expectations)** + +Point users to `lever_catalog.yaml` entries marked **`blocked`** or **`accepted-baked-in`**, including at minimum: + +- **`torch.compile`** on HydraGNN MLIP MACE — blocked (double-backward / AOT autograd). +- **`torch.jit.script`** — blocked (keyword-arg expansion at MACE interaction block; see lever catalog). +- **TunableOp live or warmup-then-use** — blocked on this stack (documented GPU faults during tuning); do not re-propose without a stack update and catalog change. + +**Q7. Partition and account** + +- **Provide manually** / **Auto-discover** — `sinfo` + `sacctmgr` as in `/run-orbit2`. + +**Q8. (perf-optimizer-loop only) Iteration budget** + +Default **5**; run inside **`tmux`**. Optional: `export ANTHROPIC_API_KEY=...` for Claude Code CLI orchestration inside `run_optimizer_loop.sh`. + +--- + +## Step 2 — Auto-discovery + +Reuse SIF/overlay/SLURM snippets from `.claude/commands/run-orbit2.md`. + +**HydraGNN repo:** + +```bash +find "$HOME" /scratch /projects /opt -maxdepth 5 -type d -name "HydraGNN" 2>/dev/null | head -20 +``` + +Confirm branch **`Predictive_GFM_2024`** if user cares about upstream alignment (see model README). + +--- + +## Step 3 — Run + +```bash +export AI4S_SHARED_DIR= +export PERF_TOOLS_DIR= +# Optional overrides: +# export HG_SIF=... HG_OVERLAY=... HG_REPO_DIR=... HG_DATASETS=... HG_DATA_DIR=... +``` + +Edit **`#SBATCH`** in **`material_science/models/HydraGNN/examples/sbatch_train_perf_amd.sh`** for partition/account. + +**Default perf-analysis job is 2-node** — ensure `#SBATCH --nodes=2` (or user override) matches intent. + +### perf-analysis + +```bash +sbatch material_science/models/HydraGNN/examples/sbatch_train_perf_amd.sh +``` + +After completion, drive subagents per **`ai4science-perf-analysis` SKILL** using prompts under `material_science/models/HydraGNN/recipes/perf-analysis/agents/`. + +### perf-optimizer-loop + +Preflight only (smoke; no Claude invocation): + +```bash +bash material_science/models/HydraGNN/examples/run_optimizer_loop.sh "$(uuidgen)" 5 --preflight-only +``` + +Full loop: + +```bash +tmux new -s hg-perf-loop +bash material_science/models/HydraGNN/examples/run_optimizer_loop.sh "$(uuidgen)" 5 +``` + +Graceful stop: `touch "$AI4S_SHARED_DIR/models/HydraGNN/perf-runs/loop-/STOP"`. + +--- + +## Step 4 — Monitor and artifacts + +```bash +squeue -j +tail -f "$AI4S_SHARED_DIR/models/HydraGNN/perf-runs//"*.out +``` + +**perf-analysis:** `$AI4S_SHARED_DIR/models/HydraGNN/perf-runs//` — `combined_report.md`, `manifest.json`, `tracelens/`, `omnistat/`, traces under `logs/`. + +**perf-optimizer-loop:** `$AI4S_SHARED_DIR/models/HydraGNN/perf-runs/loop-/` — `STATUS.txt`, `foms.csv`, `story.md`, per-iteration job subdirs. + +--- + +## Arguments + +$ARGUMENTS diff --git a/.claude/commands/run-perf-orbit2.md b/.claude/commands/run-perf-orbit2.md new file mode 100644 index 0000000..899d2d3 --- /dev/null +++ b/.claude/commands/run-perf-orbit2.md @@ -0,0 +1,164 @@ +# Run ORBIT-2 perf-analysis or perf-optimizer-loop on an AMD cluster + +Guide the user through **performance engineering** workflows for ORBIT-2 (Bayes-CAST EDM / same-dir timing configs): one-shot **perf-analysis** (TraceLens + Omnistat diagnosis) or iterative **perf-optimizer-loop** (accept/revert on throughput). This is **not** inference/visualization — use `/run-orbit2` for that. + +**Recipe docs:** [earth_science/models/ORBIT-2/recipes/perf-analysis/README.md](../earth_science/models/ORBIT-2/recipes/perf-analysis/README.md), [earth_science/models/ORBIT-2/recipes/perf-optimizer-loop/README.md](../earth_science/models/ORBIT-2/recipes/perf-optimizer-loop/README.md). + +**Orchestration reference (Cursor / any agent):** read `.cursor/skills/ai4science-perf-analysis/SKILL.md` for the launcher → analyst → verifier → synthesizer loop. + +## Step 0 — Cluster config check (required) + +Read `.cluster-config.yaml` (repo root) or `~/.config/ai4science-studio/cluster.yaml`. + +**Required for perf recipes:** + +- `paths.scratch` (or equivalent) → use as **`AI4S_SHARED_DIR`** (same convention as other studio scripts). +- `perf_tools.dir` → use as **`PERF_TOOLS_DIR`** (layout: `perf-inspect/`, `omnistat-src/`, `victoriametrics/`). If missing or empty, stop and tell the user to re-run **`/init-cluster`** and answer **Q9 (Performance tooling)**. + +Pre-fill SLURM partition/account from `slurm.partition` / `slurm.account` when present. + +**Runtime note:** Perf workflows here assume **Apptainer + SLURM** (same as `sbatch_train_perf_amd.sh`). There is no Docker perf path in the recipes. + +--- + +## Step 1 — Questionnaire (ask ALL questions before doing anything) + +Ask the following. Do not assume defaults. Wait for answers before proceeding. + +**Q0. Mode** + +- **perf-analysis** — Submit one training perf job; after it finishes, drive subagents per `ai4science-perf-analysis` SKILL (`agents/launcher.md`, analysts, verifiers, `synthesizer.md`) to produce `combined_report.md`. +- **perf-optimizer-loop** — Iterative loop: one lever per iteration, accept/revert on **`throughput_samples_per_s`** with **`loss_sanity_pass`** as a control (see perf-optimizer-loop README). + +**Q1. Code root (`ORBIT2_ROOT`)** + +Bayes-CAST checkout used for training (must contain `launch/train_edm.py` for the default perf path). + +- **Yes** — user provides full path. +- **No** — generate clone instructions (user may use Bayes-CAST fork or path documented in perf-analysis README). +- **Auto-discover** — search for a directory containing `launch/train_edm.py`. + +**Q2. SIF** + +Apptainer image for ROCm PyTorch `rocm7.2.2_ubuntu24.04_py3.12_pytorch_release_2.10.0` (or user override). + +- **Yes** / **No** (pull command) / **Auto-discover** — same pattern as `/run-orbit2`. + +**Q3. ORBIT-2 overlay** + +Pre-built `orbit2-overlay.img` (recommended). + +- **Yes** / **No, build one** (`sbatch earth_science/models/ORBIT-2/examples/build_overlay_amd.sh`) / **Auto-discover** — same pattern as `/run-orbit2`. + +**Q4. Data** + +- **`ORBIT2_DATA_ROOT`** — Staged ERA5 NPZ tree (perf defaults use Bayes-CAST template + ERA5; see perf-analysis README). +- **`ORBIT2_CONFIG_TEMPLATE`** — Optional; default **`edm_8m_era5_1x8.yaml`** in perf docs. For HBM saturation / staging notes: [STAGING_ERA5_FOR_HBM.md](../earth_science/models/ORBIT-2/recipes/perf-optimizer-loop/STAGING_ERA5_FOR_HBM.md). +- **`interm_8m_prism.yaml` / `interm_8m_era5.yaml`** — Public ORBIT-2 same-dir configs (timing-only); document if used instead of Bayes-CAST EDM. + +**Q5. Topology** + +- **1 node × 8 GPUs** (default for phase-1 loop and perf-analysis quick start). +- **Multi-node** — e.g. `sbatch --nodes=2 ...`; RCCL / GID / `LD_LIBRARY_PATH` caveats in perf-analysis README and `sbatch_train_perf_amd.sh`. After one-node loop converges, see perf-optimizer-loop README **Two-node gate**; optional driver: `earth_science/models/ORBIT-2/examples/run_2node_scaleout_loop.sh`. + +**Q6. Partition and account** + +- **Provide manually** / **Auto-discover** (same as `/run-orbit2`). + +**Q7. (perf-optimizer-loop only) Iteration budget** + +Max iterations (default **5**). Confirm use of **`tmux`** on the login node so SSH disconnect does not kill the driver. + +**Q8. (perf-optimizer-loop only) Claude Code CLI driver** + +`run_optimizer_loop.sh` can invoke the orchestrator via Claude Code CLI when **`ANTHROPIC_API_KEY`** is set. If the user has no API key, explain they can still run **preflight-only** steps and manual `sbatch` per iteration using `recipes/perf-optimizer-loop/agents/orchestrator.md` as a checklist. + +--- + +## Step 2 — Auto-discovery (when chosen) + +**Code root containing `launch/train_edm.py`:** + +```bash +find "$HOME" /scratch /projects /opt -maxdepth 6 -type f -path "*/launch/train_edm.py" 2>/dev/null | head -20 +``` + +Derive `ORBIT2_ROOT` as the parent of `launch/`. + +**SIF / overlay / SLURM:** reuse the same `find` / `sinfo` / `sacctmgr` snippets as in `.claude/commands/run-orbit2.md` (SIF: `*.sif` filtered for rocm/pytorch; overlay: `*orbit2*overlay*`). + +Always confirm discovered paths with the user before `sbatch`. + +--- + +## Step 3 — Run + +Export from cluster config + user answers: + +```bash +export AI4S_SHARED_DIR= +export PERF_TOOLS_DIR= +export ORBIT2_ROOT= +export ORBIT2_SIF= # if not using default under $AI4S_SHARED_DIR/images/... +export ORBIT2_OVERLAY= # optional if script defaults suffice +export ORBIT2_DATA_ROOT= +# export ORBIT2_CONFIG_TEMPLATE=edm_8m_era5_1x8.yaml # optional +``` + +Edit **`#SBATCH`** lines in **`earth_science/models/ORBIT-2/examples/sbatch_train_perf_amd.sh`** for partition/account (and `--nodes` if multi-node). + +### perf-analysis + +```bash +# From ai4science-studio repo root +sbatch earth_science/models/ORBIT-2/examples/sbatch_train_perf_amd.sh +``` + +After the job completes, follow **`.cursor/skills/ai4science-perf-analysis/SKILL.md`**: dispatch launcher (if not already folded into job), `tracelens_analyst` + `omnistat_analyst` in parallel, verifiers in parallel, then `synthesizer`. Prompts live under `earth_science/models/ORBIT-2/recipes/perf-analysis/agents/`. + +### perf-optimizer-loop + +Optional repo file check (does not validate cluster env): + +```bash +bash earth_science/models/ORBIT-2/examples/validate_orbit2_optimizer_loop_recipe.sh +``` + +Preflight only (recommended first): + +```bash +bash earth_science/models/ORBIT-2/examples/run_optimizer_loop.sh "$(uuidgen)" 5 --preflight-only +``` + +Full loop (inside **`tmux`**): + +```bash +tmux new -s orbit2-perf-loop +bash earth_science/models/ORBIT-2/examples/run_optimizer_loop.sh "$(uuidgen)" 5 +# detach: Ctrl-b d +``` + +Graceful stop: `touch "$AI4S_SHARED_DIR/models/ORBIT-2/perf-runs/loop-/STOP"`. + +--- + +## Step 4 — Monitor and artifacts + +**SLURM job:** + +```bash +squeue -j +tail -f "$AI4S_SHARED_DIR/models/ORBIT-2/perf-runs//"*.out +``` + +**perf-analysis deliverables** (under `$AI4S_SHARED_DIR/models/ORBIT-2/perf-runs//`): `manifest.json`, `omnistat-db/`, `tracelens/`, `combined_report.md`, `foms.json` (when extractor run), traces under `traces/`. + +**perf-optimizer-loop** (under `$AI4S_SHARED_DIR/models/ORBIT-2/perf-runs/loop-/`): `STATUS.txt`, `foms.csv`, `story.md`, `do_not_retry.json`, `known_bad_nodes.txt`; per-iter job dirs mirror perf-analysis layout. + +**Primary FOM (loop):** `throughput_samples_per_s` (higher is better); keep **`loss_sanity_pass`**. GEMM / lever evidence: [gemm-attribution.md](../earth_science/models/ORBIT-2/recipes/perf-optimizer-loop/gemm-attribution.md), [lever_catalog.yaml](../earth_science/models/ORBIT-2/recipes/perf-optimizer-loop/lever_catalog.yaml). + +--- + +## Arguments + +$ARGUMENTS 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 f713f60..1bce7c5 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. **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,6 +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` 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_era5.yaml` 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). **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 **[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_prism.yaml`/`interm_8m_era5.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. IB/ionic). ### ArchesWeather (geoarches) 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-perf-analysis/SKILL.md b/.cursor/skills/ai4science-perf-analysis/SKILL.md new file mode 100644 index 0000000..1c780a9 --- /dev/null +++ b/.cursor/skills/ai4science-perf-analysis/SKILL.md @@ -0,0 +1,154 @@ +--- +name: ai4science-perf-analysis +description: Runs the AMD AI agents bottleneck-analysis workflow on a HydraGNN or ORBIT-2 training job using TraceLens + Omnistat user-mode + paired analyst/verifier subagents, or drives the iterative perf-optimizer-loop (accept/revert on throughput or epoch time). Use when the user wants to analyze a perf run, analyze a 2-node HydraGNN run, find bottlenecks with TraceLens/Omnistat, run the AMD AI agents on a job, run the perf-optimizer-loop, tune ORBIT-2 throughput, or optimize HydraGNN training performance. +--- + +# 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 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 + +### 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 [`dispatch-attribution.md`](../../../material_science/models/HydraGNN/recipes/perf-optimizer-loop/dispatch-attribution.md) (current-best attribution) and [`lever_catalog.yaml`](../../../material_science/models/HydraGNN/recipes/perf-optimizer-loop/lever_catalog.yaml) (levers tried / blocked with evidence)** +- 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 [`README.md`](../../../earth_science/models/ORBIT-2/recipes/perf-analysis/README.md) (FOM contract, defaults, landmines) and [`gemm-attribution.md`](../../../earth_science/models/ORBIT-2/recipes/perf-optimizer-loop/gemm-attribution.md) (GEMM-time finding + levers) +- **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) — 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), 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_prism.yaml`) or ERA5 1.0° same-dir (`interm_8m_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` (res_slimvit path; Bayes EDM defaults to no hook unless overridden) + +## 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 -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 (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. +- 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 + +- 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 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/...`. +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="..."})) + ``` + +## 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-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. +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 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` 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 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. **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. **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. + +## 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). + +## 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 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 7355375..34e0b47 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)). + **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 @@ -313,7 +315,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). 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 docs 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) @@ -349,6 +351,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. Watch for it whenever you add `python3 -c "..."` scaffolding around an `srun` that uses `--mpi=pmix`. + --- # Docker-specific patterns @@ -464,3 +518,261 @@ 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.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 = $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 + +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 +``` + +**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. + +### 8. Never read configs from a `/shared` "staged copy" — always from the in-repo template + +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. +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 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. +- 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 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 +[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. + +### 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.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. `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 ${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 ${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: + +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 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. + +### 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. + +### 14. Cluster-state checks before queueing build/probe jobs + +`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 -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} \ + "$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). + +### 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. + +### 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_prism.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. +- **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** (the **`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 perf-analysis [README.md](../../../earth_science/models/ORBIT-2/recipes/perf-analysis/README.md)). +- **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. This pattern was hit during the ORBIT-2 CK + **`bfloat16`** SIGSEGV investigation. +- **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 79e9c27..4659a86 100755 --- a/.gitignore +++ b/.gitignore @@ -45,6 +45,18 @@ 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 +**/examples/scaling_study* +**/examples/foms-*.json # Model outputs & generated data **/examples/outputs/ @@ -52,9 +64,15 @@ secrets/ # 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 # IDE .idea/ .vscode/ + +# Stray SLURM stdout from perf runs (real logs live under perf-runs//) +orbit2-train-*.out diff --git a/CLAUDE.md b/CLAUDE.md index 968db65..2e5dc7b 100644 --- a/CLAUDE.md +++ b/CLAUDE.md @@ -67,4 +67,4 @@ When you fix anything in a model's scripts, do **all** of the following in the s ## Agent Skills - **Cursor**: `.cursor/skills/` — domain conventions (earth science, healthcare, material science, protein folding, physics simulation, run-models, discover) -- **Claude Code**: `.claude/commands/` — slash commands for add-model, add-recipe, check-model, list-models, audit-models, init-cluster, and run-* for each model +- **Claude Code**: `.claude/commands/` — slash commands for add-model, add-recipe, check-model, list-models, audit-models, init-cluster, run-* for each model, and **`/run-perf-orbit2`** / **`/run-perf-hydragnn`** for perf-analysis + perf-optimizer-loop recipes diff --git a/earth_science/models/ORBIT-2/README.md b/earth_science/models/ORBIT-2/README.md index 235143c..04ddc0e 100755 --- a/earth_science/models/ORBIT-2/README.md +++ b/earth_science/models/ORBIT-2/README.md @@ -14,6 +14,8 @@ Training, inference, visualization, and SLURM examples are maintained in the **G - [`recipes/train/`](recipes/train/) — large-scale / Frontier-oriented training summary - [`recipes/data/`](recipes/data/) — Constellation / Globus / OLCF data context and staging - [`recipes/compute/`](recipes/compute/) — Frontier, AMD cloud / HPC Fund, institutional clusters +- [`recipes/perf-analysis/`](recipes/perf-analysis/) — multi-subagent bottleneck analysis (TraceLens + Omnistat) on AMD Instinct +- [`recipes/perf-optimizer-loop/`](recipes/perf-optimizer-loop/) — iterative throughput optimizer loop + lever catalog - [`examples/`](examples/) — ready-to-run Docker, SLURM, and smoke-test scripts ## Installation (reference) diff --git a/earth_science/models/ORBIT-2/examples/README.md b/earth_science/models/ORBIT-2/examples/README.md index 6283a86..c518ff5 100644 --- a/earth_science/models/ORBIT-2/examples/README.md +++ b/earth_science/models/ORBIT-2/examples/README.md @@ -6,15 +6,111 @@ with AMD/ROCm container launch, overlay builds, and synthetic smoke tests. ## Files +The folder is flat (every script resolves siblings, configs, and the +`/examples` container bind-mount relative to its own directory), but the files +fall into the groups below. + +### Setup & inference + | Script | Purpose | |--------|---------| | [`preflight_orbit2.py`](preflight_orbit2.py) | Check your environment before submitting jobs | | [`run_visualize.py`](run_visualize.py) | Run upstream `visualize.py` with Studio env-var overrides | -| [`make_synthetic_data.py`](make_synthetic_data.py) | Generate ~2 MB synthetic dataset for smoke tests | + +### Container & launch drivers + +| Script | Purpose | +|--------|---------| | [`docker_run.sh`](docker_run.sh) | Docker launcher for local workstations and interactive nodes | -| [`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) | +| [`sbatch_train_amd.sh`](sbatch_train_amd.sh) | Multi-node training (Apptainer + MPI); defaults PRISM 10.0_arcmin + `interm_8m_prism.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 …` | +| [`sbatch_infer_amd.sh`](sbatch_infer_amd.sh) | SLURM driver for inference on AMD Instinct (MI250X/MI300X/MI350X) | +| [`sbatch_infer_docker.sh`](sbatch_infer_docker.sh) | SLURM inference driver using Docker instead of Apptainer | + +### In-container launch glue + +Run on the host (`render_*`) or inside the container via the `/examples` +bind-mount; the sbatch drivers depend on these by name, so they stay flat. + +| Script | Purpose | +|--------|---------| +| [`render_orbit2_config.py`](render_orbit2_config.py) | Render a per-job YAML from a template (parallelism, data root, trainer caps); invoked by the sbatch drivers | +| [`run_orbit2_train.py`](run_orbit2_train.py) | Studio training entry run inside the container (wraps upstream `intermediate_downscaling.py`; gptl4py stub, FusedAttn fallback, batch cap) | +| [`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 | +| [`orbit2_profiler_hook.py`](orbit2_profiler_hook.py) | Optional per-rank PyTorch-profiler hook (`res_slimvit` path) invoked via `ORBIT2_RANK_PRE_TRAIN_HOOK` | + +### Configs (YAML) + +Kept beside the sbatch drivers — they are resolved as `${SCRIPT_DIR}/${ORBIT2_CONFIG_TEMPLATE}`. + +| Script | Purpose | +|--------|---------| +| [`interm_8m_prism.yaml`](interm_8m_prism.yaml) | PRISM config template (10.0_arcmin same-dir) | +| [`interm_8m_era5.yaml`](interm_8m_era5.yaml) | ERA5 1.0° config template (same-dir sanity / pipeline test) | | [`interm_8m_synthetic.yaml`](interm_8m_synthetic.yaml) | Template YAML for synthetic mode | +| [`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) | + +### Data prep + +| Script | Purpose | +|--------|---------| +| [`make_synthetic_data.py`](make_synthetic_data.py) | Generate ~2 MB synthetic dataset for smoke tests | +| [`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 | + +### Perf & scaling orchestration + +| Script | Purpose | +|--------|---------| +| [`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_scaling_study.sh`](run_scaling_study.sh) | Submit matched 1/2/4/8-node scaling sweep | +| [`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` | +| [`run_gemm_analysis.sh`](run_gemm_analysis.sh) | Unattended GEMM-time bottleneck analysis (TraceLens + Omnistat analyst/verifier) at 1 and 2 nodes | +| [`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 | + +### Post-processing & analysis + +Host-side; run after a job against its logs / traces. + +| Script | Purpose | +|--------|---------| +| [`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 | +| [`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)) | +| [`collate_scaling_study.py`](collate_scaling_study.py) | Aggregate a multi-node scaling sweep into `scaling_study` outputs | +| [`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) | +| [`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) + +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_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/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_era5.yaml b/earth_science/models/ORBIT-2/examples/interm_8m_era5.yaml new file mode 100644 index 0000000..dc42c6e --- /dev/null +++ b/earth_science/models/ORBIT-2/examples/interm_8m_era5.yaml @@ -0,0 +1,145 @@ +# ORBIT-2 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 (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: + 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/interm_8m_prism.yaml b/earth_science/models/ORBIT-2/examples/interm_8m_prism.yaml new file mode 100644 index 0000000..198615d --- /dev/null +++ b/earth_science/models/ORBIT-2/examples/interm_8m_prism.yaml @@ -0,0 +1,129 @@ +# ORBIT-2 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 (same intent as upstream `configs/interm_8m.yaml` / +# Bayes-CAST `edm_8m_era5.yaml`: **bfloat16**; override float32 only for debug) + +#----------- 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/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 new file mode 100755 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/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 new file mode 100755 index 0000000..7624573 --- /dev/null +++ b/earth_science/models/ORBIT-2/examples/parse_training_log.py @@ -0,0 +1,404 @@ +#!/usr/bin/env python3 +"""Parse ORBIT-2 training metrics from SLURM stdout (orbit2-train-*.out). + +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 + +import argparse +import csv +import json +import re +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: + line = line.strip() + + m = EPOCH_START_RE.match(line) + if m: + 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( + { + "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)) + saw_upstream_epoch = True + records.append( + { + "kind": "epoch", + "batch": None, + "batch_time_s": None, + "epoch": epoch, + "loss": float(m.group(2)), + } + ) + 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 + + +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, + global_batch_size: int | None = None, + nominal_per_rank_batch: int | None = None, +) -> dict: + """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] + + 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) + + 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] + final_loss = ( + steady_epoch_losses[-1]["loss"] + if steady_epoch_losses + 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": 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, + "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)", + ) + 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(): + 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", "batch_samples"], + extrasaction="ignore", + ) + 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, + 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) + 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..917eabd --- /dev/null +++ b/earth_science/models/ORBIT-2/examples/render_orbit2_config.py @@ -0,0 +1,145 @@ +#!/usr/bin/env python3 +"""Render a per-job ORBIT-2 training YAML from a Studio template. + +Templates include same-dir (**``interm_8m_prism.yaml``** / **``interm_8m_era5.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 + 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 training config.") + parser.add_argument( + "--template", + type=Path, + default=None, + help="Template YAML (default: interm_8m_prism.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_prism.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 + + # 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: + 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", "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° ERA5 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), + "__FSDP__": str(fsdp), + "__SIMPLE_DDP__": str(simple_ddp), + "__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) + + 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/report_orbit2_gpu_baseline.py b/earth_science/models/ORBIT-2/examples/report_orbit2_gpu_baseline.py new file mode 100755 index 0000000..77e81e0 --- /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_prism_{job_id}.yaml", f"interm_8m_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 ("prism" in p.name or "era5" 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 new file mode 100755 index 0000000..e986e56 --- /dev/null +++ b/earth_science/models/ORBIT-2/examples/run_fom_extractor.py @@ -0,0 +1,205 @@ +#!/usr/bin/env python3 +"""Extract ORBIT-2 training FOMs from a perf-run job directory.""" + +from __future__ import annotations + +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, + 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) + 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}") + 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_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_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..32fa9a4 --- /dev/null +++ b/earth_science/models/ORBIT-2/examples/run_scaling_study.sh @@ -0,0 +1,84 @@ +#!/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_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="${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_prism.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}}" + +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,ORBIT2_DISABLE_CKPT=1 \ + "$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_infer_amd.sh b/earth_science/models/ORBIT-2/examples/sbatch_infer_amd.sh index 2144821..c5b55c0 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 MI355X (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 new file mode 100755 index 0000000..723df79 --- /dev/null +++ b/earth_science/models/ORBIT-2/examples/sbatch_train_amd.sh @@ -0,0 +1,344 @@ +#!/usr/bin/env bash +# ORBIT-2 multi-node training on AMD Instinct via SLURM (Apptainer + MPI). +# +# 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_prism.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_prism.yaml; +# use interm_8m_era5.yaml for new ERA5 1.0_deg data) +# 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) + +#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" +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}" +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}" + +NODES="${SLURM_JOB_NUM_NODES:-1}" +GPUS_PER_NODE=8 +TOTAL_RANKS=$((NODES * GPUS_PER_NODE)) + +# 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 + +_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 + +mkdir -p "$ORBIT2_OUTPUT_DIR" + +ORBIT2_CONFIG_TEMPLATE="${ORBIT2_CONFIG_TEMPLATE:-interm_8m_prism.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 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 " 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 " 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 "" + +# --------------------------------------------------------------------------- +# 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 + +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}" + +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" +# 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 MI355X (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_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:-}" +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" + +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..7d2f093 --- /dev/null +++ b/earth_science/models/ORBIT-2/examples/sbatch_train_perf_amd.sh @@ -0,0 +1,507 @@ +#!/usr/bin/env bash +# 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 (ERA5 same-dir): +# export AI4S_SHARED_DIR=/path/to/shared +# 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=1 +#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 + +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) +else + SCRIPT_DIR=$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd) +fi + +ORBIT2_BASE="${AI4S_SHARED_DIR:?AI4S_SHARED_DIR must be set}/models/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}" +# 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_prism.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/README.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}" + +: "${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:-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 + 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|@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 + +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" +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) + 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 '/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 + 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}" +TRACE_DIR="${ORBIT2_OUTPUT_DIR}/traces" +# Bayes EDM: default no hook (res_slimvit 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 +#!/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" +# 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 MI355X (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}" +# 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="${_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" + +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 + ) + # 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="${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="${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="${RCCL_LL128_FORCE_ENABLE:-1}" + --env NCCL_IB_PCI_RELAXED_ORDERING=1 + --env NCCL_DMABUF_ENABLE=1 + --env NCCL_DEBUG="${NCCL_DEBUG:-WARN}" + ) + # 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 ---" +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")" \ + --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:/opt/rocm/lib:/usr/lib/x86_64-linux-gnu" \ + "${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 + +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 + + +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": workload, + "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, + "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'}") +PYEOF + +echo "=== ORBIT-2 perf run complete (rc=$TRAIN_RC) ===" +exit $TRAIN_RC 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..53dfb87 --- /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° + `edm_8m_era5_1x8.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=YOUR_PARTITION --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..1e8b147 --- /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 the earth-science SKILL (bf16 Flash 65535 cap). + """ + 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 f6479be..8282913 100644 --- a/earth_science/models/ORBIT-2/model.yaml +++ b/earth_science/models/ORBIT-2/model.yaml @@ -16,12 +16,46 @@ recipes: recipe_path: recipes/inference/ script: examples/run_visualize.py slurm: examples/sbatch_infer_amd.sh + - task: perf-analysis + 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 with `interm_8m_era5.yaml` + 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: - 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; **`interm_8m_prism.yaml`** / **`interm_8m_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/README.md b/earth_science/models/ORBIT-2/recipes/README.md index 5e6a994..b1d2c4c 100755 --- a/earth_science/models/ORBIT-2/recipes/README.md +++ b/earth_science/models/ORBIT-2/recipes/README.md @@ -15,3 +15,10 @@ Studio runbooks for [ORBIT-2](../README.md) (`jychoi-hpc/ORBIT-2` on Hugging Fac |-----|---------| | [data/README.md](data/README.md) | Public dataset via Constellation + Globus; staging on shared cluster FS; OLCF Orion and copy-out notes; config paths `low_res_dir` / `high_res_dir` are yours to fill. | | [compute/README.md](compute/README.md) | Frontier access is restricted; AMD Developer Cloud, AI & HPC Fund, institutional AMD clusters. | + +## Performance engineering (AMD Instinct) + +| Doc | Summary | +|-----|---------| +| [perf-analysis/README.md](perf-analysis/README.md) | Multi-subagent bottleneck analysis (TraceLens + Omnistat) after a perf run; FOM contract, parallelism defaults, landmines. Baselines: [one-node-gpu-baseline.md](perf-analysis/one-node-gpu-baseline.md), [BASELINE_LOCKIN.md](perf-analysis/BASELINE_LOCKIN.md). | +| [perf-optimizer-loop/README.md](perf-optimizer-loop/README.md) | Iterative sysopt loop (one lever/iter, accept/revert on throughput). Levers tried/blocked in [lever_catalog.yaml](perf-optimizer-loop/lever_catalog.yaml); GEMM-time finding in [gemm-attribution.md](perf-optimizer-loop/gemm-attribution.md); HBM staging in [STAGING_ERA5_FOR_HBM.md](perf-optimizer-loop/STAGING_ERA5_FOR_HBM.md). | 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..8be9020 --- /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_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 + fix: see the earth-science SKILL (bf16 Flash 65535 cap). 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_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/README.md b/earth_science/models/ORBIT-2/recipes/perf-analysis/README.md new file mode 100644 index 0000000..4b50156 --- /dev/null +++ b/earth_science/models/ORBIT-2/recipes/perf-analysis/README.md @@ -0,0 +1,68 @@ +# ORBIT-2 Performance Analysis with AMD AI Agents + +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:** Perf defaults **Bayes-CAST** **`edm_8m_era5_1x8.yaml`** + staged **ERA5 1.0°** (`ORBIT2_DATA_ROOT`). **`interm_8m_era5.yaml`** remains for public ORBIT-2 same-dir `res_slimvit`. Loss is not meaningful without true downscaling targets. + +## Quick start (1 node × 8 GPUs, Bayes-CAST EDM template) + +```bash +export AI4S_SHARED_DIR=/path/to/shared +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 # 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 +├── baseline_report.md # from report_orbit2_gpu_baseline.py (optional) +├── baseline_report.json +├── tracelens/ # analyst outputs +└── combined_report.md # synthesizer output +``` + +## 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. Legacy **`batch_time_s`** in older notes maps to the same parser options. + +**Loss sanity:** epoch-completed losses must **strictly decrease** epoch-over-epoch (`loss_sanity_pass`). Same-dir data makes absolute loss non-physical; use as a 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 examples/parse_training_log.py --log ... --steady-epoch-start 2 --warmup-batches-per-epoch 1 +python3 examples/collate_scaling_study.py --log-dir .../logs --jobs ... --steady-epoch-start 2 --require-loss-sanity +``` + +## Defaults & landmines + +- **Parallelism (1 node × 8 GPU):** `fsdp=8`, `simple_ddp=1` unless `ORBIT2_FSDP` / `ORBIT2_SIMPLE_DDP` are set; multi-node renders `fsdp=nodes`, `simple_ddp=8`. +- **Launcher discovery:** when `$ORBIT2_ROOT/launch/train_edm.py` exists (Bayes-CAST), ranks run `train_edm.py /config/config.yaml` from `/orbit2/launch`; do **not** exec upstream `launch_diffusion.sh` (OLCF/Crusher job wrapper — nested `srun`, ignores argv). For public ORBIT-2, `launch_diffusion.sh` + argv applies when `train_edm.py` is absent. +- **Same-dir data:** set `superres_mag: 1` in `interm_8m_prism.yaml` / `interm_8m_era5.yaml`, else tensor-shape mismatch. +- **`max_epochs` must be ≥ 2** (upstream loop is `while (epoch_start + 1) < max_epochs`). +- For the **GEMM-time bottleneck**, MIOpen-flag A/B, and the tabled conv-padding / dead-end `channels_last` levers, see [`../perf-optimizer-loop/gemm-attribution.md`](../perf-optimizer-loop/gemm-attribution.md). Cross-cutting lessons (bf16 Flash 65535 cap, xFormers CK, std_delta/ERA5_1, debug-log handling) live in the earth-science and perf-analysis SKILLs. + +## 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) 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..1c3dce2 --- /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 + +- `${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. +- `$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 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 +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=${PERF_TOOLS_DIR}/perf-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 `${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/`) + +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 `interm_8m.json` from `$AI4S_SHARED_DIR/models/ORBIT-2/code/ORBIT-2/examples/multidataset_hpo_sc26/interm_8m.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_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..5ff92cf --- /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 `perf-inspect` venv at `${PERF_TOOLS_DIR}/perf-inspect/`. +- The VictoriaMetrics binary at `${PERF_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 ${PERF_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 +. ${PERF_TOOLS_DIR}/perf-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..1919bad --- /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_MAX_EPOCH=1 ORBIT2_MAX_BATCHES=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/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/synthesizer.md b/earth_science/models/ORBIT-2/recipes/perf-analysis/agents/synthesizer.md new file mode 100644 index 0000000..9dba0e0 --- /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 `earth_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..ea556d0 --- /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 +. ${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; } +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`" | +| "BF16 GEMM runs far below MI355X peak" | "Run one short discriminator job with `ORBIT2_DATA_TYPE=float32` (same batch) and compare `steady_batch_time_s`" | +| "DataLoader events occupy 18% of step time" | "Increase `ORBIT2_NUM_WORKERS` (template token `__NUM_WORKERS__`) and re-run a short job" | +| "BF16 MFMA approaches MI355X spec ceiling" | "null — system limit; only path is larger batch / wider model" | + +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..8861684 --- /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 vs float32 discriminator 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 `intermediate_downscaling.py` (or `examples/run_orbit2_train.py`) for ~10 batches inside the container with the alternate dtype. + +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..deecfd7 --- /dev/null +++ b/earth_science/models/ORBIT-2/recipes/perf-analysis/omnistat.config.template @@ -0,0 +1,147 @@ +#-- +# 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 ENABLED by default (bf16 MFMA + HBM read); override per-probe if needed +# - 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 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). +# +# 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] +# 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: +# * 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 = @PERF_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-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..f379455 --- /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 **`interm_8m_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_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 [README.md](README.md) §Defaults & landmines). 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_prism.yaml`, `interm_8m_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_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 new file mode 100644 index 0000000..f01e32a --- /dev/null +++ b/earth_science/models/ORBIT-2/recipes/perf-optimizer-loop/README.md @@ -0,0 +1,117 @@ +# ORBIT-2 — iterative systems optimizer loop (1 node × MI355X, Bayes-CAST EDM) + +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/). + +> **Audience:** AMD performance engineering. Not a scientific downscaling claim — same-dir ERA5 configs are timing-friendly. + +## Prerequisites + +- **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 [perf-analysis/README.md](../perf-analysis/README.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 [`gemm-attribution.md`](gemm-attribution.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..e4f158e --- /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 the studio SKILL `std_delta`/`ERA5_1` note). +- **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 (see the earth-science SKILL xFormers MEA note). + +## 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/gemm-attribution.md b/earth_science/models/ORBIT-2/recipes/perf-optimizer-loop/gemm-attribution.md new file mode 100644 index 0000000..89fad9a --- /dev/null +++ b/earth_science/models/ORBIT-2/recipes/perf-optimizer-loop/gemm-attribution.md @@ -0,0 +1,69 @@ +# GEMM-time attribution — Bayes-CAST EDM (bf16, ERA5 1.0°) + +ORBIT-2 analog of HydraGNN's [`dispatch-attribution.md`](../../../../../material_science/models/HydraGNN/recipes/perf-optimizer-loop/dispatch-attribution.md): the durable "where does the step time go" finding for the EDM path, plus the levers tried against it (see [`lever_catalog.yaml`](lever_catalog.yaml) for the accept/reject record). + +**Workload:** profiled EDM, bf16, batch 4096, ERA5 1.0°, 1- and 2-node (jobs 12562 / 12648). +**Driver:** [`agents/orchestrator_gemm_analysis.md`](../perf-analysis/agents/orchestrator_gemm_analysis.md) → [`examples/run_gemm_analysis.sh`](../../examples/run_gemm_analysis.sh) (TraceLens + Omnistat analyst/verifier). +**Report:** `$AI4S_SHARED_DIR/models/ORBIT-2/perf-runs/gemm-analysis-/GEMM_TIME_REPORT.md`. + +## Attribution (authoritative for this workload) + +- **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 is why TunableOp gives no uplift** — it tunes 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`](../../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 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 `res_slimvit` hook did not apply). + +## Levers tried against the GEMM bottleneck + +### ORNL Frontier MIOpen flags — performance-NEUTRAL (jobs 12997/12998 vs 12562) + +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 non-Slingshot IB/ionic): `FI_CXI_*`, `NCCL_NET="AWS Libfabric"`, `NCCL_NET_PLUGIN=librccl-net.so`, `NCCL_SOCKET_IFNAME=hsn0`, Frontier `LD_PRELOAD` host paths. + +**Result (2026-06-16): performance-NEUTRAL on MI355X** (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 — the bottleneck is shape-driven, not algorithm-selection-driven. Keeping the flags as faithful, harmless defaults; they are **not** the lever here. Report: `perf-runs/miopen-ornl-validation-/MIOPEN_ORNL_VALIDATION.md`. + +### Conv channel-padding (`ORBIT2_CONV_PAD`) — TABLED, timing-only (NOT ADOPTED, jobs 14633/14634) + +> **Status: tabled.** The −25% is a *timing-only* result. Channel-padding **changes model capacity** (wider internal convs ⇒ different params) and **has no 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. 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 attribution: 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. If revisited, the open timing question is **N=32** (top GEMM `MT16x16x128` at 27% still has headroom). Report: `perf-runs/conv-pad-validation-/CONV_PAD_VALIDATION.md`. + +### `channels_last` + implicit-GEMM — DEAD END on ROCm 7.2.2 (DISPROVEN, jobs 14645–14648) + +Tested stacking `channels_last` on the padding idea as a clean 2×2 ({NCHW, channels_last} × {pad=0, pad=16}) via 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 (~88–91% of GPU-kernel time). 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`; 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`. + +## Reproduce + +```bash +export AI4S_SHARED_DIR=/path/to/shared +export PERF_TOOLS_DIR=/path/to/perf-tools +# 1- and 2-node profiled GEMM analysis (TraceLens + Omnistat analyst/verifier) +bash earth_science/models/ORBIT-2/examples/run_gemm_analysis.sh +# Tool-independent kernel ranking from a single trace: +python3 earth_science/models/ORBIT-2/examples/compare_trace_kernels.py \ + --trace $AI4S_SHARED_DIR/models/ORBIT-2/perf-runs//traces/orbit2-epoch*-rank0.pt.trace.json +``` 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..3d8fe71 --- /dev/null +++ b/earth_science/models/ORBIT-2/recipes/perf-optimizer-loop/lever_catalog.yaml @@ -0,0 +1,174 @@ +# 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 BASELINE_LOCKIN.md and the earth-science SKILL (bf16 Flash 65535 cap). 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 gemm-attribution.md (channels_last dead end)." + + - 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-attribution.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: conv_pad_internal_width + description: "Widen path2/refine conv INTERNAL width to a tile-friendly N (3-channel I/O unchanged) via ORBIT2_CONV_PAD=N in edm.py, to kill the starved MT16x16x32 conv-bwd GEMM." + catalog_order: 10 + kind: env_only + status: "tabled" + env_vars: + ORBIT2_CONV_PAD: "16" + diagnostic_only: false + expected_payoff: "high (timing-only -25%)" + risk: "high" + revert_method: "drop_env" + citation: "gemm-attribution.md" + notes: "TABLED (NOT ADOPTED, default 0/off). Single-node A/B (bf16/4096/ERA5, N=16): steady 8.646 -> 6.498 s (-24.8%), top MT16x16x32 GEMM gone. BUT this CHANGES MODEL CAPACITY (wider internal convs) and only 'loss sanity' (finite/decreasing) was checked, NOT convergence/quality. Do NOT enable for real training until a matched convergence study (same data/schedule; final val loss/skill + steps-to-target, pad=0 vs pad=16) passes. Open timing question if revisited: N=32. See gemm-attribution.md." + + - id: kernel_trace_diag_only + description: "OMNISTAT_KERNEL_TRACE=1 for one iter — high VM cardinality." + catalog_order: 11 + 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..ad6f7eb 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 templates (`interm_8m_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). diff --git a/material_science/models/HydraGNN/README.md b/material_science/models/HydraGNN/README.md index deddfad..1106fd3 100644 --- a/material_science/models/HydraGNN/README.md +++ b/material_science/models/HydraGNN/README.md @@ -14,6 +14,8 @@ Authoritative training, HPO, ensemble, and inference scripts live in **upstream - [`recipes/train/`](recipes/train/) — Frontier-scale pretraining / HPO summary - [`recipes/data/`](recipes/data/) — Hub ADIOS bundles, Constellation, staging, and data context - [`recipes/compute/`](recipes/compute/) — OLCF, AMD cloud / HPC Fund, institutional clusters +- [`recipes/perf-analysis/`](recipes/perf-analysis/) — multi-subagent bottleneck analysis (TraceLens + Omnistat) on AMD Instinct +- [`recipes/perf-optimizer-loop/`](recipes/perf-optimizer-loop/) — iterative epoch-time optimizer loop + lever catalog - [`examples/`](examples/) — ready-to-run Docker and run scripts ## Installation (reference) diff --git a/material_science/models/HydraGNN/examples/README.md b/material_science/models/HydraGNN/examples/README.md index 12a69ca..2e3cb99 100644 --- a/material_science/models/HydraGNN/examples/README.md +++ b/material_science/models/HydraGNN/examples/README.md @@ -6,12 +6,55 @@ with AMD/ROCm container launch. ## Files +The folder is flat (scripts resolve siblings and the container bind-mount +relative to their own directory), but the files fall into the groups below. + +### Setup & diagnostics + | Script | Purpose | |--------|---------| | [`preflight_hydragnn.py`](preflight_hydragnn.py) | Check your environment before submitting jobs | -| [`run_inference.sh`](run_inference.sh) | Load a checkpoint and run predictions | +| [`microbench_node_health.sh`](microbench_node_health.sh) | ~30 s per-node MI355X (gfx950) health survey (host/GPU inventory, dual-NUMA STREAM, HIP launch latency); safe to run in a SLURM Prolog | + +### Run entrypoints + +| Script | Purpose | +|--------|---------| | [`run_train.sh`](run_train.sh) | Train on bundled or staged ADIOS datasets | +| [`run_inference.sh`](run_inference.sh) | Load a checkpoint and run predictions | + +### Container & launch drivers + +| Script | Purpose | +|--------|---------| | [`docker_run.sh`](docker_run.sh) | Docker launcher (clones HydraGNN, installs deps) | +| [`build_overlay_amd.sh`](build_overlay_amd.sh) | One-time Apptainer ext3 overlay build (pre-loads pip deps; reuse across jobs to skip the per-job pip install) | +| [`sbatch_train_amd.sh`](sbatch_train_amd.sh) | SLURM/Apptainer multi-node training (HPC) | +| [`sbatch_train_perf_amd.sh`](sbatch_train_perf_amd.sh) | 2-node training + PyTorch profiler + Omnistat telemetry (thin variant of `sbatch_train_amd.sh` for the perf-analysis recipe) | +| [`sbatch_infer_amd.sh`](sbatch_infer_amd.sh) | SLURM/Apptainer inference on AMD Instinct (needs a prebuilt overlay + downloaded weights) | + +### Perf & scaling orchestration + +| Script | Purpose | +|--------|---------| +| [`run_scaling_study.sh`](run_scaling_study.sh) | Submit matched 1/2/4/8-node strong-scaling sweep | +| [`run_optimizer_loop.sh`](run_optimizer_loop.sh) | Entrypoint for the iterative sysopt **perf-optimizer-loop** (Claude CLI; run under `tmux`) | + +### Post-processing & analysis + +Host-side; run after a job against its logs / traces. + +| Script | Purpose | +|--------|---------| +| [`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 | +| [`run_fom_extractor.py`](run_fom_extractor.py) | Compute FOMs + TraceLens↔Omnistat `kernel_correlation.csv` (login-node post-processing of a completed perf run) | + +### Patches + +| Path | Purpose | +|--------|---------| +| [`patches/`](patches/README.md) | Opt-in upstream patches applied during the overlay build (see [`patches/README.md`](patches/README.md)) | ## Quick start — Docker @@ -26,6 +69,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/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/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/microbench_node_health.sh b/material_science/models/HydraGNN/examples/microbench_node_health.sh new file mode 100755 index 0000000..a242de0 --- /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. +# 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: $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 +# 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=YOUR_GPU_PARTITION +#SBATCH --account=YOUR_ACCOUNT +#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:-${AI4S_SHARED_DIR:-/tmp}/microbench-node-health}" +HOST_OUT="$OUT_DIR/$H" +mkdir -p "$HOST_OUT" + +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}" +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| ${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 ---------- +# 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 "$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 +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/examples/parse_convergence.py b/material_science/models/HydraGNN/examples/parse_convergence.py index f5fadaa..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: @@ -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/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/run_fom_extractor.py b/material_science/models/HydraGNN/examples/run_fom_extractor.py new file mode 100755 index 0000000..7565c36 --- /dev/null +++ b/material_science/models/HydraGNN/examples/run_fom_extractor.py @@ -0,0 +1,429 @@ +#!/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/$USER + 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] + +Override the VictoriaMetrics binary directly with VICTORIA_METRICS_BIN if needed. +""" + +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['PERF_TOOLS_DIR']}/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() + + 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": "", + "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/examples/run_optimizer_loop.sh b/material_science/models/HydraGNN/examples/run_optimizer_loop.sh new file mode 100755 index 0000000..a5b046e --- /dev/null +++ b/material_science/models/HydraGNN/examples/run_optimizer_loop.sh @@ -0,0 +1,240 @@ +#!/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 $AI4S_SHARED_DIR/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 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)}" +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 -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 — check partition names in .cluster-config.yaml") + 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 \ + "${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)") + 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 # 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") + 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 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=" + +# 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} model=${HYDRAGNN_CLAUDE_MODEL}" + set +e + claude \ + --print \ + --model "$HYDRAGNN_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} 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/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 20dfe2c..ffaf998 100755 --- a/material_science/models/HydraGNN/examples/sbatch_train_amd.sh +++ b/material_science/models/HydraGNN/examples/sbatch_train_amd.sh @@ -35,8 +35,12 @@ # 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) +# HydraGNN pinned SHA: 2fb0bd0157e3c85a74f9841887155095bd163303 (main) #SBATCH --job-name=hydragnn-train #SBATCH --partition=YOUR_GPU_PARTITION @@ -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)) @@ -101,7 +109,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" @@ -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'), ] @@ -311,6 +332,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 +361,14 @@ 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 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" \ 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 new file mode 100755 index 0000000..cbf5c0e --- /dev/null +++ b/material_science/models/HydraGNN/examples/sbatch_train_perf_amd.sh @@ -0,0 +1,611 @@ +#!/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. 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. +# +# Quick start: +# export AI4S_SHARED_DIR=/shared/$USER +# sbatch material_science/models/HydraGNN/examples/sbatch_train_perf_amd.sh +# +# 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) +# +# 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=YOUR_GPU_PARTITION +#SBATCH --account=YOUR_ACCOUNT +#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" +# 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)}" +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:-${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 +# 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.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:-${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 +# 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 +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 + +# 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 + 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:-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" + 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" \ + -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 +# 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 -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 " 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. +# --------------------------------------------------------------------------- +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 " 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 " 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 "" + +# --------------------------------------------------------------------------- +# 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 "'"$HG_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.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 + +# 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:-}" + +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" +# 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}" + +# 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 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. +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 + +# 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 + +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 + +# --------------------------------------------------------------------------- +# 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 + +# --------------------------------------------------------------------------- +# 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.) +# +# 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 \ + 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 +# --------------------------------------------------------------------------- +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}" \ + "${KTRACE_BIND_ENVS[@]}" \ + --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_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" \ + "${OPT_ENVS[@]}" \ + "${HOOK_BIND_ENVS[@]}" \ + --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..30e546c 100644 --- a/material_science/models/HydraGNN/model.yaml +++ b/material_science/models/HydraGNN/model.yaml @@ -25,6 +25,15 @@ 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 + - 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/README.md b/material_science/models/HydraGNN/recipes/README.md index b13c0ac..8c0fe31 100644 --- a/material_science/models/HydraGNN/recipes/README.md +++ b/material_science/models/HydraGNN/recipes/README.md @@ -15,3 +15,10 @@ Studio runbooks for [HydraGNN (Predictive GFM 2024)](../README.md) (`mlupopa/Hyd |-----|---------| | [data/README.md](data/README.md) | `ADIOS_files/` on Hub (CLI staging), Constellation / Globus mirror, OLCF copy-out notes; large downloads; energy/force preprocessing summarized on the model card. | | [compute/README.md](compute/README.md) | Frontier access is restricted; AMD cloud / HPC Fund, institutional AMD clusters; inference vs pretraining footprint. | + +## Performance engineering (AMD Instinct) + +| Doc | Summary | +|-----|---------| +| [perf-analysis/README.md](perf-analysis/README.md) | Multi-subagent bottleneck analysis (TraceLens + Omnistat) after a 2-node perf run; artifacts, bottleneck taxonomy, telemetry knobs. | +| [perf-optimizer-loop/README.md](perf-optimizer-loop/README.md) | Iterative sysopt loop (one lever/iter, accept/revert on epoch time). Levers tried/blocked in [lever_catalog.yaml](perf-optimizer-loop/lever_catalog.yaml); current-best attribution in [dispatch-attribution.md](perf-optimizer-loop/dispatch-attribution.md). | 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..0d0405e --- /dev/null +++ b/material_science/models/HydraGNN/recipes/perf-analysis/README.md @@ -0,0 +1,125 @@ +# HydraGNN Performance Analysis with AMD AI Agents + +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:** 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`) 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/$USER +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 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: + +``` +$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) +├── 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 + +- 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}/`. +- 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` | `${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. + +## 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. +- 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 `$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 new file mode 100644 index 0000000..63ea2dc --- /dev/null +++ b/material_science/models/HydraGNN/recipes/perf-analysis/agents/launcher.md @@ -0,0 +1,215 @@ +# launcher subagent + +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 + +- 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 + +- `${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. +- `$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 + +```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, + "hydragnn_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 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 +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=${PERF_TOOLS_DIR}/perf-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 `${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/`) + +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 `@JOB_DIR@` / `@OMNIHUB_TOOLS_DIR@` placeholders, which the sbatch wrapper substitutes (via `sed`) into a per-job config under the perf-run directory. + +### 3. Generate the per-job profile config + +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"] = { + "enable": 1, + "target_epoch": int(os.environ.get("PROFILE_TARGET_EPOCH", "1")), +} +``` + +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/$USER +export HG_OUTPUT_DIR=$AI4S_SHARED_DIR/models/HydraGNN/perf-runs/PENDING-$$ +mkdir -p "$HG_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="${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=$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) +# 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..5ff92cf --- /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 `perf-inspect` venv at `${PERF_TOOLS_DIR}/perf-inspect/`. +- The VictoriaMetrics binary at `${PERF_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 ${PERF_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 +. ${PERF_TOOLS_DIR}/perf-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/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..81ab39c --- /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 -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 `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..99beea4 --- /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 + +- `$AI4S_SHARED_DIR/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 -A -N1 --time=00:05:00 ... | null" +} +``` + +## Steps + +### 1. Sanity check + +```bash +. ${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; } +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 `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..c1d668f --- /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 -A -N1 --ntasks=8 --gpus-per-node=8 --cpus-per-task=16 \ + --time=00:05:00 --pty bash -c ' +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 "$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 +' +``` + +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.config.template b/material_science/models/HydraGNN/recipes/perf-analysis/omnistat.config.template new file mode 100644 index 0000000..9d55579 --- /dev/null +++ b/material_science/models/HydraGNN/recipes/perf-analysis/omnistat.config.template @@ -0,0 +1,135 @@ +#-- +# 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: +# Differences from upstream / from the system-mode default: +# - enable_amd_smi (newer SMI), not rocm_smi +# - rocprofiler counters ENABLED by default (enable_rocprofiler=True; see [omnistat.collectors.rocprofiler]) +# - 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 +# - 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 = @PERF_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/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..5c2ac80 --- /dev/null +++ b/material_science/models/HydraGNN/recipes/perf-optimizer-loop/README.md @@ -0,0 +1,182 @@ +# 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. + +> **Picking up this work?** Start with [`dispatch-attribution.md`](dispatch-attribution.md) (current-best attribution, ~75 s/epoch) and [`lever_catalog.yaml`](lever_catalog.yaml) (every lever tried, accepted/blocked with evidence). Cross-cutting lessons live in [`.cursor/skills/ai4science-perf-analysis/SKILL.md`](../../../../.cursor/skills/ai4science-perf-analysis/SKILL.md). + +## 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 +export ANTHROPIC_API_KEY=... # user sets manually; never committed +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. +``` + +Morning workflow: + +```bash +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) + +``` +$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) +│ ├── 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 $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. 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. + +## 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..9d97925 --- /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 +. "${PERF_TOOLS_DIR}/perf-inspect/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 "$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" +``` + +`_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 "${PERF_TOOLS_DIR}/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. 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}. + +## 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..000aca6 --- /dev/null +++ b/material_science/models/HydraGNN/recipes/perf-optimizer-loop/agents/lever_picker.md @@ -0,0 +1,123 @@ +# 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). +- 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. + +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..85cc223 --- /dev/null +++ b/material_science/models/HydraGNN/recipes/perf-optimizer-loop/agents/orchestrator.md @@ -0,0 +1,238 @@ +# 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 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 `$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`. +- `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 `$AI4S_SHARED_DIR/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 `$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 + +| 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 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: + +- 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 ${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. + +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-/`. + +**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 "$REPO_ROOT"; \ + 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=` (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.) + +**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=`. + +**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 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. + +**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 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 | +| 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..f296ca4 --- /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). +**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 omnistat venv at `${PERF_TOOLS_DIR}/perf-inspect/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/dispatch-attribution.md b/material_science/models/HydraGNN/recipes/perf-optimizer-loop/dispatch-attribution.md new file mode 100644 index 0000000..ebedaf7 --- /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:** `$AI4S_SHARED_DIR/models/HydraGNN/perf-runs/7187/` +**Lever:** `rccl_high_priority` (`TORCH_NCCL_HIGH_PRIORITY=1`, `GPU_MAX_HW_QUEUES=2`) +**Nodes:** 2 × MI355X (8 GPUs each) + +## 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/$USER +export REPO_ROOT=$HOME/git/ai4science-studio +source ${PERF_TOOLS_DIR}/perf-inspect/bin/activate + +# VictoriaMetrics on the perf-run DB (login node) +${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 & + +cd $AI4S_SHARED_DIR/models/HydraGNN/perf-runs/7187/tracelens +TraceLens_generate_perf_report_pytorch \ + --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 + +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 +``` 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..2543837 --- /dev/null +++ b/material_science/models/HydraGNN/recipes/perf-optimizer-loop/lever_catalog.yaml @@ -0,0 +1,282 @@ +# 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: "$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%)." + catalog_order: 1 + kind: env_only + env_vars: + HG_BATCH_SIZE: "800" + diagnostic_only: false + expected_payoff: "none-for-MLIP-MACE" + risk: "high" + revert_method: "drop_env" + 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." + 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: "reduce-overhead" + rank_script_patch: | + # 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: "none (blocked)" + risk: "blocked" + revert_method: "drop_hook" + 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 + 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: "$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 + 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: "1.3%-realized" + 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. 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." + 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: "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" + PYTORCH_TUNABLEOP_TUNING: "0" + 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 $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 + # (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(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: "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: "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." + catalog_order: 5 + kind: env_only + env_vars: + HYDRAGNN_NUM_WORKERS: "12" + diagnostic_only: false + expected_payoff: "unreliable (variance > effect)" + risk: "low" + revert_method: "drop_env" + 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%)." + catalog_order: 6 + kind: env_only + env_vars: + NCCL_MIN_NCHANNELS: "112" + diagnostic_only: false + expected_payoff: "negative-at-N=2" + risk: "medium" + revert_method: "drop_env" + 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. 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: "negligible-but-diagnostically-valuable" + 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. 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." + 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." diff --git a/material_science/models/HydraGNN/recipes/train/README.md b/material_science/models/HydraGNN/recipes/train/README.md index a607177..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 @@ -247,7 +253,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`).