diff --git a/.claude/commands/init-cluster.md b/.claude/commands/init-cluster.md index 9bb77a9..7a52b9b 100644 --- a/.claude/commands/init-cluster.md +++ b/.claude/commands/init-cluster.md @@ -69,9 +69,9 @@ echo "HTTPS_PROXY=${HTTPS_PROXY:-}" find "$HOME" /scratch /projects /opt -maxdepth 4 -name "*.sif" 2>/dev/null | head -20 # 14. Performance tooling (perf-analysis / perf-optimizer-loop recipes — optional) -# Look for an existing shared perf-tools dir (omnistat + TraceLens venv "perf-inspect/"). +# Look for an existing shared perf-tools dir (omnistat + TraceLens venv "omnihub-inspect/"). for d in /shared/*/tools /shared/perf-tools "$HOME/perf-tools"; do - [[ -d "$d/perf-inspect" || -d "$d/omnistat-src" ]] && echo "perf_tools: $d" + [[ -d "$d/omnihub-inspect" || -d "$d/omnistat-src" ]] && echo "perf_tools: $d" done ``` @@ -127,8 +127,8 @@ If no internet, ask for proxy settings. **Q9. Performance tooling** (optional — only for `perf-analysis` / `perf-optimizer-loop` recipes) If a perf-tools dir was discovered (step 14), show it and ask to confirm; otherwise ask for the path -or leave blank. This is exported to perf scripts as `PERF_TOOLS_DIR`; expected layout under it: -`perf-inspect/` (omnistat + TraceLens venv), `omnistat-src/`, `victoriametrics/victoria-metrics-prod`. +or leave blank. This is exported to perf scripts as `OMNIHUB_TOOLS_DIR`; expected layout under it: +`omnihub-inspect/` (omnistat + TraceLens venv), `omnistat-src/`, `victoriametrics/victoria-metrics-prod`. **Q10. Config file location** - `.cluster-config.yaml` in this repo (recommended — per-checkout) @@ -173,9 +173,9 @@ network: libionic_path: "" # path to libionic.so.1 (default: /usr/lib/x86_64-linux-gnu/libionic.so.1) # Performance tooling — only if using perf-analysis / perf-optimizer-loop recipes (Q9). -# Exported to scripts as PERF_TOOLS_DIR. Leave dir blank if not used. +# Exported to scripts as OMNIHUB_TOOLS_DIR. Leave dir blank if not used. perf_tools: - dir: "" # layout: perf-inspect/, omnistat-src/, victoriametrics/ + dir: "" # layout: omnihub-inspect/, omnistat-src/, victoriametrics/ discovered_sifs: # SIF files found during init (informational — confirm they are under sif_cache) @@ -196,4 +196,4 @@ Suggest what the user can do next: - 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 +$ARGUMENTS \ No newline at end of file diff --git a/.claude/commands/run-perf-hydragnn.md b/.claude/commands/run-perf-hydragnn.md index c98abdf..f527fa2 100644 --- a/.claude/commands/run-perf-hydragnn.md +++ b/.claude/commands/run-perf-hydragnn.md @@ -15,7 +15,7 @@ Read `.cluster-config.yaml` (repo root) or `~/.config/ai4science-studio/cluster. **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). +- Set **`OMNIHUB_TOOLS_DIR`** from `omnihub.tools_dir`. If missing, stop and send the user to **`/init-cluster`** Q9 (Performance tooling). Pre-fill SLURM partition/account from config. @@ -92,7 +92,7 @@ Confirm branch **`Predictive_GFM_2024`** if user cares about upstream alignment ```bash export AI4S_SHARED_DIR= -export PERF_TOOLS_DIR= +export OMNIHUB_TOOLS_DIR= # Optional overrides: # export HG_SIF=... HG_OVERLAY=... HG_REPO_DIR=... HG_DATASETS=... HG_DATA_DIR=... ``` diff --git a/.claude/commands/run-perf-orbit2.md b/.claude/commands/run-perf-orbit2.md index 899d2d3..6cff7e5 100644 --- a/.claude/commands/run-perf-orbit2.md +++ b/.claude/commands/run-perf-orbit2.md @@ -13,7 +13,7 @@ Read `.cluster-config.yaml` (repo root) or `~/.config/ai4science-studio/cluster. **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)**. +- `omnihub.tools_dir` → use as **`OMNIHUB_TOOLS_DIR`** (layout: `omnihub-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. @@ -97,7 +97,7 @@ Export from cluster config + user answers: ```bash export AI4S_SHARED_DIR= -export PERF_TOOLS_DIR= +export OMNIHUB_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 diff --git a/.cluster-config.example.yaml b/.cluster-config.example.yaml index ea5cd7b..18d7007 100644 --- a/.cluster-config.example.yaml +++ b/.cluster-config.example.yaml @@ -41,9 +41,18 @@ network: 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 +# OmniHub integration (sibling repo — see integrations/omnihub/README.md) +omnihub: + repo_dir: "" # e.g. /home/$USER/git/omnihub + cluster: vultr # maps to omnihub config/vultr.yaml on Lux + shared_dir: /shared/omnihub + results_dir: "" # default: /shared/$USER/results/omnihub + rocm_version: "7.2.2" + # Shared perf/telemetry tools. Cluster-specific — set per site, do not hardcode + # these into recipes/scripts. Default: /tools. + # The HydraGNN/ORBIT-2 perf recipes read this value via the OMNIHUB_TOOLS_DIR + # env var (export OMNIHUB_TOOLS_DIR= before running them). + tools_dir: "" # default: /shared/omnihub/tools + # /omnistat GPU telemetry, no PMC (--tools omnistat) + # /omnistat-rocprofiler adds rocprofiler PMC counters (--tools omnistat-rocprofiler-pmc1) + # /omnihub-inspect omnistat-inspect + TraceLens venv (legacy deep-dive) diff --git a/.cursor/skills/ai4science-omnihub/SKILL.md b/.cursor/skills/ai4science-omnihub/SKILL.md new file mode 100644 index 0000000..96428f3 --- /dev/null +++ b/.cursor/skills/ai4science-omnihub/SKILL.md @@ -0,0 +1,97 @@ +--- +name: ai4science-omnihub +description: Launch AI4Science Studio models via OmniHub job generation, SLURM submission, and standardized result processing. Use when the user asks for OmniHub, train-omnihub, perf-analysis-omnihub, omnihub-generate-job, or post-processing OmniHub results for ai4science models. +--- + +# AI4Science Studio + OmniHub + +Integrates this repo with a sibling [OmniHub](https://github.com/AMDResearch/omnihub) checkout **without merging repos**. + +## Prerequisites + +1. Read `.cluster-config.yaml` (or `~/.config/ai4science-studio/cluster.yaml`) — especially `omnihub:` and `paths`. +2. Set `export OMNIHUB_DIR=/path/to/omnihub` and `export AI4S_SHARED_DIR=/shared/$USER` (or your site path). +3. HydraGNN assets: SIF, overlay, datasets, code clone under `$AI4S_SHARED_DIR/models/HydraGNN/`. + +> **Omnistat note:** OmniHub's `/shared/omnihub/tools/omnistat` has no `omnistat-inspect`, so don't expect inspect JSON from OmniHub jobs — read `omnihub-process` output under `processed-data/` instead. For deep `omnistat-inspect`/PromQL workflows use the `omnihub-inspect` build at `/shared/omnihub/tools/omnihub-inspect` (see perf-analysis recipe). + +## Workflow: HydraGNN train-omnihub + +1. Read `material_science/models/HydraGNN/model.yaml` recipe `train-omnihub`. +2. Sync app stubs and generate SLURM script: + +```bash +./integrations/omnihub/generate-job.sh \ + --num-nodes 2 \ + --partition lux \ + --time-limit 2h \ + --output /tmp/hydragnn-omnihub.slurm +``` + +3. Submit (account on `sbatch`, not in the generated script): + +```bash +sbatch -A vultr_lux /tmp/hydragnn-omnihub.slurm +``` + +4. Monitor: `squeue -j `, `tail -f /shared/$USER/results/omnihub//*.out` + +5. After completion — post-process: + +```bash +$OMNIHUB_DIR/omnihub-process --results-dir /shared/$USER/results/omnihub -j 4 +$OMNIHUB_DIR/omnihub-index --results-dir /shared/$USER/results/omnihub --output index +``` + +## Workflow: perf-analysis-omnihub (Phase 2) + +Same as train, but add `--perf` to enable `omnistat`, `pytorch-trace`, and `tracelens`: + +```bash +./integrations/omnihub/generate-job.sh --perf --num-nodes 2 --time-limit 30m \ + --output /tmp/hydragnn-perf.slurm +sbatch -A vultr_lux /tmp/hydragnn-perf.slurm +``` + +**Agent result reading order** (minimal tokens): + +1. `/processed-data/tracelens-summary.json` +2. `/processed-data/omnistat-range.yaml` +3. `/processed-data/report-card.yaml` +4. Raw traces under `/tools/` only if needed + +Legacy HydraGNN perf (`sbatch_train_perf_amd.sh` + `omnihub-inspect` + subagent claims) remains for deep PromQL/inspect workflows. See `material_science/models/HydraGNN/recipes/perf-analysis/README.md`. + +## Key paths + +Cluster-specific paths come from `.cluster-config.yaml` (`omnihub.shared_dir`, `omnihub.tools_dir`, `omnihub.results_dir`). Resolve them from there rather than assuming literals. The values below are the **Lux / vultr** defaults (`shared_dir: /shared/omnihub`, `tools_dir: /shared/omnihub/tools`). + +| Resource | Path (`` = `omnihub.tools_dir`) | +|----------|------| +| OmniHub shared | `omnihub.shared_dir` (Lux: `/shared/omnihub`) | +| Omnistat | `/omnistat/` → use `--tools omnistat` | +| Omnistat + PMC | `/omnistat-rocprofiler/` → use `--tools omnistat-rocprofiler-pmc1` for hardware counters | +| Omnistat inspect | `/omnihub-inspect/` → has `omnistat-inspect` + TraceLens CLI (legacy deep-dive venv) | +| Results | `omnihub.results_dir` (Lux: `/shared/$USER/results/omnihub/$SLURM_JOB_ID/`) | +| HydraGNN overlay | `$AI4S_SHARED_DIR/models/HydraGNN/overlays/hydragnn-overlay.img` | + +## Runner + +HydraGNN uses MPI — default `--runner manual-mpi --tasks-per-node 8` (OmniHub fork). Do **not** use `--runner manual` (that enables torch DDP init). + +## Related OmniHub skills + +When working inside `$OMNIHUB_DIR`, also use: + +- `.cursor/skills/generate-job-and-run/SKILL.md` +- `.cursor/skills/post-process-results/SKILL.md` +- `.cursor/skills/check-job-results/SKILL.md` + +## Troubleshooting + +| Issue | Fix | +|-------|-----| +| App config not found | Run `sync-to-omnihub.sh` first | +| ROCm / overlay missing | Verify `HG_OVERLAY`, `HG_SIF`; check generate-job injected binds | +| Omnistat empty | Confirm `/shared/omnihub/tools/omnistat/` exists and VictoriaMetrics binary is present | +| TraceLens skipped | Install venv at `/shared/omnihub/tools/tracelens/venv` or set `TRACELENS_VENV` | diff --git a/.cursor/skills/ai4science-run-models/SKILL.md b/.cursor/skills/ai4science-run-models/SKILL.md index 7ea7106..7b46896 100644 --- a/.cursor/skills/ai4science-run-models/SKILL.md +++ b/.cursor/skills/ai4science-run-models/SKILL.md @@ -33,6 +33,7 @@ Match the user's request to a recipe in `model.yaml`: - "run inference" / "predict" / "generate" → look for `task: inference` - "train" / "fine-tune" / "pair-tune" → look for `task: train` or `task: finetune` - "ensemble" → look for `task: ensemble` +- "omnihub" / "train-omnihub" → look for `task: train-omnihub` and use `.cursor/skills/ai4science-omnihub/SKILL.md` ## Step 3: Ask the user for required inputs diff --git a/.cursor/skills/ai4science-studio/SKILL.md b/.cursor/skills/ai4science-studio/SKILL.md index 34e0b47..1e705f4 100755 --- a/.cursor/skills/ai4science-studio/SKILL.md +++ b/.cursor/skills/ai4science-studio/SKILL.md @@ -677,18 +677,18 @@ Omnistat exposes **two** independent rocprofiler-sdk integrations, and there's a **Cross-runtime gotcha:** `enable_kernel_trace=True` without the tool library loaded does *nothing* — the omnistat collector just sits with an empty endpoint. The `OMNISTAT_KERNEL_TRACE=1` switch in `sbatch_train_perf_amd.sh` enforces this with a hard `exit 2` if `libomnistat_trace.so` isn't found at the configured path; do **not** loosen that check, otherwise you re-create the exact "silent zero counters" failure mode from §8. -**Build location:** Both omnistat artifacts (the Python extension and the kernel-trace tool library) must be built on a **compute node** inside the same SIF the workload runs in — login nodes don't have apptainer or the ROCm headers, and a host-side build will pick up the wrong libcurl / glibc. `PERF_TOOLS_DIR` is the shared tools dir from `.cluster-config.yaml` `perf_tools.dir` (e.g. `/path/to/perf-tools`); export it first. Standard build: +**Build location:** Both omnistat artifacts (the Python extension and the kernel-trace tool library) must be built on a **compute node** inside the same SIF the workload runs in — login nodes don't have apptainer or the ROCm headers, and a host-side build will pick up the wrong libcurl / glibc. `OMNIHUB_TOOLS_DIR` is the shared tools dir from `.cluster-config.yaml` `omnihub.tools_dir` (e.g. `/shared/omnihub/tools`); export it first. Standard build: ```bash salloc -p -A -N 1 --time=00:15:00 --gpus-per-node=1 \ apptainer exec --rocm \ "${AI4S_SHARED_DIR}/images/pytorch_rocm7.2.2_ubuntu24.04_py3.12_pytorch_release_2.10.0.sif" \ - bash -c "cd ${PERF_TOOLS_DIR}/omnistat-src && \ + bash -c "cd ${OMNIHUB_TOOLS_DIR}/omnistat-src && \ cmake -S rocprofiler-sdk/ -B build-trace/ -DBUILD_KERNEL_TRACE_LIB=ON && \ cmake --build build-trace/ -j 8" ``` -Verify with `ls -la ${PERF_TOOLS_DIR}/omnistat-src/build-trace/libomnistat_trace.so` (~MB-class file, not the tiny 4 kB stub you get when CMake fails halfway). +Verify with `ls -la ${OMNIHUB_TOOLS_DIR}/omnistat-src/build-trace/libomnistat_trace.so` (~MB-class file, not the tiny 4 kB stub you get when CMake fails halfway). **SIF build prerequisites (one-time gotchas):** the `pytorch_rocm7.2.2_ubuntu24.04_py3.12_pytorch_release_2.10.0.sif` ships **without** `cmake` and **without** libcurl headers (only the `.so.4` runtime). Both omnistat artifacts need cmake; the kernel-trace library additionally needs ``. Working build recipe inside the SIF: @@ -737,12 +737,12 @@ When the cluster is up but you only need a one-shot interactive build inside the srun -p -A -N 1 --time=00:15:00 --gpus-per-node=1 --cpus-per-task=8 \ --job-name= \ apptainer exec --rocm \ - --bind ${PERF_TOOLS_DIR}:${PERF_TOOLS_DIR} \ + --bind ${OMNIHUB_TOOLS_DIR}:${OMNIHUB_TOOLS_DIR} \ "$HG_SIF" \ bash -c '' ``` -Why: `salloc` doesn't auto-bind `${PERF_TOOLS_DIR}/` into the apptainer namespace. Direct `srun … apptainer exec --bind …` forces the bind explicitly and writes the build artifact to its persistent NFS path in one shot. Verified during the kernel-trace library build (job 7030). +Why: `salloc` doesn't auto-bind `${OMNIHUB_TOOLS_DIR}/` into the apptainer namespace. Direct `srun … apptainer exec --bind …` forces the bind explicitly and writes the build artifact to its persistent NFS path in one shot. Verified during the kernel-trace library build (job 7030). ### 15. Dual-failure debug pattern — separate "tool wired correctly?" from "workload happy?" diff --git a/earth_science/models/ORBIT-2/examples/interm_8m_lux.yaml b/earth_science/models/ORBIT-2/examples/interm_8m_lux.yaml new file mode 100644 index 0000000..8c03b5b --- /dev/null +++ b/earth_science/models/ORBIT-2/examples/interm_8m_lux.yaml @@ -0,0 +1,129 @@ +# ORBIT-2 Lux/MI355X training config — real 10.0_arcmin PRISM grid (same-dir mode). +# +# Both low_res_dir and high_res_dir point at the same 10.0_arcmin tree so we can +# exercise the full training pipeline for perf/scaling without 2.5_arcmin targets. +# Loss and PSNR are NOT scientifically meaningful in this mode — use for timing only. +# +# Placeholders substituted by render_orbit2_config.py at job submit time: +# __DATA_ROOT__ — ORBIT2_DATA_ROOT (10.0_arcmin directory) +# __FSDP__ — fsdp parallelism (typically = node count) +# __SIMPLE_DDP__ — simple_ddp parallelism (typically 8) +# __MAX_EPOCHS__ — ORBIT2_MAX_EPOCH override +# __BATCH_SIZE__ — ORBIT2_BATCH_SIZE override +# __DATA_TYPE__ — ORBIT2_DATA_TYPE (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/interm_8m_lux_era5.yaml b/earth_science/models/ORBIT-2/examples/interm_8m_lux_era5.yaml new file mode 100644 index 0000000..6e230db --- /dev/null +++ b/earth_science/models/ORBIT-2/examples/interm_8m_lux_era5.yaml @@ -0,0 +1,145 @@ +# ORBIT-2 Lux/MI355X — ERA5 1.0_deg sanity / pipeline test (same-dir mode). +# +# Both low_res_dir and high_res_dir point at the same ERA5 tree so we can verify +# the new dataset loads and training runs before pairing with PRISM 2.5_arcmin targets. +# Loss is not scientifically meaningful in same-dir mode — use for data/pipeline sanity only. +# +# Placeholders substituted by render_orbit2_config.py: +# __DATA_ROOT__ — ORBIT2_DATA_ROOT (era5/1.0_deg directory) +# __FSDP__ — fsdp parallelism (typically = node count) +# __SIMPLE_DDP__ — simple_ddp parallelism (typically 8) +# __MAX_EPOCHS__ — ORBIT2_MAX_EPOCH override +# __BATCH_SIZE__ — ORBIT2_BATCH_SIZE override +# __DATA_TYPE__ — ORBIT2_DATA_TYPE (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/run_2node_scaleout_loop.sh b/earth_science/models/ORBIT-2/examples/run_2node_scaleout_loop.sh index 0f68adc..d7435c8 100755 --- a/earth_science/models/ORBIT-2/examples/run_2node_scaleout_loop.sh +++ b/earth_science/models/ORBIT-2/examples/run_2node_scaleout_loop.sh @@ -26,7 +26,7 @@ SBATCH="${SCRIPT_DIR}/sbatch_train_perf_amd.sh" # ---- environment (override-able) ------------------------------------------- export AI4S_SHARED_DIR="${AI4S_SHARED_DIR:?AI4S_SHARED_DIR must be set}" -export PERF_TOOLS_DIR="${PERF_TOOLS_DIR:?PERF_TOOLS_DIR must be set (perf_tools.dir in .cluster-config.yaml)}" +export OMNIHUB_TOOLS_DIR="${OMNIHUB_TOOLS_DIR:?OMNIHUB_TOOLS_DIR must be set (omnihub.tools_dir in .cluster-config.yaml)}" export ORBIT2_ROOT="${ORBIT2_ROOT:-${AI4S_SHARED_DIR}/models/ORBIT-2/code/bayes-cast}" export ORBIT2_DATA_ROOT="${ORBIT2_DATA_ROOT:-${AI4S_SHARED_DIR}/models/ORBIT-2/data/superres/era5/1.0_deg}" export ORBIT2_CONFIG_TEMPLATE="${ORBIT2_CONFIG_TEMPLATE:-edm_8m_era5_1x8.yaml}" @@ -137,7 +137,7 @@ for i in "${!LEVER_NAME[@]}"; do EXP="ALL" EXP+=",ORBIT2_BATCH_SIZE=${BATCH},ORBIT2_NUM_WORKERS=1,ORBIT2_DATA_TYPE=bfloat16,ORBIT2_FUSED_ATTN=DEFAULT" EXP+=",ORBIT2_CONFIG_TEMPLATE=${ORBIT2_CONFIG_TEMPLATE},ORBIT2_MAX_EPOCH=${MAXEPOCH},ORBIT2_MAX_BATCHES=${MAXBATCH}" - EXP+=",ORBIT2_DATA_ROOT=${ORBIT2_DATA_ROOT},ORBIT2_ROOT=${ORBIT2_ROOT},AI4S_SHARED_DIR=${AI4S_SHARED_DIR},PERF_TOOLS_DIR=${PERF_TOOLS_DIR}" + EXP+=",ORBIT2_DATA_ROOT=${ORBIT2_DATA_ROOT},ORBIT2_ROOT=${ORBIT2_ROOT},AI4S_SHARED_DIR=${AI4S_SHARED_DIR},OMNIHUB_TOOLS_DIR=${OMNIHUB_TOOLS_DIR}" EXP+=",TORCH_NCCL_HIGH_PRIORITY=1,GPU_MAX_HW_QUEUES=2" for kv in $extra; do EXP+=",${kv}"; done @@ -242,4 +242,4 @@ print(f"wrote {report}") PY log "=== ${LOOP_ID} COMPLETE — see $REPORT ===" -cat "$REPORT" 2>/dev/null | tee -a "$STATUS" +cat "$REPORT" 2>/dev/null | tee -a "$STATUS" \ No newline at end of file diff --git a/earth_science/models/ORBIT-2/examples/run_gemm_analysis.sh b/earth_science/models/ORBIT-2/examples/run_gemm_analysis.sh index 5c96b9b..fc9af5d 100755 --- a/earth_science/models/ORBIT-2/examples/run_gemm_analysis.sh +++ b/earth_science/models/ORBIT-2/examples/run_gemm_analysis.sh @@ -7,7 +7,7 @@ # # Usage (tmux on login node): # export ANTHROPIC_API_KEY=... export AI4S_SHARED_DIR= -# export PERF_TOOLS_DIR= # perf_tools.dir in .cluster-config.yaml +# export OMNIHUB_TOOLS_DIR= # omnihub.tools_dir in .cluster-config.yaml # export EXCLUDE_NODES=... # optional: comma-separated nodes to skip # bash earth_science/models/ORBIT-2/examples/run_gemm_analysis.sh [--preflight-only] # @@ -23,8 +23,8 @@ REPO_ROOT=$(cd "${SCRIPT_DIR}/../../../.." && pwd) ORCH_PROMPT="${REPO_ROOT}/earth_science/models/ORBIT-2/recipes/perf-analysis/agents/orchestrator_gemm_analysis.md" : "${AI4S_SHARED_DIR:?AI4S_SHARED_DIR must be set}" -: "${PERF_TOOLS_DIR:?PERF_TOOLS_DIR must be set}" -export AI4S_SHARED_DIR PERF_TOOLS_DIR +: "${OMNIHUB_TOOLS_DIR:?OMNIHUB_TOOLS_DIR must be set}" +export AI4S_SHARED_DIR OMNIHUB_TOOLS_DIR ORBIT2_BASE="${AI4S_SHARED_DIR}/models/ORBIT-2" PERF_RUNS_DIR="${ORBIT2_BASE}/perf-runs" @@ -50,11 +50,11 @@ if command -v sinfo >/dev/null 2>&1; then else _notes+=("sinfo not in PATH"); fi _pct=$(df -P "$PERF_RUNS_DIR" 2>/dev/null | awk 'NR==2 {sub("%","",$5); print $5}') [[ -z "$_fail" && "${_pct:-0}" -ge 95 ]] && { _fail="disk_full"; _notes+=("perf-runs at ${_pct}%"); } -VENV="${PERF_TOOLS_DIR}/perf-inspect" +VENV="${OMNIHUB_TOOLS_DIR}/omnihub-inspect" for _f in "$ORCH_PROMPT" \ "${SCRIPT_DIR}/sbatch_train_perf_amd.sh" \ "${VENV}/bin/omnistat-usermode" \ - "${PERF_TOOLS_DIR}/victoriametrics/victoria-metrics-prod" \ + "${OMNIHUB_TOOLS_DIR}/victoriametrics/victoria-metrics-prod" \ "${AI4S_SHARED_DIR}/images/pytorch_rocm7.2.2_ubuntu24.04_py3.12_pytorch_release_2.10.0.sif"; do [[ -z "$_fail" && ! -e "$_f" ]] && { _fail="missing_file"; _notes+=("missing: $_f"); } done @@ -92,7 +92,7 @@ trap '_log "ANALYSIS_SIGNAL cleanup"; kill $_GUARD_PID 2>/dev/null || true' EXIT USER_PROMPT="Run the ORBIT-2 GEMM-time bottleneck analysis (TraceLens + Omnistat analyst/verifier) at 1 and 2 nodes. Follow ${ORCH_PROMPT}. REPO_ROOT=${REPO_ROOT} AI4S_SHARED_DIR=${AI4S_SHARED_DIR} -PERF_TOOLS_DIR=${PERF_TOOLS_DIR} +OMNIHUB_TOOLS_DIR=${OMNIHUB_TOOLS_DIR} ANALYSIS_DIR=${ANALYSIS_DIR} EXCLUDE_NODES=${EXCLUDE_NODES:-} Write GEMM_TIME_REPORT.md to ANALYSIS_DIR. Persist state for resume." @@ -118,4 +118,4 @@ if [[ -f "$REPORT" ]]; then echo "=== GEMM-time analysis complete ==="; echo " report: ${REPORT}" else _log "ANALYSIS_INCOMPLETE no report after ${_MAX} attempts"; exit 1 -fi +fi \ No newline at end of file diff --git a/earth_science/models/ORBIT-2/examples/run_optimizer_loop.sh b/earth_science/models/ORBIT-2/examples/run_optimizer_loop.sh index 7c59dda..cdad90e 100755 --- a/earth_science/models/ORBIT-2/examples/run_optimizer_loop.sh +++ b/earth_science/models/ORBIT-2/examples/run_optimizer_loop.sh @@ -4,7 +4,7 @@ # Usage (tmux on login node): # export ANTHROPIC_API_KEY=... # optional; only for Claude Code CLI driver # export AI4S_SHARED_DIR=... -# export PERF_TOOLS_DIR=... +# export OMNIHUB_TOOLS_DIR=... # bash earth_science/models/ORBIT-2/examples/run_optimizer_loop.sh [--preflight-only] # # Graceful stop: touch $AI4S_SHARED_DIR/models/ORBIT-2/perf-runs/loop-/STOP @@ -33,7 +33,7 @@ RECIPE_DIR="${REPO_ROOT}/earth_science/models/ORBIT-2/recipes/perf-optimizer-loo ORCH_PROMPT="${ORBIT2_ORCH_PROMPT:-${RECIPE_DIR}/agents/orchestrator.md}" : "${AI4S_SHARED_DIR:?AI4S_SHARED_DIR must be set}" -: "${PERF_TOOLS_DIR:?PERF_TOOLS_DIR must be set}" +: "${OMNIHUB_TOOLS_DIR:?OMNIHUB_TOOLS_DIR must be set}" ORBIT2_BASE="${AI4S_SHARED_DIR}/models/ORBIT-2" PERF_RUNS_DIR="${ORBIT2_BASE}/perf-runs" @@ -77,8 +77,8 @@ if [[ -z "$PREFLIGHT_FAIL_REASON" ]]; then fi for _bin in \ - "${PERF_TOOLS_DIR}/perf-inspect/bin/omnistat-usermode" \ - "${PERF_TOOLS_DIR}/victoriametrics/victoria-metrics-prod"; do + "${OMNIHUB_TOOLS_DIR}/omnihub-inspect/bin/omnistat-usermode" \ + "${OMNIHUB_TOOLS_DIR}/victoriametrics/victoria-metrics-prod"; do if [[ -z "$PREFLIGHT_FAIL_REASON" && ! -x "$_bin" ]]; then PREFLIGHT_FAIL_REASON="tool_missing" PREFLIGHT_NOTES+=("missing: $_bin") @@ -179,4 +179,4 @@ done _log_status "DRIVER_EXIT clean" echo "=== ORBIT-2 loop driver done ===" -echo " $LOOP_DIR" +echo " $LOOP_DIR" \ No newline at end of file diff --git a/earth_science/models/ORBIT-2/examples/sbatch_train_perf_amd.sh b/earth_science/models/ORBIT-2/examples/sbatch_train_perf_amd.sh index 751fef9..46de0c9 100755 --- a/earth_science/models/ORBIT-2/examples/sbatch_train_perf_amd.sh +++ b/earth_science/models/ORBIT-2/examples/sbatch_train_perf_amd.sh @@ -6,7 +6,7 @@ # # Quick start (ERA5 same-dir): # export AI4S_SHARED_DIR=/path/to/shared -# export PERF_TOOLS_DIR=/path/to/perf-tools +# export OMNIHUB_TOOLS_DIR=/shared/omnihub/tools # export ORBIT2_DATA_ROOT=$AI4S_SHARED_DIR/models/ORBIT-2/data/superres/era5/1.0_deg # export ORBIT2_CONFIG_TEMPLATE=edm_8m_era5_1x8.yaml # export ORBIT2_ROOT=$AI4S_SHARED_DIR/models/ORBIT-2/code/bayes-cast # EDM preset + launcher @@ -72,8 +72,8 @@ ORBIT2_DISABLE_CKPT="${ORBIT2_DISABLE_CKPT:-1}" ORBIT2_TUNABLEOP_MODE="${ORBIT2_TUNABLEOP_MODE:-off}" ORBIT2_TUNABLEOP_DIR="${ORBIT2_TUNABLEOP_DIR:-${ORBIT2_BASE}/tunableop-cache}" -: "${PERF_TOOLS_DIR:?PERF_TOOLS_DIR must be set}" -OMNISTAT_VENV="${OMNISTAT_VENV:-${PERF_TOOLS_DIR}/perf-inspect}" +: "${OMNIHUB_TOOLS_DIR:?OMNIHUB_TOOLS_DIR must be set}" +OMNISTAT_VENV="${OMNISTAT_VENV:-${OMNIHUB_TOOLS_DIR}/omnihub-inspect}" OMNISTAT_TEMPLATE="${OMNISTAT_TEMPLATE:-${SCRIPT_DIR}/../recipes/perf-analysis/omnistat.config.template}" OMNISTAT_USERMODE_INTERVAL="${OMNISTAT_USERMODE_INTERVAL:-1}" @@ -141,7 +141,7 @@ fi mkdir -p "$ORBIT2_OUTPUT_DIR" OMNISTAT_CONFIG="${ORBIT2_OUTPUT_DIR}/omnistat.config" sed -e "s|@JOB_DIR@|${ORBIT2_OUTPUT_DIR}|g" \ - -e "s|@PERF_TOOLS_DIR@|${PERF_TOOLS_DIR}|g" \ + -e "s|@OMNIHUB_TOOLS_DIR@|${OMNIHUB_TOOLS_DIR}|g" \ "$OMNISTAT_TEMPLATE" > "$OMNISTAT_CONFIG" if [[ -n "${OMNISTAT_ROCPROF_PROFILE:-}" ]]; then sed -i.bak "s/^profile = .*/profile = ${OMNISTAT_ROCPROF_PROFILE}/" "$OMNISTAT_CONFIG" && rm -f "${OMNISTAT_CONFIG}.bak" diff --git a/earth_science/models/ORBIT-2/examples/submit_perf_baseline_era5_amd.sh b/earth_science/models/ORBIT-2/examples/submit_perf_baseline_era5_amd.sh index 53dfb87..2b862bb 100755 --- a/earth_science/models/ORBIT-2/examples/submit_perf_baseline_era5_amd.sh +++ b/earth_science/models/ORBIT-2/examples/submit_perf_baseline_era5_amd.sh @@ -4,14 +4,14 @@ # # Usage (from anywhere): # export AI4S_SHARED_DIR=/path/to/shared -# export PERF_TOOLS_DIR=/path/to/perf-tools +# export OMNIHUB_TOOLS_DIR=/shared/omnihub/tools # ./submit_perf_baseline_era5_amd.sh # Extra sbatch flags: ./submit_perf_baseline_era5_amd.sh --partition=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}" +: "${OMNIHUB_TOOLS_DIR:?set OMNIHUB_TOOLS_DIR}" export ORBIT2_DATA_ROOT="${ORBIT2_DATA_ROOT:-${AI4S_SHARED_DIR}/models/ORBIT-2/data/superres/era5/1.0_deg}" export ORBIT2_CONFIG_TEMPLATE="${ORBIT2_CONFIG_TEMPLATE:-edm_8m_era5_1x8.yaml}" diff --git a/earth_science/models/ORBIT-2/examples/sweep_orbit2_batch_bf16_amd.sh b/earth_science/models/ORBIT-2/examples/sweep_orbit2_batch_bf16_amd.sh index 7c30626..3d2910b 100755 --- a/earth_science/models/ORBIT-2/examples/sweep_orbit2_batch_bf16_amd.sh +++ b/earth_science/models/ORBIT-2/examples/sweep_orbit2_batch_bf16_amd.sh @@ -3,7 +3,7 @@ # # Usage (login node, repo root): # export AI4S_SHARED_DIR=... -# export PERF_TOOLS_DIR=... +# export OMNIHUB_TOOLS_DIR=... # export ORBIT2_ROOT=$AI4S_SHARED_DIR/models/ORBIT-2/code/bayes-cast # export ORBIT2_DATA_ROOT=$AI4S_SHARED_DIR/models/ORBIT-2/data/superres/era5/1.0_deg # # Optional: export ORBIT2_ERA5_SPATIAL_RES=111 @@ -27,7 +27,7 @@ if [[ $# -lt 1 ]]; then fi : "${AI4S_SHARED_DIR:?set AI4S_SHARED_DIR}" -: "${PERF_TOOLS_DIR:?set PERF_TOOLS_DIR}" +: "${OMNIHUB_TOOLS_DIR:?set OMNIHUB_TOOLS_DIR}" export ORBIT2_DATA_TYPE="${ORBIT2_DATA_TYPE:-bfloat16}" export ORBIT2_FUSED_ATTN="${ORBIT2_FUSED_ATTN:-DEFAULT}" @@ -41,4 +41,4 @@ for B in "$@"; do ( 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" +echo "Done. Track jobs: squeue -u \$USER" \ No newline at end of file diff --git a/earth_science/models/ORBIT-2/recipes/perf-analysis/HANDOFF.md b/earth_science/models/ORBIT-2/recipes/perf-analysis/HANDOFF.md new file mode 100644 index 0000000..65a3682 --- /dev/null +++ b/earth_science/models/ORBIT-2/recipes/perf-analysis/HANDOFF.md @@ -0,0 +1,308 @@ +# ORBIT-2 perf-analysis HANDOFF + +_Last updated: 2026-06-11 (perf-optimizer-loop + bf16 omnistat profile + manifest fields)._ + +## What works + +| Item | Status | Notes | +|------|--------|-------| +| 1-node × 8-GPU training on **ERA5 1.0°** same-dir (default perf template) | **Green** | Heavier **111×111** latent vs PRISM **18×18**; better default for GPU saturation baselines. Representative steady batch time workload-dependent. | +| 1-node × 8-GPU on **PRISM** 10.0_arcmin same-dir | **Green** | Lighter forwards; override `ORBIT2_DATA_ROOT` + `interm_8m_lux.yaml` if ERA5 not staged. | +| `sbatch_train_amd.sh` / `sbatch_train_perf_amd.sh` | **Green** | Prefer `launch_diffusion.sh` under `ORBIT2_ROOT` when present; else `run_orbit2_train.py`. `orbit2_rank_hook_runner.py` runs profiler hook before shell launcher. | +| **1-node default parallelism** | **fsdp=8, simple_ddp=1** | Override with `ORBIT2_FSDP` / `ORBIT2_SIMPLE_DDP`. Multi-node: omit for `fsdp=nodes`, `simple_ddp=8`. | +| `ORBIT2_CONFIG_TEMPLATE` + `render_orbit2_config.py` | **Green** | Default perf: **`edm_8m_era5_1x8.yaml`** (Bayes-CAST `edm_8m_era5` + **`fsdp=8`/`simple_ddp=1`/`seq_par:1`**, templated `data_dir` + **`ORBIT2_ERA5_SPATIAL_RES`** default 111). Lux **`interm_8m_lux_era5.yaml`** for public ORBIT-2 `res_slimvit`. `__DATA_TYPE__` defaults **bfloat16**. | +| Node health probe (exit 42) | **Green** | Catches broken home/shared/SIF/data mounts before training | +| `superres_mag: 1` same-dir config | **Required** | Without it, 4× superres vs same-resolution targets causes tensor shape mismatch | +| Baseline lock-in + VRAM-first sweep | **Doc** | [BASELINE_LOCKIN.md](BASELINE_LOCKIN.md) — freeze topology/template; sweep `ORBIT2_BATCH_SIZE` before widening ViT | + +## Landmines + +1. **`max_epochs` must be ≥ 2** — upstream loop is `while (epoch_start + 1) < max_epochs`; `max_epochs=1` runs zero epochs. +2. **`xformers.ops` missing in overlay** — `run_orbit2_train.py` can fall back to `FusedAttn.DEFAULT` (PyTorch SDPA). Rebuild overlay per `build_overlay_amd.sh` if you need the CK attention path. +3. **`bfloat16` + xformers CK path** — needs working `xformers.ops`. Studio sbatch sets **`ORBIT2_FUSED_ATTN=DEFAULT`** so PyTorch SDPA is used when CK is unavailable. If you still see dtype-related crashes, fall back to **`ORBIT2_DATA_TYPE=float32`** for isolation. +4. **Same-dir data** — both `low_res_dir` and `high_res_dir` point at the same grid; set `superres_mag: 1` in `interm_8m_lux.yaml` / `interm_8m_lux_era5.yaml`. +5. **Batch cap** — must patch `intermediate_downscaling` before `main()` (not via `runpy.run_path` reload). +6. **Bayes-CAST `launch/launch_diffusion.sh`** — upstream file is often an OLCF/Crusher **job wrapper** (nested `srun`, ignores argv). Studio **`sbatch_train_*`** does **not** exec it when **`$ORBIT2_ROOT/launch/train_edm.py`** exists; ranks run **`python3 train_edm.py /config/config.yaml`** from **`/orbit2/launch`** (no `examples/` required). For Lux/public ORBIT-2, **`launch_diffusion.sh`** + argv still applies when **`train_edm.py`** is absent. +7. **`train_edm.py` + `ERA5_1`** — some Bayes-CAST snapshots only define **`std_delta`** for IMERG/HRRR; **`data_key=="ERA5_1"`** then hits **`UnboundLocalError`** at the first step. Fix: add an **`elif data_key == "ERA5_1"`** branch (per-channel std aligned with your staged **`normalize_std.npz`** and **`dict_out_variables` order**), or use the small patch applied on your shared clone (example values `[20.601086, 5.542934, 4.7573]` for T2m / u10 / v10 @ 1.0°). + +## FOM contract (steady-state + loss sanity) + +**Primary FOM:** `steady_batch_time_s` — mean per-batch wall time from **epoch ≥ 2**, excluding **batch index < 1** within each epoch (skips compile/warmup batch 0). Epochs 0–1 are warmup. + +**Loss sanity:** epoch-completed losses must **strictly decrease** epoch-over-epoch (`loss_sanity_pass`). Same-dir data makes absolute loss non-physical; use as crash/hang detector only. + +**Scaling runs:** use `ORBIT2_MAX_EPOCH=6` (trains epochs 0–4) so the steady window spans epochs 2–4. Parse with: + +```bash +python3 parse_training_log.py --log ... --steady-epoch-start 2 --warmup-batches-per-epoch 1 +python3 collate_scaling_study.py --log-dir .../logs --jobs ... --steady-epoch-start 2 --require-loss-sanity +``` + +## GEMM-time bottleneck + ORNL MIOpen flags (2026-06-15, jobs 12562/12648) + +Full TraceLens+Omnistat analyst/verifier run (`agents/orchestrator_gemm_analysis.md`, driver +`examples/run_gemm_analysis.sh`) on 1- and 2-node profiled EDM (bf16, batch 4096, ERA5 1.0°). Report: +`perf-runs/gemm-analysis-/GEMM_TIME_REPORT.md`. + +- **The hot kernel is NOT attention/MLP.** A single hipBLASLt GEMM `Cijk_…MT16x16x32_MI16` is + **~43 %** of the step at *both* scales (scale-invariant). It is the **im2col-lowered low-channel + conv backward** for the **3- and 4-channel stem/head projection convs**, stuck on a 16×16 macro-tile + at ~0.03 TFLOP/s. The tunable `nn.Linear` bf16 GEMMs (256×256 tiles, 600–1074 TFLOP/s) are only + ~2 %. **This explains why TunableOp gave no uplift** — it tuned already-optimal GEMMs and can't + touch the conv-lowered one. `im2col` (`Im2d2Col_v2`) adds ~14 %; conv **forward** is a `naive_conv` + with **fp64 accumulate** (`..._ushort_double_ushort`) at ~3 %. +- **TraceLens `report.xlsx:ops_summary_by_category` MISLABELS the conv-GEMM as `CONV_bwd`.** The + 2-node analyst trusted it and claimed "naive conv 65.8 %"; the verifier refuted it from raw + `traceEvents` (naive_conv_bwd = 0.46 %, GEMM = 50.7 %). **Always rank by raw kernel-name device + time, not the category rollup,** for conv-heavy models. (`examples/compare_trace_kernels.py` does + this tool-independently.) +- **~43–44 % whole-job idle at both scales** (Omnistat util mean ~56 %) is the second lever — it is + structural inter-step overhead, **not** comm-bound (exposed comm 0.13 %→0.47 %; the Omnistat + "comm-bound" claim was refuted by the trace). +- **EDM trainer now has an env-gated profiler.** `train_edm.py:_orbit2_make_profiler` writes a rank-0 + `*.pt.trace.json` for `PROFILE_TARGET_EPOCH` into `ORBIT2_PROFILE_DIR` (no-op unless set). This is + what makes "where does GEMM time go" answerable for the EDM path (the old Lux `res_slimvit` hook + did not apply). + +### ORNL Frontier MIOpen flags ported to the sbatch (validate on Lux) + +From bayes-cast `launch/launch_diffusion.sh` (ORNL, gfx90a/ROCm7.1.1, "tested many times"): ORNL +**disables Winograd** and unbounds the multi-pass Winograd workspace. Now defaults in +`sbatch_train_perf_amd.sh` / `sbatch_train_amd.sh` / `sbatch_infer_*.sh` (override to A/B): + +``` +MIOPEN_DEBUG_CONV_WINOGRAD=0 # ORBIT2_MIOPEN_CONV_WINOGRAD (0=off ORNL default, 1=on) +MIOPEN_DEBUG_AMD_WINOGRAD_MPASS_WORKSPACE_MAX=-1 # ORBIT2_MIOPEN_WINOGRAD_MPASS_WS_MAX +MIOPEN_DEBUG_AMD_MP_BD_WINOGRAD_WORKSPACE_MAX=-1 # ORBIT2_MIOPEN_MP_BD_WINOGRAD_WS_MAX +HSA_FORCE_FINE_GRAIN_PCIE=1 +``` + +`MIOPEN_DISABLE_CACHE=1` + `MIOPEN_USER_DB_PATH=/tmp/` already matched ORNL — keep them (do NOT +"persist the find-db"; that deviates from the validated setup). **Deliberately NOT ported** (Frontier +Slingshot/Cray-specific, would break Lux IB/ionic): `FI_CXI_*`, `NCCL_NET="AWS Libfabric"`, +`NCCL_NET_PLUGIN=librccl-net.so`, `NCCL_SOCKET_IFNAME=hsn0`, Frontier `LD_PRELOAD` host paths. +A/B (ORNL Winograd-off vs Winograd-on vs pre-flags +baseline) → `perf-runs/miopen-ornl-validation-/MIOPEN_ORNL_VALIDATION.md`. + +**MIOpen ORNL-flags A/B result (2026-06-16, jobs 12997/12998 vs 12562): performance-NEUTRAL on Lux** +(steady 8.62→8.66→8.74 s, within ~1-2 % noise; loss sanity OK). MIOpen never selected Winograd for +the 3/4-channel convs (0 % Winograd kernels in all arms), so toggling it changes nothing — our +bottleneck is shape-driven, not algorithm-selection-driven. Keeping the flags as defaults (faithful, +harmless, may matter elsewhere). They are **not** the lever here. + +### Conv channel-padding — TABLED candidate, timing-only (NOT ADOPTED, 2026-06-16, jobs 14633/14634) + +> **Status: tabled.** The −25 % is a *timing-only* result. Channel-padding **changes model capacity** +> (wider internal convs ⇒ different params), and **we have not run a convergence/quality study**, so it +> is **not adopted** and `ORBIT2_CONV_PAD` stays **default 0 (off)**. Treat the numbers below as a +> *potential* speedup pending a same-quality convergence check (matched final loss/skill at equal or +> fewer steps). Do not enable for real training runs until that check passes. + +`edm.py` `path2`/`refine` convs run at out_channels(=3)/cnn_ratio(=4) at full resolution → the +starved `Cijk_…MT16x16x32` GEMM @ ~0.03 TFLOP/s. New env knob **`ORBIT2_CONV_PAD=N`** (in `edm.py`, +default 0 = original) widens those convs' INTERNAL width to N (tile-friendly) keeping the 3-channel +I/O. Single-node A/B (identical bf16/4096/ERA5 config), N=16: + +| Arm | steady_batch_time_s | total GPU-kernel ms | dominant `MT16x16x32` GEMM | +|---|---|---|---| +| control (pad=0) | 8.646 | 21082 | 9245 ms (43.9 %) | +| **widened (pad=16)** | **6.498 (−24.8 %)** | 15150 (−28 %) | **gone** (top GEMM now `MT16x16x128`, 4123 ms) | + +**~25 % faster step, loss *sanity* holds (not convergence).** Confirms the GEMM report: low channel +count forced the bad tile; widening makes it tile-friendly and the efficiency gain beats the extra +FLOPs. **Why tabled:** "loss sanity" only checks the loss is finite/decreasing over a few epochs — it +is **not** evidence the wider model reaches the same quality. Before adopting, run a matched +convergence study (same data/schedule, compare final validation loss/skill and steps-to-target for +pad=0 vs pad=16); only adopt if quality is equal-or-better. `channels_last`/implicit-GEMM as a way to +kill the residual `im2col` (~21 %) is **already disproven** (see next section — naive_conv fallback, +7-10× slower). If/when revisited, the open timing question is **N=32** (top GEMM `MT16x16x128` at 27 % +still has headroom) — but only worth it after convergence is settled. +Report: `perf-runs/conv-pad-validation-/CONV_PAD_VALIDATION.md`. + +### `channels_last` + implicit-GEMM — DEAD END on ROCm 7.2.2 (DISPROVEN, 2026-06-16, jobs 14645-14648) + +Tested the "stack `channels_last` on top" idea above as a clean 2×2 ({NCHW, channels_last} × {pad=0, +pad=16}) via the env knob **`ORBIT2_CHANNELS_LAST`** in `edm.py`, +default 0; sets `MIOPEN_FIND_MODE=1` on the NHWC arms). **It is 7-10× SLOWER, not faster:** + +| Arm | CONV_PAD | channels_last | steady_batch_time_s | vs control | +|---|---|---|---|---| +| nchw_pad0 (control) | 0 | off | 8.62 | — | +| **nchw_pad16** | 16 | off | **6.38** | **−26 %** ✅ | +| cl_pad0 | 0 | on | 60.8 | +606 % 💥 | +| cl_pad16 | 16 | on | 80.1 | +829 % 💥 | + +**Root cause (kernel trace):** MIOpen on this build has **no tuned NHWC implicit-GEMM solver** for +these conv shapes. Forcing channels_last makes it fall back to the brute-force **`naive_conv_…wrw_nhwc`** +direct-convolution kernel, which dominates **~88-91 %** of GPU-kernel time (cl_pad0: 139,522 ms; +cl_pad16: 185,198 ms). The `im2col` (`Im2d2Col_v2`, 14-21 % in NCHW) is **not** wasted overhead — it is +the price of reaching the tuned **hipBLASLt GEMM** kernels (`Cijk_…`, 31-54 % in NCHW). NHWC discards +that path. **Lesson: do NOT use `channels_last` for these low-channel convs on MI355X/ROCm 7.2.2.** +**The lever is NCHW + `ORBIT2_CONV_PAD` (−26 %); channels_last is abandoned.** Knob kept (default off) +only so the dead end is reproducible. Report: `perf-runs/conv-layout-validation-/CONV_LAYOUT_VALIDATION.md`. + +## Scaling / wall-clock lessons + +- Long runs on full 180×360 grids need **enough wall time** (multi-epoch × many batches can exceed 1 h on some sites). +- Point SLURM `--output` / `--error` at a **shared project filesystem** your compute nodes can always read during outages (not only the submission directory on a login node). +- **Multi-node hangs** are often **infrastructure** (drained nodes, stuck scheduler states, shared filesystem hiccups) — confirm node state with `sinfo` / site ops before assuming a code bug. +- For exclusive GPU jobs, **avoid stacking two jobs on the same node**. Pin nodes with `--nodelist=...` only when your site policy allows it and you have confirmed the nodes are healthy. + +## Debug session handoff — Bayes-CAST EDM, `bfloat16`, xFormers CK (SIGSEGV) + +**Why this session felt confusing** + +- **Two trees:** Instrumentation and hotfixes for Bayes-CAST live on the **compute-side clone** (e.g. **`$AI4S_SHARED_DIR/models/ORBIT-2/code/bayes-cast/`**), not inside this git repo. AI4Science Studio only documents launchers and perf recipes; agents may edit `/shared/...` while you expect everything under `ai4science-studio/`. +- **Debug log path vs cluster:** Agent NDJSON defaults to **`/.cursor/debug-.log`**. Compute nodes inside Apptainer often **cannot** write that path. Unless you set **`DEBUG_AGENT_LOG`** to something bind-mounted and writable, then **copy** the file back to the workspace path, the next chat turn will show **“log not found”** even after a good repro — that is an artifact workflow gap, not necessarily “no crash.” +- **Session id:** the agent debug `sessionId` is reused as the log filename `debug-.log`. + +**What was already instrumented (do not rip out until CK is understood)** + +| Location | Purpose | +|----------|---------| +| `.../bayes-cast/launch/train_edm.py` after `print("model_kwargs", ...)` | **H5:** one NDJSON line — `torch` / `xformers` versions, HIP/CUDA strings, `FusedAttn_option`, `data_type`, `gpu_type`, `ORBIT2_FUSED_ATTN` (rank 0 only). | +| `.../bayes-cast/src/climate_learn/models/hub/components/attention.py` — `_agent_dbg` | **H1/H3:** `Attention` and `CrossAttention` **pre_ck** (shapes, dtypes, `dropout_p`, `N_a`/`N_i` for cross) immediately before `memory_efficient_attention(..., op=MemoryEfficientAttentionCkOp)`; **post_ck** after a successful return. **Rank 0 only** when `dist.is_initialized()` to avoid interleaved NDJSON from 8 ranks. | + +**Hypotheses to resolve from the next log (first evidence wins)** + +1. **H1** — Self-attention CK bad layout/shapes: last line `Attention.pre_ck`, no `Attention.post_ck`. +2. **H2** — `dropout_p` / training vs CK: correlate with fields on last `pre_ck`. +3. **H3** — Cross-attention CK when `N_a ≠ N_i`: last line `CrossAttention.pre_ck`, no `CrossAttention.post_ck`. +4. **H4** — Dtype: `CrossAttention` does `q, k = q.to(x.dtype), k.to(x.dtype)` but **not** `v`; if log shows `v_dtype != q_dtype` before SIGSEGV, add **`v = v.to(x.dtype)`** before CK (then post-fix verification run **with logs still on**). +5. **H5** — Stack skew: use `train_edm.py:post_model_kwargs` line for versions. + +**Operational checklist for the next session** + +1. In sbatch or env: `export DEBUG_AGENT_LOG=/path/on/shared/fs/agent-ck.ndjson` (must be visible inside the container). +2. Run minimal repro: `ORBIT2_DATA_TYPE=bfloat16`, `ORBIT2_FUSED_ATTN=CK`, small batch. +3. After job ends: copy `agent-ck.ndjson` → workspace **`/.cursor/debug-.log`** (or paste last ~50 lines into chat). +4. Until logs prove otherwise, **`ORBIT2_FUSED_ATTN=DEFAULT`** (SDPA) remains the safe default in Studio sbatch for AMD+bf16. + +**Symptom reminder:** CK + bf16 jobs ended with **`Segmentation fault (core dumped)`** on all ranks (no Python traceback). Example slurm text may live under `.../outputs/train/batch-cal-edm/*bf16*.out`. + +### NEW (2026-06-11) — bf16 blocked even with SDPA (`FusedAttn.DEFAULT`): `var_agg` proj GEMM `hipErrorInvalidValue` + +**This is broader than the CK SIGSEGV above:** with the *safe* default **`ORBIT2_FUSED_ATTN=DEFAULT`** (PyTorch SDPA), **bf16 EDM training still crashes** on the **first `training_step`**: + +``` +File "/orbit2/src/climate_learn/models/hub/edm.py", line 316, in aggregate_variables + x = self.var_agg(var_query, x) +File "/orbit2/src/climate_learn/models/hub/components/attention.py", line 280, in forward + x = self.proj(x) # nn.Linear(256, 256) +torch.AcceleratorError: HIP error: invalid argument (hipErrorInvalidValue) +``` + +- **Reproduced across 6 jobs (lux MI355X, `pytorch_rocm7.2.2` SIF + ORBIT-2 overlay):** + - **9228/9229/9230** — bf16 batch 1024/2048/4096 → crash. (Also confirmed restaged ERA5 lifted the old ~1704 sample cap: batch 1024 → `y.shape[0]=1024`; batch ≥2048 → caps at **1704** = total train samples.) + - **9231** — bf16 batch **256** → **same crash** (not large-M dependent). + - **9232** — bf16 batch 1024 + **`TORCH_BLAS_PREFER_HIPBLASLT=0`** (rocBLAS fallback) → **same crash** (not hipBLASLt-specific). + - **9233** — bf16 batch 256 + **`HIP_LAUNCH_BLOCKING=1` / `AMD_SERIALIZE_KERNEL=3`** → **same crash, synchronous** → the failing op is genuinely the **`var_agg` output `proj` `F.linear`**, not an async misattribution. + - **9234** — **float32** batch 256, identical config → **COMPLETED**. ✅ +- **Shape at failure:** `aggregate_variables` (edm.py:309) flattens to `(B·H·L, V, D)`; after var cross-attention `proj` runs on `(B·H·L, 1, 256)` → GEMM **M = batch·648, K=256, N=256** in bf16. fp32 identical GEMM works. +- **Initial (wrong) verdict:** looked like a stack-level bf16 GEMM bug. **Superseded — see RESOLVED below.** + +#### RESOLVED (2026-06-11) — root cause: ROCm **Flash** SDPA 65535 batch-grid cap in `var_agg`; fix: force EFFICIENT/MATH + +**Root cause (isolated with a 1-GPU op-replica probe, no data):** +- The standalone bf16 **`proj` GEMM passes** even at M=663552 → the Linear is *not* the problem (the async HIP error was misattributed to `proj` line 280). +- A **faithful `CrossAttention` replica** reproduces the crash: bf16 **PASS at batch ≤64** (B=batch·648=41472), **FAIL at batch ≥128** (B=82944). The trigger is the **collapsed attention batch dim `B = batch·tokens`**: `aggregate_variables` flattens `(B,History,L)` into the SDPA batch dim, and ROCm **Flash** attention maps that to a HIP grid dimension capped at **65535** (`65535/648 ≈ 101` → matches batch 64 pass / 128 fail). +- **Backend sweep at batch 128/256 (bf16):** `default`(Flash) **FAIL**; **`math` PASS**, **`efficient` (EFFICIENT_ATTENTION) PASS**, **`eager` PASS**. fp32 "works" only because it dispatches to math, not Flash. +- **Independently corroborated by the upstream PyTorch SDPA benchmark** ([`examples/upstream_pytorch_sdpa_benchmark.py`](../../examples/upstream_pytorch_sdpa_benchmark.py) `--orbit-varagg`, which replicates the var_agg shape `batch_size=batch*648, q_len=1, kv_len=N_i`): + + | backend | batch 64 (B=41472) | batch 128 (B=82944) | batch 256 (B=165888) | + |---|---|---|---| + | default | PASS | **FAIL** | **FAIL** | + | flash | PASS | **FAIL** | **FAIL** | + | efficient | PASS | PASS | PASS | + | math | PASS | PASS | PASS | + + Confirms `default ≡ flash` for these shapes and the Flash 65535 batch-grid cap (`64·648=41472` ok, `128·648=82944` fails). Run: `python3 upstream_pytorch_sdpa_benchmark.py --orbit-varagg --orbit-varagg-batches 64,128,256 --orbit-sweep-backends default,flash,efficient,math` (1 GPU; `pip install --target=/tmp/... tabulate tqdm` first per §"Apptainer venv"). + +**Fix (compute-side clone `src/climate_learn/models/hub/components/attention.py`):** wrap **all four** SDPA call sites — **both `Attention` (self) and `CrossAttention`**, DEFAULT + CK branches — in `with sdpa_kernel([SDPBackend.EFFICIENT_ATTENTION, SDPBackend.MATH])` to avoid Flash. MATH is fallback-only (never selected for supported shapes). + + *Why global, not just CrossAttention:* the [`--orbit-selfattn`](../../examples/upstream_pytorch_sdpa_benchmark.py) benchmark (q=kv=648, heads=8, bf16) shows **EFFICIENT within ±3–5% of Flash** at batch 64/256/1024 (efficient faster at 64; flash faster fwd by ~1–2% at large batch; backward a wash), so unifying on EFFICIENT removes the 65535 grid foot-gun at ~noise-level perf cost. **MATH is 10–15× slower** (≈15 vs ≈290 TFLOPS) so it must stay fallback-only. + + | self-attn fwd/bwd µs | batch 64 | batch 256 | batch 1024 | + |---|---|---|---| + | flash | 123 / 449 | 408 / 1516 | 1515 / 5940 | + | efficient | 111 / 408 | 419 / 1514 | 1534 / 5909 | + | math | — | — | 29220 / 30940 | + +**Validated end-to-end:** job **9238** = the exact bf16 batch-256 config that crashed as **9231** now **COMPLETED (exit 0)**, `y.shape=[256,3,108,216]`, decreasing loss (`batch_idx 0→0.971, 1→0.963, 2→0.930, …`), full 6 epochs. Global fix (both `Attention`+`CrossAttention` → EFFICIENT) re-validated by job **9257** (bf16 **b1024**, COMPLETED, no errors). + +**CK-flash flags are inert on this wheel:** setting `TORCH_ROCM_FA_PREFER_CK=1` + `USE_CK_FLASH_ATTENTION=1` does **not** change anything — the flash/efficient kernel stays AOTriton `attn_fwd`, flash still crashes at B>65535 (var_agg batch 128), and self-attn perf is unchanged. `USE_CK_FLASH_ATTENTION` is a **build-time** flag and this `pytorch_rocm7.2.2` wheel wasn't built with CK flash, so `TORCH_ROCM_FA_PREFER_CK` has nothing to prefer. Getting CK flash would require rebuilding PyTorch with `USE_CK_FLASH_ATTENTION=1` — not worth it given EFFICIENT (AOTriton) already matches Flash perf and has no grid cap. + +**Caveats / follow-ups:** +- **`run_fom_extractor.py` / `parse_training_log.py` gap:** completed EDM runs log per-step loss as `epoch: N batch_idx M loss tensor(L)` and there's **no `Batch N: … seconds`** line in this build → the parser reports `epoch_records=0`, `final_loss=null`, and no `steady_batch_time_s`. **Throughput/loss FOMs need a parser update** to read the `epoch:/batch_idx/loss` format (and to time steps from omnistat or added timing). +- **HBM saturation vs data cap:** dataset has **~1704 train samples** (batch ≥1704 caps). bf16 at the cap will use roughly half the fp32 footprint (fp32@1704 ≈ 183 GiB ≈ 64%), so **bf16 alone won't reach 85–90% at 1704 samples** — stage more ERA5 ([`STAGING_ERA5_FOR_HBM.md`](../perf-optimizer-loop/STAGING_ERA5_FOR_HBM.md)) to push higher. + +### xFormers CK probe (MI355, `pytorch_rocm7.2.2` SIF + ORBIT-2 overlay) + +**`python3 -m xformers.info`:** `memory_efficient_attention.ckF`, `.ckB`, `.ck_splitKF` report **available** (Cutlass/FA paths unavailable on AMD, expected). + +**PyTorch wheel index (`https://download.pytorch.org/whl/rocm7.2`):** `pip index versions xformers` listed only **`0.0.34`** and **`0.0.35`** — no newer wheel than overlay **0.0.35** without changing ROCm/PyTorch image. + +**A one-GPU CK probe (bf16) SIGSEGVs (exit 139) on every case** — `small`, `asym`, `long_p0`, `long_p01` — after printing the case header, i.e. **`memory_efficient_attention` + `MemoryEfficientAttentionCkOp` + bf16** fails even on tiny tensors. PyTorch SDPA on the same shapes **succeeds**. + +**Implication:** do not treat **`xformers.info` “ckF available”** as safe for training MEA on this stack; keep **`ORBIT2_FUSED_ATTN=DEFAULT`** (SDPA) unless a new wheel/image is validated. + +### PyTorch SDPA benchmark (upstream `sdpa.py`) + +Studio copy: [`examples/upstream_pytorch_sdpa_benchmark.py`](../../examples/upstream_pytorch_sdpa_benchmark.py) — exercises **`torch.nn.functional.scaled_dot_product_attention`** under `sdpa_kernel(SDPBackend)` (same idea as [PyTorch `benchmarks/transformer/sdpa.py`](https://github.com/pytorch/pytorch/blob/main/benchmarks/transformer/sdpa.py)). + +**Naming (corrected 2026-06-11):** the old CLI label **`ck`** for `--orbit-sweep-backends` was **removed** — it mapped to `SDPBackend.EFFICIENT_ATTENTION`, which on this ROCm build is **AOTriton**, *not* Composable Kernel. **Use `efficient`.** Verified by profiling: both `flash` and `efficient` dispatch to the same AOTriton kernel **`attn_fwd`**; `math` is the unfused reference (`Cijk_...` Tensile GEMMs + `softmax_warp_forward`). Real **CK** exists only in the xFormers path (`--orbit-include-xformers-ck` / `--orbit-xformers-modes ck` → `MemoryEfficientAttentionCkOp`, the one that SIGSEGVs). Flash vs efficient differ only in the **grid-launch wrapper** (flash maps the batch dim to a 65535-capped grid dim; efficient does not) — same `attn_fwd` math kernel. + +**xFormers CK on generic shapes (no ORBIT code):** pass **`--orbit-include-xformers-ck`** with **`--orbit-micro`**. Each case runs in a **fresh subprocess** so a child **SIGSEGV** is reported as `KILLED_BY_SIGNAL` / `SIGSEGV` in the summary table while PyTorch SDPA results (already printed) stay intact. Tune with **`--orbit-xformers-modes ck,dispatch`**, **`--orbit-xformers-dropout 0.0,0.1`**, **`--orbit-xformers-timeout SEC`**. + +**Validated (MI355X, SLURM job **9207**, `pytorch_rocm7.2.2` SIF + ORBIT-2 overlay, 2026-06-11):** **`--orbit-micro`** SDPA sweeps (**`default`**, **`math`**, **`efficient`**, CLI label **`ck`** → `SDPBackend.EFFICIENT_ATTENTION`, **`flash`**) all **completed** for bf16 shapes 128×128 and 648×648. **`--orbit-include-xformers-ck`** — every row (**`micro_128_sym`**, **`micro_648_sym`**, **`asym_1_vs_6`** × **`ck`** and **`dispatch`**, `dropout_p=0`) reported **`KILLED_BY_SIGNAL` / `SIGSEGV`** (standalone tensors; no Bayes-CAST). Example CSV dir: **`$AI4S_SHARED_DIR/models/ORBIT-2/perf-runs/sdpa-ck-microbench-1781204729/`**. + +**Apptainer venv:** `pip install --user` fails (“User site-packages are not visible in this virtualenv”). Use **`pip install -q --target=/tmp/orbit2_sdpa_pip tabulate tqdm`** and **`export PYTHONPATH=/tmp/orbit2_sdpa_pip:$PYTHONPATH`** before running the script. + +**One GPU, bind-mount Studio `examples/`** (adjust `AI4S_*` paths): + +```bash +export AI4S_SHARED_DIR=/path/to/shared +STUDIO_ORBIT2_EXAMPLES=/path/to/ai4science-studio/earth_science/models/ORBIT-2/examples +srun -N1 -n1 --gres=gpu:1 apptainer exec --rocm \ + --overlay "${ORBIT2_OVERLAY:?}":ro \ + --bind "$STUDIO_ORBIT2_EXAMPLES":/sdpa-bench \ + "${ORBIT2_SIF:?}" bash -lc ' + set -e + PIP_LOCAL=/tmp/orbit2_sdpa_pip + mkdir -p "$PIP_LOCAL" + python3 -c "import tabulate, tqdm" 2>/dev/null || pip install -q --target="$PIP_LOCAL" tabulate tqdm + export PYTHONPATH="$PIP_LOCAL:${PYTHONPATH:-}" + export ORBIT2_SDPA_CSV_DIR="${ORBIT2_SDPA_CSV_DIR:-/tmp/orbit2_sdpa_benchmark_results}" + mkdir -p "$ORBIT2_SDPA_CSV_DIR" + cd /sdpa-bench + python3 upstream_pytorch_sdpa_benchmark.py --orbit-micro \ + --orbit-sweep-backends default,math,efficient,ck,flash \ + --orbit-include-xformers-ck \ + --orbit-xformers-modes ck,dispatch \ + --orbit-xformers-dropout 0.0,0.1 + ' +``` + +Expect **`math`** and often **`default`** / **`efficient`** to complete on ROCm; **`flash`** may skip or error if Flash is unavailable — compare timings, not only success/fail. CSV timestamps land in **`ORBIT2_SDPA_CSV_DIR`**. + +## Next work (after scaling baseline) + +- **Run the sysopt loop:** [`recipes/perf-optimizer-loop/README.md`](../recipes/perf-optimizer-loop/README.md) — `examples/run_optimizer_loop.sh`, [`lever_catalog.yaml`](../recipes/perf-optimizer-loop/lever_catalog.yaml), staging + bf16 sweep ([`STAGING_ERA5_FOR_HBM.md`](../recipes/perf-optimizer-loop/STAGING_ERA5_FOR_HBM.md), [`examples/sweep_orbit2_batch_bf16_amd.sh`](../../examples/sweep_orbit2_batch_bf16_amd.sh)). +- Rebuild overlay; re-test `bfloat16` + CK attn when `xformers.ops` is available. +- Stage **2.5_arcmin** PRISM targets for scientifically meaningful loss (not same-dir timing). + +- **One-node GPU baseline (pre-sysopt):** [one-node-gpu-baseline.md](one-node-gpu-baseline.md) — batch sweep, Omnistat + TraceLens evidence, `report_orbit2_gpu_baseline.py` for `baseline_report.md`. + +## Key paths (parameterize with `AI4S_SHARED_DIR`) + +- Data: `$AI4S_SHARED_DIR/models/ORBIT-2/data/superres/prism/10.0_arcmin` (PRISM same-dir) or `.../era5/1.0_deg` (ERA5 same-dir sanity template). +- Overlay: `$AI4S_SHARED_DIR/models/ORBIT-2/overlays/orbit2-overlay.img` +- Perf runs: `$AI4S_SHARED_DIR/models/ORBIT-2/perf-runs//` +- Training logs (recommended): `$AI4S_SHARED_DIR/models/ORBIT-2/outputs/train/logs/` +- Omnistat / TraceLens tools: set **`OMNIHUB_TOOLS_DIR`** to your site-local checkout (see root `.cursor/skills/ai4science-perf-analysis/SKILL.md` and `.cluster-config.yaml` key `omnihub.tools_dir` if your site maintains one). + +## Operational source tree + +- Set **`ORBIT2_ROOT`** to your clone (public **ORBIT-2** or institutional **Bayes-CAST**). Scripts bind-mount it at **`/orbit2`**. +- **`launch_diffusion.sh`:** if `$ORBIT2_ROOT/launch_diffusion.sh` or `examples/launch_diffusion.sh` exists, ranks **`exec bash …/launch_diffusion.sh /config/config.yaml`** after optional `ORBIT2_RANK_PRE_TRAIN_HOOK` (via `orbit2_rank_hook_runner.py`). Override path with **`ORBIT2_LAUNCH_SCRIPT`** (must live under `ORBIT2_ROOT`). +- **`manifest.json`** (perf runs): includes `git_sha`, `git_branch`, `git_remote_origin`, `rendered_config`, `parallelism.fsdp` / `parallelism.simple_ddp`, `global_batch_size`, `runtime_seconds`, `orbit2_batch_size`, `total_ranks`, `max_epochs`, `data_type`, and `workload` (`bayes_edm` vs `intermediate_downscaling`). diff --git a/earth_science/models/ORBIT-2/recipes/perf-analysis/README.md b/earth_science/models/ORBIT-2/recipes/perf-analysis/README.md index 4b50156..900b5f3 100644 --- a/earth_science/models/ORBIT-2/recipes/perf-analysis/README.md +++ b/earth_science/models/ORBIT-2/recipes/perf-analysis/README.md @@ -10,7 +10,7 @@ Multi-subagent workflow for ORBIT-2 / Bayes-CAST training on AMD MI355X, with Om ```bash export AI4S_SHARED_DIR=/path/to/shared -export PERF_TOOLS_DIR=/path/to/perf-tools +export OMNIHUB_TOOLS_DIR=/shared/omnihub/tools # ORBIT2_ROOT defaults to .../code/bayes-cast when that directory exists export ORBIT2_DATA_ROOT=$AI4S_SHARED_DIR/models/ORBIT-2/data/superres/era5/1.0_deg # export ORBIT2_CONFIG_TEMPLATE=edm_8m_era5_1x8.yaml # default diff --git a/earth_science/models/ORBIT-2/recipes/perf-analysis/agents/launcher.md b/earth_science/models/ORBIT-2/recipes/perf-analysis/agents/launcher.md index 1c3dce2..eec4cdc 100644 --- a/earth_science/models/ORBIT-2/recipes/perf-analysis/agents/launcher.md +++ b/earth_science/models/ORBIT-2/recipes/perf-analysis/agents/launcher.md @@ -10,8 +10,8 @@ Submits a 2-node ORBIT-2 training (AMD Instinct MI355X) with PyTorch profiling a ## Outputs -- `${PERF_TOOLS_DIR}/perf-inspect/` — Python venv with omnistat (PR #271 + main merged) and TraceLens. -- `${PERF_TOOLS_DIR}/victoriametrics/victoria-metrics-prod` — VictoriaMetrics binary. +- `${OMNIHUB_TOOLS_DIR}/omnihub-inspect/` — Python venv with omnistat (PR #271 + main merged) and TraceLens. +- `${OMNIHUB_TOOLS_DIR}/victoriametrics/victoria-metrics-prod` — VictoriaMetrics binary. - `$AI4S_SHARED_DIR/models/ORBIT-2/perf-runs/omnistat.config` — Omnistat user-mode config. - `$AI4S_SHARED_DIR/models/ORBIT-2/perf-runs//manifest.json` — manifest schema below. - `$AI4S_SHARED_DIR/models/ORBIT-2/perf-runs//orbit2-train-.out` — symlinked from output dir. @@ -57,8 +57,8 @@ Submits a 2-node ORBIT-2 training (AMD Instinct MI355X) with PyTorch profiling a ```bash set -euo pipefail -# Shared tool location — read from .cluster-config.yaml perf_tools.dir. -TOOLS="${PERF_TOOLS_DIR:?set PERF_TOOLS_DIR to a persistent checkout (e.g. under \$AI4S_SHARED_DIR/tools)}" +# Shared tool location — read from .cluster-config.yaml omnihub.tools_dir. +TOOLS="${OMNIHUB_TOOLS_DIR:?set OMNIHUB_TOOLS_DIR to a persistent checkout (e.g. under \$AI4S_SHARED_DIR/tools)}" mkdir -p "$TOOLS" # 1a. Omnistat: jorda/skills branch with origin/main merged in @@ -77,7 +77,7 @@ fi OMNISTAT_COMMIT=$(git rev-parse HEAD) # 1b. Venv (py3.12 to match cluster python) -VENV=${PERF_TOOLS_DIR}/perf-inspect +VENV=${OMNIHUB_TOOLS_DIR}/omnihub-inspect if [[ ! -d $VENV ]]; then python3 -m venv "$VENV" fi @@ -123,7 +123,7 @@ if [[ "${OMNISTAT_KERNEL_TRACE:-0}" == "1" && ! -f "$TRACE_LIB" ]]; then fi ``` -The sbatch wrapper expects the `.so` at `${PERF_TOOLS_DIR}/omnistat-src/build-trace/libomnistat_trace.so` by default; override with `OMNISTAT_TRACE_LIB=/path/to/lib.so`. See `ai4science-studio` SKILL §12 for when to enable kernel tracing vs the default device-counting collector. +The sbatch wrapper expects the `.so` at `${OMNIHUB_TOOLS_DIR}/omnistat-src/build-trace/libomnistat_trace.so` by default; override with `OMNISTAT_TRACE_LIB=/path/to/lib.so`. See `ai4science-studio` SKILL §12 for when to enable kernel tracing vs the default device-counting collector. ### 2. Author the Omnistat user-mode config (once, in repo `perf-runs/`) diff --git a/earth_science/models/ORBIT-2/recipes/perf-analysis/agents/omnistat_analyst.md b/earth_science/models/ORBIT-2/recipes/perf-analysis/agents/omnistat_analyst.md index d97cd5d..e7a0003 100644 --- a/earth_science/models/ORBIT-2/recipes/perf-analysis/agents/omnistat_analyst.md +++ b/earth_science/models/ORBIT-2/recipes/perf-analysis/agents/omnistat_analyst.md @@ -6,8 +6,8 @@ Drive `omnistat-inspect` (PR #271) through the analyze-job phases on the user-mo - `/manifest.json` - The Omnistat DB at `manifest.omnistat_db`. -- The omnistat/TraceLens venv at `${OMNISTAT_VENV:-${PERF_TOOLS_DIR}/perf-inspect}/` (the venv name is site-specific; honour `OMNISTAT_VENV` when set at submit time). -- The VictoriaMetrics binary at `${PERF_TOOLS_DIR}/victoriametrics/victoria-metrics-prod`. +- The omnistat/TraceLens venv at `${OMNISTAT_VENV:-${OMNIHUB_TOOLS_DIR}/omnihub-inspect}/` (the venv name is site-specific; honour `OMNISTAT_VENV` when set at submit time). +- The VictoriaMetrics binary at `${OMNIHUB_TOOLS_DIR}/victoriametrics/victoria-metrics-prod`. ## Outputs @@ -36,7 +36,7 @@ ss -lnt "sport = :$PORT" 2>/dev/null | grep -q LISTEN && PORT=8429 VMLOG=$PERF_RUN/omnistat/vm.log mkdir -p "$PERF_RUN/omnistat/inspect_outputs" -nohup ${PERF_TOOLS_DIR}/victoriametrics/victoria-metrics-prod \ +nohup ${OMNIHUB_TOOLS_DIR}/victoriametrics/victoria-metrics-prod \ -storageDataPath="$DB" \ -httpListenAddr=127.0.0.1:$PORT \ -retentionPeriod=100y \ @@ -62,7 +62,7 @@ done ### 2. Run analyze-job phases (per the SKILL) ```bash -. "${OMNISTAT_VENV:-${PERF_TOOLS_DIR}/perf-inspect}/bin/activate" +. "${OMNISTAT_VENV:-${OMNIHUB_TOOLS_DIR}/omnihub-inspect}/bin/activate" SCRATCH=$PERF_RUN/omnistat/scratch mkdir -p "$SCRATCH" JOBID=$(jq -r .job_id ) # sbatch manifest key is job_id (not jobid) diff --git a/earth_science/models/ORBIT-2/recipes/perf-analysis/agents/orchestrator_gemm_analysis.md b/earth_science/models/ORBIT-2/recipes/perf-analysis/agents/orchestrator_gemm_analysis.md index c4d6d0d..4520c4b 100644 --- a/earth_science/models/ORBIT-2/recipes/perf-analysis/agents/orchestrator_gemm_analysis.md +++ b/earth_science/models/ORBIT-2/recipes/perf-analysis/agents/orchestrator_gemm_analysis.md @@ -12,7 +12,7 @@ GEMM time is actually spent on** (which shapes/kernels dominate, compute- vs mem and what changes from 1→2 nodes), to decide the next real lever. ## Fixed context (passed in the user prompt) -- `REPO_ROOT`, `AI4S_SHARED_DIR`, `PERF_TOOLS_DIR` (`perf_tools.dir` in `.cluster-config.yaml`). +- `REPO_ROOT`, `AI4S_SHARED_DIR`, `OMNIHUB_TOOLS_DIR` (`omnihub.tools_dir` in `.cluster-config.yaml`). - SLURM: partition and account from `.cluster-config.yaml` (`slurm.partition`, `slurm.account`). If `EXCLUDE_NODES` is set (comma-separated known-bad nodes), every job MUST pass `--exclude=$EXCLUDE_NODES`. @@ -53,7 +53,7 @@ bf16, SDPA DEFAULT, `ORBIT2_BATCH_SIZE=4096`, `ORBIT2_MAX_EPOCH=6`, `sbatch --parsable --partition= --account= --nodes= [--exclude=$EXCLUDE_NODES] --job-name=o2-gemm-n - --export=ALL,AI4S_SHARED_DIR=...,PERF_TOOLS_DIR=...,ORBIT2_DATA_ROOT=...,ORBIT2_ROOT=..., + --export=ALL,AI4S_SHARED_DIR=...,OMNIHUB_TOOLS_DIR=...,ORBIT2_DATA_ROOT=...,ORBIT2_ROOT=..., ORBIT2_CONFIG_TEMPLATE=edm_8m_era5_1x8.yaml,ORBIT2_ERA5_SPATIAL_RES=111,ORBIT2_BATCH_SIZE=4096, ORBIT2_MAX_EPOCH=6,TORCH_NCCL_HIGH_PRIORITY=1,GPU_MAX_HW_QUEUES=2, PROFILE_TARGET_EPOCH=2,PROFILE_RANK0_ONLY=1 `. @@ -113,4 +113,4 @@ Append `LOOP_COMPLETE 1node_jobid=<> 2node_jobid=<>` to STATUS.txt and print it ## 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. +- Keep `STATUS.txt` append-only: ` EVENT key=val...` per line. \ No newline at end of file diff --git a/earth_science/models/ORBIT-2/recipes/perf-analysis/agents/tracelens_analyst.md b/earth_science/models/ORBIT-2/recipes/perf-analysis/agents/tracelens_analyst.md index 0522aaf..a69349a 100644 --- a/earth_science/models/ORBIT-2/recipes/perf-analysis/agents/tracelens_analyst.md +++ b/earth_science/models/ORBIT-2/recipes/perf-analysis/agents/tracelens_analyst.md @@ -40,7 +40,7 @@ Every entry: ```bash # Venv name is site-specific; honour OMNISTAT_VENV (set at submit time), else the portable default. -. "${OMNISTAT_VENV:-${PERF_TOOLS_DIR}/perf-inspect}/bin/activate" +. "${OMNISTAT_VENV:-${OMNIHUB_TOOLS_DIR}/omnihub-inspect}/bin/activate" # sbatch-written manifest uses job_dir / trace_dir (not perf_run_dir / trace_paths[]). PERF_RUN=$(jq -r .job_dir ) TRACE=$(find "$(jq -r .trace_dir )" -name '*.pt.trace.json' 2>/dev/null | sort | head -1) diff --git a/earth_science/models/ORBIT-2/recipes/perf-analysis/omnistat.config.template b/earth_science/models/ORBIT-2/recipes/perf-analysis/omnistat.config.template index deecfd7..f638cfa 100644 --- a/earth_science/models/ORBIT-2/recipes/perf-analysis/omnistat.config.template +++ b/earth_science/models/ORBIT-2/recipes/perf-analysis/omnistat.config.template @@ -134,7 +134,7 @@ system_name = mi355x #-- [omnistat.usermode] -victoria_binary = @PERF_TOOLS_DIR@/victoriametrics/victoria-metrics-prod +victoria_binary = @OMNIHUB_TOOLS_DIR@/victoriametrics/victoria-metrics-prod victoria_datadir = @JOB_DIR@/omnistat-db victoria_logfile = vic_server.log diff --git a/earth_science/models/ORBIT-2/recipes/perf-analysis/one-node-gpu-baseline.md b/earth_science/models/ORBIT-2/recipes/perf-analysis/one-node-gpu-baseline.md index f379455..7d4c888 100644 --- a/earth_science/models/ORBIT-2/recipes/perf-analysis/one-node-gpu-baseline.md +++ b/earth_science/models/ORBIT-2/recipes/perf-analysis/one-node-gpu-baseline.md @@ -25,7 +25,7 @@ Goal: on **one MI355-class node × 8 GPUs**, pick a **per-rank batch size** and ```bash export AI4S_SHARED_DIR=... -export PERF_TOOLS_DIR=... +export OMNIHUB_TOOLS_DIR=... export ORBIT2_ROOT=... # optional; defaults to $AI4S_SHARED_DIR/models/ORBIT-2/code/ORBIT-2 # Defaults pick ERA5 path + template; omit these if data is already staged there: # export ORBIT2_DATA_ROOT=$AI4S_SHARED_DIR/models/ORBIT-2/data/superres/era5/1.0_deg diff --git a/earth_science/models/ORBIT-2/recipes/perf-optimizer-loop/README.md b/earth_science/models/ORBIT-2/recipes/perf-optimizer-loop/README.md index f01e32a..a432827 100644 --- a/earth_science/models/ORBIT-2/recipes/perf-optimizer-loop/README.md +++ b/earth_science/models/ORBIT-2/recipes/perf-optimizer-loop/README.md @@ -10,14 +10,14 @@ Agent-driven optimization loop for **ORBIT-2** training with **throughput (`thro - **Container:** `pytorch_rocm7.2.2_ubuntu24.04_py3.12_pytorch_release_2.10.0.sif` + ORBIT-2 overlay (`build_overlay_amd.sh`). - **Data:** staged ERA5 NPZ tree; for HBM saturation follow [`STAGING_ERA5_FOR_HBM.md`](STAGING_ERA5_FOR_HBM.md). - **Code:** `ORBIT2_ROOT` → Bayes-CAST with `launch/train_edm.py` (see [perf-analysis/README.md](../perf-analysis/README.md)). -- **Tools:** `PERF_TOOLS_DIR` (omnistat-usermode + VictoriaMetrics), `AI4S_SHARED_DIR`. +- **Tools:** `OMNIHUB_TOOLS_DIR` (omnistat-usermode + VictoriaMetrics), `AI4S_SHARED_DIR`. ## Environment variables (loop driver) | Variable | Required | Description | |----------|----------|-------------| | `AI4S_SHARED_DIR` | Yes | Shared storage root | -| `PERF_TOOLS_DIR` | Yes | Omnistat + VM binaries | +| `OMNIHUB_TOOLS_DIR` | Yes | Omnistat + VM binaries | | `ORBIT2_ROOT` | Recommended | Bayes-CAST checkout | | `ORBIT2_DATA_ROOT` | Yes | ERA5 staging path | | `ORBIT2_BATCH_SIZE` | Yes (baseline) | Locked after VRAM sweep | @@ -29,7 +29,7 @@ Agent-driven optimization loop for **ORBIT-2** training with **throughput (`thro ```bash export AI4S_SHARED_DIR=... -export PERF_TOOLS_DIR=... +export OMNIHUB_TOOLS_DIR=... bash earth_science/models/ORBIT-2/examples/validate_orbit2_optimizer_loop_recipe.sh tmux new -s orbit2-loop diff --git a/earth_science/models/ORBIT-2/recipes/perf-optimizer-loop/agents/orchestrator.md b/earth_science/models/ORBIT-2/recipes/perf-optimizer-loop/agents/orchestrator.md index 67fc129..91a7ef7 100644 --- a/earth_science/models/ORBIT-2/recipes/perf-optimizer-loop/agents/orchestrator.md +++ b/earth_science/models/ORBIT-2/recipes/perf-optimizer-loop/agents/orchestrator.md @@ -8,7 +8,7 @@ Drives the ORBIT-2 (Bayes-CAST EDM) optimizer loop on **MI355X-class** hardware. - Loop args: ``, ``. - `REPO_ROOT` (cwd or env). -- `.cluster-config.yaml` (partition, perf_tools.dir). +- `.cluster-config.yaml` (partition, omnihub.tools_dir). - [`../lever_catalog.yaml`](../lever_catalog.yaml). ## Outputs @@ -32,7 +32,7 @@ Per job: existing perf-analysis layout under `..//` plus `foms.json` from ## Preflight (ORBIT-2) -- `AI4S_SHARED_DIR`, `PERF_TOOLS_DIR` set. +- `AI4S_SHARED_DIR`, `OMNIHUB_TOOLS_DIR` set. - `${AI4S_SHARED_DIR}/models/ORBIT-2/overlays/orbit2-overlay.img` exists. - `${AI4S_SHARED_DIR}/images/pytorch_rocm7.2.2_ubuntu24.04_py3.12_pytorch_release_2.10.0.sif` (or `ORBIT2_SIF`) exists. - `omnistat-usermode` + `victoria-metrics-prod` binaries present. @@ -50,4 +50,4 @@ Same as HydraGNN orchestrator (`STATUS=ok|partial|fail` last line; resume from ` ``` STATUS=ok; reason=loop_complete uuid= best_iter= best_lever= best_throughput= samples/s -``` +``` \ No newline at end of file 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 index 89fad9a..d80d230 100644 --- 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 @@ -60,7 +60,7 @@ Tested stacking `channels_last` on the padding idea as a clean 2×2 ({NCHW, chan ```bash export AI4S_SHARED_DIR=/path/to/shared -export PERF_TOOLS_DIR=/path/to/perf-tools +export OMNIHUB_TOOLS_DIR=/shared/omnihub/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: diff --git a/integrations/omnihub/README.md b/integrations/omnihub/README.md new file mode 100644 index 0000000..9ca63b7 --- /dev/null +++ b/integrations/omnihub/README.md @@ -0,0 +1,76 @@ +# OmniHub integration for AI4Science Studio + +Bridge between [AI4Science Studio](https://github.com/AMDResearch/ai4science-studio) model recipes and [OmniHub](https://github.com/AMDResearch/omnihub) job generation, without merging the two repositories. + +## Prerequisites + +- OmniHub clone (sibling repo), e.g. `$HOME/git/omnihub` +- `.cluster-config.yaml` with `omnihub:` section (see `.cluster-config.example.yaml`) +- `AI4S_SHARED_DIR` set (HydraGNN assets: SIF, overlay, datasets, code) +- Lux/Vultr: OmniHub `config/vultr.yaml` (`shared-dir: /shared/omnihub`) + +## Quick start (HydraGNN train) + +```bash +export AI4S_SHARED_DIR=/shared/$USER +export OMNIHUB_DIR=$HOME/git/omnihub + +# 1. Sync application stubs into the OmniHub checkout +./integrations/omnihub/sync-to-omnihub.sh + +# 2. Generate SLURM job (wrapper sets HydraGNN binds/overlay) +./integrations/omnihub/generate-job.sh \ + --num-nodes 2 --partition lux --time-limit 2h \ + --output /tmp/hydragnn-omnihub.slurm + +# 3. Submit +sbatch -A vultr_lux /tmp/hydragnn-omnihub.slurm + +# 4. Post-process (after job completes) +$OMNIHUB_DIR/omnihub-process --results-dir /shared/$USER/results/omnihub -j 4 +$OMNIHUB_DIR/omnihub-index --results-dir /shared/$USER/results/omnihub --output index +``` + +## Layout + +``` +integrations/omnihub/ +├── README.md # this file +├── sync-to-omnihub.sh # rsync apps → $OMNIHUB_DIR/applications/ +├── generate-job.sh # sync + omnihub-generate-job + HydraGNN apptainer patches +└── applications/ + └── hydragnn-train/ # OmniHub app (synced into omnihub repo) +``` + +## Environment variables + +| Variable | Purpose | +|----------|---------| +| `OMNIHUB_DIR` | Path to OmniHub repo (required for sync/generate) | +| `AI4S_SHARED_DIR` | Shared assets root (SIF, HydraGNN overlay, datasets) | +| `AI4S_REPO_ROOT` | ai4science-studio root (default: auto-detected) | + +HydraGNN training knobs (`HG_*`, `HYDRAGNN_*`) match the standard `sbatch_train_amd.sh` recipe. + +## Results layout + +OmniHub jobs write to `/shared/$USER/results/omnihub/$SLURM_JOB_ID/` (per `config/vultr.yaml`): + +``` +/ +├── job.yaml, app.yaml, job-status.yaml +├── logs/srun-*.out +├── tools/omnihub-monitor/ +└── processed-data/ # after omnihub-process +``` + +Agents should read `processed-data/` first (token-efficient). See `.cursor/skills/ai4science-omnihub/SKILL.md`. + +## Recipes + +| task | OmniHub tools | Notes | +|------|---------------|-------| +| `train-omnihub` | `omnihub-monitor` | Phase 1 pilot | +| `perf-analysis-omnihub` | `omnihub-monitor omnistat pytorch-trace tracelens` | Phase 2; optional `omnistat-rocprofiler-pmc1` | + +Legacy HydraGNN perf (`sbatch_train_perf_amd.sh` + `omnihub-inspect` at `/shared/omnihub/tools/omnihub-inspect`) remains available for deep agent workflows. diff --git a/integrations/omnihub/applications/hydragnn-train/config-perf.yaml b/integrations/omnihub/applications/hydragnn-train/config-perf.yaml new file mode 100644 index 0000000..9a2e5a6 --- /dev/null +++ b/integrations/omnihub/applications/hydragnn-train/config-perf.yaml @@ -0,0 +1,11 @@ +entrypoint: applications/ai4science-hydragnn-train/train_wrapper.py + +# Perf-analysis variant: epoch-gated kineto on rank 0 (use with --perf in generate-job.sh) + +HydraGNN: + datasets: ANI1x,Alexandria + batch_size: 200 + num_epoch: 2 + precision: fp64 + max_num_batch: 30 + profile_target_epoch: 1 diff --git a/integrations/omnihub/applications/hydragnn-train/config.yaml b/integrations/omnihub/applications/hydragnn-train/config.yaml new file mode 100644 index 0000000..1a676b5 --- /dev/null +++ b/integrations/omnihub/applications/hydragnn-train/config.yaml @@ -0,0 +1,13 @@ +entrypoint: applications/ai4science-hydragnn-train/train_wrapper.py + +# HydraGNN multi-dataset GFM training via OmniHub. +# Paths resolve from AI4S_SHARED_DIR / HG_* env vars (see train_wrapper.py). + +HydraGNN: + datasets: ANI1x,Alexandria + batch_size: 200 + num_epoch: 1 + precision: fp64 + max_num_batch: null + valtest: 0 + profile_target_epoch: null diff --git a/integrations/omnihub/applications/hydragnn-train/parse.py b/integrations/omnihub/applications/hydragnn-train/parse.py new file mode 100755 index 0000000..9aabb0c --- /dev/null +++ b/integrations/omnihub/applications/hydragnn-train/parse.py @@ -0,0 +1,37 @@ +#!/usr/bin/env python3 +"""Post-process HydraGNN OmniHub job results into processed-data/.""" + +import json +import os +import sys +from pathlib import Path + + +def main(): + if len(sys.argv) < 2: + print(f"Usage: {sys.argv[0]} ", file=sys.stderr) + sys.exit(1) + + job_dir = Path(sys.argv[1]) + processed = job_dir / "processed-data" + processed.mkdir(parents=True, exist_ok=True) + + summary = { + "app": "ai4science-hydragnn-train", + "job_dir": str(job_dir), + } + + status = job_dir / "job-status.yaml" + if status.is_file(): + summary["job_status_file"] = str(status) + + logs = list(job_dir.glob("logs/srun-*.out")) + summary["log_count"] = len(logs) + + out = processed / "app-parser.json" + out.write_text(json.dumps(summary, indent=2)) + print(f"Wrote {out}") + + +if __name__ == "__main__": + main() diff --git a/integrations/omnihub/applications/hydragnn-train/train_wrapper.py b/integrations/omnihub/applications/hydragnn-train/train_wrapper.py new file mode 100644 index 0000000..aaff709 --- /dev/null +++ b/integrations/omnihub/applications/hydragnn-train/train_wrapper.py @@ -0,0 +1,205 @@ +#!/usr/bin/env python3 +""" +HydraGNN training wrapper for OmniHub (AI4Science Studio integration). + +Runs the upstream multidataset_hpo_sc26 entrypoint inside the container with +MPI-style rank identity from Slurm (manual-mpi runner). Mirrors the rank script +in material_science/models/HydraGNN/examples/sbatch_train_amd.sh. +""" + +from __future__ import annotations + +import json +import os +import runpy +import sys +from pathlib import Path +from typing import Any, Dict, Optional + +import omnihub.run +import omnihub.tools + +# Pinned SHA matches sbatch_train_amd.sh default. +DEFAULT_HG_SHA = "2fb0bd0157e3c85a74f9841887155095bd163303" +DEFAULT_AVG_NEIGHBORS = 13.735293601560318 + + +def _require_env(name: str) -> str: + val = os.environ.get(name) + if not val: + raise RuntimeError(f"Required environment variable not set: {name}") + return val + + +def _ai4s_base() -> Path: + return Path(_require_env("AI4S_SHARED_DIR")) / "models" / "HydraGNN" + + +def _resolve_paths(config: Dict[str, Any]) -> Dict[str, Path]: + hg = config.get("HydraGNN", {}) + base = _ai4s_base() + repo = Path(os.environ.get("HG_REPO_DIR", base / "code" / "HydraGNN")) + example = repo / "examples" / "multidataset_hpo_sc26" + data_dir = Path(os.environ.get("HG_DATA_DIR", base / "weights")) + output = Path(os.environ.get("HG_OUTPUT_DIR", os.environ.get("OMNIHUB_RESULTS_DIR", "."))) + return { + "repo": repo, + "example": example, + "data_dir": data_dir, + "output": output, + "config_json": example / "gfm_mlip.json", + "entry": example / "gfm_mlip_all_mpnn.py", + } + + +def _ensure_datasets(paths: Dict[str, Path], datasets: str) -> None: + dataset_dir = paths["example"] / "dataset" + dataset_dir.mkdir(parents=True, exist_ok=True) + for ds in datasets.split(","): + ds = ds.strip() + if not ds: + continue + target = dataset_dir / f"{ds}-v2.bp" + source = paths["data_dir"] / f"{ds}-v2.bp" + if not source.is_dir(): + raise FileNotFoundError(f"Dataset not found: {source}") + if not target.exists(): + target.symlink_to(source, target_is_directory=True) + + +def _maybe_inject_profile(config_path: Path, target_epoch: Optional[int]) -> Path: + """Inject NeuralNetwork.Profile block for epoch-gated kineto (perf-analysis).""" + if target_epoch is None: + return config_path + data = json.loads(config_path.read_text()) + nn = data.setdefault("NeuralNetwork", {}) + nn["Profile"] = { + "enable": 1, + "target_epoch": int(target_epoch), + "wait": 5, + "warmup": 3, + "active": 3, + "repeat": 1, + } + out = Path( + os.environ.get( + "OMNIHUB_RESULTS_DIR", + os.environ.get("HG_OUTPUT_DIR", str(config_path.parent)), + ) + ) / "gfm_mlip_omnihub_profile.json" + out.write_text(json.dumps(data, indent=2)) + return out + + +def _patch_adios_avg_neighbors() -> None: + avg_nn = float(os.environ.get("HYDRAGNN_AVG_NUM_NEIGHBORS", DEFAULT_AVG_NEIGHBORS)) + if avg_nn <= 0: + return + import hydragnn.utils.datasets.adiosdataset as adm + + orig_init = adm.AdiosMultiDataset.__init__ + + def patched_init(self, *args, **kwargs): + orig_init(self, *args, **kwargs) + self.avg_num_neighbors = avg_nn + + adm.AdiosMultiDataset.__init__ = patched_init + + +def _setup_runtime_env(paths: Dict[str, Path]) -> None: + scratch = os.environ.get("SCRATCH_LOCAL", "/scratch") + proc = os.environ.get("SLURM_PROCID", os.environ.get("RANK", "0")) + job = os.environ.get("SLURM_JOB_ID", "local") + scratch_rank = Path(scratch) / os.environ.get("USER", "user") / f"hydragnn-{job}" / proc + scratch_rank.mkdir(parents=True, exist_ok=True) + os.environ["TMPDIR"] = str(scratch_rank) + os.environ["MIOPEN_DISABLE_CACHE"] = "1" + os.environ["MIOPEN_USER_DB_PATH"] = str(scratch_rank / "miopen") + Path(os.environ["MIOPEN_USER_DB_PATH"]).mkdir(parents=True, exist_ok=True) + os.environ["PYTHONNOUSERSITE"] = "1" + os.environ["HYDRAGNN_USE_VARIABLE_GRAPH_SIZE"] = "1" + os.environ["HYDRAGNN_AGGR_BACKEND"] = "mpi" + os.environ["HYDRAGNN_USE_FSDP"] = "0" + os.environ["HYDRAGNN_TRACE_LEVEL"] = os.environ.get("HYDRAGNN_TRACE_LEVEL", "1") + os.environ["HSA_NO_SCRATCH_RECLAIM"] = "1" + os.environ["HG_EXAMPLE_DIR"] = str(paths["example"]) + os.environ["HG_OUTPUT_DIR"] = str(paths["output"]) + paths["output"].mkdir(parents=True, exist_ok=True) + + pkg = "/opt/hydragnn-pkgs" + if Path(pkg).is_dir(): + os.environ["PYTHONPATH"] = f"{pkg}:{os.environ.get('PYTHONPATH', '')}" + ld = f"{pkg}/adios2:/opt/ompi/lib" + os.environ["LD_LIBRARY_PATH"] = f"{ld}:{os.environ.get('LD_LIBRARY_PATH', '')}" + + +def _build_argv(paths: Dict[str, Path], hg: Dict[str, Any], config_file: Path) -> list: + datasets = os.environ.get("HG_DATASETS", hg.get("datasets", "ANI1x,Alexandria")) + batch_size = int(os.environ.get("HG_BATCH_SIZE", hg.get("batch_size", 200))) + max_batch = os.environ.get("HYDRAGNN_MAX_NUM_BATCH", hg.get("max_num_batch") or "") + num_epoch = os.environ.get("HG_NUM_EPOCH", str(hg.get("num_epoch", 1))) + precision = os.environ.get("HG_PRECISION", hg.get("precision", "fp64")) + job = os.environ.get("SLURM_JOB_ID", "0") + nodes = os.environ.get("SLURM_JOB_NUM_NODES", "1") + + argv = [ + str(paths["entry"]), + f"--inputfile={config_file}", + "--multi", + "--everyone", + f"--multi_model_list={datasets}", + f"--precision={precision}", + f"--batch_size={batch_size}", + f"--num_epoch={num_epoch}", + "--startfrom=none", + f"--log=hydragnn-train-{job}-N{nodes}", + ] + if max_batch: + n = int(max_batch) * batch_size + argv += [ + f"--num_samples={n}", + "--oversampling", + f"--oversampling_num_samples={n}", + ] + return argv + + +@omnihub.run.entrypoint +def run(extra_args, config=None): + """OmniHub entrypoint — train HydraGNN GFM ensemble.""" + hg = (config or {}).get("HydraGNN", {}) + paths = _resolve_paths(config or {}) + datasets = os.environ.get("HG_DATASETS", hg.get("datasets", "ANI1x,Alexandria")) + + if not paths["entry"].is_file(): + raise FileNotFoundError( + f"HydraGNN entry not found: {paths['entry']}. " + f"Clone upstream to HG_REPO_DIR (SHA {DEFAULT_HG_SHA})." + ) + + _ensure_datasets(paths, datasets) + profile_epoch = hg.get("profile_target_epoch") + if profile_epoch is not None: + profile_epoch = int(profile_epoch) + elif os.environ.get("PROFILE_TARGET_EPOCH"): + profile_epoch = int(os.environ["PROFILE_TARGET_EPOCH"]) + + config_file = _maybe_inject_profile(paths["config_json"], profile_epoch) + _setup_runtime_env(paths) + + rank = int(os.environ.get("SLURM_PROCID", os.environ.get("RANK", "0"))) + if os.environ.get("PROFILE_RANK0_ONLY", "1") == "1" and rank != 0: + os.environ["HYDRAGNN_DISABLE_PROFILE"] = "1" + + @omnihub.tools.profile() + def _train(): + _patch_adios_avg_neighbors() + os.chdir(paths["output"]) + sys.argv = _build_argv(paths, hg, config_file) + runpy.run_path(str(paths["entry"]), run_name="__main__") + + _train() + + +if __name__ == "__main__": + run([], {}) diff --git a/integrations/omnihub/generate-job.sh b/integrations/omnihub/generate-job.sh new file mode 100755 index 0000000..fadf70b --- /dev/null +++ b/integrations/omnihub/generate-job.sh @@ -0,0 +1,191 @@ +#!/usr/bin/env bash +# Generate an OmniHub SLURM job for ai4science HydraGNN with Lux-specific apptainer binds. +set -euo pipefail + +SCRIPT_DIR=$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd) +AI4S_REPO_ROOT="${AI4S_REPO_ROOT:-$(cd "$SCRIPT_DIR/../.." && pwd)}" +OMNIHUB_DIR="${OMNIHUB_DIR:-}" +AI4S_SHARED_DIR="${AI4S_SHARED_DIR:-}" + +NUM_NODES=1 +PARTITION="" +CLUSTER=vultr +TIME_LIMIT=2h +TOOLS=(omnihub-monitor) +RUNNER=manual-mpi +TASKS_PER_NODE=8 +OUTPUT="" +IMAGE="" +APP_CONFIG=applications/ai4science-hydragnn-train/config.yaml +ROCM_VERSION=7.2.2 +EXTRA_ARGS=() + +usage() { + cat <&2 + exit 2 +fi + +if [[ ! -x "$OMNIHUB_DIR/omnihub-generate-job" ]]; then + echo "ERROR: OMNIHUB_DIR does not look like an OmniHub repo: $OMNIHUB_DIR" >&2 + exit 2 +fi + +SRC="$SCRIPT_DIR/applications" +DEST="$OMNIHUB_DIR/applications" + +mkdir -p "$DEST" + +for app_dir in "$SRC"/*/; do + [[ -d "$app_dir" ]] || continue + name=$(basename "$app_dir") + target="$DEST/ai4science-${name}" + echo "Syncing $name -> $target" + rsync -a --delete "$app_dir" "$target/" +done + +echo "Done. Use --app-config applications/ai4science-/config.yaml with omnihub-generate-job." diff --git a/material_science/models/HydraGNN/examples/run_fom_extractor.py b/material_science/models/HydraGNN/examples/run_fom_extractor.py index 7565c36..6ec0d42 100755 --- a/material_science/models/HydraGNN/examples/run_fom_extractor.py +++ b/material_science/models/HydraGNN/examples/run_fom_extractor.py @@ -6,7 +6,7 @@ Usage: export AI4S_SHARED_DIR=/shared/$USER - export PERF_TOOLS_DIR=/path/to/perf-tools # from .cluster-config.yaml perf_tools.dir + export OMNIHUB_TOOLS_DIR=/shared/omnihub/tools # from .cluster-config.yaml omnihub.tools_dir export REPO_ROOT=/path/to/ai4science-studio python run_fom_extractor.py --manifest /path/to/perf-runs//manifest.json \\ [--vm-port 8432] [--tsdb-url http://127.0.0.1:8432] @@ -140,7 +140,7 @@ def promql_scalar(tsdb_url: str, promql: str, t: int) -> float | None: def start_vm(db_path: str, port: int, log_path: Path) -> None: vm = os.environ.get( "VICTORIA_METRICS_BIN", - f"{os.environ['PERF_TOOLS_DIR']}/victoriametrics/victoria-metrics-prod", + f"{os.environ['OMNIHUB_TOOLS_DIR']}/victoriametrics/victoria-metrics-prod", ) subprocess.Popen( [ diff --git a/material_science/models/HydraGNN/examples/run_optimizer_loop.sh b/material_science/models/HydraGNN/examples/run_optimizer_loop.sh index a5b046e..37ca554 100755 --- a/material_science/models/HydraGNN/examples/run_optimizer_loop.sh +++ b/material_science/models/HydraGNN/examples/run_optimizer_loop.sh @@ -46,8 +46,8 @@ ORCH_PROMPT="${RECIPE_DIR}/agents/orchestrator.md" # Loop-dir convention matches recipes/perf-optimizer-loop/README.md : "${AI4S_SHARED_DIR:?AI4S_SHARED_DIR must be set (e.g. export AI4S_SHARED_DIR=/shared/\$USER)}" # Shared perf-tool location (omnistat venv, VictoriaMetrics). Cluster-specific: -# set from .cluster-config.yaml perf_tools.dir (e.g. /path/to/perf-tools). -: "${PERF_TOOLS_DIR:?PERF_TOOLS_DIR must be set (from .cluster-config.yaml perf_tools.dir, e.g. /path/to/perf-tools)}" +# set from .cluster-config.yaml omnihub.tools_dir (e.g. /shared/omnihub/tools). +: "${OMNIHUB_TOOLS_DIR:?OMNIHUB_TOOLS_DIR must be set (from .cluster-config.yaml omnihub.tools_dir, e.g. /shared/omnihub/tools)}" HG_BASE="${AI4S_SHARED_DIR}/models/HydraGNN" PERF_RUNS_DIR="${HG_BASE}/perf-runs" LOOP_DIR="${PERF_RUNS_DIR}/loop-${LOOP_UUID}" @@ -96,8 +96,8 @@ fi # 3. tools present for _bin in \ - "${PERF_TOOLS_DIR}/perf-inspect/bin/omnistat-usermode" \ - "${PERF_TOOLS_DIR}/victoriametrics/victoria-metrics-prod"; do + "${OMNIHUB_TOOLS_DIR}/omnihub-inspect/bin/omnistat-usermode" \ + "${OMNIHUB_TOOLS_DIR}/victoriametrics/victoria-metrics-prod"; do if [[ -z "$PREFLIGHT_FAIL_REASON" && ! -x "$_bin" ]]; then PREFLIGHT_FAIL_REASON="tool_missing" PREFLIGHT_NOTES+=("missing: $_bin (run the perf-analysis launcher subagent first to install)") diff --git a/material_science/models/HydraGNN/examples/sbatch_train_perf_amd.sh b/material_science/models/HydraGNN/examples/sbatch_train_perf_amd.sh index cbf5c0e..313c34f 100755 --- a/material_science/models/HydraGNN/examples/sbatch_train_perf_amd.sh +++ b/material_science/models/HydraGNN/examples/sbatch_train_perf_amd.sh @@ -21,8 +21,8 @@ # # Required: # AI4S_SHARED_DIR — shared base path -# ${PERF_TOOLS_DIR}/perf-inspect/ (created by the launcher subagent) -# ${PERF_TOOLS_DIR}/victoriametrics/ (created by the launcher subagent) +# ${OMNIHUB_TOOLS_DIR}/omnihub-inspect/ (created by the launcher subagent) +# ${OMNIHUB_TOOLS_DIR}/victoriametrics/ (created by the launcher subagent) # # Optional env-var overrides (with defaults): # HG_NUM_EPOCH=2 @@ -60,8 +60,8 @@ fi # --------------------------------------------------------------------------- HG_BASE="${AI4S_SHARED_DIR:?AI4S_SHARED_DIR must be set}/models/HydraGNN" # Shared perf-tool location (omnistat venv, omnistat-src, VictoriaMetrics). -# Cluster-specific: set from .cluster-config.yaml perf_tools.dir. -: "${PERF_TOOLS_DIR:?PERF_TOOLS_DIR must be set (from .cluster-config.yaml perf_tools.dir, e.g. /path/to/perf-tools)}" +# Cluster-specific: set from .cluster-config.yaml omnihub.tools_dir. +: "${OMNIHUB_TOOLS_DIR:?OMNIHUB_TOOLS_DIR must be set (from .cluster-config.yaml omnihub.tools_dir, e.g. /shared/omnihub/tools)}" HG_SIF="${HG_SIF:-${AI4S_SHARED_DIR}/images/pytorch_rocm7.2.2_ubuntu24.04_py3.12_pytorch_release_2.10.0.sif}" HG_OVERLAY="${HG_OVERLAY:-${HG_BASE}/overlays/hydragnn-overlay.img}" HG_DATASETS="${HG_DATASETS:-ANI1x,Alexandria}" @@ -74,7 +74,7 @@ HG_REPO_DIR="${HG_REPO_DIR:-${HG_BASE}/code/HydraGNN}" HYDRAGNN_MAX_NUM_BATCH="${HYDRAGNN_MAX_NUM_BATCH:-30}" PROFILE_TARGET_EPOCH="${PROFILE_TARGET_EPOCH:-1}" -OMNISTAT_VENV="${OMNISTAT_VENV:-${PERF_TOOLS_DIR}/perf-inspect}" +OMNISTAT_VENV="${OMNISTAT_VENV:-${OMNIHUB_TOOLS_DIR}/omnihub-inspect}" # Read the omnistat config from the in-repo recipe (single source of truth). # Override with OMNISTAT_TEMPLATE=/path/to/file if you need a custom probe config. # A staged copy under ${HG_BASE}/perf-runs/ is intentionally NOT used as the @@ -90,7 +90,7 @@ OMNISTAT_USERMODE_INTERVAL="${OMNISTAT_USERMODE_INTERVAL:-1}" # See material_science/models/HydraGNN/recipes/perf-analysis/agents/launcher.md # for how to build libomnistat_trace.so on a compute node. OMNISTAT_KERNEL_TRACE="${OMNISTAT_KERNEL_TRACE:-0}" -OMNISTAT_TRACE_LIB="${OMNISTAT_TRACE_LIB:-${PERF_TOOLS_DIR}/omnistat-src/build-trace/libomnistat_trace.so}" +OMNISTAT_TRACE_LIB="${OMNISTAT_TRACE_LIB:-${OMNIHUB_TOOLS_DIR}/omnistat-src/build-trace/libomnistat_trace.so}" # IMPORTANT: the kernel-trace collector registers /kernel_trace on the SAME # Flask app as the other omnistat endpoints, so the tool library must POST to # the [omnistat.collectors] `port` (8101 here), NOT the library's own default @@ -165,7 +165,7 @@ done mkdir -p "$HG_OUTPUT_DIR" OMNISTAT_CONFIG="${HG_OUTPUT_DIR}/omnistat.config" sed -e "s|@JOB_DIR@|${HG_OUTPUT_DIR}|g" \ - -e "s|@PERF_TOOLS_DIR@|${PERF_TOOLS_DIR}|g" \ + -e "s|@OMNIHUB_TOOLS_DIR@|${OMNIHUB_TOOLS_DIR}|g" \ "$OMNISTAT_TEMPLATE" > "$OMNISTAT_CONFIG" # Opt-in kernel tracing: flip enable_kernel_trace=True in the rendered config @@ -180,7 +180,7 @@ if [[ "$OMNISTAT_KERNEL_TRACE" == "1" ]]; then echo " Build it on a compute node:" >&2 echo " salloc -p -A -N 1 --time=00:15:00 --gpus-per-node=1 \\" >&2 echo " apptainer exec --rocm \"\$HG_SIF\" bash -c '" >&2 - echo " cd ${PERF_TOOLS_DIR}/omnistat-src && \\" >&2 + echo " cd ${OMNIHUB_TOOLS_DIR}/omnistat-src && \\" >&2 echo " cmake -S rocprofiler-sdk/ -B build-trace/ -DBUILD_KERNEL_TRACE_LIB=ON && \\" >&2 echo " cmake --build build-trace/ -j 8'" >&2 exit 2 diff --git a/material_science/models/HydraGNN/recipes/perf-analysis/README.md b/material_science/models/HydraGNN/recipes/perf-analysis/README.md index 0d0405e..efa3aa0 100644 --- a/material_science/models/HydraGNN/recipes/perf-analysis/README.md +++ b/material_science/models/HydraGNN/recipes/perf-analysis/README.md @@ -68,8 +68,8 @@ The `combined_report.md` is the deliverable: ranked bottlenecks, remedies tried ## Prerequisites - Cluster access with the standard HydraGNN setup at `$AI4S_SHARED_DIR/models/HydraGNN/` — overlay, weights, code clone. See [HydraGNN/recipes/train/](../train/). -- Export `PERF_TOOLS_DIR` (the shared perf-tool location, from `.cluster-config.yaml` `perf_tools.dir`, e.g. `/path/to/perf-tools`). Scripts and subagents resolve all tool paths from this var — no cluster path is hardcoded. -- One-time install of Omnistat (PR #271 branch `jorda/skills`, with `origin/main` merged in), VictoriaMetrics, and TraceLens — performed lazily by the launcher subagent into `${PERF_TOOLS_DIR}/`. +- Export `OMNIHUB_TOOLS_DIR` (the shared perf-tool location, from `.cluster-config.yaml` `omnihub.tools_dir`, e.g. `/shared/omnihub/tools`). Scripts and subagents resolve all tool paths from this var — no cluster path is hardcoded. +- One-time install of Omnistat (PR #271 branch `jorda/skills`, with `origin/main` merged in), VictoriaMetrics, and TraceLens — performed lazily by the launcher subagent into `${OMNIHUB_TOOLS_DIR}/`. - The launcher writes `gfm_mlip_with_profile.json` at submit time (does not modify the upstream `gfm_mlip.json`). ## Telemetry knobs (set at `sbatch` submit time) @@ -77,7 +77,7 @@ The `combined_report.md` is the deliverable: ranked bottlenecks, remedies tried | Env var | Default | What it does | |---|---|---| | `OMNISTAT_KERNEL_TRACE` | `0` | When `1`, loads `libomnistat_trace.so` via `ROCP_TOOL_LIBRARIES` on every rank and turns on the omnistat kernel-trace collector. Adds per-kernel dispatch count + duration time series across all 8 cards on every node. Requires the library to be built once (see `agents/launcher.md` step 1e). Validated end-to-end on job 7034. | -| `OMNISTAT_TRACE_LIB` | `${PERF_TOOLS_DIR}/omnistat-src/build-trace/libomnistat_trace.so` | Override only for development. The wrapper hard-fails if the path is missing when `OMNISTAT_KERNEL_TRACE=1`. | +| `OMNISTAT_TRACE_LIB` | `${OMNIHUB_TOOLS_DIR}/omnistat-src/build-trace/libomnistat_trace.so` | Override only for development. The wrapper hard-fails if the path is missing when `OMNISTAT_KERNEL_TRACE=1`. | | `OMNISTAT_TRACE_LOG` | `1` | Library prints `[host][pid][omnistat] Trace summary: N/N processed records (M/M successful flushes)` on rank exit; useful as a smoke-signal that the tool initialized even when the workload crashed. | `enable_rocprofiler=True` is **on** in `omnistat.config.template` by default; the rendered config's state is printed in the sbatch banner so a silent "counters off" cannot recur. Kernel tracing and device-counting can co-exist on a single GPU but raise per-job VictoriaMetrics cardinality — pick by run length and the question you're asking. See `ai4science-studio` SKILL §12 for the device-counting vs kernel-trace vs `rocprofv3` decision matrix. diff --git a/material_science/models/HydraGNN/recipes/perf-analysis/agents/launcher.md b/material_science/models/HydraGNN/recipes/perf-analysis/agents/launcher.md index 20888c5..899f6e5 100644 --- a/material_science/models/HydraGNN/recipes/perf-analysis/agents/launcher.md +++ b/material_science/models/HydraGNN/recipes/perf-analysis/agents/launcher.md @@ -10,8 +10,8 @@ Submits a 2-node HydraGNN training (AMD Instinct MI355X) with PyTorch profiling ## Outputs -- `${PERF_TOOLS_DIR}/perf-inspect/` — Python venv with omnistat (PR #271 + main merged) and TraceLens. -- `${PERF_TOOLS_DIR}/victoriametrics/victoria-metrics-prod` — VictoriaMetrics binary. +- `${OMNIHUB_TOOLS_DIR}/omnihub-inspect/` — Python venv with omnistat (PR #271 + main merged) and TraceLens. +- `${OMNIHUB_TOOLS_DIR}/victoriametrics/victoria-metrics-prod` — VictoriaMetrics binary. - `$AI4S_SHARED_DIR/models/HydraGNN/perf-runs/omnistat.config` — Omnistat user-mode config. - `$AI4S_SHARED_DIR/models/HydraGNN/perf-runs//manifest.json` — manifest schema below. - `$AI4S_SHARED_DIR/models/HydraGNN/perf-runs//hydragnn-train-.out` — symlinked from output dir. @@ -57,8 +57,8 @@ Submits a 2-node HydraGNN training (AMD Instinct MI355X) with PyTorch profiling ```bash set -euo pipefail -# Shared tool location — read from .cluster-config.yaml perf_tools.dir. -TOOLS="${PERF_TOOLS_DIR:?set PERF_TOOLS_DIR from .cluster-config.yaml perf_tools.dir (e.g. /path/to/perf-tools)}" +# Shared tool location — read from .cluster-config.yaml omnihub.tools_dir. +TOOLS="${OMNIHUB_TOOLS_DIR:?set OMNIHUB_TOOLS_DIR from .cluster-config.yaml omnihub.tools_dir (e.g. /shared/omnihub/tools)}" mkdir -p "$TOOLS" # 1a. Omnistat: jorda/skills branch with origin/main merged in @@ -77,7 +77,7 @@ fi OMNISTAT_COMMIT=$(git rev-parse HEAD) # 1b. Venv (py3.12 to match cluster python) -VENV=${PERF_TOOLS_DIR}/perf-inspect +VENV=${OMNIHUB_TOOLS_DIR}/omnihub-inspect if [[ ! -d $VENV ]]; then python3 -m venv "$VENV" fi @@ -123,11 +123,11 @@ if [[ "${OMNISTAT_KERNEL_TRACE:-0}" == "1" && ! -f "$TRACE_LIB" ]]; then fi ``` -The sbatch wrapper expects the `.so` at `${PERF_TOOLS_DIR}/omnistat-src/build-trace/libomnistat_trace.so` by default; override with `OMNISTAT_TRACE_LIB=/path/to/lib.so`. See `ai4science-studio` SKILL §12 for when to enable kernel tracing vs the default device-counting collector. +The sbatch wrapper expects the `.so` at `${OMNIHUB_TOOLS_DIR}/omnistat-src/build-trace/libomnistat_trace.so` by default; override with `OMNISTAT_TRACE_LIB=/path/to/lib.so`. See `ai4science-studio` SKILL §12 for when to enable kernel tracing vs the default device-counting collector. ### 2. Author the Omnistat user-mode config (once, in repo `perf-runs/`) -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@` / `@PERF_TOOLS_DIR@` placeholders, which the sbatch wrapper substitutes (via `sed`) into a per-job config under the perf-run directory. +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 diff --git a/material_science/models/HydraGNN/recipes/perf-analysis/agents/omnistat_analyst.md b/material_science/models/HydraGNN/recipes/perf-analysis/agents/omnistat_analyst.md index 5ff92cf..4260912 100644 --- a/material_science/models/HydraGNN/recipes/perf-analysis/agents/omnistat_analyst.md +++ b/material_science/models/HydraGNN/recipes/perf-analysis/agents/omnistat_analyst.md @@ -6,8 +6,8 @@ Drive `omnistat-inspect` (PR #271) through the analyze-job phases on the user-mo - `/manifest.json` - The Omnistat DB at `manifest.omnistat_db_path`. -- The `perf-inspect` venv at `${PERF_TOOLS_DIR}/perf-inspect/`. -- The VictoriaMetrics binary at `${PERF_TOOLS_DIR}/victoriametrics/victoria-metrics-prod`. +- The `omnihub-inspect` venv at `${OMNIHUB_TOOLS_DIR}/omnihub-inspect/`. +- The VictoriaMetrics binary at `${OMNIHUB_TOOLS_DIR}/victoriametrics/victoria-metrics-prod`. ## Outputs @@ -36,7 +36,7 @@ ss -lnt "sport = :$PORT" 2>/dev/null | grep -q LISTEN && PORT=8429 VMLOG=$PERF_RUN/omnistat/vm.log mkdir -p "$PERF_RUN/omnistat/inspect_outputs" -nohup ${PERF_TOOLS_DIR}/victoriametrics/victoria-metrics-prod \ +nohup ${OMNIHUB_TOOLS_DIR}/victoriametrics/victoria-metrics-prod \ -storageDataPath="$DB" \ -httpListenAddr=127.0.0.1:$PORT \ -retentionPeriod=100y \ @@ -62,7 +62,7 @@ done ### 2. Run analyze-job phases (per the SKILL) ```bash -. ${PERF_TOOLS_DIR}/perf-inspect/bin/activate +. ${OMNIHUB_TOOLS_DIR}/omnihub-inspect/bin/activate SCRATCH=$PERF_RUN/omnistat/scratch mkdir -p "$SCRATCH" JOBID=$(jq -r .jobid ) @@ -140,4 +140,4 @@ The orchestrator (or the verifier, when finished) is responsible for stopping VM - 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. +- Never query `atlvrmonad01:9090` — it's IP-blocked and not the right database for this job anyway. \ No newline at end of file diff --git a/material_science/models/HydraGNN/recipes/perf-analysis/agents/tracelens_analyst.md b/material_science/models/HydraGNN/recipes/perf-analysis/agents/tracelens_analyst.md index 99beea4..2b1dd3c 100644 --- a/material_science/models/HydraGNN/recipes/perf-analysis/agents/tracelens_analyst.md +++ b/material_science/models/HydraGNN/recipes/perf-analysis/agents/tracelens_analyst.md @@ -39,7 +39,7 @@ Every entry: ### 1. Sanity check ```bash -. ${PERF_TOOLS_DIR}/perf-inspect/bin/activate +. ${OMNIHUB_TOOLS_DIR}/omnihub-inspect/bin/activate PERF_RUN=$(jq -r .perf_run_dir ) TRACE=$(jq -r '.trace_paths[0] // empty' ) [[ -z "$TRACE" ]] && { echo "STATUS=partial; reason=no_trace_in_manifest"; exit 0; } @@ -121,4 +121,4 @@ STATUS=ok; reason= claims, top=: - 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. +- Keep `evidence_paths` machine-parseable so the verifier can re-derive the number from the same source. \ No newline at end of file diff --git a/material_science/models/HydraGNN/recipes/perf-analysis/omnistat.config.template b/material_science/models/HydraGNN/recipes/perf-analysis/omnistat.config.template index 9d55579..ab32dee 100644 --- a/material_science/models/HydraGNN/recipes/perf-analysis/omnistat.config.template +++ b/material_science/models/HydraGNN/recipes/perf-analysis/omnistat.config.template @@ -122,7 +122,7 @@ system_name = mi355x #-- [omnistat.usermode] -victoria_binary = @PERF_TOOLS_DIR@/victoriametrics/victoria-metrics-prod +victoria_binary = @OMNIHUB_TOOLS_DIR@/victoriametrics/victoria-metrics-prod victoria_datadir = @JOB_DIR@/omnistat-db victoria_logfile = vic_server.log diff --git a/material_science/models/HydraGNN/recipes/perf-optimizer-loop/agents/fom_extractor.md b/material_science/models/HydraGNN/recipes/perf-optimizer-loop/agents/fom_extractor.md index 9d97925..455b6d5 100644 --- a/material_science/models/HydraGNN/recipes/perf-optimizer-loop/agents/fom_extractor.md +++ b/material_science/models/HydraGNN/recipes/perf-optimizer-loop/agents/fom_extractor.md @@ -68,7 +68,7 @@ omnistat kernel-trace collector). ```bash set -euo pipefail -. "${PERF_TOOLS_DIR}/perf-inspect/bin/activate" +. "${OMNIHUB_TOOLS_DIR}/omnihub-inspect/bin/activate" PERF_RUN=$(jq -r .perf_run_dir ) JOBID=$(jq -r .jobid ) RUNTIME=$(jq -r .runtime_seconds ) @@ -102,7 +102,7 @@ Start a local VictoriaMetrics if not already running: PORT=8428 ss -lnt "sport = :$PORT" 2>/dev/null | grep -q LISTEN && PORT=8429 if ! curl -sf "http://127.0.0.1:$PORT/api/v1/status/tsdb" > /dev/null 2>&1; then - nohup "${PERF_TOOLS_DIR}/victoriametrics/victoria-metrics-prod" \ + nohup "${OMNIHUB_TOOLS_DIR}/victoriametrics/victoria-metrics-prod" \ -storageDataPath="$PERF_RUN/omnistat-db" \ -httpListenAddr=127.0.0.1:$PORT \ -retentionPeriod=100y -fs.disableMmap \ @@ -255,4 +255,4 @@ The existing `examples/parse_convergence.py` is extended with `--json-foms /foms.csv` if it exists. If iter rows are present, set `current Run these checks in sequence; abort with `PREFLIGHT_FAIL reason=` on the first failure: -- Ensure `AI4S_SHARED_DIR` and `PERF_TOOLS_DIR` are exported (the latter from `.cluster-config.yaml` `perf_tools.dir`, e.g. `/path/to/perf-tools`); abort if either is unset. All tool paths below derive from `$PERF_TOOLS_DIR` — never hardcode a cluster path. +- Ensure `AI4S_SHARED_DIR` and `OMNIHUB_TOOLS_DIR` are exported (the latter from `.cluster-config.yaml` `omnihub.tools_dir`, e.g. `/shared/omnihub/tools`); abort if either is unset. All tool paths below derive from `$OMNIHUB_TOOLS_DIR` — never hardcode a cluster path. - `sinfo -p -h` (partition from `.cluster-config.yaml`) → if zero `IDLE` or `MIX` nodes, abort. Cluster may be in maintenance (see SKILL §14). - `df -h "$AI4S_SHARED_DIR" | tail -1` → abort if `Use%` > 95. -- `test -x ${PERF_TOOLS_DIR}/perf-inspect/bin/omnistat-usermode` → abort if missing. -- `test -x ${PERF_TOOLS_DIR}/victoriametrics/victoria-metrics-prod` → abort if missing. +- `test -x ${OMNIHUB_TOOLS_DIR}/omnihub-inspect/bin/omnistat-usermode` → abort if missing. +- `test -x ${OMNIHUB_TOOLS_DIR}/victoriametrics/victoria-metrics-prod` → abort if missing. - `test -f $AI4S_SHARED_DIR/models/HydraGNN/overlays/hydragnn-overlay.img` → abort if missing. - `test -f $AI4S_SHARED_DIR/images/pytorch_rocm7.2.2_ubuntu24.04_py3.12_pytorch_release_2.10.0.sif` → abort if missing. - If invoked via Claude Code CLI on the login node: `[[ -n "$ANTHROPIC_API_KEY" ]]` and `curl -fsS https://api.anthropic.com/v1/models > /dev/null` (5 s timeout) → abort if either fails. diff --git a/material_science/models/HydraGNN/recipes/perf-optimizer-loop/agents/story_writer.md b/material_science/models/HydraGNN/recipes/perf-optimizer-loop/agents/story_writer.md index f296ca4..74af59d 100644 --- a/material_science/models/HydraGNN/recipes/perf-optimizer-loop/agents/story_writer.md +++ b/material_science/models/HydraGNN/recipes/perf-optimizer-loop/agents/story_writer.md @@ -156,7 +156,7 @@ fig.tight_layout() fig.savefig(out_png_path, dpi=150, bbox_inches="tight") ``` -Use `matplotlib.use("Agg")` for headless rendering. Do NOT use seaborn (extra dep not in the omnistat venv). Use the omnistat venv at `${PERF_TOOLS_DIR}/perf-inspect/bin/python` which already has matplotlib via TraceLens deps. +Use `matplotlib.use("Agg")` for headless rendering. Do NOT use seaborn (extra dep not in the omnistat venv). Use the omnistat venv at `${OMNIHUB_TOOLS_DIR}/omnihub-inspect/bin/python` which already has matplotlib via TraceLens deps. ## Style rules @@ -179,4 +179,4 @@ Use `matplotlib.use("Agg")` for headless rendering. Do NOT use seaborn (extra de - 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. +- 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. \ No newline at end of file diff --git a/material_science/models/HydraGNN/recipes/perf-optimizer-loop/dispatch-attribution.md b/material_science/models/HydraGNN/recipes/perf-optimizer-loop/dispatch-attribution.md index ebedaf7..08121fb 100644 --- a/material_science/models/HydraGNN/recipes/perf-optimizer-loop/dispatch-attribution.md +++ b/material_science/models/HydraGNN/recipes/perf-optimizer-loop/dispatch-attribution.md @@ -59,10 +59,10 @@ Short-kernel study (`short_kernels_summary.csv`): thousands of `aten::fill_`, `a ```bash export AI4S_SHARED_DIR=/shared/$USER export REPO_ROOT=$HOME/git/ai4science-studio -source ${PERF_TOOLS_DIR}/perf-inspect/bin/activate +source ${OMNIHUB_TOOLS_DIR}/omnihub-inspect/bin/activate # VictoriaMetrics on the perf-run DB (login node) -${PERF_TOOLS_DIR}/victoriametrics/victoria-metrics-prod \ +${OMNIHUB_TOOLS_DIR}/victoriametrics/victoria-metrics-prod \ -storageDataPath=$AI4S_SHARED_DIR/models/HydraGNN/perf-runs/7187/omnistat-db \ -httpListenAddr=127.0.0.1:8432 -retentionPeriod=100y -fs.disableMmap & @@ -75,4 +75,4 @@ TraceLens_generate_perf_report_pytorch \ 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 -``` +``` \ No newline at end of file