Skip to content

Commit 020b438

Browse files
ashwinmacursoragent
ashwinma
and
cursoragent
committed
ORBIT-2 + HydraGNN: group examples/ README by purpose
Reorganize the flat examples/ file lists into labeled sections (setup, container/launch drivers, in-container glue, configs, data prep, perf orchestration, post-processing) and document files that were previously missing from the tables. Keep the folders flat (paths are resolved relative to each script's dir and the container bind-mount); this is docs-only. Co-authored-by: Cursor <cursoragent@cursor.com>
1 parent e8ee40a commit 020b438

2 files changed

Lines changed: 101 additions & 13 deletions

File tree

  • earth_science/models/ORBIT-2/examples
  • material_science/models/HydraGNN/examples

‎earth_science/models/ORBIT-2/examples/README.md‎

Lines changed: 60 additions & 11 deletions
Original file line numberDiff line numberDiff line change
@@ -6,33 +6,82 @@ with AMD/ROCm container launch, overlay builds, and synthetic smoke tests.
66

77
## Files
88

9+
The folder is flat (every script resolves siblings, configs, and the
10+
`/examples` container bind-mount relative to its own directory), but the files
11+
fall into the groups below.
12+
13+
### Setup & inference
14+
915
| Script | Purpose |
1016
|--------|---------|
1117
| [`preflight_orbit2.py`](preflight_orbit2.py) | Check your environment before submitting jobs |
1218
| [`run_visualize.py`](run_visualize.py) | Run upstream `visualize.py` with Studio env-var overrides |
13-
| [`make_synthetic_data.py`](make_synthetic_data.py) | Generate ~2 MB synthetic dataset for smoke tests |
19+
20+
### Container & launch drivers
21+
22+
| Script | Purpose |
23+
|--------|---------|
1424
| [`docker_run.sh`](docker_run.sh) | Docker launcher for local workstations and interactive nodes |
15-
| [`sbatch_infer_amd.sh`](sbatch_infer_amd.sh) | SLURM driver for inference on AMD Instinct (MI250X/MI300X/MI350X) |
1625
| [`build_overlay_amd.sh`](build_overlay_amd.sh) | One-time overlay build (pre-bakes pip deps, skips ~15 min per job) |
17-
| [`interm_8m_synthetic.yaml`](interm_8m_synthetic.yaml) | Template YAML for synthetic mode |
1826
| [`sbatch_train_amd.sh`](sbatch_train_amd.sh) | Multi-node training (Apptainer + MPI); defaults PRISM 10.0_arcmin + `interm_8m_prism.yaml` |
1927
| [`sbatch_train_perf_amd.sh`](sbatch_train_perf_amd.sh) | 1-node × 8-GPU perf (Omnistat + profiler); defaults **ERA5 1.0°** + `edm_8m_era5_1x8.yaml`; multi-node: `sbatch --nodes=N …` |
28+
| [`sbatch_infer_amd.sh`](sbatch_infer_amd.sh) | SLURM driver for inference on AMD Instinct (MI250X/MI300X/MI350X) |
29+
| [`sbatch_infer_docker.sh`](sbatch_infer_docker.sh) | SLURM inference driver using Docker instead of Apptainer |
30+
31+
### In-container launch glue
32+
33+
Run on the host (`render_*`) or inside the container via the `/examples`
34+
bind-mount; the sbatch drivers depend on these by name, so they stay flat.
35+
36+
| Script | Purpose |
37+
|--------|---------|
38+
| [`render_orbit2_config.py`](render_orbit2_config.py) | Render a per-job YAML from a template (parallelism, data root, trainer caps); invoked by the sbatch drivers |
39+
| [`run_orbit2_train.py`](run_orbit2_train.py) | Studio training entry run inside the container (wraps upstream `intermediate_downscaling.py`; gptl4py stub, FusedAttn fallback, batch cap) |
40+
| [`orbit2_rank_hook_runner.py`](orbit2_rank_hook_runner.py) | Runs an optional per-rank pre-train hook (`ORBIT2_RANK_PRE_TRAIN_HOOK`) before the shell launcher |
41+
| [`orbit2_profiler_hook.py`](orbit2_profiler_hook.py) | Optional per-rank PyTorch-profiler hook (`res_slimvit` path) invoked via `ORBIT2_RANK_PRE_TRAIN_HOOK` |
42+
43+
### Configs (YAML)
44+
45+
Kept beside the sbatch drivers — they are resolved as `${SCRIPT_DIR}/${ORBIT2_CONFIG_TEMPLATE}`.
46+
47+
| Script | Purpose |
48+
|--------|---------|
49+
| [`interm_8m_prism.yaml`](interm_8m_prism.yaml) | PRISM config template (10.0_arcmin same-dir) |
50+
| [`interm_8m_era5.yaml`](interm_8m_era5.yaml) | ERA5 1.0° config template (same-dir sanity / pipeline test) |
51+
| [`interm_8m_synthetic.yaml`](interm_8m_synthetic.yaml) | Template YAML for synthetic mode |
52+
| [`edm_8m_era5_1x8.yaml`](edm_8m_era5_1x8.yaml) | Bayes-CAST `edm_8m_era5.yaml` + **`fsdp=8`/`simple_ddp=1`** + `seq_par:1`; `ORBIT2_ROOT` → `code/bayes-cast` (auto if present) |
53+
54+
### Data prep
55+
56+
| Script | Purpose |
57+
|--------|---------|
58+
| [`make_synthetic_data.py`](make_synthetic_data.py) | Generate ~2 MB synthetic dataset for smoke tests |
59+
| [`stage_era5_3x_symlink.sh`](stage_era5_3x_symlink.sh) | Symlink-replicate a staged ERA5 year (Nx) to break the small-corpus per-step batch cap for HBM saturation |
60+
61+
### Perf & scaling orchestration
62+
63+
| Script | Purpose |
64+
|--------|---------|
2065
| [`submit_perf_baseline_era5_amd.sh`](submit_perf_baseline_era5_amd.sh) | Thin `sbatch` wrapper (sets ERA5 defaults, forwards extra `sbatch` flags) |
2166
| [`sweep_orbit2_batch_bf16_amd.sh`](sweep_orbit2_batch_bf16_amd.sh) | Submit multiple `ORBIT2_BATCH_SIZE` probes (bf16 + SDPA) for HBM saturation sweeps |
67+
| [`run_scaling_study.sh`](run_scaling_study.sh) | Submit matched 1/2/4/8-node scaling sweep |
68+
| [`run_2node_scaleout_loop.sh`](run_2node_scaleout_loop.sh) | Unattended 2-node scale-out lever loop; writes a 1-node-vs-2-node `REPORT.md` |
69+
| [`run_gemm_analysis.sh`](run_gemm_analysis.sh) | Unattended GEMM-time bottleneck analysis (TraceLens + Omnistat analyst/verifier) at 1 and 2 nodes |
2270
| [`run_optimizer_loop.sh`](run_optimizer_loop.sh) | Driver for the iterative **perf-optimizer-loop** (Claude CLI optional) |
2371
| [`validate_orbit2_optimizer_loop_recipe.sh`](validate_orbit2_optimizer_loop_recipe.sh) | Repo smoke: required perf-optimizer-loop files exist |
24-
| [`run_scaling_study.sh`](run_scaling_study.sh) | Submit matched 1/2/4/8-node scaling sweep |
25-
| [`orbit2_estimate_batch_from_memory.py`](orbit2_estimate_batch_from_memory.py) | Heuristic max-batch from two `memory_reserved` calibrations or SLURM logs (see [BASELINE_LOCKIN.md](../recipes/perf-analysis/BASELINE_LOCKIN.md)) |
72+
73+
### Post-processing & analysis
74+
75+
Host-side; run after a job against its logs / traces.
76+
77+
| Script | Purpose |
78+
|--------|---------|
2679
| [`parse_training_log.py`](parse_training_log.py) | Extract batch/epoch metrics from SLURM log (handles both `intermediate_downscaling.py` and Bayes-CAST `train_edm.py` `epoch:/batch_idx/tic4-tic1` formats) |
2780
| [`run_fom_extractor.py`](run_fom_extractor.py) | Write `foms.json` (throughput + steady batch time + optional Omnistat PromQL via `ORBIT2_TSDB_URL`) |
2881
| [`report_orbit2_gpu_baseline.py`](report_orbit2_gpu_baseline.py) | `baseline_report.md` / JSON for perf runs |
29-
| [`interm_8m_prism.yaml`](interm_8m_prism.yaml) | PRISM config template (10.0_arcmin same-dir) |
30-
| [`edm_8m_era5_1x8.yaml`](edm_8m_era5_1x8.yaml) | Bayes-CAST `edm_8m_era5.yaml` + **`fsdp=8`/`simple_ddp=1`** + `seq_par:1`; `ORBIT2_ROOT` → `code/bayes-cast` (auto if present) |
31-
| [`orbit2_rank_hook_runner.py`](orbit2_rank_hook_runner.py) | Runs an optional per-rank pre-train hook (`ORBIT2_RANK_PRE_TRAIN_HOOK`) before the shell launcher |
32-
| [`stage_era5_3x_symlink.sh`](stage_era5_3x_symlink.sh) | Symlink-replicate a staged ERA5 year (Nx) to break the small-corpus per-step batch cap for HBM saturation |
33-
| [`run_gemm_analysis.sh`](run_gemm_analysis.sh) | Unattended GEMM-time bottleneck analysis (TraceLens + Omnistat analyst/verifier) at 1 and 2 nodes |
82+
| [`orbit2_estimate_batch_from_memory.py`](orbit2_estimate_batch_from_memory.py) | Heuristic max-batch from two `memory_reserved` calibrations or SLURM logs (see [BASELINE_LOCKIN.md](../recipes/perf-analysis/BASELINE_LOCKIN.md)) |
83+
| [`collate_scaling_study.py`](collate_scaling_study.py) | Aggregate a multi-node scaling sweep into `scaling_study` outputs |
3484
| [`compare_trace_kernels.py`](compare_trace_kernels.py) | Tool-independent kernel aggregation from a `*.pt.trace.json` (rank by raw device time, not category rollup) |
35-
| [`run_2node_scaleout_loop.sh`](run_2node_scaleout_loop.sh) | Unattended 2-node scale-out lever loop; writes a 1-node-vs-2-node `REPORT.md` |
3685
| [`upstream_pytorch_sdpa_benchmark.py`](upstream_pytorch_sdpa_benchmark.py) | Vendored [PyTorch `benchmarks/transformer/sdpa.py`](https://github.com/pytorch/pytorch/blob/main/benchmarks/transformer/sdpa.py) with **`--orbit-micro`** + backend sweep (`ck` → `SDPBackend.EFFICIENT_ATTENTION`, **not** xFormers `MemoryEfficientAttentionCkOp`) for comparing SDPA backends off the critical path |
3786

3887
## Quick start — ERA5 1.0_deg sanity (new dataset)

‎material_science/models/HydraGNN/examples/README.md‎

Lines changed: 41 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -6,16 +6,55 @@ with AMD/ROCm container launch.
66

77
## Files
88

9+
The folder is flat (scripts resolve siblings and the container bind-mount
10+
relative to their own directory), but the files fall into the groups below.
11+
12+
### Setup & diagnostics
13+
914
| Script | Purpose |
1015
|--------|---------|
1116
| [`preflight_hydragnn.py`](preflight_hydragnn.py) | Check your environment before submitting jobs |
12-
| [`run_inference.sh`](run_inference.sh) | Load a checkpoint and run predictions |
17+
| [`microbench_node_health.sh`](microbench_node_health.sh) | ~30 s per-node MI355X (gfx950) health survey (host/GPU inventory, dual-NUMA STREAM, HIP launch latency); safe to run in a SLURM Prolog |
18+
19+
### Run entrypoints
20+
21+
| Script | Purpose |
22+
|--------|---------|
1323
| [`run_train.sh`](run_train.sh) | Train on bundled or staged ADIOS datasets |
24+
| [`run_inference.sh`](run_inference.sh) | Load a checkpoint and run predictions |
25+
26+
### Container & launch drivers
27+
28+
| Script | Purpose |
29+
|--------|---------|
30+
| [`docker_run.sh`](docker_run.sh) | Docker launcher (clones HydraGNN, installs deps) |
31+
| [`build_overlay_amd.sh`](build_overlay_amd.sh) | One-time Apptainer ext3 overlay build (pre-loads pip deps; reuse across jobs to skip the per-job pip install) |
1432
| [`sbatch_train_amd.sh`](sbatch_train_amd.sh) | SLURM/Apptainer multi-node training (HPC) |
33+
| [`sbatch_train_perf_amd.sh`](sbatch_train_perf_amd.sh) | 2-node training + PyTorch profiler + Omnistat telemetry (thin variant of `sbatch_train_amd.sh` for the perf-analysis recipe) |
34+
| [`sbatch_infer_amd.sh`](sbatch_infer_amd.sh) | SLURM/Apptainer inference on AMD Instinct (needs a prebuilt overlay + downloaded weights) |
35+
36+
### Perf & scaling orchestration
37+
38+
| Script | Purpose |
39+
|--------|---------|
1540
| [`run_scaling_study.sh`](run_scaling_study.sh) | Submit matched 1/2/4/8-node strong-scaling sweep |
41+
| [`run_optimizer_loop.sh`](run_optimizer_loop.sh) | Entrypoint for the iterative sysopt **perf-optimizer-loop** (Claude CLI; run under `tmux`) |
42+
43+
### Post-processing & analysis
44+
45+
Host-side; run after a job against its logs / traces.
46+
47+
| Script | Purpose |
48+
|--------|---------|
1649
| [`collate_scaling_study.py`](collate_scaling_study.py) | Parse SLURM logs → steady-state throughput table |
1750
| [`parse_convergence.py`](parse_convergence.py) | Extract loss/epoch metrics from training logs |
18-
| [`docker_run.sh`](docker_run.sh) | Docker launcher (clones HydraGNN, installs deps) |
51+
| [`run_fom_extractor.py`](run_fom_extractor.py) | Compute FOMs + TraceLens↔Omnistat `kernel_correlation.csv` (login-node post-processing of a completed perf run) |
52+
53+
### Patches
54+
55+
| Path | Purpose |
56+
|--------|---------|
57+
| [`patches/`](patches/README.md) | Opt-in upstream patches applied during the overlay build (see [`patches/README.md`](patches/README.md)) |
1958

2059
## Quick start — Docker
2160

0 commit comments

Comments
 (0)