From 8e08da257359da36ed1d1f317f59076a0b5b7d61 Mon Sep 17 00:00:00 2001 From: OliviaViessmann Date: Fri, 10 Apr 2026 18:06:03 +0000 Subject: [PATCH 01/13] updated tracker and slurm scripts --- .devcontainer/devcontainer.json | 3 +- examples/well/v2/active_matter/TRACKER.md | 30 ++-- requirements-dev.txt | 1 + slurm/queue_well.sh | 56 +++++++ slurm/submit_well.sh | 183 ++++++++++++++++++++++ 5 files changed, 257 insertions(+), 16 deletions(-) create mode 100644 slurm/queue_well.sh create mode 100644 slurm/submit_well.sh diff --git a/.devcontainer/devcontainer.json b/.devcontainer/devcontainer.json index 64b561c5..dd003886 100644 --- a/.devcontainer/devcontainer.json +++ b/.devcontainer/devcontainer.json @@ -17,7 +17,8 @@ "source=${localEnv:HOME}/.netrc,target=/home/ubuntu/.netrc,readonly,type=bind,consistency=cached", // Mount bash history file for persistence. Created if it doesn't exist in initializeCommand.sh "source=${localEnv:HOME}/.bash_history_devcontainer,target=/home/ubuntu/.bash_history,type=bind,consistency=cached", - "source=${localEnv:HOME}/scratch/data/well_data,target=/workspaces/nvSubquadratic-private/data/the_well,type=bind" + // Mount gitconfig so commits from inside the container use the correct identity + "source=${localEnv:HOME}/.gitconfig,target=/home/ubuntu/.gitconfig,type=bind,consistency=cached" ], "containerEnv": { "TMPDIR": "/tmp" diff --git a/examples/well/v2/active_matter/TRACKER.md b/examples/well/v2/active_matter/TRACKER.md index 0c1fa56a..d0de76ce 100644 --- a/examples/well/v2/active_matter/TRACKER.md +++ b/examples/well/v2/active_matter/TRACKER.md @@ -47,27 +47,27 @@ If OOM at batch 64, add: ## Results — CNextU-net (baseline) -| Batch/GPU | Grad Accum | GPUs | WandB Run | val/MSE | it/s | WANDB | -| --------- | ---------- | ---- | --------- | ------- | ----- | ------------------------------------------------------------------- | -| 64 | 1 | 1 | | | 1.308 | UB_examples_well_v2_active_matter_unet_convnext_2026-04-03-15-36-52 | +| Batch/GPU | Grad Accum | GPUs | WandB Run | val/MSE | it/s | WANDB | +| --------- | ---------- | ---- | --------- | ------- | ---- | --------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------- | +| 64 | 1 | 1 | | | | OV\_\_workspaces_nvSubquadratic-private_examples_well_v2_active_matter_cfg_unet_convnext_experiment_dir\_/workspace/results_lr_0.005_num_nodes_1_run_start_time_177_2026-04-10-07-57-59 | ## Results — Attention -| Patch | Tokens/img | Batch/GPU | Grad Accum | GPUs | WandB Run | val/MSE | it/s | -| ----- | ---------- | --------- | ---------- | ---- | ---------------------------------------------------------------------------- | ------- | ----- | -| 16 | 256 | 64 | 1 | 1 | UB_examples_well_v2_active_matter_attention_2026-04-03-17-41-55 | | 0.983 | -| 8 | 1,024 | 64 | 1 | 1 | UB_examples_well_v2_active_matter_attention_patch_size_8_2026-04-03-18-50-30 | 0.067 | 1.213 | -| 4 | 4,096 | 64 | 1 | 1 | UB_examples_well_v2_active_matter_attention_patch_size_4_2026-04-03-18-54-10 | | | -| 2 | 16,384 | 32 | 2 | 1 | | | | +| Patch | Tokens/img | Batch/GPU | Grad Accum | GPUs | WandB Run | val/MSE | it/s | +| ----- | ---------- | --------- | ---------- | ---- | --------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------- | ------- | ---- | +| 16 | 256 | 64 | 1 | 1 | OV\_\_workspaces_nvSubquadratic-private_examples_well_v2_active_matter_attention_experiment_dir\_/workspace/results_lr_0.005_num_nodes_1_patch_size_16_run_start_ti_2026-04-10-07-53-21 | | | +| 8 | 1,024 | 64 | 1 | 1 | OV\_\_workspaces_nvSubquadratic-private_examples_well_v2_active_matter_attention_experiment_dir\_/workspace/results_lr_0.005_num_nodes_1_patch_size_8_run_start_tim_2026-04-10-07-53-30 | | | +| 4 | 4,096 | 64 | 1 | 1 | OV\_\_workspaces_nvSubquadratic-private_examples_well_v2_active_matter_attention_experiment_dir\_/workspace/results_lr_0.005_num_nodes_1_patch_size_4_run_start_tim_2026-04-10-07-53-41 | | | +| 2 | 16,384 | 32 | 2 | 1 | OV\_\_workspaces_nvSubquadratic-private_examples_well_v2_active_matter_attention_bs_32_experiment_dir\_/workspace/results_lr_0.005_num_nodes_1_patch_size_2_run_sta_2026-04-10-07-54-08 | | | ## Results — Hyena + Gaussian Mask -| Patch | Tokens/img | Batch/GPU | Grad Accum | GPUs | WandB Run | val/MSE | it/s | -| ----- | ---------- | --------- | ---------- | ---- | --------------------------------------------------------------------------------------- | ------- | ----- | -| 16 | 256 | 64 | 1 | 1 | UB_examples_well_v2_active_matter_hyena_gaussian_mask_patch_size_16_2026-04-03-17-08-22 | 0.0198 | 1.258 | -| 8 | 1,024 | 64 | 1 | 1 | UB_examples_well_v2_active_matter_hyena_gaussian_mask_patch_size_8_2026-04-03-18-46-37 | | 1.145 | -| 4 | 4,096 | 64 | 1 | 1 | UB_examples_well_v2_active_matter_hyena_gaussian_mask_patch_size_4_2026-04-03-18-53-17 | | | -| 2 | 16,384 | 16 | 4 | 1 | | | | +| Patch | Tokens/img | Batch/GPU | Grad Accum | GPUs | WandB Run | val/MSE | it/s | +| ----- | ---------- | --------- | ---------- | ---- | --------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------- | ------- | ---- | +| 16 | 256 | 64 | 1 | 1 | OV\_\_workspaces_nvSubquadratic-private_examples_well_v2_active_matter_hyena_gaussian_mask_experiment_dir\_/workspace/results_lr_0.005_num_nodes_1_patch_size_16_ru_2026-04-10-07-41-33 | | | +| 8 | 1,024 | 64 | 1 | 1 | OV\_\_workspaces_nvSubquadratic-private_examples_well_v2_active_matter_hyena_gaussian_mask_experiment_dir\_/workspace/results_lr_0.005_num_nodes_1_patch_size_8_run_2026-04-10-07-42-04 | | | +| 4 | 4,096 | 64 | 1 | 1 | OV\_\_workspaces_nvSubquadratic-private_examples_well_v2_active_matter_hyena_gaussian_mask_experiment_dir\_/workspace/results_lr_0.005_num_nodes_1_patch_size_4_run_2026-04-10-07-42-12 | | | +| 2 | 16,384 | 16 | 1 | 1 | OV\_\_workspaces_nvSubquadratic-private_examples_well_v2_active_matter_hyena_gaussian_mask_bs_16_experiment_dir\_/workspace/results_lr_0.005_num_nodes_1_patch_size_2026-04-10-07-44-00 | | | ## Memory Probing Notes (1×H100 80GB, with torch.compile max-autotune-no-cudagraphs) diff --git a/requirements-dev.txt b/requirements-dev.txt index f3675916..fa72a64c 100644 --- a/requirements-dev.txt +++ b/requirements-dev.txt @@ -18,3 +18,4 @@ detect-secrets>=1.5.0 jupyterlab>=4.0.0 clean-fid>=0.1.35 the-well>=1.0.0 +matplotlib diff --git a/slurm/queue_well.sh b/slurm/queue_well.sh new file mode 100644 index 00000000..a97f7484 --- /dev/null +++ b/slurm/queue_well.sh @@ -0,0 +1,56 @@ +#!/bin/bash +# Submit a linear chain of Well training jobs to work around the 4-hour wall-time limit. +# Each job starts after the previous one finishes (regardless of exit code), +# and autoresume picks up the latest checkpoint automatically. +# +# Usage (from repo root): +# bash slurm/queue_well.sh [config_overrides...] +# +# Examples: +# bash slurm/queue_well.sh 6 examples/well/v2/active_matter/hyena_gaussian_mask.py +# bash slurm/queue_well.sh 6 examples/well/v2/active_matter/hyena_gaussian_mask.py net.in_proj_cfg.patch_size=8 +# bash slurm/queue_well.sh 6 examples/well/v2/active_matter/hyena_gaussian_mask.py net.in_proj_cfg.patch_size=8 start_after_jid=5039096 + +if [ $# -lt 2 ]; then + echo "Usage: $0 [config_overrides...] [start_after_jid=]" + exit 1 +fi + +SCRIPT_NAME="slurm/submit_well.sh" +NUM_JOBS="$1" +CONFIG="$2" +shift 2 + +START_AFTER="" +CONFIG_OVERRIDES=() +for arg in "$@"; do + if [[ "${arg}" == start_after_jid=* ]]; then + START_AFTER="${arg#start_after_jid=}" + else + CONFIG_OVERRIDES+=("${arg}") + fi +done + +submit_job() { + local dep_flag="${1:-}" + sbatch ${dep_flag} "${SCRIPT_NAME}" "${CONFIG}" \ + "${CONFIG_OVERRIDES[@]+"${CONFIG_OVERRIDES[@]}"}" \ + | awk '{print $4}' +} + +if [ -n "${START_AFTER}" ]; then + jid_prev=$(submit_job "--dependency=afterany:${START_AFTER}") + echo "Submitted job 1/${NUM_JOBS}: ${jid_prev} (after ${START_AFTER})" +else + jid_prev=$(submit_job) + echo "Submitted job 1/${NUM_JOBS}: ${jid_prev}" +fi + +for i in $(seq 2 "${NUM_JOBS}"); do + jid_prev=$(submit_job "--dependency=afterany:${jid_prev}") + echo "Submitted job ${i}/${NUM_JOBS}: ${jid_prev}" +done + +echo "" +echo "Chain submitted. Last job ID: ${jid_prev}" +echo "Monitor with: squeue -u oviessmann" diff --git a/slurm/submit_well.sh b/slurm/submit_well.sh new file mode 100644 index 00000000..a00372b1 --- /dev/null +++ b/slurm/submit_well.sh @@ -0,0 +1,183 @@ +#!/bin/bash +#SBATCH --account=healthcareeng_research +#SBATCH --nodes=1 +#SBATCH --partition=batch,backfill +#SBATCH --ntasks-per-node=1 +#SBATCH --time=04:00:00 +#SBATCH --mem=0 +#SBATCH --mail-type=FAIL +#SBATCH --job-name=healthcareeng_research-nvsubq.well + +# Usage (from repo root): +# sbatch slurm/submit_well.sh examples/well/v2/active_matter/hyena_gaussian_mask.py +# sbatch slurm/submit_well.sh examples/well/v2/active_matter/hyena_gaussian_mask.py \ +# net.in_proj_cfg.patch_size=8 dataset.batch_size=32 +# +# Via queue.sh for chaining: +# bash slurm/queue.sh slurm/submit_well.sh 10 \ +# examples/well/v2/active_matter/hyena_gaussian_mask.py + +set -x + +if [ -z "${1:-}" ]; then + echo "Usage: sbatch slurm/enroot/submit_8gpu.sh [overrides...]" + exit 1 +fi + +# Capture start time immediately +JOB_START_TIMESTAMP=$(date +%s) +echo "Start time captured: ${JOB_START_TIMESTAMP}" + +# Container mount paths +CONTAINER_DATA="/workspace/data" +CONTAINER_RESULTS="/workspace/results" + +# ============================================================================ +# Configuration +# ============================================================================ +TIME_LIMIT_HOURS=4 +EXPERIMENT_NAME="$(basename "${1%.py}")" +CONFIG_FILE="$1"; shift + +CONFIG_OVERRIDES="num_nodes=${SLURM_JOB_NUM_NODES}" +CONFIG_OVERRIDES="${CONFIG_OVERRIDES} experiment_dir=${CONTAINER_RESULTS}" +# Append any extra overrides passed on the command line +for arg in "$@"; do + CONFIG_OVERRIDES="${CONFIG_OVERRIDES} ${arg}" +done + +# Container image +IMAGE_NAME="${SQSH_IMAGE:-/lustre/fsw/healthcareeng_research/oviessmann/nvsubquadratic/enroot/nvsubquadratic-x86_64.sqsh}" + +# Host paths +WORKDIR="${PWD}" +RUNS_DIR="${WORKDIR}/runs" +DATA_DIR="${WELL_HOST_DIR:-/lustre/fsw/healthcareeng_research/oviessmann/nvsubquadratic/data/well_data/datasets}" + +# Create necessary directories +mkdir -p "${RUNS_DIR}" + +# ============================================================================ +# Run naming (deterministic hash → same dir across restarts) +# ============================================================================ +RUN_NAME_HASH=$(echo "${CONFIG_FILE} ${CONFIG_OVERRIDES} ${EXPERIMENT_NAME}" | md5sum | awk '{print $1}' | cut -c1-8) +RUN_NAME="run_${RUN_NAME_HASH}" + +EXPERIMENT_DIR="${RUNS_DIR}/${EXPERIMENT_NAME}" +RESULTS_PATH="${EXPERIMENT_DIR}/${RUN_NAME}" + +mkdir -p "${EXPERIMENT_DIR}" +mkdir -p "${RESULTS_PATH}" + +# ============================================================================ +# W&B run ID — persisted across restarts so resume attaches to the same run +# ============================================================================ +if [ -f "${RESULTS_PATH}/run.id" ]; then + RUN_ID=$(<"${RESULTS_PATH}/run.id") + echo "Resuming with existing W&B run ID: ${RUN_ID}" +else + array=() + for i in {a..z} {A..Z} {0..9}; do + array[$RANDOM]=$i + done + RUN_ID=$(printf %s "${array[@]::8}") + echo "${RUN_ID}" > "${RESULTS_PATH}/run.id" + echo "Generated new W&B run ID: ${RUN_ID}" +fi + +echo "$(date): Job ${SLURM_JOB_ID} started (W&B run ID: ${RUN_ID})" >> "${RESULTS_PATH}/job_chain.log" + +# ============================================================================ +# Container mounts +# ============================================================================ +MOUNTS="${DATA_DIR}:${CONTAINER_DATA}" +MOUNTS="${MOUNTS},${RESULTS_PATH}:${CONTAINER_RESULTS}" +MOUNTS="${MOUNTS},$HOME/.cache:/root/.cache" + +if [ -f "$HOME/.netrc" ]; then + MOUNTS="${MOUNTS},$HOME/.netrc:/root/.netrc" +fi + +# Code is baked into the container +WORK_DIR="/workspaces/nvSubquadratic-private" +CONFIG_PATH="${WORK_DIR}/${CONFIG_FILE}" + +echo "================================================" +echo "Experiment: ${EXPERIMENT_NAME}" +echo "Job ID: ${SLURM_JOB_ID}" +echo "Run Name: ${RUN_NAME}" +echo "Node(s): ${SLURM_NODELIST}" +echo "Num nodes: ${SLURM_JOB_NUM_NODES}" +echo "GPUs per node: ${SLURM_NTASKS_PER_NODE}" +echo "Total GPUs: $((${SLURM_JOB_NUM_NODES} * ${SLURM_NTASKS_PER_NODE}))" +echo "Config: ${CONFIG_FILE}" +echo "Overrides: ${CONFIG_OVERRIDES}" +echo "Container: ${IMAGE_NAME}" +echo "Results: ${RESULTS_PATH}" +echo "Mounts: ${MOUNTS}" +echo "================================================" + +# ============================================================================ +# Environment Setup +# ============================================================================ +export WANDB_API_KEY="${WANDB_API_KEY:-}" +export PYTHONPATH="." +export CUDA_VISIBLE_DEVICES=0 +export DALI_NO_MMAP=1 +export TRITON_CACHE_DIR="/tmp/triton_${SLURM_JOB_ID}" +export TORCHINDUCTOR_FX_GRAPH_CACHE=0 +export TORCH_NCCL_AVOID_RECORD_STREAMS=1 +export PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True + + +# Well dataset path (container-side) — read by _base.py via os.environ.get("WELL_DATA_PATH") +# DATA_DIR (host) is mounted at CONTAINER_DATA; it must contain active_matter/ directly. +# If on-cluster layout is the_well/datasets/active_matter/, set WELL_HOST_DIR=.../the_well/datasets +export WELL_DATA_PATH="${CONTAINER_DATA}" + +# ============================================================================ +# Autoresume +# ============================================================================ +if [ -f "${RESULTS_PATH}/checkpoints/last.ckpt" ]; then + echo "Found existing checkpoint — enabling autoresume" + AUTORESUME_ARG="autoresume.enabled=True" +else + echo "No existing checkpoint found, starting fresh" + AUTORESUME_ARG="wandb.run_id=${RUN_ID}" +fi + +# ============================================================================ +# Training command +# ============================================================================ +read -r -d '' PYTHON_CMD <> "${RESULTS_PATH}/job_chain.log" + +set +x +exit ${TRAIN_EXIT_CODE} From a96a74d3868efd701bb6d1fb37fc672818c86e6f Mon Sep 17 00:00:00 2001 From: OliviaViessmann Date: Thu, 16 Apr 2026 13:45:04 +0000 Subject: [PATCH 02/13] active matter final experiments results (tracker) --- .gitignore | 4 +++ .../well/v2/acoustic_scattering_maze/hyena.py | 2 +- examples/well/v2/active_matter/TRACKER.md | 30 +++++++++---------- scripts/data/download_well.sh | 2 +- slurm/submit_well.sh | 9 +++--- 5 files changed, 26 insertions(+), 21 deletions(-) diff --git a/.gitignore b/.gitignore index bdd60e25..d673dd0f 100644 --- a/.gitignore +++ b/.gitignore @@ -245,6 +245,10 @@ $RECYCLE.BIN/ # Apptainer images *.sif +# Enroot container images and packages +*.sqsh +enroot*.deb + # Development and testing .pytest_cache/ .coverage diff --git a/examples/well/v2/acoustic_scattering_maze/hyena.py b/examples/well/v2/acoustic_scattering_maze/hyena.py index 47fd02d7..9dee4d0f 100644 --- a/examples/well/v2/acoustic_scattering_maze/hyena.py +++ b/examples/well/v2/acoustic_scattering_maze/hyena.py @@ -42,7 +42,7 @@ DROPOUT_IN_RATE = 0.0 DROPOUT_RATE = 0.0 GRID_TYPE = "single" -FFT_PADDING = "zeros" # non-periodic boundary conditions (WALL/OPEN) +FFT_PADDING = "zero" # non-periodic boundary conditions (WALL/OPEN) OMEGA_0 = 30.0 diff --git a/examples/well/v2/active_matter/TRACKER.md b/examples/well/v2/active_matter/TRACKER.md index d0de76ce..d3f46746 100644 --- a/examples/well/v2/active_matter/TRACKER.md +++ b/examples/well/v2/active_matter/TRACKER.md @@ -47,27 +47,27 @@ If OOM at batch 64, add: ## Results — CNextU-net (baseline) -| Batch/GPU | Grad Accum | GPUs | WandB Run | val/MSE | it/s | WANDB | -| --------- | ---------- | ---- | --------- | ------- | ---- | --------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------- | -| 64 | 1 | 1 | | | | OV\_\_workspaces_nvSubquadratic-private_examples_well_v2_active_matter_cfg_unet_convnext_experiment_dir\_/workspace/results_lr_0.005_num_nodes_1_run_start_time_177_2026-04-10-07-57-59 | +| Batch/GPU | Grad Accum | GPUs | WandB Run | val/VRMSE | it/s | total steps | +| --------- | ---------- | ---- | ----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------- | -------------------- | ---- | ----------- | +| 64 | 1 | 1 | OV\_\_workspaces_nvSubquadratic-private_examples_well_v2_active_matter_cfg_unet_convnext_enabled_True_experiment_dir\_/workspace/results_lr_0.005_num_nodes_1_run_start_time_1775922336_run_time_limit_hour | 0.034676797688007355 | 1.39 | 110,000 | ## Results — Attention -| Patch | Tokens/img | Batch/GPU | Grad Accum | GPUs | WandB Run | val/MSE | it/s | -| ----- | ---------- | --------- | ---------- | ---- | --------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------- | ------- | ---- | -| 16 | 256 | 64 | 1 | 1 | OV\_\_workspaces_nvSubquadratic-private_examples_well_v2_active_matter_attention_experiment_dir\_/workspace/results_lr_0.005_num_nodes_1_patch_size_16_run_start_ti_2026-04-10-07-53-21 | | | -| 8 | 1,024 | 64 | 1 | 1 | OV\_\_workspaces_nvSubquadratic-private_examples_well_v2_active_matter_attention_experiment_dir\_/workspace/results_lr_0.005_num_nodes_1_patch_size_8_run_start_tim_2026-04-10-07-53-30 | | | -| 4 | 4,096 | 64 | 1 | 1 | OV\_\_workspaces_nvSubquadratic-private_examples_well_v2_active_matter_attention_experiment_dir\_/workspace/results_lr_0.005_num_nodes_1_patch_size_4_run_start_tim_2026-04-10-07-53-41 | | | -| 2 | 16,384 | 32 | 2 | 1 | OV\_\_workspaces_nvSubquadratic-private_examples_well_v2_active_matter_attention_bs_32_experiment_dir\_/workspace/results_lr_0.005_num_nodes_1_patch_size_2_run_sta_2026-04-10-07-54-08 | | | +| Patch | Tokens/img | Batch/GPU | Grad Accum | GPUs | WandB Run | val/VRMSE | it/s | total steps | +| ----- | ---------- | --------- | ---------- | ---- | ----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------- | -------------------- | ---- | ----------- | +| 16 | 256 | 64 | 1 | 1 | OV\_\_workspaces_nvSubquadratic-private_examples_well_v2_active_matter_attention_enabled_True_experiment_dir\_/workspace/results_lr_0.005_num_nodes_1_patch_size_16_run_start_time_1775918134_run_time_limi | 0.06180257350206375 | 1.40 | 110,000 | +| 8 | 1,024 | 64 | 1 | 1 | OV\_\_workspaces_nvSubquadratic-private_examples_well_v2_active_matter_attention_enabled_True_experiment_dir\_/workspace/results_lr_0.005_num_nodes_1_patch_size_8_run_start_time_1775919205_run_time_limit | 0.058588799089193344 | 1.43 | 110,000 | +| 4 | 4,096 | 64 | 1 | 1 | OV\_\_workspaces_nvSubquadratic-private_examples_well_v2_active_matter_attention_enabled_True_experiment_dir\_/workspace/results_lr_0.005_num_nodes_1_patch_size_4_run_start_time_1775922307_run_time_limit | 0.06164788082242012 | 1.33 | 102,480 | +| 2 | 16,384 | 32 | 2 | 1 | OV\_\_workspaces_nvSubquadratic-private_examples_well_v2_active_matter_attention_bs_32_enabled_True_experiment_dir\_/workspace/results_lr_0.005_num_nodes_1_patch_size_2_run_start_time_1775922311_run_time | 0.09142962843179704 | 0.43 | 36,605 | ## Results — Hyena + Gaussian Mask -| Patch | Tokens/img | Batch/GPU | Grad Accum | GPUs | WandB Run | val/MSE | it/s | -| ----- | ---------- | --------- | ---------- | ---- | --------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------- | ------- | ---- | -| 16 | 256 | 64 | 1 | 1 | OV\_\_workspaces_nvSubquadratic-private_examples_well_v2_active_matter_hyena_gaussian_mask_experiment_dir\_/workspace/results_lr_0.005_num_nodes_1_patch_size_16_ru_2026-04-10-07-41-33 | | | -| 8 | 1,024 | 64 | 1 | 1 | OV\_\_workspaces_nvSubquadratic-private_examples_well_v2_active_matter_hyena_gaussian_mask_experiment_dir\_/workspace/results_lr_0.005_num_nodes_1_patch_size_8_run_2026-04-10-07-42-04 | | | -| 4 | 4,096 | 64 | 1 | 1 | OV\_\_workspaces_nvSubquadratic-private_examples_well_v2_active_matter_hyena_gaussian_mask_experiment_dir\_/workspace/results_lr_0.005_num_nodes_1_patch_size_4_run_2026-04-10-07-42-12 | | | -| 2 | 16,384 | 16 | 1 | 1 | OV\_\_workspaces_nvSubquadratic-private_examples_well_v2_active_matter_hyena_gaussian_mask_bs_16_experiment_dir\_/workspace/results_lr_0.005_num_nodes_1_patch_size_2026-04-10-07-44-00 | | | +| Patch | Tokens/img | Batch/GPU | Grad Accum | GPUs | WandB Run | val/VRMSE | it/s | total steps | +| ----- | ---------- | --------- | ---------- | ---- | ------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------ | -------------------- | ---- | ----------- | +| 16 | 256 | 64 | 1 | 1 | OV\_\_workspaces_nvSubquadratic-private_examples_well_v2_active_matter_hyena_gaussian_mask_enabled_True_experiment_dir\_/workspace/results_lr_0.005_num_nodes_1_patch_size_16_run_start_time_1775914506_run | 0.012452345341444016 | 1.38 | 110,000 | +| 8 | 1,024 | 64 | 1 | 1 | OV\_\_workspaces_nvSubquadratic-private_examples_well_v2_active_matter_hyena_gaussian_mask_enabled_True_experiment_dir\_/workspace/results_lr_0.005_num_nodes_1_patch_size_8_run_start_time_1775918046_run\_ | 0.007334186229854822 | 1.39 | 110,000 | +| 4 | 4,096 | 64 | 1 | 1 | OV\_\_workspaces_nvSubquadratic-private_examples_well_v2_active_matter_hyena_gaussian_mask_enabled_True_experiment_dir\_/workspace/results_lr_0.005_num_nodes_1_patch_size_4_run_start_time_1775918057_run\_ | 0.007981576956808567 | 1.36 | 110,000 | +| 2 | 16,384 | 16 | 1 | 1 | OV\_\_workspaces_nvSubquadratic-private_examples_well_v2_active_matter_hyena_gaussian_mask_bs_16_enabled_True_experiment_dir\_/workspace/results_lr_0.005_num_nodes_1_patch_size_2_run_start_time_177587628 | 0.00703924847766757 | 2.63 | 110,000 | ## Memory Probing Notes (1×H100 80GB, with torch.compile max-autotune-no-cudagraphs) diff --git a/scripts/data/download_well.sh b/scripts/data/download_well.sh index c94e6236..cb8d094f 100644 --- a/scripts/data/download_well.sh +++ b/scripts/data/download_well.sh @@ -53,7 +53,7 @@ DATASET_NAME="$1" shift # Default base path from env var or fallback -BASE_PATH="${WELL_DATA_PATH:-./data/the_well}" +BASE_PATH="${WELL_DATA_PATH:-./data/well_data}" SPLIT="" # Parse optional arguments diff --git a/slurm/submit_well.sh b/slurm/submit_well.sh index a00372b1..664ef9f3 100644 --- a/slurm/submit_well.sh +++ b/slurm/submit_well.sh @@ -90,18 +90,19 @@ echo "$(date): Job ${SLURM_JOB_ID} started (W&B run ID: ${RUN_ID})" >> "${RESULT # ============================================================================ # Container mounts # ============================================================================ +REPO_DIR="/lustre/fsw/healthcareeng_research/oviessmann/nvsubquadratic" +WORK_DIR="/workspaces/nvSubquadratic-private" +CONFIG_PATH="${WORK_DIR}/${CONFIG_FILE}" + MOUNTS="${DATA_DIR}:${CONTAINER_DATA}" MOUNTS="${MOUNTS},${RESULTS_PATH}:${CONTAINER_RESULTS}" +MOUNTS="${MOUNTS},${REPO_DIR}:${WORK_DIR}" MOUNTS="${MOUNTS},$HOME/.cache:/root/.cache" if [ -f "$HOME/.netrc" ]; then MOUNTS="${MOUNTS},$HOME/.netrc:/root/.netrc" fi -# Code is baked into the container -WORK_DIR="/workspaces/nvSubquadratic-private" -CONFIG_PATH="${WORK_DIR}/${CONFIG_FILE}" - echo "================================================" echo "Experiment: ${EXPERIMENT_NAME}" echo "Job ID: ${SLURM_JOB_ID}" From 2114f85913a266a8d7efaa2204ca4795cf8d603a Mon Sep 17 00:00:00 2001 From: OliviaViessmann Date: Fri, 17 Apr 2026 18:07:24 +0000 Subject: [PATCH 03/13] updated tracker with final results for acoustic_scattering_maze --- examples/well/v2/MHD_64/_base.py | 2 +- .../v2/acoustic_scattering_maze/TRACKER.md | 39 +++++++++++++++++++ 2 files changed, 40 insertions(+), 1 deletion(-) create mode 100644 examples/well/v2/acoustic_scattering_maze/TRACKER.md diff --git a/examples/well/v2/MHD_64/_base.py b/examples/well/v2/MHD_64/_base.py index c5e1d772..28348b08 100644 --- a/examples/well/v2/MHD_64/_base.py +++ b/examples/well/v2/MHD_64/_base.py @@ -44,7 +44,7 @@ TRAINING_ITERATIONS = 110_000 WARMUP_ITERATIONS_PERCENTAGE = 0.05 -BATCH_SIZE = 2 # configs/data/MHD_64.yaml +BATCH_SIZE = 64 # configs/data/MHD_64.yaml NUM_WORKERS = 12 GRAD_CLIP = 1.0 PRECISION = "bf16-mixed" diff --git a/examples/well/v2/acoustic_scattering_maze/TRACKER.md b/examples/well/v2/acoustic_scattering_maze/TRACKER.md new file mode 100644 index 00000000..5871e2db --- /dev/null +++ b/examples/well/v2/acoustic_scattering_maze/TRACKER.md @@ -0,0 +1,39 @@ +# Acoustic Scattering Maze + +## Goal + +Compare **CNextU-net** (baseline), **Attention**, and **Hyena + Gaussian mask** + +All runs share the same training recipe (110k iterations, AdamW lr=5e-3, +wd=1e-4, cosine schedule with 5% warmup, bf16-mixed, grad_clip=1.0). +Default batch size is **64** on **1 GPU**. + +## Run order + +1. CNextU-net (baseline, no patch size) +1. Attention p16 → p8 → p4 → p2 +1. Hyena+G p16 → p8 → p4 → p2 + +## Results — CNextU-net (baseline) + +| Batch/GPU | Grad Accum | GPUs | WandB Run | val/VRMSE | it/s | total steps | +| --------- | ---------- | ---- | ----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------- | -------------------- | ------------ | ----------- | +| 64 | 1 | 1 | OV\_\_workspaces_nvSubquadratic-private_examples_well_v2_acoustic_scattering_maze_cfg_unet_convnext_enabled_True_experiment_dir\_/workspace/results_lr_0.001_num_nodes_1_run_start_time_1776227052_run_time | 0.008186010643839836 | 5.3794992175 | 110,000 | + +## Results — Attention + +| Patch | Tokens/img | Batch/GPU | Grad Accum | GPUs | WandB Run | val/VRMSE | it/s | total steps | +| ----- | ---------- | --------- | ---------- | ---- | ----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------- | -------------------- | ------------ | ----------- | +| 16 | 256 | 64 | 1 | 1 | OV\_\_workspaces_nvSubquadratic-private_examples_well_v2_acoustic_scattering_maze_attention_enabled_True_experiment_dir\_/workspace/results_lr_0.001_num_nodes_1_run_start_time_1776226723_run_time_limit_h | 0.04768187180161476 | 5.8760683761 | 110,000 | +| 8 | 1,024 | 64 | 1 | 1 | OV\_\_workspaces_nvSubquadratic-private_examples_well_v2_acoustic_scattering_maze_attention_enabled_True_experiment_dir\_/workspace/results_lr_0.001_num_nodes_1_patch_size_8_run_start_time_1776335992_run | 0.045643627643585205 | 5.7651991614 | 110,000 | +| 4 | 4,096 | 64 | 1 | 1 | OV\_\_workspaces_nvSubquadratic-private_examples_well_v2_acoustic_scattering_maze_attention_enabled_True_experiment_dir\_/workspace/results_lr_0.001_num_nodes_1_patch_size_4_run_start_time_1776340302_run | 0.056851577013731 | 2.3595023595 | 110,000 | +| 2 | 16,384 | 32 | 2 | 1 | OV\_\_workspaces_nvSubquadratic-private_examples_well_v2_acoustic_scattering_maze_attention_bs_16_enabled_True_experiment_dir\_/workspace/results_lr_0.001_num_nodes_1_patch_size_2_run_start_time_17763786 | 0.10574174672365189 | 0.8702652948 | 69,019 | + +## Results — Hyena + Gaussian Mask + +| Patch | Tokens/img | Batch/GPU | Grad Accum | GPUs | WandB Run | val/VRMSE | it/s | total steps | +| ----- | ---------- | --------- | ---------- | ---- | ----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------- | -------------------- | ------------ | ----------- | +| 16 | 256 | 64 | 1 | 1 | OV\_\_workspaces_nvSubquadratic-private_examples_well_v2_acoustic_scattering_maze_hyena_enabled_True_experiment_dir\_/workspace/results_lr_0.001_num_nodes_1_run_start_time_1776335625_run_time_limit_hours | 0.0260219257324934 | 5.8873902805 | 110,000 | +| 8 | 1,024 | 64 | 1 | 1 | OV\_\_workspaces_nvSubquadratic-private_examples_well_v2_acoustic_scattering_maze_hyena_gaussian_mask_enabled_True_experiment_dir\_/workspace/results_lr_0.001_num_nodes_1_patch_size_8_run_start_time_1776 | 0.00858149304986 | 4.2379411311 | 110,000 | +| 4 | 4,096 | 64 | 1 | 1 | OV\_\_workspaces_nvSubquadratic-private_examples_well_v2_acoustic_scattering_maze_hyena_gaussian_mask_enabled_True_experiment_dir\_/workspace/results_lr_0.001_num_nodes_1_patch_size_4_run_start_time_1776 | 0.006847195327281952 | 1.6516516517 | 110,000 | +| 2 | 16,384 | 16 | 1 | 1 | OV\_\_workspaces_nvSubquadratic-private_examples_well_v2_acoustic_scattering_maze_hyena_gaussian_mask_bs_16_enabled_True_experiment_dir\_/workspace/results_lr_0.001_num_nodes_1_patch_size_2_run_start_tim | 0.006232426967471838 | 1.5774680204 | 110,000 | From e0ce3c9fc57a6b0c7bf365d30c135d031164b9d6 Mon Sep 17 00:00:00 2001 From: OliviaViessmann Date: Mon, 27 Apr 2026 14:21:03 +0000 Subject: [PATCH 04/13] adding MHD64 Tracker results and gray scott reaction diffusion configs --- examples/well/v2/MHD_64/TRACKER.md | 39 +++++ examples/well/v2/MHD_64/attention.py | 4 +- .../v2/acoustic_scattering_maze/TRACKER.md | 32 ++-- examples/well/v2/active_matter/TRACKER.md | 49 +++--- .../gray_scott_reaction_diffusion/TRACKER.md | 51 +++++++ .../attention.py | 104 +++++++++++++ .../v2/gray_scott_reaction_diffusion/hyena.py | 141 ++++++++++++++++++ .../hyena_gaussian_mask.py | 31 ++++ 8 files changed, 403 insertions(+), 48 deletions(-) create mode 100644 examples/well/v2/MHD_64/TRACKER.md create mode 100644 examples/well/v2/gray_scott_reaction_diffusion/TRACKER.md create mode 100644 examples/well/v2/gray_scott_reaction_diffusion/attention.py create mode 100644 examples/well/v2/gray_scott_reaction_diffusion/hyena.py create mode 100644 examples/well/v2/gray_scott_reaction_diffusion/hyena_gaussian_mask.py diff --git a/examples/well/v2/MHD_64/TRACKER.md b/examples/well/v2/MHD_64/TRACKER.md new file mode 100644 index 00000000..e4cc565c --- /dev/null +++ b/examples/well/v2/MHD_64/TRACKER.md @@ -0,0 +1,39 @@ +# MHD_64 + +## Goal + +Compare **CNextU-net** (baseline), **Attention**, and **Hyena + Gaussian mask** + +All runs share the same training recipe ( AdamW lr=2-5e-3, +wd=1e-4, cosine schedule with 5% warmup, bf16-mixed, grad_clip=1.0). +Default batch size is **64** on **1 GPU**. + +## Run order + +1. CNextU-net (baseline, no patch size) +1. Attention p16 → p8 → p4 → p2 +1. Hyena+G p16 → p8 → p4 → p2 + +## Results — CNextU-net (baseline) + +| Batch/GPU | GPUs | WandB Run | val/VRMSE | it/s | total steps | +| --------- | ---- | --------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------- | ------------------- | ------------ | ----------- | +| 64 | 1 | OV\_\_workspaces_nvSubquadratic-private_examples_well_v2_MHD_64_unet_convnext_enabled_True_experiment_dir\_/workspace/results_lr_0.002_num_nodes_1_run_start_time_1776680521_run_time_limit_hours_4 | 0.21076937019824984 | 0.1066574978 | 9,151 | + +## Results — Attention + +| Patch | Batch/GPU | GPUs | WandB Run | val/VRMSE | it/s | total steps | +| ----- | --------- | ---- | ------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------ | ------------------- | ------------ | ----------- | +| 8 | 64 | 1 | OV\_\_workspaces_nvSubquadratic-private_examples_well_v2_MHD_64_attention_enabled_True_experiment_dir\_/workspace/results_lr_0.002_num_nodes_1_patch_size_8_run_start_time_1777096814_run_time_limit_hours\_ | 0.30443859100341797 | 0.7517476852 | 64951 | +| 4 | 64 | 1 | OV\_\_workspaces_nvSubquadratic-private_examples_well_v2_MHD_64_attention_enabled_True_experiment_dir\_/workspace/results_lr_0.002_num_nodes_1_patch_size_4_run_start_time_1777036508_run_time_limit_hours\_ | 0.21639029681682587 | 0.3653587963 | 31567 | +| 2 | 16 | 1 | OV\_\_workspaces_nvSubquadratic-private_examples_well_v2_MHD_64_attention_bs_16_enabled_True_experiment_dir\_/workspace/results_lr_0.002_num_nodes_1_patch_size_2_run_start_time_1777036405_run_time_limit\_ | 0.30372998118400574 | 0.0803125 | 6939 | + +## Results — Hyena + Gaussian Mask + +All experiments run with lr=2e-3 + +| Patch | Batch/GPU | GPUs | WandB Run | val/VRMSE | it/s | total steps | +| ----- | --------- | ---- | ----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------- | ------------------- | ------------ | ----------- | +| 8 | 64 | 1 | OV\_\_workspaces_nvSubquadratic-private_examples_well_v2_MHD_64_hyena_gaussian_mask_enabled_True_experiment_dir\_/workspace/results_lr_0.002_num_nodes_1_patch_size_8_run_start_time_1776680370_run_time_li | 0.2809585630893707 | 0.7427864124 | 63610 | +| 4 | 64 | 1 | OV\_\_workspaces_nvSubquadratic-private_examples_well_v2_MHD_64_hyena_gaussian_mask_enabled_True_experiment_dir\_/workspace/results_lr_0.002_num_nodes_1_patch_size_4_run_start_time_1776680398_run_time_li | 0.10876571387052536 | 0.5321398562 | 45598 | +| 2 | 16 | 1 | OV\_\_workspaces_nvSubquadratic-private_examples_well_v2_MHD_64_hyena_gaussian_mask_bs_8_enabled_True_experiment_dir\_/workspace/results_lr_0.002_num_nodes_1_patch_size_2_run_start_time_1776682943_run_ti | 0.05432766303420067 | 0.8076215108 | 69323 | diff --git a/examples/well/v2/MHD_64/attention.py b/examples/well/v2/MHD_64/attention.py index fcbea825..4f3aa112 100644 --- a/examples/well/v2/MHD_64/attention.py +++ b/examples/well/v2/MHD_64/attention.py @@ -36,7 +36,9 @@ # 4×4×4 = 64 tokens, too coarse for a 3D PDE. NUM_HIDDEN_CHANNELS = 384 NUM_BLOCKS = 12 -NUM_HEADS = 6 +# NUM_HEADS=8 (vs 6 in the 2D siblings): 3D RoPE requires head_dim % 6 == 0. +# 384/6=64 (64%6=4) fails the assert; 384/8=48 (48%6=0) passes. +NUM_HEADS = 8 PATCH_SIZE = 8 DROPOUT_IN_RATE = 0.0 diff --git a/examples/well/v2/acoustic_scattering_maze/TRACKER.md b/examples/well/v2/acoustic_scattering_maze/TRACKER.md index 5871e2db..30981b06 100644 --- a/examples/well/v2/acoustic_scattering_maze/TRACKER.md +++ b/examples/well/v2/acoustic_scattering_maze/TRACKER.md @@ -4,7 +4,7 @@ Compare **CNextU-net** (baseline), **Attention**, and **Hyena + Gaussian mask** -All runs share the same training recipe (110k iterations, AdamW lr=5e-3, +All runs share the same training recipe (24 hours, AdamW lr=5e-3, wd=1e-4, cosine schedule with 5% warmup, bf16-mixed, grad_clip=1.0). Default batch size is **64** on **1 GPU**. @@ -16,24 +16,24 @@ Default batch size is **64** on **1 GPU**. ## Results — CNextU-net (baseline) -| Batch/GPU | Grad Accum | GPUs | WandB Run | val/VRMSE | it/s | total steps | -| --------- | ---------- | ---- | ----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------- | -------------------- | ------------ | ----------- | -| 64 | 1 | 1 | OV\_\_workspaces_nvSubquadratic-private_examples_well_v2_acoustic_scattering_maze_cfg_unet_convnext_enabled_True_experiment_dir\_/workspace/results_lr_0.001_num_nodes_1_run_start_time_1776227052_run_time | 0.008186010643839836 | 5.3794992175 | 110,000 | +| Batch/GPU | GPUs | WandB Run | val/VRMSE | it/s | total steps | +| --------- | ---- | ----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------- | -------------------- | ------------ | ----------- | +| 64 | 1 | OV\_\_workspaces_nvSubquadratic-private_examples_well_v2_acoustic_scattering_maze_cfg_unet_convnext_enabled_True_experiment_dir\_/workspace/results_lr_0.001_num_nodes_1_run_start_time_1776227052_run_time | 0.008186010643839836 | 5.3794992175 | 110,000 | ## Results — Attention -| Patch | Tokens/img | Batch/GPU | Grad Accum | GPUs | WandB Run | val/VRMSE | it/s | total steps | -| ----- | ---------- | --------- | ---------- | ---- | ----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------- | -------------------- | ------------ | ----------- | -| 16 | 256 | 64 | 1 | 1 | OV\_\_workspaces_nvSubquadratic-private_examples_well_v2_acoustic_scattering_maze_attention_enabled_True_experiment_dir\_/workspace/results_lr_0.001_num_nodes_1_run_start_time_1776226723_run_time_limit_h | 0.04768187180161476 | 5.8760683761 | 110,000 | -| 8 | 1,024 | 64 | 1 | 1 | OV\_\_workspaces_nvSubquadratic-private_examples_well_v2_acoustic_scattering_maze_attention_enabled_True_experiment_dir\_/workspace/results_lr_0.001_num_nodes_1_patch_size_8_run_start_time_1776335992_run | 0.045643627643585205 | 5.7651991614 | 110,000 | -| 4 | 4,096 | 64 | 1 | 1 | OV\_\_workspaces_nvSubquadratic-private_examples_well_v2_acoustic_scattering_maze_attention_enabled_True_experiment_dir\_/workspace/results_lr_0.001_num_nodes_1_patch_size_4_run_start_time_1776340302_run | 0.056851577013731 | 2.3595023595 | 110,000 | -| 2 | 16,384 | 32 | 2 | 1 | OV\_\_workspaces_nvSubquadratic-private_examples_well_v2_acoustic_scattering_maze_attention_bs_16_enabled_True_experiment_dir\_/workspace/results_lr_0.001_num_nodes_1_patch_size_2_run_start_time_17763786 | 0.10574174672365189 | 0.8702652948 | 69,019 | +| Patch | Tokens/img | Batch/GPU | GPUs | WandB Run | val/VRMSE | it/s | total steps | +| ----- | ---------- | --------- | ---- | ----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------- | -------------------- | ------------ | ----------- | +| 16 | 256 | 64 | 1 | OV\_\_workspaces_nvSubquadratic-private_examples_well_v2_acoustic_scattering_maze_attention_enabled_True_experiment_dir\_/workspace/results_lr_0.001_num_nodes_1_run_start_time_1776226723_run_time_limit_h | 0.04768187180161476 | 5.8760683761 | 110,000 | +| 8 | 1,024 | 64 | 1 | OV\_\_workspaces_nvSubquadratic-private_examples_well_v2_acoustic_scattering_maze_attention_enabled_True_experiment_dir\_/workspace/results_lr_0.001_num_nodes_1_patch_size_8_run_start_time_1776335992_run | 0.045643627643585205 | 5.7651991614 | 110,000 | +| 4 | 4,096 | 64 | 1 | OV\_\_workspaces_nvSubquadratic-private_examples_well_v2_acoustic_scattering_maze_attention_enabled_True_experiment_dir\_/workspace/results_lr_0.001_num_nodes_1_patch_size_4_run_start_time_1776340302_run | 0.056851577013731 | 2.3595023595 | 110,000 | +| 2 | 16,384 | 16 | 1 | OV\_\_workspaces_nvSubquadratic-private_examples_well_v2_acoustic_scattering_maze_attention_bs_16_enabled_True_experiment_dir\_/workspace/results_lr_0.001_num_nodes_1_patch_size_2_run_start_time_17763786 | 0.10574174672365189 | 0.8702652948 | 69,019 | ## Results — Hyena + Gaussian Mask -| Patch | Tokens/img | Batch/GPU | Grad Accum | GPUs | WandB Run | val/VRMSE | it/s | total steps | -| ----- | ---------- | --------- | ---------- | ---- | ----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------- | -------------------- | ------------ | ----------- | -| 16 | 256 | 64 | 1 | 1 | OV\_\_workspaces_nvSubquadratic-private_examples_well_v2_acoustic_scattering_maze_hyena_enabled_True_experiment_dir\_/workspace/results_lr_0.001_num_nodes_1_run_start_time_1776335625_run_time_limit_hours | 0.0260219257324934 | 5.8873902805 | 110,000 | -| 8 | 1,024 | 64 | 1 | 1 | OV\_\_workspaces_nvSubquadratic-private_examples_well_v2_acoustic_scattering_maze_hyena_gaussian_mask_enabled_True_experiment_dir\_/workspace/results_lr_0.001_num_nodes_1_patch_size_8_run_start_time_1776 | 0.00858149304986 | 4.2379411311 | 110,000 | -| 4 | 4,096 | 64 | 1 | 1 | OV\_\_workspaces_nvSubquadratic-private_examples_well_v2_acoustic_scattering_maze_hyena_gaussian_mask_enabled_True_experiment_dir\_/workspace/results_lr_0.001_num_nodes_1_patch_size_4_run_start_time_1776 | 0.006847195327281952 | 1.6516516517 | 110,000 | -| 2 | 16,384 | 16 | 1 | 1 | OV\_\_workspaces_nvSubquadratic-private_examples_well_v2_acoustic_scattering_maze_hyena_gaussian_mask_bs_16_enabled_True_experiment_dir\_/workspace/results_lr_0.001_num_nodes_1_patch_size_2_run_start_tim | 0.006232426967471838 | 1.5774680204 | 110,000 | +| Patch | Tokens/img | Batch/GPU | GPUs | WandB Run | val/VRMSE | it/s | total steps | +| ----- | ---------- | --------- | ---- | ----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------- | -------------------- | ------------ | ----------- | +| 16 | 256 | 64 | 1 | OV\_\_workspaces_nvSubquadratic-private_examples_well_v2_acoustic_scattering_maze_hyena_enabled_True_experiment_dir\_/workspace/results_lr_0.001_num_nodes_1_run_start_time_1776335625_run_time_limit_hours | 0.0260219257324934 | 5.8873902805 | 110,000 | +| 8 | 1,024 | 64 | 1 | OV\_\_workspaces_nvSubquadratic-private_examples_well_v2_acoustic_scattering_maze_hyena_gaussian_mask_enabled_True_experiment_dir\_/workspace/results_lr_0.001_num_nodes_1_patch_size_8_run_start_time_1776 | 0.00858149304986 | 4.2379411311 | 110,000 | +| 4 | 4,096 | 64 | 1 | OV\_\_workspaces_nvSubquadratic-private_examples_well_v2_acoustic_scattering_maze_hyena_gaussian_mask_enabled_True_experiment_dir\_/workspace/results_lr_0.001_num_nodes_1_patch_size_4_run_start_time_1776 | 0.006847195327281952 | 1.6516516517 | 110,000 | +| 2 | 16,384 | 16 | 1 | OV\_\_workspaces_nvSubquadratic-private_examples_well_v2_acoustic_scattering_maze_hyena_gaussian_mask_bs_16_enabled_True_experiment_dir\_/workspace/results_lr_0.001_num_nodes_1_patch_size_2_run_start_tim | 0.006232426967471838 | 1.5774680204 | 110,000 | diff --git a/examples/well/v2/active_matter/TRACKER.md b/examples/well/v2/active_matter/TRACKER.md index d3f46746..39c4a78e 100644 --- a/examples/well/v2/active_matter/TRACKER.md +++ b/examples/well/v2/active_matter/TRACKER.md @@ -8,7 +8,9 @@ ablated across patch sizes (2, 4, 8, 16). All runs share the same training recipe (110k iterations, AdamW lr=5e-3, wd=1e-4, cosine schedule with 5% warmup, bf16-mixed, grad_clip=1.0). -Default batch size is **64** on **1 GPU**. +Default batch size is **64** on **1 GPU**. We ended up not using gradient +accumulation for patch_size 2 because the training curve was stable at batch +size 16 and we traded off for speed. ## Model sizes @@ -35,10 +37,6 @@ interpolators from `net.in_proj_cfg.patch_size`. Only one override is needed: net.in_proj_cfg.patch_size=P ``` -If OOM at batch 64, add: -`dataset.batch_size= trainer.accumulate_grad_batches=` -(keep `micro_batch * accum = 64`). - ## Run order 1. CNextU-net (baseline, no patch size) @@ -47,35 +45,24 @@ If OOM at batch 64, add: ## Results — CNextU-net (baseline) -| Batch/GPU | Grad Accum | GPUs | WandB Run | val/VRMSE | it/s | total steps | -| --------- | ---------- | ---- | ----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------- | -------------------- | ---- | ----------- | -| 64 | 1 | 1 | OV\_\_workspaces_nvSubquadratic-private_examples_well_v2_active_matter_cfg_unet_convnext_enabled_True_experiment_dir\_/workspace/results_lr_0.005_num_nodes_1_run_start_time_1775922336_run_time_limit_hour | 0.034676797688007355 | 1.39 | 110,000 | +| Batch/GPU | GPUs | WandB Run | val/VRMSE | it/s | total steps | +| --------- | ---- | ----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------- | -------------------- | ---- | ----------- | +| 64 | 1 | OV\_\_workspaces_nvSubquadratic-private_examples_well_v2_active_matter_cfg_unet_convnext_enabled_True_experiment_dir\_/workspace/results_lr_0.005_num_nodes_1_run_start_time_1775922336_run_time_limit_hour | 0.034676797688007355 | 1.39 | 110,000 | ## Results — Attention -| Patch | Tokens/img | Batch/GPU | Grad Accum | GPUs | WandB Run | val/VRMSE | it/s | total steps | -| ----- | ---------- | --------- | ---------- | ---- | ----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------- | -------------------- | ---- | ----------- | -| 16 | 256 | 64 | 1 | 1 | OV\_\_workspaces_nvSubquadratic-private_examples_well_v2_active_matter_attention_enabled_True_experiment_dir\_/workspace/results_lr_0.005_num_nodes_1_patch_size_16_run_start_time_1775918134_run_time_limi | 0.06180257350206375 | 1.40 | 110,000 | -| 8 | 1,024 | 64 | 1 | 1 | OV\_\_workspaces_nvSubquadratic-private_examples_well_v2_active_matter_attention_enabled_True_experiment_dir\_/workspace/results_lr_0.005_num_nodes_1_patch_size_8_run_start_time_1775919205_run_time_limit | 0.058588799089193344 | 1.43 | 110,000 | -| 4 | 4,096 | 64 | 1 | 1 | OV\_\_workspaces_nvSubquadratic-private_examples_well_v2_active_matter_attention_enabled_True_experiment_dir\_/workspace/results_lr_0.005_num_nodes_1_patch_size_4_run_start_time_1775922307_run_time_limit | 0.06164788082242012 | 1.33 | 102,480 | -| 2 | 16,384 | 32 | 2 | 1 | OV\_\_workspaces_nvSubquadratic-private_examples_well_v2_active_matter_attention_bs_32_enabled_True_experiment_dir\_/workspace/results_lr_0.005_num_nodes_1_patch_size_2_run_start_time_1775922311_run_time | 0.09142962843179704 | 0.43 | 36,605 | +| Patch | Tokens/img | Batch/GPU | GPUs | WandB Run | val/VRMSE | it/s | total steps | +| ----- | ---------- | --------- | ---- | ----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------- | -------------------- | ---- | ----------- | +| 16 | 256 | 64 | 1 | OV\_\_workspaces_nvSubquadratic-private_examples_well_v2_active_matter_attention_enabled_True_experiment_dir\_/workspace/results_lr_0.005_num_nodes_1_patch_size_16_run_start_time_1775918134_run_time_limi | 0.06180257350206375 | 1.40 | 110,000 | +| 8 | 1,024 | 64 | 1 | OV\_\_workspaces_nvSubquadratic-private_examples_well_v2_active_matter_attention_enabled_True_experiment_dir\_/workspace/results_lr_0.005_num_nodes_1_patch_size_8_run_start_time_1775919205_run_time_limit | 0.058588799089193344 | 1.43 | 110,000 | +| 4 | 4,096 | 64 | 1 | OV\_\_workspaces_nvSubquadratic-private_examples_well_v2_active_matter_attention_enabled_True_experiment_dir\_/workspace/results_lr_0.005_num_nodes_1_patch_size_4_run_start_time_1775922307_run_time_limit | 0.06164788082242012 | 1.33 | 102,480 | +| 2 | 16,384 | 16 | 1 | OV\_\_workspaces_nvSubquadratic-private_examples_well_v2_active_matter_attention_bs_32_enabled_True_experiment_dir\_/workspace/results_lr_0.005_num_nodes_1_patch_size_2_run_start_time_1775922311_run_time | 0.09142962843179704 | 0.43 | 36,605 | ## Results — Hyena + Gaussian Mask -| Patch | Tokens/img | Batch/GPU | Grad Accum | GPUs | WandB Run | val/VRMSE | it/s | total steps | -| ----- | ---------- | --------- | ---------- | ---- | ------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------ | -------------------- | ---- | ----------- | -| 16 | 256 | 64 | 1 | 1 | OV\_\_workspaces_nvSubquadratic-private_examples_well_v2_active_matter_hyena_gaussian_mask_enabled_True_experiment_dir\_/workspace/results_lr_0.005_num_nodes_1_patch_size_16_run_start_time_1775914506_run | 0.012452345341444016 | 1.38 | 110,000 | -| 8 | 1,024 | 64 | 1 | 1 | OV\_\_workspaces_nvSubquadratic-private_examples_well_v2_active_matter_hyena_gaussian_mask_enabled_True_experiment_dir\_/workspace/results_lr_0.005_num_nodes_1_patch_size_8_run_start_time_1775918046_run\_ | 0.007334186229854822 | 1.39 | 110,000 | -| 4 | 4,096 | 64 | 1 | 1 | OV\_\_workspaces_nvSubquadratic-private_examples_well_v2_active_matter_hyena_gaussian_mask_enabled_True_experiment_dir\_/workspace/results_lr_0.005_num_nodes_1_patch_size_4_run_start_time_1775918057_run\_ | 0.007981576956808567 | 1.36 | 110,000 | -| 2 | 16,384 | 16 | 1 | 1 | OV\_\_workspaces_nvSubquadratic-private_examples_well_v2_active_matter_hyena_gaussian_mask_bs_16_enabled_True_experiment_dir\_/workspace/results_lr_0.005_num_nodes_1_patch_size_2_run_start_time_177587628 | 0.00703924847766757 | 2.63 | 110,000 | - -## Memory Probing Notes (1×H100 80GB, with torch.compile max-autotune-no-cudagraphs) - -| Patch | Batch 64 | Batch 32 | Batch 16 | -| ----- | ---------- | ---------------- | ------------- | -| 16 | OK (both) | — | — | -| 8 | OK (both) | — | — | -| 4 | OK (both) | — | — | -| 2 | OOM (both) | Attn OK, H+G OOM | **OK (both)** | - -Grad accumulation set so `batch * accum = 64`. +| Patch | Tokens/img | Batch/GPU | GPUs | WandB Run | val/VRMSE | it/s | total steps | +| ----- | ---------- | --------- | ---- | ------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------ | -------------------- | ---- | ----------- | +| 16 | 256 | 64 | 1 | OV\_\_workspaces_nvSubquadratic-private_examples_well_v2_active_matter_hyena_gaussian_mask_enabled_True_experiment_dir\_/workspace/results_lr_0.005_num_nodes_1_patch_size_16_run_start_time_1775914506_run | 0.012452345341444016 | 1.38 | 110,000 | +| 8 | 1,024 | 64 | 1 | OV\_\_workspaces_nvSubquadratic-private_examples_well_v2_active_matter_hyena_gaussian_mask_enabled_True_experiment_dir\_/workspace/results_lr_0.005_num_nodes_1_patch_size_8_run_start_time_1775918046_run\_ | 0.007334186229854822 | 1.39 | 110,000 | +| 4 | 4,096 | 64 | 1 | OV\_\_workspaces_nvSubquadratic-private_examples_well_v2_active_matter_hyena_gaussian_mask_enabled_True_experiment_dir\_/workspace/results_lr_0.005_num_nodes_1_patch_size_4_run_start_time_1775918057_run\_ | 0.007981576956808567 | 1.36 | 110,000 | +| 2 | 16,384 | 16 | 1 | OV\_\_workspaces_nvSubquadratic-private_examples_well_v2_active_matter_hyena_gaussian_mask_bs_16_enabled_True_experiment_dir\_/workspace/results_lr_0.005_num_nodes_1_patch_size_2_run_start_time_177587628 | 0.00703924847766757 | 2.63 | 110,000 | diff --git a/examples/well/v2/gray_scott_reaction_diffusion/TRACKER.md b/examples/well/v2/gray_scott_reaction_diffusion/TRACKER.md new file mode 100644 index 00000000..9181ed2c --- /dev/null +++ b/examples/well/v2/gray_scott_reaction_diffusion/TRACKER.md @@ -0,0 +1,51 @@ +# Gray Scott Reaction Diffusion v2 Patch-Size Ablation + +## Goal + +Compare **CNextU-net** (baseline), **Attention**, and **Hyena + Gaussian mask** +on `gray_scott_reaction_diffusion` (128x128). Attention and Hyena+G are +ablated across patch sizes (2, 4, 8, 16). + +All runs share the same training recipe (24 hours, AdamW lr=5e-3, cosine schedule with 5% warmup, bf16-mixed, grad_clip=1.0). +Default batch size is **64** on **1 GPU**. We ended up not using gradient +accumulation for patch_size 2 because the training curve was stable at batch +size 16 and we traded off for speed. + +## Patch-size overrides (Attention & Hyena+G only) + +Stride, out_proj patch_size, and kernel `L_cache` are derived via OmegaConf +interpolators from `net.in_proj_cfg.patch_size`. Only one override is needed: + +``` +net.in_proj_cfg.patch_size=P +``` + +## Run order + +1. CNextU-net (baseline, no patch size) +1. Attention p16 → p8 → p4 → p2 +1. Hyena+G p16 → p8 → p4 → p2 + +## Results — CNextU-net (baseline) + +| Batch/GPU | GPUs | WandB Run | val/VRMSE | it/s | total steps | +| --------- | ---- | --------- | -------------------- | ---- | ----------- | +| 256 | 1 | | 0.034676797688007355 | 1.39 | 110,000 | + +## Results — Attention + +| Patch | Tokens/img | Batch/GPU | GPUs | WandB Run | val/VRMSE | it/s | total steps | +| ----- | ---------- | --------- | ---- | --------- | -------------------- | ---- | ----------- | +| 16 | 64 | 256 | 1 | | 0.06180257350206375 | 1.40 | 110,000 | +| 8 | 256 | 256 | 1 | | 0.058588799089193344 | 1.43 | 110,000 | +| 4 | 1,024 | 256 | 1 | | 0.06164788082242012 | 1.33 | 102,480 | +| 2 | 4,096 | 256 | 1 | | 0.09142962843179704 | 0.43 | 36,605 | + +## Results — Hyena + Gaussian Mask + +| Patch | Tokens/img | Batch/GPU | GPUs | WandB Run | val/VRMSE | it/s | total steps | +| ----- | ---------- | --------- | ---- | --------- | -------------------- | ---- | ----------- | +| 16 | 64 | 256 | 1 | | 0.012452345341444016 | 1.38 | 110,000 | +| 8 | 256 | 256 | 1 | | 0.007334186229854822 | 1.39 | 110,000 | +| 4 | 1,024 | 256 | 1 | | 0.007981576956808567 | 1.36 | 110,000 | +| 2 | 4,096 | 1256 | 1 | | 0.00703924847766757 | 2.63 | 110,000 | diff --git a/examples/well/v2/gray_scott_reaction_diffusion/attention.py b/examples/well/v2/gray_scott_reaction_diffusion/attention.py new file mode 100644 index 00000000..2d414faf --- /dev/null +++ b/examples/well/v2/gray_scott_reaction_diffusion/attention.py @@ -0,0 +1,104 @@ +"""Attention config for gray_scott_reaction_diffusion (v2). + +Uses a ResidualNetwork with multi-head self-attention (QKV + RoPE) as the +sequence mixer. With patch_size=16 the effective sequence resolution is 16×16. + +Patch-size CLI override +----------------------- +Only ``net.in_proj_cfg.patch_size=P`` is needed; stride and out_proj patch_size +are derived via OmegaConf interpolators. +""" + +import torch + +from examples.well.v2.active_matter._base import ( + DATA_DIM, + IN_CHANNELS, + OUT_CHANNELS, + get_base_config, +) +from experiments.default_cfg import ExperimentConfig +from nvsubquadratic.lazy_config import LazyConfig +from nvsubquadratic.modules.attention import Attention +from nvsubquadratic.modules.mlp import MLP +from nvsubquadratic.modules.patchify import Patchify, Unpatchify +from nvsubquadratic.modules.residual_block import ResidualBlock +from nvsubquadratic.modules.rms_norm import RMSNorm +from nvsubquadratic.modules.sequence_mixer import QKVSequenceMixer +from nvsubquadratic.networks.general_purpose_resnet import ResidualNetwork +from nvsubquadratic.utils.init import partial_wang_init_fn_with_num_layers, small_init + + +# ─── Model hyperparameters ──────────────────────────────────────────────────── +NUM_HIDDEN_CHANNELS = 384 +NUM_BLOCKS = 12 +NUM_HEADS = 6 +PATCH_SIZE = 16 + +DROPOUT_IN_RATE = 0.0 +DROPOUT_RATE = 0.0 + + +def get_config() -> ExperimentConfig: + """Build Attention experiment config for active_matter.""" + config = get_base_config() + + config.compile = True + config.compile_mode = "max-autotune-no-cudagraphs" + + norm_cfg = LazyConfig(RMSNorm)(dim=NUM_HIDDEN_CHANNELS) + + config.net = LazyConfig(ResidualNetwork)( + in_channels=IN_CHANNELS, + out_channels=OUT_CHANNELS, + num_blocks=NUM_BLOCKS, + hidden_dim=NUM_HIDDEN_CHANNELS, + data_dim=DATA_DIM, + in_proj_cfg=LazyConfig(Patchify)( + in_features=IN_CHANNELS, + out_features=NUM_HIDDEN_CHANNELS, + data_dim=DATA_DIM, + patch_size=PATCH_SIZE, + stride="${net.in_proj_cfg.patch_size}", + ), + out_proj_cfg=LazyConfig(Unpatchify)( + in_features=NUM_HIDDEN_CHANNELS, + out_features=OUT_CHANNELS, + data_dim=DATA_DIM, + patch_size="${net.in_proj_cfg.patch_size}", + stride="${net.in_proj_cfg.patch_size}", + ), + norm_cfg=norm_cfg, + block_cfg=LazyConfig(ResidualBlock)( + sequence_mixer_cfg=LazyConfig(QKVSequenceMixer)( + hidden_dim=NUM_HIDDEN_CHANNELS, + mixer_cfg=LazyConfig(Attention)( + hidden_dim=NUM_HIDDEN_CHANNELS, + num_heads=NUM_HEADS, + apply_qk_norm=True, + use_rope=True, + is_causal=False, + attn_dropout=0.0, + rope_base=10000.0, + ), + init_method_in=small_init, + init_method_out=partial_wang_init_fn_with_num_layers(num_layers=NUM_BLOCKS), + ), + sequence_mixer_norm_cfg=norm_cfg, + condition_mixer_cfg=LazyConfig(torch.nn.Identity)(), + condition_mixer_norm_cfg=LazyConfig(torch.nn.Identity)(), + mlp_cfg=LazyConfig(MLP)( + dim=NUM_HIDDEN_CHANNELS, + activation="glu", + expansion_factor=1.0, + dropout_cfg=LazyConfig(torch.nn.Dropout)(p=DROPOUT_RATE), + init_method_in=small_init, + init_method_out=partial_wang_init_fn_with_num_layers(num_layers=NUM_BLOCKS), + ), + mlp_norm_cfg=norm_cfg, + dropout_cfg=LazyConfig(torch.nn.Dropout)(p=DROPOUT_RATE), + ), + dropout_in_cfg=LazyConfig(torch.nn.Dropout)(p=DROPOUT_IN_RATE), + ) + + return config diff --git a/examples/well/v2/gray_scott_reaction_diffusion/hyena.py b/examples/well/v2/gray_scott_reaction_diffusion/hyena.py new file mode 100644 index 00000000..d8013f8c --- /dev/null +++ b/examples/well/v2/gray_scott_reaction_diffusion/hyena.py @@ -0,0 +1,141 @@ +"""Hyena config for gray_scott_reaction_diffusion (v2). + +Uses a ResidualNetwork with Hyena (QKV + CKConv global conv) as the +sequence mixer. Circular FFT padding matches the dataset's periodic +boundary conditions. With patch_size=16 the effective sequence +resolution is 16×16. + +Patch-size CLI override +----------------------- +Only ``net.in_proj_cfg.patch_size=P`` is needed; stride, out_proj patch_size, +and kernel L_cache are derived via OmegaConf interpolators. +""" + +import torch + +from examples.well.v2.gray_scott_reaction_diffusion._base import ( + DATA_DIM, + IN_CHANNELS, + OUT_CHANNELS, + get_base_config, +) +from experiments.default_cfg import ExperimentConfig +from nvsubquadratic.lazy_config import LazyConfig +from nvsubquadratic.modules.ckconv_nd import CKConvND +from nvsubquadratic.modules.hyena_nd import Hyena +from nvsubquadratic.modules.kernels_nd import SIRENKernelND +from nvsubquadratic.modules.mlp import MLP +from nvsubquadratic.modules.patchify import Patchify, Unpatchify +from nvsubquadratic.modules.residual_block import ResidualBlock +from nvsubquadratic.modules.rms_norm import RMSNorm +from nvsubquadratic.modules.sequence_mixer import QKVSequenceMixer +from nvsubquadratic.networks.general_purpose_resnet import ResidualNetwork +from nvsubquadratic.utils.init import partial_wang_init_fn_with_num_layers, small_init +from nvsubquadratic.utils.qk_norm import L2Norm + + +# ─── Model hyperparameters ──────────────────────────────────────────────────── +NUM_HIDDEN_CHANNELS = 384 +NUM_BLOCKS = 12 +PATCH_SIZE = 16 + +DROPOUT_IN_RATE = 0.0 +DROPOUT_RATE = 0.0 +GRID_TYPE = "single" +FFT_PADDING = "circular" # periodic boundary conditions +OMEGA_0 = 30.0 + + +def get_config() -> ExperimentConfig: + """Build Hyena experiment config for active_matter.""" + config = get_base_config() + + config.compile = True + config.compile_mode = "max-autotune-no-cudagraphs" + + norm_cfg = LazyConfig(RMSNorm)(dim=NUM_HIDDEN_CHANNELS) + + config.net = LazyConfig(ResidualNetwork)( + in_channels=IN_CHANNELS, + out_channels=OUT_CHANNELS, + num_blocks=NUM_BLOCKS, + hidden_dim=NUM_HIDDEN_CHANNELS, + data_dim=DATA_DIM, + in_proj_cfg=LazyConfig(Patchify)( + in_features=IN_CHANNELS, + out_features=NUM_HIDDEN_CHANNELS, + data_dim=DATA_DIM, + patch_size=PATCH_SIZE, + stride="${net.in_proj_cfg.patch_size}", + ), + out_proj_cfg=LazyConfig(Unpatchify)( + in_features=NUM_HIDDEN_CHANNELS, + out_features=OUT_CHANNELS, + data_dim=DATA_DIM, + patch_size="${net.in_proj_cfg.patch_size}", + stride="${net.in_proj_cfg.patch_size}", + ), + norm_cfg=norm_cfg, + block_cfg=LazyConfig(ResidualBlock)( + sequence_mixer_cfg=LazyConfig(QKVSequenceMixer)( + hidden_dim=NUM_HIDDEN_CHANNELS, + mixer_cfg=LazyConfig(Hyena)( + global_conv_cfg=LazyConfig(CKConvND)( + data_dim=DATA_DIM, + hidden_dim=NUM_HIDDEN_CHANNELS, + fft_padding=FFT_PADDING, + use_fp16_fft=False, + kernel_cfg=LazyConfig(SIRENKernelND)( + data_dim=DATA_DIM, + out_dim=NUM_HIDDEN_CHANNELS, + mlp_hidden_dim=64, + num_layers=3, + embedding_dim=64, + omega_0=OMEGA_0, + L_cache="${eval:'256 // ${net.in_proj_cfg.patch_size}'}", + use_bias=True, + hidden_omega_0=1.0, + ), + mask_cfg=LazyConfig(torch.nn.Identity)(), + grid_type=GRID_TYPE, + ), + short_conv_cfg=LazyConfig(torch.nn.Conv2d)( + in_channels=3 * NUM_HIDDEN_CHANNELS, + out_channels=3 * NUM_HIDDEN_CHANNELS, + kernel_size=3, + groups=3 * NUM_HIDDEN_CHANNELS, + padding=1, + bias=False, + ), + gate_nonlinear_cfg=LazyConfig(torch.nn.SiLU)(), + gate_nonlinear_2_cfg=LazyConfig(torch.nn.Sigmoid)(), + pixelhyena_norm_cfg=LazyConfig(RMSNorm)( + dim=NUM_HIDDEN_CHANNELS, + ), + output_norm_cfg=LazyConfig(RMSNorm)( + dim=NUM_HIDDEN_CHANNELS, + ), + qk_norm_cfg=LazyConfig(L2Norm)(), + use_rope=False, + ), + init_method_in=small_init, + init_method_out=partial_wang_init_fn_with_num_layers(num_layers=NUM_BLOCKS), + ), + sequence_mixer_norm_cfg=norm_cfg, + condition_mixer_cfg=LazyConfig(torch.nn.Identity)(), + condition_mixer_norm_cfg=LazyConfig(torch.nn.Identity)(), + mlp_cfg=LazyConfig(MLP)( + dim=NUM_HIDDEN_CHANNELS, + activation="glu", + expansion_factor=1.0, + dropout_cfg=LazyConfig(torch.nn.Dropout)(p=DROPOUT_RATE), + init_method_in=small_init, + init_method_out=partial_wang_init_fn_with_num_layers(num_layers=NUM_BLOCKS), + ), + mlp_norm_cfg=norm_cfg, + dropout_cfg=LazyConfig(torch.nn.Dropout)(p=DROPOUT_RATE), + ), + dropout_in_cfg=LazyConfig(torch.nn.Dropout)(p=DROPOUT_IN_RATE), + ) + + return config diff --git a/examples/well/v2/gray_scott_reaction_diffusion/hyena_gaussian_mask.py b/examples/well/v2/gray_scott_reaction_diffusion/hyena_gaussian_mask.py new file mode 100644 index 00000000..87e7ae43 --- /dev/null +++ b/examples/well/v2/gray_scott_reaction_diffusion/hyena_gaussian_mask.py @@ -0,0 +1,31 @@ +"""Hyena config with Gaussian modulation mask for gray_scott_reaction_diffusion (v2). + +Identical to ``cfg_hyena.py`` but replaces the ``nn.Identity`` mask with a +``GaussianModulationND`` mask on the CKConv global convolution kernel. +""" + +from examples.well.v2.active_matter._base import DATA_DIM +from examples.well.v2.active_matter.hyena import NUM_HIDDEN_CHANNELS +from examples.well.v2.active_matter.hyena import get_config as _get_hyena_config +from experiments.callbacks.mask_monitor import MaskMonitorCallback +from experiments.default_cfg import ExperimentConfig +from nvsubquadratic.lazy_config import LazyConfig +from nvsubquadratic.modules.masks_nd import GaussianModulationND + + +def get_config() -> ExperimentConfig: + """Build Hyena + Gaussian mask config for active_matter.""" + config = _get_hyena_config() + + config.net.block_cfg.sequence_mixer_cfg.mixer_cfg.global_conv_cfg.mask_cfg = LazyConfig(GaussianModulationND)( + data_dim=DATA_DIM, + num_channels=NUM_HIDDEN_CHANNELS, + min_attenuation_at_step=0.1, + max_attenuation_at_limit=0.95, + init_extent=1.0, + parametrization="direct", + ) + + config.callbacks.append(LazyConfig(MaskMonitorCallback)(log_every_n_steps=50)) + + return config From 581ff0445d926f2a0a2e7b1930540dda8c8911a3 Mon Sep 17 00:00:00 2001 From: OliviaViessmann Date: Fri, 1 May 2026 13:25:10 +0000 Subject: [PATCH 05/13] before pulling main --- .../TRACKER.md | 40 ++++++++++++++ .../gray_scott_reaction_diffusion/TRACKER.md | 52 ++++++++----------- 2 files changed, 62 insertions(+), 30 deletions(-) create mode 100644 examples/well/v2/euler_multi_quadrants_periodicBC/TRACKER.md diff --git a/examples/well/v2/euler_multi_quadrants_periodicBC/TRACKER.md b/examples/well/v2/euler_multi_quadrants_periodicBC/TRACKER.md new file mode 100644 index 00000000..9174727c --- /dev/null +++ b/examples/well/v2/euler_multi_quadrants_periodicBC/TRACKER.md @@ -0,0 +1,40 @@ +# Euler Multi Quadrants v2 + +## Goal + +Compare **CNextU-net** (baseline), **Attention**, and **Hyena + Gaussian mask** +on `euer_multi_quadrants` (512x512). + +All runs share the same training recipe (24 hours, AdamW, cosine schedule with 5% warmup, bf16-mixed, grad_clip=1.0). + +## Run order + +1. CNextU-net (baseline, no patch size) +1. Attention p2 +1. Hyena+G p2 + +## Results — CNextU-net (baseline) + +| Batch/GPU | GPUs | WandB Run | val/VRMSE | it/s | total steps | +| --------- | ---- | ---------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------- | ----------------- | ---- | ----------- | +| 256 | 1 | OV\_\_workspaces_nvSubquadratic-private_examples_well_v2_gray_scott_reaction_diffusion_cfg_unet_convnext_enabled_True_experiment_dir\_/workspace/results_lr_0.0001_num_nodes_1_run_start_time_1777331397_r | 0.231915682554245 | 6.00 | 110,000 | + +## Results — Attention + +Total params: 13,747,979 + +| Patch | Tokens/img | Batch/GPU | GPUs | WandB Run | val/VRMSE | it/s | total steps | +| ----- | ---------- | --------- | ---- | ----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------- | -------------------- | ---- | ----------- | +| 8 | 2,048 | 256 | 1 | OV\_\_workspaces_nvSubquadratic-private_examples_well_v2_gray_scott_reaction_diffusion_attention_enabled_True_experiment_dir\_/workspace/results_lr_0.0005_num_nodes_1_patch_size_8_run_start_time_17773947 | 0.052010852843523026 | 1.4 | 110,000 | +| 4 | 4,096 | 256 | 1 | | 0.007981576956808567 | 1.36 | 110,000 | +| 2 | 8,192 | 16 | 1 | | 0.00703924847766757 | 2.63 | 110,000 | + +## Results — Hyena + Gaussian Mask + +Total params: 14,297,099 + +| Patch | Tokens/img | Batch/GPU | GPUs | WandB Run | val/VRMSE | it/s | total steps | +| ----- | ---------- | --------- | ---- | ----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------- | -------------------- | ---- | ----------- | +| 8 | 2,048 | 256 | 1 | OV\_\_workspaces_nvSubquadratic-private_examples_well_v2_gray_scott_reaction_diffusion_hyena_gaussian_mask_enabled_True_experiment_dir\_/workspace/results_lr_0.0005_num_nodes_1_patch_size_8_run_start_tim | 0.009229527786374092 | 1.36 | 88200 | +| 4 | 4,096 | 256 | 1 | | 0.007981576956808567 | 1.36 | 110,000 | +| 2 | 8,192 | 16 | 1 | | 0.00703924847766757 | 2.63 | 110,000 | diff --git a/examples/well/v2/gray_scott_reaction_diffusion/TRACKER.md b/examples/well/v2/gray_scott_reaction_diffusion/TRACKER.md index 1d052d7c..9b22fe77 100644 --- a/examples/well/v2/gray_scott_reaction_diffusion/TRACKER.md +++ b/examples/well/v2/gray_scott_reaction_diffusion/TRACKER.md @@ -4,48 +4,40 @@ Compare **CNextU-net** (baseline), **Attention**, and **Hyena + Gaussian mask** on `gray_scott_reaction_diffusion` (128x128). Attention and Hyena+G are -ablated across patch sizes (2, 4, 8, 16). +ablated across patch sizes (2, 4, 8). -All runs share the same training recipe (24 hours, AdamW lr=1e-4, cosine schedule with 5% warmup, bf16-mixed, grad_clip=1.0). +All runs share the same training recipe (24 hours, AdamW, cosine schedule with 5% warmup, bf16-mixed, grad_clip=1.0). Default batch size is **64** on **1 GPU**. We ended up not using gradient -accumulation for patch_size 2 because the training curve was stable at batch -size 16 and we traded off for speed. - -## Patch-size overrides (Attention & Hyena+G only) - -Stride, out_proj patch_size, and kernel `L_cache` are derived via OmegaConf -interpolators from `net.in_proj_cfg.patch_size`. Only one override is needed: - -``` -net.in_proj_cfg.patch_size=P -``` +accumulation for patch_size 2 because the training curve was stable and we traded off for speed. ## Run order 1. CNextU-net (baseline, no patch size) -1. Attention p16 → p8 → p4 → p2 -1. Hyena+G p16 → p8 → p4 → p2 +1. Attention → p8 → p4 → p2 +1. Hyena+G → p8 → p4 → p2 ## Results — CNextU-net (baseline) -| Batch/GPU | GPUs | WandB Run | val/VRMSE | it/s | total steps | -| --------- | ---- | --------- | -------------------- | ---- | ----------- | -| 256 | 1 | | 0.034676797688007355 | 1.39 | 110,000 | +| Batch/GPU | GPUs | WandB Run | val/VRMSE | it/s | total steps | +| --------- | ---- | ---------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------- | ----------------- | ---- | ----------- | +| 256 | 1 | OV\_\_workspaces_nvSubquadratic-private_examples_well_v2_gray_scott_reaction_diffusion_cfg_unet_convnext_enabled_True_experiment_dir\_/workspace/results_lr_0.0001_num_nodes_1_run_start_time_1777331397_r | 0.231915682554245 | 6.00 | 110,000 | ## Results — Attention -| Patch | Tokens/img | Batch/GPU | GPUs | WandB Run | val/VRMSE | it/s | total steps | -| ----- | ---------- | --------- | ---- | --------- | -------------------- | ---- | ----------- | -| 16 | 64 | 256 | 1 | | 0.06180257350206375 | 1.40 | 110,000 | -| 8 | 256 | 256 | 1 | | 0.058588799089193344 | 1.43 | 110,000 | -| 4 | 1,024 | 256 | 1 | | 0.06164788082242012 | 1.33 | 102,480 | -| 2 | 4,096 | 256 | 1 | | 0.09142962843179704 | 0.43 | 36,605 | +Total params: 13,747,979 + +| Patch | Tokens/img | Batch/GPU | GPUs | WandB Run | val/VRMSE | it/s | total steps | +| ----- | ---------- | --------- | ---- | ----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------- | -------------------- | ---- | ----------- | +| 8 | 2,048 | 256 | 1 | OV\_\_workspaces_nvSubquadratic-private_examples_well_v2_gray_scott_reaction_diffusion_attention_enabled_True_experiment_dir\_/workspace/results_lr_0.0005_num_nodes_1_patch_size_8_run_start_time_17773947 | 0.052010852843523026 | 1.4 | 110,000 | +| 4 | 4,096 | 256 | 1 | | 0.007981576956808567 | 1.36 | 110,000 | +| 2 | 8,192 | 16 | 1 | | 0.00703924847766757 | 2.63 | 110,000 | ## Results — Hyena + Gaussian Mask -| Patch | Tokens/img | Batch/GPU | GPUs | WandB Run | val/VRMSE | it/s | total steps | -| ----- | ---------- | --------- | ---- | --------- | -------------------- | ---- | ----------- | -| 16 | 64 | 256 | 1 | | 0.012452345341444016 | 1.38 | 110,000 | -| 8 | 256 | 256 | 1 | | 0.007334186229854822 | 1.39 | 110,000 | -| 4 | 1,024 | 256 | 1 | | 0.007981576956808567 | 1.36 | 110,000 | -| 2 | 4,096 | 1256 | 1 | | 0.00703924847766757 | 2.63 | 110,000 | +Total params: 14,297,099 + +| Patch | Tokens/img | Batch/GPU | GPUs | WandB Run | val/VRMSE | it/s | total steps | +| ----- | ---------- | --------- | ---- | ----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------- | -------------------- | ---- | ----------- | +| 8 | 2,048 | 256 | 1 | OV\_\_workspaces_nvSubquadratic-private_examples_well_v2_gray_scott_reaction_diffusion_hyena_gaussian_mask_enabled_True_experiment_dir\_/workspace/results_lr_0.0005_num_nodes_1_patch_size_8_run_start_tim | 0.009229527786374092 | 1.36 | 88200 | +| 4 | 4,096 | 256 | 1 | | 0.007981576956808567 | 1.36 | 110,000 | +| 2 | 8,192 | 16 | 1 | | 0.00703924847766757 | 2.63 | 110,000 | From 785ea4c5e33f313651370965054a41e48b1d1a3e Mon Sep 17 00:00:00 2001 From: OliviaViessmann Date: Mon, 4 May 2026 17:05:25 +0000 Subject: [PATCH 06/13] trackers and flop cpount --- .../euler_multi_quadrants_periodicBC/_base.py | 2 +- .../attention.py | 107 +++++++++++++ .../euler_multi_quadrants_periodicBC/hyena.py | 144 ++++++++++++++++++ .../hyena_gaussian_mask.py | 31 ++++ .../v2/supernova_explosion_64/attention.py | 4 +- .../well/v2/supernova_explosion_64/hyena.py | 4 +- experiments/callbacks/flop_counter.py | 130 ++++++++++++++++ experiments/trainer.py | 3 + 8 files changed, 422 insertions(+), 3 deletions(-) create mode 100644 examples/well/v2/euler_multi_quadrants_periodicBC/attention.py create mode 100644 examples/well/v2/euler_multi_quadrants_periodicBC/hyena.py create mode 100644 examples/well/v2/euler_multi_quadrants_periodicBC/hyena_gaussian_mask.py create mode 100644 experiments/callbacks/flop_counter.py diff --git a/examples/well/v2/euler_multi_quadrants_periodicBC/_base.py b/examples/well/v2/euler_multi_quadrants_periodicBC/_base.py index cbfaa19f..1e1e684b 100644 --- a/examples/well/v2/euler_multi_quadrants_periodicBC/_base.py +++ b/examples/well/v2/euler_multi_quadrants_periodicBC/_base.py @@ -44,7 +44,7 @@ TRAINING_ITERATIONS = 110_000 WARMUP_ITERATIONS_PERCENTAGE = 0.05 -BATCH_SIZE = 24 # configs/data/euler_multi_quadrants_periodicBC.yaml +BATCH_SIZE = 48 # configs/data/euler_multi_quadrants_periodicBC.yaml NUM_WORKERS = 12 GRAD_CLIP = 1.0 PRECISION = "bf16-mixed" diff --git a/examples/well/v2/euler_multi_quadrants_periodicBC/attention.py b/examples/well/v2/euler_multi_quadrants_periodicBC/attention.py new file mode 100644 index 00000000..d11d683e --- /dev/null +++ b/examples/well/v2/euler_multi_quadrants_periodicBC/attention.py @@ -0,0 +1,107 @@ +"""Attention config for euler_multi_quadrants_periodicBC (v2). + +Uses a ResidualNetwork with multi-head self-attention (QKV + RoPE) as the +sequence mixer. With patch_size=8 the effective sequence resolution is 8×8×8. + +Patch-size CLI override +----------------------- +Only ``net.in_proj_cfg.patch_size=P`` is needed; stride and out_proj patch_size +are derived via OmegaConf interpolators. +""" + +import torch + +from examples.well.v2.euler_multi_quadrants_periodicBC._base import ( + DATA_DIM, + IN_CHANNELS, + OUT_CHANNELS, + get_base_config, +) +from experiments.default_cfg import ExperimentConfig +from nvsubquadratic.lazy_config import LazyConfig +from nvsubquadratic.modules.attention import Attention +from nvsubquadratic.modules.mlp import MLP +from nvsubquadratic.modules.patchify import Patchify, Unpatchify +from nvsubquadratic.modules.residual_block import ResidualBlock +from nvsubquadratic.modules.rms_norm import RMSNorm +from nvsubquadratic.modules.sequence_mixer import QKVSequenceMixer +from nvsubquadratic.networks.general_purpose_resnet import ResidualNetwork +from nvsubquadratic.utils.init import partial_wang_init_fn_with_num_layers, small_init + + +# ─── Model hyperparameters ──────────────────────────────────────────────────── +NUM_HIDDEN_CHANNELS = 384 +NUM_BLOCKS = 12 +NUM_HEADS = 8 # head_dim=48; must be divisible by 6 for 3D RoPE +PATCH_SIZE = 8 + +DROPOUT_IN_RATE = 0.0 +DROPOUT_RATE = 0.0 + +GRADIENT_CHECKPOINTING = False + + +def get_config() -> ExperimentConfig: + """Build Attention experiment config for euler_multi_quadrants_periodicBC.""" + config = get_base_config(learning_rate=1e-3, weight_decay=1e-5) + + config.compile = True + config.compile_mode = "max-autotune-no-cudagraphs" + + norm_cfg = LazyConfig(RMSNorm)(dim=NUM_HIDDEN_CHANNELS) + + config.net = LazyConfig(ResidualNetwork)( + in_channels=IN_CHANNELS, + out_channels=OUT_CHANNELS, + num_blocks=NUM_BLOCKS, + hidden_dim=NUM_HIDDEN_CHANNELS, + data_dim=DATA_DIM, + in_proj_cfg=LazyConfig(Patchify)( + in_features=IN_CHANNELS, + out_features=NUM_HIDDEN_CHANNELS, + data_dim=DATA_DIM, + patch_size=PATCH_SIZE, + stride="${net.in_proj_cfg.patch_size}", + ), + out_proj_cfg=LazyConfig(Unpatchify)( + in_features=NUM_HIDDEN_CHANNELS, + out_features=OUT_CHANNELS, + data_dim=DATA_DIM, + patch_size="${net.in_proj_cfg.patch_size}", + stride="${net.in_proj_cfg.patch_size}", + ), + norm_cfg=norm_cfg, + block_cfg=LazyConfig(ResidualBlock)( + sequence_mixer_cfg=LazyConfig(QKVSequenceMixer)( + hidden_dim=NUM_HIDDEN_CHANNELS, + mixer_cfg=LazyConfig(Attention)( + hidden_dim=NUM_HIDDEN_CHANNELS, + num_heads=NUM_HEADS, + apply_qk_norm=True, + use_rope=True, + is_causal=False, + attn_dropout=0.0, + rope_base=10000.0, + ), + init_method_in=small_init, + init_method_out=partial_wang_init_fn_with_num_layers(num_layers=NUM_BLOCKS), + ), + sequence_mixer_norm_cfg=norm_cfg, + condition_mixer_cfg=LazyConfig(torch.nn.Identity)(), + condition_mixer_norm_cfg=LazyConfig(torch.nn.Identity)(), + mlp_cfg=LazyConfig(MLP)( + dim=NUM_HIDDEN_CHANNELS, + activation="glu", + expansion_factor=1.0, + dropout_cfg=LazyConfig(torch.nn.Dropout)(p=DROPOUT_RATE), + init_method_in=small_init, + init_method_out=partial_wang_init_fn_with_num_layers(num_layers=NUM_BLOCKS), + ), + mlp_norm_cfg=norm_cfg, + dropout_cfg=LazyConfig(torch.nn.Dropout)(p=DROPOUT_RATE), + ), + dropout_in_cfg=LazyConfig(torch.nn.Dropout)(p=DROPOUT_IN_RATE), + gradient_checkpointing=GRADIENT_CHECKPOINTING, + ) + + return config diff --git a/examples/well/v2/euler_multi_quadrants_periodicBC/hyena.py b/examples/well/v2/euler_multi_quadrants_periodicBC/hyena.py new file mode 100644 index 00000000..15207d4f --- /dev/null +++ b/examples/well/v2/euler_multi_quadrants_periodicBC/hyena.py @@ -0,0 +1,144 @@ +"""Hyena config for gray_scott_reaction_diffusion (v2). + +Uses a ResidualNetwork with Hyena (QKV + CKConv global conv) as the +sequence mixer. Circular FFT padding matches the dataset's periodic +boundary conditions. With patch_size=16 the effective sequence +resolution is 16×16. + +Patch-size CLI override +----------------------- +Only ``net.in_proj_cfg.patch_size=P`` is needed; stride, out_proj patch_size, +and kernel L_cache are derived via OmegaConf interpolators. +""" + +import torch + +from examples.well.v2.euler_multi_quadrants_periodicBC._base import ( + DATA_DIM, + IN_CHANNELS, + OUT_CHANNELS, + get_base_config, +) +from experiments.default_cfg import ExperimentConfig +from nvsubquadratic.lazy_config import LazyConfig +from nvsubquadratic.modules.ckconv_nd import CKConvND +from nvsubquadratic.modules.hyena_nd import Hyena +from nvsubquadratic.modules.kernels_nd import SIRENKernelND +from nvsubquadratic.modules.mlp import MLP +from nvsubquadratic.modules.patchify import Patchify, Unpatchify +from nvsubquadratic.modules.residual_block import ResidualBlock +from nvsubquadratic.modules.rms_norm import RMSNorm +from nvsubquadratic.modules.sequence_mixer import QKVSequenceMixer +from nvsubquadratic.networks.general_purpose_resnet import ResidualNetwork +from nvsubquadratic.utils.init import partial_wang_init_fn_with_num_layers, small_init +from nvsubquadratic.utils.qk_norm import L2Norm + + +# ─── Model hyperparameters ──────────────────────────────────────────────────── +NUM_HIDDEN_CHANNELS = 384 +NUM_BLOCKS = 12 +PATCH_SIZE = 16 + +DROPOUT_IN_RATE = 0.0 +DROPOUT_RATE = 0.0 +GRID_TYPE = "single" +FFT_PADDING = "circular" # periodic boundary conditions +OMEGA_0 = 30.0 + +GRADIENT_CHECKPOINTING = True + + +def get_config() -> ExperimentConfig: + """Build Hyena experiment config for euler_multi_quadrants_periodicBC.""" + config = get_base_config() + + config.compile = True + config.compile_mode = "max-autotune-no-cudagraphs" + + norm_cfg = LazyConfig(RMSNorm)(dim=NUM_HIDDEN_CHANNELS) + + config.net = LazyConfig(ResidualNetwork)( + in_channels=IN_CHANNELS, + out_channels=OUT_CHANNELS, + num_blocks=NUM_BLOCKS, + hidden_dim=NUM_HIDDEN_CHANNELS, + data_dim=DATA_DIM, + in_proj_cfg=LazyConfig(Patchify)( + in_features=IN_CHANNELS, + out_features=NUM_HIDDEN_CHANNELS, + data_dim=DATA_DIM, + patch_size=PATCH_SIZE, + stride="${net.in_proj_cfg.patch_size}", + ), + out_proj_cfg=LazyConfig(Unpatchify)( + in_features=NUM_HIDDEN_CHANNELS, + out_features=OUT_CHANNELS, + data_dim=DATA_DIM, + patch_size="${net.in_proj_cfg.patch_size}", + stride="${net.in_proj_cfg.patch_size}", + ), + norm_cfg=norm_cfg, + block_cfg=LazyConfig(ResidualBlock)( + sequence_mixer_cfg=LazyConfig(QKVSequenceMixer)( + hidden_dim=NUM_HIDDEN_CHANNELS, + mixer_cfg=LazyConfig(Hyena)( + global_conv_cfg=LazyConfig(CKConvND)( + data_dim=DATA_DIM, + hidden_dim=NUM_HIDDEN_CHANNELS, + fft_padding=FFT_PADDING, + use_fp16_fft=False, + kernel_cfg=LazyConfig(SIRENKernelND)( + data_dim=DATA_DIM, + out_dim=NUM_HIDDEN_CHANNELS, + mlp_hidden_dim=64, + num_layers=3, + embedding_dim=64, + omega_0=OMEGA_0, + L_cache="${eval:'256 // ${net.in_proj_cfg.patch_size}'}", + use_bias=True, + hidden_omega_0=1.0, + ), + mask_cfg=LazyConfig(torch.nn.Identity)(), + grid_type=GRID_TYPE, + ), + short_conv_cfg=LazyConfig(torch.nn.Conv2d)( + in_channels=3 * NUM_HIDDEN_CHANNELS, + out_channels=3 * NUM_HIDDEN_CHANNELS, + kernel_size=3, + groups=3 * NUM_HIDDEN_CHANNELS, + padding=1, + bias=False, + ), + gate_nonlinear_cfg=LazyConfig(torch.nn.SiLU)(), + gate_nonlinear_2_cfg=LazyConfig(torch.nn.Sigmoid)(), + pixelhyena_norm_cfg=LazyConfig(RMSNorm)( + dim=NUM_HIDDEN_CHANNELS, + ), + output_norm_cfg=LazyConfig(RMSNorm)( + dim=NUM_HIDDEN_CHANNELS, + ), + qk_norm_cfg=LazyConfig(L2Norm)(), + use_rope=False, + ), + init_method_in=small_init, + init_method_out=partial_wang_init_fn_with_num_layers(num_layers=NUM_BLOCKS), + ), + sequence_mixer_norm_cfg=norm_cfg, + condition_mixer_cfg=LazyConfig(torch.nn.Identity)(), + condition_mixer_norm_cfg=LazyConfig(torch.nn.Identity)(), + mlp_cfg=LazyConfig(MLP)( + dim=NUM_HIDDEN_CHANNELS, + activation="glu", + expansion_factor=1.0, + dropout_cfg=LazyConfig(torch.nn.Dropout)(p=DROPOUT_RATE), + init_method_in=small_init, + init_method_out=partial_wang_init_fn_with_num_layers(num_layers=NUM_BLOCKS), + ), + mlp_norm_cfg=norm_cfg, + dropout_cfg=LazyConfig(torch.nn.Dropout)(p=DROPOUT_RATE), + ), + dropout_in_cfg=LazyConfig(torch.nn.Dropout)(p=DROPOUT_IN_RATE), + gradient_checkpointing=GRADIENT_CHECKPOINTING, + ) + + return config diff --git a/examples/well/v2/euler_multi_quadrants_periodicBC/hyena_gaussian_mask.py b/examples/well/v2/euler_multi_quadrants_periodicBC/hyena_gaussian_mask.py new file mode 100644 index 00000000..87e7ae43 --- /dev/null +++ b/examples/well/v2/euler_multi_quadrants_periodicBC/hyena_gaussian_mask.py @@ -0,0 +1,31 @@ +"""Hyena config with Gaussian modulation mask for gray_scott_reaction_diffusion (v2). + +Identical to ``cfg_hyena.py`` but replaces the ``nn.Identity`` mask with a +``GaussianModulationND`` mask on the CKConv global convolution kernel. +""" + +from examples.well.v2.active_matter._base import DATA_DIM +from examples.well.v2.active_matter.hyena import NUM_HIDDEN_CHANNELS +from examples.well.v2.active_matter.hyena import get_config as _get_hyena_config +from experiments.callbacks.mask_monitor import MaskMonitorCallback +from experiments.default_cfg import ExperimentConfig +from nvsubquadratic.lazy_config import LazyConfig +from nvsubquadratic.modules.masks_nd import GaussianModulationND + + +def get_config() -> ExperimentConfig: + """Build Hyena + Gaussian mask config for active_matter.""" + config = _get_hyena_config() + + config.net.block_cfg.sequence_mixer_cfg.mixer_cfg.global_conv_cfg.mask_cfg = LazyConfig(GaussianModulationND)( + data_dim=DATA_DIM, + num_channels=NUM_HIDDEN_CHANNELS, + min_attenuation_at_step=0.1, + max_attenuation_at_limit=0.95, + init_extent=1.0, + parametrization="direct", + ) + + config.callbacks.append(LazyConfig(MaskMonitorCallback)(log_every_n_steps=50)) + + return config diff --git a/examples/well/v2/supernova_explosion_64/attention.py b/examples/well/v2/supernova_explosion_64/attention.py index f32ef28d..a7085a62 100644 --- a/examples/well/v2/supernova_explosion_64/attention.py +++ b/examples/well/v2/supernova_explosion_64/attention.py @@ -38,6 +38,8 @@ DROPOUT_IN_RATE = 0.0 DROPOUT_RATE = 0.0 +GRADIENT_CHECKPOINTING = False + def get_config() -> ExperimentConfig: """Build Attention experiment config for supernova_explosion_64.""" @@ -99,7 +101,7 @@ def get_config() -> ExperimentConfig: dropout_cfg=LazyConfig(torch.nn.Dropout)(p=DROPOUT_RATE), ), dropout_in_cfg=LazyConfig(torch.nn.Dropout)(p=DROPOUT_IN_RATE), - gradient_checkpointing=False, + gradient_checkpointing=GRADIENT_CHECKPOINTING, ) return config diff --git a/examples/well/v2/supernova_explosion_64/hyena.py b/examples/well/v2/supernova_explosion_64/hyena.py index 58e214e3..05f8f302 100644 --- a/examples/well/v2/supernova_explosion_64/hyena.py +++ b/examples/well/v2/supernova_explosion_64/hyena.py @@ -45,6 +45,8 @@ FFT_PADDING = "zero" # open (non-periodic) boundary conditions OMEGA_0 = 30.0 +GRADIENT_CHECKPOINTING = False + def get_config() -> ExperimentConfig: """Build Hyena experiment config for supernova_explosion_64.""" @@ -136,7 +138,7 @@ def get_config() -> ExperimentConfig: dropout_cfg=LazyConfig(torch.nn.Dropout)(p=DROPOUT_RATE), ), dropout_in_cfg=LazyConfig(torch.nn.Dropout)(p=DROPOUT_IN_RATE), - gradient_checkpointing=False, + gradient_checkpointing=GRADIENT_CHECKPOINTING, ) return config diff --git a/experiments/callbacks/flop_counter.py b/experiments/callbacks/flop_counter.py new file mode 100644 index 00000000..8bd4948b --- /dev/null +++ b/experiments/callbacks/flop_counter.py @@ -0,0 +1,130 @@ +"""Callback to measure per-step FLOPs at fit start. + +Runs one real ``training_step`` (forward + backward) on a single batch under +``torch.utils.flop_counter.FlopCounterMode``, and logs the counts to the +attached logger as a one-shot metric: + + ``flops/fwd`` – forward FLOPs for one batch (per rank) + ``flops/bwd`` – backward FLOPs for one batch (per rank) + ``flops/step`` – total per-step FLOPs (fwd + bwd, per rank) + ``flops/step_global`` – ``flops/step * world_size`` + +``FlopCounterMode`` is a ``TorchDispatchMode`` and natively counts ``aten`` +ops including ``fft_*`` / ``conv_*`` / ``matmul`` / SDPA, which matters here +because Hyena/CKConv use FFT-based long convolutions that ``fvcore`` / +``torchinfo`` do not see. + +Notes: +----- +* If ``pl_module.network`` was wrapped by ``torch.compile`` (i.e. it has a + ``_orig_mod`` attribute), the un-compiled module is swapped in for the + measurement. Dispatch modes do not always observe ops inside compiled + graphs, so this gives a reliable count. +* The measurement runs on every rank (so the DDP gradient all-reduce stays + synchronized during ``loss.backward()``) but only rank 0 logs. +* Gradients are zeroed and the original ``pl_module.network`` is restored + before training starts, so this callback is non-invasive. +* With ``gradient_checkpointing=True`` the backward also re-runs the + checkpointed forward chunks; those re-runs are counted in ``flops/bwd``, + which is the correct accounting for true per-step training cost. +""" + +from __future__ import annotations + +import pytorch_lightning as pl +from torch.utils.flop_counter import FlopCounterMode + + +class FlopCounterCallback(pl.Callback): + """Measure FLOPs for one fwd+bwd training step at fit start.""" + + def __init__(self) -> None: # noqa: D107 + super().__init__() + self._done = False + + def on_train_start(self, trainer: pl.Trainer, pl_module: pl.LightningModule) -> None: # noqa: D102 + if self._done: + return + self._done = True + + # Fetch one training batch. Prefer the trainer's already-built + # dataloader; fall back to the datamodule. + batch = self._fetch_batch(trainer) + if batch is None: + if trainer.is_global_zero: + print("[flop_counter] could not fetch a training batch; skipping FLOP measurement.") + return + + batch = pl_module._apply_batch_transfer_handler(batch, pl_module.device, dataloader_idx=0) + + # If torch.compile wrapped the network, run the count against the + # un-compiled module so the dispatch mode sees every op. + compiled_net = getattr(pl_module, "network", None) + original_net = None + if compiled_net is not None and hasattr(compiled_net, "_orig_mod"): + original_net = compiled_net + pl_module.network = compiled_net._orig_mod + + # Suppress ``self.log`` calls that wrappers normally make inside + # ``training_step`` — we are not inside the actual training loop hook. + original_log = pl_module.log + pl_module.log = lambda *a, **k: None + + try: + with FlopCounterMode(display=False) as fc: + with trainer.precision_plugin.forward_context(): + output = pl_module.training_step(batch, 0) + fwd_flops = fc.get_total_flops() + + loss = output["loss"] if isinstance(output, dict) else output + loss.backward() + total_flops = fc.get_total_flops() + except Exception as e: + pl_module.log = original_log + if original_net is not None: + pl_module.network = original_net + pl_module.zero_grad(set_to_none=True) + if trainer.is_global_zero: + print(f"[flop_counter] measurement failed ({type(e).__name__}: {e}); skipping.") + return + + pl_module.log = original_log + if original_net is not None: + pl_module.network = original_net + pl_module.zero_grad(set_to_none=True) + + bwd_flops = total_flops - fwd_flops + world_size = trainer.world_size if trainer.world_size else 1 + + if trainer.is_global_zero: + payload = { + "flops/fwd": float(fwd_flops), + "flops/bwd": float(bwd_flops), + "flops/step": float(total_flops), + "flops/step_global": float(total_flops * world_size), + } + bs = getattr(getattr(trainer, "datamodule", None), "batch_size", "?") + print( + f"[flop_counter] per-rank batch_size={bs} " + f"fwd={fwd_flops / 1e12:.3f} TFLOPs " + f"bwd={bwd_flops / 1e12:.3f} TFLOPs " + f"step={total_flops / 1e12:.3f} TFLOPs " + f"(global step ≈ {total_flops * world_size / 1e12:.3f} TFLOPs across {world_size} rank(s))" + ) + if trainer.logger is not None: + trainer.logger.log_metrics(payload, step=0) + + @staticmethod + def _fetch_batch(trainer: pl.Trainer): + dl = getattr(trainer, "train_dataloader", None) + if dl is None and getattr(trainer, "datamodule", None) is not None: + try: + dl = trainer.datamodule.train_dataloader() + except Exception: + dl = None + if dl is None: + return None + try: + return next(iter(dl)) + except Exception: + return None diff --git a/experiments/trainer.py b/experiments/trainer.py index 242b1b8b..d559c046 100644 --- a/experiments/trainer.py +++ b/experiments/trainer.py @@ -8,6 +8,7 @@ import torch from pytorch_lightning import callbacks as pl_callbacks +from experiments.callbacks.flop_counter import FlopCounterCallback from experiments.callbacks.walltime_checkpointer import WalltimeCheckpointer from experiments.callbacks.wandb_cache_cleanup import WandbCacheCleanupCallback from experiments.default_cfg import ExperimentConfig @@ -122,6 +123,8 @@ def construct_trainer( checkpoint_callback, # Model summary callback pl_callbacks.ModelSummary(max_depth=-1), + # One-shot FLOP measurement (fwd + bwd) at fit start; logs flops/{fwd,bwd,step} + FlopCounterCallback(), # Learning rate monitor callback pl_callbacks.LearningRateMonitor(log_weight_decay=True), # Timer callback From a74a9fe57ba4179cf4796d631516cbbc0d365bfb Mon Sep 17 00:00:00 2001 From: OliviaViessmann Date: Wed, 13 May 2026 20:15:35 +0000 Subject: [PATCH 07/13] configs --- benchmarks/well/measure_flops.py | 196 ++++++++++++++++++ .../hyena_gaussian_mask.py | 4 +- .../attention.py | 6 +- slurm/submit_flops.sh | 116 +++++++++++ 4 files changed, 319 insertions(+), 3 deletions(-) create mode 100644 benchmarks/well/measure_flops.py create mode 100755 slurm/submit_flops.sh diff --git a/benchmarks/well/measure_flops.py b/benchmarks/well/measure_flops.py new file mode 100644 index 00000000..a6ed5877 --- /dev/null +++ b/benchmarks/well/measure_flops.py @@ -0,0 +1,196 @@ +"""Standalone FLOP measurement for a WELL v2 experiment. + +Loads a config the same way ``experiments/run.py`` does, instantiates the +datamodule + network + Lightning wrapper, pulls one real training batch, +and runs a single forward + backward inside +``torch.utils.flop_counter.FlopCounterMode``. No wandb, no Lightning +trainer, no torch.compile — just the FLOP numbers, printed to stdout and +written to a JSON file. + +Use this when the SLURM run cannot reach wandb (so the in-training +``FlopCounterCallback`` never gets to log) but you still need the per-step +FLOP counts for a config. + +Usage: + PYTHONPATH=. python benchmarks/well/measure_flops.py \ + --config examples/well/v2/gray_scott_reaction_diffusion/hyena_gaussian_mask.py \ + [--out flops.json] \ + [--device cuda] \ + [overrides ...] + +Notes: + * ``torch.compile`` is intentionally skipped: ``FlopCounterMode`` is a + ``TorchDispatchMode`` and is most reliable in eager mode. The + eager-mode FLOP count is the ground truth — compile changes how + ops execute, not how many FLOPs they perform. + * With ``gradient_checkpointing=True`` the backward re-runs the + checkpointed forward chunks; those re-runs are counted in + ``flops/bwd``, which is the correct accounting for true per-step + training cost. + * The script picks a single batch from ``datamodule.train_dataloader()``. + Per-rank FLOPs scale with the batch size in that dataloader. +""" + +from __future__ import annotations + +import argparse +import json +import sys +from pathlib import Path + +import pytorch_lightning as pl +import torch +from torch.utils.flop_counter import FlopCounterMode + +from experiments.utils.cli import ( + apply_config_overrides, + load_config_from_file, + verify_no_interpolator_overwrites, +) +from nvsubquadratic.lazy_config import instantiate + + +def parse_args() -> argparse.Namespace: + """Parse command-line arguments for the FLOP measurement script.""" + parser = argparse.ArgumentParser(description="Measure per-step FLOPs for a WELL v2 config") + parser.add_argument("--config", type=str, required=True, help="Path to the experiment config .py file") + parser.add_argument("--out", type=str, default=None, help="Optional path to write the JSON result") + parser.add_argument("--device", type=str, default="cuda", help="Device to run measurement on (cuda or cpu)") + parser.add_argument( + "overrides", + nargs="*", + help="Config overrides, e.g. dataset.batch_size=4 (use a small value to fit on one GPU)", + ) + args, unknown = parser.parse_known_args() + for arg in unknown: + args.overrides.append(arg[2:] if arg.startswith("--") else arg) + return args + + +def main() -> None: + """Build the model, fetch a batch, run fwd+bwd under FlopCounterMode, print + write.""" + args = parse_args() + + pl.seed_everything(0) + torch.set_float32_matmul_precision("high") + + config = load_config_from_file(args.config) + verify_no_interpolator_overwrites(config, args.overrides) + config = apply_config_overrides(config, args.overrides) + + device = torch.device(args.device if torch.cuda.is_available() or args.device == "cpu" else "cpu") + print(f"[measure_flops] config={args.config}", flush=True) + print(f"[measure_flops] device={device} precision={config.train.precision}", flush=True) + print(f"[measure_flops] batch_size={config.dataset.batch_size}", flush=True) + + # ------------------------------------------------------------------ # + # Build datamodule + network + wrapper (mirrors experiments/run.py) # + # ------------------------------------------------------------------ # + datamodule = instantiate(config.dataset) + datamodule.prepare_data() + datamodule.setup() + + network = instantiate(config.net) + + # Skip torch.compile on purpose — FlopCounterMode is most reliable in eager mode. + + wrapper_kwargs: dict = {"network": network, "cfg": config} + if hasattr(datamodule, "metadata"): + wrapper_kwargs["metadata"] = datamodule.metadata + if hasattr(datamodule, "normalization"): + wrapper_kwargs["normalization"] = datamodule.normalization + model = instantiate(config.lightning_wrapper_class, **wrapper_kwargs) + + model = model.to(device) + model.train() + + num_params = sum(p.numel() for p in model.parameters()) + num_trainable_params = sum(p.numel() for p in model.parameters() if p.requires_grad) + model_class = type(network).__name__ + try: + patch_size = int(config.net.in_proj_cfg.patch_size) + except (AttributeError, TypeError, ValueError): + patch_size = None + print( + f"[measure_flops] model={model_class} patch_size={patch_size} " + f"params: total={num_params:,} trainable={num_trainable_params:,}", + flush=True, + ) + + # ------------------------------------------------------------------ # + # Pull one batch # + # ------------------------------------------------------------------ # + print("[measure_flops] fetching one training batch …", flush=True) + loader = datamodule.train_dataloader() + batch = next(iter(loader)) + + def _to_device(obj): + if torch.is_tensor(obj): + return obj.to(device, non_blocking=True) + if isinstance(obj, dict): + return {k: _to_device(v) for k, v in obj.items()} + if isinstance(obj, (list, tuple)): + return type(obj)(_to_device(v) for v in obj) + return obj + + batch = _to_device(batch) + + # The wrapper's training_step calls ``self.log(...)``; suppress those + # since we are not inside the actual trainer loop. + model.log = lambda *a, **k: None # type: ignore[assignment] + + # ------------------------------------------------------------------ # + # Measure # + # ------------------------------------------------------------------ # + # Match the precision the actual run would use. + use_bf16 = "bf16" in config.train.precision + autocast_dtype = torch.bfloat16 if use_bf16 else torch.float32 + + print("[measure_flops] running fwd+bwd under FlopCounterMode …", flush=True) + with FlopCounterMode(display=False) as fc: + with torch.amp.autocast(device_type=device.type, dtype=autocast_dtype, enabled=use_bf16): + output = model.training_step(batch, 0) + fwd_flops = fc.get_total_flops() + + loss = output["loss"] if isinstance(output, dict) else output + loss.backward() + total_flops = fc.get_total_flops() + + bwd_flops = total_flops - fwd_flops + + # ------------------------------------------------------------------ # + # Report # + # ------------------------------------------------------------------ # + bs = config.dataset.batch_size + print() + print(f"[measure_flops] per-rank batch_size = {bs}") + print(f"[measure_flops] fwd = {fwd_flops:>20,d} ({fwd_flops / 1e12:.3f} TFLOPs)") + print(f"[measure_flops] bwd = {bwd_flops:>20,d} ({bwd_flops / 1e12:.3f} TFLOPs)") + print(f"[measure_flops] step = {total_flops:>20,d} ({total_flops / 1e12:.3f} TFLOPs)") + + payload = { + "config": args.config, + "overrides": args.overrides, + "model_class": model_class, + "patch_size": patch_size, + "batch_size": bs, + "precision": config.train.precision, + "device": str(device), + "num_params": int(num_params), + "num_trainable_params": int(num_trainable_params), + "fwd_flops": int(fwd_flops), + "bwd_flops": int(bwd_flops), + "step_flops": int(total_flops), + "fwd_tflops": fwd_flops / 1e12, + "bwd_tflops": bwd_flops / 1e12, + "step_tflops": total_flops / 1e12, + } + + out_path = Path(args.out) if args.out else Path(args.config).with_suffix(".flops.json") + out_path.parent.mkdir(parents=True, exist_ok=True) + out_path.write_text(json.dumps(payload, indent=2)) + print(f"\n[measure_flops] wrote {out_path}", flush=True) + + +if __name__ == "__main__": + sys.exit(main()) diff --git a/examples/well/v2/euler_multi_quadrants_periodicBC/hyena_gaussian_mask.py b/examples/well/v2/euler_multi_quadrants_periodicBC/hyena_gaussian_mask.py index 87e7ae43..fea9d8f3 100644 --- a/examples/well/v2/euler_multi_quadrants_periodicBC/hyena_gaussian_mask.py +++ b/examples/well/v2/euler_multi_quadrants_periodicBC/hyena_gaussian_mask.py @@ -1,4 +1,4 @@ -"""Hyena config with Gaussian modulation mask for gray_scott_reaction_diffusion (v2). +"""Hyena config with Gaussian modulation mask for euler_multi_quadrants_periodicBC (v2). Identical to ``cfg_hyena.py`` but replaces the ``nn.Identity`` mask with a ``GaussianModulationND`` mask on the CKConv global convolution kernel. @@ -14,7 +14,7 @@ def get_config() -> ExperimentConfig: - """Build Hyena + Gaussian mask config for active_matter.""" + """Build Hyena + Gaussian mask config for euler_multi_quadrants_periodicBC.""" config = _get_hyena_config() config.net.block_cfg.sequence_mixer_cfg.mixer_cfg.global_conv_cfg.mask_cfg = LazyConfig(GaussianModulationND)( diff --git a/examples/well/v2/gray_scott_reaction_diffusion/attention.py b/examples/well/v2/gray_scott_reaction_diffusion/attention.py index 2d414faf..376c0995 100644 --- a/examples/well/v2/gray_scott_reaction_diffusion/attention.py +++ b/examples/well/v2/gray_scott_reaction_diffusion/attention.py @@ -11,7 +11,7 @@ import torch -from examples.well.v2.active_matter._base import ( +from examples.well.v2.gray_scott_reaction_diffusion._base import ( DATA_DIM, IN_CHANNELS, OUT_CHANNELS, @@ -80,6 +80,10 @@ def get_config() -> ExperimentConfig: is_causal=False, attn_dropout=0.0, rope_base=10000.0, + rope_spatial_dims=( + "${eval:'128 // ${net.in_proj_cfg.patch_size}'}", + "${eval:'128 // ${net.in_proj_cfg.patch_size}'}", + ), ), init_method_in=small_init, init_method_out=partial_wang_init_fn_with_num_layers(num_layers=NUM_BLOCKS), diff --git a/slurm/submit_flops.sh b/slurm/submit_flops.sh new file mode 100755 index 00000000..d982c546 --- /dev/null +++ b/slurm/submit_flops.sh @@ -0,0 +1,116 @@ +#!/bin/bash +#SBATCH --account=healthcareeng_research +#SBATCH --nodes=1 +#SBATCH --partition=batch,backfill +#SBATCH --ntasks-per-node=1 +#SBATCH --time=00:30:00 +#SBATCH --mem=0 +#SBATCH --mail-type=FAIL +#SBATCH --job-name=healthcareeng_research-nvsubq.flops + +# Usage (from repo root): +# sbatch slurm/submit_flops.sh examples/well/v2/active_matter/hyena_gaussian_mask.py +# sbatch slurm/submit_flops.sh examples/well/v2/active_matter/hyena_gaussian_mask.py \ +# net.in_proj_cfg.patch_size=8 dataset.batch_size=4 +# +# Runs benchmarks/well/measure_flops.py inside the same container as training, +# writes flops.json into runs//flops_/. + +set -x + +if [ -z "${1:-}" ]; then + echo "Usage: sbatch slurm/submit_flops.sh [overrides...]" + exit 1 +fi + +CONTAINER_DATA="/workspace/data" +CONTAINER_RESULTS="/workspace/results" + +EXPERIMENT_NAME="$(basename "${1%.py}")" +CONFIG_FILE="$1"; shift + +CONFIG_OVERRIDES=() +PATCH_SIZE="" +for arg in "$@"; do + CONFIG_OVERRIDES+=("${arg}") + if [[ "${arg}" == net.in_proj_cfg.patch_size=* ]]; then + PATCH_SIZE="${arg#net.in_proj_cfg.patch_size=}" + fi +done + +IMAGE_NAME="${SQSH_IMAGE:-/lustre/fsw/healthcareeng_research/oviessmann/nvsubquadratic/enroot/nvsubquadratic-x86_64.sqsh}" + +WORKDIR="${PWD}" +RUNS_DIR="${WORKDIR}/runs" +DATA_DIR="${WELL_HOST_DIR:-/lustre/fsw/healthcareeng_research/oviessmann/nvsubquadratic/data/well_data/datasets}" + +mkdir -p "${RUNS_DIR}" + +EXPERIMENT_DIR="${RUNS_DIR}/${EXPERIMENT_NAME}" +RESULTS_PATH="${EXPERIMENT_DIR}/flops_${SLURM_JOB_ID}" +mkdir -p "${RESULTS_PATH}" + +REPO_DIR="/lustre/fsw/healthcareeng_research/oviessmann/nvsubquadratic" +WORK_DIR="/workspaces/nvSubquadratic-private" +CONFIG_PATH="${WORK_DIR}/${CONFIG_FILE}" +if [ -n "${PATCH_SIZE}" ]; then + OUT_FILENAME="flops_${EXPERIMENT_NAME}_patch${PATCH_SIZE}.json" +else + OUT_FILENAME="flops_${EXPERIMENT_NAME}.json" +fi +OUT_PATH="${CONTAINER_RESULTS}/${OUT_FILENAME}" + +MOUNTS="${DATA_DIR}:${CONTAINER_DATA}" +MOUNTS="${MOUNTS},${RESULTS_PATH}:${CONTAINER_RESULTS}" +MOUNTS="${MOUNTS},${REPO_DIR}:${WORK_DIR}" +MOUNTS="${MOUNTS},$HOME/.cache:/root/.cache" + +echo "================================================" +echo "Experiment: ${EXPERIMENT_NAME}" +echo "Job ID: ${SLURM_JOB_ID}" +echo "Node(s): ${SLURM_NODELIST}" +echo "Config: ${CONFIG_FILE}" +echo "Overrides: ${CONFIG_OVERRIDES[*]}" +echo "Container: ${IMAGE_NAME}" +echo "Results: ${RESULTS_PATH}" +echo "Output JSON: ${RESULTS_PATH}/${OUT_FILENAME}" +echo "Mounts: ${MOUNTS}" +echo "================================================" + +export PYTHONPATH="." +export CUDA_VISIBLE_DEVICES=0 +export DALI_NO_MMAP=1 +export TRITON_CACHE_DIR="/tmp/triton_${SLURM_JOB_ID}" +export TORCHINDUCTOR_FX_GRAPH_CACHE=0 +export TORCH_NCCL_AVOID_RECORD_STREAMS=1 +export PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True +export WELL_DATA_PATH="${CONTAINER_DATA}" + +OVERRIDES_STR="${CONFIG_OVERRIDES[*]}" + +read -r -d '' PYTHON_CMD < Date: Fri, 22 May 2026 18:45:49 +0000 Subject: [PATCH 08/13] well: euler_multi_quadrants_periodicBC --- .../v2/euler_multi_quadrants_periodicBC/attention.py | 4 ++++ .../well/v2/euler_multi_quadrants_periodicBC/hyena.py | 4 ++-- .../hyena_gaussian_mask.py | 6 +++--- experiments/run.py | 11 ++++++++--- 4 files changed, 17 insertions(+), 8 deletions(-) diff --git a/examples/well/v2/euler_multi_quadrants_periodicBC/attention.py b/examples/well/v2/euler_multi_quadrants_periodicBC/attention.py index d11d683e..952dec70 100644 --- a/examples/well/v2/euler_multi_quadrants_periodicBC/attention.py +++ b/examples/well/v2/euler_multi_quadrants_periodicBC/attention.py @@ -82,6 +82,10 @@ def get_config() -> ExperimentConfig: is_causal=False, attn_dropout=0.0, rope_base=10000.0, + rope_spatial_dims=( + "${eval:'512 // ${net.in_proj_cfg.patch_size}'}", + "${eval:'512 // ${net.in_proj_cfg.patch_size}'}", + ), ), init_method_in=small_init, init_method_out=partial_wang_init_fn_with_num_layers(num_layers=NUM_BLOCKS), diff --git a/examples/well/v2/euler_multi_quadrants_periodicBC/hyena.py b/examples/well/v2/euler_multi_quadrants_periodicBC/hyena.py index 15207d4f..677e0b44 100644 --- a/examples/well/v2/euler_multi_quadrants_periodicBC/hyena.py +++ b/examples/well/v2/euler_multi_quadrants_periodicBC/hyena.py @@ -3,7 +3,7 @@ Uses a ResidualNetwork with Hyena (QKV + CKConv global conv) as the sequence mixer. Circular FFT padding matches the dataset's periodic boundary conditions. With patch_size=16 the effective sequence -resolution is 16×16. +resolution is 32×32. Patch-size CLI override ----------------------- @@ -94,7 +94,7 @@ def get_config() -> ExperimentConfig: num_layers=3, embedding_dim=64, omega_0=OMEGA_0, - L_cache="${eval:'256 // ${net.in_proj_cfg.patch_size}'}", + L_cache="${eval:'512 // ${net.in_proj_cfg.patch_size}'}", use_bias=True, hidden_omega_0=1.0, ), diff --git a/examples/well/v2/euler_multi_quadrants_periodicBC/hyena_gaussian_mask.py b/examples/well/v2/euler_multi_quadrants_periodicBC/hyena_gaussian_mask.py index fea9d8f3..6c4c8d81 100644 --- a/examples/well/v2/euler_multi_quadrants_periodicBC/hyena_gaussian_mask.py +++ b/examples/well/v2/euler_multi_quadrants_periodicBC/hyena_gaussian_mask.py @@ -4,9 +4,9 @@ ``GaussianModulationND`` mask on the CKConv global convolution kernel. """ -from examples.well.v2.active_matter._base import DATA_DIM -from examples.well.v2.active_matter.hyena import NUM_HIDDEN_CHANNELS -from examples.well.v2.active_matter.hyena import get_config as _get_hyena_config +from examples.well.v2.euler_multi_quadrants_periodicBC._base import DATA_DIM +from examples.well.v2.euler_multi_quadrants_periodicBC.hyena import NUM_HIDDEN_CHANNELS +from examples.well.v2.euler_multi_quadrants_periodicBC.hyena import get_config as _get_hyena_config from experiments.callbacks.mask_monitor import MaskMonitorCallback from experiments.default_cfg import ExperimentConfig from nvsubquadratic.lazy_config import LazyConfig diff --git a/experiments/run.py b/experiments/run.py index d26701f7..7e33458e 100644 --- a/experiments/run.py +++ b/experiments/run.py @@ -202,9 +202,13 @@ def main() -> None: # Determine if we should attach to an existing run by name (autoresume) autoresume_ckpt_path = None attach_run_id = None + run_id_file = experiment_dir / "run.id" - # If autoresume is enabled, search W&B for existing run and download checkpoint - if config.autoresume.enabled and not offline: + # If autoresume is enabled, search W&B for existing run and download checkpoint. + # Skip the API search when run.id already exists locally — the local file is + # authoritative for chained jobs, and hitting the W&B GraphQL endpoint on every + # job start can 429 and kill the job before training begins. + if config.autoresume.enabled and not offline and not run_id_file.exists(): api = wandb.Api() runs = api.runs( path=f"{config.wandb.entity}/{config.wandb.project}", @@ -229,6 +233,8 @@ def main() -> None: print("[autoresume] Will check for local checkpoint instead.") else: print(f"[autoresume] No existing run found with name '{run_name}', starting fresh.") + elif config.autoresume.enabled and run_id_file.exists(): + print(f"[autoresume] Found local run.id at {run_id_file}, skipping W&B API search.") # If autoresume enabled but no W&B checkpoint found, check for local checkpoint if config.autoresume.enabled and autoresume_ckpt_path is None: @@ -242,7 +248,6 @@ def main() -> None: print(f"[autoresume] No last.ckpt found in {ckpt_dir}, starting from scratch.") # Generate or reuse run ID - run_id_file = experiment_dir / "run.id" if config.wandb.run_id is not None: # Explicit run_id provided via config override attach_run_id = config.wandb.run_id From ea6b22fc0e6ab0677b6bdc40492b6106d9d50f75 Mon Sep 17 00:00:00 2001 From: OliviaViessmann Date: Fri, 5 Jun 2026 17:42:04 +0000 Subject: [PATCH 09/13] euler_multiquadrants config and tracker + commening out MonitorCallback that calls wandb.API() which isn't needed --- .../well/v2/MHD_64/hyena_gaussian_mask.py | 5 +- .../hyena_gaussian_mask.py | 5 +- .../v2/active_matter/hyena_gaussian_mask.py | 5 +- .../TRACKER.md | 32 ++-- .../euler_multi_quadrants_periodicBC/hyena.py | 1 - .../hyena_gaussian_mask.py | 5 +- .../hyena_gaussian_mask.py | 5 +- .../well/v2/helmholtz_staircase/attention.py | 114 ++++++++++++++ examples/well/v2/helmholtz_staircase/hyena.py | 148 ++++++++++++++++++ .../hyena_gaussian_mask.py | 32 ++++ .../hyena_gaussian_mask.py | 5 +- experiments/default_cfg.py | 1 + experiments/run.py | 2 +- slurm/submit_well.sh | 14 +- 14 files changed, 339 insertions(+), 35 deletions(-) create mode 100644 examples/well/v2/helmholtz_staircase/attention.py create mode 100644 examples/well/v2/helmholtz_staircase/hyena.py create mode 100644 examples/well/v2/helmholtz_staircase/hyena_gaussian_mask.py diff --git a/examples/well/v2/MHD_64/hyena_gaussian_mask.py b/examples/well/v2/MHD_64/hyena_gaussian_mask.py index ae6690c8..85fa08e9 100644 --- a/examples/well/v2/MHD_64/hyena_gaussian_mask.py +++ b/examples/well/v2/MHD_64/hyena_gaussian_mask.py @@ -7,7 +7,8 @@ from examples.well.v2.MHD_64._base import DATA_DIM from examples.well.v2.MHD_64.hyena import NUM_HIDDEN_CHANNELS from examples.well.v2.MHD_64.hyena import get_config as _get_hyena_config -from experiments.callbacks.mask_monitor import MaskMonitorCallback + +# from experiments.callbacks.mask_monitor import MaskMonitorCallback # disabled: wandb artifact API retry loop blocks training from experiments.default_cfg import ExperimentConfig from nvsubquadratic.lazy_config import LazyConfig from nvsubquadratic.modules.masks_nd import GaussianModulationND @@ -26,6 +27,6 @@ def get_config() -> ExperimentConfig: parametrization="direct", ) - config.callbacks.append(LazyConfig(MaskMonitorCallback)(log_every_n_steps=50)) + # config.callbacks.append(LazyConfig(MaskMonitorCallback)(log_every_n_steps=50)) # disabled: wandb artifact API retry loop blocks training return config diff --git a/examples/well/v2/acoustic_scattering_maze/hyena_gaussian_mask.py b/examples/well/v2/acoustic_scattering_maze/hyena_gaussian_mask.py index f95c48eb..f65efd22 100644 --- a/examples/well/v2/acoustic_scattering_maze/hyena_gaussian_mask.py +++ b/examples/well/v2/acoustic_scattering_maze/hyena_gaussian_mask.py @@ -7,7 +7,8 @@ from examples.well.v2.acoustic_scattering_maze._base import DATA_DIM from examples.well.v2.acoustic_scattering_maze.hyena import NUM_HIDDEN_CHANNELS from examples.well.v2.acoustic_scattering_maze.hyena import get_config as _get_hyena_config -from experiments.callbacks.mask_monitor import MaskMonitorCallback + +# from experiments.callbacks.mask_monitor import MaskMonitorCallback # disabled: wandb artifact API retry loop blocks training from experiments.default_cfg import ExperimentConfig from nvsubquadratic.lazy_config import LazyConfig from nvsubquadratic.modules.masks_nd import GaussianModulationND @@ -26,6 +27,6 @@ def get_config() -> ExperimentConfig: parametrization="direct", ) - config.callbacks.append(LazyConfig(MaskMonitorCallback)(log_every_n_steps=50)) + # config.callbacks.append(LazyConfig(MaskMonitorCallback)(log_every_n_steps=50)) # disabled: wandb artifact API retry loop blocks training return config diff --git a/examples/well/v2/active_matter/hyena_gaussian_mask.py b/examples/well/v2/active_matter/hyena_gaussian_mask.py index e88a7e1a..7d2f3cc0 100644 --- a/examples/well/v2/active_matter/hyena_gaussian_mask.py +++ b/examples/well/v2/active_matter/hyena_gaussian_mask.py @@ -7,7 +7,8 @@ from examples.well.v2.active_matter._base import DATA_DIM from examples.well.v2.active_matter.hyena import NUM_HIDDEN_CHANNELS from examples.well.v2.active_matter.hyena import get_config as _get_hyena_config -from experiments.callbacks.mask_monitor import MaskMonitorCallback + +# from experiments.callbacks.mask_monitor import MaskMonitorCallback # disabled: wandb artifact API retry loop blocks training from experiments.default_cfg import ExperimentConfig from nvsubquadratic.lazy_config import LazyConfig from nvsubquadratic.modules.masks_nd import GaussianModulationND @@ -26,6 +27,6 @@ def get_config() -> ExperimentConfig: parametrization="direct", ) - config.callbacks.append(LazyConfig(MaskMonitorCallback)(log_every_n_steps=50)) + # config.callbacks.append(LazyConfig(MaskMonitorCallback)(log_every_n_steps=50)) # disabled: wandb artifact API retry loop blocks training return config diff --git a/examples/well/v2/euler_multi_quadrants_periodicBC/TRACKER.md b/examples/well/v2/euler_multi_quadrants_periodicBC/TRACKER.md index 9174727c..3011faf5 100644 --- a/examples/well/v2/euler_multi_quadrants_periodicBC/TRACKER.md +++ b/examples/well/v2/euler_multi_quadrants_periodicBC/TRACKER.md @@ -3,38 +3,38 @@ ## Goal Compare **CNextU-net** (baseline), **Attention**, and **Hyena + Gaussian mask** -on `euer_multi_quadrants` (512x512). +on `euler_multi_quadrants_periodicBC` (512x512). All runs share the same training recipe (24 hours, AdamW, cosine schedule with 5% warmup, bf16-mixed, grad_clip=1.0). ## Run order 1. CNextU-net (baseline, no patch size) -1. Attention p2 -1. Hyena+G p2 +1. Attention p8 +1. Hyena+G p8, p4, p2 ## Results — CNextU-net (baseline) -| Batch/GPU | GPUs | WandB Run | val/VRMSE | it/s | total steps | -| --------- | ---- | ---------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------- | ----------------- | ---- | ----------- | -| 256 | 1 | OV\_\_workspaces_nvSubquadratic-private_examples_well_v2_gray_scott_reaction_diffusion_cfg_unet_convnext_enabled_True_experiment_dir\_/workspace/results_lr_0.0001_num_nodes_1_run_start_time_1777331397_r | 0.231915682554245 | 6.00 | 110,000 | +| Batch/GPU | GPUs | WandB Run | val/VRMSE | it/s | total steps | +| --------- | ---- | ----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------- | --------- | ---- | ----------- | +| 42 | 1 | OV\_\_workspaces_nvSubquadratic-private_examples_well_v2_euler_multi_quadrants_periodicBC_cfg_unet_convnext_bs_42_enabled_True_experiment_dir\_/workspace/results_lr_0.003_num_nodes_1_run_start_time_17806 | 0.0332 | 1.52 | 110,000 | ## Results — Attention Total params: 13,747,979 -| Patch | Tokens/img | Batch/GPU | GPUs | WandB Run | val/VRMSE | it/s | total steps | -| ----- | ---------- | --------- | ---- | ----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------- | -------------------- | ---- | ----------- | -| 8 | 2,048 | 256 | 1 | OV\_\_workspaces_nvSubquadratic-private_examples_well_v2_gray_scott_reaction_diffusion_attention_enabled_True_experiment_dir\_/workspace/results_lr_0.0005_num_nodes_1_patch_size_8_run_start_time_17773947 | 0.052010852843523026 | 1.4 | 110,000 | -| 4 | 4,096 | 256 | 1 | | 0.007981576956808567 | 1.36 | 110,000 | -| 2 | 8,192 | 16 | 1 | | 0.00703924847766757 | 2.63 | 110,000 | +| Patch | Tokens/img | Batch/GPU | GPUs | WandB Run | val/VRMSE | it/s | total steps | +| ----- | ---------- | --------- | ---- | ------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------ | --------- | ---- | ----------- | +| 8 | 2,048 | 42 | 1 | OV\_\_workspaces_nvSubquadratic-private_examples_well_v2_euler_multi_quadrants_periodicBC_attention_bs_42_enabled_True_experiment_dir\_/workspace/results_lr_0.003_num_nodes_1_patch_size_8_run_start_time\_ | 0.1293 | 1.38 | 110,000 | +| 4 | 4,096 | 42 | 1 | OOM at bs=42 | — | — | — | +| 2 | 8,192 | 12 | 1 | OOM at bs=12 | — | — | — | ## Results — Hyena + Gaussian Mask Total params: 14,297,099 -| Patch | Tokens/img | Batch/GPU | GPUs | WandB Run | val/VRMSE | it/s | total steps | -| ----- | ---------- | --------- | ---- | ----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------- | -------------------- | ---- | ----------- | -| 8 | 2,048 | 256 | 1 | OV\_\_workspaces_nvSubquadratic-private_examples_well_v2_gray_scott_reaction_diffusion_hyena_gaussian_mask_enabled_True_experiment_dir\_/workspace/results_lr_0.0005_num_nodes_1_patch_size_8_run_start_tim | 0.009229527786374092 | 1.36 | 88200 | -| 4 | 4,096 | 256 | 1 | | 0.007981576956808567 | 1.36 | 110,000 | -| 2 | 8,192 | 16 | 1 | | 0.00703924847766757 | 2.63 | 110,000 | +| Patch | Tokens/img | Batch/GPU | GPUs | WandB Run | val/VRMSE | it/s | total steps | +| ----- | ---------- | --------- | ---- | ----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------- | --------- | ------ | ----------- | +| 8 | 2,048 | 42 | 1 | OV\_\_workspaces_nvSubquadratic-private_examples_well_v2_euler_multi_quadrants_periodicBC_hyena_gaussian_mask_bs_42_experiment_dir\_/workspace/results_lr_0.003_num_2026-06-03-06-49-37 | 0.0311 | 1.33 | 110,000 | +| 4 | 4,096 | 42 | 1 | OV\_\_workspaces_nvSubquadratic-private_examples_well_v2_euler_multi_quadrants_periodicBC_hyena_gaussian_mask_bs_42_enabled_True_experiment_dir\_/workspace/results_lr_0.003_num_nodes_1_patch_size_4_run_s | 0.0378 | 0.5915 | 44,787 | +| 2 | 8,192 | 12 | 1 | OV\_\_workspaces_nvSubquadratic-private_examples_well_v2_euler_multi_quadrants_periodicBC_hyena_gaussian_mask_bs_12_enabled_True_experiment_dir\_/workspace/results_lr_0.003_num_nodes_1_patch_size_2_run_s | 0.1088 | 0.4749 | 40,949 | diff --git a/examples/well/v2/euler_multi_quadrants_periodicBC/hyena.py b/examples/well/v2/euler_multi_quadrants_periodicBC/hyena.py index 677e0b44..c3ea4966 100644 --- a/examples/well/v2/euler_multi_quadrants_periodicBC/hyena.py +++ b/examples/well/v2/euler_multi_quadrants_periodicBC/hyena.py @@ -118,7 +118,6 @@ def get_config() -> ExperimentConfig: dim=NUM_HIDDEN_CHANNELS, ), qk_norm_cfg=LazyConfig(L2Norm)(), - use_rope=False, ), init_method_in=small_init, init_method_out=partial_wang_init_fn_with_num_layers(num_layers=NUM_BLOCKS), diff --git a/examples/well/v2/euler_multi_quadrants_periodicBC/hyena_gaussian_mask.py b/examples/well/v2/euler_multi_quadrants_periodicBC/hyena_gaussian_mask.py index 6c4c8d81..0d40a091 100644 --- a/examples/well/v2/euler_multi_quadrants_periodicBC/hyena_gaussian_mask.py +++ b/examples/well/v2/euler_multi_quadrants_periodicBC/hyena_gaussian_mask.py @@ -7,7 +7,8 @@ from examples.well.v2.euler_multi_quadrants_periodicBC._base import DATA_DIM from examples.well.v2.euler_multi_quadrants_periodicBC.hyena import NUM_HIDDEN_CHANNELS from examples.well.v2.euler_multi_quadrants_periodicBC.hyena import get_config as _get_hyena_config -from experiments.callbacks.mask_monitor import MaskMonitorCallback + +# from experiments.callbacks.mask_monitor import MaskMonitorCallback # disabled: wandb artifact API retry loop blocks training from experiments.default_cfg import ExperimentConfig from nvsubquadratic.lazy_config import LazyConfig from nvsubquadratic.modules.masks_nd import GaussianModulationND @@ -26,6 +27,6 @@ def get_config() -> ExperimentConfig: parametrization="direct", ) - config.callbacks.append(LazyConfig(MaskMonitorCallback)(log_every_n_steps=50)) + # config.callbacks.append(LazyConfig(MaskMonitorCallback)(log_every_n_steps=50)) # disabled: wandb artifact API retry loop blocks training return config diff --git a/examples/well/v2/gray_scott_reaction_diffusion/hyena_gaussian_mask.py b/examples/well/v2/gray_scott_reaction_diffusion/hyena_gaussian_mask.py index 87e7ae43..64165833 100644 --- a/examples/well/v2/gray_scott_reaction_diffusion/hyena_gaussian_mask.py +++ b/examples/well/v2/gray_scott_reaction_diffusion/hyena_gaussian_mask.py @@ -7,7 +7,8 @@ from examples.well.v2.active_matter._base import DATA_DIM from examples.well.v2.active_matter.hyena import NUM_HIDDEN_CHANNELS from examples.well.v2.active_matter.hyena import get_config as _get_hyena_config -from experiments.callbacks.mask_monitor import MaskMonitorCallback + +# from experiments.callbacks.mask_monitor import MaskMonitorCallback # disabled: wandb artifact API retry loop blocks training from experiments.default_cfg import ExperimentConfig from nvsubquadratic.lazy_config import LazyConfig from nvsubquadratic.modules.masks_nd import GaussianModulationND @@ -26,6 +27,6 @@ def get_config() -> ExperimentConfig: parametrization="direct", ) - config.callbacks.append(LazyConfig(MaskMonitorCallback)(log_every_n_steps=50)) + # config.callbacks.append(LazyConfig(MaskMonitorCallback)(log_every_n_steps=50)) # disabled: wandb artifact API retry loop blocks training return config diff --git a/examples/well/v2/helmholtz_staircase/attention.py b/examples/well/v2/helmholtz_staircase/attention.py new file mode 100644 index 00000000..f05e7333 --- /dev/null +++ b/examples/well/v2/helmholtz_staircase/attention.py @@ -0,0 +1,114 @@ +"""Attention config for helmholtz_staircase (v2). + +Uses a ResidualNetwork with multi-head self-attention (QKV + RoPE) as the +sequence mixer. With patch_size=8 the effective sequence resolution is 128×32. + +Patch-size CLI override +----------------------- +Only ``net.in_proj_cfg.patch_size=P`` is needed; stride and out_proj patch_size +are derived via OmegaConf interpolators. + +Note: rope_spatial_dims is per-axis (1024//P, 256//P) because the spatial +resolution is non-square (1024×256). +""" + +import torch + +from examples.well.v2.helmholtz_staircase._base import ( + DATA_DIM, + IN_CHANNELS, + OUT_CHANNELS, + get_base_config, +) +from experiments.default_cfg import ExperimentConfig +from nvsubquadratic.lazy_config import LazyConfig +from nvsubquadratic.modules.attention import Attention +from nvsubquadratic.modules.mlp import MLP +from nvsubquadratic.modules.patchify import Patchify, Unpatchify +from nvsubquadratic.modules.residual_block import ResidualBlock +from nvsubquadratic.modules.rms_norm import RMSNorm +from nvsubquadratic.modules.sequence_mixer import QKVSequenceMixer +from nvsubquadratic.networks.general_purpose_resnet import ResidualNetwork +from nvsubquadratic.utils.init import partial_wang_init_fn_with_num_layers, small_init + + +# ─── Model hyperparameters ──────────────────────────────────────────────────── +NUM_HIDDEN_CHANNELS = 384 +NUM_BLOCKS = 12 +NUM_HEADS = 8 # head_dim=48; must be divisible by 4 for 2D RoPE (2 dims × 2 axes) +PATCH_SIZE = 8 + +DROPOUT_IN_RATE = 0.0 +DROPOUT_RATE = 0.0 + +GRADIENT_CHECKPOINTING = False + + +def get_config() -> ExperimentConfig: + """Build Attention experiment config for helmholtz_staircase.""" + config = get_base_config(learning_rate=1e-3, weight_decay=1e-5) + + config.compile = True + config.compile_mode = "max-autotune-no-cudagraphs" + + norm_cfg = LazyConfig(RMSNorm)(dim=NUM_HIDDEN_CHANNELS) + + config.net = LazyConfig(ResidualNetwork)( + in_channels=IN_CHANNELS, + out_channels=OUT_CHANNELS, + num_blocks=NUM_BLOCKS, + hidden_dim=NUM_HIDDEN_CHANNELS, + data_dim=DATA_DIM, + in_proj_cfg=LazyConfig(Patchify)( + in_features=IN_CHANNELS, + out_features=NUM_HIDDEN_CHANNELS, + data_dim=DATA_DIM, + patch_size=PATCH_SIZE, + stride="${net.in_proj_cfg.patch_size}", + ), + out_proj_cfg=LazyConfig(Unpatchify)( + in_features=NUM_HIDDEN_CHANNELS, + out_features=OUT_CHANNELS, + data_dim=DATA_DIM, + patch_size="${net.in_proj_cfg.patch_size}", + stride="${net.in_proj_cfg.patch_size}", + ), + norm_cfg=norm_cfg, + block_cfg=LazyConfig(ResidualBlock)( + sequence_mixer_cfg=LazyConfig(QKVSequenceMixer)( + hidden_dim=NUM_HIDDEN_CHANNELS, + mixer_cfg=LazyConfig(Attention)( + hidden_dim=NUM_HIDDEN_CHANNELS, + num_heads=NUM_HEADS, + apply_qk_norm=True, + use_rope=True, + is_causal=False, + attn_dropout=0.0, + rope_base=10000.0, + rope_spatial_dims=( + "${eval:'1024 // ${net.in_proj_cfg.patch_size}'}", + "${eval:'256 // ${net.in_proj_cfg.patch_size}'}", + ), + ), + init_method_in=small_init, + init_method_out=partial_wang_init_fn_with_num_layers(num_layers=NUM_BLOCKS), + ), + sequence_mixer_norm_cfg=norm_cfg, + condition_mixer_cfg=LazyConfig(torch.nn.Identity)(), + condition_mixer_norm_cfg=LazyConfig(torch.nn.Identity)(), + mlp_cfg=LazyConfig(MLP)( + dim=NUM_HIDDEN_CHANNELS, + activation="glu", + expansion_factor=1.0, + dropout_cfg=LazyConfig(torch.nn.Dropout)(p=DROPOUT_RATE), + init_method_in=small_init, + init_method_out=partial_wang_init_fn_with_num_layers(num_layers=NUM_BLOCKS), + ), + mlp_norm_cfg=norm_cfg, + dropout_cfg=LazyConfig(torch.nn.Dropout)(p=DROPOUT_RATE), + ), + dropout_in_cfg=LazyConfig(torch.nn.Dropout)(p=DROPOUT_IN_RATE), + gradient_checkpointing=GRADIENT_CHECKPOINTING, + ) + + return config diff --git a/examples/well/v2/helmholtz_staircase/hyena.py b/examples/well/v2/helmholtz_staircase/hyena.py new file mode 100644 index 00000000..7fdfd187 --- /dev/null +++ b/examples/well/v2/helmholtz_staircase/hyena.py @@ -0,0 +1,148 @@ +"""Hyena config for helmholtz_staircase (v2). + +Uses a ResidualNetwork with Hyena (QKV + CKConv global conv) as the +sequence mixer. Zero-pad FFT matches the dataset's open boundary conditions. +With patch_size=8 the effective sequence resolution is 128×32. + +Patch-size CLI override +----------------------- +Only ``net.in_proj_cfg.patch_size=P`` is needed; stride, out_proj patch_size, +and kernel L_cache are derived via OmegaConf interpolators. + +Note: L_cache is a per-axis list [1024//P, 256//P] because the spatial +resolution is non-square (1024×256). +""" + +import torch + +from examples.well.v2.helmholtz_staircase._base import ( + DATA_DIM, + IN_CHANNELS, + OUT_CHANNELS, + get_base_config, +) +from experiments.default_cfg import ExperimentConfig +from nvsubquadratic.lazy_config import LazyConfig +from nvsubquadratic.modules.ckconv_nd import CKConvND +from nvsubquadratic.modules.hyena_nd import Hyena +from nvsubquadratic.modules.kernels_nd import SIRENKernelND +from nvsubquadratic.modules.mlp import MLP +from nvsubquadratic.modules.patchify import Patchify, Unpatchify +from nvsubquadratic.modules.residual_block import ResidualBlock +from nvsubquadratic.modules.rms_norm import RMSNorm +from nvsubquadratic.modules.sequence_mixer import QKVSequenceMixer +from nvsubquadratic.networks.general_purpose_resnet import ResidualNetwork +from nvsubquadratic.utils.init import partial_wang_init_fn_with_num_layers, small_init +from nvsubquadratic.utils.qk_norm import L2Norm + + +# ─── Model hyperparameters ──────────────────────────────────────────────────── +NUM_HIDDEN_CHANNELS = 384 +NUM_BLOCKS = 12 +PATCH_SIZE = 8 + +DROPOUT_IN_RATE = 0.0 +DROPOUT_RATE = 0.0 +GRID_TYPE = "single" +FFT_PADDING = "zero" # open boundary conditions +OMEGA_0 = 30.0 + +GRADIENT_CHECKPOINTING = True + + +def get_config() -> ExperimentConfig: + """Build Hyena experiment config for helmholtz_staircase.""" + config = get_base_config(learning_rate=1e-3, weight_decay=1e-5) + + config.compile = True + config.compile_mode = "max-autotune-no-cudagraphs" + + norm_cfg = LazyConfig(RMSNorm)(dim=NUM_HIDDEN_CHANNELS) + + config.net = LazyConfig(ResidualNetwork)( + in_channels=IN_CHANNELS, + out_channels=OUT_CHANNELS, + num_blocks=NUM_BLOCKS, + hidden_dim=NUM_HIDDEN_CHANNELS, + data_dim=DATA_DIM, + in_proj_cfg=LazyConfig(Patchify)( + in_features=IN_CHANNELS, + out_features=NUM_HIDDEN_CHANNELS, + data_dim=DATA_DIM, + patch_size=PATCH_SIZE, + stride="${net.in_proj_cfg.patch_size}", + ), + out_proj_cfg=LazyConfig(Unpatchify)( + in_features=NUM_HIDDEN_CHANNELS, + out_features=OUT_CHANNELS, + data_dim=DATA_DIM, + patch_size="${net.in_proj_cfg.patch_size}", + stride="${net.in_proj_cfg.patch_size}", + ), + norm_cfg=norm_cfg, + block_cfg=LazyConfig(ResidualBlock)( + sequence_mixer_cfg=LazyConfig(QKVSequenceMixer)( + hidden_dim=NUM_HIDDEN_CHANNELS, + mixer_cfg=LazyConfig(Hyena)( + global_conv_cfg=LazyConfig(CKConvND)( + data_dim=DATA_DIM, + hidden_dim=NUM_HIDDEN_CHANNELS, + fft_padding=FFT_PADDING, + use_fp16_fft=False, + kernel_cfg=LazyConfig(SIRENKernelND)( + data_dim=DATA_DIM, + out_dim=NUM_HIDDEN_CHANNELS, + mlp_hidden_dim=64, + num_layers=3, + embedding_dim=64, + omega_0=OMEGA_0, + L_cache=[ + "${eval:'1024 // ${net.in_proj_cfg.patch_size}'}", + "${eval:'256 // ${net.in_proj_cfg.patch_size}'}", + ], + use_bias=True, + hidden_omega_0=1.0, + ), + mask_cfg=LazyConfig(torch.nn.Identity)(), + grid_type=GRID_TYPE, + ), + short_conv_cfg=LazyConfig(torch.nn.Conv2d)( + in_channels=3 * NUM_HIDDEN_CHANNELS, + out_channels=3 * NUM_HIDDEN_CHANNELS, + kernel_size=3, + groups=3 * NUM_HIDDEN_CHANNELS, + padding=1, + bias=False, + ), + gate_nonlinear_cfg=LazyConfig(torch.nn.SiLU)(), + gate_nonlinear_2_cfg=LazyConfig(torch.nn.Sigmoid)(), + pixelhyena_norm_cfg=LazyConfig(RMSNorm)( + dim=NUM_HIDDEN_CHANNELS, + ), + output_norm_cfg=LazyConfig(RMSNorm)( + dim=NUM_HIDDEN_CHANNELS, + ), + qk_norm_cfg=LazyConfig(L2Norm)(), + ), + init_method_in=small_init, + init_method_out=partial_wang_init_fn_with_num_layers(num_layers=NUM_BLOCKS), + ), + sequence_mixer_norm_cfg=norm_cfg, + condition_mixer_cfg=LazyConfig(torch.nn.Identity)(), + condition_mixer_norm_cfg=LazyConfig(torch.nn.Identity)(), + mlp_cfg=LazyConfig(MLP)( + dim=NUM_HIDDEN_CHANNELS, + activation="glu", + expansion_factor=1.0, + dropout_cfg=LazyConfig(torch.nn.Dropout)(p=DROPOUT_RATE), + init_method_in=small_init, + init_method_out=partial_wang_init_fn_with_num_layers(num_layers=NUM_BLOCKS), + ), + mlp_norm_cfg=norm_cfg, + dropout_cfg=LazyConfig(torch.nn.Dropout)(p=DROPOUT_RATE), + ), + dropout_in_cfg=LazyConfig(torch.nn.Dropout)(p=DROPOUT_IN_RATE), + gradient_checkpointing=GRADIENT_CHECKPOINTING, + ) + + return config diff --git a/examples/well/v2/helmholtz_staircase/hyena_gaussian_mask.py b/examples/well/v2/helmholtz_staircase/hyena_gaussian_mask.py new file mode 100644 index 00000000..e62a4531 --- /dev/null +++ b/examples/well/v2/helmholtz_staircase/hyena_gaussian_mask.py @@ -0,0 +1,32 @@ +"""Hyena config with Gaussian modulation mask for helmholtz_staircase (v2). + +Identical to ``hyena.py`` but replaces the ``nn.Identity`` mask with a +``GaussianModulationND`` mask on the CKConv global convolution kernel. +""" + +from examples.well.v2.helmholtz_staircase._base import DATA_DIM +from examples.well.v2.helmholtz_staircase.hyena import NUM_HIDDEN_CHANNELS +from examples.well.v2.helmholtz_staircase.hyena import get_config as _get_hyena_config + +# from experiments.callbacks.mask_monitor import MaskMonitorCallback # disabled: wandb artifact API retry loop blocks training +from experiments.default_cfg import ExperimentConfig +from nvsubquadratic.lazy_config import LazyConfig +from nvsubquadratic.modules.masks_nd import GaussianModulationND + + +def get_config() -> ExperimentConfig: + """Build Hyena + Gaussian mask config for helmholtz_staircase.""" + config = _get_hyena_config() + + config.net.block_cfg.sequence_mixer_cfg.mixer_cfg.global_conv_cfg.mask_cfg = LazyConfig(GaussianModulationND)( + data_dim=DATA_DIM, + num_channels=NUM_HIDDEN_CHANNELS, + min_attenuation_at_step=0.1, + max_attenuation_at_limit=0.95, + init_extent=1.0, + parametrization="direct", + ) + + # config.callbacks.append(LazyConfig(MaskMonitorCallback)(log_every_n_steps=50)) # disabled: wandb artifact API retry loop blocks training + + return config diff --git a/examples/well/v2/supernova_explosion_64/hyena_gaussian_mask.py b/examples/well/v2/supernova_explosion_64/hyena_gaussian_mask.py index 65703a96..f87fedf1 100644 --- a/examples/well/v2/supernova_explosion_64/hyena_gaussian_mask.py +++ b/examples/well/v2/supernova_explosion_64/hyena_gaussian_mask.py @@ -7,7 +7,8 @@ from examples.well.v2.supernova_explosion_64._base import DATA_DIM from examples.well.v2.supernova_explosion_64.hyena import NUM_HIDDEN_CHANNELS from examples.well.v2.supernova_explosion_64.hyena import get_config as _get_hyena_config -from experiments.callbacks.mask_monitor import MaskMonitorCallback + +# from experiments.callbacks.mask_monitor import MaskMonitorCallback # disabled: wandb artifact API retry loop blocks training from experiments.default_cfg import ExperimentConfig from nvsubquadratic.lazy_config import LazyConfig from nvsubquadratic.modules.masks_nd import GaussianModulationND @@ -26,6 +27,6 @@ def get_config() -> ExperimentConfig: parametrization="direct", ) - config.callbacks.append(LazyConfig(MaskMonitorCallback)(log_every_n_steps=50)) + # config.callbacks.append(LazyConfig(MaskMonitorCallback)(log_every_n_steps=50)) # disabled: wandb artifact API retry loop blocks training return config diff --git a/experiments/default_cfg.py b/experiments/default_cfg.py index 4e8b1101..da47e14f 100644 --- a/experiments/default_cfg.py +++ b/experiments/default_cfg.py @@ -94,6 +94,7 @@ class WandbConfig: job_group: str = "" tags: list = field(default_factory=list) run_id: Optional[str] = None # Explicit W&B run ID for resuming or linking runs + offline: bool = False # If True, log to disk only (no remote sync). Implied when debug=True. @dataclass diff --git a/experiments/run.py b/experiments/run.py index 7e33458e..36c574c1 100644 --- a/experiments/run.py +++ b/experiments/run.py @@ -183,7 +183,7 @@ def main() -> None: else: # Avoid auto logging all checkpoints; selective uploader handles best/last log_model = False - offline = False + offline = config.wandb.offline if config.autoresume.enabled: # If run name is not provided, use the deterministic run name without timestamp diff --git a/slurm/submit_well.sh b/slurm/submit_well.sh index 664ef9f3..f1e90840 100644 --- a/slurm/submit_well.sh +++ b/slurm/submit_well.sh @@ -7,6 +7,7 @@ #SBATCH --mem=0 #SBATCH --mail-type=FAIL #SBATCH --job-name=healthcareeng_research-nvsubq.well +#SBATCH --gpus-per-node=1 # Usage (from repo root): # sbatch slurm/submit_well.sh examples/well/v2/active_matter/hyena_gaussian_mask.py @@ -47,12 +48,15 @@ for arg in "$@"; do done # Container image -IMAGE_NAME="${SQSH_IMAGE:-/lustre/fsw/healthcareeng_research/oviessmann/nvsubquadratic/enroot/nvsubquadratic-x86_64.sqsh}" +#eos lustre path: +#IMAGE_NAME="${SQSH_IMAGE:-/lustre/fsw/portfolios/healthcareeng/users/oviessmann/nvsubquadratic/enroot/nvsubquadratic-x86_64.sqsh}" +#dfw lustre path: +IMAGE_NAME="${SQSH_IMAGE:-/lustre/fsw/portfolios/healthcareeng/projects/healthcareeng_research/users/oviessmann/nvsubquadratic/enroot/nvsubquadratic-x86_64.sqsh}" # Host paths WORKDIR="${PWD}" RUNS_DIR="${WORKDIR}/runs" -DATA_DIR="${WELL_HOST_DIR:-/lustre/fsw/healthcareeng_research/oviessmann/nvsubquadratic/data/well_data/datasets}" +DATA_DIR="${WELL_HOST_DIR:-/lustre/fsw/portfolios/healthcareeng/users/oviessmann/nvsubquadratic/data/well_data/datasets}" # Create necessary directories mkdir -p "${RUNS_DIR}" @@ -90,7 +94,7 @@ echo "$(date): Job ${SLURM_JOB_ID} started (W&B run ID: ${RUN_ID})" >> "${RESULT # ============================================================================ # Container mounts # ============================================================================ -REPO_DIR="/lustre/fsw/healthcareeng_research/oviessmann/nvsubquadratic" +REPO_DIR="/lustre/fsw/portfolios/healthcareeng/users/oviessmann/nvsubquadratic" WORK_DIR="/workspaces/nvSubquadratic-private" CONFIG_PATH="${WORK_DIR}/${CONFIG_FILE}" @@ -109,8 +113,8 @@ echo "Job ID: ${SLURM_JOB_ID}" echo "Run Name: ${RUN_NAME}" echo "Node(s): ${SLURM_NODELIST}" echo "Num nodes: ${SLURM_JOB_NUM_NODES}" -echo "GPUs per node: ${SLURM_NTASKS_PER_NODE}" -echo "Total GPUs: $((${SLURM_JOB_NUM_NODES} * ${SLURM_NTASKS_PER_NODE}))" +echo "GPUs per node: ${SLURM_GPUS_PER_NODE}" +echo "Total GPUs: $((${SLURM_JOB_NUM_NODES} * ${SLURM_GPUS_PER_NODE}))" echo "Config: ${CONFIG_FILE}" echo "Overrides: ${CONFIG_OVERRIDES}" echo "Container: ${IMAGE_NAME}" From a4661f35d797855f614d9db061739a60e4e1e44c Mon Sep 17 00:00:00 2001 From: OliviaViessmann Date: Tue, 9 Jun 2026 14:38:27 +0000 Subject: [PATCH 10/13] well: add helmholtz_staircase TRACKER.md with Attention and Hyena+G results Attention p8/p4 and Hyena+G p8/p4/p2 results logged. Attention OOMs at p2. CNextU-net row left pending (run still in progress). --- .../v2/gray_scott_reaction_diffusion/hyena.py | 1 - .../well/v2/helmholtz_staircase/TRACKER.md | 55 +++++++++++++++++++ 2 files changed, 55 insertions(+), 1 deletion(-) create mode 100644 examples/well/v2/helmholtz_staircase/TRACKER.md diff --git a/examples/well/v2/gray_scott_reaction_diffusion/hyena.py b/examples/well/v2/gray_scott_reaction_diffusion/hyena.py index d8013f8c..6daa5422 100644 --- a/examples/well/v2/gray_scott_reaction_diffusion/hyena.py +++ b/examples/well/v2/gray_scott_reaction_diffusion/hyena.py @@ -116,7 +116,6 @@ def get_config() -> ExperimentConfig: dim=NUM_HIDDEN_CHANNELS, ), qk_norm_cfg=LazyConfig(L2Norm)(), - use_rope=False, ), init_method_in=small_init, init_method_out=partial_wang_init_fn_with_num_layers(num_layers=NUM_BLOCKS), diff --git a/examples/well/v2/helmholtz_staircase/TRACKER.md b/examples/well/v2/helmholtz_staircase/TRACKER.md new file mode 100644 index 00000000..904990a8 --- /dev/null +++ b/examples/well/v2/helmholtz_staircase/TRACKER.md @@ -0,0 +1,55 @@ +# Helmholtz Staircase v2 + +## Goal + +Compare **CNextU-net** (baseline), **Attention**, and **Hyena + Gaussian mask** +on `helmholtz_staircase` (1024×256, open BCs). + +All runs share the same training recipe (AdamW, cosine schedule with 5% warmup, bf16-mixed, grad_clip=1.0). + +Paper baseline: CNextU-net VRMSE = 0.02758, FNO = 0.00046 (dominant). + +## Run order + +1. CNextU-net (baseline, no patch size) ← still running +1. Attention p8, p4 (p2 OOM) +1. Hyena+G p8, p4, p2 + +## Tokens per image + +| Patch | Tokens/img | +| ----- | ---------- | +| 8 | 4,096 | +| 4 | 16,384 | +| 2 | 65,536 | + +## Results — CNextU-net (baseline) + +| Batch/GPU | GPUs | WandB Run | val/VRMSE | it/s | total steps | +| --------- | ---- | --------- | --------- | ---- | ----------- | +| 24 | 1 | running | — | — | — | + +## Results — Attention + +| Patch | Tokens/img | Batch/GPU | LR | GPUs | WandB Run | val/VRMSE | it/s | total steps | +| ----- | ---------- | --------- | ---- | ---- | ----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------- | --------- | ------ | ----------- | +| 8 | 4,096 | 68 | 3e-3 | 1 | OV\_\_workspaces_nvSubquadratic-private_examples_well_v2_helmholtz_staircase_attention_bs_64_enabled_True_experiment_dir\_/workspace/results_lr_0.003_num_nodes_1_patch_size_8_run_start_time_1780763270_ru | 0.004970 | 0.5900 | 50,973 | +| 4 | 16,384 | 36 | 3e-3 | 1 | OV\_\_workspaces_nvSubquadratic-private_examples_well_v2_helmholtz_staircase_attention_bs_36_enabled_True_experiment_dir\_/workspace/results_lr_0.003_num_nodes_1_patch_size_4_run_start_time_1780777733_ru | 0.014651 | 0.2120 | 18,318 | +| 2 | 65,536 | — | — | — | OOM at bs≤12 | — | — | — | + +## Results — Hyena + Gaussian Mask + +| Patch | Tokens/img | Batch/GPU | LR | GPUs | WandB Run | val/VRMSE | it/s | total steps | +| ----- | ---------- | --------- | ---- | ---- | ----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------- | --------- | ------ | ----------- | +| 8 | 4,096 | 48 | 3e-4 | 1 | OV\_\_workspaces_nvSubquadratic-private_examples_well_v2_helmholtz_staircase_hyena_gaussian_mask_bs_48_enabled_True_experiment_dir\_/workspace/results_lr_0.0003_num_nodes_1_patch_size_8_run_start_time_17 | 0.006125 | 0.7212 | 49,979 | +| 4 | 16,384 | 48 | 3e-4 | 1 | OV\_\_workspaces_nvSubquadratic-private_examples_well_v2_helmholtz_staircase_hyena_gaussian_mask_bs_48_enabled_True_experiment_dir\_/workspace/results_lr_0.0003_num_nodes_1_patch_size_4_run_start_time_17 | 0.011989 | 0.2432 | 17,511 | +| 2 | 65,536 | 12 | 3e-4 | 1 | OV\_\_workspaces_nvSubquadratic-private_examples_well_v2_helmholtz_staircase_hyena_gaussian_mask_bs_12_enabled_True_experiment_dir\_/workspace/results_lr_0.0003_num_nodes_1_patch_size_2_run_start_time_17 | 0.031442 | 0.3148 | 22,099 | + +## Notes + +- **Attention p2**: OOM at all tested batch sizes (≤12); not run. +- **Hyena+G uses gradient checkpointing** (`GRADIENT_CHECKPOINTING = True` in `hyena.py`) to fit larger patch sizes in memory. +- **LR difference**: Attention converged best at lr=3e-3; Hyena+G at lr=3e-4 — a 10× gap vs Attention. This mirrors the supernova finding where Hyena+G preferred lower LRs. +- **Patch-size trend (Hyena+G)**: VRMSE degrades as patch size decreases (p8: 0.006 → p4: 0.012 → p2: 0.031). This is the **opposite** of the supernova_explosion_64 trend, consistent with helmholtz being a high-frequency problem where the large effective receptive field at coarser patches is actually beneficial. Smaller patches give more tokens but reduce the kernel span relative to the wavelength. +- **Attention p4 step count** (18,318) and **Hyena+G p4** (17,511) are much lower than p8 (~50k) — these may not be fully converged; flag if comparing head-to-head. +- **CNextU-net** still running; results to be filled in once complete. From 680efdbbd3ec74c1bc680ddac535a03fe97fbb9b Mon Sep 17 00:00:00 2001 From: OliviaViessmann Date: Thu, 11 Jun 2026 13:41:49 +0000 Subject: [PATCH 11/13] final CNextU-net results added to helmholtz TRACKER --- examples/well/v2/helmholtz_staircase/TRACKER.md | 10 +++++----- 1 file changed, 5 insertions(+), 5 deletions(-) diff --git a/examples/well/v2/helmholtz_staircase/TRACKER.md b/examples/well/v2/helmholtz_staircase/TRACKER.md index 904990a8..bdcde003 100644 --- a/examples/well/v2/helmholtz_staircase/TRACKER.md +++ b/examples/well/v2/helmholtz_staircase/TRACKER.md @@ -11,7 +11,7 @@ Paper baseline: CNextU-net VRMSE = 0.02758, FNO = 0.00046 (dominant). ## Run order -1. CNextU-net (baseline, no patch size) ← still running +1. CNextU-net (baseline, no patch size) ✓ 1. Attention p8, p4 (p2 OOM) 1. Hyena+G p8, p4, p2 @@ -25,9 +25,9 @@ Paper baseline: CNextU-net VRMSE = 0.02758, FNO = 0.00046 (dominant). ## Results — CNextU-net (baseline) -| Batch/GPU | GPUs | WandB Run | val/VRMSE | it/s | total steps | -| --------- | ---- | --------- | --------- | ---- | ----------- | -| 24 | 1 | running | — | — | — | +| Batch/GPU | GPUs | WandB Run | val/VRMSE | it/s | total steps | +| --------- | ---- | ----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------- | --------- | ------ | ----------- | +| 42 | 1 | OV\_\_workspaces_nvSubquadratic-private_examples_well_v2_helmholtz_staircase_cfg_unet_convnext_bs_42_enabled_True_experiment_dir\_/workspace/results_lr_0.003_num_nodes_1_run_start_time_1781084507_run_tim | 0.004504 | 0.8176 | 72,115 | ## Results — Attention @@ -52,4 +52,4 @@ Paper baseline: CNextU-net VRMSE = 0.02758, FNO = 0.00046 (dominant). - **LR difference**: Attention converged best at lr=3e-3; Hyena+G at lr=3e-4 — a 10× gap vs Attention. This mirrors the supernova finding where Hyena+G preferred lower LRs. - **Patch-size trend (Hyena+G)**: VRMSE degrades as patch size decreases (p8: 0.006 → p4: 0.012 → p2: 0.031). This is the **opposite** of the supernova_explosion_64 trend, consistent with helmholtz being a high-frequency problem where the large effective receptive field at coarser patches is actually beneficial. Smaller patches give more tokens but reduce the kernel span relative to the wavelength. - **Attention p4 step count** (18,318) and **Hyena+G p4** (17,511) are much lower than p8 (~50k) — these may not be fully converged; flag if comparing head-to-head. -- **CNextU-net** still running; results to be filled in once complete. +- **CNextU-net** val/VRMSE = 0.004504 at 72,115 steps — slightly better than Attention p8 (0.004970) and well below the paper baseline (0.02758). From 11a3edb93f4998d5c0179cce54f5c77498a19849 Mon Sep 17 00:00:00 2001 From: OliviaViessmann Date: Thu, 11 Jun 2026 13:51:12 +0000 Subject: [PATCH 12/13] final CNextU-net results added to helmholtz TRACKER --- examples/well/v2/helmholtz_staircase/TRACKER.md | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/examples/well/v2/helmholtz_staircase/TRACKER.md b/examples/well/v2/helmholtz_staircase/TRACKER.md index bdcde003..c1262639 100644 --- a/examples/well/v2/helmholtz_staircase/TRACKER.md +++ b/examples/well/v2/helmholtz_staircase/TRACKER.md @@ -27,7 +27,7 @@ Paper baseline: CNextU-net VRMSE = 0.02758, FNO = 0.00046 (dominant). | Batch/GPU | GPUs | WandB Run | val/VRMSE | it/s | total steps | | --------- | ---- | ----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------- | --------- | ------ | ----------- | -| 42 | 1 | OV\_\_workspaces_nvSubquadratic-private_examples_well_v2_helmholtz_staircase_cfg_unet_convnext_bs_42_enabled_True_experiment_dir\_/workspace/results_lr_0.003_num_nodes_1_run_start_time_1781084507_run_tim | 0.004504 | 0.8176 | 72,115 | +| 42 | 1 | OV\_\_workspaces_nvSubquadratic-private_examples_well_v2_helmholtz_staircase_cfg_unet_convnext_bs_42_enabled_True_experiment_dir\_/workspace/results_lr_0.003_num_nodes_1_run_start_time_1781084507_run_tim | 0.004504 | 490.58 | 72,115 | ## Results — Attention From 68e8484898d8b8af740f31f93b9e6ecac0667197 Mon Sep 17 00:00:00 2001 From: OliviaViessmann Date: Thu, 11 Jun 2026 14:20:44 +0000 Subject: [PATCH 13/13] final CNextU-net results added to helmholtz TRACKER --- examples/well/v2/helmholtz_staircase/TRACKER.md | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/examples/well/v2/helmholtz_staircase/TRACKER.md b/examples/well/v2/helmholtz_staircase/TRACKER.md index c1262639..1ea18f51 100644 --- a/examples/well/v2/helmholtz_staircase/TRACKER.md +++ b/examples/well/v2/helmholtz_staircase/TRACKER.md @@ -27,7 +27,7 @@ Paper baseline: CNextU-net VRMSE = 0.02758, FNO = 0.00046 (dominant). | Batch/GPU | GPUs | WandB Run | val/VRMSE | it/s | total steps | | --------- | ---- | ----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------- | --------- | ------ | ----------- | -| 42 | 1 | OV\_\_workspaces_nvSubquadratic-private_examples_well_v2_helmholtz_staircase_cfg_unet_convnext_bs_42_enabled_True_experiment_dir\_/workspace/results_lr_0.003_num_nodes_1_run_start_time_1781084507_run_tim | 0.004504 | 490.58 | 72,115 | +| 42 | 1 | OV\_\_workspaces_nvSubquadratic-private_examples_well_v2_helmholtz_staircase_cfg_unet_convnext_bs_42_enabled_True_experiment_dir\_/workspace/results_lr_0.003_num_nodes_1_run_start_time_1781084507_run_tim | 0.004504 | 0.8347 | 72,115 | ## Results — Attention