AMDResearch · ashwinma · Jul 14, 2026 · Jul 13, 2026
diff --git a/.cursor/skills/ai4science-earth-science/SKILL.md b/.cursor/skills/ai4science-earth-science/SKILL.md
@@ -56,6 +56,7 @@ The `earth_science/` domain covers **climate**, **weather**, and broader **Earth
 - **Training landmines (2026-06):** stub `gptl4py` in `run_orbit2_train.py`; overlay may lack `xformers.ops` → **`ORBIT2_FUSED_ATTN=DEFAULT`** in sbatch. **`ORBIT2_DATA_TYPE`** defaults **`bfloat16`** (matches upstream **`configs/interm_8m.yaml`** and Bayes-CAST-style **`edm_8m_era5.yaml`**); use **`float32`** only to isolate ROCm/hipBLAS issues. `max_epochs` must be ≥ 2; batch cap patches module before `main()` not via `runpy.run_path`.
 - **bf16 + ROCm Flash SDPA 65535 grid cap (2026-06, FIXED):** Bayes-CAST EDM `var_agg`/`temporal_agg` cross-attention flattens `(B,History,L)` into the SDPA **batch** dim (`B = batch·648`), and ROCm **Flash** SDPA caps that dim at **65535** → bf16 crashes with `HIP error: invalid argument` at **batch ≥ ~128** (batch 64 OK). Fix applied **globally** to the compute-side clone: all `Attention` + `CrossAttention` SDPA calls wrapped in `sdpa_kernel([SDPBackend.EFFICIENT_ATTENTION, SDPBackend.MATH])` (MATH fallback-only). Going global is safe because `--orbit-selfattn` benchmarks EFFICIENT within ±3–5% of Flash on the main 648-token self-attention. Validated job 9238 (bf16 b256, decreasing loss). **Generalizes** to any model folding extra dims into the attention batch.
 - **Perf/scaling:** `sbatch_train_perf_amd.sh` defaults **Bayes-CAST** template **`edm_8m_era5_1x8.yaml`** + ERA5 1.0° data; **`ORBIT2_ROOT`** prefers **`…/code/bayes-cast`** when present. Bayes EDM clears **`ORBIT2_RANK_PRE_TRAIN_HOOK`** only when unset — export an absolute hook path **before** `sbatch` for sysopt (`torch.compile`, etc.). **`ORBIT2_MAX_EPOCH=6`** default. See **[one-node GPU baseline](../../../earth_science/models/ORBIT-2/recipes/perf-analysis/one-node-gpu-baseline.md)**, **[perf-optimizer-loop](../../../earth_science/models/ORBIT-2/recipes/perf-optimizer-loop/README.md)** (throughput-primary loop), and **`ORBIT2_NUM_WORKERS`** / template token **`__NUM_WORKERS__`** in `edm_8m_era5_1x8.yaml` via `render_orbit2_config.py`.
