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/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 new file mode 100644 index 00000000..3011faf5 --- /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 `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 p8 +1. Hyena+G p8, p4, p2 + +## Results — CNextU-net (baseline) + +| 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 | 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 | 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/_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..952dec70 --- /dev/null +++ b/examples/well/v2/euler_multi_quadrants_periodicBC/attention.py @@ -0,0 +1,111 @@ +"""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, + 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), + ), + 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..c3ea4966 --- /dev/null +++ b/examples/well/v2/euler_multi_quadrants_periodicBC/hyena.py @@ -0,0 +1,143 @@ +"""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 32×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. +""" + +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:'512 // ${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/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..0d40a091 --- /dev/null +++ b/examples/well/v2/euler_multi_quadrants_periodicBC/hyena_gaussian_mask.py @@ -0,0 +1,32 @@ +"""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. +""" + +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 # 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 euler_multi_quadrants_periodicBC.""" + 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/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 | 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/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/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/TRACKER.md b/examples/well/v2/helmholtz_staircase/TRACKER.md new file mode 100644 index 00000000..1ea18f51 --- /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) ✓ +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 | +| --------- | ---- | ----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------- | --------- | ------ | ----------- | +| 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 + +| 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** val/VRMSE = 0.004504 at 72,115 steps — slightly better than Attention p8 (0.004970) and well below the paper baseline (0.02758). 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/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/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/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/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 d26701f7..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 @@ -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 diff --git a/experiments/trainer.py b/experiments/trainer.py index 76cdcdf6..17236b6e 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 @@ -120,6 +121,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 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 <> "${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}"