+- **Scaling sweep = WEAK scaling + HSDP (2026-06):** `examples/run_scaling_study.sh` (1/2/4/8 nodes) holds **per-rank batch fixed** so each GPU stays equally filled at every node count (global batch = `batch x 8 x N` grows) — this is **weak** scaling; `collate_scaling_study.py`'s "efficiency" is a comm-overlap / RCCL-bandwidth metric (`t_1node/t_Nnode`), not a classic strong-scaling speedup. (HydraGNN's recipe labels the same fixed-per-rank setup "strong-scaling efficiency" — mechanically it is weak.) Parallelism is **HSDP**: set **`ORBIT2_FSDP=8`** (shard within a node over XGMI) + **`ORBIT2_SIMPLE_DDP=N`** (data-parallel replica per node over IB/ANP); the script does this per node count via `ORBIT2_SCALING_FSDP`. **Do NOT** fall back to the render default `fsdp=N, simple_ddp=8` for multi-node — it shards weights across the slow inter-node fabric. EDM compute-saturated baseline: `edm_8m_era5_1x8.yaml`, bf16, per-rank batch **1024** (~4098 samples/s, ~21% HBM on the ~1704-sample ERA5 1.0° tree; reaching ~80% HBM needs more staged ERA5 — see STAGING_ERA5_FOR_HBM.md). **Pitfall — data cap inflates weak-scaling efficiency >1 at high node counts (validated 2026-06, jobs 14826-9):** with fixed per-rank batch the global batch grows as `1024 x 8 x N`, so on a small ERA5 tree the steps/epoch collapse (e.g. 57→27→**9**→9 steady batches at 1/2/4/8 nodes); steady step time over ~3 batches/epoch is noise-dominated and `throughput = global_batch/step_time` reports **fake super-linear efficiency (1.3x)** while the loss actually gets *worse* (too few steps to train). Only trust node counts where `steady_batch_count` stays large (here 1→2 nodes: eff ≈ 1.01, the real clean result). Fix = stage more ERA5 (Phase 2), not a code change. Always check the `steady_batch_count`/steps-per-epoch before quoting multi-node throughput.
 - **Perf saturation vs ~10M params:** Default `interm_8m_prism.yaml`/`interm_8m_era5.yaml` uses `res_slimvit` @ `embed_dim` 256 / `depth` 6 → ~10M weights; **PRISM `spatial_resolution` 18** yields tiny activations per step — raise **batch**, use **ERA5 111×111** template for heavier forwards, or widen **`embed_dim`/`depth`** (with `embed_dim % num_heads == 0`) to approach MI-class MFMA; see [one-node-gpu-baseline.md](../../../earth_science/models/ORBIT-2/recipes/perf-analysis/one-node-gpu-baseline.md) §3. **VRAM-first sweep:** [BASELINE_LOCKIN.md](../../../earth_science/models/ORBIT-2/recipes/perf-analysis/BASELINE_LOCKIN.md) (batch sweep before widening ViT).
 - **Multi-node scheduling:** Before large submits, check partition state (`sinfo`, site dashboards). Prefer **idle, healthy** nodes; use `--nodelist` / `--exclude` only per your site's policy. Do not stack two exclusive full-node GPU jobs on the same host.
 - **GEMM-time bottleneck = the 3/4-channel conv stem/head, NOT attention/MLP (2026-06, MI355X):** Full TraceLens+Omnistat analyst/verifier run (`recipes/perf-analysis/agents/orchestrator_gemm_analysis.md` + `examples/run_gemm_analysis.sh`) on profiled EDM (bf16, batch 4096, ERA5) found one hipBLASLt GEMM `Cijk_…MT16x16x32` = **~43 % of the step at both 1- and 2-node** — the **im2col-lowered low-channel conv backward** (3/4-channel projections, 16×16 tile @ ~0.03 TFLOP/s). The tunable `nn.Linear` GEMMs are ~2 % and near-peak → **this is why TunableOp gave no uplift.** **Most promising lever = channel-padding in NCHW (timing-only −26 %, but TABLED pending a convergence study — see next bullet); `channels_last` is a DEAD END on ROCm 7.2.2 (also next bullet).** **Profile any EDM run** via `train_edm.py:_orbit2_make_profiler` (env-gated: set `ORBIT2_PROFILE_DIR`+`PROFILE_TARGET_EPOCH`; sbatch does this). **Trust raw kernel-name device-time, not TraceLens `report.xlsx:ops_summary_by_category`** — it mislabels conv-lowered GEMMs as `CONV_bwd` (a 2-node analyst claim was refuted on this). `examples/compare_trace_kernels.py` aggregates kernels tool-independently. Separately, **~43 % whole-job idle** (inter-step overhead, NOT comm-bound) is a second lever.

diff --git a/earth_science/models/ORBIT-2/examples/run_scaling_study.sh b/earth_science/models/ORBIT-2/examples/run_scaling_study.sh
@@ -1,17 +1,31 @@
 #!/usr/bin/env bash
-# Submit matched ORBIT-2 strong-scaling sweep (1/2/4/8 nodes).
+# Submit a matched ORBIT-2 WEAK-scaling sweep (1/2/4/8 nodes).
 #
-# Usage:
+# Weak scaling: per-rank batch is held FIXED at every node count (each GPU does the same
+# work; global batch grows as batch x 8 x N). Ideal result is constant steady step time vs N;
+# the efficiency reported by collate_scaling_study.py is then a comm-overlap (RCCL bandwidth)
+# metric, not a classic strong-scaling speedup. Parallelism is HSDP: fsdp=8 within a node
+# (XGMI), simple_ddp=N across nodes (IB/ANP) — see ORBIT2_SCALING_FSDP below.
+#
+# Usage (PRISM res_slimvit, default):
 #   export AI4S_SHARED_DIR=/path/to/shared
 #   export ORBIT2_DATA_ROOT=$AI4S_SHARED_DIR/models/ORBIT-2/data/superres/prism/10.0_arcmin
 #   ./run_scaling_study.sh
 #   ./run_scaling_study.sh --nodes 1,2
 #
-# ERA5 1.0_deg same-dir (timing only):
+# ERA5 1.0_deg same-dir res_slimvit (timing only):
 #   export ORBIT2_DATA_ROOT=$AI4S_SHARED_DIR/models/ORBIT-2/data/superres/era5/1.0_deg
 #   export ORBIT2_CONFIG_TEMPLATE=interm_8m_era5.yaml
 #   export ORBIT2_SCALING_TAG=era5
 #   ./run_scaling_study.sh
+#
+# Bayes-CAST EDM, compute-saturated (per-rank batch fills the GPU at every node count):
+#   export ORBIT2_ROOT=$AI4S_SHARED_DIR/models/ORBIT-2/code/bayes-cast
+#   export ORBIT2_CONFIG_TEMPLATE=edm_8m_era5_1x8.yaml
+#   export ORBIT2_DATA_ROOT=$AI4S_SHARED_DIR/models/ORBIT-2/data/superres/era5/1.0_deg
+#   export ORBIT2_ERA5_SPATIAL_RES=111 ORBIT2_DATA_TYPE=bfloat16 ORBIT2_FUSED_ATTN=DEFAULT
+#   export ORBIT2_BATCH_SIZE=1024 ORBIT2_SCALING_TAG=edm-era5
+#   ./run_scaling_study.sh
 
 set -euo pipefail
 
@@ -22,12 +36,21 @@ export TORCH_NCCL_HIGH_PRIORITY=1
 export GPU_MAX_HW_QUEUES=2
 export ORBIT2_DATA_TYPE="${ORBIT2_DATA_TYPE:-bfloat16}"
 # max_epochs=6 → trains epochs 0–4; collate uses steady epochs 2–4 for FOM
-export ORBIT2_MAX_EPOCH=6
-export ORBIT2_MAX_BATCHES=20
-export ORBIT2_BATCH_SIZE=4
+# All sizing knobs honor caller overrides (EDM weak-scaling uses larger batch).
+export ORBIT2_MAX_EPOCH="${ORBIT2_MAX_EPOCH:-6}"
+export ORBIT2_MAX_BATCHES="${ORBIT2_MAX_BATCHES:-20}"
+export ORBIT2_BATCH_SIZE="${ORBIT2_BATCH_SIZE:-4}"
 export ORBIT2_DATA_ROOT="${ORBIT2_DATA_ROOT:-${AI4S_SHARED_DIR}/models/ORBIT-2/data/superres/prism/10.0_arcmin}"
 export ORBIT2_CONFIG_TEMPLATE="${ORBIT2_CONFIG_TEMPLATE:-interm_8m_prism.yaml}"
-# sbatch_train_amd.sh: 1-node×8-GPU jobs default to fsdp=8 simple_ddp=1; N>1 uses fsdp=N simple_ddp=8 from render defaults.
+# Pass-through for the EDM (Bayes-CAST) path; harmless when unset for the PRISM/ERA5 res_slimvit path.
+[[ -n "${ORBIT2_ROOT:-}" ]] && export ORBIT2_ROOT
+[[ -n "${ORBIT2_ERA5_SPATIAL_RES:-}" ]] && export ORBIT2_ERA5_SPATIAL_RES
+export ORBIT2_FUSED_ATTN="${ORBIT2_FUSED_ATTN:-DEFAULT}"
+# Weak-scaling parallelism: keep FSDP sharding *within* a node (fast XGMI) and add a
+# data-parallel replica *per node* across the IB fabric → per-rank work stays fixed at
+# every node count (fsdp=8 simple_ddp=N), instead of the render default (fsdp=N simple_ddp=8)
+# which would shard weights across slow inter-node links. Override fsdp/node via ORBIT2_SCALING_FSDP.
+export ORBIT2_SCALING_FSDP="${ORBIT2_SCALING_FSDP:-8}"
 export ORBIT2_SCALING_TAG="${ORBIT2_SCALING_TAG:-prism}"
 export ORBIT2_OUTPUT_DIR="${ORBIT2_OUTPUT_DIR:-${AI4S_SHARED_DIR}/models/ORBIT-2/outputs/scaling-${ORBIT2_SCALING_TAG}}"
 
@@ -47,6 +70,9 @@ mkdir -p "$LOG_DIR"
 
 submit() {
   local n="$1" t="$2"
+  # Weak-scaling HSDP: fsdp = GPUs/node (intra-node), simple_ddp = node count (inter-node).
+  export ORBIT2_FSDP="$ORBIT2_SCALING_FSDP"
+  export ORBIT2_SIMPLE_DDP="$n"
   sbatch --parsable \
     --partition="$SBATCH_PARTITION" \
     --account="$SBATCH_ACCOUNT" \
@@ -55,7 +81,7 @@ submit() {
     --job-name="orbit2-scale-${n}N" \
     --output="${LOG_DIR}/orbit2-train-%j.out" \
     --error="${LOG_DIR}/orbit2-train-%j.out" \
-    --export=ALL,TORCH_NCCL_HIGH_PRIORITY,GPU_MAX_HW_QUEUES,ORBIT2_DATA_TYPE,ORBIT2_MAX_EPOCH,ORBIT2_MAX_BATCHES,ORBIT2_BATCH_SIZE,ORBIT2_DATA_ROOT,ORBIT2_CONFIG_TEMPLATE,ORBIT2_SCALING_TAG,ORBIT2_OUTPUT_DIR,AI4S_SHARED_DIR,ORBIT2_DISABLE_CKPT=1 \
+    --export=ALL,TORCH_NCCL_HIGH_PRIORITY,GPU_MAX_HW_QUEUES,ORBIT2_DATA_TYPE,ORBIT2_MAX_EPOCH,ORBIT2_MAX_BATCHES,ORBIT2_BATCH_SIZE,ORBIT2_DATA_ROOT,ORBIT2_CONFIG_TEMPLATE,ORBIT2_SCALING_TAG,ORBIT2_OUTPUT_DIR,AI4S_SHARED_DIR,ORBIT2_ROOT,ORBIT2_ERA5_SPATIAL_RES,ORBIT2_FUSED_ATTN,ORBIT2_FSDP,ORBIT2_SIMPLE_DDP,ORBIT2_DISABLE_CKPT=1 \
     "$SCRIPT_DIR/sbatch_train_amd.sh"
 }
 
@@ -65,7 +91,10 @@ echo "=== ORBIT-2 scaling sweep (${ORBIT2_SCALING_TAG}) ==="
 echo "  Data root    : $ORBIT2_DATA_ROOT"
 echo "  Config       : $ORBIT2_CONFIG_TEMPLATE"
 echo "  Output base  : $ORBIT2_OUTPUT_DIR"
+echo "  Per-rank batch: $ORBIT2_BATCH_SIZE (fixed → weak scaling)  dtype: $ORBIT2_DATA_TYPE"
+echo "  Parallelism  : fsdp=$ORBIT2_SCALING_FSDP/node, simple_ddp=N (HSDP)"
 echo "  Epochs: $ORBIT2_MAX_EPOCH  Max batches: $ORBIT2_MAX_BATCHES"
+[[ -n "${ORBIT2_ROOT:-}" ]] && echo "  ORBIT2_ROOT  : $ORBIT2_ROOT"
 echo ""
 
 JOB_IDS=()

diff --git a/earth_science/models/ORBIT-2/recipes/train/README.md b/earth_science/models/ORBIT-2/recipes/train/README.md
@@ -16,6 +16,23 @@ This recipe summarizes the **exascale-oriented** workflow documented in [`XiaoWa
 5. **Data paths** — Populate `low_res_dir`, `high_res_dir`, `spatial_resolution`, and variable dictionaries to point at **your** staged dataset. Paths on Orion are **project- and allocation-specific**; see [../data/README.md](../data/README.md).
 6. **Submit training** — Edit `examples/launch_intermediate.sh` (`#SBATCH -A`, conda activation, config path), then `sbatch launch_intermediate.sh`. Monitor `flash-{JOBID}.out` as documented upstream.
 
+## Weak-scaling results — EDM, ERA5 1.0°, MI355X
+
+The table below is a **weak-scaling** sweep of the Bayes-CAST **EDM** model on AMD **MI355X** nodes (8 GPU/node). It is **weak** scaling: the **per-rank batch is fixed at 1024**, so the global batch grows with GPU count and per-GPU work stays constant. Parallelism is **HSDP** (`fsdp=8` within a node, `simple_ddp=N` across nodes), precision **bf16**, dataset **ERA5 1.0°** (`spatial_resolution=111`). "Steady s/batch" is the mean batch wall time over epochs 2+ (warmup epochs 0–1 and the first batch of each epoch excluded). **Efficiency = t₁ₙₒ𝒹ₑ / t_Nₙₒ𝒹ₑ** — a weak-scaling metric that reflects communication-overlap quality / RCCL bandwidth (≈1.0 is ideal).
+
+| Nodes | GPUs | per-rank batch | global batch | steady s/batch | samples/s | efficiency | steady batches | trust |
+|---:|---:|---:|---:|---:|---:|---:|---:|:---:|
+| 1 | 8  | 1024 | 8192  | 2.639 | 3103.7  | 1.00 | 57 | ✅ |
+| 2 | 16 | 1024 | 16384 | 2.626 | 6240.2  | 1.01 | 27 | ✅ |
+| 4 | 32 | 1024 | 32768 | 2.021 | 16212.7 | 1.31 | 9  | ⚠️ |
+| 8 | 64 | 1024 | 65536 | 2.060 | 31820.6 | 1.28 | 9  | ⚠️ |
+
+> **Read this before quoting the table.** Clean weak scaling is demonstrated **only at 1→2 nodes**: steady step time is essentially flat (2.639 s → 2.626 s, efficiency ≈ 1.01), confirming RCCL/IB communication is fully overlapped. The **4- and 8-node efficiencies (>1.0) are a measurement artifact, NOT super-linear speedup.** At per-rank batch 1024 the global batch reaches 32768 / 65536, which is large relative to the staged ERA5 1.0° tree, so only **~3 batches/epoch** run at 4N/8N (`steady batches` column collapses from 57 → 9). Steady step time over so few batches is noise-dominated, and `throughput = global_batch / step_time` inflates mechanically. The loss curves corroborate this: 4N/8N final loss (0.90 / 0.97) is **worse** than 1N/2N (0.81 / 0.76) because each rank sees too few steps to train. To get trustworthy 4/8-node weak-scaling numbers, the dataset must be enlarged so big global batches still yield many steps per epoch — that is **Phase 2** (stage more ERA5, see [`STAGING_ERA5_FOR_HBM.md`](../perf-optimizer-loop/STAGING_ERA5_FOR_HBM.md)).
+
+- **Loss sanity:** all four pass the monotonic-decrease check; note this only verifies loss goes down each epoch, not absolute convergence (see caveat above).
+- **Artifacts:** [`collate_scaling_study.py`](../../examples/collate_scaling_study.py) writes `scaling_study.{md,csv,json}` under `$ORBIT2_OUTPUT_DIR`.
+- **Reproduce:** [`run_scaling_study.sh`](../../examples/run_scaling_study.sh) `--nodes 1` then `--nodes 2,4,8` with `ORBIT2_CONFIG_TEMPLATE=edm_8m_era5_1x8.yaml`, `ORBIT2_BATCH_SIZE=1024`, `ORBIT2_DATA_TYPE=bfloat16`, `ORBIT2_MAX_EPOCH=6`, `ORBIT2_MAX_BATCHES=20`.
+
 ## Fine-tuning from a pretrained checkpoint
 
 Use `pretrain:` / `checkpoint:` in the YAML as described in the upstream "Hyperparameter Configuration" section when continuing or fine-tuning.