diff --git a/AGENTS.md b/AGENTS.md
index 1b2d858..2698812 100644
--- a/AGENTS.md
+++ b/AGENTS.md
@@ -92,6 +92,20 @@ below the ensemble member-to-member spread (docs/validation.md).
- **Triton IEEE fp32 dots do not use tensor cores** (~10× slower on sm_120); under `highest`
the kernel uses 3xTF32 (`strict_fp32=True` for IEEE). Turing (sm_75) has no TF32 → the
XLA attention path is selected automatically.
+- **Attention backward: precision and shared memory**: all attention kernels use one dot precision
+ (the backward recomputes P from the forward's log-sum-exp). 3xTF32 keeps hi/lo operand copies in
+ shared memory: the forward (2 dots per tile) fits sm_120's 99 KB in 64-row programs, the backward
+ kernels (3–4 dots) asked for 147 KB, so under `highest` they run in 32-row programs
+ (`bwd_split=2`): 115 ms/layer forward+backward (797 ms with IEEE dots). Whole-model gradients at
+ `highest` vs the official path (H200): 3xTF32 attention ~1.7× the official-vs-official floor
+ with a smooth loss, ~5× with the CRPS; `strict_fp32=True` at the floor. Matching forward and
+ backward precision did not measurably change this; the forward's precision did. In Pallas
+ interpret mode, `pl.program_id` is unavailable inside `fori_loop` bodies: read it at kernel top
+ level.
+- **CRPS gradients are noisy references**: the fair CRPS hits non-differentiable points exactly
+ (zero precipitation, land-filled SST), so rounding-level forward differences flip signs. Between
+ two runs of the same comparison (official vs XLA-autodiff path, H200, `highest`) the gradient
+ difference went from 1.1e-3 to 4.3e-3. Judge implementations with the smooth `--loss linear`.
- **XLA autotuning is timing-based**: kernel picks (and last bits) differ between runs; once, an
A100 fp32 run deviated at TF32 level (1 of 6, not reproduced). Use `--autotune-cache` for
reproducible runs and `--no-autotune` for strict fp32.
@@ -99,6 +113,13 @@ below the ensemble member-to-member spread (docs/validation.md).
`dynamic_update_slice`, or XLA fuses the read into the update and copies the whole buffer.
- **Transformer layers as a loop** (`hk.scan` + `hk.switch`) avoid XLA heap fragmentation
(~2 GiB with 24 unrolled layers).
+- **Remat inside raw JAX transforms**: `hk.remat` threads Haiku state (the RNG key) and leaks it
+ out of a raw `jax.lax.map`/`lax.scan` (UnexpectedTracerError on a jit call; with AOT
+ `.lower().compile()` it surfaced as "compiled for N inputs but called with M"). Use
+ `jax.checkpoint` for pure-JAX bodies there; `hk.remat` only inside `hk.scan`/`hk.map`.
+- **Independent remat regions run side by side**: two rematerialised calls with no data dependence
+ (e.g. one per ensemble member) did not lower peak memory, because XLA overlapped their backward
+ passes (46.1 vs 45.8 GiB). An `hk.scan` over them forces one at a time (25.5 GiB).
- **Nsight Compute**: JAX 0.11's CUDA 13.4 runtime crashes ncu 2025.3; use ncu ≥ 2026.3.
- **eccodes** must be imported after pyproj (heap corruption at exit otherwise): keep it lazy.
- **Compile time**: ~1 min on a fast desktop CPU, 2–3 min on cloud vCPUs; building the tile
diff --git a/CHANGELOG.md b/CHANGELOG.md
index 229c1f7..8ef67fc 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -2,6 +2,22 @@
## Unreleased
+- Head-only fine-tuning (prototype): `faster_weathernext.finetune` trains the mesh→grid GNN and
+ grid decoder (5.6M of WeatherNext 2's 184M parameters) with the official training loss (fair
+ CRPS, 2 members) while the encoder, grid→mesh GNN and transformer stay frozen. Gradients never
+ enter the frozen part, so no attention backward pass is needed. New option `remat` (off by
+ default; inference unchanged) rematerialises each GNN / encoder / decoder block in the backward
+ pass. `ifs.training_example` builds one-step examples from IFS open data, and
+ `scripts/finetune.py` / `scripts/verify_grads.py` train and validate. New `[train]`
+ extra (optax). Validation and memory: docs/validation.md §8.
+- The fused attention kernel is differentiable (FlashAttention-2 backward as two more tiled
+ Pallas kernels, dQ and dK/dV, over the same tile lists), so every parameter can be trained:
+ `finetune.split(params, "all")`, `scripts/finetune.py --modules all`. `remat` now also
+ rematerialises each transformer layer and each XLA attention chunk, and the new option
+ `sequential_samples` runs the 2 ensemble members one at a time in training (same noise; half
+ the activation memory). Forward and inference unchanged. Whole-model gradients checked against
+ the official path on an H200 (`cloud/modal_grads.py`; the official backward needs 102 GiB even
+ for the 1° Mini model). Validation: docs/validation.md §9.
- IFS open-data inputs fall back across mirrors: AWS, then Google Cloud
(`storage.googleapis.com/ecmwf-open-data`), then ECMWF (`data.ecmwf.int`). The AWS bucket
throttles a freshly published run with 503 SlowDown on every request, sometimes for hours;
diff --git a/README.md b/README.md
index 912456b..8c0cfd8 100644
--- a/README.md
+++ b/README.md
@@ -136,6 +136,27 @@ flowchart LR
Details and measurements: [docs/how-it-works.md](docs/how-it-works.md).
+## Fine-tuning (experimental)
+
+The same patch makes WeatherNext 2 trainable on one GPU. The fused attention kernel has a backward pass, and the blocked graph networks recompute their blocks instead of storing them, so gradients fit where the official code needs ~100 GiB even for the 1° Mini model.
+
+| what is trained | parameters | memory (0.25°, one step, 2 ensemble members) | time per step (RTX 5090) |
+|---|---|---|---|
+| head: mesh→grid GNN + decoder | 5.6M | 14.1 GiB (within a 24 GB card's pool ¹) | 4.0 s |
+| everything | 184M | 28.6 GiB (32 GB card, `XLA_PYTHON_CLIENT_MEM_FRACTION=0.95`) | 8.5 s |
+
+¹ Measured by capping the JAX pool on an RTX 5090 at what a 24 GB card gets at `XLA_PYTHON_CLIENT_MEM_FRACTION=0.9`; no real 24 GB card tried yet.
+
+```bash
+pip install "faster-weathernext[cuda12,weathernext,train] @ git+https://github.com/Raymondlol/Faster-WeatherNext.git"
+python scripts/finetune.py --train 2026101912 2026101918 2026102000 --eval 2026103012 \
+ --steps 30 --mask sea_surface_temperature --out head.npz
+```
+
+The loss is the official training loss (fair CRPS over two ensemble members); examples come from ECMWF open data. Gradients match the official implementation as closely as two official attention implementations match each other (checked end to end on an H200, where the official backward fits). The saved checkpoint loads with the official code.
+
+This is a prototype: one-step (6 h) training only, one example per step, no rollout evaluation. A fine-tuned model is a different model — the "same model" claims above cover inference with the official checkpoints, not the skill of a fine-tuned one. Details: [docs/validation.md](docs/validation.md#8-head-only-fine-tuning-prototype).
+
## FAQ
@@ -177,7 +198,7 @@ Yes. Install `faster-weathernext` (without the `[weathernext]` extra) into the e
Can I train or fine-tune with this?
-Not yet — the fused kernel and blocked graph networks implement the forward pass only.
+Yes, experimentally — see [Fine-tuning](#fine-tuning-experimental). The head (5.6M parameters) fits a 24 GB card's memory pool, the whole model (184M) a 32 GB card. Training from scratch is out of reach of a single GPU (the paper reports ~490 TPU-days per model).
## Citation and licenses
diff --git a/cloud/modal_canary.py b/cloud/modal_canary.py
index 9ce3dbb..8554d68 100644
--- a/cloud/modal_canary.py
+++ b/cloud/modal_canary.py
@@ -30,7 +30,7 @@
.apt_install("git")
.add_local_dir(REPO, "/opt/faster-weathernext", copy=True,
ignore=[".git", "**/__pycache__", "docs/results", "*.nc", "*.npz"])
- .run_commands("pip install '/opt/faster-weathernext[cuda12]'")
+ .run_commands("pip install '/opt/faster-weathernext[cuda12,weathernext]'")
)
diff --git a/cloud/modal_grads.py b/cloud/modal_grads.py
new file mode 100644
index 0000000..ce7b18a
--- /dev/null
+++ b/cloud/modal_grads.py
@@ -0,0 +1,86 @@
+"""Whole-model gradients vs the official path, on a GPU where the official backward fits (Modal H200).
+
+On a 32 GB GPU the official path's end-to-end backward does not fit even for the Mini model at 1°
+(compiled temp 82-116 GiB with the fair CRPS, 33-57 GiB with a one-member loss; docs/validation.md
+§9). On an H200 (141 GB) it does, so scripts/verify_grads.py --modules all can compare
+faster-weathernext (Pallas backward kernels) with the official GPU path directly, with a second
+official implementation as the floor. Each loss (crps, linear) runs in its own container; the first
+also runs the kernel tests (the Pallas kernel's first run on Hopper, sm_90).
+
+ modal run --detach cloud/modal_grads.py --tag grads01
+ modal volume get fwn-grads results/grads01 ./grads01
+ modal app stop -y # detached apps can linger after finishing
+
+Cost (Modal, 2026-11): H200 $4.54/h + 8 cores and 64 GiB ~ $0.9/h; ~40 min per container.
+"""
+import os
+import subprocess
+import time
+import urllib.parse
+import urllib.request
+
+import modal
+
+REPO = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
+BUCKET = "https://storage.googleapis.com/dm_graphcast/weathernext2"
+WEIGHTS = "params/WeatherNextCyclones_Mini_<2024.npz"
+DATA = "dataset/source-hres_forecast_init-2024-10-07 00:00:00_res-1.0_levels-13_steps-04.nc"
+
+app = modal.App("fwn-grads")
+vol = modal.Volume.from_name("fwn-grads", create_if_missing=True)
+image = (
+ modal.Image.debian_slim(python_version="3.12")
+ .apt_install("git")
+ .add_local_dir(REPO, "/opt/faster-weathernext", copy=True,
+ ignore=[".git", "**/__pycache__", "docs/results", "*.nc", "*.npz", ".pytest_cache"])
+ .run_commands("pip install '/opt/faster-weathernext[cuda12,weathernext,test]'")
+)
+
+
+@app.function(image=image, volumes={"/data": vol}, timeout=1800, cpu=2)
+def fetch():
+ for rel in (WEIGHTS, DATA):
+ dst = os.path.join("/data", os.path.basename(rel))
+ if not os.path.exists(dst):
+ urllib.request.urlretrieve(f"{BUCKET}/{urllib.parse.quote(rel)}", dst + ".part")
+ os.replace(dst + ".part", dst)
+ vol.commit()
+
+
+def _run(cmd, log, env=None):
+ t0 = time.time()
+ proc = subprocess.run(cmd, cwd="/opt/faster-weathernext", capture_output=True, text=True,
+ env={**os.environ, **(env or {})})
+ with open(log, "w") as f:
+ f.write(proc.stdout + proc.stderr + f"\nexit {proc.returncode} wall {time.time() - t0:.0f}s\n")
+ vol.commit()
+ return proc.returncode
+
+
+@app.function(image=image, gpu="H200", volumes={"/data": vol}, timeout=9000, cpu=8, memory=65536)
+def grads(loss: str, tag: str, run_tests: bool, precisions: str) -> str:
+ out = f"/data/results/{tag}"
+ os.makedirs(out, exist_ok=True)
+ smi = subprocess.run(["nvidia-smi", "--query-gpu=name,memory.total,compute_cap", "--format=csv,noheader"],
+ capture_output=True, text=True).stdout.strip()
+ lines = [f"[{loss}] {smi}"]
+ if run_tests:
+ rc = _run(["python", "-m", "pytest", "-q", "-rs", "tests/test_pallas_kernel.py"], f"{out}/pytest_h200.txt")
+ lines.append(f"kernel tests exit {rc}: " + open(f"{out}/pytest_h200.txt").read().strip().splitlines()[-2])
+ # The official mha path needs ~116 GiB with the CRPS: above JAX's default 75 % pool.
+ rc = _run(["python", "scripts/verify_grads.py", "--modules", "all", "--loss", loss,
+ "--data", f"/data/{os.path.basename(DATA)}", "--weights-dir", "/data",
+ "--precisions", precisions, "--out", f"{out}/full_grads_mini_{loss}_h200.json"],
+ f"{out}/full_grads_mini_{loss}_h200.log", env={"XLA_PYTHON_CLIENT_MEM_FRACTION": "0.95"})
+ log = open(f"{out}/full_grads_mini_{loss}_h200.log").read().splitlines()
+ lines += [f"verify exit {rc}"] + [l for l in log if l.startswith(("[", " ")) and "WARNING" not in l][-40:]
+ return "\n".join(lines)
+
+
+@app.local_entrypoint()
+def main(tag: str = "grads01", losses: str = "crps,linear", precisions: str = "highest,default",
+ tests: bool = True):
+ fetch.remote()
+ calls = [grads.spawn(loss, tag, tests and i == 0, precisions) for i, loss in enumerate(losses.split(","))]
+ for call in calls:
+ print(call.get(), flush=True)
diff --git a/docs/results/attn_bwd/attn_bwd_default_rtx5090.txt b/docs/results/attn_bwd/attn_bwd_default_rtx5090.txt
new file mode 100644
index 0000000..429a81b
--- /dev/null
+++ b/docs/results/attn_bwd/attn_bwd_default_rtx5090.txt
@@ -0,0 +1,10 @@
+/home/raymxnd/Workspace/WeatherForcast/.venv/lib/python3.12/site-packages/scipy/sparse/_index.py:216: SparseEfficiencyWarning: Changing the sparsity structure of a csr_matrix is expensive. lil and dok are more efficient.
+ self._set_arrayXarray(i, j, x)
+mask: 40962 nodes, 3101 keys/query, built in 35s
+hilbert tiles: 5480 keys/query-block
+xla 98.0 ms/layer (x24 = 2.35 s/step) forward+backward 228.8 ms/layer (x24 = 5.49 s)
+hilbert 9.3 ms/layer (x24 = 0.22 s/step) forward+backward 42.2 ms/layer (x24 = 1.01 s)
+hilbert_w4s2 9.3 ms/layer (x24 = 0.22 s/step) forward+backward 43.7 ms/layer (x24 = 1.05 s)
+hilbert vs xla: output rel RMS 2.82e-04 dq rel RMS 2.98e-04 dk rel RMS 2.94e-04 dv rel RMS 1.82e-05
+hilbert_w4s2 vs xla: output rel RMS 2.82e-04 dq rel RMS 2.98e-04 dk rel RMS 2.94e-04 dv rel RMS 1.82e-05
+exit 0
diff --git a/docs/results/attn_bwd/attn_bwd_highest_rtx5090.txt b/docs/results/attn_bwd/attn_bwd_highest_rtx5090.txt
new file mode 100644
index 0000000..fef6cce
--- /dev/null
+++ b/docs/results/attn_bwd/attn_bwd_highest_rtx5090.txt
@@ -0,0 +1,8 @@
+/home/raymxnd/Workspace/WeatherForcast/.venv/lib/python3.12/site-packages/scipy/sparse/_index.py:216: SparseEfficiencyWarning: Changing the sparsity structure of a csr_matrix is expensive. lil and dok are more efficient.
+ self._set_arrayXarray(i, j, x)
+mask: 40962 nodes, 3101 keys/query, built in 35s
+hilbert tiles: 5480 keys/query-block
+xla 128.4 ms/layer (x24 = 3.08 s/step) forward+backward 258.5 ms/layer (x24 = 6.20 s)
+hilbert 31.0 ms/layer (x24 = 0.74 s/step) forward+backward 115.4 ms/layer (x24 = 2.77 s)
+hilbert vs xla: output rel RMS 8.68e-07 dq rel RMS 1.14e-06 dk rel RMS 1.08e-06 dv rel RMS 8.76e-07
+exit 0
diff --git a/docs/results/attn_bwd/pytest_rtx5090.txt b/docs/results/attn_bwd/pytest_rtx5090.txt
new file mode 100644
index 0000000..e653986
--- /dev/null
+++ b/docs/results/attn_bwd/pytest_rtx5090.txt
@@ -0,0 +1,9 @@
+..................... [100%]
+=============================== warnings summary ===============================
+tests/test_pallas_kernel.py: 55 warnings
+ /home/raymxnd/Workspace/WeatherForcast/.venv/lib/python3.12/site-packages/jax/_src/pallas/pallas_call.py:926: DeprecationWarning: The Pallas Triton backend is deprecated and will be removed in a future JAX version. To keep using Pallas on GPU, please migrate to the Mosaic GPU backend. To keep using Triton, switch to the official Triton bindings and jax_triton.
+ return rule(ctx, *in_nodes, compiler_params=compiler_params, **params)
+
+-- Docs: https://docs.pytest.org/en/stable/how-to/capture-warnings.html
+21 passed, 55 warnings in 16.25s
+pytest exit 0
diff --git a/docs/results/canary_modal/v0.2/a100_cap8.json b/docs/results/canary_modal/v0.2/a100_cap8.json
new file mode 100644
index 0000000..83ad760
--- /dev/null
+++ b/docs/results/canary_modal/v0.2/a100_cap8.json
@@ -0,0 +1,26 @@
+{
+ "label": "a100_cap8",
+ "env": {
+ "XLA_PYTHON_CLIENT_MEM_FRACTION": "0.1750"
+ },
+ "ok": true,
+ "device": "NVIDIA A100-SXM4-40GB",
+ "jax": "0.11.2",
+ "faster_weathernext": "0.1.1",
+ "compute_capability": "8.0",
+ "attention": "pallas",
+ "pool_gib": 6.916015625,
+ "step_s": [
+ 121.82319569587708,
+ 1.1012904644012451,
+ 1.101973295211792,
+ 1.0997393131256104
+ ],
+ "finite": [
+ true,
+ true,
+ true,
+ true
+ ],
+ "peak_gib": 6.335763216018677
+}
\ No newline at end of file
diff --git a/docs/results/canary_modal/v0.2/a100_uncapped.json b/docs/results/canary_modal/v0.2/a100_uncapped.json
new file mode 100644
index 0000000..5a15935
--- /dev/null
+++ b/docs/results/canary_modal/v0.2/a100_uncapped.json
@@ -0,0 +1,27 @@
+{
+ "label": "a100_uncapped",
+ "env": {
+ "XLA_PYTHON_CLIENT_PREALLOCATE": "false"
+ },
+ "ok": true,
+ "device": "NVIDIA A100-SXM4-40GB",
+ "jax": "0.11.2",
+ "faster_weathernext": "0.1.1",
+ "compute_capability": "8.0",
+ "attention": "pallas",
+ "pool_gib": 29.63671875,
+ "highest_first_call_s": 134.05399680137634,
+ "step_s": [
+ 123.92502570152283,
+ 1.1289727687835693,
+ 1.102010726928711,
+ 1.1039471626281738
+ ],
+ "finite": [
+ true,
+ true,
+ true,
+ true
+ ],
+ "peak_gib": 6.1756885051727295
+}
\ No newline at end of file
diff --git a/docs/results/canary_modal/v0.2/a10g_cap8.json b/docs/results/canary_modal/v0.2/a10g_cap8.json
new file mode 100644
index 0000000..d9e9700
--- /dev/null
+++ b/docs/results/canary_modal/v0.2/a10g_cap8.json
@@ -0,0 +1,26 @@
+{
+ "label": "a10g_cap8",
+ "env": {
+ "XLA_PYTHON_CLIENT_MEM_FRACTION": "0.3113"
+ },
+ "ok": true,
+ "device": "NVIDIA A10",
+ "jax": "0.11.2",
+ "faster_weathernext": "0.1.1",
+ "compute_capability": "8.6",
+ "attention": "pallas",
+ "pool_gib": 6.8671875,
+ "step_s": [
+ 129.14016127586365,
+ 2.4310555458068848,
+ 2.4261364936828613,
+ 2.4280853271484375
+ ],
+ "finite": [
+ true,
+ true,
+ true,
+ true
+ ],
+ "peak_gib": 6.335774183273315
+}
\ No newline at end of file
diff --git a/docs/results/canary_modal/v0.2/a10g_uncapped.json b/docs/results/canary_modal/v0.2/a10g_uncapped.json
new file mode 100644
index 0000000..289f8b8
--- /dev/null
+++ b/docs/results/canary_modal/v0.2/a10g_uncapped.json
@@ -0,0 +1,27 @@
+{
+ "label": "a10g_uncapped",
+ "env": {
+ "XLA_PYTHON_CLIENT_PREALLOCATE": "false"
+ },
+ "ok": true,
+ "device": "NVIDIA A10",
+ "jax": "0.11.2",
+ "faster_weathernext": "0.1.1",
+ "compute_capability": "8.6",
+ "attention": "pallas",
+ "pool_gib": 16.544921875,
+ "highest_first_call_s": 126.86918115615845,
+ "step_s": [
+ 116.79271531105042,
+ 2.430650472640991,
+ 2.420816659927368,
+ 2.4241998195648193
+ ],
+ "finite": [
+ true,
+ true,
+ true,
+ true
+ ],
+ "peak_gib": 6.527284860610962
+}
\ No newline at end of file
diff --git a/docs/results/canary_modal/v0.2/cross_card_fp32.txt b/docs/results/canary_modal/v0.2/cross_card_fp32.txt
new file mode 100644
index 0000000..73bc473
--- /dev/null
+++ b/docs/results/canary_modal/v0.2/cross_card_fp32.txt
@@ -0,0 +1,21 @@
+# fp32-highest one-step outputs (every 4th grid point, 101 fields) vs RTX 5090 (same code: branch head-only-finetune, with the attention backward)
+== a100 (canary03)
+101 fields min corr 0.999913071645 max rel RMS 1.32e-02 median rel RMS 6.21e-04 NaN masks equal: True
+ cyclone_usa_rmw_disc corr 0.9999130716 rel RMS 1.32e-02
+ vertical_velocity@50 corr 0.9999802992 rel RMS 6.28e-03
+ vertical_velocity@100 corr 0.9999813400 rel RMS 6.11e-03
+== a10g (canary03)
+101 fields min corr 0.999999999878 max rel RMS 1.67e-05 median rel RMS 1.64e-06 NaN masks equal: True
+ cyclone_exists_gaussian_unit_mode corr 0.9999999999 rel RMS 1.67e-05
+ vertical_velocity@50 corr 0.9999999999 rel RMS 1.56e-05
+ vertical_velocity@100 corr 0.9999999999 rel RMS 1.22e-05
+== l4 (canary03)
+101 fields min corr 0.999999999829 max rel RMS 1.85e-05 median rel RMS 1.99e-06 NaN masks equal: True
+ vertical_velocity@50 corr 0.9999999998 rel RMS 1.85e-05
+ vertical_velocity@100 corr 0.9999999999 rel RMS 1.46e-05
+ vertical_velocity@200 corr 0.9999999999 rel RMS 1.41e-05
+== t4 (canary03)
+101 fields min corr 0.999999995324 max rel RMS 9.67e-05 median rel RMS 1.08e-05 NaN masks equal: True
+ vertical_velocity@50 corr 0.9999999953 rel RMS 9.67e-05
+ vertical_velocity@100 corr 0.9999999972 rel RMS 7.51e-05
+ vertical_velocity@150 corr 0.9999999974 rel RMS 7.21e-05
diff --git a/docs/results/canary_modal/v0.2/l4_cap8.json b/docs/results/canary_modal/v0.2/l4_cap8.json
new file mode 100644
index 0000000..ffb29f5
--- /dev/null
+++ b/docs/results/canary_modal/v0.2/l4_cap8.json
@@ -0,0 +1,26 @@
+{
+ "label": "l4_cap8",
+ "env": {
+ "XLA_PYTHON_CLIENT_MEM_FRACTION": "0.3112"
+ },
+ "ok": true,
+ "device": "NVIDIA L4",
+ "jax": "0.11.2",
+ "faster_weathernext": "0.1.1",
+ "compute_capability": "8.9",
+ "attention": "pallas",
+ "pool_gib": 6.8671875,
+ "step_s": [
+ 172.45196914672852,
+ 3.5930335521698,
+ 3.576650619506836,
+ 3.612750291824341
+ ],
+ "finite": [
+ true,
+ true,
+ true,
+ true
+ ],
+ "peak_gib": 6.335777282714844
+}
\ No newline at end of file
diff --git a/docs/results/canary_modal/v0.2/l4_uncapped.json b/docs/results/canary_modal/v0.2/l4_uncapped.json
new file mode 100644
index 0000000..4362751
--- /dev/null
+++ b/docs/results/canary_modal/v0.2/l4_uncapped.json
@@ -0,0 +1,27 @@
+{
+ "label": "l4_uncapped",
+ "env": {
+ "XLA_PYTHON_CLIENT_PREALLOCATE": "false"
+ },
+ "ok": true,
+ "device": "NVIDIA L4",
+ "jax": "0.11.2",
+ "faster_weathernext": "0.1.1",
+ "compute_capability": "8.9",
+ "attention": "pallas",
+ "pool_gib": 16.546875,
+ "highest_first_call_s": 188.54865312576294,
+ "step_s": [
+ 173.26529955863953,
+ 3.565401315689087,
+ 3.563096761703491,
+ 3.5803046226501465
+ ],
+ "finite": [
+ true,
+ true,
+ true,
+ true
+ ],
+ "peak_gib": 6.52728796005249
+}
\ No newline at end of file
diff --git a/docs/results/canary_modal/v0.2/rtx5090.json b/docs/results/canary_modal/v0.2/rtx5090.json
new file mode 100644
index 0000000..56c51fd
--- /dev/null
+++ b/docs/results/canary_modal/v0.2/rtx5090.json
@@ -0,0 +1,27 @@
+{
+ "label": "rtx5090",
+ "env": {
+ "XLA_PYTHON_CLIENT_PREALLOCATE": "false"
+ },
+ "ok": true,
+ "device": "NVIDIA GeForce RTX 5090",
+ "jax": "0.11.2",
+ "faster_weathernext": "0.1.1",
+ "compute_capability": "12.0",
+ "attention": "pallas",
+ "pool_gib": 23.494140625,
+ "highest_first_call_s": 81.99931263923645,
+ "step_s": [
+ 67.39950156211853,
+ 0.7325632572174072,
+ 0.7271003723144531,
+ 0.7285873889923096
+ ],
+ "finite": [
+ true,
+ true,
+ true,
+ true
+ ],
+ "peak_gib": 6.400768518447876
+}
\ No newline at end of file
diff --git a/docs/results/canary_modal/v0.2/t4_cap8.json b/docs/results/canary_modal/v0.2/t4_cap8.json
new file mode 100644
index 0000000..f3e1b71
--- /dev/null
+++ b/docs/results/canary_modal/v0.2/t4_cap8.json
@@ -0,0 +1,26 @@
+{
+ "label": "t4_cap8",
+ "env": {
+ "XLA_PYTHON_CLIENT_MEM_FRACTION": "0.4667"
+ },
+ "ok": true,
+ "device": "Tesla T4",
+ "jax": "0.11.2",
+ "faster_weathernext": "0.1.1",
+ "compute_capability": "7.5",
+ "attention": "xla",
+ "pool_gib": 6.80078125,
+ "step_s": [
+ 190.8113465309143,
+ 42.88824415206909,
+ 43.10090494155884,
+ 43.2635600566864
+ ],
+ "finite": [
+ true,
+ true,
+ true,
+ true
+ ],
+ "peak_gib": 6.296582221984863
+}
\ No newline at end of file
diff --git a/docs/results/canary_modal/v0.2/t4_uncapped.json b/docs/results/canary_modal/v0.2/t4_uncapped.json
new file mode 100644
index 0000000..afff13e
--- /dev/null
+++ b/docs/results/canary_modal/v0.2/t4_uncapped.json
@@ -0,0 +1,27 @@
+{
+ "label": "t4_uncapped",
+ "env": {
+ "XLA_PYTHON_CLIENT_PREALLOCATE": "false"
+ },
+ "ok": true,
+ "device": "Tesla T4",
+ "jax": "0.11.2",
+ "faster_weathernext": "0.1.1",
+ "compute_capability": "7.5",
+ "attention": "xla",
+ "pool_gib": 10.9296875,
+ "highest_first_call_s": 191.93480825424194,
+ "step_s": [
+ 191.03654861450195,
+ 42.91545248031616,
+ 42.96103096008301,
+ 43.121235609054565
+ ],
+ "finite": [
+ true,
+ true,
+ true,
+ true
+ ],
+ "peak_gib": 6.48809289932251
+}
\ No newline at end of file
diff --git a/docs/results/finetune/cpu_compile_noremat.json b/docs/results/finetune/cpu_compile_noremat.json
new file mode 100644
index 0000000..b1ac536
--- /dev/null
+++ b/docs/results/finetune/cpu_compile_noremat.json
@@ -0,0 +1,55 @@
+{
+ "model": "WeatherNext2",
+ "checkpoint": 1,
+ "device": "cpu",
+ "jax": "0.11.2",
+ "faster_weathernext": "0.1.1",
+ "args": {
+ "model": "WeatherNext2",
+ "checkpoint": 1,
+ "train": [
+ "2026102200"
+ ],
+ "eval": [],
+ "steps": 10,
+ "lr": 1e-05,
+ "weight_decay": 0.0,
+ "clip": 1.0,
+ "eval_keys": 2,
+ "precision": "default",
+ "seed": 0,
+ "no_precip": false,
+ "compile_only": true,
+ "no_remat": true,
+ "weights_dir": "/home/raymxnd/Workspace/WeatherForcast/weights",
+ "out": null,
+ "log": "/tmp/claude-1000/-home-raymxnd-Workspace-WeatherForcast/d4a93a6d-df87-4c4c-b7ce-cc8b5b5265d2/scratchpad/cpu_compile_noremat.json"
+ },
+ "masked": [
+ "cyclone_exists_gaussian_unit_mode",
+ "cyclone_all_wind_disc",
+ "cyclone_usa_wind_disc",
+ "cyclone_usa_r34_ne_radius_disc",
+ "cyclone_usa_r34_se_radius_disc",
+ "cyclone_usa_r34_sw_radius_disc",
+ "cyclone_usa_r34_nw_radius_disc",
+ "cyclone_usa_r50_ne_radius_disc",
+ "cyclone_usa_r50_se_radius_disc",
+ "cyclone_usa_r50_sw_radius_disc",
+ "cyclone_usa_r50_nw_radius_disc",
+ "cyclone_usa_r64_ne_radius_disc",
+ "cyclone_usa_r64_se_radius_disc",
+ "cyclone_usa_r64_sw_radius_disc",
+ "cyclone_usa_r64_nw_radius_disc",
+ "cyclone_usa_rmw_disc",
+ "cyclone_usa_pres_disc"
+ ],
+ "head_params": 5572967,
+ "total_params": 183785575,
+ "compile": {
+ "seconds": 59.72055411338806,
+ "temp_gib": 249.57595263421535,
+ "argument_gib": 1.7898386642336845,
+ "output_gib": 0.062285151332616806
+ }
+}
\ No newline at end of file
diff --git a/docs/results/finetune/cpu_compile_noremat.log b/docs/results/finetune/cpu_compile_noremat.log
new file mode 100644
index 0000000..fe375ba
--- /dev/null
+++ b/docs/results/finetune/cpu_compile_noremat.log
@@ -0,0 +1,32 @@
+example 2026102200 built in 2.0s
+WeatherNext2 ckpt 1: training 5.57M of 183.8M params; masked targets: 17 (cyclone channels only)
+WARNING:root:No normalization location found for cyclone_exists_gaussian_unit_mode
+WARNING:root:No normalization scale found for cyclone_exists_gaussian_unit_mode
+WARNING:absl:No coordinates found for dim noise_channels of array Size: 256B
+/home/raymxnd/Workspace/WeatherForcast/.venv/lib/python3.12/site-packages/scipy/sparse/_index.py:216: SparseEfficiencyWarning: Changing the sparsity structure of a csr_matrix is expensive. lil and dok are more efficient.
+/home/raymxnd/Workspace/WeatherForcast/weathernext/weathernext/weathernext2/architecture.py:286: FutureWarning: Deleting a single level of a MultiIndex is deprecated. Previously, this deleted all levels of a MultiIndex. Please also drop the following variables: {'batch'} to avoid an error in the future.
+train step compiled in 60s: temp 249.58 GiB, arguments 1.79 GiB
+wrote /tmp/claude-1000/-home-raymxnd-Workspace-WeatherForcast/d4a93a6d-df87-4c4c-b7ce-cc8b5b5265d2/scratchpad/cpu_compile_noremat.json
+ Command being timed: "/home/raymxnd/Workspace/WeatherForcast/.venv/bin/python scripts/finetune_head.py --train 2026102200 --compile-only --no-remat --weights-dir /home/raymxnd/Workspace/WeatherForcast/weights --log /tmp/claude-1000/-home-raymxnd-Workspace-WeatherForcast/d4a93a6d-df87-4c4c-b7ce-cc8b5b5265d2/scratchpad/cpu_compile_noremat.json"
+ User time (seconds): 89.82
+ System time (seconds): 1.95
+ Percent of CPU this job got: 143%
+ Elapsed (wall clock) time (h:mm:ss or m:ss): 1:04.03
+ Average shared text size (kbytes): 0
+ Average unshared data size (kbytes): 0
+ Average stack size (kbytes): 0
+ Average total size (kbytes): 0
+ Maximum resident set size (kbytes): 6856064
+ Average resident set size (kbytes): 0
+ Major (requiring I/O) page faults: 2
+ Minor (reclaiming a frame) page faults: 1448051
+ Voluntary context switches: 13672
+ Involuntary context switches: 3771
+ Swaps: 0
+ File system inputs: 232
+ File system outputs: 0
+ Socket messages sent: 0
+ Socket messages received: 0
+ Signals delivered: 0
+ Page size (bytes): 4096
+ Exit status: 0
diff --git a/docs/results/finetune/cpu_compile_remat.json b/docs/results/finetune/cpu_compile_remat.json
new file mode 100644
index 0000000..4e7329d
--- /dev/null
+++ b/docs/results/finetune/cpu_compile_remat.json
@@ -0,0 +1,55 @@
+{
+ "model": "WeatherNext2",
+ "checkpoint": 1,
+ "device": "cpu",
+ "jax": "0.11.2",
+ "faster_weathernext": "0.1.1",
+ "args": {
+ "model": "WeatherNext2",
+ "checkpoint": 1,
+ "train": [
+ "2026102200"
+ ],
+ "eval": [],
+ "steps": 10,
+ "lr": 1e-05,
+ "weight_decay": 0.0,
+ "clip": 1.0,
+ "eval_keys": 2,
+ "precision": "default",
+ "seed": 0,
+ "no_precip": false,
+ "compile_only": true,
+ "no_remat": false,
+ "weights_dir": "/home/raymxnd/Workspace/WeatherForcast/weights",
+ "out": null,
+ "log": "/tmp/claude-1000/-home-raymxnd-Workspace-WeatherForcast/d4a93a6d-df87-4c4c-b7ce-cc8b5b5265d2/scratchpad/cpu_compile_remat.json"
+ },
+ "masked": [
+ "cyclone_exists_gaussian_unit_mode",
+ "cyclone_all_wind_disc",
+ "cyclone_usa_wind_disc",
+ "cyclone_usa_r34_ne_radius_disc",
+ "cyclone_usa_r34_se_radius_disc",
+ "cyclone_usa_r34_sw_radius_disc",
+ "cyclone_usa_r34_nw_radius_disc",
+ "cyclone_usa_r50_ne_radius_disc",
+ "cyclone_usa_r50_se_radius_disc",
+ "cyclone_usa_r50_sw_radius_disc",
+ "cyclone_usa_r50_nw_radius_disc",
+ "cyclone_usa_r64_ne_radius_disc",
+ "cyclone_usa_r64_se_radius_disc",
+ "cyclone_usa_r64_sw_radius_disc",
+ "cyclone_usa_r64_nw_radius_disc",
+ "cyclone_usa_rmw_disc",
+ "cyclone_usa_pres_disc"
+ ],
+ "head_params": 5572967,
+ "total_params": 183785575,
+ "compile": {
+ "seconds": 60.59584712982178,
+ "temp_gib": 34.339843556284904,
+ "argument_gib": 1.7898386642336845,
+ "output_gib": 0.062285151332616806
+ }
+}
\ No newline at end of file
diff --git a/docs/results/finetune/cpu_compile_remat.log b/docs/results/finetune/cpu_compile_remat.log
new file mode 100644
index 0000000..ce36c27
--- /dev/null
+++ b/docs/results/finetune/cpu_compile_remat.log
@@ -0,0 +1,32 @@
+example 2026102200 built in 2.1s
+WeatherNext2 ckpt 1: training 5.57M of 183.8M params; masked targets: 17 (cyclone channels only)
+WARNING:root:No normalization location found for cyclone_exists_gaussian_unit_mode
+WARNING:root:No normalization scale found for cyclone_exists_gaussian_unit_mode
+WARNING:absl:No coordinates found for dim noise_channels of array Size: 256B
+/home/raymxnd/Workspace/WeatherForcast/.venv/lib/python3.12/site-packages/scipy/sparse/_index.py:216: SparseEfficiencyWarning: Changing the sparsity structure of a csr_matrix is expensive. lil and dok are more efficient.
+/home/raymxnd/Workspace/WeatherForcast/weathernext/weathernext/weathernext2/architecture.py:286: FutureWarning: Deleting a single level of a MultiIndex is deprecated. Previously, this deleted all levels of a MultiIndex. Please also drop the following variables: {'batch'} to avoid an error in the future.
+train step compiled in 61s: temp 34.34 GiB, arguments 1.79 GiB
+wrote /tmp/claude-1000/-home-raymxnd-Workspace-WeatherForcast/d4a93a6d-df87-4c4c-b7ce-cc8b5b5265d2/scratchpad/cpu_compile_remat.json
+ Command being timed: "/home/raymxnd/Workspace/WeatherForcast/.venv/bin/python scripts/finetune_head.py --train 2026102200 --compile-only --weights-dir /home/raymxnd/Workspace/WeatherForcast/weights --log /tmp/claude-1000/-home-raymxnd-Workspace-WeatherForcast/d4a93a6d-df87-4c4c-b7ce-cc8b5b5265d2/scratchpad/cpu_compile_remat.json"
+ User time (seconds): 91.80
+ System time (seconds): 2.48
+ Percent of CPU this job got: 144%
+ Elapsed (wall clock) time (h:mm:ss or m:ss): 1:05.15
+ Average shared text size (kbytes): 0
+ Average unshared data size (kbytes): 0
+ Average stack size (kbytes): 0
+ Average total size (kbytes): 0
+ Maximum resident set size (kbytes): 6840028
+ Average resident set size (kbytes): 0
+ Major (requiring I/O) page faults: 294
+ Minor (reclaiming a frame) page faults: 1550482
+ Voluntary context switches: 16666
+ Involuntary context switches: 4858
+ Swaps: 0
+ File system inputs: 1759536
+ File system outputs: 0
+ Socket messages sent: 0
+ Socket messages received: 0
+ Signals delivered: 0
+ Page size (bytes): 4096
+ Exit status: 0
diff --git a/docs/results/finetune/equivalence_rtx5090_attention_bwd.json b/docs/results/finetune/equivalence_rtx5090_attention_bwd.json
new file mode 100644
index 0000000..27d6db3
--- /dev/null
+++ b/docs/results/finetune/equivalence_rtx5090_attention_bwd.json
@@ -0,0 +1,85 @@
+{
+ "device": "NVIDIA GeForce RTX 5090",
+ "jax": "0.11.2",
+ "faster_weathernext": "0.1.1",
+ "cases": {
+ "highest/reference": {
+ "temp_gib": 20.793040931224823,
+ "step_s": 9.101388454437256,
+ "compile_s": 61.71154594421387,
+ "attention": "xla"
+ },
+ "highest/faster_weathernext": {
+ "temp_gib": 4.142392627894878,
+ "step_s": 1.491499662399292,
+ "compile_s": 77.88819169998169,
+ "attention": "pallas",
+ "vs_reference": {
+ "min_corr": 0.9999999998066557,
+ "max_rel_rms": 1.966784832823028e-05,
+ "median_rel_rms": 2.1372561506665216e-06,
+ "nan_masks_equal": true,
+ "n_fields": 101,
+ "worst": [
+ {
+ "field": "vertical_velocity@50",
+ "corr": 0.9999999998066557,
+ "rel_rms": 1.966784832823028e-05,
+ "nan_mask_equal": true
+ },
+ {
+ "field": "cyclone_exists_gaussian_unit_mode",
+ "corr": 0.9999999998370072,
+ "rel_rms": 1.9118753530698537e-05,
+ "nan_mask_equal": true
+ },
+ {
+ "field": "vertical_velocity@100",
+ "corr": 0.9999999998729742,
+ "rel_rms": 1.5948352947738497e-05,
+ "nan_mask_equal": true
+ }
+ ]
+ }
+ },
+ "default/reference": {
+ "temp_gib": 20.793040931224823,
+ "step_s": 5.41904616355896,
+ "compile_s": 60.78113007545471,
+ "attention": "xla"
+ },
+ "default/faster_weathernext": {
+ "temp_gib": 3.7908003255724907,
+ "step_s": 0.7766594886779785,
+ "compile_s": 66.80423951148987,
+ "attention": "pallas",
+ "vs_reference": {
+ "min_corr": 0.9999801944178773,
+ "max_rel_rms": 0.006535481648120584,
+ "median_rel_rms": 0.0005235910731192411,
+ "nan_masks_equal": true,
+ "n_fields": 101,
+ "worst": [
+ {
+ "field": "cyclone_exists_gaussian_unit_mode",
+ "corr": 0.9999801944178773,
+ "rel_rms": 0.006535481648120584,
+ "nan_mask_equal": true
+ },
+ {
+ "field": "vertical_velocity@50",
+ "corr": 0.9999878725708148,
+ "rel_rms": 0.00492621730693891,
+ "nan_mask_equal": true
+ },
+ {
+ "field": "vertical_velocity@100",
+ "corr": 0.9999925636089264,
+ "rel_rms": 0.0038569065213644085,
+ "nan_mask_equal": true
+ }
+ ]
+ }
+ }
+ }
+}
\ No newline at end of file
diff --git a/docs/results/finetune/equivalence_rtx5090_attention_bwd.log b/docs/results/finetune/equivalence_rtx5090_attention_bwd.log
new file mode 100644
index 0000000..28a3ed4
--- /dev/null
+++ b/docs/results/finetune/equivalence_rtx5090_attention_bwd.log
@@ -0,0 +1,57 @@
+WARNING:absl:Skipping gradient checkpointing for sequence length of 1
+WARNING:absl:No coordinates found for dim noise_channels of array Size: 128B
+JitTracer(float32[1,1,32])
+Dimensions without coordinates: batch, channels, noise_channels., using index for param name. This may make model surgery more error prone if changing coordinates at fine-tuning time.
+/home/raymxnd/Workspace/WeatherForcast/.venv/lib/python3.12/site-packages/scipy/sparse/_index.py:216: SparseEfficiencyWarning: Changing the sparsity structure of a csr_matrix is expensive. lil and dok are more efficient.
+ self._set_arrayXarray(i, j, x)
+WARNING:absl:Only the first input of the static gather operation is being used, others are being ignored at runtime.
+WARNING:absl:Only the first input of the static gather operation is being used, others are being ignored at runtime.
+WARNING:absl:Only the first input of the sorted segment sum operation is being used, others are being ignored at runtime.
+WARNING:absl:Only the first input of the static gather operation is being used, others are being ignored at runtime.
+WARNING:absl:Only the first input of the static gather operation is being used, others are being ignored at runtime.
+WARNING:absl:Only the first input of the static gather operation is being used, others are being ignored at runtime.
+WARNING:absl:Only the first input of the sorted segment sum operation is being used, others are being ignored at runtime.
+WARNING:root:No normalization scale found for cyclone_exists_gaussian_unit_mode
+WARNING:root:No normalization location found for cyclone_exists_gaussian_unit_mode
+W1101 01:09:22.845838 176910 hlo_rematerialization.cc:3363] Can't reduce memory use below 21.08GiB (22629749965 bytes) by rematerialization; only reduced to 24.29GiB (26085976296 bytes), down from 24.29GiB (26086000776 bytes) originally
+highest reference temp 20.79 GiB step 9.10s compile 62s attention=xla
+WARNING:absl:Skipping gradient checkpointing for sequence length of 1
+WARNING:absl:No coordinates found for dim noise_channels of array Size: 128B
+JitTracer(float32[1,1,32])
+Dimensions without coordinates: batch, channels, noise_channels., using index for param name. This may make model surgery more error prone if changing coordinates at fine-tuning time.
+/home/raymxnd/Workspace/WeatherForcast/.venv/lib/python3.12/site-packages/scipy/sparse/_index.py:216: SparseEfficiencyWarning: Changing the sparsity structure of a csr_matrix is expensive. lil and dok are more efficient.
+ self._set_arrayXarray(i, j, x)
+WARNING:root:No normalization scale found for cyclone_exists_gaussian_unit_mode
+WARNING:root:No normalization location found for cyclone_exists_gaussian_unit_mode
+W1101 01:10:58.924915 176910 hlo_rematerialization.cc:3363] Can't reduce memory use below 21.08GiB (22629749965 bytes) by rematerialization; only reduced to 28.82GiB (30943716076 bytes), down from 28.82GiB (30943728380 bytes) originally
+highest faster_weathernext temp 4.14 GiB step 1.49s compile 78s attention=pallas
+highest faster-weathernext vs reference: min corr 0.9999999998 max rel RMS 1.97e-05 median 2.14e-06 NaN masks equal True (101 fields)
+WARNING:absl:Skipping gradient checkpointing for sequence length of 1
+WARNING:absl:No coordinates found for dim noise_channels of array Size: 128B
+JitTracer(float32[1,1,32])
+Dimensions without coordinates: batch, channels, noise_channels., using index for param name. This may make model surgery more error prone if changing coordinates at fine-tuning time.
+/home/raymxnd/Workspace/WeatherForcast/.venv/lib/python3.12/site-packages/scipy/sparse/_index.py:216: SparseEfficiencyWarning: Changing the sparsity structure of a csr_matrix is expensive. lil and dok are more efficient.
+ self._set_arrayXarray(i, j, x)
+WARNING:absl:Only the first input of the static gather operation is being used, others are being ignored at runtime.
+WARNING:absl:Only the first input of the static gather operation is being used, others are being ignored at runtime.
+WARNING:absl:Only the first input of the sorted segment sum operation is being used, others are being ignored at runtime.
+WARNING:absl:Only the first input of the static gather operation is being used, others are being ignored at runtime.
+WARNING:absl:Only the first input of the static gather operation is being used, others are being ignored at runtime.
+WARNING:absl:Only the first input of the static gather operation is being used, others are being ignored at runtime.
+WARNING:absl:Only the first input of the sorted segment sum operation is being used, others are being ignored at runtime.
+WARNING:root:No normalization scale found for cyclone_exists_gaussian_unit_mode
+WARNING:root:No normalization location found for cyclone_exists_gaussian_unit_mode
+W1101 01:12:04.476291 176910 hlo_rematerialization.cc:3363] Can't reduce memory use below 21.08GiB (22629749965 bytes) by rematerialization; only reduced to 24.29GiB (26085976296 bytes), down from 24.29GiB (26086000776 bytes) originally
+default reference temp 20.79 GiB step 5.42s compile 61s attention=xla
+WARNING:absl:Skipping gradient checkpointing for sequence length of 1
+WARNING:absl:No coordinates found for dim noise_channels of array Size: 128B
+JitTracer(float32[1,1,32])
+Dimensions without coordinates: batch, channels, noise_channels., using index for param name. This may make model surgery more error prone if changing coordinates at fine-tuning time.
+/home/raymxnd/Workspace/WeatherForcast/.venv/lib/python3.12/site-packages/scipy/sparse/_index.py:216: SparseEfficiencyWarning: Changing the sparsity structure of a csr_matrix is expensive. lil and dok are more efficient.
+ self._set_arrayXarray(i, j, x)
+WARNING:root:No normalization scale found for cyclone_exists_gaussian_unit_mode
+WARNING:root:No normalization location found for cyclone_exists_gaussian_unit_mode
+W1101 01:13:22.252941 176910 hlo_rematerialization.cc:3363] Can't reduce memory use below 21.08GiB (22629749965 bytes) by rematerialization; only reduced to 28.82GiB (30943716076 bytes), down from 28.82GiB (30943728380 bytes) originally
+default faster_weathernext temp 3.79 GiB step 0.78s compile 67s attention=pallas
+default faster-weathernext vs reference: min corr 0.9999801944 max rel RMS 6.54e-03 median 5.24e-04 NaN masks equal True (101 fields)
+wrote docs/results/finetune/equivalence_rtx5090_attention_bwd.json
diff --git a/docs/results/finetune/equivalence_rtx5090_attention_bwd_default_repeat.json b/docs/results/finetune/equivalence_rtx5090_attention_bwd_default_repeat.json
new file mode 100644
index 0000000..9b7c504
--- /dev/null
+++ b/docs/results/finetune/equivalence_rtx5090_attention_bwd_default_repeat.json
@@ -0,0 +1,46 @@
+{
+ "device": "NVIDIA GeForce RTX 5090",
+ "jax": "0.11.2",
+ "faster_weathernext": "0.1.1",
+ "cases": {
+ "default/reference": {
+ "temp_gib": 20.793040931224823,
+ "step_s": 5.419211387634277,
+ "compile_s": 60.77202224731445,
+ "attention": "xla"
+ },
+ "default/faster_weathernext": {
+ "temp_gib": 3.7908000871539116,
+ "step_s": 0.9195766448974609,
+ "compile_s": 67.179283618927,
+ "attention": "pallas",
+ "vs_reference": {
+ "min_corr": 0.9999887671752711,
+ "max_rel_rms": 0.004740084327257787,
+ "median_rel_rms": 0.0005069677876871889,
+ "nan_masks_equal": true,
+ "n_fields": 101,
+ "worst": [
+ {
+ "field": "vertical_velocity@50",
+ "corr": 0.9999887671752711,
+ "rel_rms": 0.004740084327257787,
+ "nan_mask_equal": true
+ },
+ {
+ "field": "vertical_velocity@100",
+ "corr": 0.9999927712337858,
+ "rel_rms": 0.003802333470595173,
+ "nan_mask_equal": true
+ },
+ {
+ "field": "vertical_velocity@150",
+ "corr": 0.9999933002979728,
+ "rel_rms": 0.003660543185140332,
+ "nan_mask_equal": true
+ }
+ ]
+ }
+ }
+ }
+}
\ No newline at end of file
diff --git a/docs/results/finetune/equivalence_rtx5090_attention_bwd_default_repeat.log b/docs/results/finetune/equivalence_rtx5090_attention_bwd_default_repeat.log
new file mode 100644
index 0000000..ef17e45
--- /dev/null
+++ b/docs/results/finetune/equivalence_rtx5090_attention_bwd_default_repeat.log
@@ -0,0 +1,29 @@
+WARNING:absl:Skipping gradient checkpointing for sequence length of 1
+WARNING:absl:No coordinates found for dim noise_channels of array Size: 128B
+JitTracer(float32[1,1,32])
+Dimensions without coordinates: batch, channels, noise_channels., using index for param name. This may make model surgery more error prone if changing coordinates at fine-tuning time.
+/home/raymxnd/Workspace/WeatherForcast/.venv/lib/python3.12/site-packages/scipy/sparse/_index.py:216: SparseEfficiencyWarning: Changing the sparsity structure of a csr_matrix is expensive. lil and dok are more efficient.
+ self._set_arrayXarray(i, j, x)
+WARNING:absl:Only the first input of the static gather operation is being used, others are being ignored at runtime.
+WARNING:absl:Only the first input of the static gather operation is being used, others are being ignored at runtime.
+WARNING:absl:Only the first input of the sorted segment sum operation is being used, others are being ignored at runtime.
+WARNING:absl:Only the first input of the static gather operation is being used, others are being ignored at runtime.
+WARNING:absl:Only the first input of the static gather operation is being used, others are being ignored at runtime.
+WARNING:absl:Only the first input of the static gather operation is being used, others are being ignored at runtime.
+WARNING:absl:Only the first input of the sorted segment sum operation is being used, others are being ignored at runtime.
+WARNING:root:No normalization scale found for cyclone_exists_gaussian_unit_mode
+WARNING:root:No normalization location found for cyclone_exists_gaussian_unit_mode
+W1101 01:15:02.817995 183688 hlo_rematerialization.cc:3363] Can't reduce memory use below 21.08GiB (22629749965 bytes) by rematerialization; only reduced to 24.29GiB (26085976296 bytes), down from 24.29GiB (26086000776 bytes) originally
+default reference temp 20.79 GiB step 5.42s compile 61s attention=xla
+WARNING:absl:Skipping gradient checkpointing for sequence length of 1
+WARNING:absl:No coordinates found for dim noise_channels of array Size: 128B
+JitTracer(float32[1,1,32])
+Dimensions without coordinates: batch, channels, noise_channels., using index for param name. This may make model surgery more error prone if changing coordinates at fine-tuning time.
+/home/raymxnd/Workspace/WeatherForcast/.venv/lib/python3.12/site-packages/scipy/sparse/_index.py:216: SparseEfficiencyWarning: Changing the sparsity structure of a csr_matrix is expensive. lil and dok are more efficient.
+ self._set_arrayXarray(i, j, x)
+WARNING:root:No normalization scale found for cyclone_exists_gaussian_unit_mode
+WARNING:root:No normalization location found for cyclone_exists_gaussian_unit_mode
+W1101 01:16:21.302678 183688 hlo_rematerialization.cc:3363] Can't reduce memory use below 21.08GiB (22629749965 bytes) by rematerialization; only reduced to 28.82GiB (30943716076 bytes), down from 28.82GiB (30943728380 bytes) originally
+default faster_weathernext temp 3.79 GiB step 0.92s compile 67s attention=pallas
+default faster-weathernext vs reference: min corr 0.9999887672 max rel RMS 4.74e-03 median 5.07e-04 NaN masks equal True (101 fields)
+wrote /tmp/claude-1000/-home-raymxnd-Workspace-WeatherForcast/d4a93a6d-df87-4c4c-b7ce-cc8b5b5265d2/scratchpad/equiv_run_b.json
diff --git a/docs/results/finetune/equivalence_rtx5090_remat_option.json b/docs/results/finetune/equivalence_rtx5090_remat_option.json
new file mode 100644
index 0000000..0c62fd2
--- /dev/null
+++ b/docs/results/finetune/equivalence_rtx5090_remat_option.json
@@ -0,0 +1,85 @@
+{
+ "device": "NVIDIA GeForce RTX 5090",
+ "jax": "0.11.2",
+ "faster_weathernext": "0.1.1",
+ "cases": {
+ "highest/reference": {
+ "temp_gib": 20.793040931224823,
+ "step_s": 9.111310243606567,
+ "compile_s": 62.11035084724426,
+ "attention": "xla"
+ },
+ "highest/faster_weathernext": {
+ "temp_gib": 4.14239215105772,
+ "step_s": 1.4934382438659668,
+ "compile_s": 78.03873538970947,
+ "attention": "pallas",
+ "vs_reference": {
+ "min_corr": 0.9999999998049606,
+ "max_rel_rms": 1.975972130091474e-05,
+ "median_rel_rms": 2.098829744870732e-06,
+ "nan_masks_equal": true,
+ "n_fields": 101,
+ "worst": [
+ {
+ "field": "vertical_velocity@50",
+ "corr": 0.9999999998049606,
+ "rel_rms": 1.975972130091474e-05,
+ "nan_mask_equal": true
+ },
+ {
+ "field": "cyclone_exists_gaussian_unit_mode",
+ "corr": 0.9999999998550044,
+ "rel_rms": 1.7947416447884274e-05,
+ "nan_mask_equal": true
+ },
+ {
+ "field": "vertical_velocity@100",
+ "corr": 0.9999999998757233,
+ "rel_rms": 1.5772170452897742e-05,
+ "nan_mask_equal": true
+ }
+ ]
+ }
+ },
+ "default/reference": {
+ "temp_gib": 20.793040931224823,
+ "step_s": 5.4155449867248535,
+ "compile_s": 61.008819818496704,
+ "attention": "xla"
+ },
+ "default/faster_weathernext": {
+ "temp_gib": 4.0158998891711235,
+ "step_s": 0.7735049724578857,
+ "compile_s": 67.40848541259766,
+ "attention": "pallas",
+ "vs_reference": {
+ "min_corr": 0.9999881616485653,
+ "max_rel_rms": 0.004865874054835816,
+ "median_rel_rms": 0.0005143238789506393,
+ "nan_masks_equal": true,
+ "n_fields": 101,
+ "worst": [
+ {
+ "field": "vertical_velocity@50",
+ "corr": 0.9999881616485653,
+ "rel_rms": 0.004865874054835816,
+ "nan_mask_equal": true
+ },
+ {
+ "field": "vertical_velocity@100",
+ "corr": 0.9999926451431069,
+ "rel_rms": 0.0038353257136195897,
+ "nan_mask_equal": true
+ },
+ {
+ "field": "vertical_velocity@150",
+ "corr": 0.9999932103539011,
+ "rel_rms": 0.0036850448483533465,
+ "nan_mask_equal": true
+ }
+ ]
+ }
+ }
+ }
+}
\ No newline at end of file
diff --git a/docs/results/finetune/equivalence_rtx5090_remat_option.log b/docs/results/finetune/equivalence_rtx5090_remat_option.log
new file mode 100644
index 0000000..8222913
--- /dev/null
+++ b/docs/results/finetune/equivalence_rtx5090_remat_option.log
@@ -0,0 +1,90 @@
+gpu-run: queued behind /home/raymxnd/Workspace/WeatherForcast/.venv/bin/python scripts/verify_head_grad
+gpu-run: queued behind /home/raymxnd/Workspace/WeatherForcast/.venv/bin/python scripts/verify_head_grad
+gpu-run: queued behind /home/raymxnd/Workspace/WeatherForcast/.venv/bin/python scripts/verify_head_grad
+gpu-run: queued behind /home/raymxnd/Workspace/WeatherForcast/.venv/bin/python scripts/verify_head_grad
+gpu-run: queued behind /home/raymxnd/Workspace/WeatherForcast/.venv/bin/python scripts/verify_head_grad
+gpu-run: queued behind /home/raymxnd/Workspace/WeatherForcast/.venv/bin/python scripts/verify_head_grad
+gpu-run: queued behind /home/raymxnd/Workspace/WeatherForcast/.venv/bin/python scripts/verify_head_grad
+gpu-run: queued behind /home/raymxnd/Workspace/WeatherForcast/.venv/bin/python scripts/verify_head_grad
+gpu-run: queued behind /home/raymxnd/Workspace/WeatherForcast/.venv/bin/python scripts/verify_head_grad
+gpu-run: queued behind /home/raymxnd/Workspace/WeatherForcast/.venv/bin/python scripts/verify_head_grad
+gpu-run: queued behind /home/raymxnd/Workspace/WeatherForcast/.venv/bin/python scripts/verify_head_grad
+gpu-run: queued behind /home/raymxnd/Workspace/WeatherForcast/.venv/bin/python scripts/verify_head_grad
+gpu-run: queued behind /home/raymxnd/Workspace/WeatherForcast/.venv/bin/python scripts/verify_head_grad
+gpu-run: queued behind /home/raymxnd/Workspace/WeatherForcast/.venv/bin/python scripts/verify_head_grad
+gpu-run: queued behind /home/raymxnd/Workspace/WeatherForcast/.venv/bin/python scripts/verify_head_grad
+gpu-run: queued behind /home/raymxnd/Workspace/WeatherForcast/.venv/bin/python scripts/verify_head_grad
+gpu-run: queued behind /home/raymxnd/Workspace/WeatherForcast/.venv/bin/python scripts/verify_head_grad
+gpu-run: queued behind /home/raymxnd/Workspace/WeatherForcast/.venv/bin/python scripts/verify_head_grad
+gpu-run: queued behind /home/raymxnd/Workspace/WeatherForcast/.venv/bin/python scripts/verify_head_grad
+gpu-run: queued behind /home/raymxnd/Workspace/WeatherForcast/.venv/bin/python scripts/verify_head_grad
+gpu-run: queued behind /home/raymxnd/Workspace/WeatherForcast/.venv/bin/python scripts/verify_head_grad
+gpu-run: queued behind /home/raymxnd/Workspace/WeatherForcast/.venv/bin/python scripts/verify_head_grad
+gpu-run: queued behind /home/raymxnd/Workspace/WeatherForcast/.venv/bin/python scripts/verify_head_grad
+gpu-run: queued behind /home/raymxnd/Workspace/WeatherForcast/.venv/bin/python scripts/verify_head_grad
+gpu-run: queued behind /home/raymxnd/Workspace/WeatherForcast/.venv/bin/python scripts/verify_head_grad
+gpu-run: queued behind /home/raymxnd/Workspace/WeatherForcast/.venv/bin/python scripts/verify_head_grad
+gpu-run: queued behind /home/raymxnd/Workspace/WeatherForcast/.venv/bin/python scripts/verify_head_grad
+gpu-run: queued behind /home/raymxnd/Workspace/WeatherForcast/.venv/bin/python scripts/verify_head_grad
+gpu-run: queued behind /home/raymxnd/Workspace/WeatherForcast/.venv/bin/python scripts/verify_head_grad
+gpu-run: queued behind /home/raymxnd/Workspace/WeatherForcast/.venv/bin/python scripts/verify_head_grad
+gpu-run: queued behind /home/raymxnd/Workspace/WeatherForcast/.venv/bin/python scripts/finetune_head.py
+gpu-run: queued behind /home/raymxnd/Workspace/WeatherForcast/.venv/bin/python scripts/finetune_head.py
+gpu-run: queued behind /home/raymxnd/Workspace/WeatherForcast/.venv/bin/python scripts/finetune_head.py
+WARNING:absl:Skipping gradient checkpointing for sequence length of 1
+WARNING:absl:No coordinates found for dim noise_channels of array Size: 128B
+JitTracer(float32[1,1,32])
+Dimensions without coordinates: batch, channels, noise_channels., using index for param name. This may make model surgery more error prone if changing coordinates at fine-tuning time.
+/home/raymxnd/Workspace/WeatherForcast/.venv/lib/python3.12/site-packages/scipy/sparse/_index.py:216: SparseEfficiencyWarning: Changing the sparsity structure of a csr_matrix is expensive. lil and dok are more efficient.
+ self._set_arrayXarray(i, j, x)
+WARNING:absl:Only the first input of the static gather operation is being used, others are being ignored at runtime.
+WARNING:absl:Only the first input of the static gather operation is being used, others are being ignored at runtime.
+WARNING:absl:Only the first input of the sorted segment sum operation is being used, others are being ignored at runtime.
+WARNING:absl:Only the first input of the static gather operation is being used, others are being ignored at runtime.
+WARNING:absl:Only the first input of the static gather operation is being used, others are being ignored at runtime.
+WARNING:absl:Only the first input of the static gather operation is being used, others are being ignored at runtime.
+WARNING:absl:Only the first input of the sorted segment sum operation is being used, others are being ignored at runtime.
+WARNING:root:No normalization scale found for cyclone_exists_gaussian_unit_mode
+WARNING:root:No normalization location found for cyclone_exists_gaussian_unit_mode
+W1031 12:51:22.018697 1588239 hlo_rematerialization.cc:3363] Can't reduce memory use below 21.08GiB (22629749965 bytes) by rematerialization; only reduced to 24.29GiB (26085976296 bytes), down from 24.29GiB (26086000776 bytes) originally
+highest reference temp 20.79 GiB step 9.11s compile 62s attention=xla
+WARNING:absl:Skipping gradient checkpointing for sequence length of 1
+WARNING:absl:No coordinates found for dim noise_channels of array Size: 128B
+JitTracer(float32[1,1,32])
+Dimensions without coordinates: batch, channels, noise_channels., using index for param name. This may make model surgery more error prone if changing coordinates at fine-tuning time.
+/home/raymxnd/Workspace/WeatherForcast/.venv/lib/python3.12/site-packages/scipy/sparse/_index.py:216: SparseEfficiencyWarning: Changing the sparsity structure of a csr_matrix is expensive. lil and dok are more efficient.
+ self._set_arrayXarray(i, j, x)
+WARNING:root:No normalization scale found for cyclone_exists_gaussian_unit_mode
+WARNING:root:No normalization location found for cyclone_exists_gaussian_unit_mode
+W1031 12:52:58.344189 1588239 hlo_rematerialization.cc:3363] Can't reduce memory use below 21.08GiB (22629749965 bytes) by rematerialization; only reduced to 28.82GiB (30943716076 bytes), down from 28.82GiB (30943728380 bytes) originally
+highest faster_weathernext temp 4.14 GiB step 1.49s compile 78s attention=pallas
+highest faster-weathernext vs reference: min corr 0.9999999998 max rel RMS 1.98e-05 median 2.10e-06 NaN masks equal True (101 fields)
+WARNING:absl:Skipping gradient checkpointing for sequence length of 1
+WARNING:absl:No coordinates found for dim noise_channels of array Size: 128B
+JitTracer(float32[1,1,32])
+Dimensions without coordinates: batch, channels, noise_channels., using index for param name. This may make model surgery more error prone if changing coordinates at fine-tuning time.
+/home/raymxnd/Workspace/WeatherForcast/.venv/lib/python3.12/site-packages/scipy/sparse/_index.py:216: SparseEfficiencyWarning: Changing the sparsity structure of a csr_matrix is expensive. lil and dok are more efficient.
+ self._set_arrayXarray(i, j, x)
+WARNING:absl:Only the first input of the static gather operation is being used, others are being ignored at runtime.
+WARNING:absl:Only the first input of the static gather operation is being used, others are being ignored at runtime.
+WARNING:absl:Only the first input of the sorted segment sum operation is being used, others are being ignored at runtime.
+WARNING:absl:Only the first input of the static gather operation is being used, others are being ignored at runtime.
+WARNING:absl:Only the first input of the static gather operation is being used, others are being ignored at runtime.
+WARNING:absl:Only the first input of the static gather operation is being used, others are being ignored at runtime.
+WARNING:absl:Only the first input of the sorted segment sum operation is being used, others are being ignored at runtime.
+WARNING:root:No normalization scale found for cyclone_exists_gaussian_unit_mode
+WARNING:root:No normalization location found for cyclone_exists_gaussian_unit_mode
+W1031 12:54:04.159116 1588239 hlo_rematerialization.cc:3363] Can't reduce memory use below 21.08GiB (22629749965 bytes) by rematerialization; only reduced to 24.29GiB (26085976296 bytes), down from 24.29GiB (26086000776 bytes) originally
+default reference temp 20.79 GiB step 5.42s compile 61s attention=xla
+WARNING:absl:Skipping gradient checkpointing for sequence length of 1
+WARNING:absl:No coordinates found for dim noise_channels of array Size: 128B
+JitTracer(float32[1,1,32])
+Dimensions without coordinates: batch, channels, noise_channels., using index for param name. This may make model surgery more error prone if changing coordinates at fine-tuning time.
+/home/raymxnd/Workspace/WeatherForcast/.venv/lib/python3.12/site-packages/scipy/sparse/_index.py:216: SparseEfficiencyWarning: Changing the sparsity structure of a csr_matrix is expensive. lil and dok are more efficient.
+ self._set_arrayXarray(i, j, x)
+WARNING:root:No normalization scale found for cyclone_exists_gaussian_unit_mode
+WARNING:root:No normalization location found for cyclone_exists_gaussian_unit_mode
+W1031 12:55:22.552492 1588239 hlo_rematerialization.cc:3363] Can't reduce memory use below 21.08GiB (22629749965 bytes) by rematerialization; only reduced to 28.82GiB (30943716076 bytes), down from 28.82GiB (30943728380 bytes) originally
+default faster_weathernext temp 4.02 GiB step 0.77s compile 67s attention=pallas
+default faster-weathernext vs reference: min corr 0.9999881616 max rel RMS 4.87e-03 median 5.14e-04 NaN masks equal True (101 fields)
+wrote docs/results/finetune/equivalence_rtx5090_remat_option.json
diff --git a/docs/results/finetune/finetune_full_batched_oom_rtx5090.log b/docs/results/finetune/finetune_full_batched_oom_rtx5090.log
new file mode 100644
index 0000000..4e209b4
--- /dev/null
+++ b/docs/results/finetune/finetune_full_batched_oom_rtx5090.log
@@ -0,0 +1,64 @@
+example 2026100412 built in 2.3s
+WeatherNext2 ckpt 1: training 183.79M of 183.8M params; masked targets: 18 (cyclone channels, sea_surface_temperature)
+WARNING:root:No normalization location found for cyclone_exists_gaussian_unit_mode
+WARNING:root:No normalization scale found for cyclone_exists_gaussian_unit_mode
+WARNING:absl:No coordinates found for dim noise_channels of array Size: 256B
+JitTracer(float32[2,1,32])
+Coordinates:
+ * batch (batch) object 16B MultiIndex
+ * _batch (batch) int64 16B 0 0
+ * sample (batch) int64 16B 0 1
+Dimensions without coordinates: channels, noise_channels., using index for param name. This may make model surgery more error prone if changing coordinates at fine-tuning time.
+/home/raymxnd/Workspace/WeatherForcast/.venv/lib/python3.12/site-packages/scipy/sparse/_index.py:216: SparseEfficiencyWarning: Changing the sparsity structure of a csr_matrix is expensive. lil and dok are more efficient.
+ self._set_arrayXarray(i, j, x)
+/home/raymxnd/Workspace/WeatherForcast/weathernext/weathernext/weathernext2/architecture.py:286: FutureWarning: Deleting a single level of a MultiIndex is deprecated. Previously, this deleted all levels of a MultiIndex. Please also drop the following variables: {'batch'} to avoid an error in the future.
+ predictions[var] = xarray_jax.apply_ufunc(activation_fn, predictions[var])
+W1101 00:35:53.629394 37746 hlo_rematerialization.cc:3363] Can't reduce memory use below 25.87GiB (27782055658 bytes) by rematerialization; only reduced to 112.56GiB (120859557177 bytes), down from 118.66GiB (127407407273 bytes) originally
+train step compiled in 125s: temp 45.80 GiB, arguments 3.12 GiB
+WARNING:root:No normalization location found for cyclone_exists_gaussian_unit_mode
+WARNING:root:No normalization scale found for cyclone_exists_gaussian_unit_mode
+WARNING:absl:No coordinates found for dim noise_channels of array Size: 256B
+JitTracer(float32[2,1,32])
+Coordinates:
+ * batch (batch) object 16B MultiIndex
+ * _batch (batch) int64 16B 0 0
+ * sample (batch) int64 16B 0 1
+Dimensions without coordinates: channels, noise_channels., using index for param name. This may make model surgery more error prone if changing coordinates at fine-tuning time.
+/home/raymxnd/Workspace/WeatherForcast/.venv/lib/python3.12/site-packages/scipy/sparse/_index.py:216: SparseEfficiencyWarning: Changing the sparsity structure of a csr_matrix is expensive. lil and dok are more efficient.
+ self._set_arrayXarray(i, j, x)
+/home/raymxnd/Workspace/WeatherForcast/weathernext/weathernext/weathernext2/architecture.py:286: FutureWarning: Deleting a single level of a MultiIndex is deprecated. Previously, this deleted all levels of a MultiIndex. Please also drop the following variables: {'batch'} to avoid an error in the future.
+ predictions[var] = xarray_jax.apply_ufunc(activation_fn, predictions[var])
+W1101 00:37:41.541453 37746 hlo_rematerialization.cc:3363] Can't reduce memory use below 25.12GiB (26978079655 bytes) by rematerialization; only reduced to 56.83GiB (61019999108 bytes), down from 56.83GiB (61020035972 bytes) originally
+eval before (104s): 2026102800 1.04093 2026102812 1.05581 2026103000 1.00724 2026103012 1.02891
+W1101 00:38:10.689671 37746 bfc_allocator.cc:598] Allocator (GPU_0_bfc) ran out of memory trying to allocate 45.80GiB (rounded to 49173576448) from the lower end requested by op
+If this is caused by memory fragmentation, the environment variable 'TF_GPU_ALLOCATOR=cuda_malloc_async' may improve the situation.
+Current allocation summary follows.
+W1101 00:38:10.690844 37746 bfc_allocator.cc:609] ************________________________________________________________________________________________
+Traceback (most recent call last):
+ File "/home/raymxnd/Workspace/faster-weathernext/scripts/finetune_head.py", line 178, in
+ main()
+ File "/home/raymxnd/Workspace/faster-weathernext/scripts/finetune_head.py", line 146, in main
+ metrics = jax.device_get(metrics)
+ ^^^^^^^^^^^^^^^^^^^^^^^
+ File "/home/raymxnd/Workspace/WeatherForcast/.venv/lib/python3.12/site-packages/jax/_src/api.py", line 2592, in device_get
+ return tree_map(_device_get, x)
+ ^^^^^^^^^^^^^^^^^^^^^^^^
+ File "/home/raymxnd/Workspace/WeatherForcast/.venv/lib/python3.12/site-packages/jax/_src/tree_util.py", line 406, in tree_map
+ return treedef.unflatten(f(*xs) for xs in zip(*all_leaves))
+ ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
+ File "/home/raymxnd/Workspace/WeatherForcast/.venv/lib/python3.12/site-packages/jax/_src/tree_util.py", line 406, in
+ return treedef.unflatten(f(*xs) for xs in zip(*all_leaves))
+ ^^^^^^
+ File "/home/raymxnd/Workspace/WeatherForcast/.venv/lib/python3.12/site-packages/jax/_src/api.py", line 2553, in _device_get
+ return toarray()
+ ^^^^^^^^^
+ File "/home/raymxnd/Workspace/WeatherForcast/.venv/lib/python3.12/site-packages/jax/_src/array.py", line 409, in __array__
+ return np.asarray(self._value, dtype=dtype, **kwds) # pyrefly: ignore[no-matching-overload]
+ ^^^^^^^^^^^
+ File "/home/raymxnd/Workspace/WeatherForcast/.venv/lib/python3.12/site-packages/jax/_src/profiler.py", line 420, in wrapper
+ return func(*args, **kwargs)
+ ^^^^^^^^^^^^^^^^^^^^^
+ File "/home/raymxnd/Workspace/WeatherForcast/.venv/lib/python3.12/site-packages/jax/_src/array.py", line 637, in _value
+ npy_value, did_copy = self._single_device_array_to_np_array_did_copy()
+ ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
+jax.errors.JaxRuntimeError: RESOURCE_EXHAUSTED: Out of memory while trying to allocate 45.80GiB with allocator GPU_0_bfc on device 0. [executable_name='jit_step'] [tf-allocator-allocation-error='']
diff --git a/docs/results/finetune/finetune_full_scan_pool90_oom_rtx5090.log b/docs/results/finetune/finetune_full_scan_pool90_oom_rtx5090.log
new file mode 100644
index 0000000..7e55be5
--- /dev/null
+++ b/docs/results/finetune/finetune_full_scan_pool90_oom_rtx5090.log
@@ -0,0 +1,52 @@
+example 2026100412 built in 2.2s
+WeatherNext2 ckpt 1: training 183.79M of 183.8M params; masked targets: 18 (cyclone channels, sea_surface_temperature)
+WARNING:root:No normalization location found for cyclone_exists_gaussian_unit_mode
+WARNING:root:No normalization scale found for cyclone_exists_gaussian_unit_mode
+WARNING:absl:No coordinates found for dim noise_channels of array Size: 128B
+JitTracer(float32[1,1,32])
+Dimensions without coordinates: batch, channels, noise_channels., using index for param name. This may make model surgery more error prone if changing coordinates at fine-tuning time.
+/home/raymxnd/Workspace/WeatherForcast/.venv/lib/python3.12/site-packages/scipy/sparse/_index.py:216: SparseEfficiencyWarning: Changing the sparsity structure of a csr_matrix is expensive. lil and dok are more efficient.
+ self._set_arrayXarray(i, j, x)
+W1101 01:14:04.838198 134248 hlo_rematerialization.cc:3363] Can't reduce memory use below 25.87GiB (27782055658 bytes) by rematerialization; only reduced to 60.88GiB (65367275125 bytes), down from 64.04GiB (68764898983 bytes) originally
+train step compiled in 122s: temp 25.45 GiB, arguments 3.12 GiB
+WARNING:root:No normalization location found for cyclone_exists_gaussian_unit_mode
+WARNING:root:No normalization scale found for cyclone_exists_gaussian_unit_mode
+WARNING:absl:No coordinates found for dim noise_channels of array Size: 128B
+JitTracer(float32[1,1,32])
+Dimensions without coordinates: batch, channels, noise_channels., using index for param name. This may make model surgery more error prone if changing coordinates at fine-tuning time.
+/home/raymxnd/Workspace/WeatherForcast/.venv/lib/python3.12/site-packages/scipy/sparse/_index.py:216: SparseEfficiencyWarning: Changing the sparsity structure of a csr_matrix is expensive. lil and dok are more efficient.
+ self._set_arrayXarray(i, j, x)
+W1101 01:15:23.382879 134248 hlo_rematerialization.cc:3363] Can't reduce memory use below 25.12GiB (26978079655 bytes) by rematerialization; only reduced to 31.11GiB (33406662528 bytes), down from 31.99GiB (34353549912 bytes) originally
+eval before (74s): 2026102800 1.04095 2026102812 1.05584 2026103000 1.00725 2026103012 1.02891
+W1101 01:15:49.189661 134248 bfc_allocator.cc:598] Allocator (GPU_0_bfc) ran out of memory trying to allocate 25.45GiB (rounded to 27324783104) from the lower end requested by op
+If this is caused by memory fragmentation, the environment variable 'TF_GPU_ALLOCATOR=cuda_malloc_async' may improve the situation.
+Current allocation summary follows.
+W1101 01:15:49.190765 134248 bfc_allocator.cc:609] ************________________________________________________________________________________________
+Traceback (most recent call last):
+ File "/home/raymxnd/Workspace/faster-weathernext/scripts/finetune_head.py", line 182, in
+ main()
+ File "/home/raymxnd/Workspace/faster-weathernext/scripts/finetune_head.py", line 150, in main
+ metrics = jax.device_get(metrics)
+ ^^^^^^^^^^^^^^^^^^^^^^^
+ File "/home/raymxnd/Workspace/WeatherForcast/.venv/lib/python3.12/site-packages/jax/_src/api.py", line 2592, in device_get
+ return tree_map(_device_get, x)
+ ^^^^^^^^^^^^^^^^^^^^^^^^
+ File "/home/raymxnd/Workspace/WeatherForcast/.venv/lib/python3.12/site-packages/jax/_src/tree_util.py", line 406, in tree_map
+ return treedef.unflatten(f(*xs) for xs in zip(*all_leaves))
+ ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
+ File "/home/raymxnd/Workspace/WeatherForcast/.venv/lib/python3.12/site-packages/jax/_src/tree_util.py", line 406, in
+ return treedef.unflatten(f(*xs) for xs in zip(*all_leaves))
+ ^^^^^^
+ File "/home/raymxnd/Workspace/WeatherForcast/.venv/lib/python3.12/site-packages/jax/_src/api.py", line 2553, in _device_get
+ return toarray()
+ ^^^^^^^^^
+ File "/home/raymxnd/Workspace/WeatherForcast/.venv/lib/python3.12/site-packages/jax/_src/array.py", line 409, in __array__
+ return np.asarray(self._value, dtype=dtype, **kwds) # pyrefly: ignore[no-matching-overload]
+ ^^^^^^^^^^^
+ File "/home/raymxnd/Workspace/WeatherForcast/.venv/lib/python3.12/site-packages/jax/_src/profiler.py", line 420, in wrapper
+ return func(*args, **kwargs)
+ ^^^^^^^^^^^^^^^^^^^^^
+ File "/home/raymxnd/Workspace/WeatherForcast/.venv/lib/python3.12/site-packages/jax/_src/array.py", line 637, in _value
+ npy_value, did_copy = self._single_device_array_to_np_array_did_copy()
+ ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
+jax.errors.JaxRuntimeError: RESOURCE_EXHAUSTED: Out of memory while trying to allocate 25.45GiB with allocator GPU_0_bfc on device 0. [executable_name='jit_step'] [tf-allocator-allocation-error='']
diff --git a/docs/results/finetune/finetune_full_smoke_rtx5090.json b/docs/results/finetune/finetune_full_smoke_rtx5090.json
new file mode 100644
index 0000000..9582c7b
--- /dev/null
+++ b/docs/results/finetune/finetune_full_smoke_rtx5090.json
@@ -0,0 +1,573 @@
+{
+ "model": "WeatherNext2",
+ "checkpoint": 1,
+ "device": "NVIDIA GeForce RTX 5090",
+ "jax": "0.11.2",
+ "faster_weathernext": "0.1.1",
+ "args": {
+ "model": "WeatherNext2",
+ "checkpoint": 1,
+ "train": [
+ "2026100412",
+ "2026100418",
+ "2026100500",
+ "2026100506",
+ "2026100512",
+ "2026100518",
+ "2026100600",
+ "2026100606",
+ "2026100612",
+ "2026100618",
+ "2026100700",
+ "2026100706",
+ "2026100712",
+ "2026100718",
+ "2026100800",
+ "2026100806",
+ "2026100912",
+ "2026100918",
+ "2026101000",
+ "2026101006",
+ "2026101112",
+ "2026101118",
+ "2026101200",
+ "2026101206",
+ "2026101212",
+ "2026101218",
+ "2026101300",
+ "2026101306",
+ "2026101412",
+ "2026101418",
+ "2026101500",
+ "2026101506",
+ "2026101512",
+ "2026101518",
+ "2026101600",
+ "2026101606",
+ "2026101612",
+ "2026101618",
+ "2026101700",
+ "2026101706",
+ "2026101712",
+ "2026101718",
+ "2026101800",
+ "2026101806",
+ "2026101812",
+ "2026101818",
+ "2026101900",
+ "2026101906",
+ "2026101912",
+ "2026101918",
+ "2026102000",
+ "2026102006",
+ "2026102012",
+ "2026102018",
+ "2026102100",
+ "2026102106",
+ "2026102112",
+ "2026102118",
+ "2026102200",
+ "2026102206",
+ "2026102212",
+ "2026102218",
+ "2026102300",
+ "2026102306",
+ "2026102312",
+ "2026102318",
+ "2026102400",
+ "2026102406"
+ ],
+ "eval": [
+ "2026102800",
+ "2026102812",
+ "2026103000",
+ "2026103012"
+ ],
+ "steps": 20,
+ "lr": 1e-06,
+ "warmup": 5,
+ "weight_decay": 0.0,
+ "clip": 1.0,
+ "eval_keys": 2,
+ "precision": "default",
+ "seed": 0,
+ "no_precip": false,
+ "modules": [
+ "all"
+ ],
+ "sequential_samples": "auto",
+ "mask": [
+ "sea_surface_temperature"
+ ],
+ "compile_only": false,
+ "no_remat": false,
+ "weights_dir": "/home/raymxnd/Workspace/WeatherForcast/weights",
+ "out": null,
+ "log": "docs/results/finetune/finetune_full_smoke_rtx5090.json"
+ },
+ "masked": [
+ "cyclone_all_wind_disc",
+ "cyclone_exists_gaussian_unit_mode",
+ "cyclone_usa_pres_disc",
+ "cyclone_usa_r34_ne_radius_disc",
+ "cyclone_usa_r34_nw_radius_disc",
+ "cyclone_usa_r34_se_radius_disc",
+ "cyclone_usa_r34_sw_radius_disc",
+ "cyclone_usa_r50_ne_radius_disc",
+ "cyclone_usa_r50_nw_radius_disc",
+ "cyclone_usa_r50_se_radius_disc",
+ "cyclone_usa_r50_sw_radius_disc",
+ "cyclone_usa_r64_ne_radius_disc",
+ "cyclone_usa_r64_nw_radius_disc",
+ "cyclone_usa_r64_se_radius_disc",
+ "cyclone_usa_r64_sw_radius_disc",
+ "cyclone_usa_rmw_disc",
+ "cyclone_usa_wind_disc",
+ "sea_surface_temperature"
+ ],
+ "modules": "all",
+ "head_params": 183785575,
+ "total_params": 183785575,
+ "compile": {
+ "seconds": 121.53278493881226,
+ "temp_gib": 25.46281883865595,
+ "argument_gib": 3.117626067250967,
+ "output_gib": 2.0539751574397087
+ },
+ "eval_before": {
+ "2026102800": {
+ "loss": 1.0409454107284546,
+ "per_var": {
+ "100m_u_component_of_wind": 0.13943615555763245,
+ "100m_v_component_of_wind": 0.12799286842346191,
+ "10m_u_component_of_wind": 0.14852318167686462,
+ "10m_v_component_of_wind": 0.13639746606349945,
+ "2m_temperature": 0.09907721728086472,
+ "cyclone_all_wind_disc": 0.0,
+ "cyclone_exists_gaussian_unit_mode": 0.0,
+ "cyclone_usa_pres_disc": 0.0,
+ "cyclone_usa_r34_ne_radius_disc": 0.0,
+ "cyclone_usa_r34_nw_radius_disc": 0.0,
+ "cyclone_usa_r34_se_radius_disc": 0.0,
+ "cyclone_usa_r34_sw_radius_disc": 0.0,
+ "cyclone_usa_r50_ne_radius_disc": 0.0,
+ "cyclone_usa_r50_nw_radius_disc": 0.0,
+ "cyclone_usa_r50_se_radius_disc": 0.0,
+ "cyclone_usa_r50_sw_radius_disc": 0.0,
+ "cyclone_usa_r64_ne_radius_disc": 0.0,
+ "cyclone_usa_r64_nw_radius_disc": 0.0,
+ "cyclone_usa_r64_se_radius_disc": 0.0,
+ "cyclone_usa_r64_sw_radius_disc": 0.0,
+ "cyclone_usa_rmw_disc": 0.0,
+ "cyclone_usa_wind_disc": 0.0,
+ "geopotential": 0.06217734143137932,
+ "mean_sea_level_pressure": 0.06354136019945145,
+ "sea_surface_temperature": 1.588570475578308,
+ "specific_humidity": 0.21853762865066528,
+ "temperature": 0.12121966481208801,
+ "total_precipitation_6hr": 0.127057746052742,
+ "u_component_of_wind": 0.13052859902381897,
+ "v_component_of_wind": 0.11552776396274567,
+ "vertical_velocity": 0.27828842401504517
+ }
+ },
+ "2026102812": {
+ "loss": 1.055836796760559,
+ "per_var": {
+ "100m_u_component_of_wind": 0.14223718643188477,
+ "100m_v_component_of_wind": 0.13159044086933136,
+ "10m_u_component_of_wind": 0.14976435899734497,
+ "10m_v_component_of_wind": 0.1381080597639084,
+ "2m_temperature": 0.10236847400665283,
+ "cyclone_all_wind_disc": 0.0,
+ "cyclone_exists_gaussian_unit_mode": 0.0,
+ "cyclone_usa_pres_disc": 0.0,
+ "cyclone_usa_r34_ne_radius_disc": 0.0,
+ "cyclone_usa_r34_nw_radius_disc": 0.0,
+ "cyclone_usa_r34_se_radius_disc": 0.0,
+ "cyclone_usa_r34_sw_radius_disc": 0.0,
+ "cyclone_usa_r50_ne_radius_disc": 0.0,
+ "cyclone_usa_r50_nw_radius_disc": 0.0,
+ "cyclone_usa_r50_se_radius_disc": 0.0,
+ "cyclone_usa_r50_sw_radius_disc": 0.0,
+ "cyclone_usa_r64_ne_radius_disc": 0.0,
+ "cyclone_usa_r64_nw_radius_disc": 0.0,
+ "cyclone_usa_r64_se_radius_disc": 0.0,
+ "cyclone_usa_r64_sw_radius_disc": 0.0,
+ "cyclone_usa_rmw_disc": 0.0,
+ "cyclone_usa_wind_disc": 0.0,
+ "geopotential": 0.05947577953338623,
+ "mean_sea_level_pressure": 0.06730344891548157,
+ "sea_surface_temperature": 1.6101081371307373,
+ "specific_humidity": 0.22229136526584625,
+ "temperature": 0.12098716199398041,
+ "total_precipitation_6hr": 0.13529032468795776,
+ "u_component_of_wind": 0.13082465529441833,
+ "v_component_of_wind": 0.11493195593357086,
+ "vertical_velocity": 0.28635090589523315
+ }
+ },
+ "2026103000": {
+ "loss": 1.0072450637817383,
+ "per_var": {
+ "100m_u_component_of_wind": 0.13534832000732422,
+ "100m_v_component_of_wind": 0.12464036047458649,
+ "10m_u_component_of_wind": 0.14479902386665344,
+ "10m_v_component_of_wind": 0.1337926834821701,
+ "2m_temperature": 0.09720027446746826,
+ "cyclone_all_wind_disc": 0.0,
+ "cyclone_exists_gaussian_unit_mode": 0.0,
+ "cyclone_usa_pres_disc": 0.0,
+ "cyclone_usa_r34_ne_radius_disc": 0.0,
+ "cyclone_usa_r34_nw_radius_disc": 0.0,
+ "cyclone_usa_r34_se_radius_disc": 0.0,
+ "cyclone_usa_r34_sw_radius_disc": 0.0,
+ "cyclone_usa_r50_ne_radius_disc": 0.0,
+ "cyclone_usa_r50_nw_radius_disc": 0.0,
+ "cyclone_usa_r50_se_radius_disc": 0.0,
+ "cyclone_usa_r50_sw_radius_disc": 0.0,
+ "cyclone_usa_r64_ne_radius_disc": 0.0,
+ "cyclone_usa_r64_nw_radius_disc": 0.0,
+ "cyclone_usa_r64_se_radius_disc": 0.0,
+ "cyclone_usa_r64_sw_radius_disc": 0.0,
+ "cyclone_usa_rmw_disc": 0.0,
+ "cyclone_usa_wind_disc": 0.0,
+ "geopotential": 0.06039592996239662,
+ "mean_sea_level_pressure": 0.06657945364713669,
+ "sea_surface_temperature": 1.5575478076934814,
+ "specific_humidity": 0.21349915862083435,
+ "temperature": 0.11542226374149323,
+ "total_precipitation_6hr": 0.12492947280406952,
+ "u_component_of_wind": 0.12509451806545258,
+ "v_component_of_wind": 0.11082319915294647,
+ "vertical_velocity": 0.26892778277397156
+ }
+ },
+ "2026103012": {
+ "loss": 1.028918743133545,
+ "per_var": {
+ "100m_u_component_of_wind": 0.13895925879478455,
+ "100m_v_component_of_wind": 0.1291230469942093,
+ "10m_u_component_of_wind": 0.14677700400352478,
+ "10m_v_component_of_wind": 0.1362985223531723,
+ "2m_temperature": 0.09880942106246948,
+ "cyclone_all_wind_disc": 0.0,
+ "cyclone_exists_gaussian_unit_mode": 0.0,
+ "cyclone_usa_pres_disc": 0.0,
+ "cyclone_usa_r34_ne_radius_disc": 0.0,
+ "cyclone_usa_r34_nw_radius_disc": 0.0,
+ "cyclone_usa_r34_se_radius_disc": 0.0,
+ "cyclone_usa_r34_sw_radius_disc": 0.0,
+ "cyclone_usa_r50_ne_radius_disc": 0.0,
+ "cyclone_usa_r50_nw_radius_disc": 0.0,
+ "cyclone_usa_r50_se_radius_disc": 0.0,
+ "cyclone_usa_r50_sw_radius_disc": 0.0,
+ "cyclone_usa_r64_ne_radius_disc": 0.0,
+ "cyclone_usa_r64_nw_radius_disc": 0.0,
+ "cyclone_usa_r64_se_radius_disc": 0.0,
+ "cyclone_usa_r64_sw_radius_disc": 0.0,
+ "cyclone_usa_rmw_disc": 0.0,
+ "cyclone_usa_wind_disc": 0.0,
+ "geopotential": 0.05989161133766174,
+ "mean_sea_level_pressure": 0.06429021060466766,
+ "sea_surface_temperature": 1.6431214809417725,
+ "specific_humidity": 0.2159145027399063,
+ "temperature": 0.11675936728715897,
+ "total_precipitation_6hr": 0.14073090255260468,
+ "u_component_of_wind": 0.12529711425304413,
+ "v_component_of_wind": 0.11095214635133743,
+ "vertical_velocity": 0.2831803858280182
+ }
+ }
+ },
+ "steps": [
+ {
+ "step": 0,
+ "init": "2026101006",
+ "loss": 1.119167685508728,
+ "grad_norm": 16.9317684173584,
+ "seconds": 8.796973943710327
+ },
+ {
+ "step": 1,
+ "init": "2026102406",
+ "loss": 1.0820095539093018,
+ "grad_norm": 18.239730834960938,
+ "seconds": 8.495733976364136
+ },
+ {
+ "step": 2,
+ "init": "2026101206",
+ "loss": 1.1618558168411255,
+ "grad_norm": 18.61575698852539,
+ "seconds": 8.504448413848877
+ },
+ {
+ "step": 3,
+ "init": "2026100612",
+ "loss": 0.9954826831817627,
+ "grad_norm": 13.9835844039917,
+ "seconds": 8.507168292999268
+ },
+ {
+ "step": 4,
+ "init": "2026101900",
+ "loss": 1.0272438526153564,
+ "grad_norm": 15.795987129211426,
+ "seconds": 8.509703636169434
+ },
+ {
+ "step": 5,
+ "init": "2026101906",
+ "loss": 1.074432134628296,
+ "grad_norm": 16.4752197265625,
+ "seconds": 8.511326313018799
+ },
+ {
+ "step": 6,
+ "init": "2026102300",
+ "loss": 1.1275370121002197,
+ "grad_norm": 16.28327178955078,
+ "seconds": 8.511404037475586
+ },
+ {
+ "step": 7,
+ "init": "2026100512",
+ "loss": 1.0502113103866577,
+ "grad_norm": 14.333346366882324,
+ "seconds": 8.510565280914307
+ },
+ {
+ "step": 8,
+ "init": "2026101806",
+ "loss": 1.091897964477539,
+ "grad_norm": 12.178189277648926,
+ "seconds": 8.510094404220581
+ },
+ {
+ "step": 9,
+ "init": "2026102200",
+ "loss": 0.9743238687515259,
+ "grad_norm": 13.208512306213379,
+ "seconds": 8.511714220046997
+ },
+ {
+ "step": 10,
+ "init": "2026101500",
+ "loss": 1.0952866077423096,
+ "grad_norm": 16.87466049194336,
+ "seconds": 8.51449728012085
+ },
+ {
+ "step": 11,
+ "init": "2026100700",
+ "loss": 1.0515720844268799,
+ "grad_norm": 14.492918014526367,
+ "seconds": 8.511247873306274
+ },
+ {
+ "step": 12,
+ "init": "2026100912",
+ "loss": 1.0031788349151611,
+ "grad_norm": 13.45763874053955,
+ "seconds": 8.514005661010742
+ },
+ {
+ "step": 13,
+ "init": "2026100500",
+ "loss": 0.9974920153617859,
+ "grad_norm": 13.286833763122559,
+ "seconds": 8.515864849090576
+ },
+ {
+ "step": 14,
+ "init": "2026101812",
+ "loss": 1.060226321220398,
+ "grad_norm": 18.3127498626709,
+ "seconds": 8.512723684310913
+ },
+ {
+ "step": 15,
+ "init": "2026101612",
+ "loss": 0.9534765481948853,
+ "grad_norm": 11.353670120239258,
+ "seconds": 8.513659477233887
+ },
+ {
+ "step": 16,
+ "init": "2026101112",
+ "loss": 0.9863734245300293,
+ "grad_norm": 15.268701553344727,
+ "seconds": 8.514719009399414
+ },
+ {
+ "step": 17,
+ "init": "2026102006",
+ "loss": 1.0501787662506104,
+ "grad_norm": 15.60908317565918,
+ "seconds": 8.516311407089233
+ },
+ {
+ "step": 18,
+ "init": "2026100706",
+ "loss": 1.1040557622909546,
+ "grad_norm": 15.235260963439941,
+ "seconds": 8.51640248298645
+ },
+ {
+ "step": 19,
+ "init": "2026102312",
+ "loss": 0.9601024389266968,
+ "grad_norm": 13.556150436401367,
+ "seconds": 8.510567665100098
+ }
+ ],
+ "peak_gib": 28.580492973327637,
+ "median_step_s": 8.511404037475586,
+ "eval_after": {
+ "2026102800": {
+ "loss": 1.0185554027557373,
+ "per_var": {
+ "100m_u_component_of_wind": 0.13732294738292694,
+ "100m_v_component_of_wind": 0.1271643191576004,
+ "10m_u_component_of_wind": 0.14614742994308472,
+ "10m_v_component_of_wind": 0.13554558157920837,
+ "2m_temperature": 0.08816573768854141,
+ "cyclone_all_wind_disc": 0.0,
+ "cyclone_exists_gaussian_unit_mode": 0.0,
+ "cyclone_usa_pres_disc": 0.0,
+ "cyclone_usa_r34_ne_radius_disc": 0.0,
+ "cyclone_usa_r34_nw_radius_disc": 0.0,
+ "cyclone_usa_r34_se_radius_disc": 0.0,
+ "cyclone_usa_r34_sw_radius_disc": 0.0,
+ "cyclone_usa_r50_ne_radius_disc": 0.0,
+ "cyclone_usa_r50_nw_radius_disc": 0.0,
+ "cyclone_usa_r50_se_radius_disc": 0.0,
+ "cyclone_usa_r50_sw_radius_disc": 0.0,
+ "cyclone_usa_r64_ne_radius_disc": 0.0,
+ "cyclone_usa_r64_nw_radius_disc": 0.0,
+ "cyclone_usa_r64_se_radius_disc": 0.0,
+ "cyclone_usa_r64_sw_radius_disc": 0.0,
+ "cyclone_usa_rmw_disc": 0.0,
+ "cyclone_usa_wind_disc": 0.0,
+ "geopotential": 0.05828139930963516,
+ "mean_sea_level_pressure": 0.062490686774253845,
+ "sea_surface_temperature": 1.5847361087799072,
+ "specific_humidity": 0.21661026775836945,
+ "temperature": 0.11655979603528976,
+ "total_precipitation_6hr": 0.12064985930919647,
+ "u_component_of_wind": 0.1295609176158905,
+ "v_component_of_wind": 0.11548873782157898,
+ "vertical_velocity": 0.277406245470047
+ }
+ },
+ "2026102812": {
+ "loss": 1.0334675312042236,
+ "per_var": {
+ "100m_u_component_of_wind": 0.13979710638523102,
+ "100m_v_component_of_wind": 0.1298966407775879,
+ "10m_u_component_of_wind": 0.14719809591770172,
+ "10m_v_component_of_wind": 0.13603812456130981,
+ "2m_temperature": 0.09320610761642456,
+ "cyclone_all_wind_disc": 0.0,
+ "cyclone_exists_gaussian_unit_mode": 0.0,
+ "cyclone_usa_pres_disc": 0.0,
+ "cyclone_usa_r34_ne_radius_disc": 0.0,
+ "cyclone_usa_r34_nw_radius_disc": 0.0,
+ "cyclone_usa_r34_se_radius_disc": 0.0,
+ "cyclone_usa_r34_sw_radius_disc": 0.0,
+ "cyclone_usa_r50_ne_radius_disc": 0.0,
+ "cyclone_usa_r50_nw_radius_disc": 0.0,
+ "cyclone_usa_r50_se_radius_disc": 0.0,
+ "cyclone_usa_r50_sw_radius_disc": 0.0,
+ "cyclone_usa_r64_ne_radius_disc": 0.0,
+ "cyclone_usa_r64_nw_radius_disc": 0.0,
+ "cyclone_usa_r64_se_radius_disc": 0.0,
+ "cyclone_usa_r64_sw_radius_disc": 0.0,
+ "cyclone_usa_rmw_disc": 0.0,
+ "cyclone_usa_wind_disc": 0.0,
+ "geopotential": 0.056234586983919144,
+ "mean_sea_level_pressure": 0.06361456215381622,
+ "sea_surface_temperature": 1.6099262237548828,
+ "specific_humidity": 0.21956267952919006,
+ "temperature": 0.11630144715309143,
+ "total_precipitation_6hr": 0.12750405073165894,
+ "u_component_of_wind": 0.1300998330116272,
+ "v_component_of_wind": 0.11420255154371262,
+ "vertical_velocity": 0.28521928191185
+ }
+ },
+ "2026103000": {
+ "loss": 0.9829198122024536,
+ "per_var": {
+ "100m_u_component_of_wind": 0.13370296359062195,
+ "100m_v_component_of_wind": 0.12361572682857513,
+ "10m_u_component_of_wind": 0.14286479353904724,
+ "10m_v_component_of_wind": 0.13239677250385284,
+ "2m_temperature": 0.08690063655376434,
+ "cyclone_all_wind_disc": 0.0,
+ "cyclone_exists_gaussian_unit_mode": 0.0,
+ "cyclone_usa_pres_disc": 0.0,
+ "cyclone_usa_r34_ne_radius_disc": 0.0,
+ "cyclone_usa_r34_nw_radius_disc": 0.0,
+ "cyclone_usa_r34_se_radius_disc": 0.0,
+ "cyclone_usa_r34_sw_radius_disc": 0.0,
+ "cyclone_usa_r50_ne_radius_disc": 0.0,
+ "cyclone_usa_r50_nw_radius_disc": 0.0,
+ "cyclone_usa_r50_se_radius_disc": 0.0,
+ "cyclone_usa_r50_sw_radius_disc": 0.0,
+ "cyclone_usa_r64_ne_radius_disc": 0.0,
+ "cyclone_usa_r64_nw_radius_disc": 0.0,
+ "cyclone_usa_r64_se_radius_disc": 0.0,
+ "cyclone_usa_r64_sw_radius_disc": 0.0,
+ "cyclone_usa_rmw_disc": 0.0,
+ "cyclone_usa_wind_disc": 0.0,
+ "geopotential": 0.05628117173910141,
+ "mean_sea_level_pressure": 0.06131090223789215,
+ "sea_surface_temperature": 1.5555641651153564,
+ "specific_humidity": 0.21118010580539703,
+ "temperature": 0.10979911684989929,
+ "total_precipitation_6hr": 0.11862312257289886,
+ "u_component_of_wind": 0.12408313155174255,
+ "v_component_of_wind": 0.1102541834115982,
+ "vertical_velocity": 0.26793956756591797
+ }
+ },
+ "2026103012": {
+ "loss": 1.0100600719451904,
+ "per_var": {
+ "100m_u_component_of_wind": 0.1371956467628479,
+ "100m_v_component_of_wind": 0.12772268056869507,
+ "10m_u_component_of_wind": 0.14464063942432404,
+ "10m_v_component_of_wind": 0.13459379971027374,
+ "2m_temperature": 0.09210316836833954,
+ "cyclone_all_wind_disc": 0.0,
+ "cyclone_exists_gaussian_unit_mode": 0.0,
+ "cyclone_usa_pres_disc": 0.0,
+ "cyclone_usa_r34_ne_radius_disc": 0.0,
+ "cyclone_usa_r34_nw_radius_disc": 0.0,
+ "cyclone_usa_r34_se_radius_disc": 0.0,
+ "cyclone_usa_r34_sw_radius_disc": 0.0,
+ "cyclone_usa_r50_ne_radius_disc": 0.0,
+ "cyclone_usa_r50_nw_radius_disc": 0.0,
+ "cyclone_usa_r50_se_radius_disc": 0.0,
+ "cyclone_usa_r50_sw_radius_disc": 0.0,
+ "cyclone_usa_r64_ne_radius_disc": 0.0,
+ "cyclone_usa_r64_nw_radius_disc": 0.0,
+ "cyclone_usa_r64_se_radius_disc": 0.0,
+ "cyclone_usa_r64_sw_radius_disc": 0.0,
+ "cyclone_usa_rmw_disc": 0.0,
+ "cyclone_usa_wind_disc": 0.0,
+ "geopotential": 0.05651574581861496,
+ "mean_sea_level_pressure": 0.06226147711277008,
+ "sea_surface_temperature": 1.643027424812317,
+ "specific_humidity": 0.21331024169921875,
+ "temperature": 0.11321263015270233,
+ "total_precipitation_6hr": 0.13325899839401245,
+ "u_component_of_wind": 0.12395166605710983,
+ "v_component_of_wind": 0.11031177639961243,
+ "vertical_velocity": 0.2821529805660248
+ }
+ }
+ }
+}
\ No newline at end of file
diff --git a/docs/results/finetune/finetune_full_smoke_rtx5090.log b/docs/results/finetune/finetune_full_smoke_rtx5090.log
new file mode 100644
index 0000000..724576e
--- /dev/null
+++ b/docs/results/finetune/finetune_full_smoke_rtx5090.log
@@ -0,0 +1,46 @@
+example 2026100412 built in 2.2s
+WeatherNext2 ckpt 1: training 183.79M of 183.8M params; masked targets: 18 (cyclone channels, sea_surface_temperature)
+WARNING:root:No normalization location found for cyclone_exists_gaussian_unit_mode
+WARNING:root:No normalization scale found for cyclone_exists_gaussian_unit_mode
+WARNING:absl:No coordinates found for dim noise_channels of array Size: 128B
+JitTracer(float32[1,1,32])
+Dimensions without coordinates: batch, channels, noise_channels., using index for param name. This may make model surgery more error prone if changing coordinates at fine-tuning time.
+/home/raymxnd/Workspace/WeatherForcast/.venv/lib/python3.12/site-packages/scipy/sparse/_index.py:216: SparseEfficiencyWarning: Changing the sparsity structure of a csr_matrix is expensive. lil and dok are more efficient.
+ self._set_arrayXarray(i, j, x)
+W1101 01:23:18.923607 149577 hlo_rematerialization.cc:3363] Can't reduce memory use below 27.36GiB (29379875767 bytes) by rematerialization; only reduced to 60.88GiB (65367275125 bytes), down from 64.04GiB (68764898983 bytes) originally
+train step compiled in 122s: temp 25.46 GiB, arguments 3.12 GiB
+WARNING:root:No normalization location found for cyclone_exists_gaussian_unit_mode
+WARNING:root:No normalization scale found for cyclone_exists_gaussian_unit_mode
+WARNING:absl:No coordinates found for dim noise_channels of array Size: 128B
+JitTracer(float32[1,1,32])
+Dimensions without coordinates: batch, channels, noise_channels., using index for param name. This may make model surgery more error prone if changing coordinates at fine-tuning time.
+/home/raymxnd/Workspace/WeatherForcast/.venv/lib/python3.12/site-packages/scipy/sparse/_index.py:216: SparseEfficiencyWarning: Changing the sparsity structure of a csr_matrix is expensive. lil and dok are more efficient.
+ self._set_arrayXarray(i, j, x)
+W1101 01:24:37.340890 149577 hlo_rematerialization.cc:3363] Can't reduce memory use below 26.61GiB (28575899763 bytes) by rematerialization; only reduced to 31.11GiB (33406662528 bytes), down from 31.99GiB (34353549912 bytes) originally
+eval before (73s): 2026102800 1.04095 2026102812 1.05584 2026103000 1.00725 2026103012 1.02892
+step 0 2026101006 loss 1.11917 |g| 1.693e+01 8.80s
+step 1 2026102406 loss 1.08201 |g| 1.824e+01 8.50s
+step 2 2026101206 loss 1.16186 |g| 1.862e+01 8.50s
+step 3 2026100612 loss 0.99548 |g| 1.398e+01 8.51s
+step 4 2026101900 loss 1.02724 |g| 1.580e+01 8.51s
+step 5 2026101906 loss 1.07443 |g| 1.648e+01 8.51s
+step 6 2026102300 loss 1.12754 |g| 1.628e+01 8.51s
+step 7 2026100512 loss 1.05021 |g| 1.433e+01 8.51s
+step 8 2026101806 loss 1.09190 |g| 1.218e+01 8.51s
+step 9 2026102200 loss 0.97432 |g| 1.321e+01 8.51s
+step 10 2026101500 loss 1.09529 |g| 1.687e+01 8.51s
+step 11 2026100700 loss 1.05157 |g| 1.449e+01 8.51s
+step 12 2026100912 loss 1.00318 |g| 1.346e+01 8.51s
+step 13 2026100500 loss 0.99749 |g| 1.329e+01 8.52s
+step 14 2026101812 loss 1.06023 |g| 1.831e+01 8.51s
+step 15 2026101612 loss 0.95348 |g| 1.135e+01 8.51s
+step 16 2026101112 loss 0.98637 |g| 1.527e+01 8.51s
+step 17 2026102006 loss 1.05018 |g| 1.561e+01 8.52s
+step 18 2026100706 loss 1.10406 |g| 1.524e+01 8.52s
+step 19 2026102312 loss 0.96010 |g| 1.356e+01 8.51s
+median step 8.51s, device peak 28.58 GiB
+eval 2026102800: 1.04095 -> 1.01856 (-2.15 %)
+eval 2026102812: 1.05584 -> 1.03347 (-2.12 %)
+eval 2026103000: 1.00725 -> 0.98292 (-2.42 %)
+eval 2026103012: 1.02892 -> 1.01006 (-1.83 %)
+wrote docs/results/finetune/finetune_full_smoke_rtx5090.json
diff --git a/docs/results/finetune/finetune_full_unrolled_oom_rtx5090.log b/docs/results/finetune/finetune_full_unrolled_oom_rtx5090.log
new file mode 100644
index 0000000..bc513ff
--- /dev/null
+++ b/docs/results/finetune/finetune_full_unrolled_oom_rtx5090.log
@@ -0,0 +1,78 @@
+example 2026100412 built in 2.2s
+WeatherNext2 ckpt 1: training 183.79M of 183.8M params; masked targets: 18 (cyclone channels, sea_surface_temperature)
+WARNING:root:No normalization location found for cyclone_exists_gaussian_unit_mode
+WARNING:root:No normalization scale found for cyclone_exists_gaussian_unit_mode
+WARNING:absl:No coordinates found for dim noise_channels of array Size: 128B
+JitTracer(float32[1,1,32])
+Coordinates:
+ * batch (batch) object 8B MultiIndex
+ * _batch (batch) int64 8B 0
+ * sample (batch) int64 8B 0
+Dimensions without coordinates: channels, noise_channels., using index for param name. This may make model surgery more error prone if changing coordinates at fine-tuning time.
+/home/raymxnd/Workspace/WeatherForcast/.venv/lib/python3.12/site-packages/scipy/sparse/_index.py:216: SparseEfficiencyWarning: Changing the sparsity structure of a csr_matrix is expensive. lil and dok are more efficient.
+ self._set_arrayXarray(i, j, x)
+/home/raymxnd/Workspace/WeatherForcast/weathernext/weathernext/weathernext2/architecture.py:286: FutureWarning: Deleting a single level of a MultiIndex is deprecated. Previously, this deleted all levels of a MultiIndex. Please also drop the following variables: {'batch'} to avoid an error in the future.
+ predictions[var] = xarray_jax.apply_ufunc(activation_fn, predictions[var])
+WARNING:absl:No coordinates found for dim noise_channels of array Size: 128B
+JitTracer(float32[1,1,32])
+Coordinates:
+ * batch (batch) object 8B MultiIndex
+ * _batch (batch) int64 8B 0
+ * sample (batch) int64 8B 1
+Dimensions without coordinates: channels, noise_channels., using index for param name. This may make model surgery more error prone if changing coordinates at fine-tuning time.
+W1101 01:04:26.853586 122882 hlo_rematerialization.cc:3363] Can't reduce memory use below 25.87GiB (27782055658 bytes) by rematerialization; only reduced to 77.72GiB (83449192183 bytes), down from 81.23GiB (87223108895 bytes) originally
+train step compiled in 199s: temp 46.10 GiB, arguments 3.12 GiB
+WARNING:root:No normalization location found for cyclone_exists_gaussian_unit_mode
+WARNING:root:No normalization scale found for cyclone_exists_gaussian_unit_mode
+WARNING:absl:No coordinates found for dim noise_channels of array Size: 128B
+JitTracer(float32[1,1,32])
+Coordinates:
+ * batch (batch) object 8B MultiIndex
+ * _batch (batch) int64 8B 0
+ * sample (batch) int64 8B 0
+Dimensions without coordinates: channels, noise_channels., using index for param name. This may make model surgery more error prone if changing coordinates at fine-tuning time.
+/home/raymxnd/Workspace/WeatherForcast/.venv/lib/python3.12/site-packages/scipy/sparse/_index.py:216: SparseEfficiencyWarning: Changing the sparsity structure of a csr_matrix is expensive. lil and dok are more efficient.
+ self._set_arrayXarray(i, j, x)
+/home/raymxnd/Workspace/WeatherForcast/weathernext/weathernext/weathernext2/architecture.py:286: FutureWarning: Deleting a single level of a MultiIndex is deprecated. Previously, this deleted all levels of a MultiIndex. Please also drop the following variables: {'batch'} to avoid an error in the future.
+ predictions[var] = xarray_jax.apply_ufunc(activation_fn, predictions[var])
+WARNING:absl:No coordinates found for dim noise_channels of array Size: 128B
+JitTracer(float32[1,1,32])
+Coordinates:
+ * batch (batch) object 8B MultiIndex
+ * _batch (batch) int64 8B 0
+ * sample (batch) int64 8B 1
+Dimensions without coordinates: channels, noise_channels., using index for param name. This may make model surgery more error prone if changing coordinates at fine-tuning time.
+W1101 01:07:06.543671 122882 hlo_rematerialization.cc:3363] Can't reduce memory use below 25.12GiB (26978079655 bytes) by rematerialization; only reduced to 33.12GiB (35568612140 bytes), down from 33.12GiB (35568636748 bytes) originally
+eval before (155s): 2026102800 1.04096 2026102812 1.05583 2026103000 1.00725 2026103012 1.02892
+W1101 01:07:37.227851 122882 bfc_allocator.cc:598] Allocator (GPU_0_bfc) ran out of memory trying to allocate 46.10GiB (rounded to 49502263552) from the lower end requested by op
+If this is caused by memory fragmentation, the environment variable 'TF_GPU_ALLOCATOR=cuda_malloc_async' may improve the situation.
+Current allocation summary follows.
+W1101 01:07:37.229280 122882 bfc_allocator.cc:609] ************________________________________________________________________________________________
+Traceback (most recent call last):
+ File "/home/raymxnd/Workspace/faster-weathernext/scripts/finetune_head.py", line 182, in
+ main()
+ File "/home/raymxnd/Workspace/faster-weathernext/scripts/finetune_head.py", line 150, in main
+ metrics = jax.device_get(metrics)
+ ^^^^^^^^^^^^^^^^^^^^^^^
+ File "/home/raymxnd/Workspace/WeatherForcast/.venv/lib/python3.12/site-packages/jax/_src/api.py", line 2592, in device_get
+ return tree_map(_device_get, x)
+ ^^^^^^^^^^^^^^^^^^^^^^^^
+ File "/home/raymxnd/Workspace/WeatherForcast/.venv/lib/python3.12/site-packages/jax/_src/tree_util.py", line 406, in tree_map
+ return treedef.unflatten(f(*xs) for xs in zip(*all_leaves))
+ ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
+ File "/home/raymxnd/Workspace/WeatherForcast/.venv/lib/python3.12/site-packages/jax/_src/tree_util.py", line 406, in
+ return treedef.unflatten(f(*xs) for xs in zip(*all_leaves))
+ ^^^^^^
+ File "/home/raymxnd/Workspace/WeatherForcast/.venv/lib/python3.12/site-packages/jax/_src/api.py", line 2553, in _device_get
+ return toarray()
+ ^^^^^^^^^
+ File "/home/raymxnd/Workspace/WeatherForcast/.venv/lib/python3.12/site-packages/jax/_src/array.py", line 409, in __array__
+ return np.asarray(self._value, dtype=dtype, **kwds) # pyrefly: ignore[no-matching-overload]
+ ^^^^^^^^^^^
+ File "/home/raymxnd/Workspace/WeatherForcast/.venv/lib/python3.12/site-packages/jax/_src/profiler.py", line 420, in wrapper
+ return func(*args, **kwargs)
+ ^^^^^^^^^^^^^^^^^^^^^
+ File "/home/raymxnd/Workspace/WeatherForcast/.venv/lib/python3.12/site-packages/jax/_src/array.py", line 637, in _value
+ npy_value, did_copy = self._single_device_array_to_np_array_did_copy()
+ ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
+jax.errors.JaxRuntimeError: RESOURCE_EXHAUSTED: Out of memory while trying to allocate 46.10GiB with allocator GPU_0_bfc on device 0. [executable_name='jit_step'] [tf-allocator-allocation-error='']
diff --git a/docs/results/finetune/finetune_lr3e-6_rtx5090.json b/docs/results/finetune/finetune_lr3e-6_rtx5090.json
new file mode 100644
index 0000000..44e5437
--- /dev/null
+++ b/docs/results/finetune/finetune_lr3e-6_rtx5090.json
@@ -0,0 +1,848 @@
+{
+ "model": "WeatherNext2",
+ "checkpoint": 1,
+ "device": "NVIDIA GeForce RTX 5090",
+ "jax": "0.11.2",
+ "faster_weathernext": "0.1.1",
+ "args": {
+ "model": "WeatherNext2",
+ "checkpoint": 1,
+ "train": [
+ "2026100412",
+ "2026100418",
+ "2026100500",
+ "2026100506",
+ "2026100512",
+ "2026100518",
+ "2026100600",
+ "2026100606",
+ "2026100612",
+ "2026100618",
+ "2026100700",
+ "2026100706",
+ "2026100712",
+ "2026100718",
+ "2026100800",
+ "2026100806",
+ "2026100912",
+ "2026100918",
+ "2026101000",
+ "2026101006",
+ "2026101112",
+ "2026101118",
+ "2026101200",
+ "2026101206",
+ "2026101212",
+ "2026101218",
+ "2026101300",
+ "2026101306",
+ "2026101412",
+ "2026101418",
+ "2026101500",
+ "2026101506",
+ "2026101512",
+ "2026101518",
+ "2026101600",
+ "2026101606",
+ "2026101612",
+ "2026101618",
+ "2026101700",
+ "2026101706",
+ "2026101712",
+ "2026101718",
+ "2026101800",
+ "2026101806",
+ "2026101812",
+ "2026101818",
+ "2026101900",
+ "2026101906",
+ "2026101912",
+ "2026101918",
+ "2026102000",
+ "2026102006",
+ "2026102012",
+ "2026102018",
+ "2026102100",
+ "2026102106",
+ "2026102112",
+ "2026102118",
+ "2026102200",
+ "2026102206",
+ "2026102212",
+ "2026102218",
+ "2026102300",
+ "2026102306",
+ "2026102312",
+ "2026102318",
+ "2026102400",
+ "2026102406"
+ ],
+ "eval": [
+ "2026102800",
+ "2026102812",
+ "2026103000",
+ "2026103012"
+ ],
+ "steps": 60,
+ "lr": 3e-06,
+ "warmup": 10,
+ "weight_decay": 0.0,
+ "clip": 1.0,
+ "eval_keys": 2,
+ "precision": "default",
+ "seed": 0,
+ "no_precip": false,
+ "mask": [
+ "sea_surface_temperature"
+ ],
+ "compile_only": false,
+ "no_remat": false,
+ "weights_dir": "/home/raymxnd/Workspace/WeatherForcast/weights",
+ "out": null,
+ "log": "docs/results/finetune/finetune_lr3e-6_rtx5090.json"
+ },
+ "masked": [
+ "cyclone_all_wind_disc",
+ "cyclone_exists_gaussian_unit_mode",
+ "cyclone_usa_pres_disc",
+ "cyclone_usa_r34_ne_radius_disc",
+ "cyclone_usa_r34_nw_radius_disc",
+ "cyclone_usa_r34_se_radius_disc",
+ "cyclone_usa_r34_sw_radius_disc",
+ "cyclone_usa_r50_ne_radius_disc",
+ "cyclone_usa_r50_nw_radius_disc",
+ "cyclone_usa_r50_se_radius_disc",
+ "cyclone_usa_r50_sw_radius_disc",
+ "cyclone_usa_r64_ne_radius_disc",
+ "cyclone_usa_r64_nw_radius_disc",
+ "cyclone_usa_r64_se_radius_disc",
+ "cyclone_usa_r64_sw_radius_disc",
+ "cyclone_usa_rmw_disc",
+ "cyclone_usa_wind_disc",
+ "sea_surface_temperature"
+ ],
+ "head_params": 5572967,
+ "total_params": 183785575,
+ "compile": {
+ "seconds": 84.99111080169678,
+ "temp_gib": 20.03102134913206,
+ "argument_gib": 1.7898386679589748,
+ "output_gib": 0.0622851625084877
+ },
+ "eval_before": {
+ "2026102800": {
+ "loss": 1.040929913520813,
+ "per_var": {
+ "100m_u_component_of_wind": 0.13943792879581451,
+ "100m_v_component_of_wind": 0.12798184156417847,
+ "10m_u_component_of_wind": 0.14852693676948547,
+ "10m_v_component_of_wind": 0.13638706505298615,
+ "2m_temperature": 0.09907613694667816,
+ "cyclone_all_wind_disc": 0.0,
+ "cyclone_exists_gaussian_unit_mode": 0.0,
+ "cyclone_usa_pres_disc": 0.0,
+ "cyclone_usa_r34_ne_radius_disc": 0.0,
+ "cyclone_usa_r34_nw_radius_disc": 0.0,
+ "cyclone_usa_r34_se_radius_disc": 0.0,
+ "cyclone_usa_r34_sw_radius_disc": 0.0,
+ "cyclone_usa_r50_ne_radius_disc": 0.0,
+ "cyclone_usa_r50_nw_radius_disc": 0.0,
+ "cyclone_usa_r50_se_radius_disc": 0.0,
+ "cyclone_usa_r50_sw_radius_disc": 0.0,
+ "cyclone_usa_r64_ne_radius_disc": 0.0,
+ "cyclone_usa_r64_nw_radius_disc": 0.0,
+ "cyclone_usa_r64_se_radius_disc": 0.0,
+ "cyclone_usa_r64_sw_radius_disc": 0.0,
+ "cyclone_usa_rmw_disc": 0.0,
+ "cyclone_usa_wind_disc": 0.0,
+ "geopotential": 0.06217382848262787,
+ "mean_sea_level_pressure": 0.06353482604026794,
+ "sea_surface_temperature": 1.5886213779449463,
+ "specific_humidity": 0.21852953732013702,
+ "temperature": 0.12121714651584625,
+ "total_precipitation_6hr": 0.12706182897090912,
+ "u_component_of_wind": 0.1305268108844757,
+ "v_component_of_wind": 0.11552518606185913,
+ "vertical_velocity": 0.2782917618751526
+ }
+ },
+ "2026102812": {
+ "loss": 1.0558066368103027,
+ "per_var": {
+ "100m_u_component_of_wind": 0.14223292469978333,
+ "100m_v_component_of_wind": 0.13159583508968353,
+ "10m_u_component_of_wind": 0.14976075291633606,
+ "10m_v_component_of_wind": 0.1381133496761322,
+ "2m_temperature": 0.10236659646034241,
+ "cyclone_all_wind_disc": 0.0,
+ "cyclone_exists_gaussian_unit_mode": 0.0,
+ "cyclone_usa_pres_disc": 0.0,
+ "cyclone_usa_r34_ne_radius_disc": 0.0,
+ "cyclone_usa_r34_nw_radius_disc": 0.0,
+ "cyclone_usa_r34_se_radius_disc": 0.0,
+ "cyclone_usa_r34_sw_radius_disc": 0.0,
+ "cyclone_usa_r50_ne_radius_disc": 0.0,
+ "cyclone_usa_r50_nw_radius_disc": 0.0,
+ "cyclone_usa_r50_se_radius_disc": 0.0,
+ "cyclone_usa_r50_sw_radius_disc": 0.0,
+ "cyclone_usa_r64_ne_radius_disc": 0.0,
+ "cyclone_usa_r64_nw_radius_disc": 0.0,
+ "cyclone_usa_r64_se_radius_disc": 0.0,
+ "cyclone_usa_r64_sw_radius_disc": 0.0,
+ "cyclone_usa_rmw_disc": 0.0,
+ "cyclone_usa_wind_disc": 0.0,
+ "geopotential": 0.05947427451610565,
+ "mean_sea_level_pressure": 0.06730848550796509,
+ "sea_surface_temperature": 1.6101053953170776,
+ "specific_humidity": 0.22227931022644043,
+ "temperature": 0.12098470330238342,
+ "total_precipitation_6hr": 0.13528835773468018,
+ "u_component_of_wind": 0.13082078099250793,
+ "v_component_of_wind": 0.11493003368377686,
+ "vertical_velocity": 0.2863432765007019
+ }
+ },
+ "2026103000": {
+ "loss": 1.007240891456604,
+ "per_var": {
+ "100m_u_component_of_wind": 0.13535135984420776,
+ "100m_v_component_of_wind": 0.12463852018117905,
+ "10m_u_component_of_wind": 0.14480358362197876,
+ "10m_v_component_of_wind": 0.13379094004631042,
+ "2m_temperature": 0.09720313549041748,
+ "cyclone_all_wind_disc": 0.0,
+ "cyclone_exists_gaussian_unit_mode": 0.0,
+ "cyclone_usa_pres_disc": 0.0,
+ "cyclone_usa_r34_ne_radius_disc": 0.0,
+ "cyclone_usa_r34_nw_radius_disc": 0.0,
+ "cyclone_usa_r34_se_radius_disc": 0.0,
+ "cyclone_usa_r34_sw_radius_disc": 0.0,
+ "cyclone_usa_r50_ne_radius_disc": 0.0,
+ "cyclone_usa_r50_nw_radius_disc": 0.0,
+ "cyclone_usa_r50_se_radius_disc": 0.0,
+ "cyclone_usa_r50_sw_radius_disc": 0.0,
+ "cyclone_usa_r64_ne_radius_disc": 0.0,
+ "cyclone_usa_r64_nw_radius_disc": 0.0,
+ "cyclone_usa_r64_se_radius_disc": 0.0,
+ "cyclone_usa_r64_sw_radius_disc": 0.0,
+ "cyclone_usa_rmw_disc": 0.0,
+ "cyclone_usa_wind_disc": 0.0,
+ "geopotential": 0.06037849187850952,
+ "mean_sea_level_pressure": 0.06656663864850998,
+ "sea_surface_temperature": 1.5575722455978394,
+ "specific_humidity": 0.21349863708019257,
+ "temperature": 0.11542174220085144,
+ "total_precipitation_6hr": 0.12492766231298447,
+ "u_component_of_wind": 0.12509292364120483,
+ "v_component_of_wind": 0.11082534492015839,
+ "vertical_velocity": 0.2689312696456909
+ }
+ },
+ "2026103012": {
+ "loss": 1.0289032459259033,
+ "per_var": {
+ "100m_u_component_of_wind": 0.13896289467811584,
+ "100m_v_component_of_wind": 0.1291162669658661,
+ "10m_u_component_of_wind": 0.14678022265434265,
+ "10m_v_component_of_wind": 0.13629581034183502,
+ "2m_temperature": 0.09881523251533508,
+ "cyclone_all_wind_disc": 0.0,
+ "cyclone_exists_gaussian_unit_mode": 0.0,
+ "cyclone_usa_pres_disc": 0.0,
+ "cyclone_usa_r34_ne_radius_disc": 0.0,
+ "cyclone_usa_r34_nw_radius_disc": 0.0,
+ "cyclone_usa_r34_se_radius_disc": 0.0,
+ "cyclone_usa_r34_sw_radius_disc": 0.0,
+ "cyclone_usa_r50_ne_radius_disc": 0.0,
+ "cyclone_usa_r50_nw_radius_disc": 0.0,
+ "cyclone_usa_r50_se_radius_disc": 0.0,
+ "cyclone_usa_r50_sw_radius_disc": 0.0,
+ "cyclone_usa_r64_ne_radius_disc": 0.0,
+ "cyclone_usa_r64_nw_radius_disc": 0.0,
+ "cyclone_usa_r64_se_radius_disc": 0.0,
+ "cyclone_usa_r64_sw_radius_disc": 0.0,
+ "cyclone_usa_rmw_disc": 0.0,
+ "cyclone_usa_wind_disc": 0.0,
+ "geopotential": 0.05989237129688263,
+ "mean_sea_level_pressure": 0.06428378820419312,
+ "sea_surface_temperature": 1.6431248188018799,
+ "specific_humidity": 0.21590983867645264,
+ "temperature": 0.11675645411014557,
+ "total_precipitation_6hr": 0.14074179530143738,
+ "u_component_of_wind": 0.12529444694519043,
+ "v_component_of_wind": 0.11094464361667633,
+ "vertical_velocity": 0.28317612409591675
+ }
+ }
+ },
+ "steps": [
+ {
+ "step": 0,
+ "init": "2026101006",
+ "loss": 1.1192269325256348,
+ "grad_norm": 3.3067853450775146,
+ "seconds": 3.989969491958618
+ },
+ {
+ "step": 1,
+ "init": "2026102406",
+ "loss": 1.0819909572601318,
+ "grad_norm": 3.9273650646209717,
+ "seconds": 3.8303279876708984
+ },
+ {
+ "step": 2,
+ "init": "2026101206",
+ "loss": 1.1632181406021118,
+ "grad_norm": 4.088436603546143,
+ "seconds": 3.829003095626831
+ },
+ {
+ "step": 3,
+ "init": "2026100612",
+ "loss": 0.9969218969345093,
+ "grad_norm": 3.07970929145813,
+ "seconds": 3.8305106163024902
+ },
+ {
+ "step": 4,
+ "init": "2026101900",
+ "loss": 1.029414176940918,
+ "grad_norm": 3.816864252090454,
+ "seconds": 3.8323919773101807
+ },
+ {
+ "step": 5,
+ "init": "2026101906",
+ "loss": 1.0838956832885742,
+ "grad_norm": 3.5104336738586426,
+ "seconds": 3.833014488220215
+ },
+ {
+ "step": 6,
+ "init": "2026102300",
+ "loss": 1.1333730220794678,
+ "grad_norm": 4.801530838012695,
+ "seconds": 3.8334226608276367
+ },
+ {
+ "step": 7,
+ "init": "2026100512",
+ "loss": 1.059767484664917,
+ "grad_norm": 3.5471932888031006,
+ "seconds": 3.8337552547454834
+ },
+ {
+ "step": 8,
+ "init": "2026101806",
+ "loss": 1.1058920621871948,
+ "grad_norm": 3.2356626987457275,
+ "seconds": 3.8325774669647217
+ },
+ {
+ "step": 9,
+ "init": "2026102200",
+ "loss": 0.9849255084991455,
+ "grad_norm": 3.2249393463134766,
+ "seconds": 3.8339390754699707
+ },
+ {
+ "step": 10,
+ "init": "2026101500",
+ "loss": 1.1045509576797485,
+ "grad_norm": 4.718138694763184,
+ "seconds": 3.834264039993286
+ },
+ {
+ "step": 11,
+ "init": "2026100700",
+ "loss": 1.0667481422424316,
+ "grad_norm": 4.425965785980225,
+ "seconds": 3.8351261615753174
+ },
+ {
+ "step": 12,
+ "init": "2026100912",
+ "loss": 1.011527419090271,
+ "grad_norm": 2.7919561862945557,
+ "seconds": 3.8342647552490234
+ },
+ {
+ "step": 13,
+ "init": "2026100500",
+ "loss": 1.0096163749694824,
+ "grad_norm": 3.288435459136963,
+ "seconds": 3.8350107669830322
+ },
+ {
+ "step": 14,
+ "init": "2026101812",
+ "loss": 1.0679728984832764,
+ "grad_norm": 4.00501823425293,
+ "seconds": 3.8349153995513916
+ },
+ {
+ "step": 15,
+ "init": "2026101612",
+ "loss": 0.9622900485992432,
+ "grad_norm": 2.382941722869873,
+ "seconds": 3.835181474685669
+ },
+ {
+ "step": 16,
+ "init": "2026101112",
+ "loss": 0.9930404424667358,
+ "grad_norm": 4.04895544052124,
+ "seconds": 3.8349363803863525
+ },
+ {
+ "step": 17,
+ "init": "2026102006",
+ "loss": 1.0657817125320435,
+ "grad_norm": 3.5417263507843018,
+ "seconds": 3.833918333053589
+ },
+ {
+ "step": 18,
+ "init": "2026100706",
+ "loss": 1.1219922304153442,
+ "grad_norm": 3.346857786178589,
+ "seconds": 3.8370466232299805
+ },
+ {
+ "step": 19,
+ "init": "2026102312",
+ "loss": 0.9704004526138306,
+ "grad_norm": 2.805772304534912,
+ "seconds": 3.833247184753418
+ },
+ {
+ "step": 20,
+ "init": "2026102106",
+ "loss": 1.0882494449615479,
+ "grad_norm": 3.0195047855377197,
+ "seconds": 3.835040807723999
+ },
+ {
+ "step": 21,
+ "init": "2026101606",
+ "loss": 1.1693618297576904,
+ "grad_norm": 3.9549331665039062,
+ "seconds": 3.8342206478118896
+ },
+ {
+ "step": 22,
+ "init": "2026101000",
+ "loss": 0.997355580329895,
+ "grad_norm": 2.5960371494293213,
+ "seconds": 3.834469795227051
+ },
+ {
+ "step": 23,
+ "init": "2026101600",
+ "loss": 0.9624165296554565,
+ "grad_norm": 3.098451614379883,
+ "seconds": 3.833658456802368
+ },
+ {
+ "step": 24,
+ "init": "2026100506",
+ "loss": 1.057349443435669,
+ "grad_norm": 2.3970205783843994,
+ "seconds": 3.8359148502349854
+ },
+ {
+ "step": 25,
+ "init": "2026100418",
+ "loss": 1.1982371807098389,
+ "grad_norm": 3.652320623397827,
+ "seconds": 3.835576057434082
+ },
+ {
+ "step": 26,
+ "init": "2026101412",
+ "loss": 0.9693777561187744,
+ "grad_norm": 2.619764804840088,
+ "seconds": 3.8372726440429688
+ },
+ {
+ "step": 27,
+ "init": "2026102306",
+ "loss": 1.070088267326355,
+ "grad_norm": 3.091752529144287,
+ "seconds": 3.8364055156707764
+ },
+ {
+ "step": 28,
+ "init": "2026101212",
+ "loss": 0.9595528841018677,
+ "grad_norm": 2.85148286819458,
+ "seconds": 3.8347113132476807
+ },
+ {
+ "step": 29,
+ "init": "2026100918",
+ "loss": 1.117027759552002,
+ "grad_norm": 2.7948904037475586,
+ "seconds": 3.8349485397338867
+ },
+ {
+ "step": 30,
+ "init": "2026101306",
+ "loss": 1.0593416690826416,
+ "grad_norm": 3.2000768184661865,
+ "seconds": 3.8370249271392822
+ },
+ {
+ "step": 31,
+ "init": "2026101800",
+ "loss": 0.9512765407562256,
+ "grad_norm": 2.3337655067443848,
+ "seconds": 3.837232828140259
+ },
+ {
+ "step": 32,
+ "init": "2026101118",
+ "loss": 1.0578129291534424,
+ "grad_norm": 2.445765495300293,
+ "seconds": 3.8347885608673096
+ },
+ {
+ "step": 33,
+ "init": "2026100412",
+ "loss": 1.0441654920578003,
+ "grad_norm": 3.415700674057007,
+ "seconds": 3.835824966430664
+ },
+ {
+ "step": 34,
+ "init": "2026100600",
+ "loss": 1.1035103797912598,
+ "grad_norm": 3.813368558883667,
+ "seconds": 3.8349099159240723
+ },
+ {
+ "step": 35,
+ "init": "2026101200",
+ "loss": 1.0035624504089355,
+ "grad_norm": 3.292027473449707,
+ "seconds": 3.8359906673431396
+ },
+ {
+ "step": 36,
+ "init": "2026101300",
+ "loss": 0.9259215593338013,
+ "grad_norm": 2.8280227184295654,
+ "seconds": 3.8350348472595215
+ },
+ {
+ "step": 37,
+ "init": "2026101618",
+ "loss": 1.0876407623291016,
+ "grad_norm": 3.2235662937164307,
+ "seconds": 3.83548903465271
+ },
+ {
+ "step": 38,
+ "init": "2026102318",
+ "loss": 1.084805965423584,
+ "grad_norm": 2.56314754486084,
+ "seconds": 3.83660888671875
+ },
+ {
+ "step": 39,
+ "init": "2026102012",
+ "loss": 1.0008708238601685,
+ "grad_norm": 2.4713516235351562,
+ "seconds": 3.835223913192749
+ },
+ {
+ "step": 40,
+ "init": "2026100806",
+ "loss": 1.1869910955429077,
+ "grad_norm": 3.684384346008301,
+ "seconds": 3.8367152214050293
+ },
+ {
+ "step": 41,
+ "init": "2026101512",
+ "loss": 0.9650352001190186,
+ "grad_norm": 3.136996269226074,
+ "seconds": 3.8366951942443848
+ },
+ {
+ "step": 42,
+ "init": "2026101218",
+ "loss": 1.0908368825912476,
+ "grad_norm": 3.1978204250335693,
+ "seconds": 3.8358519077301025
+ },
+ {
+ "step": 43,
+ "init": "2026102212",
+ "loss": 1.0293278694152832,
+ "grad_norm": 2.6473119258880615,
+ "seconds": 3.8365278244018555
+ },
+ {
+ "step": 44,
+ "init": "2026100618",
+ "loss": 1.0616509914398193,
+ "grad_norm": 2.3891847133636475,
+ "seconds": 3.8358805179595947
+ },
+ {
+ "step": 45,
+ "init": "2026101700",
+ "loss": 1.0147740840911865,
+ "grad_norm": 3.1602120399475098,
+ "seconds": 3.8354666233062744
+ },
+ {
+ "step": 46,
+ "init": "2026102018",
+ "loss": 1.0929865837097168,
+ "grad_norm": 2.8981800079345703,
+ "seconds": 3.8365137577056885
+ },
+ {
+ "step": 47,
+ "init": "2026102218",
+ "loss": 1.0778894424438477,
+ "grad_norm": 2.439805507659912,
+ "seconds": 3.834641933441162
+ },
+ {
+ "step": 48,
+ "init": "2026102000",
+ "loss": 1.0053120851516724,
+ "grad_norm": 3.6493818759918213,
+ "seconds": 3.834256649017334
+ },
+ {
+ "step": 49,
+ "init": "2026101712",
+ "loss": 0.9970908164978027,
+ "grad_norm": 2.562253713607788,
+ "seconds": 3.8366057872772217
+ },
+ {
+ "step": 50,
+ "init": "2026101912",
+ "loss": 1.0124934911727905,
+ "grad_norm": 2.590519428253174,
+ "seconds": 3.836496114730835
+ },
+ {
+ "step": 51,
+ "init": "2026100718",
+ "loss": 1.0638058185577393,
+ "grad_norm": 3.040156841278076,
+ "seconds": 3.8346924781799316
+ },
+ {
+ "step": 52,
+ "init": "2026100712",
+ "loss": 1.060908317565918,
+ "grad_norm": 3.2956037521362305,
+ "seconds": 3.836531639099121
+ },
+ {
+ "step": 53,
+ "init": "2026100606",
+ "loss": 1.0745906829833984,
+ "grad_norm": 2.741349935531616,
+ "seconds": 3.8374931812286377
+ },
+ {
+ "step": 54,
+ "init": "2026101818",
+ "loss": 1.1122087240219116,
+ "grad_norm": 2.9382705688476562,
+ "seconds": 3.8357670307159424
+ },
+ {
+ "step": 55,
+ "init": "2026101706",
+ "loss": 1.0591838359832764,
+ "grad_norm": 2.7150232791900635,
+ "seconds": 3.834744453430176
+ },
+ {
+ "step": 56,
+ "init": "2026101506",
+ "loss": 1.118042230606079,
+ "grad_norm": 3.073740005493164,
+ "seconds": 3.836676597595215
+ },
+ {
+ "step": 57,
+ "init": "2026100518",
+ "loss": 1.0759599208831787,
+ "grad_norm": 2.7621521949768066,
+ "seconds": 3.835617780685425
+ },
+ {
+ "step": 58,
+ "init": "2026101918",
+ "loss": 1.0416721105575562,
+ "grad_norm": 2.201932430267334,
+ "seconds": 3.83652663230896
+ },
+ {
+ "step": 59,
+ "init": "2026102118",
+ "loss": 1.1190178394317627,
+ "grad_norm": 3.2081642150878906,
+ "seconds": 3.831350803375244
+ }
+ ],
+ "peak_gib": 21.82089877128601,
+ "median_step_s": 3.835040807723999,
+ "eval_after": {
+ "2026102800": {
+ "loss": 1.0250498056411743,
+ "per_var": {
+ "100m_u_component_of_wind": 0.1381281614303589,
+ "100m_v_component_of_wind": 0.12788733839988708,
+ "10m_u_component_of_wind": 0.14678412675857544,
+ "10m_v_component_of_wind": 0.13595493137836456,
+ "2m_temperature": 0.08931844681501389,
+ "cyclone_all_wind_disc": 0.0,
+ "cyclone_exists_gaussian_unit_mode": 0.0,
+ "cyclone_usa_pres_disc": 0.0,
+ "cyclone_usa_r34_ne_radius_disc": 0.0,
+ "cyclone_usa_r34_nw_radius_disc": 0.0,
+ "cyclone_usa_r34_se_radius_disc": 0.0,
+ "cyclone_usa_r34_sw_radius_disc": 0.0,
+ "cyclone_usa_r50_ne_radius_disc": 0.0,
+ "cyclone_usa_r50_nw_radius_disc": 0.0,
+ "cyclone_usa_r50_se_radius_disc": 0.0,
+ "cyclone_usa_r50_sw_radius_disc": 0.0,
+ "cyclone_usa_r64_ne_radius_disc": 0.0,
+ "cyclone_usa_r64_nw_radius_disc": 0.0,
+ "cyclone_usa_r64_se_radius_disc": 0.0,
+ "cyclone_usa_r64_sw_radius_disc": 0.0,
+ "cyclone_usa_rmw_disc": 0.0,
+ "cyclone_usa_wind_disc": 0.0,
+ "geopotential": 0.06130439043045044,
+ "mean_sea_level_pressure": 0.06384440511465073,
+ "sea_surface_temperature": 1.5873539447784424,
+ "specific_humidity": 0.2173776924610138,
+ "temperature": 0.11828567087650299,
+ "total_precipitation_6hr": 0.12172940373420715,
+ "u_component_of_wind": 0.13013312220573425,
+ "v_component_of_wind": 0.11572330445051193,
+ "vertical_velocity": 0.27756422758102417
+ }
+ },
+ "2026102812": {
+ "loss": 1.0390284061431885,
+ "per_var": {
+ "100m_u_component_of_wind": 0.14036166667938232,
+ "100m_v_component_of_wind": 0.1308678686618805,
+ "10m_u_component_of_wind": 0.1474299430847168,
+ "10m_v_component_of_wind": 0.1369265615940094,
+ "2m_temperature": 0.09393928200006485,
+ "cyclone_all_wind_disc": 0.0,
+ "cyclone_exists_gaussian_unit_mode": 0.0,
+ "cyclone_usa_pres_disc": 0.0,
+ "cyclone_usa_r34_ne_radius_disc": 0.0,
+ "cyclone_usa_r34_nw_radius_disc": 0.0,
+ "cyclone_usa_r34_se_radius_disc": 0.0,
+ "cyclone_usa_r34_sw_radius_disc": 0.0,
+ "cyclone_usa_r50_ne_radius_disc": 0.0,
+ "cyclone_usa_r50_nw_radius_disc": 0.0,
+ "cyclone_usa_r50_se_radius_disc": 0.0,
+ "cyclone_usa_r50_sw_radius_disc": 0.0,
+ "cyclone_usa_r64_ne_radius_disc": 0.0,
+ "cyclone_usa_r64_nw_radius_disc": 0.0,
+ "cyclone_usa_r64_se_radius_disc": 0.0,
+ "cyclone_usa_r64_sw_radius_disc": 0.0,
+ "cyclone_usa_rmw_disc": 0.0,
+ "cyclone_usa_wind_disc": 0.0,
+ "geopotential": 0.058046795427799225,
+ "mean_sea_level_pressure": 0.06548143178224564,
+ "sea_surface_temperature": 1.6100666522979736,
+ "specific_humidity": 0.22043463587760925,
+ "temperature": 0.11776454746723175,
+ "total_precipitation_6hr": 0.12899558246135712,
+ "u_component_of_wind": 0.130298912525177,
+ "v_component_of_wind": 0.11487138271331787,
+ "vertical_velocity": 0.2854691743850708
+ }
+ },
+ "2026103000": {
+ "loss": 0.9917856454849243,
+ "per_var": {
+ "100m_u_component_of_wind": 0.13411208987236023,
+ "100m_v_component_of_wind": 0.1244078055024147,
+ "10m_u_component_of_wind": 0.14330054819583893,
+ "10m_v_component_of_wind": 0.13336080312728882,
+ "2m_temperature": 0.08904600143432617,
+ "cyclone_all_wind_disc": 0.0,
+ "cyclone_exists_gaussian_unit_mode": 0.0,
+ "cyclone_usa_pres_disc": 0.0,
+ "cyclone_usa_r34_ne_radius_disc": 0.0,
+ "cyclone_usa_r34_nw_radius_disc": 0.0,
+ "cyclone_usa_r34_se_radius_disc": 0.0,
+ "cyclone_usa_r34_sw_radius_disc": 0.0,
+ "cyclone_usa_r50_ne_radius_disc": 0.0,
+ "cyclone_usa_r50_nw_radius_disc": 0.0,
+ "cyclone_usa_r50_se_radius_disc": 0.0,
+ "cyclone_usa_r50_sw_radius_disc": 0.0,
+ "cyclone_usa_r64_ne_radius_disc": 0.0,
+ "cyclone_usa_r64_nw_radius_disc": 0.0,
+ "cyclone_usa_r64_se_radius_disc": 0.0,
+ "cyclone_usa_r64_sw_radius_disc": 0.0,
+ "cyclone_usa_rmw_disc": 0.0,
+ "cyclone_usa_wind_disc": 0.0,
+ "geopotential": 0.05909130722284317,
+ "mean_sea_level_pressure": 0.06461087614297867,
+ "sea_surface_temperature": 1.5539729595184326,
+ "specific_humidity": 0.2121933400630951,
+ "temperature": 0.1121445894241333,
+ "total_precipitation_6hr": 0.11939429491758347,
+ "u_component_of_wind": 0.1246003732085228,
+ "v_component_of_wind": 0.11093677580356598,
+ "vertical_velocity": 0.26815927028656006
+ }
+ },
+ "2026103012": {
+ "loss": 1.013411283493042,
+ "per_var": {
+ "100m_u_component_of_wind": 0.1376926302909851,
+ "100m_v_component_of_wind": 0.12838691473007202,
+ "10m_u_component_of_wind": 0.1449083685874939,
+ "10m_v_component_of_wind": 0.13539353013038635,
+ "2m_temperature": 0.09186872094869614,
+ "cyclone_all_wind_disc": 0.0,
+ "cyclone_exists_gaussian_unit_mode": 0.0,
+ "cyclone_usa_pres_disc": 0.0,
+ "cyclone_usa_r34_ne_radius_disc": 0.0,
+ "cyclone_usa_r34_nw_radius_disc": 0.0,
+ "cyclone_usa_r34_se_radius_disc": 0.0,
+ "cyclone_usa_r34_sw_radius_disc": 0.0,
+ "cyclone_usa_r50_ne_radius_disc": 0.0,
+ "cyclone_usa_r50_nw_radius_disc": 0.0,
+ "cyclone_usa_r50_se_radius_disc": 0.0,
+ "cyclone_usa_r50_sw_radius_disc": 0.0,
+ "cyclone_usa_r64_ne_radius_disc": 0.0,
+ "cyclone_usa_r64_nw_radius_disc": 0.0,
+ "cyclone_usa_r64_se_radius_disc": 0.0,
+ "cyclone_usa_r64_sw_radius_disc": 0.0,
+ "cyclone_usa_rmw_disc": 0.0,
+ "cyclone_usa_wind_disc": 0.0,
+ "geopotential": 0.05839185416698456,
+ "mean_sea_level_pressure": 0.06337311118841171,
+ "sea_surface_temperature": 1.6430537700653076,
+ "specific_humidity": 0.21391956508159637,
+ "temperature": 0.11382471024990082,
+ "total_precipitation_6hr": 0.13409870862960815,
+ "u_component_of_wind": 0.12452033907175064,
+ "v_component_of_wind": 0.11072681844234467,
+ "vertical_velocity": 0.28228896856307983
+ }
+ }
+ }
+}
\ No newline at end of file
diff --git a/docs/results/finetune/finetune_lr3e-6_rtx5090.log b/docs/results/finetune/finetune_lr3e-6_rtx5090.log
new file mode 100644
index 0000000..d580f29
--- /dev/null
+++ b/docs/results/finetune/finetune_lr3e-6_rtx5090.log
@@ -0,0 +1,98 @@
+example 2026100412 built in 2.2s
+WeatherNext2 ckpt 1: training 5.57M of 183.8M params; masked targets: 18 (cyclone channels, sea_surface_temperature)
+WARNING:root:No normalization location found for cyclone_exists_gaussian_unit_mode
+WARNING:root:No normalization scale found for cyclone_exists_gaussian_unit_mode
+WARNING:absl:No coordinates found for dim noise_channels of array Size: 256B
+JitTracer(float32[2,1,32])
+Coordinates:
+ * batch (batch) object 16B MultiIndex
+ * _batch (batch) int64 16B 0 0
+ * sample (batch) int64 16B 0 1
+Dimensions without coordinates: channels, noise_channels., using index for param name. This may make model surgery more error prone if changing coordinates at fine-tuning time.
+/home/raymxnd/Workspace/WeatherForcast/.venv/lib/python3.12/site-packages/scipy/sparse/_index.py:216: SparseEfficiencyWarning: Changing the sparsity structure of a csr_matrix is expensive. lil and dok are more efficient.
+ self._set_arrayXarray(i, j, x)
+/home/raymxnd/Workspace/WeatherForcast/weathernext/weathernext/weathernext2/architecture.py:286: FutureWarning: Deleting a single level of a MultiIndex is deprecated. Previously, this deleted all levels of a MultiIndex. Please also drop the following variables: {'batch'} to avoid an error in the future.
+ predictions[var] = xarray_jax.apply_ufunc(activation_fn, predictions[var])
+W1031 13:24:23.786936 1649993 hlo_rematerialization.cc:3363] Can't reduce memory use below 25.14GiB (26997919706 bytes) by rematerialization; only reduced to 63.07GiB (67717567756 bytes), down from 63.07GiB (67717567756 bytes) originally
+train step compiled in 85s: temp 20.03 GiB, arguments 1.79 GiB
+WARNING:root:No normalization location found for cyclone_exists_gaussian_unit_mode
+WARNING:root:No normalization scale found for cyclone_exists_gaussian_unit_mode
+WARNING:absl:No coordinates found for dim noise_channels of array Size: 256B
+JitTracer(float32[2,1,32])
+Coordinates:
+ * batch (batch) object 16B MultiIndex
+ * _batch (batch) int64 16B 0 0
+ * sample (batch) int64 16B 0 1
+Dimensions without coordinates: channels, noise_channels., using index for param name. This may make model surgery more error prone if changing coordinates at fine-tuning time.
+/home/raymxnd/Workspace/WeatherForcast/.venv/lib/python3.12/site-packages/scipy/sparse/_index.py:216: SparseEfficiencyWarning: Changing the sparsity structure of a csr_matrix is expensive. lil and dok are more efficient.
+ self._set_arrayXarray(i, j, x)
+/home/raymxnd/Workspace/WeatherForcast/weathernext/weathernext/weathernext2/architecture.py:286: FutureWarning: Deleting a single level of a MultiIndex is deprecated. Previously, this deleted all levels of a MultiIndex. Please also drop the following variables: {'batch'} to avoid an error in the future.
+ predictions[var] = xarray_jax.apply_ufunc(activation_fn, predictions[var])
+W1031 13:25:33.534081 1649993 hlo_rematerialization.cc:3363] Can't reduce memory use below 25.12GiB (26978079655 bytes) by rematerialization; only reduced to 56.83GiB (61019999108 bytes), down from 56.83GiB (61020035972 bytes) originally
+eval before (77s): 2026102800 1.04093 2026102812 1.05581 2026103000 1.00724 2026103012 1.02890
+step 0 2026101006 loss 1.11923 |g| 3.307e+00 3.99s
+step 1 2026102406 loss 1.08199 |g| 3.927e+00 3.83s
+step 2 2026101206 loss 1.16322 |g| 4.088e+00 3.83s
+step 3 2026100612 loss 0.99692 |g| 3.080e+00 3.83s
+step 4 2026101900 loss 1.02941 |g| 3.817e+00 3.83s
+step 5 2026101906 loss 1.08390 |g| 3.510e+00 3.83s
+step 6 2026102300 loss 1.13337 |g| 4.802e+00 3.83s
+step 7 2026100512 loss 1.05977 |g| 3.547e+00 3.83s
+step 8 2026101806 loss 1.10589 |g| 3.236e+00 3.83s
+step 9 2026102200 loss 0.98493 |g| 3.225e+00 3.83s
+step 10 2026101500 loss 1.10455 |g| 4.718e+00 3.83s
+step 11 2026100700 loss 1.06675 |g| 4.426e+00 3.84s
+step 12 2026100912 loss 1.01153 |g| 2.792e+00 3.83s
+step 13 2026100500 loss 1.00962 |g| 3.288e+00 3.84s
+step 14 2026101812 loss 1.06797 |g| 4.005e+00 3.83s
+step 15 2026101612 loss 0.96229 |g| 2.383e+00 3.84s
+step 16 2026101112 loss 0.99304 |g| 4.049e+00 3.83s
+step 17 2026102006 loss 1.06578 |g| 3.542e+00 3.83s
+step 18 2026100706 loss 1.12199 |g| 3.347e+00 3.84s
+step 19 2026102312 loss 0.97040 |g| 2.806e+00 3.83s
+step 20 2026102106 loss 1.08825 |g| 3.020e+00 3.84s
+step 21 2026101606 loss 1.16936 |g| 3.955e+00 3.83s
+step 22 2026101000 loss 0.99736 |g| 2.596e+00 3.83s
+step 23 2026101600 loss 0.96242 |g| 3.098e+00 3.83s
+step 24 2026100506 loss 1.05735 |g| 2.397e+00 3.84s
+step 25 2026100418 loss 1.19824 |g| 3.652e+00 3.84s
+step 26 2026101412 loss 0.96938 |g| 2.620e+00 3.84s
+step 27 2026102306 loss 1.07009 |g| 3.092e+00 3.84s
+step 28 2026101212 loss 0.95955 |g| 2.851e+00 3.83s
+step 29 2026100918 loss 1.11703 |g| 2.795e+00 3.83s
+step 30 2026101306 loss 1.05934 |g| 3.200e+00 3.84s
+step 31 2026101800 loss 0.95128 |g| 2.334e+00 3.84s
+step 32 2026101118 loss 1.05781 |g| 2.446e+00 3.83s
+step 33 2026100412 loss 1.04417 |g| 3.416e+00 3.84s
+step 34 2026100600 loss 1.10351 |g| 3.813e+00 3.83s
+step 35 2026101200 loss 1.00356 |g| 3.292e+00 3.84s
+step 36 2026101300 loss 0.92592 |g| 2.828e+00 3.84s
+step 37 2026101618 loss 1.08764 |g| 3.224e+00 3.84s
+step 38 2026102318 loss 1.08481 |g| 2.563e+00 3.84s
+step 39 2026102012 loss 1.00087 |g| 2.471e+00 3.84s
+step 40 2026100806 loss 1.18699 |g| 3.684e+00 3.84s
+step 41 2026101512 loss 0.96504 |g| 3.137e+00 3.84s
+step 42 2026101218 loss 1.09084 |g| 3.198e+00 3.84s
+step 43 2026102212 loss 1.02933 |g| 2.647e+00 3.84s
+step 44 2026100618 loss 1.06165 |g| 2.389e+00 3.84s
+step 45 2026101700 loss 1.01477 |g| 3.160e+00 3.84s
+step 46 2026102018 loss 1.09299 |g| 2.898e+00 3.84s
+step 47 2026102218 loss 1.07789 |g| 2.440e+00 3.83s
+step 48 2026102000 loss 1.00531 |g| 3.649e+00 3.83s
+step 49 2026101712 loss 0.99709 |g| 2.562e+00 3.84s
+step 50 2026101912 loss 1.01249 |g| 2.591e+00 3.84s
+step 51 2026100718 loss 1.06381 |g| 3.040e+00 3.83s
+step 52 2026100712 loss 1.06091 |g| 3.296e+00 3.84s
+step 53 2026100606 loss 1.07459 |g| 2.741e+00 3.84s
+step 54 2026101818 loss 1.11221 |g| 2.938e+00 3.84s
+step 55 2026101706 loss 1.05918 |g| 2.715e+00 3.83s
+step 56 2026101506 loss 1.11804 |g| 3.074e+00 3.84s
+step 57 2026100518 loss 1.07596 |g| 2.762e+00 3.84s
+step 58 2026101918 loss 1.04167 |g| 2.202e+00 3.84s
+step 59 2026102118 loss 1.11902 |g| 3.208e+00 3.83s
+median step 3.84s, device peak 21.82 GiB
+eval 2026102800: 1.04093 -> 1.02505 (-1.53 %)
+eval 2026102812: 1.05581 -> 1.03903 (-1.59 %)
+eval 2026103000: 1.00724 -> 0.99179 (-1.53 %)
+eval 2026103012: 1.02890 -> 1.01341 (-1.51 %)
+wrote docs/results/finetune/finetune_lr3e-6_rtx5090.json
diff --git a/docs/results/finetune/finetune_masksst_rtx5090.json b/docs/results/finetune/finetune_masksst_rtx5090.json
new file mode 100644
index 0000000..92c72a2
--- /dev/null
+++ b/docs/results/finetune/finetune_masksst_rtx5090.json
@@ -0,0 +1,847 @@
+{
+ "model": "WeatherNext2",
+ "checkpoint": 1,
+ "device": "NVIDIA GeForce RTX 5090",
+ "jax": "0.11.2",
+ "faster_weathernext": "0.1.1",
+ "args": {
+ "model": "WeatherNext2",
+ "checkpoint": 1,
+ "train": [
+ "2026100412",
+ "2026100418",
+ "2026100500",
+ "2026100506",
+ "2026100512",
+ "2026100518",
+ "2026100600",
+ "2026100606",
+ "2026100612",
+ "2026100618",
+ "2026100700",
+ "2026100706",
+ "2026100712",
+ "2026100718",
+ "2026100800",
+ "2026100806",
+ "2026100912",
+ "2026100918",
+ "2026101000",
+ "2026101006",
+ "2026101112",
+ "2026101118",
+ "2026101200",
+ "2026101206",
+ "2026101212",
+ "2026101218",
+ "2026101300",
+ "2026101306",
+ "2026101412",
+ "2026101418",
+ "2026101500",
+ "2026101506",
+ "2026101512",
+ "2026101518",
+ "2026101600",
+ "2026101606",
+ "2026101612",
+ "2026101618",
+ "2026101700",
+ "2026101706",
+ "2026101712",
+ "2026101718",
+ "2026101800",
+ "2026101806",
+ "2026101812",
+ "2026101818",
+ "2026101900",
+ "2026101906",
+ "2026101912",
+ "2026101918",
+ "2026102000",
+ "2026102006",
+ "2026102012",
+ "2026102018",
+ "2026102100",
+ "2026102106",
+ "2026102112",
+ "2026102118",
+ "2026102200",
+ "2026102206",
+ "2026102212",
+ "2026102218",
+ "2026102300",
+ "2026102306",
+ "2026102312",
+ "2026102318",
+ "2026102400",
+ "2026102406"
+ ],
+ "eval": [
+ "2026102800",
+ "2026102812",
+ "2026103000",
+ "2026103012"
+ ],
+ "steps": 60,
+ "lr": 3e-05,
+ "weight_decay": 0.0,
+ "clip": 1.0,
+ "eval_keys": 2,
+ "precision": "default",
+ "seed": 0,
+ "no_precip": false,
+ "mask": [
+ "sea_surface_temperature"
+ ],
+ "compile_only": false,
+ "no_remat": false,
+ "weights_dir": "/home/raymxnd/Workspace/WeatherForcast/weights",
+ "out": null,
+ "log": "docs/results/finetune/finetune_masksst_rtx5090.json"
+ },
+ "masked": [
+ "cyclone_all_wind_disc",
+ "cyclone_exists_gaussian_unit_mode",
+ "cyclone_usa_pres_disc",
+ "cyclone_usa_r34_ne_radius_disc",
+ "cyclone_usa_r34_nw_radius_disc",
+ "cyclone_usa_r34_se_radius_disc",
+ "cyclone_usa_r34_sw_radius_disc",
+ "cyclone_usa_r50_ne_radius_disc",
+ "cyclone_usa_r50_nw_radius_disc",
+ "cyclone_usa_r50_se_radius_disc",
+ "cyclone_usa_r50_sw_radius_disc",
+ "cyclone_usa_r64_ne_radius_disc",
+ "cyclone_usa_r64_nw_radius_disc",
+ "cyclone_usa_r64_se_radius_disc",
+ "cyclone_usa_r64_sw_radius_disc",
+ "cyclone_usa_rmw_disc",
+ "cyclone_usa_wind_disc",
+ "sea_surface_temperature"
+ ],
+ "head_params": 5572967,
+ "total_params": 183785575,
+ "compile": {
+ "seconds": 85.42077422142029,
+ "temp_gib": 20.031021110713482,
+ "argument_gib": 1.7898386642336845,
+ "output_gib": 0.062285151332616806
+ },
+ "eval_before": {
+ "2026102800": {
+ "loss": 1.0409315824508667,
+ "per_var": {
+ "100m_u_component_of_wind": 0.13943974673748016,
+ "100m_v_component_of_wind": 0.1279769241809845,
+ "10m_u_component_of_wind": 0.14852724969387054,
+ "10m_v_component_of_wind": 0.13638198375701904,
+ "2m_temperature": 0.0990767851471901,
+ "cyclone_all_wind_disc": 0.0,
+ "cyclone_exists_gaussian_unit_mode": 0.0,
+ "cyclone_usa_pres_disc": 0.0,
+ "cyclone_usa_r34_ne_radius_disc": 0.0,
+ "cyclone_usa_r34_nw_radius_disc": 0.0,
+ "cyclone_usa_r34_se_radius_disc": 0.0,
+ "cyclone_usa_r34_sw_radius_disc": 0.0,
+ "cyclone_usa_r50_ne_radius_disc": 0.0,
+ "cyclone_usa_r50_nw_radius_disc": 0.0,
+ "cyclone_usa_r50_se_radius_disc": 0.0,
+ "cyclone_usa_r50_sw_radius_disc": 0.0,
+ "cyclone_usa_r64_ne_radius_disc": 0.0,
+ "cyclone_usa_r64_nw_radius_disc": 0.0,
+ "cyclone_usa_r64_se_radius_disc": 0.0,
+ "cyclone_usa_r64_sw_radius_disc": 0.0,
+ "cyclone_usa_rmw_disc": 0.0,
+ "cyclone_usa_wind_disc": 0.0,
+ "geopotential": 0.06217426806688309,
+ "mean_sea_level_pressure": 0.06353595107793808,
+ "sea_surface_temperature": 1.588623285293579,
+ "specific_humidity": 0.21853123605251312,
+ "temperature": 0.12121684104204178,
+ "total_precipitation_6hr": 0.12706270813941956,
+ "u_component_of_wind": 0.13052809238433838,
+ "v_component_of_wind": 0.11552539467811584,
+ "vertical_velocity": 0.27828991413116455
+ }
+ },
+ "2026102812": {
+ "loss": 1.0558079481124878,
+ "per_var": {
+ "100m_u_component_of_wind": 0.1422329545021057,
+ "100m_v_component_of_wind": 0.13159945607185364,
+ "10m_u_component_of_wind": 0.1497618854045868,
+ "10m_v_component_of_wind": 0.13811764121055603,
+ "2m_temperature": 0.1023666113615036,
+ "cyclone_all_wind_disc": 0.0,
+ "cyclone_exists_gaussian_unit_mode": 0.0,
+ "cyclone_usa_pres_disc": 0.0,
+ "cyclone_usa_r34_ne_radius_disc": 0.0,
+ "cyclone_usa_r34_nw_radius_disc": 0.0,
+ "cyclone_usa_r34_se_radius_disc": 0.0,
+ "cyclone_usa_r34_sw_radius_disc": 0.0,
+ "cyclone_usa_r50_ne_radius_disc": 0.0,
+ "cyclone_usa_r50_nw_radius_disc": 0.0,
+ "cyclone_usa_r50_se_radius_disc": 0.0,
+ "cyclone_usa_r50_sw_radius_disc": 0.0,
+ "cyclone_usa_r64_ne_radius_disc": 0.0,
+ "cyclone_usa_r64_nw_radius_disc": 0.0,
+ "cyclone_usa_r64_se_radius_disc": 0.0,
+ "cyclone_usa_r64_sw_radius_disc": 0.0,
+ "cyclone_usa_rmw_disc": 0.0,
+ "cyclone_usa_wind_disc": 0.0,
+ "geopotential": 0.059477370232343674,
+ "mean_sea_level_pressure": 0.06731019914150238,
+ "sea_surface_temperature": 1.6101070642471313,
+ "specific_humidity": 0.22228027880191803,
+ "temperature": 0.12098342925310135,
+ "total_precipitation_6hr": 0.13529044389724731,
+ "u_component_of_wind": 0.13082188367843628,
+ "v_component_of_wind": 0.11493133008480072,
+ "vertical_velocity": 0.2863400876522064
+ }
+ },
+ "2026103000": {
+ "loss": 1.0072473287582397,
+ "per_var": {
+ "100m_u_component_of_wind": 0.13535073399543762,
+ "100m_v_component_of_wind": 0.12464098632335663,
+ "10m_u_component_of_wind": 0.14480449259281158,
+ "10m_v_component_of_wind": 0.13379283249378204,
+ "2m_temperature": 0.09720295667648315,
+ "cyclone_all_wind_disc": 0.0,
+ "cyclone_exists_gaussian_unit_mode": 0.0,
+ "cyclone_usa_pres_disc": 0.0,
+ "cyclone_usa_r34_ne_radius_disc": 0.0,
+ "cyclone_usa_r34_nw_radius_disc": 0.0,
+ "cyclone_usa_r34_se_radius_disc": 0.0,
+ "cyclone_usa_r34_sw_radius_disc": 0.0,
+ "cyclone_usa_r50_ne_radius_disc": 0.0,
+ "cyclone_usa_r50_nw_radius_disc": 0.0,
+ "cyclone_usa_r50_se_radius_disc": 0.0,
+ "cyclone_usa_r50_sw_radius_disc": 0.0,
+ "cyclone_usa_r64_ne_radius_disc": 0.0,
+ "cyclone_usa_r64_nw_radius_disc": 0.0,
+ "cyclone_usa_r64_se_radius_disc": 0.0,
+ "cyclone_usa_r64_sw_radius_disc": 0.0,
+ "cyclone_usa_rmw_disc": 0.0,
+ "cyclone_usa_wind_disc": 0.0,
+ "geopotential": 0.060380153357982635,
+ "mean_sea_level_pressure": 0.06656627357006073,
+ "sea_surface_temperature": 1.5575697422027588,
+ "specific_humidity": 0.21350151300430298,
+ "temperature": 0.11542178690433502,
+ "total_precipitation_6hr": 0.12492793798446655,
+ "u_component_of_wind": 0.1250949203968048,
+ "v_component_of_wind": 0.11082611978054047,
+ "vertical_velocity": 0.26893115043640137
+ }
+ },
+ "2026103012": {
+ "loss": 1.0289126634597778,
+ "per_var": {
+ "100m_u_component_of_wind": 0.138963520526886,
+ "100m_v_component_of_wind": 0.12912282347679138,
+ "10m_u_component_of_wind": 0.14678165316581726,
+ "10m_v_component_of_wind": 0.13630428910255432,
+ "2m_temperature": 0.09881569445133209,
+ "cyclone_all_wind_disc": 0.0,
+ "cyclone_exists_gaussian_unit_mode": 0.0,
+ "cyclone_usa_pres_disc": 0.0,
+ "cyclone_usa_r34_ne_radius_disc": 0.0,
+ "cyclone_usa_r34_nw_radius_disc": 0.0,
+ "cyclone_usa_r34_se_radius_disc": 0.0,
+ "cyclone_usa_r34_sw_radius_disc": 0.0,
+ "cyclone_usa_r50_ne_radius_disc": 0.0,
+ "cyclone_usa_r50_nw_radius_disc": 0.0,
+ "cyclone_usa_r50_se_radius_disc": 0.0,
+ "cyclone_usa_r50_sw_radius_disc": 0.0,
+ "cyclone_usa_r64_ne_radius_disc": 0.0,
+ "cyclone_usa_r64_nw_radius_disc": 0.0,
+ "cyclone_usa_r64_se_radius_disc": 0.0,
+ "cyclone_usa_r64_sw_radius_disc": 0.0,
+ "cyclone_usa_rmw_disc": 0.0,
+ "cyclone_usa_wind_disc": 0.0,
+ "geopotential": 0.05989435315132141,
+ "mean_sea_level_pressure": 0.06428782641887665,
+ "sea_surface_temperature": 1.6431219577789307,
+ "specific_humidity": 0.21590931713581085,
+ "temperature": 0.11675760895013809,
+ "total_precipitation_6hr": 0.14073795080184937,
+ "u_component_of_wind": 0.12529563903808594,
+ "v_component_of_wind": 0.11094693839550018,
+ "vertical_velocity": 0.2831791341304779
+ }
+ }
+ },
+ "steps": [
+ {
+ "step": 0,
+ "init": "2026101006",
+ "loss": 1.1192289590835571,
+ "grad_norm": 3.3050899505615234,
+ "seconds": 3.985809326171875
+ },
+ {
+ "step": 1,
+ "init": "2026102406",
+ "loss": 1.0843896865844727,
+ "grad_norm": 5.549214839935303,
+ "seconds": 3.8309130668640137
+ },
+ {
+ "step": 2,
+ "init": "2026101206",
+ "loss": 1.1607990264892578,
+ "grad_norm": 4.803843975067139,
+ "seconds": 3.831921339035034
+ },
+ {
+ "step": 3,
+ "init": "2026100612",
+ "loss": 1.0067498683929443,
+ "grad_norm": 5.488168239593506,
+ "seconds": 3.830737352371216
+ },
+ {
+ "step": 4,
+ "init": "2026101900",
+ "loss": 1.0461229085922241,
+ "grad_norm": 6.701242923736572,
+ "seconds": 3.8333256244659424
+ },
+ {
+ "step": 5,
+ "init": "2026101906",
+ "loss": 1.0750306844711304,
+ "grad_norm": 4.78563117980957,
+ "seconds": 3.8339226245880127
+ },
+ {
+ "step": 6,
+ "init": "2026102300",
+ "loss": 1.1482946872711182,
+ "grad_norm": 7.357048034667969,
+ "seconds": 3.83500337600708
+ },
+ {
+ "step": 7,
+ "init": "2026100512",
+ "loss": 1.0596120357513428,
+ "grad_norm": 4.752439022064209,
+ "seconds": 3.834677219390869
+ },
+ {
+ "step": 8,
+ "init": "2026101806",
+ "loss": 1.1010043621063232,
+ "grad_norm": 4.621790409088135,
+ "seconds": 3.8346645832061768
+ },
+ {
+ "step": 9,
+ "init": "2026102200",
+ "loss": 0.9847313761711121,
+ "grad_norm": 5.491550445556641,
+ "seconds": 3.8364791870117188
+ },
+ {
+ "step": 10,
+ "init": "2026101500",
+ "loss": 1.1056504249572754,
+ "grad_norm": 5.595157146453857,
+ "seconds": 3.836566209793091
+ },
+ {
+ "step": 11,
+ "init": "2026100700",
+ "loss": 1.0602140426635742,
+ "grad_norm": 4.758091449737549,
+ "seconds": 3.8355607986450195
+ },
+ {
+ "step": 12,
+ "init": "2026100912",
+ "loss": 1.014237880706787,
+ "grad_norm": 4.617791652679443,
+ "seconds": 3.8346872329711914
+ },
+ {
+ "step": 13,
+ "init": "2026100500",
+ "loss": 1.0038481950759888,
+ "grad_norm": 4.671819686889648,
+ "seconds": 3.835641860961914
+ },
+ {
+ "step": 14,
+ "init": "2026101812",
+ "loss": 1.074454426765442,
+ "grad_norm": 5.496978759765625,
+ "seconds": 3.8377368450164795
+ },
+ {
+ "step": 15,
+ "init": "2026101612",
+ "loss": 0.965660810470581,
+ "grad_norm": 4.952380180358887,
+ "seconds": 3.8387444019317627
+ },
+ {
+ "step": 16,
+ "init": "2026101112",
+ "loss": 1.0006184577941895,
+ "grad_norm": 6.517866134643555,
+ "seconds": 3.8362984657287598
+ },
+ {
+ "step": 17,
+ "init": "2026102006",
+ "loss": 1.0640418529510498,
+ "grad_norm": 5.840461730957031,
+ "seconds": 3.8375892639160156
+ },
+ {
+ "step": 18,
+ "init": "2026100706",
+ "loss": 1.119278073310852,
+ "grad_norm": 4.642729759216309,
+ "seconds": 3.8378124237060547
+ },
+ {
+ "step": 19,
+ "init": "2026102312",
+ "loss": 0.9690228700637817,
+ "grad_norm": 5.099659442901611,
+ "seconds": 3.837205648422241
+ },
+ {
+ "step": 20,
+ "init": "2026102106",
+ "loss": 1.07974112033844,
+ "grad_norm": 3.5158801078796387,
+ "seconds": 3.8357462882995605
+ },
+ {
+ "step": 21,
+ "init": "2026101606",
+ "loss": 1.1835076808929443,
+ "grad_norm": 6.615167617797852,
+ "seconds": 3.8376126289367676
+ },
+ {
+ "step": 22,
+ "init": "2026101000",
+ "loss": 1.0275977849960327,
+ "grad_norm": 8.481230735778809,
+ "seconds": 3.837085008621216
+ },
+ {
+ "step": 23,
+ "init": "2026101600",
+ "loss": 0.9580624103546143,
+ "grad_norm": 6.061651229858398,
+ "seconds": 3.8359456062316895
+ },
+ {
+ "step": 24,
+ "init": "2026100506",
+ "loss": 1.0645805597305298,
+ "grad_norm": 7.788820266723633,
+ "seconds": 3.83670973777771
+ },
+ {
+ "step": 25,
+ "init": "2026100418",
+ "loss": 1.2330278158187866,
+ "grad_norm": 8.857479095458984,
+ "seconds": 3.837568759918213
+ },
+ {
+ "step": 26,
+ "init": "2026101412",
+ "loss": 0.9605545997619629,
+ "grad_norm": 5.47609806060791,
+ "seconds": 3.8366336822509766
+ },
+ {
+ "step": 27,
+ "init": "2026102306",
+ "loss": 1.0533347129821777,
+ "grad_norm": 5.438162326812744,
+ "seconds": 3.838200569152832
+ },
+ {
+ "step": 28,
+ "init": "2026101212",
+ "loss": 0.9402004480361938,
+ "grad_norm": 4.736769199371338,
+ "seconds": 3.837575912475586
+ },
+ {
+ "step": 29,
+ "init": "2026100918",
+ "loss": 1.1300545930862427,
+ "grad_norm": 5.309725284576416,
+ "seconds": 3.835589647293091
+ },
+ {
+ "step": 30,
+ "init": "2026101306",
+ "loss": 1.027353048324585,
+ "grad_norm": 5.323559284210205,
+ "seconds": 3.8376591205596924
+ },
+ {
+ "step": 31,
+ "init": "2026101800",
+ "loss": 0.9222906827926636,
+ "grad_norm": 5.552150726318359,
+ "seconds": 3.835721492767334
+ },
+ {
+ "step": 32,
+ "init": "2026101118",
+ "loss": 1.0264267921447754,
+ "grad_norm": 4.69077730178833,
+ "seconds": 3.8362841606140137
+ },
+ {
+ "step": 33,
+ "init": "2026100412",
+ "loss": 1.0472629070281982,
+ "grad_norm": 5.718715667724609,
+ "seconds": 3.837327003479004
+ },
+ {
+ "step": 34,
+ "init": "2026100600",
+ "loss": 1.1754722595214844,
+ "grad_norm": 13.511990547180176,
+ "seconds": 3.836961030960083
+ },
+ {
+ "step": 35,
+ "init": "2026101200",
+ "loss": 1.0436232089996338,
+ "grad_norm": 9.135140419006348,
+ "seconds": 3.836928129196167
+ },
+ {
+ "step": 36,
+ "init": "2026101300",
+ "loss": 0.8742632865905762,
+ "grad_norm": 6.3150458335876465,
+ "seconds": 3.8381576538085938
+ },
+ {
+ "step": 37,
+ "init": "2026101618",
+ "loss": 1.0396636724472046,
+ "grad_norm": 7.777122497558594,
+ "seconds": 3.8379366397857666
+ },
+ {
+ "step": 38,
+ "init": "2026102318",
+ "loss": 1.0375274419784546,
+ "grad_norm": 6.237412929534912,
+ "seconds": 3.837449312210083
+ },
+ {
+ "step": 39,
+ "init": "2026102012",
+ "loss": 0.9661576747894287,
+ "grad_norm": 9.837079048156738,
+ "seconds": 3.837475538253784
+ },
+ {
+ "step": 40,
+ "init": "2026100806",
+ "loss": 1.4604496955871582,
+ "grad_norm": 20.2987003326416,
+ "seconds": 3.8390016555786133
+ },
+ {
+ "step": 41,
+ "init": "2026101512",
+ "loss": 0.8626143932342529,
+ "grad_norm": 8.777162551879883,
+ "seconds": 3.839458465576172
+ },
+ {
+ "step": 42,
+ "init": "2026101218",
+ "loss": 1.315075397491455,
+ "grad_norm": 18.344627380371094,
+ "seconds": 3.836446523666382
+ },
+ {
+ "step": 43,
+ "init": "2026102212",
+ "loss": 1.2463264465332031,
+ "grad_norm": 16.68795394897461,
+ "seconds": 3.8377444744110107
+ },
+ {
+ "step": 44,
+ "init": "2026100618",
+ "loss": 0.9707945585250854,
+ "grad_norm": 7.846063613891602,
+ "seconds": 3.837914228439331
+ },
+ {
+ "step": 45,
+ "init": "2026101700",
+ "loss": 1.184424877166748,
+ "grad_norm": 15.520648956298828,
+ "seconds": 3.8365509510040283
+ },
+ {
+ "step": 46,
+ "init": "2026102018",
+ "loss": 1.0190258026123047,
+ "grad_norm": 7.7886962890625,
+ "seconds": 3.838561773300171
+ },
+ {
+ "step": 47,
+ "init": "2026102218",
+ "loss": 0.9902896881103516,
+ "grad_norm": 6.681366920471191,
+ "seconds": 3.8372209072113037
+ },
+ {
+ "step": 48,
+ "init": "2026102000",
+ "loss": 1.327329397201538,
+ "grad_norm": 20.562637329101562,
+ "seconds": 3.83795166015625
+ },
+ {
+ "step": 49,
+ "init": "2026101712",
+ "loss": 0.9382866621017456,
+ "grad_norm": 9.62611198425293,
+ "seconds": 3.838214635848999
+ },
+ {
+ "step": 50,
+ "init": "2026101912",
+ "loss": 1.1421862840652466,
+ "grad_norm": 14.932365417480469,
+ "seconds": 3.839059591293335
+ },
+ {
+ "step": 51,
+ "init": "2026100718",
+ "loss": 0.9972145557403564,
+ "grad_norm": 9.965677261352539,
+ "seconds": 3.8382794857025146
+ },
+ {
+ "step": 52,
+ "init": "2026100712",
+ "loss": 1.2084567546844482,
+ "grad_norm": 15.09177303314209,
+ "seconds": 3.8382749557495117
+ },
+ {
+ "step": 53,
+ "init": "2026100606",
+ "loss": 0.9696911573410034,
+ "grad_norm": 7.181429386138916,
+ "seconds": 3.840169906616211
+ },
+ {
+ "step": 54,
+ "init": "2026101818",
+ "loss": 1.2891335487365723,
+ "grad_norm": 14.908720016479492,
+ "seconds": 3.837639808654785
+ },
+ {
+ "step": 55,
+ "init": "2026101706",
+ "loss": 0.9645771980285645,
+ "grad_norm": 7.332891464233398,
+ "seconds": 3.837045907974243
+ },
+ {
+ "step": 56,
+ "init": "2026101506",
+ "loss": 1.1777693033218384,
+ "grad_norm": 10.248461723327637,
+ "seconds": 3.8372743129730225
+ },
+ {
+ "step": 57,
+ "init": "2026100518",
+ "loss": 0.9776817560195923,
+ "grad_norm": 7.030047416687012,
+ "seconds": 3.8379814624786377
+ },
+ {
+ "step": 58,
+ "init": "2026101918",
+ "loss": 0.9206633567810059,
+ "grad_norm": 7.427570343017578,
+ "seconds": 3.838114023208618
+ },
+ {
+ "step": 59,
+ "init": "2026102118",
+ "loss": 1.184816598892212,
+ "grad_norm": 12.046716690063477,
+ "seconds": 3.834653615951538
+ }
+ ],
+ "peak_gib": 21.820898294448853,
+ "median_step_s": 3.8372209072113037,
+ "eval_after": {
+ "2026102800": {
+ "loss": 1.160541296005249,
+ "per_var": {
+ "100m_u_component_of_wind": 0.14024949073791504,
+ "100m_v_component_of_wind": 0.12892509996891022,
+ "10m_u_component_of_wind": 0.14835524559020996,
+ "10m_v_component_of_wind": 0.13673800230026245,
+ "2m_temperature": 0.09061406552791595,
+ "cyclone_all_wind_disc": 0.0,
+ "cyclone_exists_gaussian_unit_mode": 0.0,
+ "cyclone_usa_pres_disc": 0.0,
+ "cyclone_usa_r34_ne_radius_disc": 0.0,
+ "cyclone_usa_r34_nw_radius_disc": 0.0,
+ "cyclone_usa_r34_se_radius_disc": 0.0,
+ "cyclone_usa_r34_sw_radius_disc": 0.0,
+ "cyclone_usa_r50_ne_radius_disc": 0.0,
+ "cyclone_usa_r50_nw_radius_disc": 0.0,
+ "cyclone_usa_r50_se_radius_disc": 0.0,
+ "cyclone_usa_r50_sw_radius_disc": 0.0,
+ "cyclone_usa_r64_ne_radius_disc": 0.0,
+ "cyclone_usa_r64_nw_radius_disc": 0.0,
+ "cyclone_usa_r64_se_radius_disc": 0.0,
+ "cyclone_usa_r64_sw_radius_disc": 0.0,
+ "cyclone_usa_rmw_disc": 0.0,
+ "cyclone_usa_wind_disc": 0.0,
+ "geopotential": 0.06798189878463745,
+ "mean_sea_level_pressure": 0.0649193525314331,
+ "sea_surface_temperature": 1.5749677419662476,
+ "specific_humidity": 0.24997997283935547,
+ "temperature": 0.13978080451488495,
+ "total_precipitation_6hr": 0.11433413624763489,
+ "u_component_of_wind": 0.14928951859474182,
+ "v_component_of_wind": 0.13590960204601288,
+ "vertical_velocity": 0.3153376281261444
+ }
+ },
+ "2026102812": {
+ "loss": 1.168492078781128,
+ "per_var": {
+ "100m_u_component_of_wind": 0.14291316270828247,
+ "100m_v_component_of_wind": 0.1328369826078415,
+ "10m_u_component_of_wind": 0.14936333894729614,
+ "10m_v_component_of_wind": 0.13860516250133514,
+ "2m_temperature": 0.09147855639457703,
+ "cyclone_all_wind_disc": 0.0,
+ "cyclone_exists_gaussian_unit_mode": 0.0,
+ "cyclone_usa_pres_disc": 0.0,
+ "cyclone_usa_r34_ne_radius_disc": 0.0,
+ "cyclone_usa_r34_nw_radius_disc": 0.0,
+ "cyclone_usa_r34_se_radius_disc": 0.0,
+ "cyclone_usa_r34_sw_radius_disc": 0.0,
+ "cyclone_usa_r50_ne_radius_disc": 0.0,
+ "cyclone_usa_r50_nw_radius_disc": 0.0,
+ "cyclone_usa_r50_se_radius_disc": 0.0,
+ "cyclone_usa_r50_sw_radius_disc": 0.0,
+ "cyclone_usa_r64_ne_radius_disc": 0.0,
+ "cyclone_usa_r64_nw_radius_disc": 0.0,
+ "cyclone_usa_r64_se_radius_disc": 0.0,
+ "cyclone_usa_r64_sw_radius_disc": 0.0,
+ "cyclone_usa_rmw_disc": 0.0,
+ "cyclone_usa_wind_disc": 0.0,
+ "geopotential": 0.06400833278894424,
+ "mean_sea_level_pressure": 0.06866370886564255,
+ "sea_surface_temperature": 1.6096842288970947,
+ "specific_humidity": 0.25394320487976074,
+ "temperature": 0.13538530468940735,
+ "total_precipitation_6hr": 0.12010172009468079,
+ "u_component_of_wind": 0.15061742067337036,
+ "v_component_of_wind": 0.13568633794784546,
+ "vertical_velocity": 0.3223147988319397
+ }
+ },
+ "2026103000": {
+ "loss": 1.1248246431350708,
+ "per_var": {
+ "100m_u_component_of_wind": 0.13557814061641693,
+ "100m_v_component_of_wind": 0.1254660040140152,
+ "10m_u_component_of_wind": 0.1437649130821228,
+ "10m_v_component_of_wind": 0.13414902985095978,
+ "2m_temperature": 0.08912350237369537,
+ "cyclone_all_wind_disc": 0.0,
+ "cyclone_exists_gaussian_unit_mode": 0.0,
+ "cyclone_usa_pres_disc": 0.0,
+ "cyclone_usa_r34_ne_radius_disc": 0.0,
+ "cyclone_usa_r34_nw_radius_disc": 0.0,
+ "cyclone_usa_r34_se_radius_disc": 0.0,
+ "cyclone_usa_r34_sw_radius_disc": 0.0,
+ "cyclone_usa_r50_ne_radius_disc": 0.0,
+ "cyclone_usa_r50_nw_radius_disc": 0.0,
+ "cyclone_usa_r50_se_radius_disc": 0.0,
+ "cyclone_usa_r50_sw_radius_disc": 0.0,
+ "cyclone_usa_r64_ne_radius_disc": 0.0,
+ "cyclone_usa_r64_nw_radius_disc": 0.0,
+ "cyclone_usa_r64_se_radius_disc": 0.0,
+ "cyclone_usa_r64_sw_radius_disc": 0.0,
+ "cyclone_usa_rmw_disc": 0.0,
+ "cyclone_usa_wind_disc": 0.0,
+ "geopotential": 0.06485021114349365,
+ "mean_sea_level_pressure": 0.06466376781463623,
+ "sea_surface_temperature": 1.538600206375122,
+ "specific_humidity": 0.24350053071975708,
+ "temperature": 0.13352540135383606,
+ "total_precipitation_6hr": 0.11074448376893997,
+ "u_component_of_wind": 0.14339186251163483,
+ "v_component_of_wind": 0.13192090392112732,
+ "vertical_velocity": 0.3064485192298889
+ }
+ },
+ "2026103012": {
+ "loss": 1.140930414199829,
+ "per_var": {
+ "100m_u_component_of_wind": 0.13896304368972778,
+ "100m_v_component_of_wind": 0.12947028875350952,
+ "10m_u_component_of_wind": 0.14556588232517242,
+ "10m_v_component_of_wind": 0.13630226254463196,
+ "2m_temperature": 0.09075110405683517,
+ "cyclone_all_wind_disc": 0.0,
+ "cyclone_exists_gaussian_unit_mode": 0.0,
+ "cyclone_usa_pres_disc": 0.0,
+ "cyclone_usa_r34_ne_radius_disc": 0.0,
+ "cyclone_usa_r34_nw_radius_disc": 0.0,
+ "cyclone_usa_r34_se_radius_disc": 0.0,
+ "cyclone_usa_r34_sw_radius_disc": 0.0,
+ "cyclone_usa_r50_ne_radius_disc": 0.0,
+ "cyclone_usa_r50_nw_radius_disc": 0.0,
+ "cyclone_usa_r50_se_radius_disc": 0.0,
+ "cyclone_usa_r50_sw_radius_disc": 0.0,
+ "cyclone_usa_r64_ne_radius_disc": 0.0,
+ "cyclone_usa_r64_nw_radius_disc": 0.0,
+ "cyclone_usa_r64_se_radius_disc": 0.0,
+ "cyclone_usa_r64_sw_radius_disc": 0.0,
+ "cyclone_usa_rmw_disc": 0.0,
+ "cyclone_usa_wind_disc": 0.0,
+ "geopotential": 0.062267035245895386,
+ "mean_sea_level_pressure": 0.06503040343523026,
+ "sea_surface_temperature": 1.6432154178619385,
+ "specific_humidity": 0.24740362167358398,
+ "temperature": 0.132475346326828,
+ "total_precipitation_6hr": 0.12259252369403839,
+ "u_component_of_wind": 0.14381173253059387,
+ "v_component_of_wind": 0.13080573081970215,
+ "vertical_velocity": 0.3182721734046936
+ }
+ }
+ }
+}
\ No newline at end of file
diff --git a/docs/results/finetune/finetune_masksst_rtx5090.log b/docs/results/finetune/finetune_masksst_rtx5090.log
new file mode 100644
index 0000000..64bb6cd
--- /dev/null
+++ b/docs/results/finetune/finetune_masksst_rtx5090.log
@@ -0,0 +1,98 @@
+example 2026100412 built in 2.2s
+WeatherNext2 ckpt 1: training 5.57M of 183.8M params; masked targets: 18 (cyclone channels, sea_surface_temperature)
+WARNING:root:No normalization location found for cyclone_exists_gaussian_unit_mode
+WARNING:root:No normalization scale found for cyclone_exists_gaussian_unit_mode
+WARNING:absl:No coordinates found for dim noise_channels of array Size: 256B
+JitTracer(float32[2,1,32])
+Coordinates:
+ * batch (batch) object 16B MultiIndex
+ * _batch (batch) int64 16B 0 0
+ * sample (batch) int64 16B 0 1
+Dimensions without coordinates: channels, noise_channels., using index for param name. This may make model surgery more error prone if changing coordinates at fine-tuning time.
+/home/raymxnd/Workspace/WeatherForcast/.venv/lib/python3.12/site-packages/scipy/sparse/_index.py:216: SparseEfficiencyWarning: Changing the sparsity structure of a csr_matrix is expensive. lil and dok are more efficient.
+ self._set_arrayXarray(i, j, x)
+/home/raymxnd/Workspace/WeatherForcast/weathernext/weathernext/weathernext2/architecture.py:286: FutureWarning: Deleting a single level of a MultiIndex is deprecated. Previously, this deleted all levels of a MultiIndex. Please also drop the following variables: {'batch'} to avoid an error in the future.
+ predictions[var] = xarray_jax.apply_ufunc(activation_fn, predictions[var])
+W1031 13:16:55.643671 1640757 hlo_rematerialization.cc:3363] Can't reduce memory use below 25.14GiB (26997919705 bytes) by rematerialization; only reduced to 63.07GiB (67717567744 bytes), down from 63.07GiB (67717567744 bytes) originally
+train step compiled in 85s: temp 20.03 GiB, arguments 1.79 GiB
+WARNING:root:No normalization location found for cyclone_exists_gaussian_unit_mode
+WARNING:root:No normalization scale found for cyclone_exists_gaussian_unit_mode
+WARNING:absl:No coordinates found for dim noise_channels of array Size: 256B
+JitTracer(float32[2,1,32])
+Coordinates:
+ * batch (batch) object 16B MultiIndex
+ * _batch (batch) int64 16B 0 0
+ * sample (batch) int64 16B 0 1
+Dimensions without coordinates: channels, noise_channels., using index for param name. This may make model surgery more error prone if changing coordinates at fine-tuning time.
+/home/raymxnd/Workspace/WeatherForcast/.venv/lib/python3.12/site-packages/scipy/sparse/_index.py:216: SparseEfficiencyWarning: Changing the sparsity structure of a csr_matrix is expensive. lil and dok are more efficient.
+ self._set_arrayXarray(i, j, x)
+/home/raymxnd/Workspace/WeatherForcast/weathernext/weathernext/weathernext2/architecture.py:286: FutureWarning: Deleting a single level of a MultiIndex is deprecated. Previously, this deleted all levels of a MultiIndex. Please also drop the following variables: {'batch'} to avoid an error in the future.
+ predictions[var] = xarray_jax.apply_ufunc(activation_fn, predictions[var])
+W1031 13:18:05.723884 1640757 hlo_rematerialization.cc:3363] Can't reduce memory use below 25.12GiB (26978079655 bytes) by rematerialization; only reduced to 56.83GiB (61019999108 bytes), down from 56.83GiB (61020035972 bytes) originally
+eval before (77s): 2026102800 1.04093 2026102812 1.05581 2026103000 1.00725 2026103012 1.02891
+step 0 2026101006 loss 1.11923 |g| 3.305e+00 3.99s
+step 1 2026102406 loss 1.08439 |g| 5.549e+00 3.83s
+step 2 2026101206 loss 1.16080 |g| 4.804e+00 3.83s
+step 3 2026100612 loss 1.00675 |g| 5.488e+00 3.83s
+step 4 2026101900 loss 1.04612 |g| 6.701e+00 3.83s
+step 5 2026101906 loss 1.07503 |g| 4.786e+00 3.83s
+step 6 2026102300 loss 1.14829 |g| 7.357e+00 3.84s
+step 7 2026100512 loss 1.05961 |g| 4.752e+00 3.83s
+step 8 2026101806 loss 1.10100 |g| 4.622e+00 3.83s
+step 9 2026102200 loss 0.98473 |g| 5.492e+00 3.84s
+step 10 2026101500 loss 1.10565 |g| 5.595e+00 3.84s
+step 11 2026100700 loss 1.06021 |g| 4.758e+00 3.84s
+step 12 2026100912 loss 1.01424 |g| 4.618e+00 3.83s
+step 13 2026100500 loss 1.00385 |g| 4.672e+00 3.84s
+step 14 2026101812 loss 1.07445 |g| 5.497e+00 3.84s
+step 15 2026101612 loss 0.96566 |g| 4.952e+00 3.84s
+step 16 2026101112 loss 1.00062 |g| 6.518e+00 3.84s
+step 17 2026102006 loss 1.06404 |g| 5.840e+00 3.84s
+step 18 2026100706 loss 1.11928 |g| 4.643e+00 3.84s
+step 19 2026102312 loss 0.96902 |g| 5.100e+00 3.84s
+step 20 2026102106 loss 1.07974 |g| 3.516e+00 3.84s
+step 21 2026101606 loss 1.18351 |g| 6.615e+00 3.84s
+step 22 2026101000 loss 1.02760 |g| 8.481e+00 3.84s
+step 23 2026101600 loss 0.95806 |g| 6.062e+00 3.84s
+step 24 2026100506 loss 1.06458 |g| 7.789e+00 3.84s
+step 25 2026100418 loss 1.23303 |g| 8.857e+00 3.84s
+step 26 2026101412 loss 0.96055 |g| 5.476e+00 3.84s
+step 27 2026102306 loss 1.05333 |g| 5.438e+00 3.84s
+step 28 2026101212 loss 0.94020 |g| 4.737e+00 3.84s
+step 29 2026100918 loss 1.13005 |g| 5.310e+00 3.84s
+step 30 2026101306 loss 1.02735 |g| 5.324e+00 3.84s
+step 31 2026101800 loss 0.92229 |g| 5.552e+00 3.84s
+step 32 2026101118 loss 1.02643 |g| 4.691e+00 3.84s
+step 33 2026100412 loss 1.04726 |g| 5.719e+00 3.84s
+step 34 2026100600 loss 1.17547 |g| 1.351e+01 3.84s
+step 35 2026101200 loss 1.04362 |g| 9.135e+00 3.84s
+step 36 2026101300 loss 0.87426 |g| 6.315e+00 3.84s
+step 37 2026101618 loss 1.03966 |g| 7.777e+00 3.84s
+step 38 2026102318 loss 1.03753 |g| 6.237e+00 3.84s
+step 39 2026102012 loss 0.96616 |g| 9.837e+00 3.84s
+step 40 2026100806 loss 1.46045 |g| 2.030e+01 3.84s
+step 41 2026101512 loss 0.86261 |g| 8.777e+00 3.84s
+step 42 2026101218 loss 1.31508 |g| 1.834e+01 3.84s
+step 43 2026102212 loss 1.24633 |g| 1.669e+01 3.84s
+step 44 2026100618 loss 0.97079 |g| 7.846e+00 3.84s
+step 45 2026101700 loss 1.18442 |g| 1.552e+01 3.84s
+step 46 2026102018 loss 1.01903 |g| 7.789e+00 3.84s
+step 47 2026102218 loss 0.99029 |g| 6.681e+00 3.84s
+step 48 2026102000 loss 1.32733 |g| 2.056e+01 3.84s
+step 49 2026101712 loss 0.93829 |g| 9.626e+00 3.84s
+step 50 2026101912 loss 1.14219 |g| 1.493e+01 3.84s
+step 51 2026100718 loss 0.99721 |g| 9.966e+00 3.84s
+step 52 2026100712 loss 1.20846 |g| 1.509e+01 3.84s
+step 53 2026100606 loss 0.96969 |g| 7.181e+00 3.84s
+step 54 2026101818 loss 1.28913 |g| 1.491e+01 3.84s
+step 55 2026101706 loss 0.96458 |g| 7.333e+00 3.84s
+step 56 2026101506 loss 1.17777 |g| 1.025e+01 3.84s
+step 57 2026100518 loss 0.97768 |g| 7.030e+00 3.84s
+step 58 2026101918 loss 0.92066 |g| 7.428e+00 3.84s
+step 59 2026102118 loss 1.18482 |g| 1.205e+01 3.83s
+median step 3.84s, device peak 21.82 GiB
+eval 2026102800: 1.04093 -> 1.16054 (+11.49 %)
+eval 2026102812: 1.05581 -> 1.16849 (+10.67 %)
+eval 2026103000: 1.00725 -> 1.12482 (+11.67 %)
+eval 2026103012: 1.02891 -> 1.14093 (+10.89 %)
+wrote docs/results/finetune/finetune_masksst_rtx5090.json
diff --git a/docs/results/finetune/finetune_smoke_rtx5090.json b/docs/results/finetune/finetune_smoke_rtx5090.json
new file mode 100644
index 0000000..7ba3e49
--- /dev/null
+++ b/docs/results/finetune/finetune_smoke_rtx5090.json
@@ -0,0 +1,843 @@
+{
+ "model": "WeatherNext2",
+ "checkpoint": 1,
+ "device": "NVIDIA GeForce RTX 5090",
+ "jax": "0.11.2",
+ "faster_weathernext": "0.1.1",
+ "args": {
+ "model": "WeatherNext2",
+ "checkpoint": 1,
+ "train": [
+ "2026100412",
+ "2026100418",
+ "2026100500",
+ "2026100506",
+ "2026100512",
+ "2026100518",
+ "2026100600",
+ "2026100606",
+ "2026100612",
+ "2026100618",
+ "2026100700",
+ "2026100706",
+ "2026100712",
+ "2026100718",
+ "2026100800",
+ "2026100806",
+ "2026100912",
+ "2026100918",
+ "2026101000",
+ "2026101006",
+ "2026101112",
+ "2026101118",
+ "2026101200",
+ "2026101206",
+ "2026101212",
+ "2026101218",
+ "2026101300",
+ "2026101306",
+ "2026101412",
+ "2026101418",
+ "2026101500",
+ "2026101506",
+ "2026101512",
+ "2026101518",
+ "2026101600",
+ "2026101606",
+ "2026101612",
+ "2026101618",
+ "2026101700",
+ "2026101706",
+ "2026101712",
+ "2026101718",
+ "2026101800",
+ "2026101806",
+ "2026101812",
+ "2026101818",
+ "2026101900",
+ "2026101906",
+ "2026101912",
+ "2026101918",
+ "2026102000",
+ "2026102006",
+ "2026102012",
+ "2026102018",
+ "2026102100",
+ "2026102106",
+ "2026102112",
+ "2026102118",
+ "2026102200",
+ "2026102206",
+ "2026102212",
+ "2026102218",
+ "2026102300",
+ "2026102306",
+ "2026102312",
+ "2026102318",
+ "2026102400",
+ "2026102406"
+ ],
+ "eval": [
+ "2026102800",
+ "2026102812",
+ "2026103000",
+ "2026103012"
+ ],
+ "steps": 60,
+ "lr": 3e-05,
+ "weight_decay": 0.0,
+ "clip": 1.0,
+ "eval_keys": 2,
+ "precision": "default",
+ "seed": 0,
+ "no_precip": false,
+ "compile_only": false,
+ "no_remat": false,
+ "weights_dir": "/home/raymxnd/Workspace/WeatherForcast/weights",
+ "out": "/tmp/claude-1000/-home-raymxnd-Workspace-WeatherForcast/d4a93a6d-df87-4c4c-b7ce-cc8b5b5265d2/scratchpad/head_smoke.npz",
+ "log": "docs/results/finetune/finetune_smoke_rtx5090.json"
+ },
+ "masked": [
+ "cyclone_exists_gaussian_unit_mode",
+ "cyclone_all_wind_disc",
+ "cyclone_usa_wind_disc",
+ "cyclone_usa_r34_ne_radius_disc",
+ "cyclone_usa_r34_se_radius_disc",
+ "cyclone_usa_r34_sw_radius_disc",
+ "cyclone_usa_r34_nw_radius_disc",
+ "cyclone_usa_r50_ne_radius_disc",
+ "cyclone_usa_r50_se_radius_disc",
+ "cyclone_usa_r50_sw_radius_disc",
+ "cyclone_usa_r50_nw_radius_disc",
+ "cyclone_usa_r64_ne_radius_disc",
+ "cyclone_usa_r64_se_radius_disc",
+ "cyclone_usa_r64_sw_radius_disc",
+ "cyclone_usa_r64_nw_radius_disc",
+ "cyclone_usa_rmw_disc",
+ "cyclone_usa_pres_disc"
+ ],
+ "head_params": 5572967,
+ "total_params": 183785575,
+ "compile": {
+ "seconds": 85.62767791748047,
+ "temp_gib": 20.03102158755064,
+ "argument_gib": 1.7898386642336845,
+ "output_gib": 0.062285151332616806
+ },
+ "eval_before": {
+ "2026102800": {
+ "loss": 1.199784517288208,
+ "per_var": {
+ "100m_u_component_of_wind": 0.13943889737129211,
+ "100m_v_component_of_wind": 0.1279805600643158,
+ "10m_u_component_of_wind": 0.1485271155834198,
+ "10m_v_component_of_wind": 0.1363852322101593,
+ "2m_temperature": 0.09907478094100952,
+ "cyclone_all_wind_disc": 0.0,
+ "cyclone_exists_gaussian_unit_mode": 0.0,
+ "cyclone_usa_pres_disc": 0.0,
+ "cyclone_usa_r34_ne_radius_disc": 0.0,
+ "cyclone_usa_r34_nw_radius_disc": 0.0,
+ "cyclone_usa_r34_se_radius_disc": 0.0,
+ "cyclone_usa_r34_sw_radius_disc": 0.0,
+ "cyclone_usa_r50_ne_radius_disc": 0.0,
+ "cyclone_usa_r50_nw_radius_disc": 0.0,
+ "cyclone_usa_r50_se_radius_disc": 0.0,
+ "cyclone_usa_r50_sw_radius_disc": 0.0,
+ "cyclone_usa_r64_ne_radius_disc": 0.0,
+ "cyclone_usa_r64_nw_radius_disc": 0.0,
+ "cyclone_usa_r64_se_radius_disc": 0.0,
+ "cyclone_usa_r64_sw_radius_disc": 0.0,
+ "cyclone_usa_rmw_disc": 0.0,
+ "cyclone_usa_wind_disc": 0.0,
+ "geopotential": 0.0621756836771965,
+ "mean_sea_level_pressure": 0.06353693455457687,
+ "sea_surface_temperature": 1.5886144638061523,
+ "specific_humidity": 0.21853001415729523,
+ "temperature": 0.12121587991714478,
+ "total_precipitation_6hr": 0.12706361711025238,
+ "u_component_of_wind": 0.13052740693092346,
+ "v_component_of_wind": 0.1155254989862442,
+ "vertical_velocity": 0.27828487753868103
+ }
+ },
+ "2026102812": {
+ "loss": 1.2168025970458984,
+ "per_var": {
+ "100m_u_component_of_wind": 0.14223329722881317,
+ "100m_v_component_of_wind": 0.13159146904945374,
+ "10m_u_component_of_wind": 0.14976154267787933,
+ "10m_v_component_of_wind": 0.13810990750789642,
+ "2m_temperature": 0.10236281156539917,
+ "cyclone_all_wind_disc": 0.0,
+ "cyclone_exists_gaussian_unit_mode": 0.0,
+ "cyclone_usa_pres_disc": 0.0,
+ "cyclone_usa_r34_ne_radius_disc": 0.0,
+ "cyclone_usa_r34_nw_radius_disc": 0.0,
+ "cyclone_usa_r34_se_radius_disc": 0.0,
+ "cyclone_usa_r34_sw_radius_disc": 0.0,
+ "cyclone_usa_r50_ne_radius_disc": 0.0,
+ "cyclone_usa_r50_nw_radius_disc": 0.0,
+ "cyclone_usa_r50_se_radius_disc": 0.0,
+ "cyclone_usa_r50_sw_radius_disc": 0.0,
+ "cyclone_usa_r64_ne_radius_disc": 0.0,
+ "cyclone_usa_r64_nw_radius_disc": 0.0,
+ "cyclone_usa_r64_se_radius_disc": 0.0,
+ "cyclone_usa_r64_sw_radius_disc": 0.0,
+ "cyclone_usa_rmw_disc": 0.0,
+ "cyclone_usa_wind_disc": 0.0,
+ "geopotential": 0.059474069625139236,
+ "mean_sea_level_pressure": 0.06730681657791138,
+ "sea_surface_temperature": 1.6101062297821045,
+ "specific_humidity": 0.22227859497070312,
+ "temperature": 0.12098146229982376,
+ "total_precipitation_6hr": 0.13528618216514587,
+ "u_component_of_wind": 0.1308210790157318,
+ "v_component_of_wind": 0.11492675542831421,
+ "vertical_velocity": 0.28633999824523926
+ }
+ },
+ "2026103000": {
+ "loss": 1.163001298904419,
+ "per_var": {
+ "100m_u_component_of_wind": 0.13535326719284058,
+ "100m_v_component_of_wind": 0.12463682144880295,
+ "10m_u_component_of_wind": 0.14480635523796082,
+ "10m_v_component_of_wind": 0.1337897628545761,
+ "2m_temperature": 0.09720341861248016,
+ "cyclone_all_wind_disc": 0.0,
+ "cyclone_exists_gaussian_unit_mode": 0.0,
+ "cyclone_usa_pres_disc": 0.0,
+ "cyclone_usa_r34_ne_radius_disc": 0.0,
+ "cyclone_usa_r34_nw_radius_disc": 0.0,
+ "cyclone_usa_r34_se_radius_disc": 0.0,
+ "cyclone_usa_r34_sw_radius_disc": 0.0,
+ "cyclone_usa_r50_ne_radius_disc": 0.0,
+ "cyclone_usa_r50_nw_radius_disc": 0.0,
+ "cyclone_usa_r50_se_radius_disc": 0.0,
+ "cyclone_usa_r50_sw_radius_disc": 0.0,
+ "cyclone_usa_r64_ne_radius_disc": 0.0,
+ "cyclone_usa_r64_nw_radius_disc": 0.0,
+ "cyclone_usa_r64_se_radius_disc": 0.0,
+ "cyclone_usa_r64_sw_radius_disc": 0.0,
+ "cyclone_usa_rmw_disc": 0.0,
+ "cyclone_usa_wind_disc": 0.0,
+ "geopotential": 0.06037704274058342,
+ "mean_sea_level_pressure": 0.06656266748905182,
+ "sea_surface_temperature": 1.5575823783874512,
+ "specific_humidity": 0.2134988009929657,
+ "temperature": 0.11542285978794098,
+ "total_precipitation_6hr": 0.12492990493774414,
+ "u_component_of_wind": 0.1250942051410675,
+ "v_component_of_wind": 0.11082634329795837,
+ "vertical_velocity": 0.26893025636672974
+ }
+ },
+ "2026103012": {
+ "loss": 1.1932249069213867,
+ "per_var": {
+ "100m_u_component_of_wind": 0.1389583945274353,
+ "100m_v_component_of_wind": 0.12912210822105408,
+ "10m_u_component_of_wind": 0.14677539467811584,
+ "10m_v_component_of_wind": 0.13630259037017822,
+ "2m_temperature": 0.09881369024515152,
+ "cyclone_all_wind_disc": 0.0,
+ "cyclone_exists_gaussian_unit_mode": 0.0,
+ "cyclone_usa_pres_disc": 0.0,
+ "cyclone_usa_r34_ne_radius_disc": 0.0,
+ "cyclone_usa_r34_nw_radius_disc": 0.0,
+ "cyclone_usa_r34_se_radius_disc": 0.0,
+ "cyclone_usa_r34_sw_radius_disc": 0.0,
+ "cyclone_usa_r50_ne_radius_disc": 0.0,
+ "cyclone_usa_r50_nw_radius_disc": 0.0,
+ "cyclone_usa_r50_se_radius_disc": 0.0,
+ "cyclone_usa_r50_sw_radius_disc": 0.0,
+ "cyclone_usa_r64_ne_radius_disc": 0.0,
+ "cyclone_usa_r64_nw_radius_disc": 0.0,
+ "cyclone_usa_r64_se_radius_disc": 0.0,
+ "cyclone_usa_r64_sw_radius_disc": 0.0,
+ "cyclone_usa_rmw_disc": 0.0,
+ "cyclone_usa_wind_disc": 0.0,
+ "geopotential": 0.0598960742354393,
+ "mean_sea_level_pressure": 0.06429050117731094,
+ "sea_surface_temperature": 1.6431231498718262,
+ "specific_humidity": 0.21590930223464966,
+ "temperature": 0.11675658822059631,
+ "total_precipitation_6hr": 0.14074113965034485,
+ "u_component_of_wind": 0.12529529631137848,
+ "v_component_of_wind": 0.11094654351472855,
+ "vertical_velocity": 0.28318196535110474
+ }
+ }
+ },
+ "steps": [
+ {
+ "step": 0,
+ "init": "2026101006",
+ "loss": 1.2540011405944824,
+ "grad_norm": 3.307534694671631,
+ "seconds": 3.9916579723358154
+ },
+ {
+ "step": 1,
+ "init": "2026102406",
+ "loss": 1.2528198957443237,
+ "grad_norm": 5.553346157073975,
+ "seconds": 3.835658550262451
+ },
+ {
+ "step": 2,
+ "init": "2026101206",
+ "loss": 1.2906982898712158,
+ "grad_norm": 4.807368278503418,
+ "seconds": 3.8387439250946045
+ },
+ {
+ "step": 3,
+ "init": "2026100612",
+ "loss": 1.1421146392822266,
+ "grad_norm": 5.481176376342773,
+ "seconds": 3.838911533355713
+ },
+ {
+ "step": 4,
+ "init": "2026101900",
+ "loss": 1.1787699460983276,
+ "grad_norm": 7.076934337615967,
+ "seconds": 3.8412792682647705
+ },
+ {
+ "step": 5,
+ "init": "2026101906",
+ "loss": 1.1965928077697754,
+ "grad_norm": 4.851121425628662,
+ "seconds": 3.8408048152923584
+ },
+ {
+ "step": 6,
+ "init": "2026102300",
+ "loss": 1.2888104915618896,
+ "grad_norm": 7.435774803161621,
+ "seconds": 3.841110944747925
+ },
+ {
+ "step": 7,
+ "init": "2026100512",
+ "loss": 1.1931320428848267,
+ "grad_norm": 4.797940731048584,
+ "seconds": 3.840949773788452
+ },
+ {
+ "step": 8,
+ "init": "2026101806",
+ "loss": 1.220227599143982,
+ "grad_norm": 4.656756401062012,
+ "seconds": 3.8431828022003174
+ },
+ {
+ "step": 9,
+ "init": "2026102200",
+ "loss": 1.1203042268753052,
+ "grad_norm": 5.656060218811035,
+ "seconds": 3.842087984085083
+ },
+ {
+ "step": 10,
+ "init": "2026101500",
+ "loss": 1.2381117343902588,
+ "grad_norm": 5.655872344970703,
+ "seconds": 3.842684268951416
+ },
+ {
+ "step": 11,
+ "init": "2026100700",
+ "loss": 1.2100061178207397,
+ "grad_norm": 5.068142414093018,
+ "seconds": 3.8414623737335205
+ },
+ {
+ "step": 12,
+ "init": "2026100912",
+ "loss": 1.1400065422058105,
+ "grad_norm": 4.636743545532227,
+ "seconds": 3.841668128967285
+ },
+ {
+ "step": 13,
+ "init": "2026100500",
+ "loss": 1.1386815309524536,
+ "grad_norm": 4.768000602722168,
+ "seconds": 3.841740846633911
+ },
+ {
+ "step": 14,
+ "init": "2026101812",
+ "loss": 1.1979527473449707,
+ "grad_norm": 5.5894622802734375,
+ "seconds": 3.84316086769104
+ },
+ {
+ "step": 15,
+ "init": "2026101612",
+ "loss": 1.091132402420044,
+ "grad_norm": 4.96360445022583,
+ "seconds": 3.841217041015625
+ },
+ {
+ "step": 16,
+ "init": "2026101112",
+ "loss": 1.1214052438735962,
+ "grad_norm": 6.539325714111328,
+ "seconds": 3.8457374572753906
+ },
+ {
+ "step": 17,
+ "init": "2026102006",
+ "loss": 1.2042696475982666,
+ "grad_norm": 5.8669281005859375,
+ "seconds": 3.8424293994903564
+ },
+ {
+ "step": 18,
+ "init": "2026100706",
+ "loss": 1.251413106918335,
+ "grad_norm": 4.718742847442627,
+ "seconds": 3.8434340953826904
+ },
+ {
+ "step": 19,
+ "init": "2026102312",
+ "loss": 1.098642110824585,
+ "grad_norm": 5.150508880615234,
+ "seconds": 3.843902111053467
+ },
+ {
+ "step": 20,
+ "init": "2026102106",
+ "loss": 1.2313324213027954,
+ "grad_norm": 3.6204631328582764,
+ "seconds": 3.84494686126709
+ },
+ {
+ "step": 21,
+ "init": "2026101606",
+ "loss": 1.3077311515808105,
+ "grad_norm": 6.663534641265869,
+ "seconds": 3.844694137573242
+ },
+ {
+ "step": 22,
+ "init": "2026101000",
+ "loss": 1.1597702503204346,
+ "grad_norm": 8.556142807006836,
+ "seconds": 3.8438608646392822
+ },
+ {
+ "step": 23,
+ "init": "2026101600",
+ "loss": 1.0857586860656738,
+ "grad_norm": 6.318863868713379,
+ "seconds": 3.8430252075195312
+ },
+ {
+ "step": 24,
+ "init": "2026100506",
+ "loss": 1.1833932399749756,
+ "grad_norm": 7.858797073364258,
+ "seconds": 3.8436009883880615
+ },
+ {
+ "step": 25,
+ "init": "2026100418",
+ "loss": 1.3728383779525757,
+ "grad_norm": 8.900500297546387,
+ "seconds": 3.8415753841400146
+ },
+ {
+ "step": 26,
+ "init": "2026101412",
+ "loss": 1.0840479135513306,
+ "grad_norm": 5.338094234466553,
+ "seconds": 3.8430850505828857
+ },
+ {
+ "step": 27,
+ "init": "2026102306",
+ "loss": 1.1958286762237549,
+ "grad_norm": 5.533114910125732,
+ "seconds": 3.8434629440307617
+ },
+ {
+ "step": 28,
+ "init": "2026101212",
+ "loss": 1.0530054569244385,
+ "grad_norm": 4.882535934448242,
+ "seconds": 3.8447678089141846
+ },
+ {
+ "step": 29,
+ "init": "2026100918",
+ "loss": 1.259422779083252,
+ "grad_norm": 5.659214019775391,
+ "seconds": 3.845486640930176
+ },
+ {
+ "step": 30,
+ "init": "2026101306",
+ "loss": 1.142958641052246,
+ "grad_norm": 5.379838466644287,
+ "seconds": 3.843970537185669
+ },
+ {
+ "step": 31,
+ "init": "2026101800",
+ "loss": 1.0334651470184326,
+ "grad_norm": 5.858162879943848,
+ "seconds": 3.845567226409912
+ },
+ {
+ "step": 32,
+ "init": "2026101118",
+ "loss": 1.1819490194320679,
+ "grad_norm": 4.815415382385254,
+ "seconds": 3.8434898853302
+ },
+ {
+ "step": 33,
+ "init": "2026100412",
+ "loss": 1.1697465181350708,
+ "grad_norm": 5.7807393074035645,
+ "seconds": 3.8459458351135254
+ },
+ {
+ "step": 34,
+ "init": "2026100600",
+ "loss": 1.3095823526382446,
+ "grad_norm": 13.68796157836914,
+ "seconds": 3.846566677093506
+ },
+ {
+ "step": 35,
+ "init": "2026101200",
+ "loss": 1.1649742126464844,
+ "grad_norm": 9.568260192871094,
+ "seconds": 3.843423366546631
+ },
+ {
+ "step": 36,
+ "init": "2026101300",
+ "loss": 0.996849536895752,
+ "grad_norm": 6.9232563972473145,
+ "seconds": 3.844700574874878
+ },
+ {
+ "step": 37,
+ "init": "2026101618",
+ "loss": 1.2032718658447266,
+ "grad_norm": 8.203710556030273,
+ "seconds": 3.8442206382751465
+ },
+ {
+ "step": 38,
+ "init": "2026102318",
+ "loss": 1.1737384796142578,
+ "grad_norm": 6.496024131774902,
+ "seconds": 3.846515655517578
+ },
+ {
+ "step": 39,
+ "init": "2026102012",
+ "loss": 1.0846221446990967,
+ "grad_norm": 9.564035415649414,
+ "seconds": 3.8442516326904297
+ },
+ {
+ "step": 40,
+ "init": "2026100806",
+ "loss": 1.5846751928329468,
+ "grad_norm": 20.830310821533203,
+ "seconds": 3.8430240154266357
+ },
+ {
+ "step": 41,
+ "init": "2026101512",
+ "loss": 0.9770716428756714,
+ "grad_norm": 8.755279541015625,
+ "seconds": 3.8429415225982666
+ },
+ {
+ "step": 42,
+ "init": "2026101218",
+ "loss": 1.4300265312194824,
+ "grad_norm": 18.510704040527344,
+ "seconds": 3.8459880352020264
+ },
+ {
+ "step": 43,
+ "init": "2026102212",
+ "loss": 1.367002248764038,
+ "grad_norm": 16.45145034790039,
+ "seconds": 3.8435235023498535
+ },
+ {
+ "step": 44,
+ "init": "2026100618",
+ "loss": 1.1192086935043335,
+ "grad_norm": 8.375672340393066,
+ "seconds": 3.8441457748413086
+ },
+ {
+ "step": 45,
+ "init": "2026101700",
+ "loss": 1.2902436256408691,
+ "grad_norm": 15.637262344360352,
+ "seconds": 3.8445262908935547
+ },
+ {
+ "step": 46,
+ "init": "2026102018",
+ "loss": 1.1420459747314453,
+ "grad_norm": 8.363740921020508,
+ "seconds": 3.8457651138305664
+ },
+ {
+ "step": 47,
+ "init": "2026102218",
+ "loss": 1.1178910732269287,
+ "grad_norm": 7.237462997436523,
+ "seconds": 3.843132972717285
+ },
+ {
+ "step": 48,
+ "init": "2026102000",
+ "loss": 1.4383147954940796,
+ "grad_norm": 20.650054931640625,
+ "seconds": 3.8438050746917725
+ },
+ {
+ "step": 49,
+ "init": "2026101712",
+ "loss": 1.0471484661102295,
+ "grad_norm": 9.736106872558594,
+ "seconds": 3.8448727130889893
+ },
+ {
+ "step": 50,
+ "init": "2026101912",
+ "loss": 1.236140489578247,
+ "grad_norm": 14.224828720092773,
+ "seconds": 3.8442611694335938
+ },
+ {
+ "step": 51,
+ "init": "2026100718",
+ "loss": 1.1087342500686646,
+ "grad_norm": 10.197408676147461,
+ "seconds": 3.843858003616333
+ },
+ {
+ "step": 52,
+ "init": "2026100712",
+ "loss": 1.321321725845337,
+ "grad_norm": 14.604911804199219,
+ "seconds": 3.8441758155822754
+ },
+ {
+ "step": 53,
+ "init": "2026100606",
+ "loss": 1.0760518312454224,
+ "grad_norm": 7.338120460510254,
+ "seconds": 3.8448946475982666
+ },
+ {
+ "step": 54,
+ "init": "2026101818",
+ "loss": 1.3763049840927124,
+ "grad_norm": 14.582459449768066,
+ "seconds": 3.8428680896759033
+ },
+ {
+ "step": 55,
+ "init": "2026101706",
+ "loss": 1.0634742975234985,
+ "grad_norm": 7.361107349395752,
+ "seconds": 3.84375262260437
+ },
+ {
+ "step": 56,
+ "init": "2026101506",
+ "loss": 1.2820205688476562,
+ "grad_norm": 10.158431053161621,
+ "seconds": 3.8456127643585205
+ },
+ {
+ "step": 57,
+ "init": "2026100518",
+ "loss": 1.0959548950195312,
+ "grad_norm": 7.1181440353393555,
+ "seconds": 3.8461825847625732
+ },
+ {
+ "step": 58,
+ "init": "2026101918",
+ "loss": 1.0273889303207397,
+ "grad_norm": 7.616166114807129,
+ "seconds": 3.8438289165496826
+ },
+ {
+ "step": 59,
+ "init": "2026102118",
+ "loss": 1.312725305557251,
+ "grad_norm": 12.041414260864258,
+ "seconds": 3.839787244796753
+ }
+ ],
+ "peak_gib": 21.82089877128601,
+ "median_step_s": 3.8436009883880615,
+ "eval_after": {
+ "2026102800": {
+ "loss": 1.2839086055755615,
+ "per_var": {
+ "100m_u_component_of_wind": 0.1400538682937622,
+ "100m_v_component_of_wind": 0.12908464670181274,
+ "10m_u_component_of_wind": 0.14807644486427307,
+ "10m_v_component_of_wind": 0.13678115606307983,
+ "2m_temperature": 0.09090624749660492,
+ "cyclone_all_wind_disc": 0.0,
+ "cyclone_exists_gaussian_unit_mode": 0.0,
+ "cyclone_usa_pres_disc": 0.0,
+ "cyclone_usa_r34_ne_radius_disc": 0.0,
+ "cyclone_usa_r34_nw_radius_disc": 0.0,
+ "cyclone_usa_r34_se_radius_disc": 0.0,
+ "cyclone_usa_r34_sw_radius_disc": 0.0,
+ "cyclone_usa_r50_ne_radius_disc": 0.0,
+ "cyclone_usa_r50_nw_radius_disc": 0.0,
+ "cyclone_usa_r50_se_radius_disc": 0.0,
+ "cyclone_usa_r50_sw_radius_disc": 0.0,
+ "cyclone_usa_r64_ne_radius_disc": 0.0,
+ "cyclone_usa_r64_nw_radius_disc": 0.0,
+ "cyclone_usa_r64_se_radius_disc": 0.0,
+ "cyclone_usa_r64_sw_radius_disc": 0.0,
+ "cyclone_usa_rmw_disc": 0.0,
+ "cyclone_usa_wind_disc": 0.0,
+ "geopotential": 0.06772905588150024,
+ "mean_sea_level_pressure": 0.06537915766239166,
+ "sea_surface_temperature": 1.322561264038086,
+ "specific_humidity": 0.24895167350769043,
+ "temperature": 0.13783489167690277,
+ "total_precipitation_6hr": 0.11464792490005493,
+ "u_component_of_wind": 0.14825189113616943,
+ "v_component_of_wind": 0.13493132591247559,
+ "vertical_velocity": 0.3112092912197113
+ }
+ },
+ "2026102812": {
+ "loss": 1.2979035377502441,
+ "per_var": {
+ "100m_u_component_of_wind": 0.14278148114681244,
+ "100m_v_component_of_wind": 0.13293878734111786,
+ "10m_u_component_of_wind": 0.14954328536987305,
+ "10m_v_component_of_wind": 0.13885104656219482,
+ "2m_temperature": 0.09137065708637238,
+ "cyclone_all_wind_disc": 0.0,
+ "cyclone_exists_gaussian_unit_mode": 0.0,
+ "cyclone_usa_pres_disc": 0.0,
+ "cyclone_usa_r34_ne_radius_disc": 0.0,
+ "cyclone_usa_r34_nw_radius_disc": 0.0,
+ "cyclone_usa_r34_se_radius_disc": 0.0,
+ "cyclone_usa_r34_sw_radius_disc": 0.0,
+ "cyclone_usa_r50_ne_radius_disc": 0.0,
+ "cyclone_usa_r50_nw_radius_disc": 0.0,
+ "cyclone_usa_r50_se_radius_disc": 0.0,
+ "cyclone_usa_r50_sw_radius_disc": 0.0,
+ "cyclone_usa_r64_ne_radius_disc": 0.0,
+ "cyclone_usa_r64_nw_radius_disc": 0.0,
+ "cyclone_usa_r64_se_radius_disc": 0.0,
+ "cyclone_usa_r64_sw_radius_disc": 0.0,
+ "cyclone_usa_rmw_disc": 0.0,
+ "cyclone_usa_wind_disc": 0.0,
+ "geopotential": 0.06315603852272034,
+ "mean_sea_level_pressure": 0.06843462586402893,
+ "sea_surface_temperature": 1.3908898830413818,
+ "specific_humidity": 0.25252288579940796,
+ "temperature": 0.13354532420635223,
+ "total_precipitation_6hr": 0.12029478698968887,
+ "u_component_of_wind": 0.14959406852722168,
+ "v_component_of_wind": 0.13453276455402374,
+ "vertical_velocity": 0.3185918927192688
+ }
+ },
+ "2026103000": {
+ "loss": 1.2457571029663086,
+ "per_var": {
+ "100m_u_component_of_wind": 0.13579148054122925,
+ "100m_v_component_of_wind": 0.12533296644687653,
+ "10m_u_component_of_wind": 0.14383071660995483,
+ "10m_v_component_of_wind": 0.1339038461446762,
+ "2m_temperature": 0.08920268714427948,
+ "cyclone_all_wind_disc": 0.0,
+ "cyclone_exists_gaussian_unit_mode": 0.0,
+ "cyclone_usa_pres_disc": 0.0,
+ "cyclone_usa_r34_ne_radius_disc": 0.0,
+ "cyclone_usa_r34_nw_radius_disc": 0.0,
+ "cyclone_usa_r34_se_radius_disc": 0.0,
+ "cyclone_usa_r34_sw_radius_disc": 0.0,
+ "cyclone_usa_r50_ne_radius_disc": 0.0,
+ "cyclone_usa_r50_nw_radius_disc": 0.0,
+ "cyclone_usa_r50_se_radius_disc": 0.0,
+ "cyclone_usa_r50_sw_radius_disc": 0.0,
+ "cyclone_usa_r64_ne_radius_disc": 0.0,
+ "cyclone_usa_r64_nw_radius_disc": 0.0,
+ "cyclone_usa_r64_se_radius_disc": 0.0,
+ "cyclone_usa_r64_sw_radius_disc": 0.0,
+ "cyclone_usa_rmw_disc": 0.0,
+ "cyclone_usa_wind_disc": 0.0,
+ "geopotential": 0.06498272716999054,
+ "mean_sea_level_pressure": 0.06535518169403076,
+ "sea_surface_temperature": 1.2963556051254272,
+ "specific_humidity": 0.2425110638141632,
+ "temperature": 0.13159547746181488,
+ "total_precipitation_6hr": 0.11105230450630188,
+ "u_component_of_wind": 0.14253294467926025,
+ "v_component_of_wind": 0.13086187839508057,
+ "vertical_velocity": 0.3023425340652466
+ }
+ },
+ "2026103012": {
+ "loss": 1.271973967552185,
+ "per_var": {
+ "100m_u_component_of_wind": 0.1388230323791504,
+ "100m_v_component_of_wind": 0.12956681847572327,
+ "10m_u_component_of_wind": 0.14567837119102478,
+ "10m_v_component_of_wind": 0.1365891993045807,
+ "2m_temperature": 0.09053972363471985,
+ "cyclone_all_wind_disc": 0.0,
+ "cyclone_exists_gaussian_unit_mode": 0.0,
+ "cyclone_usa_pres_disc": 0.0,
+ "cyclone_usa_r34_ne_radius_disc": 0.0,
+ "cyclone_usa_r34_nw_radius_disc": 0.0,
+ "cyclone_usa_r34_se_radius_disc": 0.0,
+ "cyclone_usa_r34_sw_radius_disc": 0.0,
+ "cyclone_usa_r50_ne_radius_disc": 0.0,
+ "cyclone_usa_r50_nw_radius_disc": 0.0,
+ "cyclone_usa_r50_se_radius_disc": 0.0,
+ "cyclone_usa_r50_sw_radius_disc": 0.0,
+ "cyclone_usa_r64_ne_radius_disc": 0.0,
+ "cyclone_usa_r64_nw_radius_disc": 0.0,
+ "cyclone_usa_r64_se_radius_disc": 0.0,
+ "cyclone_usa_r64_sw_radius_disc": 0.0,
+ "cyclone_usa_rmw_disc": 0.0,
+ "cyclone_usa_wind_disc": 0.0,
+ "geopotential": 0.06138981133699417,
+ "mean_sea_level_pressure": 0.06494763493537903,
+ "sea_surface_temperature": 1.4065196514129639,
+ "specific_humidity": 0.24593773484230042,
+ "temperature": 0.13067081570625305,
+ "total_precipitation_6hr": 0.12287452071905136,
+ "u_component_of_wind": 0.1428370177745819,
+ "v_component_of_wind": 0.12972843647003174,
+ "vertical_velocity": 0.3145983815193176
+ }
+ }
+ }
+}
\ No newline at end of file
diff --git a/docs/results/finetune/finetune_smoke_rtx5090.log b/docs/results/finetune/finetune_smoke_rtx5090.log
new file mode 100644
index 0000000..eb01f7f
--- /dev/null
+++ b/docs/results/finetune/finetune_smoke_rtx5090.log
@@ -0,0 +1,104 @@
+gpu-run: queued behind /home/raymxnd/Workspace/WeatherForcast/.venv/bin/python scripts/verify_head_grad
+gpu-run: queued behind /home/raymxnd/Workspace/WeatherForcast/.venv/bin/python scripts/verify_head_grad
+gpu-run: queued behind /home/raymxnd/Workspace/WeatherForcast/.venv/bin/python scripts/verify_head_grad
+gpu-run: queued behind /home/raymxnd/Workspace/WeatherForcast/.venv/bin/python scripts/verify_head_grad
+gpu-run: queued behind /home/raymxnd/Workspace/WeatherForcast/.venv/bin/python scripts/verify_head_grad
+example 2026100412 built in 2.2s
+WeatherNext2 ckpt 1: training 5.57M of 183.8M params; masked targets: 17 (cyclone channels only)
+WARNING:root:No normalization location found for cyclone_exists_gaussian_unit_mode
+WARNING:root:No normalization scale found for cyclone_exists_gaussian_unit_mode
+WARNING:absl:No coordinates found for dim noise_channels of array Size: 256B
+JitTracer(float32[2,1,32])
+Coordinates:
+ * batch (batch) object 16B MultiIndex
+ * _batch (batch) int64 16B 0 0
+ * sample (batch) int64 16B 0 1
+Dimensions without coordinates: channels, noise_channels., using index for param name. This may make model surgery more error prone if changing coordinates at fine-tuning time.
+/home/raymxnd/Workspace/WeatherForcast/.venv/lib/python3.12/site-packages/scipy/sparse/_index.py:216: SparseEfficiencyWarning: Changing the sparsity structure of a csr_matrix is expensive. lil and dok are more efficient.
+ self._set_arrayXarray(i, j, x)
+/home/raymxnd/Workspace/WeatherForcast/weathernext/weathernext/weathernext2/architecture.py:286: FutureWarning: Deleting a single level of a MultiIndex is deprecated. Previously, this deleted all levels of a MultiIndex. Please also drop the following variables: {'batch'} to avoid an error in the future.
+ predictions[var] = xarray_jax.apply_ufunc(activation_fn, predictions[var])
+W1031 13:00:32.606690 1621781 hlo_rematerialization.cc:3363] Can't reduce memory use below 25.14GiB (26997919705 bytes) by rematerialization; only reduced to 63.07GiB (67717567744 bytes), down from 63.07GiB (67717567744 bytes) originally
+train step compiled in 86s: temp 20.03 GiB, arguments 1.79 GiB
+WARNING:root:No normalization location found for cyclone_exists_gaussian_unit_mode
+WARNING:root:No normalization scale found for cyclone_exists_gaussian_unit_mode
+WARNING:absl:No coordinates found for dim noise_channels of array Size: 256B
+JitTracer(float32[2,1,32])
+Coordinates:
+ * batch (batch) object 16B MultiIndex
+ * _batch (batch) int64 16B 0 0
+ * sample (batch) int64 16B 0 1
+Dimensions without coordinates: channels, noise_channels., using index for param name. This may make model surgery more error prone if changing coordinates at fine-tuning time.
+/home/raymxnd/Workspace/WeatherForcast/.venv/lib/python3.12/site-packages/scipy/sparse/_index.py:216: SparseEfficiencyWarning: Changing the sparsity structure of a csr_matrix is expensive. lil and dok are more efficient.
+ self._set_arrayXarray(i, j, x)
+/home/raymxnd/Workspace/WeatherForcast/weathernext/weathernext/weathernext2/architecture.py:286: FutureWarning: Deleting a single level of a MultiIndex is deprecated. Previously, this deleted all levels of a MultiIndex. Please also drop the following variables: {'batch'} to avoid an error in the future.
+ predictions[var] = xarray_jax.apply_ufunc(activation_fn, predictions[var])
+W1031 13:01:45.906469 1621781 hlo_rematerialization.cc:3363] Can't reduce memory use below 25.12GiB (26978079655 bytes) by rematerialization; only reduced to 56.83GiB (61019999108 bytes), down from 56.83GiB (61020035972 bytes) originally
+eval before (80s): 2026102800 1.19978 2026102812 1.21680 2026103000 1.16300 2026103012 1.19322
+step 0 2026101006 loss 1.25400 |g| 3.308e+00 3.99s
+step 1 2026102406 loss 1.25282 |g| 5.553e+00 3.84s
+step 2 2026101206 loss 1.29070 |g| 4.807e+00 3.84s
+step 3 2026100612 loss 1.14211 |g| 5.481e+00 3.84s
+step 4 2026101900 loss 1.17877 |g| 7.077e+00 3.84s
+step 5 2026101906 loss 1.19659 |g| 4.851e+00 3.84s
+step 6 2026102300 loss 1.28881 |g| 7.436e+00 3.84s
+step 7 2026100512 loss 1.19313 |g| 4.798e+00 3.84s
+step 8 2026101806 loss 1.22023 |g| 4.657e+00 3.84s
+step 9 2026102200 loss 1.12030 |g| 5.656e+00 3.84s
+step 10 2026101500 loss 1.23811 |g| 5.656e+00 3.84s
+step 11 2026100700 loss 1.21001 |g| 5.068e+00 3.84s
+step 12 2026100912 loss 1.14001 |g| 4.637e+00 3.84s
+step 13 2026100500 loss 1.13868 |g| 4.768e+00 3.84s
+step 14 2026101812 loss 1.19795 |g| 5.589e+00 3.84s
+step 15 2026101612 loss 1.09113 |g| 4.964e+00 3.84s
+step 16 2026101112 loss 1.12141 |g| 6.539e+00 3.85s
+step 17 2026102006 loss 1.20427 |g| 5.867e+00 3.84s
+step 18 2026100706 loss 1.25141 |g| 4.719e+00 3.84s
+step 19 2026102312 loss 1.09864 |g| 5.151e+00 3.84s
+step 20 2026102106 loss 1.23133 |g| 3.620e+00 3.84s
+step 21 2026101606 loss 1.30773 |g| 6.664e+00 3.84s
+step 22 2026101000 loss 1.15977 |g| 8.556e+00 3.84s
+step 23 2026101600 loss 1.08576 |g| 6.319e+00 3.84s
+step 24 2026100506 loss 1.18339 |g| 7.859e+00 3.84s
+step 25 2026100418 loss 1.37284 |g| 8.901e+00 3.84s
+step 26 2026101412 loss 1.08405 |g| 5.338e+00 3.84s
+step 27 2026102306 loss 1.19583 |g| 5.533e+00 3.84s
+step 28 2026101212 loss 1.05301 |g| 4.883e+00 3.84s
+step 29 2026100918 loss 1.25942 |g| 5.659e+00 3.85s
+step 30 2026101306 loss 1.14296 |g| 5.380e+00 3.84s
+step 31 2026101800 loss 1.03347 |g| 5.858e+00 3.85s
+step 32 2026101118 loss 1.18195 |g| 4.815e+00 3.84s
+step 33 2026100412 loss 1.16975 |g| 5.781e+00 3.85s
+step 34 2026100600 loss 1.30958 |g| 1.369e+01 3.85s
+step 35 2026101200 loss 1.16497 |g| 9.568e+00 3.84s
+step 36 2026101300 loss 0.99685 |g| 6.923e+00 3.84s
+step 37 2026101618 loss 1.20327 |g| 8.204e+00 3.84s
+step 38 2026102318 loss 1.17374 |g| 6.496e+00 3.85s
+step 39 2026102012 loss 1.08462 |g| 9.564e+00 3.84s
+step 40 2026100806 loss 1.58468 |g| 2.083e+01 3.84s
+step 41 2026101512 loss 0.97707 |g| 8.755e+00 3.84s
+step 42 2026101218 loss 1.43003 |g| 1.851e+01 3.85s
+step 43 2026102212 loss 1.36700 |g| 1.645e+01 3.84s
+step 44 2026100618 loss 1.11921 |g| 8.376e+00 3.84s
+step 45 2026101700 loss 1.29024 |g| 1.564e+01 3.84s
+step 46 2026102018 loss 1.14205 |g| 8.364e+00 3.85s
+step 47 2026102218 loss 1.11789 |g| 7.237e+00 3.84s
+step 48 2026102000 loss 1.43831 |g| 2.065e+01 3.84s
+step 49 2026101712 loss 1.04715 |g| 9.736e+00 3.84s
+step 50 2026101912 loss 1.23614 |g| 1.422e+01 3.84s
+step 51 2026100718 loss 1.10873 |g| 1.020e+01 3.84s
+step 52 2026100712 loss 1.32132 |g| 1.460e+01 3.84s
+step 53 2026100606 loss 1.07605 |g| 7.338e+00 3.84s
+step 54 2026101818 loss 1.37630 |g| 1.458e+01 3.84s
+step 55 2026101706 loss 1.06347 |g| 7.361e+00 3.84s
+step 56 2026101506 loss 1.28202 |g| 1.016e+01 3.85s
+step 57 2026100518 loss 1.09595 |g| 7.118e+00 3.85s
+step 58 2026101918 loss 1.02739 |g| 7.616e+00 3.84s
+step 59 2026102118 loss 1.31273 |g| 1.204e+01 3.84s
+median step 3.84s, device peak 21.82 GiB
+eval 2026102800: 1.19978 -> 1.28391 (+7.01 %)
+eval 2026102812: 1.21680 -> 1.29790 (+6.67 %)
+eval 2026103000: 1.16300 -> 1.24576 (+7.12 %)
+eval 2026103012: 1.19322 -> 1.27197 (+6.60 %)
+wrote /tmp/claude-1000/-home-raymxnd-Workspace-WeatherForcast/d4a93a6d-df87-4c4c-b7ce-cc8b5b5265d2/scratchpad/head_smoke.npz
+wrote docs/results/finetune/finetune_smoke_rtx5090.json
diff --git a/docs/results/finetune/finetune_smoke_rtx5090_oom_growing_pool.log b/docs/results/finetune/finetune_smoke_rtx5090_oom_growing_pool.log
new file mode 100644
index 0000000..62f4023
--- /dev/null
+++ b/docs/results/finetune/finetune_smoke_rtx5090_oom_growing_pool.log
@@ -0,0 +1,94 @@
+gpu-run: queued behind /home/raymxnd/Workspace/WeatherForcast/.venv/bin/python scripts/verify_head_grad
+gpu-run: queued behind /home/raymxnd/Workspace/WeatherForcast/.venv/bin/python scripts/verify_head_grad
+gpu-run: queued behind /home/raymxnd/Workspace/WeatherForcast/.venv/bin/python scripts/verify_head_grad
+gpu-run: queued behind /home/raymxnd/Workspace/WeatherForcast/.venv/bin/python scripts/verify_head_grad
+gpu-run: queued behind /home/raymxnd/Workspace/WeatherForcast/.venv/bin/python scripts/verify_head_grad
+gpu-run: queued behind /home/raymxnd/Workspace/WeatherForcast/.venv/bin/python scripts/verify_head_grad
+gpu-run: queued behind /home/raymxnd/Workspace/WeatherForcast/.venv/bin/python scripts/verify_head_grad
+gpu-run: queued behind /home/raymxnd/Workspace/WeatherForcast/.venv/bin/python scripts/verify_head_grad
+gpu-run: queued behind /home/raymxnd/Workspace/WeatherForcast/.venv/bin/python scripts/verify_head_grad
+gpu-run: queued behind /home/raymxnd/Workspace/WeatherForcast/.venv/bin/python scripts/verify_head_grad
+gpu-run: queued behind /home/raymxnd/Workspace/WeatherForcast/.venv/bin/python scripts/verify_head_grad
+gpu-run: queued behind /home/raymxnd/Workspace/WeatherForcast/.venv/bin/python scripts/verify_head_grad
+gpu-run: queued behind /home/raymxnd/Workspace/WeatherForcast/.venv/bin/python scripts/verify_head_grad
+gpu-run: queued behind /home/raymxnd/Workspace/WeatherForcast/.venv/bin/python scripts/verify_head_grad
+gpu-run: queued behind /home/raymxnd/Workspace/WeatherForcast/.venv/bin/python scripts/verify_head_grad
+gpu-run: queued behind /home/raymxnd/Workspace/WeatherForcast/.venv/bin/python scripts/verify_head_grad
+gpu-run: queued behind /home/raymxnd/Workspace/WeatherForcast/.venv/bin/python scripts/verify_head_grad
+gpu-run: queued behind /home/raymxnd/Workspace/WeatherForcast/.venv/bin/python scripts/verify_head_grad
+gpu-run: queued behind /home/raymxnd/Workspace/WeatherForcast/.venv/bin/python scripts/verify_head_grad
+gpu-run: queued behind /home/raymxnd/Workspace/WeatherForcast/.venv/bin/python scripts/verify_head_grad
+gpu-run: queued behind /home/raymxnd/Workspace/WeatherForcast/.venv/bin/python scripts/verify_head_grad
+gpu-run: queued behind /home/raymxnd/Workspace/WeatherForcast/.venv/bin/python scripts/verify_head_grad
+gpu-run: queued behind /home/raymxnd/Workspace/WeatherForcast/.venv/bin/python scripts/verify_head_grad
+gpu-run: queued behind /home/raymxnd/Workspace/WeatherForcast/.venv/bin/python scripts/verify_head_grad
+gpu-run: queued behind /home/raymxnd/Workspace/WeatherForcast/.venv/bin/python scripts/verify_head_grad
+gpu-run: queued behind /home/raymxnd/Workspace/WeatherForcast/.venv/bin/python scripts/verify_head_grad
+gpu-run: queued behind /home/raymxnd/Workspace/WeatherForcast/.venv/bin/python scripts/verify_head_grad
+gpu-run: queued behind /home/raymxnd/Workspace/WeatherForcast/.venv/bin/python scripts/verify_head_grad
+gpu-run: queued behind /home/raymxnd/Workspace/WeatherForcast/.venv/bin/python scripts/verify_head_grad
+gpu-run: queued behind /home/raymxnd/Workspace/WeatherForcast/.venv/bin/python scripts/verify_head_grad
+example 2026100412 built in 3.4s
+WeatherNext2 ckpt 1: training 5.57M of 183.8M params; masked targets: 17 (cyclone channels only)
+WARNING:root:No normalization location found for cyclone_exists_gaussian_unit_mode
+WARNING:root:No normalization scale found for cyclone_exists_gaussian_unit_mode
+WARNING:absl:No coordinates found for dim noise_channels of array Size: 256B
+JitTracer(float32[2,1,32])
+Coordinates:
+ * batch (batch) object 16B MultiIndex
+ * _batch (batch) int64 16B 0 0
+ * sample (batch) int64 16B 0 1
+Dimensions without coordinates: channels, noise_channels., using index for param name. This may make model surgery more error prone if changing coordinates at fine-tuning time.
+/home/raymxnd/Workspace/WeatherForcast/.venv/lib/python3.12/site-packages/scipy/sparse/_index.py:216: SparseEfficiencyWarning: Changing the sparsity structure of a csr_matrix is expensive. lil and dok are more efficient.
+ self._set_arrayXarray(i, j, x)
+/home/raymxnd/Workspace/WeatherForcast/weathernext/weathernext/weathernext2/architecture.py:286: FutureWarning: Deleting a single level of a MultiIndex is deprecated. Previously, this deleted all levels of a MultiIndex. Please also drop the following variables: {'batch'} to avoid an error in the future.
+ predictions[var] = xarray_jax.apply_ufunc(activation_fn, predictions[var])
+W1031 12:48:30.191703 1578596 hlo_rematerialization.cc:3363] Can't reduce memory use below 25.74GiB (27637446207 bytes) by rematerialization; only reduced to 63.07GiB (67717567744 bytes), down from 63.07GiB (67717567744 bytes) originally
+train step compiled in 86s: temp 20.03 GiB, arguments 1.79 GiB
+WARNING:root:No normalization location found for cyclone_exists_gaussian_unit_mode
+WARNING:root:No normalization scale found for cyclone_exists_gaussian_unit_mode
+WARNING:absl:No coordinates found for dim noise_channels of array Size: 256B
+JitTracer(float32[2,1,32])
+Coordinates:
+ * batch (batch) object 16B MultiIndex
+ * _batch (batch) int64 16B 0 0
+ * sample (batch) int64 16B 0 1
+Dimensions without coordinates: channels, noise_channels., using index for param name. This may make model surgery more error prone if changing coordinates at fine-tuning time.
+/home/raymxnd/Workspace/WeatherForcast/.venv/lib/python3.12/site-packages/scipy/sparse/_index.py:216: SparseEfficiencyWarning: Changing the sparsity structure of a csr_matrix is expensive. lil and dok are more efficient.
+ self._set_arrayXarray(i, j, x)
+/home/raymxnd/Workspace/WeatherForcast/weathernext/weathernext/weathernext2/architecture.py:286: FutureWarning: Deleting a single level of a MultiIndex is deprecated. Previously, this deleted all levels of a MultiIndex. Please also drop the following variables: {'batch'} to avoid an error in the future.
+ predictions[var] = xarray_jax.apply_ufunc(activation_fn, predictions[var])
+W1031 12:49:47.404300 1578596 hlo_rematerialization.cc:3363] Can't reduce memory use below 25.72GiB (27617606157 bytes) by rematerialization; only reduced to 56.83GiB (61019999108 bytes), down from 56.83GiB (61020035972 bytes) originally
+eval before (80s): 2026102800 1.19979 2026102812 1.21682 2026103000 1.16299 2026103012 1.19323
+W1031 12:50:17.608056 1578596 bfc_allocator.cc:598] Allocator (GPU_0_bfc) ran out of memory trying to allocate 20.03GiB (rounded to 21508145408) from the lower end requested by op
+If this is caused by memory fragmentation, the environment variable 'TF_GPU_ALLOCATOR=cuda_malloc_async' may improve the situation.
+Current allocation summary follows.
+W1031 12:50:17.608581 1578596 bfc_allocator.cc:609] *********__________________________________________________________________________________________*
+Traceback (most recent call last):
+ File "/home/raymxnd/Workspace/faster-weathernext/scripts/finetune_head.py", line 162, in
+ main()
+ File "/home/raymxnd/Workspace/faster-weathernext/scripts/finetune_head.py", line 130, in main
+ metrics = jax.device_get(metrics)
+ ^^^^^^^^^^^^^^^^^^^^^^^
+ File "/home/raymxnd/Workspace/WeatherForcast/.venv/lib/python3.12/site-packages/jax/_src/api.py", line 2592, in device_get
+ return tree_map(_device_get, x)
+ ^^^^^^^^^^^^^^^^^^^^^^^^
+ File "/home/raymxnd/Workspace/WeatherForcast/.venv/lib/python3.12/site-packages/jax/_src/tree_util.py", line 406, in tree_map
+ return treedef.unflatten(f(*xs) for xs in zip(*all_leaves))
+ ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
+ File "/home/raymxnd/Workspace/WeatherForcast/.venv/lib/python3.12/site-packages/jax/_src/tree_util.py", line 406, in
+ return treedef.unflatten(f(*xs) for xs in zip(*all_leaves))
+ ^^^^^^
+ File "/home/raymxnd/Workspace/WeatherForcast/.venv/lib/python3.12/site-packages/jax/_src/api.py", line 2553, in _device_get
+ return toarray()
+ ^^^^^^^^^
+ File "/home/raymxnd/Workspace/WeatherForcast/.venv/lib/python3.12/site-packages/jax/_src/array.py", line 409, in __array__
+ return np.asarray(self._value, dtype=dtype, **kwds) # pyrefly: ignore[no-matching-overload]
+ ^^^^^^^^^^^
+ File "/home/raymxnd/Workspace/WeatherForcast/.venv/lib/python3.12/site-packages/jax/_src/profiler.py", line 420, in wrapper
+ return func(*args, **kwargs)
+ ^^^^^^^^^^^^^^^^^^^^^
+ File "/home/raymxnd/Workspace/WeatherForcast/.venv/lib/python3.12/site-packages/jax/_src/array.py", line 637, in _value
+ npy_value, did_copy = self._single_device_array_to_np_array_did_copy()
+ ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
+jax.errors.JaxRuntimeError: RESOURCE_EXHAUSTED: Out of memory while trying to allocate 20.03GiB with allocator GPU_0_bfc on device 0. [executable_name='jit_step'] [tf-allocator-allocation-error='']
diff --git a/docs/results/finetune/full_grads_mini_crps_rtx5090.json b/docs/results/finetune/full_grads_mini_crps_rtx5090.json
new file mode 100644
index 0000000..98a3ece
--- /dev/null
+++ b/docs/results/finetune/full_grads_mini_crps_rtx5090.json
@@ -0,0 +1,259 @@
+{
+ "model": "WeatherNextCyclones_Mini",
+ "device": "NVIDIA GeForce RTX 5090",
+ "jax": "0.11.2",
+ "faster_weathernext": "0.1.1",
+ "blocks": {
+ "grid_chunk": 8192,
+ "edge_chunk": 16384
+ },
+ "masked": [],
+ "modules": "all",
+ "loss": "crps",
+ "head_params": 56684133,
+ "precisions": {
+ "highest": {
+ "official": {
+ "temp_gib": 83.65751740336418,
+ "compile_s": 82.9094603061676,
+ "oom": "RESOURCE_EXHAUSTED: Out of memory while trying to allocate 83.66GiB with allocator GPU_0_bfc on device 0. [tf-allocator-allocation-error=''] [executable_name='jit_fn']"
+ },
+ "official_mha": {
+ "temp_gib": 115.73732301592827,
+ "compile_s": 40.991973876953125,
+ "oom": "RESOURCE_EXHAUSTED: Out of memory while trying to allocate 115.74GiB with allocator GPU_0_bfc on device 0. [tf-allocator-allocation-error=''] [executable_name='jit_fn']"
+ },
+ "fwn": {
+ "temp_gib": 3.9292222931981087,
+ "compile_s": 57.98857641220093,
+ "step_s": 0.8979263305664062
+ },
+ "fwn_xla": {
+ "temp_gib": 4.046475827693939,
+ "compile_s": 32.29124569892883,
+ "step_s": 1.1193935871124268
+ },
+ "fwn_strict": {
+ "temp_gib": 3.9292222931981087,
+ "compile_s": 20.617359161376953,
+ "step_s": 3.3168978691101074
+ },
+ "fwn_noremat": {
+ "temp_gib": 38.95218490809202,
+ "compile_s": 25.933741569519043
+ },
+ "reference": "fwn_xla",
+ "checks": {
+ "fwn vs fwn_xla": {
+ "loss": {
+ "official": 1.2029576301574707,
+ "other": 1.2029577493667603,
+ "rel_diff": 9.909683147791032e-08,
+ "max_per_var_rel_diff": 8.181437953570938e-06
+ },
+ "grads": {
+ "global_rel_l2": 0.001152629461458604,
+ "global_cos": 0.9999993373931454,
+ "max_rel_l2": 0.005396667036540233,
+ "n_params": 376,
+ "n_compared": 335,
+ "n_zero_both": 6,
+ "n_zero_ref_only": 0,
+ "all_finite": true,
+ "worst": [
+ {
+ "param": "multimodality_forward/mesh_transformer/transformer/block_04/mha_proj_k/w",
+ "size": 262144,
+ "ref_norm": 0.15890459470408802,
+ "got_norm": 0.1586603189070189,
+ "finite": true,
+ "rel_l2": 0.005396667036540233
+ },
+ {
+ "param": "multimodality_forward/mesh_encoder/split_input_matmul/w_spatial_feature=sin_lat,sin_lon,cos_lon",
+ "size": 1536,
+ "ref_norm": 0.36128790869721605,
+ "got_norm": 0.3616580282111332,
+ "finite": true,
+ "rel_l2": 0.005370125279633845
+ },
+ {
+ "param": "multimodality_forward/mesh_transformer/transformer/block_01/block_01_norm_conditioning/linear/w",
+ "size": 32768,
+ "ref_norm": 0.05497876121409077,
+ "got_norm": 0.054763677476260694,
+ "finite": true,
+ "rel_l2": 0.005128952604055275
+ },
+ {
+ "param": "multimodality_forward/mesh_transformer/transformer/block_04/mha_proj_q/w",
+ "size": 262144,
+ "ref_norm": 0.16822486075147752,
+ "got_norm": 0.16798405898640578,
+ "finite": true,
+ "rel_l2": 0.004761338941207865
+ },
+ {
+ "param": "multimodality_forward/mesh_transformer/transformer/block_01/block_01_norm_conditioning/linear/b",
+ "size": 1024,
+ "ref_norm": 0.03666123682218743,
+ "got_norm": 0.03655480898292301,
+ "finite": true,
+ "rel_l2": 0.0043501316204721495
+ }
+ ]
+ }
+ },
+ "fwn_strict vs fwn_xla": {
+ "loss": {
+ "official": 1.2029576301574707,
+ "other": 1.2029575109481812,
+ "rel_diff": 9.909683147791032e-08,
+ "max_per_var_rel_diff": 1.2444388266032497e-05
+ },
+ "grads": {
+ "global_rel_l2": 0.0007352196909427228,
+ "global_cos": 0.9999997309546174,
+ "max_rel_l2": 0.0023494882522781283,
+ "n_params": 376,
+ "n_compared": 335,
+ "n_zero_both": 6,
+ "n_zero_ref_only": 0,
+ "all_finite": true,
+ "worst": [
+ {
+ "param": "multimodality_forward/mesh_transformer/transformer/block_07/block_07_norm_conditioning/linear/w",
+ "size": 32768,
+ "ref_norm": 0.01885554025288714,
+ "got_norm": 0.018866430856829746,
+ "finite": true,
+ "rel_l2": 0.0023494882522781283
+ },
+ {
+ "param": "multimodality_forward/grid_decoder/split_output_linear/w_specific_humidity_level=50,100,150,200,250,300,400,500,600,700,850,925,1000_time=21600",
+ "size": 6656,
+ "ref_norm": 0.06549569836271832,
+ "got_norm": 0.06549132880851923,
+ "finite": true,
+ "rel_l2": 0.0022983905501291883
+ },
+ {
+ "param": "multimodality_forward/mesh_transformer/transformer/block_07/block_07_norm_conditioning/linear/b",
+ "size": 1024,
+ "ref_norm": 0.011977532343040461,
+ "got_norm": 0.011976822276668487,
+ "finite": true,
+ "rel_l2": 0.0019974693193459117
+ },
+ {
+ "param": "multimodality_forward/mesh_transformer/transformer/block_04/block_04_norm_conditioning_1/linear/w",
+ "size": 32768,
+ "ref_norm": 0.018693075444248905,
+ "got_norm": 0.01869524516062543,
+ "finite": true,
+ "rel_l2": 0.0017746622603545985
+ },
+ {
+ "param": "multimodality_forward/mesh_transformer/transformer/block_06/mha_proj_k/w",
+ "size": 262144,
+ "ref_norm": 0.08019847974205342,
+ "got_norm": 0.08021715963376236,
+ "finite": true,
+ "rel_l2": 0.0015166800874850351
+ }
+ ]
+ }
+ }
+ }
+ },
+ "default": {
+ "official": {
+ "temp_gib": 82.17185494303703,
+ "compile_s": 47.40474319458008,
+ "oom": "RESOURCE_EXHAUSTED: Out of memory while trying to allocate 82.17GiB with allocator GPU_0_bfc on device 0. [tf-allocator-allocation-error=''] [executable_name='jit_fn']"
+ },
+ "official_mha": {
+ "temp_gib": 106.66842457652092,
+ "compile_s": 28.64564847946167,
+ "oom": "RESOURCE_EXHAUSTED: Out of memory while trying to allocate 106.67GiB with allocator GPU_0_bfc on device 0. [tf-allocator-allocation-error=''] [executable_name='jit_fn']"
+ },
+ "fwn": {
+ "temp_gib": 4.007309623062611,
+ "compile_s": 36.81330490112305,
+ "step_s": 0.4123044013977051
+ },
+ "fwn_xla": {
+ "temp_gib": 4.162897050380707,
+ "compile_s": 27.714158058166504,
+ "step_s": 0.8910126686096191
+ },
+ "fwn_noremat": {
+ "temp_gib": 38.952152244746685,
+ "compile_s": 20.89156484603882
+ },
+ "reference": "fwn_xla",
+ "checks": {
+ "fwn vs fwn_xla": {
+ "loss": {
+ "official": 1.2031620740890503,
+ "other": 1.2031536102294922,
+ "rel_diff": 7.034679483654526e-06,
+ "max_per_var_rel_diff": 0.0028837358039602396
+ },
+ "grads": {
+ "global_rel_l2": 0.044134135651816876,
+ "global_cos": 0.9990259515397509,
+ "max_rel_l2": 0.18180455021771114,
+ "n_params": 376,
+ "n_compared": 335,
+ "n_zero_both": 6,
+ "n_zero_ref_only": 0,
+ "all_finite": true,
+ "worst": [
+ {
+ "param": "multimodality_forward/mesh_transformer/transformer/block_07/block_07_norm_conditioning/linear/b",
+ "size": 1024,
+ "ref_norm": 0.01253671756174312,
+ "got_norm": 0.012857313685722162,
+ "finite": true,
+ "rel_l2": 0.18180455021771114
+ },
+ {
+ "param": "multimodality_forward/mesh_transformer/transformer/block_01/block_01_norm_conditioning/linear/b",
+ "size": 1024,
+ "ref_norm": 0.03141792291414201,
+ "got_norm": 0.03467389298263372,
+ "finite": true,
+ "rel_l2": 0.17160785316338278
+ },
+ {
+ "param": "multimodality_forward/grid_to_mesh_gnn/edge_encoder/shared_dense/mlp/linear_0/b",
+ "size": 32,
+ "ref_norm": 0.2904886811726889,
+ "got_norm": 0.25693576863629586,
+ "finite": true,
+ "rel_l2": 0.15031464004877826
+ },
+ {
+ "param": "multimodality_forward/mesh_encoder/split_input_matmul/w_forcing_year_progress_cos_time=21600",
+ "size": 512,
+ "ref_norm": 0.029876893921651937,
+ "got_norm": 0.03038230986338254,
+ "finite": true,
+ "rel_l2": 0.14808206528931833
+ },
+ {
+ "param": "multimodality_forward/mesh_encoder/split_input_matmul/w_forcing_year_progress_sin_time=21600",
+ "size": 512,
+ "ref_norm": 0.26616761879382167,
+ "got_norm": 0.2706703007440032,
+ "finite": true,
+ "rel_l2": 0.1480820277472271
+ }
+ ]
+ }
+ }
+ }
+ }
+ }
+}
\ No newline at end of file
diff --git a/docs/results/finetune/full_grads_mini_crps_rtx5090.log b/docs/results/finetune/full_grads_mini_crps_rtx5090.log
new file mode 100644
index 0000000..d970b08
--- /dev/null
+++ b/docs/results/finetune/full_grads_mini_crps_rtx5090.log
@@ -0,0 +1,220 @@
+WeatherNextCyclones_Mini: differentiating 56.68M params (all); masked targets: []
+[highest]
+WARNING:root:No normalization location found for cyclone_exists_gaussian_unit_mode
+WARNING:root:No normalization scale found for cyclone_exists_gaussian_unit_mode
+WARNING:absl:No coordinates found for dim noise_channels of array Size: 256B
+JitTracer(float32[2,1,32])
+Coordinates:
+ * batch (batch) object 16B MultiIndex
+ * _batch (batch) int64 16B 0 0
+ * sample (batch) int64 16B 0 1
+Dimensions without coordinates: channels, noise_channels., using index for param name. This may make model surgery more error prone if changing coordinates at fine-tuning time.
+/home/raymxnd/Workspace/WeatherForcast/.venv/lib/python3.12/site-packages/scipy/sparse/_index.py:216: SparseEfficiencyWarning: Changing the sparsity structure of a csr_matrix is expensive. lil and dok are more efficient.
+ self._set_arrayXarray(i, j, x)
+WARNING:absl:Only the first input of the static gather operation is being used, others are being ignored at runtime.
+WARNING:absl:Only the first input of the static gather operation is being used, others are being ignored at runtime.
+WARNING:absl:Only the first input of the sorted segment sum operation is being used, others are being ignored at runtime.
+WARNING:absl:Only the first input of the static gather operation is being used, others are being ignored at runtime.
+WARNING:absl:Only the first input of the static gather operation is being used, others are being ignored at runtime.
+WARNING:absl:Only the first input of the static gather operation is being used, others are being ignored at runtime.
+WARNING:absl:Only the first input of the sorted segment sum operation is being used, others are being ignored at runtime.
+/home/raymxnd/Workspace/WeatherForcast/weathernext/weathernext/weathernext2/architecture.py:286: FutureWarning: Deleting a single level of a MultiIndex is deprecated. Previously, this deleted all levels of a MultiIndex. Please also drop the following variables: {'batch'} to avoid an error in the future.
+ predictions[var] = xarray_jax.apply_ufunc(activation_fn, predictions[var])
+W1101 00:47:13.322365 53506 hlo_rematerialization.cc:3363] Can't reduce memory use below 22.07GiB (23696756432 bytes) by rematerialization; only reduced to 77.95GiB (83697687032 bytes), down from 101.77GiB (109270673384 bytes) originally
+W1101 00:47:27.622833 53506 bfc_allocator.cc:598] Allocator (GPU_0_bfc) ran out of memory trying to allocate 83.66GiB (rounded to 89826575360) from the lower end requested by op
+If this is caused by memory fragmentation, the environment variable 'TF_GPU_ALLOCATOR=cuda_malloc_async' may improve the situation.
+Current allocation summary follows.
+W1101 00:47:27.623379 53506 bfc_allocator.cc:609] ************_______________________________________________________________________________________*
+ official highest temp 83.66 GiB OUT OF MEMORY
+WARNING:root:No normalization location found for cyclone_exists_gaussian_unit_mode
+WARNING:root:No normalization scale found for cyclone_exists_gaussian_unit_mode
+WARNING:absl:No coordinates found for dim noise_channels of array Size: 256B
+JitTracer(float32[2,1,32])
+Coordinates:
+ * batch (batch) object 16B MultiIndex
+ * _batch (batch) int64 16B 0 0
+ * sample (batch) int64 16B 0 1
+Dimensions without coordinates: channels, noise_channels., using index for param name. This may make model surgery more error prone if changing coordinates at fine-tuning time.
+/home/raymxnd/Workspace/WeatherForcast/.venv/lib/python3.12/site-packages/scipy/sparse/_index.py:216: SparseEfficiencyWarning: Changing the sparsity structure of a csr_matrix is expensive. lil and dok are more efficient.
+ self._set_arrayXarray(i, j, x)
+WARNING:absl:Only the first input of the static gather operation is being used, others are being ignored at runtime.
+WARNING:absl:Only the first input of the static gather operation is being used, others are being ignored at runtime.
+WARNING:absl:Only the first input of the sorted segment sum operation is being used, others are being ignored at runtime.
+WARNING:absl:Only the first input of the static gather operation is being used, others are being ignored at runtime.
+WARNING:absl:Only the first input of the static gather operation is being used, others are being ignored at runtime.
+WARNING:absl:Only the first input of the static gather operation is being used, others are being ignored at runtime.
+WARNING:absl:Only the first input of the sorted segment sum operation is being used, others are being ignored at runtime.
+/home/raymxnd/Workspace/WeatherForcast/weathernext/weathernext/weathernext2/architecture.py:286: FutureWarning: Deleting a single level of a MultiIndex is deprecated. Previously, this deleted all levels of a MultiIndex. Please also drop the following variables: {'batch'} to avoid an error in the future.
+ predictions[var] = xarray_jax.apply_ufunc(activation_fn, predictions[var])
+W1101 00:48:04.836398 53506 hlo_rematerialization.cc:3363] Can't reduce memory use below 22.07GiB (23696756432 bytes) by rematerialization; only reduced to 110.82GiB (118988009376 bytes), down from 123.27GiB (132364702304 bytes) originally
+W1101 00:48:19.113686 53506 bfc_allocator.cc:598] Allocator (GPU_0_bfc) ran out of memory trying to allocate 115.74GiB (rounded to 124272004352) from the lower end requested by op
+If this is caused by memory fragmentation, the environment variable 'TF_GPU_ALLOCATOR=cuda_malloc_async' may improve the situation.
+Current allocation summary follows.
+W1101 00:48:19.114732 53506 bfc_allocator.cc:609] *****______________________________________________________________________________________________*
+ official_mha highest temp 115.74 GiB OUT OF MEMORY
+WARNING:root:No normalization location found for cyclone_exists_gaussian_unit_mode
+WARNING:root:No normalization scale found for cyclone_exists_gaussian_unit_mode
+WARNING:absl:No coordinates found for dim noise_channels of array Size: 256B
+JitTracer(float32[2,1,32])
+Coordinates:
+ * batch (batch) object 16B MultiIndex
+ * _batch (batch) int64 16B 0 0
+ * sample (batch) int64 16B 0 1
+Dimensions without coordinates: channels, noise_channels., using index for param name. This may make model surgery more error prone if changing coordinates at fine-tuning time.
+/home/raymxnd/Workspace/WeatherForcast/.venv/lib/python3.12/site-packages/scipy/sparse/_index.py:216: SparseEfficiencyWarning: Changing the sparsity structure of a csr_matrix is expensive. lil and dok are more efficient.
+ self._set_arrayXarray(i, j, x)
+/home/raymxnd/Workspace/WeatherForcast/weathernext/weathernext/weathernext2/architecture.py:286: FutureWarning: Deleting a single level of a MultiIndex is deprecated. Previously, this deleted all levels of a MultiIndex. Please also drop the following variables: {'batch'} to avoid an error in the future.
+ predictions[var] = xarray_jax.apply_ufunc(activation_fn, predictions[var])
+ fwn highest temp 3.93 GiB step 0.90s loss 1.202958
+WARNING:root:No normalization location found for cyclone_exists_gaussian_unit_mode
+WARNING:root:No normalization scale found for cyclone_exists_gaussian_unit_mode
+WARNING:absl:No coordinates found for dim noise_channels of array Size: 256B
+JitTracer(float32[2,1,32])
+Coordinates:
+ * batch (batch) object 16B MultiIndex
+ * _batch (batch) int64 16B 0 0
+ * sample (batch) int64 16B 0 1
+Dimensions without coordinates: channels, noise_channels., using index for param name. This may make model surgery more error prone if changing coordinates at fine-tuning time.
+/home/raymxnd/Workspace/WeatherForcast/.venv/lib/python3.12/site-packages/scipy/sparse/_index.py:216: SparseEfficiencyWarning: Changing the sparsity structure of a csr_matrix is expensive. lil and dok are more efficient.
+ self._set_arrayXarray(i, j, x)
+/home/raymxnd/Workspace/WeatherForcast/weathernext/weathernext/weathernext2/architecture.py:286: FutureWarning: Deleting a single level of a MultiIndex is deprecated. Previously, this deleted all levels of a MultiIndex. Please also drop the following variables: {'batch'} to avoid an error in the future.
+ predictions[var] = xarray_jax.apply_ufunc(activation_fn, predictions[var])
+ fwn_xla highest temp 4.05 GiB step 1.12s loss 1.202958
+WARNING:root:No normalization location found for cyclone_exists_gaussian_unit_mode
+WARNING:root:No normalization scale found for cyclone_exists_gaussian_unit_mode
+WARNING:absl:No coordinates found for dim noise_channels of array Size: 256B
+JitTracer(float32[2,1,32])
+Coordinates:
+ * batch (batch) object 16B MultiIndex
+ * _batch (batch) int64 16B 0 0
+ * sample (batch) int64 16B 0 1
+Dimensions without coordinates: channels, noise_channels., using index for param name. This may make model surgery more error prone if changing coordinates at fine-tuning time.
+/home/raymxnd/Workspace/WeatherForcast/.venv/lib/python3.12/site-packages/scipy/sparse/_index.py:216: SparseEfficiencyWarning: Changing the sparsity structure of a csr_matrix is expensive. lil and dok are more efficient.
+ self._set_arrayXarray(i, j, x)
+/home/raymxnd/Workspace/WeatherForcast/weathernext/weathernext/weathernext2/architecture.py:286: FutureWarning: Deleting a single level of a MultiIndex is deprecated. Previously, this deleted all levels of a MultiIndex. Please also drop the following variables: {'batch'} to avoid an error in the future.
+ predictions[var] = xarray_jax.apply_ufunc(activation_fn, predictions[var])
+ fwn_strict highest temp 3.93 GiB step 3.32s loss 1.202958
+WARNING:root:No normalization location found for cyclone_exists_gaussian_unit_mode
+WARNING:root:No normalization scale found for cyclone_exists_gaussian_unit_mode
+WARNING:absl:No coordinates found for dim noise_channels of array Size: 256B
+JitTracer(float32[2,1,32])
+Coordinates:
+ * batch (batch) object 16B MultiIndex
+ * _batch (batch) int64 16B 0 0
+ * sample (batch) int64 16B 0 1
+Dimensions without coordinates: channels, noise_channels., using index for param name. This may make model surgery more error prone if changing coordinates at fine-tuning time.
+/home/raymxnd/Workspace/WeatherForcast/.venv/lib/python3.12/site-packages/scipy/sparse/_index.py:216: SparseEfficiencyWarning: Changing the sparsity structure of a csr_matrix is expensive. lil and dok are more efficient.
+ self._set_arrayXarray(i, j, x)
+/home/raymxnd/Workspace/WeatherForcast/weathernext/weathernext/weathernext2/architecture.py:286: FutureWarning: Deleting a single level of a MultiIndex is deprecated. Previously, this deleted all levels of a MultiIndex. Please also drop the following variables: {'batch'} to avoid an error in the future.
+ predictions[var] = xarray_jax.apply_ufunc(activation_fn, predictions[var])
+W1101 00:50:43.898644 53506 hlo_rematerialization.cc:3363] Can't reduce memory use below 22.07GiB (23696756432 bytes) by rematerialization; only reduced to 39.12GiB (42009086092 bytes), down from 39.67GiB (42600728697 bytes) originally
+ fwn_noremat highest temp 38.95 GiB (compile only)
+ fwn vs fwn_xla loss rel diff 9.91e-08 (per var max 8.18e-06) grads: global rel L2 1.15e-03 cos 0.9999993374 max per-param rel L2 5.40e-03 (335/376 params; zero in both 6, zero only in official 0) finite True
+ 5.40e-03 multimodality_forward/mesh_transformer/transformer/block_04/mha_proj_k/w
+ 5.37e-03 multimodality_forward/mesh_encoder/split_input_matmul/w_spatial_feature=sin_lat,sin_lon,cos_lon
+ 5.13e-03 multimodality_forward/mesh_transformer/transformer/block_01/block_01_norm_conditioning/linear/w
+ fwn_strict vs fwn_xla loss rel diff 9.91e-08 (per var max 1.24e-05) grads: global rel L2 7.35e-04 cos 0.9999997310 max per-param rel L2 2.35e-03 (335/376 params; zero in both 6, zero only in official 0) finite True
+ 2.35e-03 multimodality_forward/mesh_transformer/transformer/block_07/block_07_norm_conditioning/linear/w
+ 2.30e-03 multimodality_forward/grid_decoder/split_output_linear/w_specific_humidity_level=50,100,150,200,250,300,400,500,600,700,850,925,1000_time=21600
+ 2.00e-03 multimodality_forward/mesh_transformer/transformer/block_07/block_07_norm_conditioning/linear/b
+[default]
+WARNING:root:No normalization location found for cyclone_exists_gaussian_unit_mode
+WARNING:root:No normalization scale found for cyclone_exists_gaussian_unit_mode
+WARNING:absl:No coordinates found for dim noise_channels of array Size: 256B
+JitTracer(float32[2,1,32])
+Coordinates:
+ * batch (batch) object 16B MultiIndex
+ * _batch (batch) int64 16B 0 0
+ * sample (batch) int64 16B 0 1
+Dimensions without coordinates: channels, noise_channels., using index for param name. This may make model surgery more error prone if changing coordinates at fine-tuning time.
+/home/raymxnd/Workspace/WeatherForcast/.venv/lib/python3.12/site-packages/scipy/sparse/_index.py:216: SparseEfficiencyWarning: Changing the sparsity structure of a csr_matrix is expensive. lil and dok are more efficient.
+ self._set_arrayXarray(i, j, x)
+WARNING:absl:Only the first input of the static gather operation is being used, others are being ignored at runtime.
+WARNING:absl:Only the first input of the static gather operation is being used, others are being ignored at runtime.
+WARNING:absl:Only the first input of the sorted segment sum operation is being used, others are being ignored at runtime.
+WARNING:absl:Only the first input of the static gather operation is being used, others are being ignored at runtime.
+WARNING:absl:Only the first input of the static gather operation is being used, others are being ignored at runtime.
+WARNING:absl:Only the first input of the static gather operation is being used, others are being ignored at runtime.
+WARNING:absl:Only the first input of the sorted segment sum operation is being used, others are being ignored at runtime.
+/home/raymxnd/Workspace/WeatherForcast/weathernext/weathernext/weathernext2/architecture.py:286: FutureWarning: Deleting a single level of a MultiIndex is deprecated. Previously, this deleted all levels of a MultiIndex. Please also drop the following variables: {'batch'} to avoid an error in the future.
+ predictions[var] = xarray_jax.apply_ufunc(activation_fn, predictions[var])
+W1101 00:51:32.245035 53506 hlo_rematerialization.cc:3363] Can't reduce memory use below 22.07GiB (23696756432 bytes) by rematerialization; only reduced to 80.89GiB (86860081712 bytes), down from 101.77GiB (109270673384 bytes) originally
+W1101 00:51:46.658586 53506 bfc_allocator.cc:598] Allocator (GPU_0_bfc) ran out of memory trying to allocate 82.17GiB (rounded to 88231357440) from the lower end requested by op
+If this is caused by memory fragmentation, the environment variable 'TF_GPU_ALLOCATOR=cuda_malloc_async' may improve the situation.
+Current allocation summary follows.
+W1101 00:51:46.659841 53506 bfc_allocator.cc:609] ******_____________________________________________________________________________________________*
+ official default temp 82.17 GiB OUT OF MEMORY
+WARNING:root:No normalization location found for cyclone_exists_gaussian_unit_mode
+WARNING:root:No normalization scale found for cyclone_exists_gaussian_unit_mode
+WARNING:absl:No coordinates found for dim noise_channels of array Size: 256B
+JitTracer(float32[2,1,32])
+Coordinates:
+ * batch (batch) object 16B MultiIndex
+ * _batch (batch) int64 16B 0 0
+ * sample (batch) int64 16B 0 1
+Dimensions without coordinates: channels, noise_channels., using index for param name. This may make model surgery more error prone if changing coordinates at fine-tuning time.
+/home/raymxnd/Workspace/WeatherForcast/.venv/lib/python3.12/site-packages/scipy/sparse/_index.py:216: SparseEfficiencyWarning: Changing the sparsity structure of a csr_matrix is expensive. lil and dok are more efficient.
+ self._set_arrayXarray(i, j, x)
+WARNING:absl:Only the first input of the static gather operation is being used, others are being ignored at runtime.
+WARNING:absl:Only the first input of the static gather operation is being used, others are being ignored at runtime.
+WARNING:absl:Only the first input of the sorted segment sum operation is being used, others are being ignored at runtime.
+WARNING:absl:Only the first input of the static gather operation is being used, others are being ignored at runtime.
+WARNING:absl:Only the first input of the static gather operation is being used, others are being ignored at runtime.
+WARNING:absl:Only the first input of the static gather operation is being used, others are being ignored at runtime.
+WARNING:absl:Only the first input of the sorted segment sum operation is being used, others are being ignored at runtime.
+/home/raymxnd/Workspace/WeatherForcast/weathernext/weathernext/weathernext2/architecture.py:286: FutureWarning: Deleting a single level of a MultiIndex is deprecated. Previously, this deleted all levels of a MultiIndex. Please also drop the following variables: {'batch'} to avoid an error in the future.
+ predictions[var] = xarray_jax.apply_ufunc(activation_fn, predictions[var])
+W1101 00:52:11.409137 53506 hlo_rematerialization.cc:3363] Can't reduce memory use below 22.07GiB (23696756432 bytes) by rematerialization; only reduced to 105.97GiB (113781324928 bytes), down from 123.27GiB (132364702304 bytes) originally
+W1101 00:52:25.861104 53506 bfc_allocator.cc:598] Allocator (GPU_0_bfc) ran out of memory trying to allocate 106.67GiB (rounded to 114534348800) from the lower end requested by op
+If this is caused by memory fragmentation, the environment variable 'TF_GPU_ALLOCATOR=cuda_malloc_async' may improve the situation.
+Current allocation summary follows.
+W1101 00:52:25.862822 53506 bfc_allocator.cc:609] *******____________________________________________________________________________________________*
+ official_mha default temp 106.67 GiB OUT OF MEMORY
+WARNING:root:No normalization location found for cyclone_exists_gaussian_unit_mode
+WARNING:root:No normalization scale found for cyclone_exists_gaussian_unit_mode
+WARNING:absl:No coordinates found for dim noise_channels of array Size: 256B
+JitTracer(float32[2,1,32])
+Coordinates:
+ * batch (batch) object 16B MultiIndex
+ * _batch (batch) int64 16B 0 0
+ * sample (batch) int64 16B 0 1
+Dimensions without coordinates: channels, noise_channels., using index for param name. This may make model surgery more error prone if changing coordinates at fine-tuning time.
+/home/raymxnd/Workspace/WeatherForcast/.venv/lib/python3.12/site-packages/scipy/sparse/_index.py:216: SparseEfficiencyWarning: Changing the sparsity structure of a csr_matrix is expensive. lil and dok are more efficient.
+ self._set_arrayXarray(i, j, x)
+/home/raymxnd/Workspace/WeatherForcast/weathernext/weathernext/weathernext2/architecture.py:286: FutureWarning: Deleting a single level of a MultiIndex is deprecated. Previously, this deleted all levels of a MultiIndex. Please also drop the following variables: {'batch'} to avoid an error in the future.
+ predictions[var] = xarray_jax.apply_ufunc(activation_fn, predictions[var])
+ fwn default temp 4.01 GiB step 0.41s loss 1.203154
+WARNING:root:No normalization location found for cyclone_exists_gaussian_unit_mode
+WARNING:root:No normalization scale found for cyclone_exists_gaussian_unit_mode
+WARNING:absl:No coordinates found for dim noise_channels of array Size: 256B
+JitTracer(float32[2,1,32])
+Coordinates:
+ * batch (batch) object 16B MultiIndex
+ * _batch (batch) int64 16B 0 0
+ * sample (batch) int64 16B 0 1
+Dimensions without coordinates: channels, noise_channels., using index for param name. This may make model surgery more error prone if changing coordinates at fine-tuning time.
+/home/raymxnd/Workspace/WeatherForcast/.venv/lib/python3.12/site-packages/scipy/sparse/_index.py:216: SparseEfficiencyWarning: Changing the sparsity structure of a csr_matrix is expensive. lil and dok are more efficient.
+ self._set_arrayXarray(i, j, x)
+/home/raymxnd/Workspace/WeatherForcast/weathernext/weathernext/weathernext2/architecture.py:286: FutureWarning: Deleting a single level of a MultiIndex is deprecated. Previously, this deleted all levels of a MultiIndex. Please also drop the following variables: {'batch'} to avoid an error in the future.
+ predictions[var] = xarray_jax.apply_ufunc(activation_fn, predictions[var])
+ fwn_xla default temp 4.16 GiB step 0.89s loss 1.203162
+WARNING:root:No normalization location found for cyclone_exists_gaussian_unit_mode
+WARNING:root:No normalization scale found for cyclone_exists_gaussian_unit_mode
+WARNING:absl:No coordinates found for dim noise_channels of array Size: 256B
+JitTracer(float32[2,1,32])
+Coordinates:
+ * batch (batch) object 16B MultiIndex
+ * _batch (batch) int64 16B 0 0
+ * sample (batch) int64 16B 0 1
+Dimensions without coordinates: channels, noise_channels., using index for param name. This may make model surgery more error prone if changing coordinates at fine-tuning time.
+/home/raymxnd/Workspace/WeatherForcast/.venv/lib/python3.12/site-packages/scipy/sparse/_index.py:216: SparseEfficiencyWarning: Changing the sparsity structure of a csr_matrix is expensive. lil and dok are more efficient.
+ self._set_arrayXarray(i, j, x)
+/home/raymxnd/Workspace/WeatherForcast/weathernext/weathernext/weathernext2/architecture.py:286: FutureWarning: Deleting a single level of a MultiIndex is deprecated. Previously, this deleted all levels of a MultiIndex. Please also drop the following variables: {'batch'} to avoid an error in the future.
+ predictions[var] = xarray_jax.apply_ufunc(activation_fn, predictions[var])
+W1101 00:53:51.440238 53506 hlo_rematerialization.cc:3363] Can't reduce memory use below 22.07GiB (23696756432 bytes) by rematerialization; only reduced to 39.12GiB (42009086092 bytes), down from 39.17GiB (42063857529 bytes) originally
+ fwn_noremat default temp 38.95 GiB (compile only)
+ fwn vs fwn_xla loss rel diff 7.03e-06 (per var max 2.88e-03) grads: global rel L2 4.41e-02 cos 0.9990259515 max per-param rel L2 1.82e-01 (335/376 params; zero in both 6, zero only in official 0) finite True
+ 1.82e-01 multimodality_forward/mesh_transformer/transformer/block_07/block_07_norm_conditioning/linear/b
+ 1.72e-01 multimodality_forward/mesh_transformer/transformer/block_01/block_01_norm_conditioning/linear/b
+ 1.50e-01 multimodality_forward/grid_to_mesh_gnn/edge_encoder/shared_dense/mlp/linear_0/b
+wrote docs/results/finetune/full_grads_mini_crps_rtx5090.json
diff --git a/docs/results/finetune/full_grads_mini_linear_rtx5090.json b/docs/results/finetune/full_grads_mini_linear_rtx5090.json
new file mode 100644
index 0000000..570a454
--- /dev/null
+++ b/docs/results/finetune/full_grads_mini_linear_rtx5090.json
@@ -0,0 +1,259 @@
+{
+ "model": "WeatherNextCyclones_Mini",
+ "device": "NVIDIA GeForce RTX 5090",
+ "jax": "0.11.2",
+ "faster_weathernext": "0.1.1",
+ "blocks": {
+ "grid_chunk": 8192,
+ "edge_chunk": 16384
+ },
+ "masked": [],
+ "modules": "all",
+ "loss": "linear",
+ "head_params": 56684133,
+ "precisions": {
+ "highest": {
+ "official": {
+ "temp_gib": 45.37405942380428,
+ "compile_s": 75.64562153816223,
+ "oom": "RESOURCE_EXHAUSTED: Out of memory while trying to allocate 45.37GiB with allocator GPU_0_bfc on device 0. [tf-allocator-allocation-error=''] [executable_name='jit_fn']"
+ },
+ "official_mha": {
+ "temp_gib": 56.799837335944176,
+ "compile_s": 32.98060917854309,
+ "oom": "RESOURCE_EXHAUSTED: Out of memory while trying to allocate 56.80GiB with allocator GPU_0_bfc on device 0. [tf-allocator-allocation-error=''] [executable_name='jit_fn']"
+ },
+ "fwn": {
+ "temp_gib": 3.258644051849842,
+ "compile_s": 60.61225605010986,
+ "step_s": 0.4557206630706787
+ },
+ "fwn_xla": {
+ "temp_gib": 3.383711516857147,
+ "compile_s": 30.882630586624146,
+ "step_s": 0.6089086532592773
+ },
+ "fwn_strict": {
+ "temp_gib": 3.258644051849842,
+ "compile_s": 19.047070741653442,
+ "step_s": 1.497154951095581
+ },
+ "fwn_noremat": {
+ "temp_gib": 19.638427458703518,
+ "compile_s": 20.204537868499756
+ },
+ "reference": "fwn_xla",
+ "checks": {
+ "fwn vs fwn_xla": {
+ "loss": {
+ "official": 506.6861572265625,
+ "other": 506.6861877441406,
+ "rel_diff": 6.022974515831148e-08,
+ "max_per_var_rel_diff": 0.0003032060775157625
+ },
+ "grads": {
+ "global_rel_l2": 0.0002508788890675034,
+ "global_cos": 0.9999999764855787,
+ "max_rel_l2": 0.0012514083634297562,
+ "n_params": 376,
+ "n_compared": 315,
+ "n_zero_both": 6,
+ "n_zero_ref_only": 0,
+ "all_finite": true,
+ "worst": [
+ {
+ "param": "multimodality_forward/mesh_transformer/transformer/block_06/mha_proj_k/w",
+ "size": 262144,
+ "ref_norm": 2.7120831628078528,
+ "got_norm": 2.7136186826774487,
+ "finite": true,
+ "rel_l2": 0.0012514083634297562
+ },
+ {
+ "param": "multimodality_forward/mesh_transformer/transformer/block_02/block_02_norm_conditioning_1/linear/w",
+ "size": 32768,
+ "ref_norm": 0.6124617058595232,
+ "got_norm": 0.6129337448134701,
+ "finite": true,
+ "rel_l2": 0.0008403606200358301
+ },
+ {
+ "param": "multimodality_forward/mesh_transformer/transformer/block_02/block_02_norm_conditioning_1/linear/b",
+ "size": 1024,
+ "ref_norm": 0.7096776081980437,
+ "got_norm": 0.7102245703081098,
+ "finite": true,
+ "rel_l2": 0.0008403565468236138
+ },
+ {
+ "param": "multimodality_forward/global_norm_conditioning_encoder/split_input_matmul/w_input_noise_noise_channels=range(32)",
+ "size": 1024,
+ "ref_norm": 23.969418629655788,
+ "got_norm": 23.98620849168166,
+ "finite": true,
+ "rel_l2": 0.0008002067917544567
+ },
+ {
+ "param": "multimodality_forward/mesh_transformer/transformer/block_07/block_07_norm_conditioning/linear/w",
+ "size": 32768,
+ "ref_norm": 0.49396599475904734,
+ "got_norm": 0.49413754179569436,
+ "finite": true,
+ "rel_l2": 0.0007814520189398472
+ }
+ ]
+ }
+ },
+ "fwn_strict vs fwn_xla": {
+ "loss": {
+ "official": 506.6861572265625,
+ "other": 506.6861877441406,
+ "rel_diff": 6.022974515831148e-08,
+ "max_per_var_rel_diff": 7.215591727129653e-05
+ },
+ "grads": {
+ "global_rel_l2": 7.00702748801237e-05,
+ "global_cos": 0.9999999975569918,
+ "max_rel_l2": 0.00028294166195617026,
+ "n_params": 376,
+ "n_compared": 315,
+ "n_zero_both": 6,
+ "n_zero_ref_only": 0,
+ "all_finite": true,
+ "worst": [
+ {
+ "param": "multimodality_forward/mesh_transformer/transformer/block_07/block_07_norm_conditioning/linear/b",
+ "size": 1024,
+ "ref_norm": 0.5723730995779387,
+ "got_norm": 0.5723946947170323,
+ "finite": true,
+ "rel_l2": 0.00028294166195617026
+ },
+ {
+ "param": "multimodality_forward/mesh_transformer/transformer/block_07/block_07_norm_conditioning/linear/w",
+ "size": 32768,
+ "ref_norm": 0.49396599475904734,
+ "got_norm": 0.49398463229563155,
+ "finite": true,
+ "rel_l2": 0.0002829416078077201
+ },
+ {
+ "param": "multimodality_forward/mesh_transformer/transformer/block_04/block_04_norm_conditioning/linear/b",
+ "size": 1024,
+ "ref_norm": 0.7267677300240337,
+ "got_norm": 0.7267119887138631,
+ "finite": true,
+ "rel_l2": 0.00027242850582178893
+ },
+ {
+ "param": "multimodality_forward/mesh_transformer/transformer/block_04/block_04_norm_conditioning/linear/w",
+ "size": 32768,
+ "ref_norm": 0.6272107223622873,
+ "got_norm": 0.6271626172304601,
+ "finite": true,
+ "rel_l2": 0.00027242841682077263
+ },
+ {
+ "param": "multimodality_forward/mesh_transformer/transformer/block_04/block_04_norm_conditioning_1/linear/b",
+ "size": 1024,
+ "ref_norm": 0.5775905981194651,
+ "got_norm": 0.5775921518521128,
+ "finite": true,
+ "rel_l2": 0.00022276233266231747
+ }
+ ]
+ }
+ }
+ }
+ },
+ "default": {
+ "official": {
+ "temp_gib": 33.92704151570797,
+ "compile_s": 44.175304651260376,
+ "oom": "RESOURCE_EXHAUSTED: Out of memory while trying to allocate 33.93GiB with allocator GPU_0_bfc on device 0. [tf-allocator-allocation-error=''] [executable_name='jit_fn']"
+ },
+ "official_mha": {
+ "temp_gib": 32.61181591451168,
+ "compile_s": 22.591861963272095,
+ "oom": "RESOURCE_EXHAUSTED: Out of memory while trying to allocate 32.61GiB with allocator GPU_0_bfc on device 0. [executable_name='jit_fn'] [tf-allocator-allocation-error='']"
+ },
+ "fwn": {
+ "temp_gib": 3.258623071014881,
+ "compile_s": 35.82840013504028,
+ "step_s": 0.2399282455444336
+ },
+ "fwn_xla": {
+ "temp_gib": 3.3798927664756775,
+ "compile_s": 25.616586446762085,
+ "step_s": 0.4442577362060547
+ },
+ "fwn_noremat": {
+ "temp_gib": 19.62287374585867,
+ "compile_s": 17.91448426246643
+ },
+ "reference": "fwn_xla",
+ "checks": {
+ "fwn vs fwn_xla": {
+ "loss": {
+ "official": 506.6859130859375,
+ "other": 506.68756103515625,
+ "rel_diff": 3.252407805682366e-06,
+ "max_per_var_rel_diff": 0.07683305534178157
+ },
+ "grads": {
+ "global_rel_l2": 0.028688181598863615,
+ "global_cos": 0.9995909616423241,
+ "max_rel_l2": 0.0820666610018653,
+ "n_params": 376,
+ "n_compared": 313,
+ "n_zero_both": 6,
+ "n_zero_ref_only": 0,
+ "all_finite": true,
+ "worst": [
+ {
+ "param": "multimodality_forward/mesh_encoder/split_input_matmul/w_input_year_progress_cos_time=0",
+ "size": 512,
+ "ref_norm": 2.351501000231752,
+ "got_norm": 2.3303718952293266,
+ "finite": true,
+ "rel_l2": 0.0820666610018653
+ },
+ {
+ "param": "multimodality_forward/mesh_encoder/split_input_matmul/w_input_year_progress_sin_time=-21600",
+ "size": 512,
+ "ref_norm": 21.78868875294578,
+ "got_norm": 21.592909116932475,
+ "finite": true,
+ "rel_l2": 0.08206663218437225
+ },
+ {
+ "param": "multimodality_forward/mesh_encoder/split_input_matmul/w_input_year_progress_sin_time=0",
+ "size": 512,
+ "ref_norm": 21.78868875294578,
+ "got_norm": 21.592909116932475,
+ "finite": true,
+ "rel_l2": 0.08206663218437225
+ },
+ {
+ "param": "multimodality_forward/mesh_encoder/split_input_matmul/w_input_year_progress_cos_time=-21600",
+ "size": 512,
+ "ref_norm": 2.257063190882897,
+ "got_norm": 2.236782741491796,
+ "finite": true,
+ "rel_l2": 0.08206662769641662
+ },
+ {
+ "param": "multimodality_forward/mesh_encoder/split_input_matmul/w_forcing_year_progress_cos_time=21600",
+ "size": 512,
+ "ref_norm": 2.4440499873566526,
+ "got_norm": 2.42208932735907,
+ "finite": true,
+ "rel_l2": 0.08206660768292522
+ }
+ ]
+ }
+ }
+ }
+ }
+ }
+}
\ No newline at end of file
diff --git a/docs/results/finetune/full_grads_mini_linear_rtx5090.log b/docs/results/finetune/full_grads_mini_linear_rtx5090.log
new file mode 100644
index 0000000..6eca5d0
--- /dev/null
+++ b/docs/results/finetune/full_grads_mini_linear_rtx5090.log
@@ -0,0 +1,171 @@
+WeatherNextCyclones_Mini: differentiating 56.68M params (all); masked targets: []
+[highest]
+WARNING:absl:Skipping gradient checkpointing for sequence length of 1
+WARNING:absl:No coordinates found for dim noise_channels of array Size: 128B
+JitTracer(float32[1,1,32])
+Dimensions without coordinates: batch, channels, noise_channels., using index for param name. This may make model surgery more error prone if changing coordinates at fine-tuning time.
+/home/raymxnd/Workspace/WeatherForcast/.venv/lib/python3.12/site-packages/scipy/sparse/_index.py:216: SparseEfficiencyWarning: Changing the sparsity structure of a csr_matrix is expensive. lil and dok are more efficient.
+ self._set_arrayXarray(i, j, x)
+WARNING:absl:Only the first input of the static gather operation is being used, others are being ignored at runtime.
+WARNING:absl:Only the first input of the static gather operation is being used, others are being ignored at runtime.
+WARNING:absl:Only the first input of the sorted segment sum operation is being used, others are being ignored at runtime.
+WARNING:absl:Only the first input of the static gather operation is being used, others are being ignored at runtime.
+WARNING:absl:Only the first input of the static gather operation is being used, others are being ignored at runtime.
+WARNING:absl:Only the first input of the static gather operation is being used, others are being ignored at runtime.
+WARNING:absl:Only the first input of the sorted segment sum operation is being used, others are being ignored at runtime.
+WARNING:root:No normalization scale found for cyclone_exists_gaussian_unit_mode
+WARNING:root:No normalization location found for cyclone_exists_gaussian_unit_mode
+WARNING:absl:Only the first input of the sorted segment sum operation is being used, others are being ignored at runtime.
+WARNING:absl:Only the first input of the sorted segment sum operation is being used, others are being ignored at runtime.
+W1101 00:55:11.031797 88980 hlo_rematerialization.cc:3363] Can't reduce memory use below 22.09GiB (23721764840 bytes) by rematerialization; only reduced to 45.19GiB (48518250448 bytes), down from 49.91GiB (53588660032 bytes) originally
+W1101 00:55:23.747208 88980 bfc_allocator.cc:598] Allocator (GPU_0_bfc) ran out of memory trying to allocate 45.37GiB (rounded to 48720025344) from the lower end requested by op
+If this is caused by memory fragmentation, the environment variable 'TF_GPU_ALLOCATOR=cuda_malloc_async' may improve the situation.
+Current allocation summary follows.
+W1101 00:55:23.747813 88980 bfc_allocator.cc:609] ************************___________________________________________________________________________*
+ official highest temp 45.37 GiB OUT OF MEMORY
+WARNING:absl:Skipping gradient checkpointing for sequence length of 1
+WARNING:absl:No coordinates found for dim noise_channels of array Size: 128B
+JitTracer(float32[1,1,32])
+Dimensions without coordinates: batch, channels, noise_channels., using index for param name. This may make model surgery more error prone if changing coordinates at fine-tuning time.
+/home/raymxnd/Workspace/WeatherForcast/.venv/lib/python3.12/site-packages/scipy/sparse/_index.py:216: SparseEfficiencyWarning: Changing the sparsity structure of a csr_matrix is expensive. lil and dok are more efficient.
+ self._set_arrayXarray(i, j, x)
+WARNING:absl:Only the first input of the static gather operation is being used, others are being ignored at runtime.
+WARNING:absl:Only the first input of the static gather operation is being used, others are being ignored at runtime.
+WARNING:absl:Only the first input of the sorted segment sum operation is being used, others are being ignored at runtime.
+WARNING:absl:Only the first input of the static gather operation is being used, others are being ignored at runtime.
+WARNING:absl:Only the first input of the static gather operation is being used, others are being ignored at runtime.
+WARNING:absl:Only the first input of the static gather operation is being used, others are being ignored at runtime.
+WARNING:absl:Only the first input of the sorted segment sum operation is being used, others are being ignored at runtime.
+WARNING:root:No normalization scale found for cyclone_exists_gaussian_unit_mode
+WARNING:root:No normalization location found for cyclone_exists_gaussian_unit_mode
+WARNING:absl:Only the first input of the sorted segment sum operation is being used, others are being ignored at runtime.
+WARNING:absl:Only the first input of the sorted segment sum operation is being used, others are being ignored at runtime.
+W1101 00:55:55.026165 88980 hlo_rematerialization.cc:3363] Can't reduce memory use below 22.09GiB (23721764840 bytes) by rematerialization; only reduced to 57.14GiB (61357027880 bytes), down from 60.68GiB (65150375352 bytes) originally
+W1101 00:56:07.194758 88980 bfc_allocator.cc:598] Allocator (GPU_0_bfc) ran out of memory trying to allocate 56.80GiB (rounded to 60988360960) from the lower end requested by op
+If this is caused by memory fragmentation, the environment variable 'TF_GPU_ALLOCATOR=cuda_malloc_async' may improve the situation.
+Current allocation summary follows.
+W1101 00:56:07.195664 88980 bfc_allocator.cc:609] *********__________________________________________________________________________________________*
+ official_mha highest temp 56.80 GiB OUT OF MEMORY
+WARNING:absl:Skipping gradient checkpointing for sequence length of 1
+WARNING:absl:No coordinates found for dim noise_channels of array Size: 128B
+JitTracer(float32[1,1,32])
+Dimensions without coordinates: batch, channels, noise_channels., using index for param name. This may make model surgery more error prone if changing coordinates at fine-tuning time.
+/home/raymxnd/Workspace/WeatherForcast/.venv/lib/python3.12/site-packages/scipy/sparse/_index.py:216: SparseEfficiencyWarning: Changing the sparsity structure of a csr_matrix is expensive. lil and dok are more efficient.
+ self._set_arrayXarray(i, j, x)
+WARNING:root:No normalization scale found for cyclone_exists_gaussian_unit_mode
+WARNING:root:No normalization location found for cyclone_exists_gaussian_unit_mode
+ fwn highest temp 3.26 GiB step 0.46s loss 506.686188
+WARNING:absl:Skipping gradient checkpointing for sequence length of 1
+WARNING:absl:No coordinates found for dim noise_channels of array Size: 128B
+JitTracer(float32[1,1,32])
+Dimensions without coordinates: batch, channels, noise_channels., using index for param name. This may make model surgery more error prone if changing coordinates at fine-tuning time.
+/home/raymxnd/Workspace/WeatherForcast/.venv/lib/python3.12/site-packages/scipy/sparse/_index.py:216: SparseEfficiencyWarning: Changing the sparsity structure of a csr_matrix is expensive. lil and dok are more efficient.
+ self._set_arrayXarray(i, j, x)
+WARNING:root:No normalization scale found for cyclone_exists_gaussian_unit_mode
+WARNING:root:No normalization location found for cyclone_exists_gaussian_unit_mode
+ fwn_xla highest temp 3.38 GiB step 0.61s loss 506.686157
+WARNING:absl:Skipping gradient checkpointing for sequence length of 1
+WARNING:absl:No coordinates found for dim noise_channels of array Size: 128B
+JitTracer(float32[1,1,32])
+Dimensions without coordinates: batch, channels, noise_channels., using index for param name. This may make model surgery more error prone if changing coordinates at fine-tuning time.
+/home/raymxnd/Workspace/WeatherForcast/.venv/lib/python3.12/site-packages/scipy/sparse/_index.py:216: SparseEfficiencyWarning: Changing the sparsity structure of a csr_matrix is expensive. lil and dok are more efficient.
+ self._set_arrayXarray(i, j, x)
+WARNING:root:No normalization scale found for cyclone_exists_gaussian_unit_mode
+WARNING:root:No normalization location found for cyclone_exists_gaussian_unit_mode
+ fwn_strict highest temp 3.26 GiB step 1.50s loss 506.686188
+WARNING:absl:Skipping gradient checkpointing for sequence length of 1
+WARNING:absl:No coordinates found for dim noise_channels of array Size: 128B
+JitTracer(float32[1,1,32])
+Dimensions without coordinates: batch, channels, noise_channels., using index for param name. This may make model surgery more error prone if changing coordinates at fine-tuning time.
+/home/raymxnd/Workspace/WeatherForcast/.venv/lib/python3.12/site-packages/scipy/sparse/_index.py:216: SparseEfficiencyWarning: Changing the sparsity structure of a csr_matrix is expensive. lil and dok are more efficient.
+ self._set_arrayXarray(i, j, x)
+WARNING:root:No normalization scale found for cyclone_exists_gaussian_unit_mode
+WARNING:root:No normalization location found for cyclone_exists_gaussian_unit_mode
+ fwn_noremat highest temp 19.64 GiB (compile only)
+ fwn vs fwn_xla loss rel diff 6.02e-08 (per var max 3.03e-04) grads: global rel L2 2.51e-04 cos 0.9999999765 max per-param rel L2 1.25e-03 (315/376 params; zero in both 6, zero only in official 0) finite True
+ 1.25e-03 multimodality_forward/mesh_transformer/transformer/block_06/mha_proj_k/w
+ 8.40e-04 multimodality_forward/mesh_transformer/transformer/block_02/block_02_norm_conditioning_1/linear/w
+ 8.40e-04 multimodality_forward/mesh_transformer/transformer/block_02/block_02_norm_conditioning_1/linear/b
+ fwn_strict vs fwn_xla loss rel diff 6.02e-08 (per var max 7.22e-05) grads: global rel L2 7.01e-05 cos 0.9999999976 max per-param rel L2 2.83e-04 (315/376 params; zero in both 6, zero only in official 0) finite True
+ 2.83e-04 multimodality_forward/mesh_transformer/transformer/block_07/block_07_norm_conditioning/linear/b
+ 2.83e-04 multimodality_forward/mesh_transformer/transformer/block_07/block_07_norm_conditioning/linear/w
+ 2.72e-04 multimodality_forward/mesh_transformer/transformer/block_04/block_04_norm_conditioning/linear/b
+[default]
+WARNING:absl:Skipping gradient checkpointing for sequence length of 1
+WARNING:absl:No coordinates found for dim noise_channels of array Size: 128B
+JitTracer(float32[1,1,32])
+Dimensions without coordinates: batch, channels, noise_channels., using index for param name. This may make model surgery more error prone if changing coordinates at fine-tuning time.
+/home/raymxnd/Workspace/WeatherForcast/.venv/lib/python3.12/site-packages/scipy/sparse/_index.py:216: SparseEfficiencyWarning: Changing the sparsity structure of a csr_matrix is expensive. lil and dok are more efficient.
+ self._set_arrayXarray(i, j, x)
+WARNING:absl:Only the first input of the static gather operation is being used, others are being ignored at runtime.
+WARNING:absl:Only the first input of the static gather operation is being used, others are being ignored at runtime.
+WARNING:absl:Only the first input of the sorted segment sum operation is being used, others are being ignored at runtime.
+WARNING:absl:Only the first input of the static gather operation is being used, others are being ignored at runtime.
+WARNING:absl:Only the first input of the static gather operation is being used, others are being ignored at runtime.
+WARNING:absl:Only the first input of the static gather operation is being used, others are being ignored at runtime.
+WARNING:absl:Only the first input of the sorted segment sum operation is being used, others are being ignored at runtime.
+WARNING:root:No normalization scale found for cyclone_exists_gaussian_unit_mode
+WARNING:root:No normalization location found for cyclone_exists_gaussian_unit_mode
+WARNING:absl:Only the first input of the sorted segment sum operation is being used, others are being ignored at runtime.
+WARNING:absl:Only the first input of the sorted segment sum operation is being used, others are being ignored at runtime.
+W1101 00:59:07.509626 88980 hlo_rematerialization.cc:3363] Can't reduce memory use below 22.09GiB (23721764840 bytes) by rematerialization; only reduced to 33.71GiB (36192906848 bytes), down from 49.91GiB (53588660032 bytes) originally
+W1101 00:59:20.385376 88980 bfc_allocator.cc:598] Allocator (GPU_0_bfc) ran out of memory trying to allocate 33.93GiB (rounded to 36428883456) from the lower end requested by op
+If this is caused by memory fragmentation, the environment variable 'TF_GPU_ALLOCATOR=cuda_malloc_async' may improve the situation.
+Current allocation summary follows.
+W1101 00:59:20.386618 88980 bfc_allocator.cc:609] ************_______________________________________________________________________________________*
+ official default temp 33.93 GiB OUT OF MEMORY
+WARNING:absl:Skipping gradient checkpointing for sequence length of 1
+WARNING:absl:No coordinates found for dim noise_channels of array Size: 128B
+JitTracer(float32[1,1,32])
+Dimensions without coordinates: batch, channels, noise_channels., using index for param name. This may make model surgery more error prone if changing coordinates at fine-tuning time.
+/home/raymxnd/Workspace/WeatherForcast/.venv/lib/python3.12/site-packages/scipy/sparse/_index.py:216: SparseEfficiencyWarning: Changing the sparsity structure of a csr_matrix is expensive. lil and dok are more efficient.
+ self._set_arrayXarray(i, j, x)
+WARNING:absl:Only the first input of the static gather operation is being used, others are being ignored at runtime.
+WARNING:absl:Only the first input of the static gather operation is being used, others are being ignored at runtime.
+WARNING:absl:Only the first input of the sorted segment sum operation is being used, others are being ignored at runtime.
+WARNING:absl:Only the first input of the static gather operation is being used, others are being ignored at runtime.
+WARNING:absl:Only the first input of the static gather operation is being used, others are being ignored at runtime.
+WARNING:absl:Only the first input of the static gather operation is being used, others are being ignored at runtime.
+WARNING:absl:Only the first input of the sorted segment sum operation is being used, others are being ignored at runtime.
+WARNING:root:No normalization scale found for cyclone_exists_gaussian_unit_mode
+WARNING:root:No normalization location found for cyclone_exists_gaussian_unit_mode
+WARNING:absl:Only the first input of the sorted segment sum operation is being used, others are being ignored at runtime.
+WARNING:absl:Only the first input of the sorted segment sum operation is being used, others are being ignored at runtime.
+W1101 00:59:41.361694 88980 hlo_rematerialization.cc:3363] Can't reduce memory use below 22.09GiB (23721764840 bytes) by rematerialization; only reduced to 32.15GiB (34523315672 bytes), down from 60.68GiB (65150375352 bytes) originally
+W1101 00:59:53.477253 88980 bfc_allocator.cc:598] Allocator (GPU_0_bfc) ran out of memory trying to allocate 32.61GiB (rounded to 35016670720) from the lower end requested by op
+If this is caused by memory fragmentation, the environment variable 'TF_GPU_ALLOCATOR=cuda_malloc_async' may improve the situation.
+Current allocation summary follows.
+W1101 00:59:53.478832 88980 bfc_allocator.cc:609] **************_____________________________________________________________________________________*
+ official_mha default temp 32.61 GiB OUT OF MEMORY
+WARNING:absl:Skipping gradient checkpointing for sequence length of 1
+WARNING:absl:No coordinates found for dim noise_channels of array Size: 128B
+JitTracer(float32[1,1,32])
+Dimensions without coordinates: batch, channels, noise_channels., using index for param name. This may make model surgery more error prone if changing coordinates at fine-tuning time.
+/home/raymxnd/Workspace/WeatherForcast/.venv/lib/python3.12/site-packages/scipy/sparse/_index.py:216: SparseEfficiencyWarning: Changing the sparsity structure of a csr_matrix is expensive. lil and dok are more efficient.
+ self._set_arrayXarray(i, j, x)
+WARNING:root:No normalization scale found for cyclone_exists_gaussian_unit_mode
+WARNING:root:No normalization location found for cyclone_exists_gaussian_unit_mode
+ fwn default temp 3.26 GiB step 0.24s loss 506.687561
+WARNING:absl:Skipping gradient checkpointing for sequence length of 1
+WARNING:absl:No coordinates found for dim noise_channels of array Size: 128B
+JitTracer(float32[1,1,32])
+Dimensions without coordinates: batch, channels, noise_channels., using index for param name. This may make model surgery more error prone if changing coordinates at fine-tuning time.
+/home/raymxnd/Workspace/WeatherForcast/.venv/lib/python3.12/site-packages/scipy/sparse/_index.py:216: SparseEfficiencyWarning: Changing the sparsity structure of a csr_matrix is expensive. lil and dok are more efficient.
+ self._set_arrayXarray(i, j, x)
+WARNING:root:No normalization scale found for cyclone_exists_gaussian_unit_mode
+WARNING:root:No normalization location found for cyclone_exists_gaussian_unit_mode
+ fwn_xla default temp 3.38 GiB step 0.44s loss 506.685913
+WARNING:absl:Skipping gradient checkpointing for sequence length of 1
+WARNING:absl:No coordinates found for dim noise_channels of array Size: 128B
+JitTracer(float32[1,1,32])
+Dimensions without coordinates: batch, channels, noise_channels., using index for param name. This may make model surgery more error prone if changing coordinates at fine-tuning time.
+/home/raymxnd/Workspace/WeatherForcast/.venv/lib/python3.12/site-packages/scipy/sparse/_index.py:216: SparseEfficiencyWarning: Changing the sparsity structure of a csr_matrix is expensive. lil and dok are more efficient.
+ self._set_arrayXarray(i, j, x)
+WARNING:root:No normalization scale found for cyclone_exists_gaussian_unit_mode
+WARNING:root:No normalization location found for cyclone_exists_gaussian_unit_mode
+ fwn_noremat default temp 19.62 GiB (compile only)
+ fwn vs fwn_xla loss rel diff 3.25e-06 (per var max 7.68e-02) grads: global rel L2 2.87e-02 cos 0.9995909616 max per-param rel L2 8.21e-02 (313/376 params; zero in both 6, zero only in official 0) finite True
+ 8.21e-02 multimodality_forward/mesh_encoder/split_input_matmul/w_input_year_progress_cos_time=0
+ 8.21e-02 multimodality_forward/mesh_encoder/split_input_matmul/w_input_year_progress_sin_time=-21600
+ 8.21e-02 multimodality_forward/mesh_encoder/split_input_matmul/w_input_year_progress_sin_time=0
+wrote docs/results/finetune/full_grads_mini_linear_rtx5090.json
diff --git a/docs/results/finetune/h200/grads01/full_grads_mini_crps_h200.json b/docs/results/finetune/h200/grads01/full_grads_mini_crps_h200.json
new file mode 100644
index 0000000..1c740b0
--- /dev/null
+++ b/docs/results/finetune/h200/grads01/full_grads_mini_crps_h200.json
@@ -0,0 +1,499 @@
+{
+ "model": "WeatherNextCyclones_Mini",
+ "device": "NVIDIA H200",
+ "jax": "0.11.2",
+ "faster_weathernext": "0.1.1",
+ "blocks": {
+ "grid_chunk": 8192,
+ "edge_chunk": 16384
+ },
+ "masked": [],
+ "modules": "all",
+ "loss": "crps",
+ "head_params": 56684133,
+ "precisions": {
+ "highest": {
+ "official": {
+ "temp_gib": 102.06442984193563,
+ "compile_s": 72.90124654769897,
+ "step_s": 0.8720457553863525
+ },
+ "official_mha": {
+ "temp_gib": 123.4894026145339,
+ "compile_s": 40.70241355895996,
+ "step_s": 1.066147804260254
+ },
+ "fwn": {
+ "temp_gib": 3.9678899720311165,
+ "compile_s": 64.52860021591187,
+ "step_s": 1.1164429187774658
+ },
+ "fwn_xla": {
+ "temp_gib": 4.092704474925995,
+ "compile_s": 45.27368521690369,
+ "step_s": 1.2497541904449463
+ },
+ "fwn_strict": {
+ "temp_gib": 3.9678899720311165,
+ "compile_s": 32.91136646270752,
+ "step_s": 1.585862636566162
+ },
+ "fwn_noremat": {
+ "temp_gib": 39.10420750826597,
+ "compile_s": 35.25035381317139
+ },
+ "reference": "official",
+ "checks": {
+ "fwn vs official": {
+ "loss": {
+ "official": 1.2029576301574707,
+ "other": 1.2029577493667603,
+ "rel_diff": 9.909683147791032e-08,
+ "max_per_var_rel_diff": 1.5409268496342365e-05
+ },
+ "grads": {
+ "global_rel_l2": 0.004745261834459455,
+ "global_cos": 0.9999887413153253,
+ "max_rel_l2": 0.028065166999121486,
+ "n_params": 376,
+ "n_compared": 335,
+ "n_zero_both": 6,
+ "n_zero_ref_only": 0,
+ "all_finite": true,
+ "worst": [
+ {
+ "param": "multimodality_forward/mesh_transformer/transformer/block_06/mha_proj_k/w",
+ "size": 262144,
+ "ref_norm": 0.0802299924765987,
+ "got_norm": 0.07943947643459336,
+ "finite": true,
+ "rel_l2": 0.028065166999121486
+ },
+ {
+ "param": "multimodality_forward/mesh_transformer/transformer/block_06/mha_proj_q/w",
+ "size": 262144,
+ "ref_norm": 0.05680076313263051,
+ "got_norm": 0.05661605399626702,
+ "finite": true,
+ "rel_l2": 0.020579371017718732
+ },
+ {
+ "param": "multimodality_forward/mesh_transformer/transformer/block_04/mha_proj_k/w",
+ "size": 262144,
+ "ref_norm": 0.15894120122711766,
+ "got_norm": 0.15996246893525268,
+ "finite": true,
+ "rel_l2": 0.019283808264596927
+ },
+ {
+ "param": "multimodality_forward/mesh_transformer/transformer/block_04/block_04_norm_conditioning/linear/b",
+ "size": 1024,
+ "ref_norm": 0.022045770859437264,
+ "got_norm": 0.022266419049186654,
+ "finite": true,
+ "rel_l2": 0.01699978804304756
+ },
+ {
+ "param": "multimodality_forward/mesh_transformer/transformer/block_07/block_07_norm_conditioning/linear/b",
+ "size": 1024,
+ "ref_norm": 0.011979358507098819,
+ "got_norm": 0.011974244395912798,
+ "finite": true,
+ "rel_l2": 0.01698993963066057
+ }
+ ]
+ }
+ },
+ "fwn_strict vs official": {
+ "loss": {
+ "official": 1.2029576301574707,
+ "other": 1.2029577493667603,
+ "rel_diff": 9.909683147791032e-08,
+ "max_per_var_rel_diff": 1.5945883103335674e-05
+ },
+ "grads": {
+ "global_rel_l2": 0.0010636326160204544,
+ "global_cos": 0.9999994351273652,
+ "max_rel_l2": 0.002938606042946185,
+ "n_params": 376,
+ "n_compared": 335,
+ "n_zero_both": 6,
+ "n_zero_ref_only": 0,
+ "all_finite": true,
+ "worst": [
+ {
+ "param": "multimodality_forward/mesh_transformer/transformer/block_06/block_06_norm_conditioning_1/linear/b",
+ "size": 1024,
+ "ref_norm": 0.011992401205959081,
+ "got_norm": 0.011981994482399536,
+ "finite": true,
+ "rel_l2": 0.002938606042946185
+ },
+ {
+ "param": "multimodality_forward/mesh_transformer/transformer/block_07/block_07_norm_conditioning/linear/b",
+ "size": 1024,
+ "ref_norm": 0.011979358507098819,
+ "got_norm": 0.011980907581789827,
+ "finite": true,
+ "rel_l2": 0.0027818204199755846
+ },
+ {
+ "param": "multimodality_forward/mesh_transformer/transformer/block_06/block_06_norm_conditioning_1/linear/w",
+ "size": 32768,
+ "ref_norm": 0.0146504758882563,
+ "got_norm": 0.01464255189852539,
+ "finite": true,
+ "rel_l2": 0.002694811886616018
+ },
+ {
+ "param": "multimodality_forward/mesh_encoder/split_input_matmul/w_spatial_feature=sin_lat,sin_lon,cos_lon",
+ "size": 1536,
+ "ref_norm": 0.3616984147465042,
+ "got_norm": 0.36149161557186815,
+ "finite": true,
+ "rel_l2": 0.002547896638937096
+ },
+ {
+ "param": "multimodality_forward/mesh_transformer/transformer/block_04/block_04_norm_conditioning/linear/b",
+ "size": 1024,
+ "ref_norm": 0.022045770859437264,
+ "got_norm": 0.02201467289504321,
+ "finite": true,
+ "rel_l2": 0.0025140784603846
+ }
+ ]
+ }
+ },
+ "fwn_xla vs official": {
+ "loss": {
+ "official": 1.2029576301574707,
+ "other": 1.2029576301574707,
+ "rel_diff": 0.0,
+ "max_per_var_rel_diff": 8.18611584022639e-06
+ },
+ "grads": {
+ "global_rel_l2": 0.0010794160625910145,
+ "global_cos": 0.9999994194356018,
+ "max_rel_l2": 0.004213919475637777,
+ "n_params": 376,
+ "n_compared": 335,
+ "n_zero_both": 6,
+ "n_zero_ref_only": 0,
+ "all_finite": true,
+ "worst": [
+ {
+ "param": "multimodality_forward/mesh_transformer/transformer/block_01/block_01_norm_conditioning/linear/w",
+ "size": 32768,
+ "ref_norm": 0.05493953766112778,
+ "got_norm": 0.05476829099232772,
+ "finite": true,
+ "rel_l2": 0.004213919475637777
+ },
+ {
+ "param": "multimodality_forward/mesh_transformer/transformer/block_04/mha_proj_k/w",
+ "size": 262144,
+ "ref_norm": 0.15894120122711766,
+ "got_norm": 0.15882039209608542,
+ "finite": true,
+ "rel_l2": 0.004180901456984088
+ },
+ {
+ "param": "multimodality_forward/mesh_encoder/split_input_matmul/w_spatial_feature=sin_lat,sin_lon,cos_lon",
+ "size": 1536,
+ "ref_norm": 0.3616984147465042,
+ "got_norm": 0.36162292202956914,
+ "finite": true,
+ "rel_l2": 0.004064530582597025
+ },
+ {
+ "param": "multimodality_forward/mesh_transformer/transformer/block_07/block_07_norm_conditioning/linear/w",
+ "size": 32768,
+ "ref_norm": 0.01887583447986264,
+ "got_norm": 0.018844949178039097,
+ "finite": true,
+ "rel_l2": 0.003917899917061477
+ },
+ {
+ "param": "multimodality_forward/mesh_transformer/transformer/block_04/block_04_norm_conditioning/linear/w",
+ "size": 32768,
+ "ref_norm": 0.029279132855036107,
+ "got_norm": 0.02930957443257014,
+ "finite": true,
+ "rel_l2": 0.0036730124992645137
+ }
+ ]
+ }
+ },
+ "official_mha vs official (floor)": {
+ "loss": {
+ "official": 1.2029576301574707,
+ "other": 1.2029576301574707,
+ "rel_diff": 0.0,
+ "max_per_var_rel_diff": 1.020726765846818e-05
+ },
+ "grads": {
+ "global_rel_l2": 0.0009136954060176284,
+ "global_cos": 0.9999995826302444,
+ "max_rel_l2": 0.002287502960632796,
+ "n_params": 376,
+ "n_compared": 335,
+ "n_zero_both": 6,
+ "n_zero_ref_only": 0,
+ "all_finite": true,
+ "worst": [
+ {
+ "param": "multimodality_forward/mesh_transformer/transformer/block_01/block_01_norm_conditioning/linear/w",
+ "size": 32768,
+ "ref_norm": 0.05493953766112778,
+ "got_norm": 0.0550357834894235,
+ "finite": true,
+ "rel_l2": 0.002287502960632796
+ },
+ {
+ "param": "multimodality_forward/mesh_transformer/transformer/block_03/block_03_norm_conditioning/linear/w",
+ "size": 32768,
+ "ref_norm": 0.028151286105652047,
+ "got_norm": 0.02818322210931235,
+ "finite": true,
+ "rel_l2": 0.0022356896195140238
+ },
+ {
+ "param": "multimodality_forward/mesh_transformer/transformer/block_04/block_04_norm_conditioning/linear/w",
+ "size": 32768,
+ "ref_norm": 0.029279132855036107,
+ "got_norm": 0.029262401991651917,
+ "finite": true,
+ "rel_l2": 0.0021159026462803125
+ },
+ {
+ "param": "multimodality_forward/mesh_transformer/transformer/block_03/mha_proj_k/w",
+ "size": 262144,
+ "ref_norm": 0.0648913028134106,
+ "got_norm": 0.06492808556114596,
+ "finite": true,
+ "rel_l2": 0.0020859323581702205
+ },
+ {
+ "param": "multimodality_forward/grid_encoder/split_input_matmul/w_input_vertical_velocity_level=50,100,150,200,250,300,400,500,600,700,850,925,1000_time=0",
+ "size": 6656,
+ "ref_norm": 4.034688790686174,
+ "got_norm": 4.0357864243215005,
+ "finite": true,
+ "rel_l2": 0.002035561559082506
+ }
+ ]
+ }
+ }
+ }
+ },
+ "default": {
+ "official": {
+ "temp_gib": 102.06438549607992,
+ "compile_s": 39.20219326019287,
+ "step_s": 0.4702334403991699
+ },
+ "official_mha": {
+ "temp_gib": 123.48937734216452,
+ "compile_s": 31.0665500164032,
+ "step_s": 0.5561277866363525
+ },
+ "fwn": {
+ "temp_gib": 3.967879243195057,
+ "compile_s": 39.15290355682373,
+ "step_s": 0.3507387638092041
+ },
+ "fwn_xla": {
+ "temp_gib": 4.0852667689323425,
+ "compile_s": 38.490339279174805,
+ "step_s": 0.6661107540130615
+ },
+ "fwn_noremat": {
+ "temp_gib": 39.10419701784849,
+ "compile_s": 29.292128086090088
+ },
+ "reference": "official",
+ "checks": {
+ "fwn vs official": {
+ "loss": {
+ "official": 1.2031211853027344,
+ "other": 1.2031347751617432,
+ "rel_diff": 1.1295503042255486e-05,
+ "max_per_var_rel_diff": 0.004589995272017043
+ },
+ "grads": {
+ "global_rel_l2": 0.052050369166098155,
+ "global_cos": 0.9986517385107002,
+ "max_rel_l2": 0.3054103458795701,
+ "n_params": 376,
+ "n_compared": 336,
+ "n_zero_both": 6,
+ "n_zero_ref_only": 0,
+ "all_finite": true,
+ "worst": [
+ {
+ "param": "multimodality_forward/mesh_encoder/shared_dense/mlp/linear_0/b",
+ "size": 512,
+ "ref_norm": 0.17328117597699777,
+ "got_norm": 0.20136888182868512,
+ "finite": true,
+ "rel_l2": 0.3054103458795701
+ },
+ {
+ "param": "multimodality_forward/mesh_encoder/split_input_matmul/w_forcing_year_progress_sin_time=21600",
+ "size": 512,
+ "ref_norm": 0.2438398696311633,
+ "got_norm": 0.28335680961155346,
+ "finite": true,
+ "rel_l2": 0.3053115199485669
+ },
+ {
+ "param": "multimodality_forward/mesh_encoder/split_input_matmul/w_input_year_progress_sin_time=-21600",
+ "size": 512,
+ "ref_norm": 0.2440091187993347,
+ "got_norm": 0.2835534226858345,
+ "finite": true,
+ "rel_l2": 0.30531145141750204
+ },
+ {
+ "param": "multimodality_forward/mesh_encoder/split_input_matmul/w_input_year_progress_sin_time=0",
+ "size": 512,
+ "ref_norm": 0.2440091187993347,
+ "got_norm": 0.2835534226858345,
+ "finite": true,
+ "rel_l2": 0.30531145141750204
+ },
+ {
+ "param": "multimodality_forward/mesh_encoder/split_input_matmul/w_input_year_progress_cos_time=-21600",
+ "size": 512,
+ "ref_norm": 0.02527659164456755,
+ "got_norm": 0.02937293561273626,
+ "finite": true,
+ "rel_l2": 0.3053114345326403
+ }
+ ]
+ }
+ },
+ "fwn_xla vs official": {
+ "loss": {
+ "official": 1.2031211853027344,
+ "other": 1.2031413316726685,
+ "rel_diff": 1.674508784334366e-05,
+ "max_per_var_rel_diff": 0.005970923640568816
+ },
+ "grads": {
+ "global_rel_l2": 0.04908203235403402,
+ "global_cos": 0.9987950153960955,
+ "max_rel_l2": 0.18898177194156154,
+ "n_params": 376,
+ "n_compared": 336,
+ "n_zero_both": 6,
+ "n_zero_ref_only": 0,
+ "all_finite": true,
+ "worst": [
+ {
+ "param": "multimodality_forward/mesh_transformer/transformer/block_02/block_02_norm_conditioning_1/linear/b",
+ "size": 1024,
+ "ref_norm": 0.01852919951159923,
+ "got_norm": 0.015847351534108614,
+ "finite": true,
+ "rel_l2": 0.18898177194156154
+ },
+ {
+ "param": "multimodality_forward/mesh_transformer/transformer/block_06/block_06_norm_conditioning_1/linear/b",
+ "size": 1024,
+ "ref_norm": 0.012167384672283505,
+ "got_norm": 0.011667806929174525,
+ "finite": true,
+ "rel_l2": 0.14962574177037477
+ },
+ {
+ "param": "multimodality_forward/mesh_transformer/transformer/block_06/block_06_norm_conditioning_1/linear/w",
+ "size": 32768,
+ "ref_norm": 0.014739708444369387,
+ "got_norm": 0.014546923992896364,
+ "finite": true,
+ "rel_l2": 0.14544199266689398
+ },
+ {
+ "param": "multimodality_forward/mesh_transformer/transformer/block_03/mha_proj_k/w",
+ "size": 262144,
+ "ref_norm": 0.06519183358974383,
+ "got_norm": 0.06592257060411956,
+ "finite": true,
+ "rel_l2": 0.14507685565384099
+ },
+ {
+ "param": "multimodality_forward/mesh_transformer/transformer/block_07/block_07_norm_conditioning/linear/b",
+ "size": 1024,
+ "ref_norm": 0.012845859690973977,
+ "got_norm": 0.012721933333675807,
+ "finite": true,
+ "rel_l2": 0.1407804109830163
+ }
+ ]
+ }
+ },
+ "official_mha vs official (floor)": {
+ "loss": {
+ "official": 1.2031211853027344,
+ "other": 1.2031679153442383,
+ "rel_diff": 3.884067712775571e-05,
+ "max_per_var_rel_diff": 0.008008474768098293
+ },
+ "grads": {
+ "global_rel_l2": 0.05450181891557605,
+ "global_cos": 0.9985183258357095,
+ "max_rel_l2": 0.24367268550547957,
+ "n_params": 376,
+ "n_compared": 336,
+ "n_zero_both": 6,
+ "n_zero_ref_only": 0,
+ "all_finite": true,
+ "worst": [
+ {
+ "param": "multimodality_forward/mesh_encoder/shared_dense/mlp/linear_0/b",
+ "size": 512,
+ "ref_norm": 0.17328117597699777,
+ "got_norm": 0.14728499400526054,
+ "finite": true,
+ "rel_l2": 0.24367268550547957
+ },
+ {
+ "param": "multimodality_forward/mesh_encoder/split_input_matmul/w_input_year_progress_cos_time=0",
+ "size": 512,
+ "ref_norm": 0.026334191823866752,
+ "got_norm": 0.022383602098585937,
+ "finite": true,
+ "rel_l2": 0.24362613355389529
+ },
+ {
+ "param": "multimodality_forward/mesh_encoder/split_input_matmul/w_input_year_progress_sin_time=-21600",
+ "size": 512,
+ "ref_norm": 0.2440091187993347,
+ "got_norm": 0.20740348389788038,
+ "finite": true,
+ "rel_l2": 0.24362609072188443
+ },
+ {
+ "param": "multimodality_forward/mesh_encoder/split_input_matmul/w_input_year_progress_sin_time=0",
+ "size": 512,
+ "ref_norm": 0.2440091187993347,
+ "got_norm": 0.20740348389788038,
+ "finite": true,
+ "rel_l2": 0.24362609072188443
+ },
+ {
+ "param": "multimodality_forward/mesh_encoder/split_input_matmul/w_input_year_progress_cos_time=-21600",
+ "size": 512,
+ "ref_norm": 0.02527659164456755,
+ "got_norm": 0.021484659711007298,
+ "finite": true,
+ "rel_l2": 0.2436260807252004
+ }
+ ]
+ }
+ }
+ }
+ }
+ }
+}
\ No newline at end of file
diff --git a/docs/results/finetune/h200/grads01/full_grads_mini_crps_h200.log b/docs/results/finetune/h200/grads01/full_grads_mini_crps_h200.log
new file mode 100644
index 0000000..c961de7
--- /dev/null
+++ b/docs/results/finetune/h200/grads01/full_grads_mini_crps_h200.log
@@ -0,0 +1,216 @@
+WeatherNextCyclones_Mini: differentiating 56.68M params (all); masked targets: []
+[highest]
+ official highest temp 102.06 GiB step 0.87s loss 1.202958
+ official_mha highest temp 123.49 GiB step 1.07s loss 1.202958
+ fwn highest temp 3.97 GiB step 1.12s loss 1.202958
+ fwn_xla highest temp 4.09 GiB step 1.25s loss 1.202958
+ fwn_strict highest temp 3.97 GiB step 1.59s loss 1.202958
+ fwn_noremat highest temp 39.10 GiB (compile only)
+ fwn vs official loss rel diff 9.91e-08 (per var max 1.54e-05) grads: global rel L2 4.75e-03 cos 0.9999887413 max per-param rel L2 2.81e-02 (335/376 params; zero in both 6, zero only in official 0) finite True
+ 2.81e-02 multimodality_forward/mesh_transformer/transformer/block_06/mha_proj_k/w
+ 2.06e-02 multimodality_forward/mesh_transformer/transformer/block_06/mha_proj_q/w
+ 1.93e-02 multimodality_forward/mesh_transformer/transformer/block_04/mha_proj_k/w
+ fwn_strict vs official loss rel diff 9.91e-08 (per var max 1.59e-05) grads: global rel L2 1.06e-03 cos 0.9999994351 max per-param rel L2 2.94e-03 (335/376 params; zero in both 6, zero only in official 0) finite True
+ 2.94e-03 multimodality_forward/mesh_transformer/transformer/block_06/block_06_norm_conditioning_1/linear/b
+ 2.78e-03 multimodality_forward/mesh_transformer/transformer/block_07/block_07_norm_conditioning/linear/b
+ 2.69e-03 multimodality_forward/mesh_transformer/transformer/block_06/block_06_norm_conditioning_1/linear/w
+ fwn_xla vs official loss rel diff 0.00e+00 (per var max 8.19e-06) grads: global rel L2 1.08e-03 cos 0.9999994194 max per-param rel L2 4.21e-03 (335/376 params; zero in both 6, zero only in official 0) finite True
+ 4.21e-03 multimodality_forward/mesh_transformer/transformer/block_01/block_01_norm_conditioning/linear/w
+ 4.18e-03 multimodality_forward/mesh_transformer/transformer/block_04/mha_proj_k/w
+ 4.06e-03 multimodality_forward/mesh_encoder/split_input_matmul/w_spatial_feature=sin_lat,sin_lon,cos_lon
+ official_mha vs official (floor) loss rel diff 0.00e+00 (per var max 1.02e-05) grads: global rel L2 9.14e-04 cos 0.9999995826 max per-param rel L2 2.29e-03 (335/376 params; zero in both 6, zero only in official 0) finite True
+ 2.29e-03 multimodality_forward/mesh_transformer/transformer/block_01/block_01_norm_conditioning/linear/w
+ 2.24e-03 multimodality_forward/mesh_transformer/transformer/block_03/block_03_norm_conditioning/linear/w
+ 2.12e-03 multimodality_forward/mesh_transformer/transformer/block_04/block_04_norm_conditioning/linear/w
+[default]
+ official default temp 102.06 GiB step 0.47s loss 1.203121
+ official_mha default temp 123.49 GiB step 0.56s loss 1.203168
+ fwn default temp 3.97 GiB step 0.35s loss 1.203135
+ fwn_xla default temp 4.09 GiB step 0.67s loss 1.203141
+ fwn_noremat default temp 39.10 GiB (compile only)
+ fwn vs official loss rel diff 1.13e-05 (per var max 4.59e-03) grads: global rel L2 5.21e-02 cos 0.9986517385 max per-param rel L2 3.05e-01 (336/376 params; zero in both 6, zero only in official 0) finite True
+ 3.05e-01 multimodality_forward/mesh_encoder/shared_dense/mlp/linear_0/b
+ 3.05e-01 multimodality_forward/mesh_encoder/split_input_matmul/w_forcing_year_progress_sin_time=21600
+ 3.05e-01 multimodality_forward/mesh_encoder/split_input_matmul/w_input_year_progress_sin_time=-21600
+ fwn_xla vs official loss rel diff 1.67e-05 (per var max 5.97e-03) grads: global rel L2 4.91e-02 cos 0.9987950154 max per-param rel L2 1.89e-01 (336/376 params; zero in both 6, zero only in official 0) finite True
+ 1.89e-01 multimodality_forward/mesh_transformer/transformer/block_02/block_02_norm_conditioning_1/linear/b
+ 1.50e-01 multimodality_forward/mesh_transformer/transformer/block_06/block_06_norm_conditioning_1/linear/b
+ 1.45e-01 multimodality_forward/mesh_transformer/transformer/block_06/block_06_norm_conditioning_1/linear/w
+ official_mha vs official (floor) loss rel diff 3.88e-05 (per var max 8.01e-03) grads: global rel L2 5.45e-02 cos 0.9985183258 max per-param rel L2 2.44e-01 (336/376 params; zero in both 6, zero only in official 0) finite True
+ 2.44e-01 multimodality_forward/mesh_encoder/shared_dense/mlp/linear_0/b
+ 2.44e-01 multimodality_forward/mesh_encoder/split_input_matmul/w_input_year_progress_cos_time=0
+ 2.44e-01 multimodality_forward/mesh_encoder/split_input_matmul/w_input_year_progress_sin_time=-21600
+wrote /data/results/grads01/full_grads_mini_crps_h200.json
+WARNING:root:No normalization location found for cyclone_exists_gaussian_unit_mode
+WARNING:root:No normalization scale found for cyclone_exists_gaussian_unit_mode
+WARNING:absl:No coordinates found for dim noise_channels of array Size: 256B
+JitTracer(float32[2,1,32])
+Coordinates:
+ * batch (batch) object 16B MultiIndex
+ * _batch (batch) int64 16B 0 0
+ * sample (batch) int64 16B 0 1
+Dimensions without coordinates: channels, noise_channels., using index for param name. This may make model surgery more error prone if changing coordinates at fine-tuning time.
+/usr/local/lib/python3.12/site-packages/scipy/sparse/_index.py:216: SparseEfficiencyWarning: Changing the sparsity structure of a csr_matrix is expensive. lil and dok are more efficient.
+ self._set_arrayXarray(i, j, x)
+WARNING:absl:Only the first input of the static gather operation is being used, others are being ignored at runtime.
+WARNING:absl:Only the first input of the static gather operation is being used, others are being ignored at runtime.
+WARNING:absl:Only the first input of the sorted segment sum operation is being used, others are being ignored at runtime.
+WARNING:absl:Only the first input of the static gather operation is being used, others are being ignored at runtime.
+WARNING:absl:Only the first input of the static gather operation is being used, others are being ignored at runtime.
+WARNING:absl:Only the first input of the static gather operation is being used, others are being ignored at runtime.
+WARNING:absl:Only the first input of the sorted segment sum operation is being used, others are being ignored at runtime.
+/usr/local/lib/python3.12/site-packages/weathernext/weathernext2/architecture.py:286: FutureWarning: Deleting a single level of a MultiIndex is deprecated. Previously, this deleted all levels of a MultiIndex. Please also drop the following variables: {'batch'} to avoid an error in the future.
+ predictions[var] = xarray_jax.apply_ufunc(activation_fn, predictions[var])
+WARNING:root:No normalization location found for cyclone_exists_gaussian_unit_mode
+WARNING:root:No normalization scale found for cyclone_exists_gaussian_unit_mode
+WARNING:absl:No coordinates found for dim noise_channels of array Size: 256B
+JitTracer(float32[2,1,32])
+Coordinates:
+ * batch (batch) object 16B MultiIndex
+ * _batch (batch) int64 16B 0 0
+ * sample (batch) int64 16B 0 1
+Dimensions without coordinates: channels, noise_channels., using index for param name. This may make model surgery more error prone if changing coordinates at fine-tuning time.
+/usr/local/lib/python3.12/site-packages/scipy/sparse/_index.py:216: SparseEfficiencyWarning: Changing the sparsity structure of a csr_matrix is expensive. lil and dok are more efficient.
+ self._set_arrayXarray(i, j, x)
+WARNING:absl:Only the first input of the static gather operation is being used, others are being ignored at runtime.
+WARNING:absl:Only the first input of the static gather operation is being used, others are being ignored at runtime.
+WARNING:absl:Only the first input of the sorted segment sum operation is being used, others are being ignored at runtime.
+WARNING:absl:Only the first input of the static gather operation is being used, others are being ignored at runtime.
+WARNING:absl:Only the first input of the static gather operation is being used, others are being ignored at runtime.
+WARNING:absl:Only the first input of the static gather operation is being used, others are being ignored at runtime.
+WARNING:absl:Only the first input of the sorted segment sum operation is being used, others are being ignored at runtime.
+/usr/local/lib/python3.12/site-packages/weathernext/weathernext2/architecture.py:286: FutureWarning: Deleting a single level of a MultiIndex is deprecated. Previously, this deleted all levels of a MultiIndex. Please also drop the following variables: {'batch'} to avoid an error in the future.
+ predictions[var] = xarray_jax.apply_ufunc(activation_fn, predictions[var])
+WARNING:root:No normalization location found for cyclone_exists_gaussian_unit_mode
+WARNING:root:No normalization scale found for cyclone_exists_gaussian_unit_mode
+WARNING:absl:No coordinates found for dim noise_channels of array Size: 256B
+JitTracer(float32[2,1,32])
+Coordinates:
+ * batch (batch) object 16B MultiIndex
+ * _batch (batch) int64 16B 0 0
+ * sample (batch) int64 16B 0 1
+Dimensions without coordinates: channels, noise_channels., using index for param name. This may make model surgery more error prone if changing coordinates at fine-tuning time.
+/usr/local/lib/python3.12/site-packages/scipy/sparse/_index.py:216: SparseEfficiencyWarning: Changing the sparsity structure of a csr_matrix is expensive. lil and dok are more efficient.
+ self._set_arrayXarray(i, j, x)
+/usr/local/lib/python3.12/site-packages/weathernext/weathernext2/architecture.py:286: FutureWarning: Deleting a single level of a MultiIndex is deprecated. Previously, this deleted all levels of a MultiIndex. Please also drop the following variables: {'batch'} to avoid an error in the future.
+ predictions[var] = xarray_jax.apply_ufunc(activation_fn, predictions[var])
+WARNING:root:No normalization location found for cyclone_exists_gaussian_unit_mode
+WARNING:root:No normalization scale found for cyclone_exists_gaussian_unit_mode
+WARNING:absl:No coordinates found for dim noise_channels of array Size: 256B
+JitTracer(float32[2,1,32])
+Coordinates:
+ * batch (batch) object 16B MultiIndex
+ * _batch (batch) int64 16B 0 0
+ * sample (batch) int64 16B 0 1
+Dimensions without coordinates: channels, noise_channels., using index for param name. This may make model surgery more error prone if changing coordinates at fine-tuning time.
+/usr/local/lib/python3.12/site-packages/scipy/sparse/_index.py:216: SparseEfficiencyWarning: Changing the sparsity structure of a csr_matrix is expensive. lil and dok are more efficient.
+ self._set_arrayXarray(i, j, x)
+/usr/local/lib/python3.12/site-packages/weathernext/weathernext2/architecture.py:286: FutureWarning: Deleting a single level of a MultiIndex is deprecated. Previously, this deleted all levels of a MultiIndex. Please also drop the following variables: {'batch'} to avoid an error in the future.
+ predictions[var] = xarray_jax.apply_ufunc(activation_fn, predictions[var])
+WARNING:root:No normalization location found for cyclone_exists_gaussian_unit_mode
+WARNING:root:No normalization scale found for cyclone_exists_gaussian_unit_mode
+WARNING:absl:No coordinates found for dim noise_channels of array Size: 256B
+JitTracer(float32[2,1,32])
+Coordinates:
+ * batch (batch) object 16B MultiIndex
+ * _batch (batch) int64 16B 0 0
+ * sample (batch) int64 16B 0 1
+Dimensions without coordinates: channels, noise_channels., using index for param name. This may make model surgery more error prone if changing coordinates at fine-tuning time.
+/usr/local/lib/python3.12/site-packages/scipy/sparse/_index.py:216: SparseEfficiencyWarning: Changing the sparsity structure of a csr_matrix is expensive. lil and dok are more efficient.
+ self._set_arrayXarray(i, j, x)
+/usr/local/lib/python3.12/site-packages/weathernext/weathernext2/architecture.py:286: FutureWarning: Deleting a single level of a MultiIndex is deprecated. Previously, this deleted all levels of a MultiIndex. Please also drop the following variables: {'batch'} to avoid an error in the future.
+ predictions[var] = xarray_jax.apply_ufunc(activation_fn, predictions[var])
+WARNING:root:No normalization location found for cyclone_exists_gaussian_unit_mode
+WARNING:root:No normalization scale found for cyclone_exists_gaussian_unit_mode
+WARNING:absl:No coordinates found for dim noise_channels of array Size: 256B
+JitTracer(float32[2,1,32])
+Coordinates:
+ * batch (batch) object 16B MultiIndex
+ * _batch (batch) int64 16B 0 0
+ * sample (batch) int64 16B 0 1
+Dimensions without coordinates: channels, noise_channels., using index for param name. This may make model surgery more error prone if changing coordinates at fine-tuning time.
+/usr/local/lib/python3.12/site-packages/scipy/sparse/_index.py:216: SparseEfficiencyWarning: Changing the sparsity structure of a csr_matrix is expensive. lil and dok are more efficient.
+ self._set_arrayXarray(i, j, x)
+/usr/local/lib/python3.12/site-packages/weathernext/weathernext2/architecture.py:286: FutureWarning: Deleting a single level of a MultiIndex is deprecated. Previously, this deleted all levels of a MultiIndex. Please also drop the following variables: {'batch'} to avoid an error in the future.
+ predictions[var] = xarray_jax.apply_ufunc(activation_fn, predictions[var])
+WARNING:root:No normalization location found for cyclone_exists_gaussian_unit_mode
+WARNING:root:No normalization scale found for cyclone_exists_gaussian_unit_mode
+WARNING:absl:No coordinates found for dim noise_channels of array Size: 256B
+JitTracer(float32[2,1,32])
+Coordinates:
+ * batch (batch) object 16B MultiIndex
+ * _batch (batch) int64 16B 0 0
+ * sample (batch) int64 16B 0 1
+Dimensions without coordinates: channels, noise_channels., using index for param name. This may make model surgery more error prone if changing coordinates at fine-tuning time.
+/usr/local/lib/python3.12/site-packages/scipy/sparse/_index.py:216: SparseEfficiencyWarning: Changing the sparsity structure of a csr_matrix is expensive. lil and dok are more efficient.
+ self._set_arrayXarray(i, j, x)
+WARNING:absl:Only the first input of the static gather operation is being used, others are being ignored at runtime.
+WARNING:absl:Only the first input of the static gather operation is being used, others are being ignored at runtime.
+WARNING:absl:Only the first input of the sorted segment sum operation is being used, others are being ignored at runtime.
+WARNING:absl:Only the first input of the static gather operation is being used, others are being ignored at runtime.
+WARNING:absl:Only the first input of the static gather operation is being used, others are being ignored at runtime.
+WARNING:absl:Only the first input of the static gather operation is being used, others are being ignored at runtime.
+WARNING:absl:Only the first input of the sorted segment sum operation is being used, others are being ignored at runtime.
+/usr/local/lib/python3.12/site-packages/weathernext/weathernext2/architecture.py:286: FutureWarning: Deleting a single level of a MultiIndex is deprecated. Previously, this deleted all levels of a MultiIndex. Please also drop the following variables: {'batch'} to avoid an error in the future.
+ predictions[var] = xarray_jax.apply_ufunc(activation_fn, predictions[var])
+WARNING:root:No normalization location found for cyclone_exists_gaussian_unit_mode
+WARNING:root:No normalization scale found for cyclone_exists_gaussian_unit_mode
+WARNING:absl:No coordinates found for dim noise_channels of array Size: 256B
+JitTracer(float32[2,1,32])
+Coordinates:
+ * batch (batch) object 16B MultiIndex
+ * _batch (batch) int64 16B 0 0
+ * sample (batch) int64 16B 0 1
+Dimensions without coordinates: channels, noise_channels., using index for param name. This may make model surgery more error prone if changing coordinates at fine-tuning time.
+/usr/local/lib/python3.12/site-packages/scipy/sparse/_index.py:216: SparseEfficiencyWarning: Changing the sparsity structure of a csr_matrix is expensive. lil and dok are more efficient.
+ self._set_arrayXarray(i, j, x)
+WARNING:absl:Only the first input of the static gather operation is being used, others are being ignored at runtime.
+WARNING:absl:Only the first input of the static gather operation is being used, others are being ignored at runtime.
+WARNING:absl:Only the first input of the sorted segment sum operation is being used, others are being ignored at runtime.
+WARNING:absl:Only the first input of the static gather operation is being used, others are being ignored at runtime.
+WARNING:absl:Only the first input of the static gather operation is being used, others are being ignored at runtime.
+WARNING:absl:Only the first input of the static gather operation is being used, others are being ignored at runtime.
+WARNING:absl:Only the first input of the sorted segment sum operation is being used, others are being ignored at runtime.
+/usr/local/lib/python3.12/site-packages/weathernext/weathernext2/architecture.py:286: FutureWarning: Deleting a single level of a MultiIndex is deprecated. Previously, this deleted all levels of a MultiIndex. Please also drop the following variables: {'batch'} to avoid an error in the future.
+ predictions[var] = xarray_jax.apply_ufunc(activation_fn, predictions[var])
+WARNING:root:No normalization location found for cyclone_exists_gaussian_unit_mode
+WARNING:root:No normalization scale found for cyclone_exists_gaussian_unit_mode
+WARNING:absl:No coordinates found for dim noise_channels of array Size: 256B
+JitTracer(float32[2,1,32])
+Coordinates:
+ * batch (batch) object 16B MultiIndex
+ * _batch (batch) int64 16B 0 0
+ * sample (batch) int64 16B 0 1
+Dimensions without coordinates: channels, noise_channels., using index for param name. This may make model surgery more error prone if changing coordinates at fine-tuning time.
+/usr/local/lib/python3.12/site-packages/scipy/sparse/_index.py:216: SparseEfficiencyWarning: Changing the sparsity structure of a csr_matrix is expensive. lil and dok are more efficient.
+ self._set_arrayXarray(i, j, x)
+/usr/local/lib/python3.12/site-packages/weathernext/weathernext2/architecture.py:286: FutureWarning: Deleting a single level of a MultiIndex is deprecated. Previously, this deleted all levels of a MultiIndex. Please also drop the following variables: {'batch'} to avoid an error in the future.
+ predictions[var] = xarray_jax.apply_ufunc(activation_fn, predictions[var])
+WARNING:root:No normalization location found for cyclone_exists_gaussian_unit_mode
+WARNING:root:No normalization scale found for cyclone_exists_gaussian_unit_mode
+WARNING:absl:No coordinates found for dim noise_channels of array Size: 256B
+JitTracer(float32[2,1,32])
+Coordinates:
+ * batch (batch) object 16B MultiIndex
+ * _batch (batch) int64 16B 0 0
+ * sample (batch) int64 16B 0 1
+Dimensions without coordinates: channels, noise_channels., using index for param name. This may make model surgery more error prone if changing coordinates at fine-tuning time.
+/usr/local/lib/python3.12/site-packages/scipy/sparse/_index.py:216: SparseEfficiencyWarning: Changing the sparsity structure of a csr_matrix is expensive. lil and dok are more efficient.
+ self._set_arrayXarray(i, j, x)
+/usr/local/lib/python3.12/site-packages/weathernext/weathernext2/architecture.py:286: FutureWarning: Deleting a single level of a MultiIndex is deprecated. Previously, this deleted all levels of a MultiIndex. Please also drop the following variables: {'batch'} to avoid an error in the future.
+ predictions[var] = xarray_jax.apply_ufunc(activation_fn, predictions[var])
+WARNING:root:No normalization location found for cyclone_exists_gaussian_unit_mode
+WARNING:root:No normalization scale found for cyclone_exists_gaussian_unit_mode
+WARNING:absl:No coordinates found for dim noise_channels of array Size: 256B
+JitTracer(float32[2,1,32])
+Coordinates:
+ * batch (batch) object 16B MultiIndex
+ * _batch (batch) int64 16B 0 0
+ * sample (batch) int64 16B 0 1
+Dimensions without coordinates: channels, noise_channels., using index for param name. This may make model surgery more error prone if changing coordinates at fine-tuning time.
+/usr/local/lib/python3.12/site-packages/scipy/sparse/_index.py:216: SparseEfficiencyWarning: Changing the sparsity structure of a csr_matrix is expensive. lil and dok are more efficient.
+ self._set_arrayXarray(i, j, x)
+/usr/local/lib/python3.12/site-packages/weathernext/weathernext2/architecture.py:286: FutureWarning: Deleting a single level of a MultiIndex is deprecated. Previously, this deleted all levels of a MultiIndex. Please also drop the following variables: {'batch'} to avoid an error in the future.
+ predictions[var] = xarray_jax.apply_ufunc(activation_fn, predictions[var])
+
+exit 0 wall 505s
diff --git a/docs/results/finetune/h200/grads01/full_grads_mini_linear_h200.json b/docs/results/finetune/h200/grads01/full_grads_mini_linear_h200.json
new file mode 100644
index 0000000..7ad2d08
--- /dev/null
+++ b/docs/results/finetune/h200/grads01/full_grads_mini_linear_h200.json
@@ -0,0 +1,499 @@
+{
+ "model": "WeatherNextCyclones_Mini",
+ "device": "NVIDIA H200",
+ "jax": "0.11.2",
+ "faster_weathernext": "0.1.1",
+ "blocks": {
+ "grid_chunk": 8192,
+ "edge_chunk": 16384
+ },
+ "masked": [],
+ "modules": "all",
+ "loss": "linear",
+ "head_params": 56684133,
+ "precisions": {
+ "highest": {
+ "official": {
+ "temp_gib": 51.412283167243004,
+ "compile_s": 66.12185215950012,
+ "step_s": 0.5073494911193848
+ },
+ "official_mha": {
+ "temp_gib": 61.79655335843563,
+ "compile_s": 34.418456077575684,
+ "step_s": 0.5334434509277344
+ },
+ "fwn": {
+ "temp_gib": 3.299081988632679,
+ "compile_s": 59.250096797943115,
+ "step_s": 0.6724720001220703
+ },
+ "fwn_xla": {
+ "temp_gib": 3.397009789943695,
+ "compile_s": 38.46831464767456,
+ "step_s": 0.7373757362365723
+ },
+ "fwn_strict": {
+ "temp_gib": 3.299081988632679,
+ "compile_s": 27.501341581344604,
+ "step_s": 0.8596265316009521
+ },
+ "fwn_noremat": {
+ "temp_gib": 19.658946715295315,
+ "compile_s": 29.531001091003418
+ },
+ "reference": "official",
+ "checks": {
+ "fwn vs official": {
+ "loss": {
+ "official": 506.6861267089844,
+ "other": 506.6861877441406,
+ "rel_diff": 1.204594975718678e-07,
+ "max_per_var_rel_diff": 8.950924192513153e-05
+ },
+ "grads": {
+ "global_rel_l2": 0.000244263538666221,
+ "global_cos": 0.99999997710036,
+ "max_rel_l2": 0.0012532542184612884,
+ "n_params": 376,
+ "n_compared": 315,
+ "n_zero_both": 6,
+ "n_zero_ref_only": 0,
+ "all_finite": true,
+ "worst": [
+ {
+ "param": "multimodality_forward/mesh_transformer/transformer/block_06/mha_proj_k/w",
+ "size": 262144,
+ "ref_norm": 2.712203724691543,
+ "got_norm": 2.713650678091154,
+ "finite": true,
+ "rel_l2": 0.0012532542184612884
+ },
+ {
+ "param": "multimodality_forward/mesh_transformer/transformer/block_02/block_02_norm_conditioning_1/linear/w",
+ "size": 32768,
+ "ref_norm": 0.6123149424207489,
+ "got_norm": 0.6127483205407561,
+ "finite": true,
+ "rel_l2": 0.0008091081465632556
+ },
+ {
+ "param": "multimodality_forward/mesh_transformer/transformer/block_02/block_02_norm_conditioning_1/linear/b",
+ "size": 1024,
+ "ref_norm": 0.7095075457226626,
+ "got_norm": 0.710009713073801,
+ "finite": true,
+ "rel_l2": 0.0008091076612630029
+ },
+ {
+ "param": "multimodality_forward/mesh_transformer/transformer/block_07/block_07_norm_conditioning/linear/b",
+ "size": 1024,
+ "ref_norm": 0.5724014032691926,
+ "got_norm": 0.5726062563797046,
+ "finite": true,
+ "rel_l2": 0.0007411735720147313
+ },
+ {
+ "param": "multimodality_forward/mesh_transformer/transformer/block_07/block_07_norm_conditioning/linear/w",
+ "size": 32768,
+ "ref_norm": 0.49399042167722107,
+ "got_norm": 0.49416721330425006,
+ "finite": true,
+ "rel_l2": 0.0007411733653940433
+ }
+ ]
+ }
+ },
+ "fwn_strict vs official": {
+ "loss": {
+ "official": 506.6861267089844,
+ "other": 506.6861572265625,
+ "rel_diff": 6.02297487859339e-08,
+ "max_per_var_rel_diff": 0.00011490809404406726
+ },
+ "grads": {
+ "global_rel_l2": 0.00013237493740884075,
+ "global_cos": 0.9999999915392175,
+ "max_rel_l2": 0.0005106505240174063,
+ "n_params": 376,
+ "n_compared": 315,
+ "n_zero_both": 6,
+ "n_zero_ref_only": 0,
+ "all_finite": true,
+ "worst": [
+ {
+ "param": "multimodality_forward/grid_to_mesh_gnn/edge_encoder/shared_dense/mlp/linear_0/b",
+ "size": 32,
+ "ref_norm": 5.1305898520590665,
+ "got_norm": 5.1326292854528335,
+ "finite": true,
+ "rel_l2": 0.0005106505240174063
+ },
+ {
+ "param": "multimodality_forward/mesh_transformer/transformer/block_01/block_01_norm_conditioning/linear/b",
+ "size": 1024,
+ "ref_norm": 0.9516581900491706,
+ "got_norm": 0.9514998551362546,
+ "finite": true,
+ "rel_l2": 0.000506067362620046
+ },
+ {
+ "param": "multimodality_forward/mesh_transformer/transformer/block_01/block_01_norm_conditioning/linear/w",
+ "size": 32768,
+ "ref_norm": 0.8212943374960904,
+ "got_norm": 0.8211576926596661,
+ "finite": true,
+ "rel_l2": 0.0005060658502279332
+ },
+ {
+ "param": "multimodality_forward/mesh_transformer/transformer/block_04/mha_proj_k/w",
+ "size": 262144,
+ "ref_norm": 7.420814979710817,
+ "got_norm": 7.418814862699009,
+ "finite": true,
+ "rel_l2": 0.00046718644322505293
+ },
+ {
+ "param": "multimodality_forward/mesh_encoder/split_input_matmul/w_spatial_feature=sin_lat,sin_lon,cos_lon",
+ "size": 1536,
+ "ref_norm": 15.961098272402646,
+ "got_norm": 15.955721873132726,
+ "finite": true,
+ "rel_l2": 0.00045288540773069256
+ }
+ ]
+ }
+ },
+ "fwn_xla vs official": {
+ "loss": {
+ "official": 506.6861267089844,
+ "other": 506.6861267089844,
+ "rel_diff": 0.0,
+ "max_per_var_rel_diff": 0.00011274816746429155
+ },
+ "grads": {
+ "global_rel_l2": 0.00011762606096105297,
+ "global_cos": 0.9999999933447786,
+ "max_rel_l2": 0.0007699412451052978,
+ "n_params": 376,
+ "n_compared": 315,
+ "n_zero_both": 6,
+ "n_zero_ref_only": 0,
+ "all_finite": true,
+ "worst": [
+ {
+ "param": "multimodality_forward/mesh_transformer/transformer/block_04/block_04_norm_conditioning/linear/b",
+ "size": 1024,
+ "ref_norm": 0.7266956977515365,
+ "got_norm": 0.7268171072822894,
+ "finite": true,
+ "rel_l2": 0.0007699412451052978
+ },
+ {
+ "param": "multimodality_forward/mesh_transformer/transformer/block_04/block_04_norm_conditioning/linear/w",
+ "size": 32768,
+ "ref_norm": 0.627148557590409,
+ "got_norm": 0.6272533354652043,
+ "finite": true,
+ "rel_l2": 0.0007699393066329524
+ },
+ {
+ "param": "multimodality_forward/mesh_encoder/split_input_matmul/w_spatial_feature=sin_lat,sin_lon,cos_lon",
+ "size": 1536,
+ "ref_norm": 15.961098272402646,
+ "got_norm": 15.954829243577963,
+ "finite": true,
+ "rel_l2": 0.0005351237024151522
+ },
+ {
+ "param": "multimodality_forward/mesh_transformer/transformer/block_01/block_01_norm_conditioning/linear/b",
+ "size": 1024,
+ "ref_norm": 0.9516581900491706,
+ "got_norm": 0.9515533263912913,
+ "finite": true,
+ "rel_l2": 0.0005235584090649549
+ },
+ {
+ "param": "multimodality_forward/mesh_transformer/transformer/block_01/block_01_norm_conditioning/linear/w",
+ "size": 32768,
+ "ref_norm": 0.8212943374960904,
+ "got_norm": 0.8212038346395856,
+ "finite": true,
+ "rel_l2": 0.0005235580009717514
+ }
+ ]
+ }
+ },
+ "official_mha vs official (floor)": {
+ "loss": {
+ "official": 506.6861267089844,
+ "other": 506.68609619140625,
+ "rel_diff": 6.02297487859339e-08,
+ "max_per_var_rel_diff": 3.446105814117564e-05
+ },
+ "grads": {
+ "global_rel_l2": 0.00012651098564520421,
+ "global_cos": 0.9999999921673326,
+ "max_rel_l2": 0.0008345211820504494,
+ "n_params": 376,
+ "n_compared": 315,
+ "n_zero_both": 6,
+ "n_zero_ref_only": 0,
+ "all_finite": true,
+ "worst": [
+ {
+ "param": "multimodality_forward/mesh_transformer/transformer/block_04/block_04_norm_conditioning/linear/b",
+ "size": 1024,
+ "ref_norm": 0.7266956977515365,
+ "got_norm": 0.7268135631996993,
+ "finite": true,
+ "rel_l2": 0.0008345211820504494
+ },
+ {
+ "param": "multimodality_forward/mesh_transformer/transformer/block_04/block_04_norm_conditioning/linear/w",
+ "size": 32768,
+ "ref_norm": 0.627148557590409,
+ "got_norm": 0.6272502768693304,
+ "finite": true,
+ "rel_l2": 0.0008345193395550328
+ },
+ {
+ "param": "multimodality_forward/mesh_encoder/split_input_matmul/w_spatial_feature=sin_lat,sin_lon,cos_lon",
+ "size": 1536,
+ "ref_norm": 15.961098272402646,
+ "got_norm": 15.95528512486742,
+ "finite": true,
+ "rel_l2": 0.0005456942304003171
+ },
+ {
+ "param": "multimodality_forward/mesh_transformer/transformer/block_04/mha_proj_k/w",
+ "size": 262144,
+ "ref_norm": 7.420814979710817,
+ "got_norm": 7.418794954355552,
+ "finite": true,
+ "rel_l2": 0.000537533317943353
+ },
+ {
+ "param": "multimodality_forward/mesh_transformer/transformer/block_04/mha_proj_q/w",
+ "size": 262144,
+ "ref_norm": 7.8224458328154975,
+ "got_norm": 7.820655928372033,
+ "finite": true,
+ "rel_l2": 0.0004883501030633923
+ }
+ ]
+ }
+ }
+ }
+ },
+ "default": {
+ "official": {
+ "temp_gib": 51.38774202764034,
+ "compile_s": 32.51545262336731,
+ "step_s": 0.2796788215637207
+ },
+ "official_mha": {
+ "temp_gib": 61.772008404135704,
+ "compile_s": 22.377246379852295,
+ "step_s": 0.31437110900878906
+ },
+ "fwn": {
+ "temp_gib": 3.2990810349583626,
+ "compile_s": 34.207173109054565,
+ "step_s": 0.24166154861450195
+ },
+ "fwn_xla": {
+ "temp_gib": 3.396993577480316,
+ "compile_s": 31.51302194595337,
+ "step_s": 0.4486827850341797
+ },
+ "fwn_noremat": {
+ "temp_gib": 19.658945761621,
+ "compile_s": 25.49704933166504
+ },
+ "reference": "official",
+ "checks": {
+ "fwn vs official": {
+ "loss": {
+ "official": 506.6874084472656,
+ "other": 506.6866760253906,
+ "rel_diff": 1.4455103142280436e-06,
+ "max_per_var_rel_diff": 0.11206645425551767
+ },
+ "grads": {
+ "global_rel_l2": 0.044117534887844354,
+ "global_cos": 0.9990361813716523,
+ "max_rel_l2": 0.22349852491558603,
+ "n_params": 376,
+ "n_compared": 315,
+ "n_zero_both": 6,
+ "n_zero_ref_only": 0,
+ "all_finite": true,
+ "worst": [
+ {
+ "param": "multimodality_forward/mesh_transformer/transformer/block_01/block_01_norm_conditioning/linear/b",
+ "size": 1024,
+ "ref_norm": 1.0117770591990807,
+ "got_norm": 1.156332303299948,
+ "finite": true,
+ "rel_l2": 0.22349852491558603
+ },
+ {
+ "param": "multimodality_forward/mesh_transformer/transformer/block_01/block_01_norm_conditioning/linear/w",
+ "size": 32768,
+ "ref_norm": 0.8731777646209737,
+ "got_norm": 0.9979309599064247,
+ "finite": true,
+ "rel_l2": 0.22349851854323546
+ },
+ {
+ "param": "multimodality_forward/mesh_transformer/transformer/block_02/block_02_norm_conditioning/linear/b",
+ "size": 1024,
+ "ref_norm": 0.5947899211484378,
+ "got_norm": 0.6545359508285181,
+ "finite": true,
+ "rel_l2": 0.1701816422646175
+ },
+ {
+ "param": "multimodality_forward/mesh_transformer/transformer/block_02/block_02_norm_conditioning/linear/w",
+ "size": 32768,
+ "ref_norm": 0.5133120270896127,
+ "got_norm": 0.5648736861383704,
+ "finite": true,
+ "rel_l2": 0.17018164137891015
+ },
+ {
+ "param": "multimodality_forward/mesh_transformer/transformer/block_03/block_03_norm_conditioning/linear/b",
+ "size": 1024,
+ "ref_norm": 0.7167292689831924,
+ "got_norm": 0.796003374445032,
+ "finite": true,
+ "rel_l2": 0.16284349983095633
+ }
+ ]
+ }
+ },
+ "fwn_xla vs official": {
+ "loss": {
+ "official": 506.6874084472656,
+ "other": 506.6881103515625,
+ "rel_diff": 1.385280717801875e-06,
+ "max_per_var_rel_diff": 0.08084677316850619
+ },
+ "grads": {
+ "global_rel_l2": 0.019679887946770405,
+ "global_cos": 0.9998077940339737,
+ "max_rel_l2": 0.059476224600919315,
+ "n_params": 376,
+ "n_compared": 315,
+ "n_zero_both": 6,
+ "n_zero_ref_only": 0,
+ "all_finite": true,
+ "worst": [
+ {
+ "param": "multimodality_forward/grid_to_mesh_gnn/edge_encoder/shared_dense/mlp/linear_0/b",
+ "size": 32,
+ "ref_norm": 5.441638744526304,
+ "got_norm": 5.68977024596661,
+ "finite": true,
+ "rel_l2": 0.059476224600919315
+ },
+ {
+ "param": "multimodality_forward/grid_to_mesh_gnn/edge_encoder/shared_dense/normalization/linear_norm_conditioning/linear/w",
+ "size": 2048,
+ "ref_norm": 0.23800722621596107,
+ "got_norm": 0.24467928565487893,
+ "finite": true,
+ "rel_l2": 0.05608277617712862
+ },
+ {
+ "param": "multimodality_forward/grid_to_mesh_gnn/edge_encoder/shared_dense/normalization/linear_norm_conditioning/linear/b",
+ "size": 64,
+ "ref_norm": 0.2757860559578531,
+ "got_norm": 0.2835171695923799,
+ "finite": true,
+ "rel_l2": 0.05608275600530208
+ },
+ {
+ "param": "multimodality_forward/grid_to_mesh_gnn/edge_encoder/shared_dense/mlp/linear_1/b",
+ "size": 32,
+ "ref_norm": 24.851766528252575,
+ "got_norm": 25.61888271096121,
+ "finite": true,
+ "rel_l2": 0.05511299698737055
+ },
+ {
+ "param": "multimodality_forward/mesh_transformer/transformer/block_06/mha_proj_k/w",
+ "size": 262144,
+ "ref_norm": 2.8108629946526866,
+ "got_norm": 2.787529759983267,
+ "finite": true,
+ "rel_l2": 0.0542160190231521
+ }
+ ]
+ }
+ },
+ "official_mha vs official (floor)": {
+ "loss": {
+ "official": 506.6874084472656,
+ "other": 506.6858215332031,
+ "rel_diff": 3.1319390141607612e-06,
+ "max_per_var_rel_diff": 0.1407590322703549
+ },
+ "grads": {
+ "global_rel_l2": 0.044800067284372115,
+ "global_cos": 0.9990057146368683,
+ "max_rel_l2": 0.1750765677236349,
+ "n_params": 376,
+ "n_compared": 315,
+ "n_zero_both": 6,
+ "n_zero_ref_only": 0,
+ "all_finite": true,
+ "worst": [
+ {
+ "param": "multimodality_forward/grid_to_mesh_gnn/edge_encoder/shared_dense/mlp/linear_0/b",
+ "size": 32,
+ "ref_norm": 5.441638744526304,
+ "got_norm": 6.210838741861778,
+ "finite": true,
+ "rel_l2": 0.1750765677236349
+ },
+ {
+ "param": "multimodality_forward/mesh_transformer/transformer/block_07/block_07_norm_conditioning/linear/b",
+ "size": 1024,
+ "ref_norm": 0.5897101140845485,
+ "got_norm": 0.5897953663132253,
+ "finite": true,
+ "rel_l2": 0.16793897090657076
+ },
+ {
+ "param": "multimodality_forward/mesh_transformer/transformer/block_07/block_07_norm_conditioning/linear/w",
+ "size": 32768,
+ "ref_norm": 0.508928081428834,
+ "got_norm": 0.5090016550054227,
+ "finite": true,
+ "rel_l2": 0.16793897087111273
+ },
+ {
+ "param": "multimodality_forward/mesh_transformer/transformer/block_04/block_04_norm_conditioning/linear/w",
+ "size": 32768,
+ "ref_norm": 0.672737338310035,
+ "got_norm": 0.7184029175034351,
+ "finite": true,
+ "rel_l2": 0.1673825597292103
+ },
+ {
+ "param": "multimodality_forward/mesh_transformer/transformer/block_04/block_04_norm_conditioning/linear/b",
+ "size": 1024,
+ "ref_norm": 0.7795207748498905,
+ "got_norm": 0.8324348396954653,
+ "finite": true,
+ "rel_l2": 0.16738255742938324
+ }
+ ]
+ }
+ }
+ }
+ }
+ }
+}
\ No newline at end of file
diff --git a/docs/results/finetune/h200/grads01/full_grads_mini_linear_h200.log b/docs/results/finetune/h200/grads01/full_grads_mini_linear_h200.log
new file mode 100644
index 0000000..4861223
--- /dev/null
+++ b/docs/results/finetune/h200/grads01/full_grads_mini_linear_h200.log
@@ -0,0 +1,169 @@
+WeatherNextCyclones_Mini: differentiating 56.68M params (all); masked targets: []
+[highest]
+ official highest temp 51.41 GiB step 0.51s loss 506.686127
+ official_mha highest temp 61.80 GiB step 0.53s loss 506.686096
+ fwn highest temp 3.30 GiB step 0.67s loss 506.686188
+ fwn_xla highest temp 3.40 GiB step 0.74s loss 506.686127
+ fwn_strict highest temp 3.30 GiB step 0.86s loss 506.686157
+ fwn_noremat highest temp 19.66 GiB (compile only)
+ fwn vs official loss rel diff 1.20e-07 (per var max 8.95e-05) grads: global rel L2 2.44e-04 cos 0.9999999771 max per-param rel L2 1.25e-03 (315/376 params; zero in both 6, zero only in official 0) finite True
+ 1.25e-03 multimodality_forward/mesh_transformer/transformer/block_06/mha_proj_k/w
+ 8.09e-04 multimodality_forward/mesh_transformer/transformer/block_02/block_02_norm_conditioning_1/linear/w
+ 8.09e-04 multimodality_forward/mesh_transformer/transformer/block_02/block_02_norm_conditioning_1/linear/b
+ fwn_strict vs official loss rel diff 6.02e-08 (per var max 1.15e-04) grads: global rel L2 1.32e-04 cos 0.9999999915 max per-param rel L2 5.11e-04 (315/376 params; zero in both 6, zero only in official 0) finite True
+ 5.11e-04 multimodality_forward/grid_to_mesh_gnn/edge_encoder/shared_dense/mlp/linear_0/b
+ 5.06e-04 multimodality_forward/mesh_transformer/transformer/block_01/block_01_norm_conditioning/linear/b
+ 5.06e-04 multimodality_forward/mesh_transformer/transformer/block_01/block_01_norm_conditioning/linear/w
+ fwn_xla vs official loss rel diff 0.00e+00 (per var max 1.13e-04) grads: global rel L2 1.18e-04 cos 0.9999999933 max per-param rel L2 7.70e-04 (315/376 params; zero in both 6, zero only in official 0) finite True
+ 7.70e-04 multimodality_forward/mesh_transformer/transformer/block_04/block_04_norm_conditioning/linear/b
+ 7.70e-04 multimodality_forward/mesh_transformer/transformer/block_04/block_04_norm_conditioning/linear/w
+ 5.35e-04 multimodality_forward/mesh_encoder/split_input_matmul/w_spatial_feature=sin_lat,sin_lon,cos_lon
+ official_mha vs official (floor) loss rel diff 6.02e-08 (per var max 3.45e-05) grads: global rel L2 1.27e-04 cos 0.9999999922 max per-param rel L2 8.35e-04 (315/376 params; zero in both 6, zero only in official 0) finite True
+ 8.35e-04 multimodality_forward/mesh_transformer/transformer/block_04/block_04_norm_conditioning/linear/b
+ 8.35e-04 multimodality_forward/mesh_transformer/transformer/block_04/block_04_norm_conditioning/linear/w
+ 5.46e-04 multimodality_forward/mesh_encoder/split_input_matmul/w_spatial_feature=sin_lat,sin_lon,cos_lon
+[default]
+ official default temp 51.39 GiB step 0.28s loss 506.687408
+ official_mha default temp 61.77 GiB step 0.31s loss 506.685822
+ fwn default temp 3.30 GiB step 0.24s loss 506.686676
+ fwn_xla default temp 3.40 GiB step 0.45s loss 506.688110
+ fwn_noremat default temp 19.66 GiB (compile only)
+ fwn vs official loss rel diff 1.45e-06 (per var max 1.12e-01) grads: global rel L2 4.41e-02 cos 0.9990361814 max per-param rel L2 2.23e-01 (315/376 params; zero in both 6, zero only in official 0) finite True
+ 2.23e-01 multimodality_forward/mesh_transformer/transformer/block_01/block_01_norm_conditioning/linear/b
+ 2.23e-01 multimodality_forward/mesh_transformer/transformer/block_01/block_01_norm_conditioning/linear/w
+ 1.70e-01 multimodality_forward/mesh_transformer/transformer/block_02/block_02_norm_conditioning/linear/b
+ fwn_xla vs official loss rel diff 1.39e-06 (per var max 8.08e-02) grads: global rel L2 1.97e-02 cos 0.9998077940 max per-param rel L2 5.95e-02 (315/376 params; zero in both 6, zero only in official 0) finite True
+ 5.95e-02 multimodality_forward/grid_to_mesh_gnn/edge_encoder/shared_dense/mlp/linear_0/b
+ 5.61e-02 multimodality_forward/grid_to_mesh_gnn/edge_encoder/shared_dense/normalization/linear_norm_conditioning/linear/w
+ 5.61e-02 multimodality_forward/grid_to_mesh_gnn/edge_encoder/shared_dense/normalization/linear_norm_conditioning/linear/b
+ official_mha vs official (floor) loss rel diff 3.13e-06 (per var max 1.41e-01) grads: global rel L2 4.48e-02 cos 0.9990057146 max per-param rel L2 1.75e-01 (315/376 params; zero in both 6, zero only in official 0) finite True
+ 1.75e-01 multimodality_forward/grid_to_mesh_gnn/edge_encoder/shared_dense/mlp/linear_0/b
+ 1.68e-01 multimodality_forward/mesh_transformer/transformer/block_07/block_07_norm_conditioning/linear/b
+ 1.68e-01 multimodality_forward/mesh_transformer/transformer/block_07/block_07_norm_conditioning/linear/w
+wrote /data/results/grads01/full_grads_mini_linear_h200.json
+WARNING:absl:Skipping gradient checkpointing for sequence length of 1
+WARNING:absl:No coordinates found for dim noise_channels of array Size: 128B
+JitTracer(float32[1,1,32])
+Dimensions without coordinates: batch, channels, noise_channels., using index for param name. This may make model surgery more error prone if changing coordinates at fine-tuning time.
+/usr/local/lib/python3.12/site-packages/scipy/sparse/_index.py:216: SparseEfficiencyWarning: Changing the sparsity structure of a csr_matrix is expensive. lil and dok are more efficient.
+ self._set_arrayXarray(i, j, x)
+WARNING:absl:Only the first input of the static gather operation is being used, others are being ignored at runtime.
+WARNING:absl:Only the first input of the static gather operation is being used, others are being ignored at runtime.
+WARNING:absl:Only the first input of the sorted segment sum operation is being used, others are being ignored at runtime.
+WARNING:absl:Only the first input of the static gather operation is being used, others are being ignored at runtime.
+WARNING:absl:Only the first input of the static gather operation is being used, others are being ignored at runtime.
+WARNING:absl:Only the first input of the static gather operation is being used, others are being ignored at runtime.
+WARNING:absl:Only the first input of the sorted segment sum operation is being used, others are being ignored at runtime.
+WARNING:root:No normalization scale found for cyclone_exists_gaussian_unit_mode
+WARNING:root:No normalization location found for cyclone_exists_gaussian_unit_mode
+WARNING:absl:Only the first input of the sorted segment sum operation is being used, others are being ignored at runtime.
+WARNING:absl:Only the first input of the sorted segment sum operation is being used, others are being ignored at runtime.
+WARNING:absl:Skipping gradient checkpointing for sequence length of 1
+WARNING:absl:No coordinates found for dim noise_channels of array Size: 128B
+JitTracer(float32[1,1,32])
+Dimensions without coordinates: batch, channels, noise_channels., using index for param name. This may make model surgery more error prone if changing coordinates at fine-tuning time.
+/usr/local/lib/python3.12/site-packages/scipy/sparse/_index.py:216: SparseEfficiencyWarning: Changing the sparsity structure of a csr_matrix is expensive. lil and dok are more efficient.
+ self._set_arrayXarray(i, j, x)
+WARNING:absl:Only the first input of the static gather operation is being used, others are being ignored at runtime.
+WARNING:absl:Only the first input of the static gather operation is being used, others are being ignored at runtime.
+WARNING:absl:Only the first input of the sorted segment sum operation is being used, others are being ignored at runtime.
+WARNING:absl:Only the first input of the static gather operation is being used, others are being ignored at runtime.
+WARNING:absl:Only the first input of the static gather operation is being used, others are being ignored at runtime.
+WARNING:absl:Only the first input of the static gather operation is being used, others are being ignored at runtime.
+WARNING:absl:Only the first input of the sorted segment sum operation is being used, others are being ignored at runtime.
+WARNING:root:No normalization scale found for cyclone_exists_gaussian_unit_mode
+WARNING:root:No normalization location found for cyclone_exists_gaussian_unit_mode
+WARNING:absl:Only the first input of the sorted segment sum operation is being used, others are being ignored at runtime.
+WARNING:absl:Only the first input of the sorted segment sum operation is being used, others are being ignored at runtime.
+WARNING:absl:Skipping gradient checkpointing for sequence length of 1
+WARNING:absl:No coordinates found for dim noise_channels of array Size: 128B
+JitTracer(float32[1,1,32])
+Dimensions without coordinates: batch, channels, noise_channels., using index for param name. This may make model surgery more error prone if changing coordinates at fine-tuning time.
+/usr/local/lib/python3.12/site-packages/scipy/sparse/_index.py:216: SparseEfficiencyWarning: Changing the sparsity structure of a csr_matrix is expensive. lil and dok are more efficient.
+ self._set_arrayXarray(i, j, x)
+WARNING:root:No normalization scale found for cyclone_exists_gaussian_unit_mode
+WARNING:root:No normalization location found for cyclone_exists_gaussian_unit_mode
+WARNING:absl:Skipping gradient checkpointing for sequence length of 1
+WARNING:absl:No coordinates found for dim noise_channels of array Size: 128B
+JitTracer(float32[1,1,32])
+Dimensions without coordinates: batch, channels, noise_channels., using index for param name. This may make model surgery more error prone if changing coordinates at fine-tuning time.
+/usr/local/lib/python3.12/site-packages/scipy/sparse/_index.py:216: SparseEfficiencyWarning: Changing the sparsity structure of a csr_matrix is expensive. lil and dok are more efficient.
+ self._set_arrayXarray(i, j, x)
+WARNING:root:No normalization scale found for cyclone_exists_gaussian_unit_mode
+WARNING:root:No normalization location found for cyclone_exists_gaussian_unit_mode
+WARNING:absl:Skipping gradient checkpointing for sequence length of 1
+WARNING:absl:No coordinates found for dim noise_channels of array Size: 128B
+JitTracer(float32[1,1,32])
+Dimensions without coordinates: batch, channels, noise_channels., using index for param name. This may make model surgery more error prone if changing coordinates at fine-tuning time.
+/usr/local/lib/python3.12/site-packages/scipy/sparse/_index.py:216: SparseEfficiencyWarning: Changing the sparsity structure of a csr_matrix is expensive. lil and dok are more efficient.
+ self._set_arrayXarray(i, j, x)
+WARNING:root:No normalization scale found for cyclone_exists_gaussian_unit_mode
+WARNING:root:No normalization location found for cyclone_exists_gaussian_unit_mode
+WARNING:absl:Skipping gradient checkpointing for sequence length of 1
+WARNING:absl:No coordinates found for dim noise_channels of array Size: 128B
+JitTracer(float32[1,1,32])
+Dimensions without coordinates: batch, channels, noise_channels., using index for param name. This may make model surgery more error prone if changing coordinates at fine-tuning time.
+/usr/local/lib/python3.12/site-packages/scipy/sparse/_index.py:216: SparseEfficiencyWarning: Changing the sparsity structure of a csr_matrix is expensive. lil and dok are more efficient.
+ self._set_arrayXarray(i, j, x)
+WARNING:root:No normalization scale found for cyclone_exists_gaussian_unit_mode
+WARNING:root:No normalization location found for cyclone_exists_gaussian_unit_mode
+WARNING:absl:Skipping gradient checkpointing for sequence length of 1
+WARNING:absl:No coordinates found for dim noise_channels of array Size: 128B
+JitTracer(float32[1,1,32])
+Dimensions without coordinates: batch, channels, noise_channels., using index for param name. This may make model surgery more error prone if changing coordinates at fine-tuning time.
+/usr/local/lib/python3.12/site-packages/scipy/sparse/_index.py:216: SparseEfficiencyWarning: Changing the sparsity structure of a csr_matrix is expensive. lil and dok are more efficient.
+ self._set_arrayXarray(i, j, x)
+WARNING:absl:Only the first input of the static gather operation is being used, others are being ignored at runtime.
+WARNING:absl:Only the first input of the static gather operation is being used, others are being ignored at runtime.
+WARNING:absl:Only the first input of the sorted segment sum operation is being used, others are being ignored at runtime.
+WARNING:absl:Only the first input of the static gather operation is being used, others are being ignored at runtime.
+WARNING:absl:Only the first input of the static gather operation is being used, others are being ignored at runtime.
+WARNING:absl:Only the first input of the static gather operation is being used, others are being ignored at runtime.
+WARNING:absl:Only the first input of the sorted segment sum operation is being used, others are being ignored at runtime.
+WARNING:root:No normalization scale found for cyclone_exists_gaussian_unit_mode
+WARNING:root:No normalization location found for cyclone_exists_gaussian_unit_mode
+WARNING:absl:Only the first input of the sorted segment sum operation is being used, others are being ignored at runtime.
+WARNING:absl:Only the first input of the sorted segment sum operation is being used, others are being ignored at runtime.
+WARNING:absl:Skipping gradient checkpointing for sequence length of 1
+WARNING:absl:No coordinates found for dim noise_channels of array Size: 128B
+JitTracer(float32[1,1,32])
+Dimensions without coordinates: batch, channels, noise_channels., using index for param name. This may make model surgery more error prone if changing coordinates at fine-tuning time.
+/usr/local/lib/python3.12/site-packages/scipy/sparse/_index.py:216: SparseEfficiencyWarning: Changing the sparsity structure of a csr_matrix is expensive. lil and dok are more efficient.
+ self._set_arrayXarray(i, j, x)
+WARNING:absl:Only the first input of the static gather operation is being used, others are being ignored at runtime.
+WARNING:absl:Only the first input of the static gather operation is being used, others are being ignored at runtime.
+WARNING:absl:Only the first input of the sorted segment sum operation is being used, others are being ignored at runtime.
+WARNING:absl:Only the first input of the static gather operation is being used, others are being ignored at runtime.
+WARNING:absl:Only the first input of the static gather operation is being used, others are being ignored at runtime.
+WARNING:absl:Only the first input of the static gather operation is being used, others are being ignored at runtime.
+WARNING:absl:Only the first input of the sorted segment sum operation is being used, others are being ignored at runtime.
+WARNING:root:No normalization scale found for cyclone_exists_gaussian_unit_mode
+WARNING:root:No normalization location found for cyclone_exists_gaussian_unit_mode
+WARNING:absl:Only the first input of the sorted segment sum operation is being used, others are being ignored at runtime.
+WARNING:absl:Only the first input of the sorted segment sum operation is being used, others are being ignored at runtime.
+WARNING:absl:Skipping gradient checkpointing for sequence length of 1
+WARNING:absl:No coordinates found for dim noise_channels of array Size: 128B
+JitTracer(float32[1,1,32])
+Dimensions without coordinates: batch, channels, noise_channels., using index for param name. This may make model surgery more error prone if changing coordinates at fine-tuning time.
+/usr/local/lib/python3.12/site-packages/scipy/sparse/_index.py:216: SparseEfficiencyWarning: Changing the sparsity structure of a csr_matrix is expensive. lil and dok are more efficient.
+ self._set_arrayXarray(i, j, x)
+WARNING:root:No normalization scale found for cyclone_exists_gaussian_unit_mode
+WARNING:root:No normalization location found for cyclone_exists_gaussian_unit_mode
+WARNING:absl:Skipping gradient checkpointing for sequence length of 1
+WARNING:absl:No coordinates found for dim noise_channels of array Size: 128B
+JitTracer(float32[1,1,32])
+Dimensions without coordinates: batch, channels, noise_channels., using index for param name. This may make model surgery more error prone if changing coordinates at fine-tuning time.
+/usr/local/lib/python3.12/site-packages/scipy/sparse/_index.py:216: SparseEfficiencyWarning: Changing the sparsity structure of a csr_matrix is expensive. lil and dok are more efficient.
+ self._set_arrayXarray(i, j, x)
+WARNING:root:No normalization scale found for cyclone_exists_gaussian_unit_mode
+WARNING:root:No normalization location found for cyclone_exists_gaussian_unit_mode
+WARNING:absl:Skipping gradient checkpointing for sequence length of 1
+WARNING:absl:No coordinates found for dim noise_channels of array Size: 128B
+JitTracer(float32[1,1,32])
+Dimensions without coordinates: batch, channels, noise_channels., using index for param name. This may make model surgery more error prone if changing coordinates at fine-tuning time.
+/usr/local/lib/python3.12/site-packages/scipy/sparse/_index.py:216: SparseEfficiencyWarning: Changing the sparsity structure of a csr_matrix is expensive. lil and dok are more efficient.
+ self._set_arrayXarray(i, j, x)
+WARNING:root:No normalization scale found for cyclone_exists_gaussian_unit_mode
+WARNING:root:No normalization location found for cyclone_exists_gaussian_unit_mode
+
+exit 0 wall 444s
diff --git a/docs/results/finetune/h200/grads01/pytest_h200.txt b/docs/results/finetune/h200/grads01/pytest_h200.txt
new file mode 100644
index 0000000..4ded940
--- /dev/null
+++ b/docs/results/finetune/h200/grads01/pytest_h200.txt
@@ -0,0 +1,10 @@
+..................... [100%]
+=============================== warnings summary ===============================
+tests/test_pallas_kernel.py: 55 warnings
+ /usr/local/lib/python3.12/site-packages/jax/_src/pallas/pallas_call.py:926: DeprecationWarning: The Pallas Triton backend is deprecated and will be removed in a future JAX version. To keep using Pallas on GPU, please migrate to the Mosaic GPU backend. To keep using Triton, switch to the official Triton bindings and jax_triton.
+ return rule(ctx, *in_nodes, compiler_params=compiler_params, **params)
+
+-- Docs: https://docs.pytest.org/en/stable/how-to/capture-warnings.html
+21 passed, 55 warnings in 34.91s
+
+exit 0 wall 37s
diff --git a/docs/results/finetune/h200/grads02/full_grads_mini_crps_h200.json b/docs/results/finetune/h200/grads02/full_grads_mini_crps_h200.json
new file mode 100644
index 0000000..6121329
--- /dev/null
+++ b/docs/results/finetune/h200/grads02/full_grads_mini_crps_h200.json
@@ -0,0 +1,290 @@
+{
+ "model": "WeatherNextCyclones_Mini",
+ "device": "NVIDIA H200",
+ "jax": "0.11.2",
+ "faster_weathernext": "0.1.1",
+ "blocks": {
+ "grid_chunk": 8192,
+ "edge_chunk": 16384
+ },
+ "masked": [],
+ "modules": "all",
+ "loss": "crps",
+ "head_params": 56684133,
+ "precisions": {
+ "highest": {
+ "official": {
+ "temp_gib": 102.06442984193563,
+ "compile_s": 82.44018316268921,
+ "step_s": 0.8911960124969482
+ },
+ "official_mha": {
+ "temp_gib": 123.4894026145339,
+ "compile_s": 40.56615233421326,
+ "step_s": 1.0314137935638428
+ },
+ "fwn": {
+ "temp_gib": 3.9678902104496956,
+ "compile_s": 66.66005754470825,
+ "step_s": 1.6664817333221436
+ },
+ "fwn_xla": {
+ "temp_gib": 4.092704713344574,
+ "compile_s": 44.69057846069336,
+ "step_s": 1.2443792819976807
+ },
+ "fwn_strict": {
+ "temp_gib": 3.9678902104496956,
+ "compile_s": 34.29314994812012,
+ "step_s": 1.5496768951416016
+ },
+ "fwn_noremat": {
+ "temp_gib": 39.10420774668455,
+ "compile_s": 36.4343626499176
+ },
+ "reference": "official",
+ "checks": {
+ "fwn vs official": {
+ "loss": {
+ "official": 1.2029575109481812,
+ "other": 1.2029576301574707,
+ "rel_diff": 9.909684129809331e-08,
+ "max_per_var_rel_diff": 1.4516024574512987e-05
+ },
+ "grads": {
+ "global_rel_l2": 0.0014249121348155792,
+ "global_cos": 0.9999989879096962,
+ "max_rel_l2": 0.004777954681938258,
+ "n_params": 376,
+ "n_compared": 335,
+ "n_zero_both": 6,
+ "n_zero_ref_only": 0,
+ "all_finite": true,
+ "worst": [
+ {
+ "param": "multimodality_forward/mesh_transformer/transformer/block_06/block_06_norm_conditioning_1/linear/b",
+ "size": 1024,
+ "ref_norm": 0.01195641972747539,
+ "got_norm": 0.01194987731837765,
+ "finite": true,
+ "rel_l2": 0.004777954681938258
+ },
+ {
+ "param": "multimodality_forward/mesh_transformer/transformer/block_04/mha_proj_k/w",
+ "size": 262144,
+ "ref_norm": 0.1598132224108344,
+ "got_norm": 0.16002500815732007,
+ "finite": true,
+ "rel_l2": 0.004576637809058942
+ },
+ {
+ "param": "multimodality_forward/mesh_transformer/transformer/block_07/mha_proj_k/w",
+ "size": 262144,
+ "ref_norm": 0.08604880457360886,
+ "got_norm": 0.08604577721067563,
+ "finite": true,
+ "rel_l2": 0.003971343797541177
+ },
+ {
+ "param": "multimodality_forward/mesh_transformer/transformer/block_07/block_07_norm_conditioning/linear/b",
+ "size": 1024,
+ "ref_norm": 0.011988146018327376,
+ "got_norm": 0.011985096670967681,
+ "finite": true,
+ "rel_l2": 0.0039110883371691
+ },
+ {
+ "param": "multimodality_forward/mesh_transformer/transformer/block_04/block_04_norm_conditioning/linear/b",
+ "size": 1024,
+ "ref_norm": 0.022245279121563027,
+ "got_norm": 0.022209942462725216,
+ "finite": true,
+ "rel_l2": 0.003792244858850051
+ }
+ ]
+ }
+ },
+ "fwn_strict vs official": {
+ "loss": {
+ "official": 1.2029575109481812,
+ "other": 1.2029576301574707,
+ "rel_diff": 9.909684129809331e-08,
+ "max_per_var_rel_diff": 9.720004779955291e-06
+ },
+ "grads": {
+ "global_rel_l2": 0.00082379129584239,
+ "global_cos": 0.999999660689066,
+ "max_rel_l2": 0.001948416906426157,
+ "n_params": 376,
+ "n_compared": 335,
+ "n_zero_both": 6,
+ "n_zero_ref_only": 0,
+ "all_finite": true,
+ "worst": [
+ {
+ "param": "multimodality_forward/grid_decoder/split_output_linear/w_u_component_of_wind_level=50,100,150,200,250,300,400,500,600,700,850,925,1000_time=21600",
+ "size": 6656,
+ "ref_norm": 0.06688884643267988,
+ "got_norm": 0.06687694200991241,
+ "finite": true,
+ "rel_l2": 0.001948416906426157
+ },
+ {
+ "param": "multimodality_forward/mesh_transformer/transformer/block_07/block_07_norm_conditioning/linear/b",
+ "size": 1024,
+ "ref_norm": 0.011988146018327376,
+ "got_norm": 0.011984917021023214,
+ "finite": true,
+ "rel_l2": 0.0018331874153464593
+ },
+ {
+ "param": "multimodality_forward/mesh_transformer/transformer/block_07/block_07_norm_conditioning/linear/w",
+ "size": 32768,
+ "ref_norm": 0.01886520837203217,
+ "got_norm": 0.018867645515592897,
+ "finite": true,
+ "rel_l2": 0.0017034554033325666
+ },
+ {
+ "param": "multimodality_forward/mesh_transformer/transformer/block_06/block_06_norm_conditioning_1/linear/w",
+ "size": 32768,
+ "ref_norm": 0.014629163223468801,
+ "got_norm": 0.014630775591435398,
+ "finite": true,
+ "rel_l2": 0.0016952843201718496
+ },
+ {
+ "param": "multimodality_forward/mesh_transformer/transformer/block_03/block_03_norm_conditioning/linear/b",
+ "size": 1024,
+ "ref_norm": 0.01752154986441303,
+ "got_norm": 0.017536811007932272,
+ "finite": true,
+ "rel_l2": 0.0016951421608641167
+ }
+ ]
+ }
+ },
+ "fwn_xla vs official": {
+ "loss": {
+ "official": 1.2029575109481812,
+ "other": 1.2029575109481812,
+ "rel_diff": 0.0,
+ "max_per_var_rel_diff": 5.806409829805195e-06
+ },
+ "grads": {
+ "global_rel_l2": 0.004334912614489594,
+ "global_cos": 0.9999906138683067,
+ "max_rel_l2": 0.02473884469355349,
+ "n_params": 376,
+ "n_compared": 335,
+ "n_zero_both": 6,
+ "n_zero_ref_only": 0,
+ "all_finite": true,
+ "worst": [
+ {
+ "param": "multimodality_forward/mesh_transformer/transformer/block_06/mha_proj_k/w",
+ "size": 262144,
+ "ref_norm": 0.07950285521643445,
+ "got_norm": 0.08024062065162438,
+ "finite": true,
+ "rel_l2": 0.02473884469355349
+ },
+ {
+ "param": "multimodality_forward/mesh_transformer/transformer/block_06/mha_proj_q/w",
+ "size": 262144,
+ "ref_norm": 0.05665707278120561,
+ "got_norm": 0.05680514404944998,
+ "finite": true,
+ "rel_l2": 0.018132354048167147
+ },
+ {
+ "param": "multimodality_forward/mesh_transformer/transformer/block_04/mha_proj_k/w",
+ "size": 262144,
+ "ref_norm": 0.1598132224108344,
+ "got_norm": 0.1586774491252515,
+ "finite": true,
+ "rel_l2": 0.017213455613771984
+ },
+ {
+ "param": "multimodality_forward/mesh_transformer/transformer/block_04/mha_proj_q/w",
+ "size": 262144,
+ "ref_norm": 0.1691768000001496,
+ "got_norm": 0.16816582677615846,
+ "finite": true,
+ "rel_l2": 0.015298984593562914
+ },
+ {
+ "param": "multimodality_forward/mesh_transformer/transformer/block_04/block_04_norm_conditioning/linear/b",
+ "size": 1024,
+ "ref_norm": 0.022245279121563027,
+ "got_norm": 0.02204641328770928,
+ "finite": true,
+ "rel_l2": 0.015150999685193739
+ }
+ ]
+ }
+ },
+ "official_mha vs official (floor)": {
+ "loss": {
+ "official": 1.2029575109481812,
+ "other": 1.2029576301574707,
+ "rel_diff": 9.909684129809331e-08,
+ "max_per_var_rel_diff": 1.998744922182942e-05
+ },
+ "grads": {
+ "global_rel_l2": 0.0007462398071607881,
+ "global_cos": 0.9999997216372897,
+ "max_rel_l2": 0.002071401419693404,
+ "n_params": 376,
+ "n_compared": 335,
+ "n_zero_both": 6,
+ "n_zero_ref_only": 0,
+ "all_finite": true,
+ "worst": [
+ {
+ "param": "multimodality_forward/mesh_transformer/transformer/block_03/block_03_norm_conditioning/linear/b",
+ "size": 1024,
+ "ref_norm": 0.01752154986441303,
+ "got_norm": 0.017540590315363827,
+ "finite": true,
+ "rel_l2": 0.002071401419693404
+ },
+ {
+ "param": "multimodality_forward/grid_decoder/split_output_linear/w_specific_humidity_level=50,100,150,200,250,300,400,500,600,700,850,925,1000_time=21600",
+ "size": 6656,
+ "ref_norm": 0.06548196123220153,
+ "got_norm": 0.06549546041100494,
+ "finite": true,
+ "rel_l2": 0.0019956467371860325
+ },
+ {
+ "param": "multimodality_forward/mesh_transformer/transformer/block_06/block_06_norm_conditioning_1/linear/w",
+ "size": 32768,
+ "ref_norm": 0.014629163223468801,
+ "got_norm": 0.014631738075070958,
+ "finite": true,
+ "rel_l2": 0.0019862609637362125
+ },
+ {
+ "param": "multimodality_forward/mesh_transformer/transformer/block_01/block_01_norm_conditioning/linear/b",
+ "size": 1024,
+ "ref_norm": 0.036950158188136205,
+ "got_norm": 0.037005963950206194,
+ "finite": true,
+ "rel_l2": 0.001872436446575433
+ },
+ {
+ "param": "multimodality_forward/mesh_transformer/transformer/block_01/block_01_norm_conditioning/linear/w",
+ "size": 32768,
+ "ref_norm": 0.05516810633438539,
+ "got_norm": 0.05522818914018945,
+ "finite": true,
+ "rel_l2": 0.0017419385949483557
+ }
+ ]
+ }
+ }
+ }
+ }
+ }
+}
\ No newline at end of file
diff --git a/docs/results/finetune/h200/grads02/full_grads_mini_crps_h200.log b/docs/results/finetune/h200/grads02/full_grads_mini_crps_h200.log
new file mode 100644
index 0000000..beddec9
--- /dev/null
+++ b/docs/results/finetune/h200/grads02/full_grads_mini_crps_h200.log
@@ -0,0 +1,119 @@
+WeatherNextCyclones_Mini: differentiating 56.68M params (all); masked targets: []
+[highest]
+ official highest temp 102.06 GiB step 0.89s loss 1.202958
+ official_mha highest temp 123.49 GiB step 1.03s loss 1.202958
+ fwn highest temp 3.97 GiB step 1.67s loss 1.202958
+ fwn_xla highest temp 4.09 GiB step 1.24s loss 1.202958
+ fwn_strict highest temp 3.97 GiB step 1.55s loss 1.202958
+ fwn_noremat highest temp 39.10 GiB (compile only)
+ fwn vs official loss rel diff 9.91e-08 (per var max 1.45e-05) grads: global rel L2 1.42e-03 cos 0.9999989879 max per-param rel L2 4.78e-03 (335/376 params; zero in both 6, zero only in official 0) finite True
+ 4.78e-03 multimodality_forward/mesh_transformer/transformer/block_06/block_06_norm_conditioning_1/linear/b
+ 4.58e-03 multimodality_forward/mesh_transformer/transformer/block_04/mha_proj_k/w
+ 3.97e-03 multimodality_forward/mesh_transformer/transformer/block_07/mha_proj_k/w
+ fwn_strict vs official loss rel diff 9.91e-08 (per var max 9.72e-06) grads: global rel L2 8.24e-04 cos 0.9999996607 max per-param rel L2 1.95e-03 (335/376 params; zero in both 6, zero only in official 0) finite True
+ 1.95e-03 multimodality_forward/grid_decoder/split_output_linear/w_u_component_of_wind_level=50,100,150,200,250,300,400,500,600,700,850,925,1000_time=21600
+ 1.83e-03 multimodality_forward/mesh_transformer/transformer/block_07/block_07_norm_conditioning/linear/b
+ 1.70e-03 multimodality_forward/mesh_transformer/transformer/block_07/block_07_norm_conditioning/linear/w
+ fwn_xla vs official loss rel diff 0.00e+00 (per var max 5.81e-06) grads: global rel L2 4.33e-03 cos 0.9999906139 max per-param rel L2 2.47e-02 (335/376 params; zero in both 6, zero only in official 0) finite True
+ 2.47e-02 multimodality_forward/mesh_transformer/transformer/block_06/mha_proj_k/w
+ 1.81e-02 multimodality_forward/mesh_transformer/transformer/block_06/mha_proj_q/w
+ 1.72e-02 multimodality_forward/mesh_transformer/transformer/block_04/mha_proj_k/w
+ official_mha vs official (floor) loss rel diff 9.91e-08 (per var max 2.00e-05) grads: global rel L2 7.46e-04 cos 0.9999997216 max per-param rel L2 2.07e-03 (335/376 params; zero in both 6, zero only in official 0) finite True
+ 2.07e-03 multimodality_forward/mesh_transformer/transformer/block_03/block_03_norm_conditioning/linear/b
+ 2.00e-03 multimodality_forward/grid_decoder/split_output_linear/w_specific_humidity_level=50,100,150,200,250,300,400,500,600,700,850,925,1000_time=21600
+ 1.99e-03 multimodality_forward/mesh_transformer/transformer/block_06/block_06_norm_conditioning_1/linear/w
+wrote /data/results/grads02/full_grads_mini_crps_h200.json
+WARNING:root:No normalization location found for cyclone_exists_gaussian_unit_mode
+WARNING:root:No normalization scale found for cyclone_exists_gaussian_unit_mode
+WARNING:absl:No coordinates found for dim noise_channels of array Size: 256B
+JitTracer(float32[2,1,32])
+Coordinates:
+ * batch (batch) object 16B MultiIndex
+ * _batch (batch) int64 16B 0 0
+ * sample (batch) int64 16B 0 1
+Dimensions without coordinates: channels, noise_channels., using index for param name. This may make model surgery more error prone if changing coordinates at fine-tuning time.
+/usr/local/lib/python3.12/site-packages/scipy/sparse/_index.py:216: SparseEfficiencyWarning: Changing the sparsity structure of a csr_matrix is expensive. lil and dok are more efficient.
+ self._set_arrayXarray(i, j, x)
+WARNING:absl:Only the first input of the static gather operation is being used, others are being ignored at runtime.
+WARNING:absl:Only the first input of the static gather operation is being used, others are being ignored at runtime.
+WARNING:absl:Only the first input of the sorted segment sum operation is being used, others are being ignored at runtime.
+WARNING:absl:Only the first input of the static gather operation is being used, others are being ignored at runtime.
+WARNING:absl:Only the first input of the static gather operation is being used, others are being ignored at runtime.
+WARNING:absl:Only the first input of the static gather operation is being used, others are being ignored at runtime.
+WARNING:absl:Only the first input of the sorted segment sum operation is being used, others are being ignored at runtime.
+/usr/local/lib/python3.12/site-packages/weathernext/weathernext2/architecture.py:286: FutureWarning: Deleting a single level of a MultiIndex is deprecated. Previously, this deleted all levels of a MultiIndex. Please also drop the following variables: {'batch'} to avoid an error in the future.
+ predictions[var] = xarray_jax.apply_ufunc(activation_fn, predictions[var])
+WARNING:root:No normalization location found for cyclone_exists_gaussian_unit_mode
+WARNING:root:No normalization scale found for cyclone_exists_gaussian_unit_mode
+WARNING:absl:No coordinates found for dim noise_channels of array Size: 256B
+JitTracer(float32[2,1,32])
+Coordinates:
+ * batch (batch) object 16B MultiIndex
+ * _batch (batch) int64 16B 0 0
+ * sample (batch) int64 16B 0 1
+Dimensions without coordinates: channels, noise_channels., using index for param name. This may make model surgery more error prone if changing coordinates at fine-tuning time.
+/usr/local/lib/python3.12/site-packages/scipy/sparse/_index.py:216: SparseEfficiencyWarning: Changing the sparsity structure of a csr_matrix is expensive. lil and dok are more efficient.
+ self._set_arrayXarray(i, j, x)
+WARNING:absl:Only the first input of the static gather operation is being used, others are being ignored at runtime.
+WARNING:absl:Only the first input of the static gather operation is being used, others are being ignored at runtime.
+WARNING:absl:Only the first input of the sorted segment sum operation is being used, others are being ignored at runtime.
+WARNING:absl:Only the first input of the static gather operation is being used, others are being ignored at runtime.
+WARNING:absl:Only the first input of the static gather operation is being used, others are being ignored at runtime.
+WARNING:absl:Only the first input of the static gather operation is being used, others are being ignored at runtime.
+WARNING:absl:Only the first input of the sorted segment sum operation is being used, others are being ignored at runtime.
+/usr/local/lib/python3.12/site-packages/weathernext/weathernext2/architecture.py:286: FutureWarning: Deleting a single level of a MultiIndex is deprecated. Previously, this deleted all levels of a MultiIndex. Please also drop the following variables: {'batch'} to avoid an error in the future.
+ predictions[var] = xarray_jax.apply_ufunc(activation_fn, predictions[var])
+WARNING:root:No normalization location found for cyclone_exists_gaussian_unit_mode
+WARNING:root:No normalization scale found for cyclone_exists_gaussian_unit_mode
+WARNING:absl:No coordinates found for dim noise_channels of array Size: 256B
+JitTracer(float32[2,1,32])
+Coordinates:
+ * batch (batch) object 16B MultiIndex
+ * _batch (batch) int64 16B 0 0
+ * sample (batch) int64 16B 0 1
+Dimensions without coordinates: channels, noise_channels., using index for param name. This may make model surgery more error prone if changing coordinates at fine-tuning time.
+/usr/local/lib/python3.12/site-packages/scipy/sparse/_index.py:216: SparseEfficiencyWarning: Changing the sparsity structure of a csr_matrix is expensive. lil and dok are more efficient.
+ self._set_arrayXarray(i, j, x)
+/usr/local/lib/python3.12/site-packages/weathernext/weathernext2/architecture.py:286: FutureWarning: Deleting a single level of a MultiIndex is deprecated. Previously, this deleted all levels of a MultiIndex. Please also drop the following variables: {'batch'} to avoid an error in the future.
+ predictions[var] = xarray_jax.apply_ufunc(activation_fn, predictions[var])
+WARNING:root:No normalization location found for cyclone_exists_gaussian_unit_mode
+WARNING:root:No normalization scale found for cyclone_exists_gaussian_unit_mode
+WARNING:absl:No coordinates found for dim noise_channels of array Size: 256B
+JitTracer(float32[2,1,32])
+Coordinates:
+ * batch (batch) object 16B MultiIndex
+ * _batch (batch) int64 16B 0 0
+ * sample (batch) int64 16B 0 1
+Dimensions without coordinates: channels, noise_channels., using index for param name. This may make model surgery more error prone if changing coordinates at fine-tuning time.
+/usr/local/lib/python3.12/site-packages/scipy/sparse/_index.py:216: SparseEfficiencyWarning: Changing the sparsity structure of a csr_matrix is expensive. lil and dok are more efficient.
+ self._set_arrayXarray(i, j, x)
+/usr/local/lib/python3.12/site-packages/weathernext/weathernext2/architecture.py:286: FutureWarning: Deleting a single level of a MultiIndex is deprecated. Previously, this deleted all levels of a MultiIndex. Please also drop the following variables: {'batch'} to avoid an error in the future.
+ predictions[var] = xarray_jax.apply_ufunc(activation_fn, predictions[var])
+WARNING:root:No normalization location found for cyclone_exists_gaussian_unit_mode
+WARNING:root:No normalization scale found for cyclone_exists_gaussian_unit_mode
+WARNING:absl:No coordinates found for dim noise_channels of array Size: 256B
+JitTracer(float32[2,1,32])
+Coordinates:
+ * batch (batch) object 16B MultiIndex
+ * _batch (batch) int64 16B 0 0
+ * sample (batch) int64 16B 0 1
+Dimensions without coordinates: channels, noise_channels., using index for param name. This may make model surgery more error prone if changing coordinates at fine-tuning time.
+/usr/local/lib/python3.12/site-packages/scipy/sparse/_index.py:216: SparseEfficiencyWarning: Changing the sparsity structure of a csr_matrix is expensive. lil and dok are more efficient.
+ self._set_arrayXarray(i, j, x)
+/usr/local/lib/python3.12/site-packages/weathernext/weathernext2/architecture.py:286: FutureWarning: Deleting a single level of a MultiIndex is deprecated. Previously, this deleted all levels of a MultiIndex. Please also drop the following variables: {'batch'} to avoid an error in the future.
+ predictions[var] = xarray_jax.apply_ufunc(activation_fn, predictions[var])
+WARNING:root:No normalization location found for cyclone_exists_gaussian_unit_mode
+WARNING:root:No normalization scale found for cyclone_exists_gaussian_unit_mode
+WARNING:absl:No coordinates found for dim noise_channels of array Size: 256B
+JitTracer(float32[2,1,32])
+Coordinates:
+ * batch (batch) object 16B MultiIndex
+ * _batch (batch) int64 16B 0 0
+ * sample (batch) int64 16B 0 1
+Dimensions without coordinates: channels, noise_channels., using index for param name. This may make model surgery more error prone if changing coordinates at fine-tuning time.
+/usr/local/lib/python3.12/site-packages/scipy/sparse/_index.py:216: SparseEfficiencyWarning: Changing the sparsity structure of a csr_matrix is expensive. lil and dok are more efficient.
+ self._set_arrayXarray(i, j, x)
+/usr/local/lib/python3.12/site-packages/weathernext/weathernext2/architecture.py:286: FutureWarning: Deleting a single level of a MultiIndex is deprecated. Previously, this deleted all levels of a MultiIndex. Please also drop the following variables: {'batch'} to avoid an error in the future.
+ predictions[var] = xarray_jax.apply_ufunc(activation_fn, predictions[var])
+
+exit 0 wall 332s
diff --git a/docs/results/finetune/h200/grads02/full_grads_mini_linear_h200.json b/docs/results/finetune/h200/grads02/full_grads_mini_linear_h200.json
new file mode 100644
index 0000000..c8ed071
--- /dev/null
+++ b/docs/results/finetune/h200/grads02/full_grads_mini_linear_h200.json
@@ -0,0 +1,290 @@
+{
+ "model": "WeatherNextCyclones_Mini",
+ "device": "NVIDIA H200",
+ "jax": "0.11.2",
+ "faster_weathernext": "0.1.1",
+ "blocks": {
+ "grid_chunk": 8192,
+ "edge_chunk": 16384
+ },
+ "masked": [],
+ "modules": "all",
+ "loss": "linear",
+ "head_params": 56684133,
+ "precisions": {
+ "highest": {
+ "official": {
+ "temp_gib": 51.41228340566158,
+ "compile_s": 71.35185122489929,
+ "step_s": 0.5169267654418945
+ },
+ "official_mha": {
+ "temp_gib": 61.79655359685421,
+ "compile_s": 41.12771034240723,
+ "step_s": 0.5860462188720703
+ },
+ "fwn": {
+ "temp_gib": 3.299081988632679,
+ "compile_s": 63.962586402893066,
+ "step_s": 0.952063798904419
+ },
+ "fwn_xla": {
+ "temp_gib": 3.397009789943695,
+ "compile_s": 41.45921754837036,
+ "step_s": 0.7416081428527832
+ },
+ "fwn_strict": {
+ "temp_gib": 3.299081988632679,
+ "compile_s": 29.10938858985901,
+ "step_s": 0.8634288311004639
+ },
+ "fwn_noremat": {
+ "temp_gib": 19.658946715295315,
+ "compile_s": 32.132792949676514
+ },
+ "reference": "official",
+ "checks": {
+ "fwn vs official": {
+ "loss": {
+ "official": 506.6861267089844,
+ "other": 506.68609619140625,
+ "rel_diff": 6.02297487859339e-08,
+ "max_per_var_rel_diff": 5.983790193194335e-05
+ },
+ "grads": {
+ "global_rel_l2": 6.747983738157028e-05,
+ "global_cos": 0.9999999977251447,
+ "max_rel_l2": 0.00040668020635455663,
+ "n_params": 376,
+ "n_compared": 315,
+ "n_zero_both": 6,
+ "n_zero_ref_only": 0,
+ "all_finite": true,
+ "worst": [
+ {
+ "param": "multimodality_forward/mesh_transformer/transformer/block_01/block_01_norm_conditioning/linear/w",
+ "size": 32768,
+ "ref_norm": 0.8212241099797718,
+ "got_norm": 0.8211282246485626,
+ "finite": true,
+ "rel_l2": 0.00040668020635455663
+ },
+ {
+ "param": "multimodality_forward/mesh_transformer/transformer/block_01/block_01_norm_conditioning/linear/b",
+ "size": 1024,
+ "ref_norm": 0.9515768185772578,
+ "got_norm": 0.9514657150469644,
+ "finite": true,
+ "rel_l2": 0.0004066777278866067
+ },
+ {
+ "param": "multimodality_forward/mesh_transformer/transformer/block_02/block_02_norm_conditioning/linear/b",
+ "size": 1024,
+ "ref_norm": 0.570128324689503,
+ "got_norm": 0.570091227624361,
+ "finite": true,
+ "rel_l2": 0.00031003748064335117
+ },
+ {
+ "param": "multimodality_forward/mesh_transformer/transformer/block_02/block_02_norm_conditioning/linear/w",
+ "size": 32768,
+ "ref_norm": 0.4920287234745772,
+ "got_norm": 0.4919967082159768,
+ "finite": true,
+ "rel_l2": 0.0003100369667402391
+ },
+ {
+ "param": "multimodality_forward/mesh_transformer/transformer/block_03/block_03_norm_conditioning/linear/b",
+ "size": 1024,
+ "ref_norm": 0.6554303125710794,
+ "got_norm": 0.6553247150737159,
+ "finite": true,
+ "rel_l2": 0.0002989693185339831
+ }
+ ]
+ }
+ },
+ "fwn_strict vs official": {
+ "loss": {
+ "official": 506.6861267089844,
+ "other": 506.6861572265625,
+ "rel_diff": 6.02297487859339e-08,
+ "max_per_var_rel_diff": 8.986486355122179e-05
+ },
+ "grads": {
+ "global_rel_l2": 7.769181962751524e-05,
+ "global_cos": 0.9999999969823907,
+ "max_rel_l2": 0.00024482507727029086,
+ "n_params": 376,
+ "n_compared": 315,
+ "n_zero_both": 6,
+ "n_zero_ref_only": 0,
+ "all_finite": true,
+ "worst": [
+ {
+ "param": "multimodality_forward/grid_to_mesh_gnn/edge_encoder/shared_dense/mlp/linear_0/b",
+ "size": 32,
+ "ref_norm": 5.131798694299954,
+ "got_norm": 5.132980508586088,
+ "finite": true,
+ "rel_l2": 0.00024482507727029086
+ },
+ {
+ "param": "multimodality_forward/mesh_transformer/transformer/block_06/mha_proj_k/w",
+ "size": 262144,
+ "ref_norm": 2.7121599826891596,
+ "got_norm": 2.712255287913129,
+ "finite": true,
+ "rel_l2": 0.00023651406176779463
+ },
+ {
+ "param": "multimodality_forward/mesh_transformer/transformer/block_02/block_02_norm_conditioning_1/linear/w",
+ "size": 32768,
+ "ref_norm": 0.6123730296586022,
+ "got_norm": 0.6124950330649098,
+ "finite": true,
+ "rel_l2": 0.00022509363453914556
+ },
+ {
+ "param": "multimodality_forward/mesh_transformer/transformer/block_02/block_02_norm_conditioning_1/linear/b",
+ "size": 1024,
+ "ref_norm": 0.7095748580288647,
+ "got_norm": 0.709716221120885,
+ "finite": true,
+ "rel_l2": 0.00022508454435158547
+ },
+ {
+ "param": "multimodality_forward/mesh_transformer/transformer/block_04/block_04_norm_conditioning_1/linear/b",
+ "size": 1024,
+ "ref_norm": 0.5775796758362444,
+ "got_norm": 0.5776165145753557,
+ "finite": true,
+ "rel_l2": 0.00020761502928520162
+ }
+ ]
+ }
+ },
+ "fwn_xla vs official": {
+ "loss": {
+ "official": 506.6861267089844,
+ "other": 506.6861572265625,
+ "rel_diff": 6.02297487859339e-08,
+ "max_per_var_rel_diff": 2.55985248337014e-05
+ },
+ "grads": {
+ "global_rel_l2": 8.365192541712793e-05,
+ "global_cos": 0.9999999965319865,
+ "max_rel_l2": 0.0004991382707796397,
+ "n_params": 376,
+ "n_compared": 315,
+ "n_zero_both": 6,
+ "n_zero_ref_only": 0,
+ "all_finite": true,
+ "worst": [
+ {
+ "param": "multimodality_forward/grid_to_mesh_gnn/edge_encoder/shared_dense/mlp/linear_0/b",
+ "size": 32,
+ "ref_norm": 5.131798694299954,
+ "got_norm": 5.133924904310215,
+ "finite": true,
+ "rel_l2": 0.0004991382707796397
+ },
+ {
+ "param": "multimodality_forward/grid_to_mesh_gnn/edge_encoder/split_input_matmul/w_spatial_feature=distance,rel_x,rel_y,rel_z",
+ "size": 128,
+ "ref_norm": 5.380476093448445,
+ "got_norm": 5.38267862004577,
+ "finite": true,
+ "rel_l2": 0.0004279359698555622
+ },
+ {
+ "param": "multimodality_forward/grid_to_mesh_gnn/edge_encoder/shared_dense/normalization/linear_norm_conditioning/linear/w",
+ "size": 2048,
+ "ref_norm": 0.21243608807573908,
+ "got_norm": 0.21250627394911525,
+ "finite": true,
+ "rel_l2": 0.00041425873259760295
+ },
+ {
+ "param": "multimodality_forward/grid_to_mesh_gnn/edge_encoder/shared_dense/normalization/linear_norm_conditioning/linear/b",
+ "size": 64,
+ "ref_norm": 0.2461560175103653,
+ "got_norm": 0.2462373353025157,
+ "finite": true,
+ "rel_l2": 0.0004142047554172557
+ },
+ {
+ "param": "multimodality_forward/grid_to_mesh_gnn/edge_encoder/shared_dense/mlp/linear_1/b",
+ "size": 32,
+ "ref_norm": 23.439334198251515,
+ "got_norm": 23.444884591868277,
+ "finite": true,
+ "rel_l2": 0.00037849919702644623
+ }
+ ]
+ }
+ },
+ "official_mha vs official (floor)": {
+ "loss": {
+ "official": 506.6861267089844,
+ "other": 506.6861267089844,
+ "rel_diff": 0.0,
+ "max_per_var_rel_diff": 6.210612565138524e-05
+ },
+ "grads": {
+ "global_rel_l2": 7.588123000222357e-05,
+ "global_cos": 0.9999999972161001,
+ "max_rel_l2": 0.0006613859159092648,
+ "n_params": 376,
+ "n_compared": 315,
+ "n_zero_both": 6,
+ "n_zero_ref_only": 0,
+ "all_finite": true,
+ "worst": [
+ {
+ "param": "multimodality_forward/grid_to_mesh_gnn/edge_encoder/shared_dense/mlp/linear_0/b",
+ "size": 32,
+ "ref_norm": 5.131798694299954,
+ "got_norm": 5.129150950775782,
+ "finite": true,
+ "rel_l2": 0.0006613859159092648
+ },
+ {
+ "param": "multimodality_forward/grid_to_mesh_gnn/edge_encoder/split_input_matmul/w_spatial_feature=distance,rel_x,rel_y,rel_z",
+ "size": 128,
+ "ref_norm": 5.380476093448445,
+ "got_norm": 5.378281485478416,
+ "finite": true,
+ "rel_l2": 0.0004965656272527077
+ },
+ {
+ "param": "multimodality_forward/grid_to_mesh_gnn/edge_encoder/shared_dense/mlp/linear_1/b",
+ "size": 32,
+ "ref_norm": 23.439334198251515,
+ "got_norm": 23.43352481180425,
+ "finite": true,
+ "rel_l2": 0.0004930205840060552
+ },
+ {
+ "param": "multimodality_forward/grid_to_mesh_gnn/edge_encoder/shared_dense/normalization/linear_norm_conditioning/linear/w",
+ "size": 2048,
+ "ref_norm": 0.21243608807573908,
+ "got_norm": 0.21238719761427038,
+ "finite": true,
+ "rel_l2": 0.00047779131280530426
+ },
+ {
+ "param": "multimodality_forward/grid_to_mesh_gnn/edge_encoder/shared_dense/normalization/linear_norm_conditioning/linear/b",
+ "size": 64,
+ "ref_norm": 0.2461560175103653,
+ "got_norm": 0.24609936721373887,
+ "finite": true,
+ "rel_l2": 0.0004777884657070471
+ }
+ ]
+ }
+ }
+ }
+ }
+ }
+}
\ No newline at end of file
diff --git a/docs/results/finetune/h200/grads02/full_grads_mini_linear_h200.log b/docs/results/finetune/h200/grads02/full_grads_mini_linear_h200.log
new file mode 100644
index 0000000..f9df8e1
--- /dev/null
+++ b/docs/results/finetune/h200/grads02/full_grads_mini_linear_h200.log
@@ -0,0 +1,93 @@
+WeatherNextCyclones_Mini: differentiating 56.68M params (all); masked targets: []
+[highest]
+ official highest temp 51.41 GiB step 0.52s loss 506.686127
+ official_mha highest temp 61.80 GiB step 0.59s loss 506.686127
+ fwn highest temp 3.30 GiB step 0.95s loss 506.686096
+ fwn_xla highest temp 3.40 GiB step 0.74s loss 506.686157
+ fwn_strict highest temp 3.30 GiB step 0.86s loss 506.686157
+ fwn_noremat highest temp 19.66 GiB (compile only)
+ fwn vs official loss rel diff 6.02e-08 (per var max 5.98e-05) grads: global rel L2 6.75e-05 cos 0.9999999977 max per-param rel L2 4.07e-04 (315/376 params; zero in both 6, zero only in official 0) finite True
+ 4.07e-04 multimodality_forward/mesh_transformer/transformer/block_01/block_01_norm_conditioning/linear/w
+ 4.07e-04 multimodality_forward/mesh_transformer/transformer/block_01/block_01_norm_conditioning/linear/b
+ 3.10e-04 multimodality_forward/mesh_transformer/transformer/block_02/block_02_norm_conditioning/linear/b
+ fwn_strict vs official loss rel diff 6.02e-08 (per var max 8.99e-05) grads: global rel L2 7.77e-05 cos 0.9999999970 max per-param rel L2 2.45e-04 (315/376 params; zero in both 6, zero only in official 0) finite True
+ 2.45e-04 multimodality_forward/grid_to_mesh_gnn/edge_encoder/shared_dense/mlp/linear_0/b
+ 2.37e-04 multimodality_forward/mesh_transformer/transformer/block_06/mha_proj_k/w
+ 2.25e-04 multimodality_forward/mesh_transformer/transformer/block_02/block_02_norm_conditioning_1/linear/w
+ fwn_xla vs official loss rel diff 6.02e-08 (per var max 2.56e-05) grads: global rel L2 8.37e-05 cos 0.9999999965 max per-param rel L2 4.99e-04 (315/376 params; zero in both 6, zero only in official 0) finite True
+ 4.99e-04 multimodality_forward/grid_to_mesh_gnn/edge_encoder/shared_dense/mlp/linear_0/b
+ 4.28e-04 multimodality_forward/grid_to_mesh_gnn/edge_encoder/split_input_matmul/w_spatial_feature=distance,rel_x,rel_y,rel_z
+ 4.14e-04 multimodality_forward/grid_to_mesh_gnn/edge_encoder/shared_dense/normalization/linear_norm_conditioning/linear/w
+ official_mha vs official (floor) loss rel diff 0.00e+00 (per var max 6.21e-05) grads: global rel L2 7.59e-05 cos 0.9999999972 max per-param rel L2 6.61e-04 (315/376 params; zero in both 6, zero only in official 0) finite True
+ 6.61e-04 multimodality_forward/grid_to_mesh_gnn/edge_encoder/shared_dense/mlp/linear_0/b
+ 4.97e-04 multimodality_forward/grid_to_mesh_gnn/edge_encoder/split_input_matmul/w_spatial_feature=distance,rel_x,rel_y,rel_z
+ 4.93e-04 multimodality_forward/grid_to_mesh_gnn/edge_encoder/shared_dense/mlp/linear_1/b
+wrote /data/results/grads02/full_grads_mini_linear_h200.json
+WARNING:absl:Skipping gradient checkpointing for sequence length of 1
+WARNING:absl:No coordinates found for dim noise_channels of array Size: 128B
+JitTracer(float32[1,1,32])
+Dimensions without coordinates: batch, channels, noise_channels., using index for param name. This may make model surgery more error prone if changing coordinates at fine-tuning time.
+/usr/local/lib/python3.12/site-packages/scipy/sparse/_index.py:216: SparseEfficiencyWarning: Changing the sparsity structure of a csr_matrix is expensive. lil and dok are more efficient.
+ self._set_arrayXarray(i, j, x)
+WARNING:absl:Only the first input of the static gather operation is being used, others are being ignored at runtime.
+WARNING:absl:Only the first input of the static gather operation is being used, others are being ignored at runtime.
+WARNING:absl:Only the first input of the sorted segment sum operation is being used, others are being ignored at runtime.
+WARNING:absl:Only the first input of the static gather operation is being used, others are being ignored at runtime.
+WARNING:absl:Only the first input of the static gather operation is being used, others are being ignored at runtime.
+WARNING:absl:Only the first input of the static gather operation is being used, others are being ignored at runtime.
+WARNING:absl:Only the first input of the sorted segment sum operation is being used, others are being ignored at runtime.
+WARNING:root:No normalization scale found for cyclone_exists_gaussian_unit_mode
+WARNING:root:No normalization location found for cyclone_exists_gaussian_unit_mode
+WARNING:absl:Only the first input of the sorted segment sum operation is being used, others are being ignored at runtime.
+WARNING:absl:Only the first input of the sorted segment sum operation is being used, others are being ignored at runtime.
+WARNING:absl:Skipping gradient checkpointing for sequence length of 1
+WARNING:absl:No coordinates found for dim noise_channels of array Size: 128B
+JitTracer(float32[1,1,32])
+Dimensions without coordinates: batch, channels, noise_channels., using index for param name. This may make model surgery more error prone if changing coordinates at fine-tuning time.
+/usr/local/lib/python3.12/site-packages/scipy/sparse/_index.py:216: SparseEfficiencyWarning: Changing the sparsity structure of a csr_matrix is expensive. lil and dok are more efficient.
+ self._set_arrayXarray(i, j, x)
+WARNING:absl:Only the first input of the static gather operation is being used, others are being ignored at runtime.
+WARNING:absl:Only the first input of the static gather operation is being used, others are being ignored at runtime.
+WARNING:absl:Only the first input of the sorted segment sum operation is being used, others are being ignored at runtime.
+WARNING:absl:Only the first input of the static gather operation is being used, others are being ignored at runtime.
+WARNING:absl:Only the first input of the static gather operation is being used, others are being ignored at runtime.
+WARNING:absl:Only the first input of the static gather operation is being used, others are being ignored at runtime.
+WARNING:absl:Only the first input of the sorted segment sum operation is being used, others are being ignored at runtime.
+WARNING:root:No normalization scale found for cyclone_exists_gaussian_unit_mode
+WARNING:root:No normalization location found for cyclone_exists_gaussian_unit_mode
+WARNING:absl:Only the first input of the sorted segment sum operation is being used, others are being ignored at runtime.
+WARNING:absl:Only the first input of the sorted segment sum operation is being used, others are being ignored at runtime.
+WARNING:absl:Skipping gradient checkpointing for sequence length of 1
+WARNING:absl:No coordinates found for dim noise_channels of array Size: 128B
+JitTracer(float32[1,1,32])
+Dimensions without coordinates: batch, channels, noise_channels., using index for param name. This may make model surgery more error prone if changing coordinates at fine-tuning time.
+/usr/local/lib/python3.12/site-packages/scipy/sparse/_index.py:216: SparseEfficiencyWarning: Changing the sparsity structure of a csr_matrix is expensive. lil and dok are more efficient.
+ self._set_arrayXarray(i, j, x)
+WARNING:root:No normalization scale found for cyclone_exists_gaussian_unit_mode
+WARNING:root:No normalization location found for cyclone_exists_gaussian_unit_mode
+WARNING:absl:Skipping gradient checkpointing for sequence length of 1
+WARNING:absl:No coordinates found for dim noise_channels of array Size: 128B
+JitTracer(float32[1,1,32])
+Dimensions without coordinates: batch, channels, noise_channels., using index for param name. This may make model surgery more error prone if changing coordinates at fine-tuning time.
+/usr/local/lib/python3.12/site-packages/scipy/sparse/_index.py:216: SparseEfficiencyWarning: Changing the sparsity structure of a csr_matrix is expensive. lil and dok are more efficient.
+ self._set_arrayXarray(i, j, x)
+WARNING:root:No normalization scale found for cyclone_exists_gaussian_unit_mode
+WARNING:root:No normalization location found for cyclone_exists_gaussian_unit_mode
+WARNING:absl:Skipping gradient checkpointing for sequence length of 1
+WARNING:absl:No coordinates found for dim noise_channels of array Size: 128B
+JitTracer(float32[1,1,32])
+Dimensions without coordinates: batch, channels, noise_channels., using index for param name. This may make model surgery more error prone if changing coordinates at fine-tuning time.
+/usr/local/lib/python3.12/site-packages/scipy/sparse/_index.py:216: SparseEfficiencyWarning: Changing the sparsity structure of a csr_matrix is expensive. lil and dok are more efficient.
+ self._set_arrayXarray(i, j, x)
+WARNING:root:No normalization scale found for cyclone_exists_gaussian_unit_mode
+WARNING:root:No normalization location found for cyclone_exists_gaussian_unit_mode
+WARNING:absl:Skipping gradient checkpointing for sequence length of 1
+WARNING:absl:No coordinates found for dim noise_channels of array Size: 128B
+JitTracer(float32[1,1,32])
+Dimensions without coordinates: batch, channels, noise_channels., using index for param name. This may make model surgery more error prone if changing coordinates at fine-tuning time.
+/usr/local/lib/python3.12/site-packages/scipy/sparse/_index.py:216: SparseEfficiencyWarning: Changing the sparsity structure of a csr_matrix is expensive. lil and dok are more efficient.
+ self._set_arrayXarray(i, j, x)
+WARNING:root:No normalization scale found for cyclone_exists_gaussian_unit_mode
+WARNING:root:No normalization location found for cyclone_exists_gaussian_unit_mode
+
+exit 0 wall 301s
diff --git a/docs/results/finetune/h200/grads03/full_grads_mini_crps_h200.json b/docs/results/finetune/h200/grads03/full_grads_mini_crps_h200.json
new file mode 100644
index 0000000..d7cf5cd
--- /dev/null
+++ b/docs/results/finetune/h200/grads03/full_grads_mini_crps_h200.json
@@ -0,0 +1,290 @@
+{
+ "model": "WeatherNextCyclones_Mini",
+ "device": "NVIDIA H200",
+ "jax": "0.11.2",
+ "faster_weathernext": "0.1.1",
+ "blocks": {
+ "grid_chunk": 8192,
+ "edge_chunk": 16384
+ },
+ "masked": [],
+ "modules": "all",
+ "loss": "crps",
+ "head_params": 56684133,
+ "precisions": {
+ "highest": {
+ "official": {
+ "temp_gib": 102.06442984193563,
+ "compile_s": 75.40727353096008,
+ "step_s": 0.9012751579284668
+ },
+ "official_mha": {
+ "temp_gib": 123.4894026145339,
+ "compile_s": 39.643322706222534,
+ "step_s": 1.0274059772491455
+ },
+ "fwn": {
+ "temp_gib": 3.9678902104496956,
+ "compile_s": 63.50120997428894,
+ "step_s": 0.6766211986541748
+ },
+ "fwn_xla": {
+ "temp_gib": 4.085288941860199,
+ "compile_s": 45.142993688583374,
+ "step_s": 1.2472524642944336
+ },
+ "fwn_strict": {
+ "temp_gib": 3.9678902104496956,
+ "compile_s": 33.57709002494812,
+ "step_s": 1.5615160465240479
+ },
+ "fwn_noremat": {
+ "temp_gib": 39.10420774668455,
+ "compile_s": 34.61651587486267
+ },
+ "reference": "official",
+ "checks": {
+ "fwn vs official": {
+ "loss": {
+ "official": 1.2029576301574707,
+ "other": 1.2029576301574707,
+ "rel_diff": 0.0,
+ "max_per_var_rel_diff": 1.7419235324601308e-05
+ },
+ "grads": {
+ "global_rel_l2": 0.00473814728321582,
+ "global_cos": 0.9999887984066298,
+ "max_rel_l2": 0.023736378522844125,
+ "n_params": 376,
+ "n_compared": 335,
+ "n_zero_both": 6,
+ "n_zero_ref_only": 0,
+ "all_finite": true,
+ "worst": [
+ {
+ "param": "multimodality_forward/mesh_transformer/transformer/block_06/mha_proj_k/w",
+ "size": 262144,
+ "ref_norm": 0.08019620001056632,
+ "got_norm": 0.07947691365714485,
+ "finite": true,
+ "rel_l2": 0.023736378522844125
+ },
+ {
+ "param": "multimodality_forward/mesh_transformer/transformer/block_06/mha_proj_q/w",
+ "size": 262144,
+ "ref_norm": 0.056794953153429265,
+ "got_norm": 0.05664967498444661,
+ "finite": true,
+ "rel_l2": 0.018143519862337053
+ },
+ {
+ "param": "multimodality_forward/mesh_transformer/transformer/block_04/mha_proj_k/w",
+ "size": 262144,
+ "ref_norm": 0.15881242287265113,
+ "got_norm": 0.1599207868886793,
+ "finite": true,
+ "rel_l2": 0.0170244757017093
+ },
+ {
+ "param": "multimodality_forward/mesh_transformer/transformer/block_04/mha_proj_q/w",
+ "size": 262144,
+ "ref_norm": 0.1681339269725539,
+ "got_norm": 0.16911101226269407,
+ "finite": true,
+ "rel_l2": 0.015736561548825554
+ },
+ {
+ "param": "multimodality_forward/mesh_transformer/transformer/block_01/block_01_norm_conditioning/linear/b",
+ "size": 1024,
+ "ref_norm": 0.03658263549895236,
+ "got_norm": 0.03698980673862434,
+ "finite": true,
+ "rel_l2": 0.014712151310122266
+ }
+ ]
+ }
+ },
+ "fwn_strict vs official": {
+ "loss": {
+ "official": 1.2029576301574707,
+ "other": 1.2029578685760498,
+ "rel_diff": 1.9819366295582064e-07,
+ "max_per_var_rel_diff": 7.146352953682588e-06
+ },
+ "grads": {
+ "global_rel_l2": 0.001186370925268374,
+ "global_cos": 0.9999992979475228,
+ "max_rel_l2": 0.004945716845074838,
+ "n_params": 376,
+ "n_compared": 335,
+ "n_zero_both": 6,
+ "n_zero_ref_only": 0,
+ "all_finite": true,
+ "worst": [
+ {
+ "param": "multimodality_forward/mesh_encoder/split_input_matmul/w_spatial_feature=sin_lat,sin_lon,cos_lon",
+ "size": 1536,
+ "ref_norm": 0.3616832367875008,
+ "got_norm": 0.36165296196995134,
+ "finite": true,
+ "rel_l2": 0.004945716845074838
+ },
+ {
+ "param": "multimodality_forward/mesh_transformer/transformer/block_04/block_04_norm_conditioning/linear/w",
+ "size": 32768,
+ "ref_norm": 0.029310623408391093,
+ "got_norm": 0.029254179045911516,
+ "finite": true,
+ "rel_l2": 0.004820102464930829
+ },
+ {
+ "param": "multimodality_forward/mesh_transformer/transformer/block_01/block_01_norm_conditioning/linear/w",
+ "size": 32768,
+ "ref_norm": 0.054780597316956656,
+ "got_norm": 0.05497381215959125,
+ "finite": true,
+ "rel_l2": 0.004731905598605231
+ },
+ {
+ "param": "multimodality_forward/mesh_transformer/transformer/block_04/mha_proj_k/w",
+ "size": 262144,
+ "ref_norm": 0.15881242287265113,
+ "got_norm": 0.15895061747879408,
+ "finite": true,
+ "rel_l2": 0.0046290889919252225
+ },
+ {
+ "param": "multimodality_forward/mesh_transformer/transformer/block_02/block_02_norm_conditioning_1/linear/b",
+ "size": 1024,
+ "ref_norm": 0.015309892317326048,
+ "got_norm": 0.01535524848336557,
+ "finite": true,
+ "rel_l2": 0.0041600937433161666
+ }
+ ]
+ }
+ },
+ "fwn_xla vs official": {
+ "loss": {
+ "official": 1.2029576301574707,
+ "other": 1.2029575109481812,
+ "rel_diff": 9.909683147791032e-08,
+ "max_per_var_rel_diff": 2.13461627159536e-05
+ },
+ "grads": {
+ "global_rel_l2": 0.0011752381884746527,
+ "global_cos": 0.9999993145870859,
+ "max_rel_l2": 0.004903469398417187,
+ "n_params": 376,
+ "n_compared": 335,
+ "n_zero_both": 6,
+ "n_zero_ref_only": 0,
+ "all_finite": true,
+ "worst": [
+ {
+ "param": "multimodality_forward/mesh_transformer/transformer/block_04/block_04_norm_conditioning/linear/b",
+ "size": 1024,
+ "ref_norm": 0.022064837227258206,
+ "got_norm": 0.02201287283376861,
+ "finite": true,
+ "rel_l2": 0.004903469398417187
+ },
+ {
+ "param": "multimodality_forward/mesh_transformer/transformer/block_04/block_04_norm_conditioning/linear/w",
+ "size": 32768,
+ "ref_norm": 0.029310623408391093,
+ "got_norm": 0.02925572395502102,
+ "finite": true,
+ "rel_l2": 0.0047625610178021836
+ },
+ {
+ "param": "multimodality_forward/mesh_transformer/transformer/block_01/block_01_norm_conditioning/linear/w",
+ "size": 32768,
+ "ref_norm": 0.054780597316956656,
+ "got_norm": 0.05495864690955887,
+ "finite": true,
+ "rel_l2": 0.004350894474859587
+ },
+ {
+ "param": "multimodality_forward/mesh_encoder/split_input_matmul/w_spatial_feature=sin_lat,sin_lon,cos_lon",
+ "size": 1536,
+ "ref_norm": 0.3616832367875008,
+ "got_norm": 0.36118063180984145,
+ "finite": true,
+ "rel_l2": 0.004120612528605233
+ },
+ {
+ "param": "multimodality_forward/mesh_transformer/transformer/block_06/mha_proj_k/w",
+ "size": 262144,
+ "ref_norm": 0.08019620001056632,
+ "got_norm": 0.08018028202867657,
+ "finite": true,
+ "rel_l2": 0.0037842113410096667
+ }
+ ]
+ }
+ },
+ "official_mha vs official (floor)": {
+ "loss": {
+ "official": 1.2029576301574707,
+ "other": 1.2029576301574707,
+ "rel_diff": 0.0,
+ "max_per_var_rel_diff": 1.62002127755005e-05
+ },
+ "grads": {
+ "global_rel_l2": 0.0009516327318550146,
+ "global_cos": 0.9999995476794832,
+ "max_rel_l2": 0.004529760028810296,
+ "n_params": 376,
+ "n_compared": 335,
+ "n_zero_both": 6,
+ "n_zero_ref_only": 0,
+ "all_finite": true,
+ "worst": [
+ {
+ "param": "multimodality_forward/mesh_transformer/transformer/block_01/block_01_norm_conditioning/linear/w",
+ "size": 32768,
+ "ref_norm": 0.054780597316956656,
+ "got_norm": 0.05497245286861988,
+ "finite": true,
+ "rel_l2": 0.004529760028810296
+ },
+ {
+ "param": "multimodality_forward/mesh_transformer/transformer/block_04/mha_proj_k/w",
+ "size": 262144,
+ "ref_norm": 0.15881242287265113,
+ "got_norm": 0.15895169169953804,
+ "finite": true,
+ "rel_l2": 0.003606907061109118
+ },
+ {
+ "param": "multimodality_forward/mesh_transformer/transformer/block_01/block_01_norm_conditioning/linear/b",
+ "size": 1024,
+ "ref_norm": 0.03658263549895236,
+ "got_norm": 0.036674713312749546,
+ "finite": true,
+ "rel_l2": 0.0035763524088118616
+ },
+ {
+ "param": "multimodality_forward/mesh_transformer/transformer/block_04/block_04_norm_conditioning/linear/w",
+ "size": 32768,
+ "ref_norm": 0.029310623408391093,
+ "got_norm": 0.029280675810622207,
+ "finite": true,
+ "rel_l2": 0.003526815620035536
+ },
+ {
+ "param": "multimodality_forward/mesh_transformer/transformer/block_02/block_02_norm_conditioning/linear/w",
+ "size": 32768,
+ "ref_norm": 0.026617538682338283,
+ "got_norm": 0.026678329184390507,
+ "finite": true,
+ "rel_l2": 0.003514245233028868
+ }
+ ]
+ }
+ }
+ }
+ }
+ }
+}
\ No newline at end of file
diff --git a/docs/results/finetune/h200/grads03/full_grads_mini_crps_h200.log b/docs/results/finetune/h200/grads03/full_grads_mini_crps_h200.log
new file mode 100644
index 0000000..c45a69b
--- /dev/null
+++ b/docs/results/finetune/h200/grads03/full_grads_mini_crps_h200.log
@@ -0,0 +1,119 @@
+WeatherNextCyclones_Mini: differentiating 56.68M params (all); masked targets: []
+[highest]
+ official highest temp 102.06 GiB step 0.90s loss 1.202958
+ official_mha highest temp 123.49 GiB step 1.03s loss 1.202958
+ fwn highest temp 3.97 GiB step 0.68s loss 1.202958
+ fwn_xla highest temp 4.09 GiB step 1.25s loss 1.202958
+ fwn_strict highest temp 3.97 GiB step 1.56s loss 1.202958
+ fwn_noremat highest temp 39.10 GiB (compile only)
+ fwn vs official loss rel diff 0.00e+00 (per var max 1.74e-05) grads: global rel L2 4.74e-03 cos 0.9999887984 max per-param rel L2 2.37e-02 (335/376 params; zero in both 6, zero only in official 0) finite True
+ 2.37e-02 multimodality_forward/mesh_transformer/transformer/block_06/mha_proj_k/w
+ 1.81e-02 multimodality_forward/mesh_transformer/transformer/block_06/mha_proj_q/w
+ 1.70e-02 multimodality_forward/mesh_transformer/transformer/block_04/mha_proj_k/w
+ fwn_strict vs official loss rel diff 1.98e-07 (per var max 7.15e-06) grads: global rel L2 1.19e-03 cos 0.9999992979 max per-param rel L2 4.95e-03 (335/376 params; zero in both 6, zero only in official 0) finite True
+ 4.95e-03 multimodality_forward/mesh_encoder/split_input_matmul/w_spatial_feature=sin_lat,sin_lon,cos_lon
+ 4.82e-03 multimodality_forward/mesh_transformer/transformer/block_04/block_04_norm_conditioning/linear/w
+ 4.73e-03 multimodality_forward/mesh_transformer/transformer/block_01/block_01_norm_conditioning/linear/w
+ fwn_xla vs official loss rel diff 9.91e-08 (per var max 2.13e-05) grads: global rel L2 1.18e-03 cos 0.9999993146 max per-param rel L2 4.90e-03 (335/376 params; zero in both 6, zero only in official 0) finite True
+ 4.90e-03 multimodality_forward/mesh_transformer/transformer/block_04/block_04_norm_conditioning/linear/b
+ 4.76e-03 multimodality_forward/mesh_transformer/transformer/block_04/block_04_norm_conditioning/linear/w
+ 4.35e-03 multimodality_forward/mesh_transformer/transformer/block_01/block_01_norm_conditioning/linear/w
+ official_mha vs official (floor) loss rel diff 0.00e+00 (per var max 1.62e-05) grads: global rel L2 9.52e-04 cos 0.9999995477 max per-param rel L2 4.53e-03 (335/376 params; zero in both 6, zero only in official 0) finite True
+ 4.53e-03 multimodality_forward/mesh_transformer/transformer/block_01/block_01_norm_conditioning/linear/w
+ 3.61e-03 multimodality_forward/mesh_transformer/transformer/block_04/mha_proj_k/w
+ 3.58e-03 multimodality_forward/mesh_transformer/transformer/block_01/block_01_norm_conditioning/linear/b
+wrote /data/results/grads03/full_grads_mini_crps_h200.json
+WARNING:root:No normalization location found for cyclone_exists_gaussian_unit_mode
+WARNING:root:No normalization scale found for cyclone_exists_gaussian_unit_mode
+WARNING:absl:No coordinates found for dim noise_channels of array Size: 256B
+JitTracer(float32[2,1,32])
+Coordinates:
+ * batch (batch) object 16B MultiIndex
+ * _batch (batch) int64 16B 0 0
+ * sample (batch) int64 16B 0 1
+Dimensions without coordinates: channels, noise_channels., using index for param name. This may make model surgery more error prone if changing coordinates at fine-tuning time.
+/usr/local/lib/python3.12/site-packages/scipy/sparse/_index.py:216: SparseEfficiencyWarning: Changing the sparsity structure of a csr_matrix is expensive. lil and dok are more efficient.
+ self._set_arrayXarray(i, j, x)
+WARNING:absl:Only the first input of the static gather operation is being used, others are being ignored at runtime.
+WARNING:absl:Only the first input of the static gather operation is being used, others are being ignored at runtime.
+WARNING:absl:Only the first input of the sorted segment sum operation is being used, others are being ignored at runtime.
+WARNING:absl:Only the first input of the static gather operation is being used, others are being ignored at runtime.
+WARNING:absl:Only the first input of the static gather operation is being used, others are being ignored at runtime.
+WARNING:absl:Only the first input of the static gather operation is being used, others are being ignored at runtime.
+WARNING:absl:Only the first input of the sorted segment sum operation is being used, others are being ignored at runtime.
+/usr/local/lib/python3.12/site-packages/weathernext/weathernext2/architecture.py:286: FutureWarning: Deleting a single level of a MultiIndex is deprecated. Previously, this deleted all levels of a MultiIndex. Please also drop the following variables: {'batch'} to avoid an error in the future.
+ predictions[var] = xarray_jax.apply_ufunc(activation_fn, predictions[var])
+WARNING:root:No normalization location found for cyclone_exists_gaussian_unit_mode
+WARNING:root:No normalization scale found for cyclone_exists_gaussian_unit_mode
+WARNING:absl:No coordinates found for dim noise_channels of array Size: 256B
+JitTracer(float32[2,1,32])
+Coordinates:
+ * batch (batch) object 16B MultiIndex
+ * _batch (batch) int64 16B 0 0
+ * sample (batch) int64 16B 0 1
+Dimensions without coordinates: channels, noise_channels., using index for param name. This may make model surgery more error prone if changing coordinates at fine-tuning time.
+/usr/local/lib/python3.12/site-packages/scipy/sparse/_index.py:216: SparseEfficiencyWarning: Changing the sparsity structure of a csr_matrix is expensive. lil and dok are more efficient.
+ self._set_arrayXarray(i, j, x)
+WARNING:absl:Only the first input of the static gather operation is being used, others are being ignored at runtime.
+WARNING:absl:Only the first input of the static gather operation is being used, others are being ignored at runtime.
+WARNING:absl:Only the first input of the sorted segment sum operation is being used, others are being ignored at runtime.
+WARNING:absl:Only the first input of the static gather operation is being used, others are being ignored at runtime.
+WARNING:absl:Only the first input of the static gather operation is being used, others are being ignored at runtime.
+WARNING:absl:Only the first input of the static gather operation is being used, others are being ignored at runtime.
+WARNING:absl:Only the first input of the sorted segment sum operation is being used, others are being ignored at runtime.
+/usr/local/lib/python3.12/site-packages/weathernext/weathernext2/architecture.py:286: FutureWarning: Deleting a single level of a MultiIndex is deprecated. Previously, this deleted all levels of a MultiIndex. Please also drop the following variables: {'batch'} to avoid an error in the future.
+ predictions[var] = xarray_jax.apply_ufunc(activation_fn, predictions[var])
+WARNING:root:No normalization location found for cyclone_exists_gaussian_unit_mode
+WARNING:root:No normalization scale found for cyclone_exists_gaussian_unit_mode
+WARNING:absl:No coordinates found for dim noise_channels of array Size: 256B
+JitTracer(float32[2,1,32])
+Coordinates:
+ * batch (batch) object 16B MultiIndex
+ * _batch (batch) int64 16B 0 0
+ * sample (batch) int64 16B 0 1
+Dimensions without coordinates: channels, noise_channels., using index for param name. This may make model surgery more error prone if changing coordinates at fine-tuning time.
+/usr/local/lib/python3.12/site-packages/scipy/sparse/_index.py:216: SparseEfficiencyWarning: Changing the sparsity structure of a csr_matrix is expensive. lil and dok are more efficient.
+ self._set_arrayXarray(i, j, x)
+/usr/local/lib/python3.12/site-packages/weathernext/weathernext2/architecture.py:286: FutureWarning: Deleting a single level of a MultiIndex is deprecated. Previously, this deleted all levels of a MultiIndex. Please also drop the following variables: {'batch'} to avoid an error in the future.
+ predictions[var] = xarray_jax.apply_ufunc(activation_fn, predictions[var])
+WARNING:root:No normalization location found for cyclone_exists_gaussian_unit_mode
+WARNING:root:No normalization scale found for cyclone_exists_gaussian_unit_mode
+WARNING:absl:No coordinates found for dim noise_channels of array Size: 256B
+JitTracer(float32[2,1,32])
+Coordinates:
+ * batch (batch) object 16B MultiIndex
+ * _batch (batch) int64 16B 0 0
+ * sample (batch) int64 16B 0 1
+Dimensions without coordinates: channels, noise_channels., using index for param name. This may make model surgery more error prone if changing coordinates at fine-tuning time.
+/usr/local/lib/python3.12/site-packages/scipy/sparse/_index.py:216: SparseEfficiencyWarning: Changing the sparsity structure of a csr_matrix is expensive. lil and dok are more efficient.
+ self._set_arrayXarray(i, j, x)
+/usr/local/lib/python3.12/site-packages/weathernext/weathernext2/architecture.py:286: FutureWarning: Deleting a single level of a MultiIndex is deprecated. Previously, this deleted all levels of a MultiIndex. Please also drop the following variables: {'batch'} to avoid an error in the future.
+ predictions[var] = xarray_jax.apply_ufunc(activation_fn, predictions[var])
+WARNING:root:No normalization location found for cyclone_exists_gaussian_unit_mode
+WARNING:root:No normalization scale found for cyclone_exists_gaussian_unit_mode
+WARNING:absl:No coordinates found for dim noise_channels of array Size: 256B
+JitTracer(float32[2,1,32])
+Coordinates:
+ * batch (batch) object 16B MultiIndex
+ * _batch (batch) int64 16B 0 0
+ * sample (batch) int64 16B 0 1
+Dimensions without coordinates: channels, noise_channels., using index for param name. This may make model surgery more error prone if changing coordinates at fine-tuning time.
+/usr/local/lib/python3.12/site-packages/scipy/sparse/_index.py:216: SparseEfficiencyWarning: Changing the sparsity structure of a csr_matrix is expensive. lil and dok are more efficient.
+ self._set_arrayXarray(i, j, x)
+/usr/local/lib/python3.12/site-packages/weathernext/weathernext2/architecture.py:286: FutureWarning: Deleting a single level of a MultiIndex is deprecated. Previously, this deleted all levels of a MultiIndex. Please also drop the following variables: {'batch'} to avoid an error in the future.
+ predictions[var] = xarray_jax.apply_ufunc(activation_fn, predictions[var])
+WARNING:root:No normalization location found for cyclone_exists_gaussian_unit_mode
+WARNING:root:No normalization scale found for cyclone_exists_gaussian_unit_mode
+WARNING:absl:No coordinates found for dim noise_channels of array Size: 256B
+JitTracer(float32[2,1,32])
+Coordinates:
+ * batch (batch) object 16B MultiIndex
+ * _batch (batch) int64 16B 0 0
+ * sample (batch) int64 16B 0 1
+Dimensions without coordinates: channels, noise_channels., using index for param name. This may make model surgery more error prone if changing coordinates at fine-tuning time.
+/usr/local/lib/python3.12/site-packages/scipy/sparse/_index.py:216: SparseEfficiencyWarning: Changing the sparsity structure of a csr_matrix is expensive. lil and dok are more efficient.
+ self._set_arrayXarray(i, j, x)
+/usr/local/lib/python3.12/site-packages/weathernext/weathernext2/architecture.py:286: FutureWarning: Deleting a single level of a MultiIndex is deprecated. Previously, this deleted all levels of a MultiIndex. Please also drop the following variables: {'batch'} to avoid an error in the future.
+ predictions[var] = xarray_jax.apply_ufunc(activation_fn, predictions[var])
+
+exit 0 wall 317s
diff --git a/docs/results/finetune/h200/grads03/full_grads_mini_linear_h200.json b/docs/results/finetune/h200/grads03/full_grads_mini_linear_h200.json
new file mode 100644
index 0000000..600bb3c
--- /dev/null
+++ b/docs/results/finetune/h200/grads03/full_grads_mini_linear_h200.json
@@ -0,0 +1,290 @@
+{
+ "model": "WeatherNextCyclones_Mini",
+ "device": "NVIDIA H200",
+ "jax": "0.11.2",
+ "faster_weathernext": "0.1.1",
+ "blocks": {
+ "grid_chunk": 8192,
+ "edge_chunk": 16384
+ },
+ "masked": [],
+ "modules": "all",
+ "loss": "linear",
+ "head_params": 56684133,
+ "precisions": {
+ "highest": {
+ "official": {
+ "temp_gib": 51.412283167243004,
+ "compile_s": 72.90220785140991,
+ "step_s": 0.4801790714263916
+ },
+ "official_mha": {
+ "temp_gib": 61.79655335843563,
+ "compile_s": 39.80091977119446,
+ "step_s": 0.5431849956512451
+ },
+ "fwn": {
+ "temp_gib": 3.299081988632679,
+ "compile_s": 64.95161771774292,
+ "step_s": 0.4908900260925293
+ },
+ "fwn_xla": {
+ "temp_gib": 3.397009789943695,
+ "compile_s": 42.87857627868652,
+ "step_s": 0.7673413753509521
+ },
+ "fwn_strict": {
+ "temp_gib": 3.299081988632679,
+ "compile_s": 31.026937246322632,
+ "step_s": 0.8436541557312012
+ },
+ "fwn_noremat": {
+ "temp_gib": 19.658946715295315,
+ "compile_s": 37.25039529800415
+ },
+ "reference": "official",
+ "checks": {
+ "fwn vs official": {
+ "loss": {
+ "official": 506.6860656738281,
+ "other": 506.68609619140625,
+ "rel_diff": 6.022975604118005e-08,
+ "max_per_var_rel_diff": 9.50762886513566e-05
+ },
+ "grads": {
+ "global_rel_l2": 0.00013561236967894033,
+ "global_cos": 0.9999999908334908,
+ "max_rel_l2": 0.0005556341620280397,
+ "n_params": 376,
+ "n_compared": 315,
+ "n_zero_both": 6,
+ "n_zero_ref_only": 0,
+ "all_finite": true,
+ "worst": [
+ {
+ "param": "multimodality_forward/mesh_encoder/split_input_matmul/w_spatial_feature=sin_lat,sin_lon,cos_lon",
+ "size": 1536,
+ "ref_norm": 15.954467520832804,
+ "got_norm": 15.961123680523185,
+ "finite": true,
+ "rel_l2": 0.0005556341620280397
+ },
+ {
+ "param": "multimodality_forward/mesh_transformer/transformer/block_04/mha_proj_k/w",
+ "size": 262144,
+ "ref_norm": 7.4186753365914475,
+ "got_norm": 7.420224717492963,
+ "finite": true,
+ "rel_l2": 0.0005231898832359417
+ },
+ {
+ "param": "multimodality_forward/mesh_transformer/transformer/block_04/block_04_norm_conditioning/linear/w",
+ "size": 32768,
+ "ref_norm": 0.6272398645339914,
+ "got_norm": 0.6271567224001288,
+ "finite": true,
+ "rel_l2": 0.0005090704423694277
+ },
+ {
+ "param": "multimodality_forward/mesh_transformer/transformer/block_04/block_04_norm_conditioning/linear/b",
+ "size": 1024,
+ "ref_norm": 0.7268014987796344,
+ "got_norm": 0.7267051595257081,
+ "finite": true,
+ "rel_l2": 0.0005090700209365709
+ },
+ {
+ "param": "multimodality_forward/grid_to_mesh_gnn/edge_encoder/shared_dense/mlp/linear_0/b",
+ "size": 32,
+ "ref_norm": 5.130661840892781,
+ "got_norm": 5.132819695062869,
+ "finite": true,
+ "rel_l2": 0.0004911391326115861
+ }
+ ]
+ }
+ },
+ "fwn_strict vs official": {
+ "loss": {
+ "official": 506.6860656738281,
+ "other": 506.68609619140625,
+ "rel_diff": 6.022975604118005e-08,
+ "max_per_var_rel_diff": 0.00012948800400019655
+ },
+ "grads": {
+ "global_rel_l2": 9.006578893174544e-05,
+ "global_cos": 0.9999999959443298,
+ "max_rel_l2": 0.000499279010213555,
+ "n_params": 376,
+ "n_compared": 315,
+ "n_zero_both": 6,
+ "n_zero_ref_only": 0,
+ "all_finite": true,
+ "worst": [
+ {
+ "param": "multimodality_forward/mesh_transformer/transformer/block_04/block_04_norm_conditioning/linear/b",
+ "size": 1024,
+ "ref_norm": 0.7268014987796344,
+ "got_norm": 0.7267032752261265,
+ "finite": true,
+ "rel_l2": 0.000499279010213555
+ },
+ {
+ "param": "multimodality_forward/mesh_transformer/transformer/block_04/block_04_norm_conditioning/linear/w",
+ "size": 32768,
+ "ref_norm": 0.6272398645339914,
+ "got_norm": 0.6271550965850112,
+ "finite": true,
+ "rel_l2": 0.0004992778603369683
+ },
+ {
+ "param": "multimodality_forward/mesh_encoder/split_input_matmul/w_spatial_feature=sin_lat,sin_lon,cos_lon",
+ "size": 1536,
+ "ref_norm": 15.954467520832804,
+ "got_norm": 15.959611876156147,
+ "finite": true,
+ "rel_l2": 0.0004196505980078108
+ },
+ {
+ "param": "multimodality_forward/mesh_transformer/transformer/block_04/mha_proj_k/w",
+ "size": 262144,
+ "ref_norm": 7.4186753365914475,
+ "got_norm": 7.420028762553728,
+ "finite": true,
+ "rel_l2": 0.00035633045711653356
+ },
+ {
+ "param": "multimodality_forward/mesh_transformer/transformer/block_04/mha_proj_q/w",
+ "size": 262144,
+ "ref_norm": 7.82067851789785,
+ "got_norm": 7.8218450754639575,
+ "finite": true,
+ "rel_l2": 0.0003357190227201953
+ }
+ ]
+ }
+ },
+ "fwn_xla vs official": {
+ "loss": {
+ "official": 506.6860656738281,
+ "other": 506.68609619140625,
+ "rel_diff": 6.022975604118005e-08,
+ "max_per_var_rel_diff": 0.0001860782576249697
+ },
+ "grads": {
+ "global_rel_l2": 7.603367593274188e-05,
+ "global_cos": 0.9999999971998349,
+ "max_rel_l2": 0.00027426703264950373,
+ "n_params": 376,
+ "n_compared": 315,
+ "n_zero_both": 6,
+ "n_zero_ref_only": 0,
+ "all_finite": true,
+ "worst": [
+ {
+ "param": "multimodality_forward/mesh_encoder/split_input_matmul/w_spatial_feature=sin_lat,sin_lon,cos_lon",
+ "size": 1536,
+ "ref_norm": 15.954467520832804,
+ "got_norm": 15.957882086805508,
+ "finite": true,
+ "rel_l2": 0.00027426703264950373
+ },
+ {
+ "param": "multimodality_forward/mesh_transformer/transformer/block_06/mha_proj_k/w",
+ "size": 262144,
+ "ref_norm": 2.7121996407746347,
+ "got_norm": 2.7121559746479087,
+ "finite": true,
+ "rel_l2": 0.00025277962627656063
+ },
+ {
+ "param": "multimodality_forward/mesh_transformer/transformer/block_04/mha_proj_k/w",
+ "size": 262144,
+ "ref_norm": 7.4186753365914475,
+ "got_norm": 7.419211031589528,
+ "finite": true,
+ "rel_l2": 0.0002313781395283307
+ },
+ {
+ "param": "multimodality_forward/mesh_encoder/shared_dense/mlp/linear_1/w",
+ "size": 262144,
+ "ref_norm": 32.62593666203144,
+ "got_norm": 32.629571295138255,
+ "finite": true,
+ "rel_l2": 0.0001933310536733495
+ },
+ {
+ "param": "multimodality_forward/mesh_transformer/transformer/block_04/mha_proj_q/w",
+ "size": 262144,
+ "ref_norm": 7.82067851789785,
+ "got_norm": 7.821128374594056,
+ "finite": true,
+ "rel_l2": 0.00019046206223078755
+ }
+ ]
+ }
+ },
+ "official_mha vs official (floor)": {
+ "loss": {
+ "official": 506.6860656738281,
+ "other": 506.6860656738281,
+ "rel_diff": 0.0,
+ "max_per_var_rel_diff": 8.87732604571823e-05
+ },
+ "grads": {
+ "global_rel_l2": 8.128736974067304e-05,
+ "global_cos": 0.9999999966982622,
+ "max_rel_l2": 0.0002867339859766047,
+ "n_params": 376,
+ "n_compared": 315,
+ "n_zero_both": 6,
+ "n_zero_ref_only": 0,
+ "all_finite": true,
+ "worst": [
+ {
+ "param": "multimodality_forward/mesh_transformer/transformer/block_04/block_04_norm_conditioning/linear/b",
+ "size": 1024,
+ "ref_norm": 0.7268014987796344,
+ "got_norm": 0.7267656387203192,
+ "finite": true,
+ "rel_l2": 0.0002867339859766047
+ },
+ {
+ "param": "multimodality_forward/mesh_transformer/transformer/block_04/block_04_norm_conditioning/linear/w",
+ "size": 32768,
+ "ref_norm": 0.6272398645339914,
+ "got_norm": 0.6272089176261173,
+ "finite": true,
+ "rel_l2": 0.00028673213888048794
+ },
+ {
+ "param": "multimodality_forward/mesh_encoder/split_input_matmul/w_spatial_feature=sin_lat,sin_lon,cos_lon",
+ "size": 1536,
+ "ref_norm": 15.954467520832804,
+ "got_norm": 15.957502491142428,
+ "finite": true,
+ "rel_l2": 0.0002836680292599463
+ },
+ {
+ "param": "multimodality_forward/grid_to_mesh_gnn/edge_encoder/shared_dense/mlp/linear_0/b",
+ "size": 32,
+ "ref_norm": 5.130661840892781,
+ "got_norm": 5.131607474124301,
+ "finite": true,
+ "rel_l2": 0.0002580676854218195
+ },
+ {
+ "param": "multimodality_forward/mesh_transformer/transformer/block_04/mha_proj_k/w",
+ "size": 262144,
+ "ref_norm": 7.4186753365914475,
+ "got_norm": 7.419418815102285,
+ "finite": true,
+ "rel_l2": 0.00024943632571534447
+ }
+ ]
+ }
+ }
+ }
+ }
+ }
+}
\ No newline at end of file
diff --git a/docs/results/finetune/h200/grads03/full_grads_mini_linear_h200.log b/docs/results/finetune/h200/grads03/full_grads_mini_linear_h200.log
new file mode 100644
index 0000000..d3d6e50
--- /dev/null
+++ b/docs/results/finetune/h200/grads03/full_grads_mini_linear_h200.log
@@ -0,0 +1,93 @@
+WeatherNextCyclones_Mini: differentiating 56.68M params (all); masked targets: []
+[highest]
+ official highest temp 51.41 GiB step 0.48s loss 506.686066
+ official_mha highest temp 61.80 GiB step 0.54s loss 506.686066
+ fwn highest temp 3.30 GiB step 0.49s loss 506.686096
+ fwn_xla highest temp 3.40 GiB step 0.77s loss 506.686096
+ fwn_strict highest temp 3.30 GiB step 0.84s loss 506.686096
+ fwn_noremat highest temp 19.66 GiB (compile only)
+ fwn vs official loss rel diff 6.02e-08 (per var max 9.51e-05) grads: global rel L2 1.36e-04 cos 0.9999999908 max per-param rel L2 5.56e-04 (315/376 params; zero in both 6, zero only in official 0) finite True
+ 5.56e-04 multimodality_forward/mesh_encoder/split_input_matmul/w_spatial_feature=sin_lat,sin_lon,cos_lon
+ 5.23e-04 multimodality_forward/mesh_transformer/transformer/block_04/mha_proj_k/w
+ 5.09e-04 multimodality_forward/mesh_transformer/transformer/block_04/block_04_norm_conditioning/linear/w
+ fwn_strict vs official loss rel diff 6.02e-08 (per var max 1.29e-04) grads: global rel L2 9.01e-05 cos 0.9999999959 max per-param rel L2 4.99e-04 (315/376 params; zero in both 6, zero only in official 0) finite True
+ 4.99e-04 multimodality_forward/mesh_transformer/transformer/block_04/block_04_norm_conditioning/linear/b
+ 4.99e-04 multimodality_forward/mesh_transformer/transformer/block_04/block_04_norm_conditioning/linear/w
+ 4.20e-04 multimodality_forward/mesh_encoder/split_input_matmul/w_spatial_feature=sin_lat,sin_lon,cos_lon
+ fwn_xla vs official loss rel diff 6.02e-08 (per var max 1.86e-04) grads: global rel L2 7.60e-05 cos 0.9999999972 max per-param rel L2 2.74e-04 (315/376 params; zero in both 6, zero only in official 0) finite True
+ 2.74e-04 multimodality_forward/mesh_encoder/split_input_matmul/w_spatial_feature=sin_lat,sin_lon,cos_lon
+ 2.53e-04 multimodality_forward/mesh_transformer/transformer/block_06/mha_proj_k/w
+ 2.31e-04 multimodality_forward/mesh_transformer/transformer/block_04/mha_proj_k/w
+ official_mha vs official (floor) loss rel diff 0.00e+00 (per var max 8.88e-05) grads: global rel L2 8.13e-05 cos 0.9999999967 max per-param rel L2 2.87e-04 (315/376 params; zero in both 6, zero only in official 0) finite True
+ 2.87e-04 multimodality_forward/mesh_transformer/transformer/block_04/block_04_norm_conditioning/linear/b
+ 2.87e-04 multimodality_forward/mesh_transformer/transformer/block_04/block_04_norm_conditioning/linear/w
+ 2.84e-04 multimodality_forward/mesh_encoder/split_input_matmul/w_spatial_feature=sin_lat,sin_lon,cos_lon
+wrote /data/results/grads03/full_grads_mini_linear_h200.json
+WARNING:absl:Skipping gradient checkpointing for sequence length of 1
+WARNING:absl:No coordinates found for dim noise_channels of array Size: 128B
+JitTracer(float32[1,1,32])
+Dimensions without coordinates: batch, channels, noise_channels., using index for param name. This may make model surgery more error prone if changing coordinates at fine-tuning time.
+/usr/local/lib/python3.12/site-packages/scipy/sparse/_index.py:216: SparseEfficiencyWarning: Changing the sparsity structure of a csr_matrix is expensive. lil and dok are more efficient.
+ self._set_arrayXarray(i, j, x)
+WARNING:absl:Only the first input of the static gather operation is being used, others are being ignored at runtime.
+WARNING:absl:Only the first input of the static gather operation is being used, others are being ignored at runtime.
+WARNING:absl:Only the first input of the sorted segment sum operation is being used, others are being ignored at runtime.
+WARNING:absl:Only the first input of the static gather operation is being used, others are being ignored at runtime.
+WARNING:absl:Only the first input of the static gather operation is being used, others are being ignored at runtime.
+WARNING:absl:Only the first input of the static gather operation is being used, others are being ignored at runtime.
+WARNING:absl:Only the first input of the sorted segment sum operation is being used, others are being ignored at runtime.
+WARNING:root:No normalization scale found for cyclone_exists_gaussian_unit_mode
+WARNING:root:No normalization location found for cyclone_exists_gaussian_unit_mode
+WARNING:absl:Only the first input of the sorted segment sum operation is being used, others are being ignored at runtime.
+WARNING:absl:Only the first input of the sorted segment sum operation is being used, others are being ignored at runtime.
+WARNING:absl:Skipping gradient checkpointing for sequence length of 1
+WARNING:absl:No coordinates found for dim noise_channels of array Size: 128B
+JitTracer(float32[1,1,32])
+Dimensions without coordinates: batch, channels, noise_channels., using index for param name. This may make model surgery more error prone if changing coordinates at fine-tuning time.
+/usr/local/lib/python3.12/site-packages/scipy/sparse/_index.py:216: SparseEfficiencyWarning: Changing the sparsity structure of a csr_matrix is expensive. lil and dok are more efficient.
+ self._set_arrayXarray(i, j, x)
+WARNING:absl:Only the first input of the static gather operation is being used, others are being ignored at runtime.
+WARNING:absl:Only the first input of the static gather operation is being used, others are being ignored at runtime.
+WARNING:absl:Only the first input of the sorted segment sum operation is being used, others are being ignored at runtime.
+WARNING:absl:Only the first input of the static gather operation is being used, others are being ignored at runtime.
+WARNING:absl:Only the first input of the static gather operation is being used, others are being ignored at runtime.
+WARNING:absl:Only the first input of the static gather operation is being used, others are being ignored at runtime.
+WARNING:absl:Only the first input of the sorted segment sum operation is being used, others are being ignored at runtime.
+WARNING:root:No normalization scale found for cyclone_exists_gaussian_unit_mode
+WARNING:root:No normalization location found for cyclone_exists_gaussian_unit_mode
+WARNING:absl:Only the first input of the sorted segment sum operation is being used, others are being ignored at runtime.
+WARNING:absl:Only the first input of the sorted segment sum operation is being used, others are being ignored at runtime.
+WARNING:absl:Skipping gradient checkpointing for sequence length of 1
+WARNING:absl:No coordinates found for dim noise_channels of array Size: 128B
+JitTracer(float32[1,1,32])
+Dimensions without coordinates: batch, channels, noise_channels., using index for param name. This may make model surgery more error prone if changing coordinates at fine-tuning time.
+/usr/local/lib/python3.12/site-packages/scipy/sparse/_index.py:216: SparseEfficiencyWarning: Changing the sparsity structure of a csr_matrix is expensive. lil and dok are more efficient.
+ self._set_arrayXarray(i, j, x)
+WARNING:root:No normalization scale found for cyclone_exists_gaussian_unit_mode
+WARNING:root:No normalization location found for cyclone_exists_gaussian_unit_mode
+WARNING:absl:Skipping gradient checkpointing for sequence length of 1
+WARNING:absl:No coordinates found for dim noise_channels of array Size: 128B
+JitTracer(float32[1,1,32])
+Dimensions without coordinates: batch, channels, noise_channels., using index for param name. This may make model surgery more error prone if changing coordinates at fine-tuning time.
+/usr/local/lib/python3.12/site-packages/scipy/sparse/_index.py:216: SparseEfficiencyWarning: Changing the sparsity structure of a csr_matrix is expensive. lil and dok are more efficient.
+ self._set_arrayXarray(i, j, x)
+WARNING:root:No normalization scale found for cyclone_exists_gaussian_unit_mode
+WARNING:root:No normalization location found for cyclone_exists_gaussian_unit_mode
+WARNING:absl:Skipping gradient checkpointing for sequence length of 1
+WARNING:absl:No coordinates found for dim noise_channels of array Size: 128B
+JitTracer(float32[1,1,32])
+Dimensions without coordinates: batch, channels, noise_channels., using index for param name. This may make model surgery more error prone if changing coordinates at fine-tuning time.
+/usr/local/lib/python3.12/site-packages/scipy/sparse/_index.py:216: SparseEfficiencyWarning: Changing the sparsity structure of a csr_matrix is expensive. lil and dok are more efficient.
+ self._set_arrayXarray(i, j, x)
+WARNING:root:No normalization scale found for cyclone_exists_gaussian_unit_mode
+WARNING:root:No normalization location found for cyclone_exists_gaussian_unit_mode
+WARNING:absl:Skipping gradient checkpointing for sequence length of 1
+WARNING:absl:No coordinates found for dim noise_channels of array Size: 128B
+JitTracer(float32[1,1,32])
+Dimensions without coordinates: batch, channels, noise_channels., using index for param name. This may make model surgery more error prone if changing coordinates at fine-tuning time.
+/usr/local/lib/python3.12/site-packages/scipy/sparse/_index.py:216: SparseEfficiencyWarning: Changing the sparsity structure of a csr_matrix is expensive. lil and dok are more efficient.
+ self._set_arrayXarray(i, j, x)
+WARNING:root:No normalization scale found for cyclone_exists_gaussian_unit_mode
+WARNING:root:No normalization location found for cyclone_exists_gaussian_unit_mode
+
+exit 0 wall 309s
diff --git a/docs/results/finetune/head_grads_mini_rtx5090.json b/docs/results/finetune/head_grads_mini_rtx5090.json
new file mode 100644
index 0000000..d194751
--- /dev/null
+++ b/docs/results/finetune/head_grads_mini_rtx5090.json
@@ -0,0 +1,365 @@
+{
+ "model": "WeatherNextCyclones_Mini",
+ "device": "NVIDIA GeForce RTX 5090",
+ "jax": "0.11.2",
+ "faster_weathernext": "0.1.1",
+ "blocks": {
+ "grid_chunk": 8192,
+ "edge_chunk": 16384
+ },
+ "masked": [],
+ "head_params": 2535781,
+ "precisions": {
+ "highest": {
+ "official": {
+ "temp_gib": 7.358244687318802,
+ "compile_s": 47.37564778327942,
+ "step_s": 0.4156680107116699
+ },
+ "official_mha": {
+ "temp_gib": 7.644339352846146,
+ "compile_s": 19.233081579208374,
+ "step_s": 0.33963489532470703
+ },
+ "fwn": {
+ "temp_gib": 1.4815475270152092,
+ "compile_s": 39.05863094329834,
+ "step_s": 0.2230541706085205
+ },
+ "fwn_noremat": {
+ "temp_gib": 9.911023788154125,
+ "compile_s": 14.838613748550415
+ },
+ "fwn_strict": {
+ "temp_gib": 1.4815475270152092,
+ "compile_s": 10.349909782409668,
+ "step_s": 1.4333524703979492
+ },
+ "checks": {
+ "fwn vs official": {
+ "loss": {
+ "official": 1.2029576301574707,
+ "other": 1.2029577493667603,
+ "rel_diff": 9.909683147791032e-08,
+ "max_per_var_rel_diff": 1.6190872747221507e-05
+ },
+ "grads": {
+ "global_rel_l2": 0.0010030226145246349,
+ "global_cos": 0.9999995042909823,
+ "max_rel_l2": 0.0022973047709093984,
+ "n_params": 86,
+ "n_compared": 58,
+ "n_zero_both": 6,
+ "n_zero_ref_only": 0,
+ "all_finite": true,
+ "worst": [
+ {
+ "param": "multimodality_forward/grid_decoder/split_output_linear/w_specific_humidity_level=50,100,150,200,250,300,400,500,600,700,850,925,1000_time=21600",
+ "size": 6656,
+ "ref_norm": 0.06548207335001195,
+ "got_norm": 0.06550444935382058,
+ "finite": true,
+ "rel_l2": 0.0022973047709093984
+ },
+ {
+ "param": "multimodality_forward/mesh_to_grid_gnn/deep_gnn/processor_edges_0_sender_mesh_to_points_nodes/w",
+ "size": 262144,
+ "ref_norm": 1.2613371802792945,
+ "got_norm": 1.2613660209182214,
+ "finite": true,
+ "rel_l2": 0.0013841510865577958
+ },
+ {
+ "param": "multimodality_forward/grid_decoder/split_output_linear/w_vertical_velocity_level=50,100,150,200,250,300,400,500,600,700,850,925,1000_time=21600",
+ "size": 6656,
+ "ref_norm": 0.1465138721253444,
+ "got_norm": 0.14647541911521744,
+ "finite": true,
+ "rel_l2": 0.0013532309792546483
+ },
+ {
+ "param": "multimodality_forward/grid_decoder/split_output_linear/w_v_component_of_wind_level=50,100,150,200,250,300,400,500,600,700,850,925,1000_time=21600",
+ "size": 6656,
+ "ref_norm": 0.07094115823576923,
+ "got_norm": 0.07093730177106626,
+ "finite": true,
+ "rel_l2": 0.0012752194701343797
+ },
+ {
+ "param": "multimodality_forward/mesh_to_grid_gnn/deep_gnn/processor_edges_0_mesh_to_points_nodes/mlp/linear_1/w",
+ "size": 262144,
+ "ref_norm": 1.5409242901568447,
+ "got_norm": 1.540869624167333,
+ "finite": true,
+ "rel_l2": 0.0011536157639586834
+ }
+ ]
+ }
+ },
+ "fwn strict_fp32 vs official": {
+ "loss": {
+ "official": 1.2029576301574707,
+ "other": 1.2029576301574707,
+ "rel_diff": 0.0,
+ "max_per_var_rel_diff": 1.1745324452003711e-05
+ },
+ "grads": {
+ "global_rel_l2": 0.000884629527042391,
+ "global_cos": 0.9999996163129573,
+ "max_rel_l2": 0.0014099353925182002,
+ "n_params": 86,
+ "n_compared": 58,
+ "n_zero_both": 6,
+ "n_zero_ref_only": 0,
+ "all_finite": true,
+ "worst": [
+ {
+ "param": "multimodality_forward/grid_decoder/split_output_linear/w_vertical_velocity_level=50,100,150,200,250,300,400,500,600,700,850,925,1000_time=21600",
+ "size": 6656,
+ "ref_norm": 0.1465138721253444,
+ "got_norm": 0.146483859000436,
+ "finite": true,
+ "rel_l2": 0.0014099353925182002
+ },
+ {
+ "param": "multimodality_forward/mesh_to_grid_gnn/deep_gnn/processor_edges_0_sender_mesh_to_points_nodes/w",
+ "size": 262144,
+ "ref_norm": 1.2613371802792945,
+ "got_norm": 1.261372757727361,
+ "finite": true,
+ "rel_l2": 0.0013158979378149725
+ },
+ {
+ "param": "multimodality_forward/grid_decoder/split_output_linear/w_u_component_of_wind_level=50,100,150,200,250,300,400,500,600,700,850,925,1000_time=21600",
+ "size": 6656,
+ "ref_norm": 0.06688905072453233,
+ "got_norm": 0.06687706374181301,
+ "finite": true,
+ "rel_l2": 0.0012721444743921035
+ },
+ {
+ "param": "multimodality_forward/grid_decoder/split_output_linear/w_specific_humidity_level=50,100,150,200,250,300,400,500,600,700,850,925,1000_time=21600",
+ "size": 6656,
+ "ref_norm": 0.06548207335001195,
+ "got_norm": 0.06548068947587456,
+ "finite": true,
+ "rel_l2": 0.0012458667873252699
+ },
+ {
+ "param": "multimodality_forward/grid_decoder/split_output_linear/w_v_component_of_wind_level=50,100,150,200,250,300,400,500,600,700,850,925,1000_time=21600",
+ "size": 6656,
+ "ref_norm": 0.07094115823576923,
+ "got_norm": 0.07093756859295751,
+ "finite": true,
+ "rel_l2": 0.001199071775250381
+ }
+ ]
+ }
+ },
+ "official_mha vs official (floor)": {
+ "loss": {
+ "official": 1.2029576301574707,
+ "other": 1.2029576301574707,
+ "rel_diff": 0.0,
+ "max_per_var_rel_diff": 1.2596844416374613e-05
+ },
+ "grads": {
+ "global_rel_l2": 0.0005659516072411567,
+ "global_cos": 0.9999998399315811,
+ "max_rel_l2": 0.002360288687931764,
+ "n_params": 86,
+ "n_compared": 58,
+ "n_zero_both": 6,
+ "n_zero_ref_only": 0,
+ "all_finite": true,
+ "worst": [
+ {
+ "param": "multimodality_forward/grid_decoder/split_output_linear/w_specific_humidity_level=50,100,150,200,250,300,400,500,600,700,850,925,1000_time=21600",
+ "size": 6656,
+ "ref_norm": 0.06548207335001195,
+ "got_norm": 0.06549532331003156,
+ "finite": true,
+ "rel_l2": 0.002360288687931764
+ },
+ {
+ "param": "multimodality_forward/grid_decoder/split_output_linear/w_10m_v_component_of_wind_time=21600",
+ "size": 512,
+ "ref_norm": 0.020083909136383356,
+ "got_norm": 0.020085930860498306,
+ "finite": true,
+ "rel_l2": 0.0013085319396149494
+ },
+ {
+ "param": "multimodality_forward/grid_decoder/split_output_linear/w_v_component_of_wind_level=50,100,150,200,250,300,400,500,600,700,850,925,1000_time=21600",
+ "size": 6656,
+ "ref_norm": 0.07094115823576923,
+ "got_norm": 0.07093550391452534,
+ "finite": true,
+ "rel_l2": 0.0012830016184828138
+ },
+ {
+ "param": "multimodality_forward/grid_decoder/split_output_linear/w_vertical_velocity_level=50,100,150,200,250,300,400,500,600,700,850,925,1000_time=21600",
+ "size": 6656,
+ "ref_norm": 0.1465138721253444,
+ "got_norm": 0.14648285970806105,
+ "finite": true,
+ "rel_l2": 0.0009228097704790526
+ },
+ {
+ "param": "multimodality_forward/mesh_to_grid_gnn/deep_gnn/processor_edges_0_mesh_to_points_nodes/mlp/linear_1/w",
+ "size": 262144,
+ "ref_norm": 1.5409242901568447,
+ "got_norm": 1.5408702641874212,
+ "finite": true,
+ "rel_l2": 0.0008855050422957197
+ }
+ ]
+ }
+ }
+ }
+ },
+ "default": {
+ "official": {
+ "temp_gib": 5.2644632160663605,
+ "compile_s": 26.895958185195923,
+ "step_s": 0.29421567916870117
+ },
+ "official_mha": {
+ "temp_gib": 7.644321233034134,
+ "compile_s": 14.141706943511963,
+ "step_s": 0.2723855972290039
+ },
+ "fwn": {
+ "temp_gib": 1.4815429970622063,
+ "compile_s": 22.251452445983887,
+ "step_s": 0.15200233459472656
+ },
+ "fwn_noremat": {
+ "temp_gib": 9.911010913550854,
+ "compile_s": 11.900716543197632
+ },
+ "checks": {
+ "fwn vs official": {
+ "loss": {
+ "official": 1.203136682510376,
+ "other": 1.2031503915786743,
+ "rel_diff": 1.1394439632357909e-05,
+ "max_per_var_rel_diff": 0.0032936017593938175
+ },
+ "grads": {
+ "global_rel_l2": 0.016157892048622315,
+ "global_cos": 0.9998696868389602,
+ "max_rel_l2": 0.21444993708984028,
+ "n_params": 86,
+ "n_compared": 58,
+ "n_zero_both": 6,
+ "n_zero_ref_only": 0,
+ "all_finite": true,
+ "worst": [
+ {
+ "param": "multimodality_forward/grid_decoder/split_output_linear/w_cyclone_usa_wind_disc_time=21600",
+ "size": 512,
+ "ref_norm": 0.07069019381708155,
+ "got_norm": 0.08089039318185927,
+ "finite": true,
+ "rel_l2": 0.21444993708984028
+ },
+ {
+ "param": "multimodality_forward/grid_decoder/split_output_linear/w_cyclone_usa_r34_nw_radius_disc_time=21600",
+ "size": 512,
+ "ref_norm": 0.006235716535720746,
+ "got_norm": 0.006682095378606407,
+ "finite": true,
+ "rel_l2": 0.11240171669931483
+ },
+ {
+ "param": "multimodality_forward/grid_decoder/split_output_linear/w_u_component_of_wind_level=50,100,150,200,250,300,400,500,600,700,850,925,1000_time=21600",
+ "size": 6656,
+ "ref_norm": 0.06715556666634775,
+ "got_norm": 0.06710388486504754,
+ "finite": true,
+ "rel_l2": 0.028785047480818227
+ },
+ {
+ "param": "multimodality_forward/grid_decoder/split_output_linear/w_10m_u_component_of_wind_time=21600",
+ "size": 512,
+ "ref_norm": 0.014901104500827768,
+ "got_norm": 0.014807208146547488,
+ "finite": true,
+ "rel_l2": 0.028554408893876525
+ },
+ {
+ "param": "multimodality_forward/grid_decoder/split_output_linear/w_specific_humidity_level=50,100,150,200,250,300,400,500,600,700,850,925,1000_time=21600",
+ "size": 6656,
+ "ref_norm": 0.06544624577926476,
+ "got_norm": 0.06550854734895428,
+ "finite": true,
+ "rel_l2": 0.02476407440764222
+ }
+ ]
+ }
+ },
+ "official_mha vs official (floor)": {
+ "loss": {
+ "official": 1.203136682510376,
+ "other": 1.2031195163726807,
+ "rel_diff": 1.4267820061387295e-05,
+ "max_per_var_rel_diff": 0.007055680412290095
+ },
+ "grads": {
+ "global_rel_l2": 0.015840275108020836,
+ "global_cos": 0.9998757501997555,
+ "max_rel_l2": 0.21433064957118533,
+ "n_params": 86,
+ "n_compared": 58,
+ "n_zero_both": 6,
+ "n_zero_ref_only": 0,
+ "all_finite": true,
+ "worst": [
+ {
+ "param": "multimodality_forward/grid_decoder/split_output_linear/w_cyclone_usa_wind_disc_time=21600",
+ "size": 512,
+ "ref_norm": 0.07069019381708155,
+ "got_norm": 0.0808980205235208,
+ "finite": true,
+ "rel_l2": 0.21433064957118533
+ },
+ {
+ "param": "multimodality_forward/grid_decoder/split_output_linear/w_cyclone_usa_r64_nw_radius_disc_time=21600",
+ "size": 512,
+ "ref_norm": 0.005471387150023349,
+ "got_norm": 0.005483216299158898,
+ "finite": true,
+ "rel_l2": 0.06063645607587121
+ },
+ {
+ "param": "multimodality_forward/grid_decoder/split_output_linear/w_10m_u_component_of_wind_time=21600",
+ "size": 512,
+ "ref_norm": 0.014901104500827768,
+ "got_norm": 0.014838145361377278,
+ "finite": true,
+ "rel_l2": 0.036706715477910504
+ },
+ {
+ "param": "multimodality_forward/grid_decoder/split_output_linear/w_10m_v_component_of_wind_time=21600",
+ "size": 512,
+ "ref_norm": 0.019737312712881868,
+ "got_norm": 0.01965782792484462,
+ "finite": true,
+ "rel_l2": 0.03295148927873153
+ },
+ {
+ "param": "multimodality_forward/grid_decoder/split_output_linear/w_u_component_of_wind_level=50,100,150,200,250,300,400,500,600,700,850,925,1000_time=21600",
+ "size": 6656,
+ "ref_norm": 0.06715556666634775,
+ "got_norm": 0.06710532828785284,
+ "finite": true,
+ "rel_l2": 0.026323949572432923
+ }
+ ]
+ }
+ }
+ }
+ }
+ }
+}
\ No newline at end of file
diff --git a/docs/results/finetune/head_grads_mini_rtx5090.log b/docs/results/finetune/head_grads_mini_rtx5090.log
new file mode 100644
index 0000000..2d67c61
--- /dev/null
+++ b/docs/results/finetune/head_grads_mini_rtx5090.log
@@ -0,0 +1,181 @@
+gpu-run: queued behind /home/raymxnd/Workspace/WeatherForcast/.venv/bin/python scripts/verify_equivalen
+gpu-run: queued behind /home/raymxnd/Workspace/WeatherForcast/.venv/bin/python scripts/verify_equivalen
+gpu-run: queued behind /home/raymxnd/Workspace/WeatherForcast/.venv/bin/python scripts/verify_equivalen
+WeatherNextCyclones_Mini: head 2.54M params; masked targets: []
+[highest]
+WARNING:root:No normalization location found for cyclone_exists_gaussian_unit_mode
+WARNING:root:No normalization scale found for cyclone_exists_gaussian_unit_mode
+WARNING:absl:No coordinates found for dim noise_channels of array Size: 256B
+JitTracer(float32[2,1,32])
+Coordinates:
+ * batch (batch) object 16B MultiIndex
+ * _batch (batch) int64 16B 0 0
+ * sample (batch) int64 16B 0 1
+Dimensions without coordinates: channels, noise_channels., using index for param name. This may make model surgery more error prone if changing coordinates at fine-tuning time.
+/home/raymxnd/Workspace/WeatherForcast/.venv/lib/python3.12/site-packages/scipy/sparse/_index.py:216: SparseEfficiencyWarning: Changing the sparsity structure of a csr_matrix is expensive. lil and dok are more efficient.
+ self._set_arrayXarray(i, j, x)
+WARNING:absl:Only the first input of the static gather operation is being used, others are being ignored at runtime.
+WARNING:absl:Only the first input of the static gather operation is being used, others are being ignored at runtime.
+WARNING:absl:Only the first input of the sorted segment sum operation is being used, others are being ignored at runtime.
+WARNING:absl:Only the first input of the static gather operation is being used, others are being ignored at runtime.
+WARNING:absl:Only the first input of the static gather operation is being used, others are being ignored at runtime.
+WARNING:absl:Only the first input of the static gather operation is being used, others are being ignored at runtime.
+WARNING:absl:Only the first input of the sorted segment sum operation is being used, others are being ignored at runtime.
+/home/raymxnd/Workspace/WeatherForcast/weathernext/weathernext/weathernext2/architecture.py:286: FutureWarning: Deleting a single level of a MultiIndex is deprecated. Previously, this deleted all levels of a MultiIndex. Please also drop the following variables: {'batch'} to avoid an error in the future.
+ predictions[var] = xarray_jax.apply_ufunc(activation_fn, predictions[var])
+ official highest temp 7.36 GiB step 0.42s loss 1.202958
+WARNING:root:No normalization location found for cyclone_exists_gaussian_unit_mode
+WARNING:root:No normalization scale found for cyclone_exists_gaussian_unit_mode
+WARNING:absl:No coordinates found for dim noise_channels of array Size: 256B
+JitTracer(float32[2,1,32])
+Coordinates:
+ * batch (batch) object 16B MultiIndex
+ * _batch (batch) int64 16B 0 0
+ * sample (batch) int64 16B 0 1
+Dimensions without coordinates: channels, noise_channels., using index for param name. This may make model surgery more error prone if changing coordinates at fine-tuning time.
+/home/raymxnd/Workspace/WeatherForcast/.venv/lib/python3.12/site-packages/scipy/sparse/_index.py:216: SparseEfficiencyWarning: Changing the sparsity structure of a csr_matrix is expensive. lil and dok are more efficient.
+ self._set_arrayXarray(i, j, x)
+WARNING:absl:Only the first input of the static gather operation is being used, others are being ignored at runtime.
+WARNING:absl:Only the first input of the static gather operation is being used, others are being ignored at runtime.
+WARNING:absl:Only the first input of the sorted segment sum operation is being used, others are being ignored at runtime.
+WARNING:absl:Only the first input of the static gather operation is being used, others are being ignored at runtime.
+WARNING:absl:Only the first input of the static gather operation is being used, others are being ignored at runtime.
+WARNING:absl:Only the first input of the static gather operation is being used, others are being ignored at runtime.
+WARNING:absl:Only the first input of the sorted segment sum operation is being used, others are being ignored at runtime.
+/home/raymxnd/Workspace/WeatherForcast/weathernext/weathernext/weathernext2/architecture.py:286: FutureWarning: Deleting a single level of a MultiIndex is deprecated. Previously, this deleted all levels of a MultiIndex. Please also drop the following variables: {'batch'} to avoid an error in the future.
+ predictions[var] = xarray_jax.apply_ufunc(activation_fn, predictions[var])
+ official_mha highest temp 7.64 GiB step 0.34s loss 1.202958
+WARNING:root:No normalization location found for cyclone_exists_gaussian_unit_mode
+WARNING:root:No normalization scale found for cyclone_exists_gaussian_unit_mode
+WARNING:absl:No coordinates found for dim noise_channels of array Size: 256B
+JitTracer(float32[2,1,32])
+Coordinates:
+ * batch (batch) object 16B MultiIndex
+ * _batch (batch) int64 16B 0 0
+ * sample (batch) int64 16B 0 1
+Dimensions without coordinates: channels, noise_channels., using index for param name. This may make model surgery more error prone if changing coordinates at fine-tuning time.
+/home/raymxnd/Workspace/WeatherForcast/.venv/lib/python3.12/site-packages/scipy/sparse/_index.py:216: SparseEfficiencyWarning: Changing the sparsity structure of a csr_matrix is expensive. lil and dok are more efficient.
+ self._set_arrayXarray(i, j, x)
+/home/raymxnd/Workspace/WeatherForcast/weathernext/weathernext/weathernext2/architecture.py:286: FutureWarning: Deleting a single level of a MultiIndex is deprecated. Previously, this deleted all levels of a MultiIndex. Please also drop the following variables: {'batch'} to avoid an error in the future.
+ predictions[var] = xarray_jax.apply_ufunc(activation_fn, predictions[var])
+ fwn highest temp 1.48 GiB step 0.22s loss 1.202958
+WARNING:root:No normalization location found for cyclone_exists_gaussian_unit_mode
+WARNING:root:No normalization scale found for cyclone_exists_gaussian_unit_mode
+WARNING:absl:No coordinates found for dim noise_channels of array Size: 256B
+JitTracer(float32[2,1,32])
+Coordinates:
+ * batch (batch) object 16B MultiIndex
+ * _batch (batch) int64 16B 0 0
+ * sample (batch) int64 16B 0 1
+Dimensions without coordinates: channels, noise_channels., using index for param name. This may make model surgery more error prone if changing coordinates at fine-tuning time.
+/home/raymxnd/Workspace/WeatherForcast/.venv/lib/python3.12/site-packages/scipy/sparse/_index.py:216: SparseEfficiencyWarning: Changing the sparsity structure of a csr_matrix is expensive. lil and dok are more efficient.
+ self._set_arrayXarray(i, j, x)
+/home/raymxnd/Workspace/WeatherForcast/weathernext/weathernext/weathernext2/architecture.py:286: FutureWarning: Deleting a single level of a MultiIndex is deprecated. Previously, this deleted all levels of a MultiIndex. Please also drop the following variables: {'batch'} to avoid an error in the future.
+ predictions[var] = xarray_jax.apply_ufunc(activation_fn, predictions[var])
+ fwn_noremat highest temp 9.91 GiB (compile only)
+WARNING:root:No normalization location found for cyclone_exists_gaussian_unit_mode
+WARNING:root:No normalization scale found for cyclone_exists_gaussian_unit_mode
+WARNING:absl:No coordinates found for dim noise_channels of array Size: 256B
+JitTracer(float32[2,1,32])
+Coordinates:
+ * batch (batch) object 16B MultiIndex
+ * _batch (batch) int64 16B 0 0
+ * sample (batch) int64 16B 0 1
+Dimensions without coordinates: channels, noise_channels., using index for param name. This may make model surgery more error prone if changing coordinates at fine-tuning time.
+/home/raymxnd/Workspace/WeatherForcast/.venv/lib/python3.12/site-packages/scipy/sparse/_index.py:216: SparseEfficiencyWarning: Changing the sparsity structure of a csr_matrix is expensive. lil and dok are more efficient.
+ self._set_arrayXarray(i, j, x)
+/home/raymxnd/Workspace/WeatherForcast/weathernext/weathernext/weathernext2/architecture.py:286: FutureWarning: Deleting a single level of a MultiIndex is deprecated. Previously, this deleted all levels of a MultiIndex. Please also drop the following variables: {'batch'} to avoid an error in the future.
+ predictions[var] = xarray_jax.apply_ufunc(activation_fn, predictions[var])
+ fwn_strict highest temp 1.48 GiB step 1.43s loss 1.202958
+ fwn vs official loss rel diff 9.91e-08 (per var max 1.62e-05) grads: global rel L2 1.00e-03 cos 0.9999995043 max per-param rel L2 2.30e-03 (58/86 params; zero in both 6, zero only in official 0) finite True
+ 2.30e-03 multimodality_forward/grid_decoder/split_output_linear/w_specific_humidity_level=50,100,150,200,250,300,400,500,600,700,850,925,1000_time=21600
+ 1.38e-03 multimodality_forward/mesh_to_grid_gnn/deep_gnn/processor_edges_0_sender_mesh_to_points_nodes/w
+ 1.35e-03 multimodality_forward/grid_decoder/split_output_linear/w_vertical_velocity_level=50,100,150,200,250,300,400,500,600,700,850,925,1000_time=21600
+ fwn strict_fp32 vs official loss rel diff 0.00e+00 (per var max 1.17e-05) grads: global rel L2 8.85e-04 cos 0.9999996163 max per-param rel L2 1.41e-03 (58/86 params; zero in both 6, zero only in official 0) finite True
+ 1.41e-03 multimodality_forward/grid_decoder/split_output_linear/w_vertical_velocity_level=50,100,150,200,250,300,400,500,600,700,850,925,1000_time=21600
+ 1.32e-03 multimodality_forward/mesh_to_grid_gnn/deep_gnn/processor_edges_0_sender_mesh_to_points_nodes/w
+ 1.27e-03 multimodality_forward/grid_decoder/split_output_linear/w_u_component_of_wind_level=50,100,150,200,250,300,400,500,600,700,850,925,1000_time=21600
+ official_mha vs official (floor) loss rel diff 0.00e+00 (per var max 1.26e-05) grads: global rel L2 5.66e-04 cos 0.9999998399 max per-param rel L2 2.36e-03 (58/86 params; zero in both 6, zero only in official 0) finite True
+ 2.36e-03 multimodality_forward/grid_decoder/split_output_linear/w_specific_humidity_level=50,100,150,200,250,300,400,500,600,700,850,925,1000_time=21600
+ 1.31e-03 multimodality_forward/grid_decoder/split_output_linear/w_10m_v_component_of_wind_time=21600
+ 1.28e-03 multimodality_forward/grid_decoder/split_output_linear/w_v_component_of_wind_level=50,100,150,200,250,300,400,500,600,700,850,925,1000_time=21600
+[default]
+WARNING:root:No normalization location found for cyclone_exists_gaussian_unit_mode
+WARNING:root:No normalization scale found for cyclone_exists_gaussian_unit_mode
+WARNING:absl:No coordinates found for dim noise_channels of array Size: 256B
+JitTracer(float32[2,1,32])
+Coordinates:
+ * batch (batch) object 16B MultiIndex
+ * _batch (batch) int64 16B 0 0
+ * sample (batch) int64 16B 0 1
+Dimensions without coordinates: channels, noise_channels., using index for param name. This may make model surgery more error prone if changing coordinates at fine-tuning time.
+/home/raymxnd/Workspace/WeatherForcast/.venv/lib/python3.12/site-packages/scipy/sparse/_index.py:216: SparseEfficiencyWarning: Changing the sparsity structure of a csr_matrix is expensive. lil and dok are more efficient.
+ self._set_arrayXarray(i, j, x)
+WARNING:absl:Only the first input of the static gather operation is being used, others are being ignored at runtime.
+WARNING:absl:Only the first input of the static gather operation is being used, others are being ignored at runtime.
+WARNING:absl:Only the first input of the sorted segment sum operation is being used, others are being ignored at runtime.
+WARNING:absl:Only the first input of the static gather operation is being used, others are being ignored at runtime.
+WARNING:absl:Only the first input of the static gather operation is being used, others are being ignored at runtime.
+WARNING:absl:Only the first input of the static gather operation is being used, others are being ignored at runtime.
+WARNING:absl:Only the first input of the sorted segment sum operation is being used, others are being ignored at runtime.
+/home/raymxnd/Workspace/WeatherForcast/weathernext/weathernext/weathernext2/architecture.py:286: FutureWarning: Deleting a single level of a MultiIndex is deprecated. Previously, this deleted all levels of a MultiIndex. Please also drop the following variables: {'batch'} to avoid an error in the future.
+ predictions[var] = xarray_jax.apply_ufunc(activation_fn, predictions[var])
+ official default temp 5.26 GiB step 0.29s loss 1.203137
+WARNING:root:No normalization location found for cyclone_exists_gaussian_unit_mode
+WARNING:root:No normalization scale found for cyclone_exists_gaussian_unit_mode
+WARNING:absl:No coordinates found for dim noise_channels of array Size: 256B
+JitTracer(float32[2,1,32])
+Coordinates:
+ * batch (batch) object 16B MultiIndex
+ * _batch (batch) int64 16B 0 0
+ * sample (batch) int64 16B 0 1
+Dimensions without coordinates: channels, noise_channels., using index for param name. This may make model surgery more error prone if changing coordinates at fine-tuning time.
+/home/raymxnd/Workspace/WeatherForcast/.venv/lib/python3.12/site-packages/scipy/sparse/_index.py:216: SparseEfficiencyWarning: Changing the sparsity structure of a csr_matrix is expensive. lil and dok are more efficient.
+ self._set_arrayXarray(i, j, x)
+WARNING:absl:Only the first input of the static gather operation is being used, others are being ignored at runtime.
+WARNING:absl:Only the first input of the static gather operation is being used, others are being ignored at runtime.
+WARNING:absl:Only the first input of the sorted segment sum operation is being used, others are being ignored at runtime.
+WARNING:absl:Only the first input of the static gather operation is being used, others are being ignored at runtime.
+WARNING:absl:Only the first input of the static gather operation is being used, others are being ignored at runtime.
+WARNING:absl:Only the first input of the static gather operation is being used, others are being ignored at runtime.
+WARNING:absl:Only the first input of the sorted segment sum operation is being used, others are being ignored at runtime.
+/home/raymxnd/Workspace/WeatherForcast/weathernext/weathernext/weathernext2/architecture.py:286: FutureWarning: Deleting a single level of a MultiIndex is deprecated. Previously, this deleted all levels of a MultiIndex. Please also drop the following variables: {'batch'} to avoid an error in the future.
+ predictions[var] = xarray_jax.apply_ufunc(activation_fn, predictions[var])
+ official_mha default temp 7.64 GiB step 0.27s loss 1.203120
+WARNING:root:No normalization location found for cyclone_exists_gaussian_unit_mode
+WARNING:root:No normalization scale found for cyclone_exists_gaussian_unit_mode
+WARNING:absl:No coordinates found for dim noise_channels of array Size: 256B
+JitTracer(float32[2,1,32])
+Coordinates:
+ * batch (batch) object 16B MultiIndex
+ * _batch (batch) int64 16B 0 0
+ * sample (batch) int64 16B 0 1
+Dimensions without coordinates: channels, noise_channels., using index for param name. This may make model surgery more error prone if changing coordinates at fine-tuning time.
+/home/raymxnd/Workspace/WeatherForcast/.venv/lib/python3.12/site-packages/scipy/sparse/_index.py:216: SparseEfficiencyWarning: Changing the sparsity structure of a csr_matrix is expensive. lil and dok are more efficient.
+ self._set_arrayXarray(i, j, x)
+/home/raymxnd/Workspace/WeatherForcast/weathernext/weathernext/weathernext2/architecture.py:286: FutureWarning: Deleting a single level of a MultiIndex is deprecated. Previously, this deleted all levels of a MultiIndex. Please also drop the following variables: {'batch'} to avoid an error in the future.
+ predictions[var] = xarray_jax.apply_ufunc(activation_fn, predictions[var])
+ fwn default temp 1.48 GiB step 0.15s loss 1.203150
+WARNING:root:No normalization location found for cyclone_exists_gaussian_unit_mode
+WARNING:root:No normalization scale found for cyclone_exists_gaussian_unit_mode
+WARNING:absl:No coordinates found for dim noise_channels of array Size: 256B
+JitTracer(float32[2,1,32])
+Coordinates:
+ * batch (batch) object 16B MultiIndex
+ * _batch (batch) int64 16B 0 0
+ * sample (batch) int64 16B 0 1
+Dimensions without coordinates: channels, noise_channels., using index for param name. This may make model surgery more error prone if changing coordinates at fine-tuning time.
+/home/raymxnd/Workspace/WeatherForcast/.venv/lib/python3.12/site-packages/scipy/sparse/_index.py:216: SparseEfficiencyWarning: Changing the sparsity structure of a csr_matrix is expensive. lil and dok are more efficient.
+ self._set_arrayXarray(i, j, x)
+/home/raymxnd/Workspace/WeatherForcast/weathernext/weathernext/weathernext2/architecture.py:286: FutureWarning: Deleting a single level of a MultiIndex is deprecated. Previously, this deleted all levels of a MultiIndex. Please also drop the following variables: {'batch'} to avoid an error in the future.
+ predictions[var] = xarray_jax.apply_ufunc(activation_fn, predictions[var])
+ fwn_noremat default temp 9.91 GiB (compile only)
+ fwn vs official loss rel diff 1.14e-05 (per var max 3.29e-03) grads: global rel L2 1.62e-02 cos 0.9998696868 max per-param rel L2 2.14e-01 (58/86 params; zero in both 6, zero only in official 0) finite True
+ 2.14e-01 multimodality_forward/grid_decoder/split_output_linear/w_cyclone_usa_wind_disc_time=21600
+ 1.12e-01 multimodality_forward/grid_decoder/split_output_linear/w_cyclone_usa_r34_nw_radius_disc_time=21600
+ 2.88e-02 multimodality_forward/grid_decoder/split_output_linear/w_u_component_of_wind_level=50,100,150,200,250,300,400,500,600,700,850,925,1000_time=21600
+ official_mha vs official (floor) loss rel diff 1.43e-05 (per var max 7.06e-03) grads: global rel L2 1.58e-02 cos 0.9998757502 max per-param rel L2 2.14e-01 (58/86 params; zero in both 6, zero only in official 0) finite True
+ 2.14e-01 multimodality_forward/grid_decoder/split_output_linear/w_cyclone_usa_wind_disc_time=21600
+ 6.06e-02 multimodality_forward/grid_decoder/split_output_linear/w_cyclone_usa_r64_nw_radius_disc_time=21600
+ 3.67e-02 multimodality_forward/grid_decoder/split_output_linear/w_10m_u_component_of_wind_time=21600
+wrote docs/results/finetune/head_grads_mini_rtx5090.json
diff --git a/docs/results/finetune/head_sequential_compile_rtx5090.json b/docs/results/finetune/head_sequential_compile_rtx5090.json
new file mode 100644
index 0000000..561e184
--- /dev/null
+++ b/docs/results/finetune/head_sequential_compile_rtx5090.json
@@ -0,0 +1,66 @@
+{
+ "model": "WeatherNext2",
+ "checkpoint": 1,
+ "device": "NVIDIA GeForce RTX 5090",
+ "jax": "0.11.2",
+ "faster_weathernext": "0.1.1",
+ "args": {
+ "model": "WeatherNext2",
+ "checkpoint": 1,
+ "train": [
+ "2026102200"
+ ],
+ "eval": [],
+ "steps": 10,
+ "lr": 3e-06,
+ "warmup": 10,
+ "weight_decay": 0.0,
+ "clip": 1.0,
+ "eval_keys": 2,
+ "precision": "default",
+ "seed": 0,
+ "no_precip": false,
+ "modules": null,
+ "sequential_samples": "on",
+ "mask": [
+ "sea_surface_temperature"
+ ],
+ "compile_only": true,
+ "no_remat": false,
+ "weights_dir": "/home/raymxnd/Workspace/WeatherForcast/weights",
+ "out": null,
+ "log": "docs/results/finetune/head_sequential_compile_rtx5090.json"
+ },
+ "masked": [
+ "cyclone_all_wind_disc",
+ "cyclone_exists_gaussian_unit_mode",
+ "cyclone_usa_pres_disc",
+ "cyclone_usa_r34_ne_radius_disc",
+ "cyclone_usa_r34_nw_radius_disc",
+ "cyclone_usa_r34_se_radius_disc",
+ "cyclone_usa_r34_sw_radius_disc",
+ "cyclone_usa_r50_ne_radius_disc",
+ "cyclone_usa_r50_nw_radius_disc",
+ "cyclone_usa_r50_se_radius_disc",
+ "cyclone_usa_r50_sw_radius_disc",
+ "cyclone_usa_r64_ne_radius_disc",
+ "cyclone_usa_r64_nw_radius_disc",
+ "cyclone_usa_r64_se_radius_disc",
+ "cyclone_usa_r64_sw_radius_disc",
+ "cyclone_usa_rmw_disc",
+ "cyclone_usa_wind_disc",
+ "sea_surface_temperature"
+ ],
+ "modules": [
+ "mesh_to_grid_gnn",
+ "grid_decoder"
+ ],
+ "head_params": 5572967,
+ "total_params": 183785575,
+ "compile": {
+ "seconds": 86.71063470840454,
+ "temp_gib": 12.27459704130888,
+ "argument_gib": 1.7898386679589748,
+ "output_gib": 0.0622851625084877
+ }
+}
\ No newline at end of file
diff --git a/docs/results/finetune/head_sequential_pool16gb_rtx5090.json b/docs/results/finetune/head_sequential_pool16gb_rtx5090.json
new file mode 100644
index 0000000..31ab085
--- /dev/null
+++ b/docs/results/finetune/head_sequential_pool16gb_rtx5090.json
@@ -0,0 +1,227 @@
+{
+ "model": "WeatherNext2",
+ "checkpoint": 1,
+ "device": "NVIDIA GeForce RTX 5090",
+ "jax": "0.11.2",
+ "faster_weathernext": "0.1.1",
+ "args": {
+ "model": "WeatherNext2",
+ "checkpoint": 1,
+ "train": [
+ "2026100412",
+ "2026100418",
+ "2026100500",
+ "2026100506",
+ "2026100512",
+ "2026100518",
+ "2026100600",
+ "2026100606",
+ "2026100612",
+ "2026100618"
+ ],
+ "eval": [
+ "2026103012"
+ ],
+ "steps": 10,
+ "lr": 3e-06,
+ "warmup": 10,
+ "weight_decay": 0.0,
+ "clip": 1.0,
+ "eval_keys": 1,
+ "precision": "default",
+ "seed": 0,
+ "no_precip": false,
+ "modules": null,
+ "sequential_samples": "on",
+ "mask": [
+ "sea_surface_temperature"
+ ],
+ "compile_only": false,
+ "no_remat": false,
+ "weights_dir": "/home/raymxnd/Workspace/WeatherForcast/weights",
+ "out": null,
+ "log": "docs/results/finetune/head_sequential_pool16gb_rtx5090.json"
+ },
+ "masked": [
+ "cyclone_all_wind_disc",
+ "cyclone_exists_gaussian_unit_mode",
+ "cyclone_usa_pres_disc",
+ "cyclone_usa_r34_ne_radius_disc",
+ "cyclone_usa_r34_nw_radius_disc",
+ "cyclone_usa_r34_se_radius_disc",
+ "cyclone_usa_r34_sw_radius_disc",
+ "cyclone_usa_r50_ne_radius_disc",
+ "cyclone_usa_r50_nw_radius_disc",
+ "cyclone_usa_r50_se_radius_disc",
+ "cyclone_usa_r50_sw_radius_disc",
+ "cyclone_usa_r64_ne_radius_disc",
+ "cyclone_usa_r64_nw_radius_disc",
+ "cyclone_usa_r64_se_radius_disc",
+ "cyclone_usa_r64_sw_radius_disc",
+ "cyclone_usa_rmw_disc",
+ "cyclone_usa_wind_disc",
+ "sea_surface_temperature"
+ ],
+ "modules": [
+ "mesh_to_grid_gnn",
+ "grid_decoder"
+ ],
+ "head_params": 5572967,
+ "total_params": 183785575,
+ "compile": {
+ "seconds": 86.93455600738525,
+ "temp_gib": 12.274611108005047,
+ "argument_gib": 1.7898386679589748,
+ "output_gib": 0.0622851625084877
+ },
+ "eval_before": {
+ "2026103012": {
+ "loss": 1.0278544425964355,
+ "per_var": {
+ "100m_u_component_of_wind": 0.13780640065670013,
+ "100m_v_component_of_wind": 0.12773455679416656,
+ "10m_u_component_of_wind": 0.14574971795082092,
+ "10m_v_component_of_wind": 0.13495780527591705,
+ "2m_temperature": 0.09912414848804474,
+ "cyclone_all_wind_disc": 0.0,
+ "cyclone_exists_gaussian_unit_mode": 0.0,
+ "cyclone_usa_pres_disc": 0.0,
+ "cyclone_usa_r34_ne_radius_disc": 0.0,
+ "cyclone_usa_r34_nw_radius_disc": 0.0,
+ "cyclone_usa_r34_se_radius_disc": 0.0,
+ "cyclone_usa_r34_sw_radius_disc": 0.0,
+ "cyclone_usa_r50_ne_radius_disc": 0.0,
+ "cyclone_usa_r50_nw_radius_disc": 0.0,
+ "cyclone_usa_r50_se_radius_disc": 0.0,
+ "cyclone_usa_r50_sw_radius_disc": 0.0,
+ "cyclone_usa_r64_ne_radius_disc": 0.0,
+ "cyclone_usa_r64_nw_radius_disc": 0.0,
+ "cyclone_usa_r64_se_radius_disc": 0.0,
+ "cyclone_usa_r64_sw_radius_disc": 0.0,
+ "cyclone_usa_rmw_disc": 0.0,
+ "cyclone_usa_wind_disc": 0.0,
+ "geopotential": 0.060890767723321915,
+ "mean_sea_level_pressure": 0.06439507007598877,
+ "sea_surface_temperature": 1.643149733543396,
+ "specific_humidity": 0.21508437395095825,
+ "temperature": 0.11762525141239166,
+ "total_precipitation_6hr": 0.14055952429771423,
+ "u_component_of_wind": 0.12581771612167358,
+ "v_component_of_wind": 0.1115986555814743,
+ "vertical_velocity": 0.28035101294517517
+ }
+ }
+ },
+ "steps": [
+ {
+ "step": 0,
+ "init": "2026100512",
+ "loss": 1.0483345985412598,
+ "grad_norm": 3.6790037155151367,
+ "seconds": 4.4708802700042725
+ },
+ {
+ "step": 1,
+ "init": "2026100600",
+ "loss": 0.9512615203857422,
+ "grad_norm": 3.759441614151001,
+ "seconds": 4.1550257205963135
+ },
+ {
+ "step": 2,
+ "init": "2026100500",
+ "loss": 1.0784831047058105,
+ "grad_norm": 5.346747398376465,
+ "seconds": 4.1635026931762695
+ },
+ {
+ "step": 3,
+ "init": "2026100606",
+ "loss": 1.0890036821365356,
+ "grad_norm": 3.225839853286743,
+ "seconds": 4.167656421661377
+ },
+ {
+ "step": 4,
+ "init": "2026100506",
+ "loss": 1.1361671686172485,
+ "grad_norm": 3.379784107208252,
+ "seconds": 4.168152570724487
+ },
+ {
+ "step": 5,
+ "init": "2026100518",
+ "loss": 1.0784902572631836,
+ "grad_norm": 3.357424259185791,
+ "seconds": 4.174056053161621
+ },
+ {
+ "step": 6,
+ "init": "2026100618",
+ "loss": 1.2269914150238037,
+ "grad_norm": 5.038262844085693,
+ "seconds": 4.176669120788574
+ },
+ {
+ "step": 7,
+ "init": "2026100412",
+ "loss": 1.0577998161315918,
+ "grad_norm": 3.0855441093444824,
+ "seconds": 4.173212051391602
+ },
+ {
+ "step": 8,
+ "init": "2026100612",
+ "loss": 1.0449318885803223,
+ "grad_norm": 3.120760440826416,
+ "seconds": 4.176803350448608
+ },
+ {
+ "step": 9,
+ "init": "2026100418",
+ "loss": 1.078608512878418,
+ "grad_norm": 2.8259177207946777,
+ "seconds": 4.174084663391113
+ }
+ ],
+ "peak_gib": 14.16015625,
+ "median_step_s": 4.173212051391602,
+ "eval_after": {
+ "2026103012": {
+ "loss": 1.0226142406463623,
+ "per_var": {
+ "100m_u_component_of_wind": 0.13774587213993073,
+ "100m_v_component_of_wind": 0.12740558385849,
+ "10m_u_component_of_wind": 0.1455879658460617,
+ "10m_v_component_of_wind": 0.1344604343175888,
+ "2m_temperature": 0.09660765528678894,
+ "cyclone_all_wind_disc": 0.0,
+ "cyclone_exists_gaussian_unit_mode": 0.0,
+ "cyclone_usa_pres_disc": 0.0,
+ "cyclone_usa_r34_ne_radius_disc": 0.0,
+ "cyclone_usa_r34_nw_radius_disc": 0.0,
+ "cyclone_usa_r34_se_radius_disc": 0.0,
+ "cyclone_usa_r34_sw_radius_disc": 0.0,
+ "cyclone_usa_r50_ne_radius_disc": 0.0,
+ "cyclone_usa_r50_nw_radius_disc": 0.0,
+ "cyclone_usa_r50_se_radius_disc": 0.0,
+ "cyclone_usa_r50_sw_radius_disc": 0.0,
+ "cyclone_usa_r64_ne_radius_disc": 0.0,
+ "cyclone_usa_r64_nw_radius_disc": 0.0,
+ "cyclone_usa_r64_se_radius_disc": 0.0,
+ "cyclone_usa_r64_sw_radius_disc": 0.0,
+ "cyclone_usa_rmw_disc": 0.0,
+ "cyclone_usa_wind_disc": 0.0,
+ "geopotential": 0.06051589921116829,
+ "mean_sea_level_pressure": 0.06435436010360718,
+ "sea_surface_temperature": 1.6431504487991333,
+ "specific_humidity": 0.2144441306591034,
+ "temperature": 0.11650193482637405,
+ "total_precipitation_6hr": 0.13957354426383972,
+ "u_component_of_wind": 0.1255229115486145,
+ "v_component_of_wind": 0.11156248301267624,
+ "vertical_velocity": 0.2800644338130951
+ }
+ }
+ }
+}
\ No newline at end of file
diff --git a/docs/results/finetune/head_sequential_pool16gb_rtx5090.log b/docs/results/finetune/head_sequential_pool16gb_rtx5090.log
new file mode 100644
index 0000000..ed3c1f9
--- /dev/null
+++ b/docs/results/finetune/head_sequential_pool16gb_rtx5090.log
@@ -0,0 +1,33 @@
+example 2026100412 built in 2.2s
+WeatherNext2 ckpt 1: training 5.57M of 183.8M params; masked targets: 18 (cyclone channels, sea_surface_temperature)
+WARNING:root:No normalization location found for cyclone_exists_gaussian_unit_mode
+WARNING:root:No normalization scale found for cyclone_exists_gaussian_unit_mode
+WARNING:absl:No coordinates found for dim noise_channels of array Size: 128B
+JitTracer(float32[1,1,32])
+Dimensions without coordinates: batch, channels, noise_channels., using index for param name. This may make model surgery more error prone if changing coordinates at fine-tuning time.
+/home/raymxnd/Workspace/WeatherForcast/.venv/lib/python3.12/site-packages/scipy/sparse/_index.py:216: SparseEfficiencyWarning: Changing the sparsity structure of a csr_matrix is expensive. lil and dok are more efficient.
+ self._set_arrayXarray(i, j, x)
+W1101 01:44:17.257984 210154 hlo_rematerialization.cc:3363] Can't reduce memory use below 11.81GiB (12685276736 bytes) by rematerialization; only reduced to 34.18GiB (36699004316 bytes), down from 34.18GiB (36699028940 bytes) originally
+train step compiled in 87s: temp 12.27 GiB, arguments 1.79 GiB
+WARNING:root:No normalization location found for cyclone_exists_gaussian_unit_mode
+WARNING:root:No normalization scale found for cyclone_exists_gaussian_unit_mode
+WARNING:absl:No coordinates found for dim noise_channels of array Size: 128B
+JitTracer(float32[1,1,32])
+Dimensions without coordinates: batch, channels, noise_channels., using index for param name. This may make model surgery more error prone if changing coordinates at fine-tuning time.
+/home/raymxnd/Workspace/WeatherForcast/.venv/lib/python3.12/site-packages/scipy/sparse/_index.py:216: SparseEfficiencyWarning: Changing the sparsity structure of a csr_matrix is expensive. lil and dok are more efficient.
+ self._set_arrayXarray(i, j, x)
+W1101 01:45:23.025802 210154 hlo_rematerialization.cc:3363] Can't reduce memory use below 11.79GiB (12665436685 bytes) by rematerialization; only reduced to 31.11GiB (33406662528 bytes), down from 31.99GiB (34353549912 bytes) originally
+eval before (63s): 2026103012 1.02785
+step 0 2026100512 loss 1.04833 |g| 3.679e+00 4.47s
+step 1 2026100600 loss 0.95126 |g| 3.759e+00 4.16s
+step 2 2026100500 loss 1.07848 |g| 5.347e+00 4.16s
+step 3 2026100606 loss 1.08900 |g| 3.226e+00 4.17s
+step 4 2026100506 loss 1.13617 |g| 3.380e+00 4.17s
+step 5 2026100518 loss 1.07849 |g| 3.357e+00 4.17s
+step 6 2026100618 loss 1.22699 |g| 5.038e+00 4.18s
+step 7 2026100412 loss 1.05780 |g| 3.086e+00 4.17s
+step 8 2026100612 loss 1.04493 |g| 3.121e+00 4.18s
+step 9 2026100418 loss 1.07861 |g| 2.826e+00 4.17s
+median step 4.17s, device peak 14.16 GiB
+eval 2026103012: 1.02785 -> 1.02261 (-0.51 %)
+wrote docs/results/finetune/head_sequential_pool16gb_rtx5090.json
diff --git a/docs/results/finetune/head_sequential_pool24gb_rtx5090.json b/docs/results/finetune/head_sequential_pool24gb_rtx5090.json
new file mode 100644
index 0000000..844e3d9
--- /dev/null
+++ b/docs/results/finetune/head_sequential_pool24gb_rtx5090.json
@@ -0,0 +1,227 @@
+{
+ "model": "WeatherNext2",
+ "checkpoint": 1,
+ "device": "NVIDIA GeForce RTX 5090",
+ "jax": "0.11.2",
+ "faster_weathernext": "0.1.1",
+ "args": {
+ "model": "WeatherNext2",
+ "checkpoint": 1,
+ "train": [
+ "2026100412",
+ "2026100418",
+ "2026100500",
+ "2026100506",
+ "2026100512",
+ "2026100518",
+ "2026100600",
+ "2026100606",
+ "2026100612",
+ "2026100618"
+ ],
+ "eval": [
+ "2026103012"
+ ],
+ "steps": 10,
+ "lr": 3e-06,
+ "warmup": 10,
+ "weight_decay": 0.0,
+ "clip": 1.0,
+ "eval_keys": 1,
+ "precision": "default",
+ "seed": 0,
+ "no_precip": false,
+ "modules": null,
+ "sequential_samples": "on",
+ "mask": [
+ "sea_surface_temperature"
+ ],
+ "compile_only": false,
+ "no_remat": false,
+ "weights_dir": "/home/raymxnd/Workspace/WeatherForcast/weights",
+ "out": null,
+ "log": "docs/results/finetune/head_sequential_pool24gb_rtx5090.json"
+ },
+ "masked": [
+ "cyclone_all_wind_disc",
+ "cyclone_exists_gaussian_unit_mode",
+ "cyclone_usa_pres_disc",
+ "cyclone_usa_r34_ne_radius_disc",
+ "cyclone_usa_r34_nw_radius_disc",
+ "cyclone_usa_r34_se_radius_disc",
+ "cyclone_usa_r34_sw_radius_disc",
+ "cyclone_usa_r50_ne_radius_disc",
+ "cyclone_usa_r50_nw_radius_disc",
+ "cyclone_usa_r50_se_radius_disc",
+ "cyclone_usa_r50_sw_radius_disc",
+ "cyclone_usa_r64_ne_radius_disc",
+ "cyclone_usa_r64_nw_radius_disc",
+ "cyclone_usa_r64_se_radius_disc",
+ "cyclone_usa_r64_sw_radius_disc",
+ "cyclone_usa_rmw_disc",
+ "cyclone_usa_wind_disc",
+ "sea_surface_temperature"
+ ],
+ "modules": [
+ "mesh_to_grid_gnn",
+ "grid_decoder"
+ ],
+ "head_params": 5572967,
+ "total_params": 183785575,
+ "compile": {
+ "seconds": 86.74354243278503,
+ "temp_gib": 12.274599187076092,
+ "argument_gib": 1.7898386679589748,
+ "output_gib": 0.0622851625084877
+ },
+ "eval_before": {
+ "2026103012": {
+ "loss": 1.0278562307357788,
+ "per_var": {
+ "100m_u_component_of_wind": 0.1378091424703598,
+ "100m_v_component_of_wind": 0.12773606181144714,
+ "10m_u_component_of_wind": 0.145751953125,
+ "10m_v_component_of_wind": 0.13495990633964539,
+ "2m_temperature": 0.09912233799695969,
+ "cyclone_all_wind_disc": 0.0,
+ "cyclone_exists_gaussian_unit_mode": 0.0,
+ "cyclone_usa_pres_disc": 0.0,
+ "cyclone_usa_r34_ne_radius_disc": 0.0,
+ "cyclone_usa_r34_nw_radius_disc": 0.0,
+ "cyclone_usa_r34_se_radius_disc": 0.0,
+ "cyclone_usa_r34_sw_radius_disc": 0.0,
+ "cyclone_usa_r50_ne_radius_disc": 0.0,
+ "cyclone_usa_r50_nw_radius_disc": 0.0,
+ "cyclone_usa_r50_se_radius_disc": 0.0,
+ "cyclone_usa_r50_sw_radius_disc": 0.0,
+ "cyclone_usa_r64_ne_radius_disc": 0.0,
+ "cyclone_usa_r64_nw_radius_disc": 0.0,
+ "cyclone_usa_r64_se_radius_disc": 0.0,
+ "cyclone_usa_r64_sw_radius_disc": 0.0,
+ "cyclone_usa_rmw_disc": 0.0,
+ "cyclone_usa_wind_disc": 0.0,
+ "geopotential": 0.06089600920677185,
+ "mean_sea_level_pressure": 0.06440246105194092,
+ "sea_surface_temperature": 1.6431478261947632,
+ "specific_humidity": 0.21508638560771942,
+ "temperature": 0.11762493848800659,
+ "total_precipitation_6hr": 0.1405649185180664,
+ "u_component_of_wind": 0.12581828236579895,
+ "v_component_of_wind": 0.1115967333316803,
+ "vertical_velocity": 0.2803499400615692
+ }
+ }
+ },
+ "steps": [
+ {
+ "step": 0,
+ "init": "2026100512",
+ "loss": 1.0483423471450806,
+ "grad_norm": 3.6788527965545654,
+ "seconds": 4.32775616645813
+ },
+ {
+ "step": 1,
+ "init": "2026100600",
+ "loss": 0.9512643814086914,
+ "grad_norm": 3.7587673664093018,
+ "seconds": 4.013657331466675
+ },
+ {
+ "step": 2,
+ "init": "2026100500",
+ "loss": 1.0784754753112793,
+ "grad_norm": 5.3454389572143555,
+ "seconds": 4.020414113998413
+ },
+ {
+ "step": 3,
+ "init": "2026100606",
+ "loss": 1.0890121459960938,
+ "grad_norm": 3.226203680038452,
+ "seconds": 4.022968053817749
+ },
+ {
+ "step": 4,
+ "init": "2026100506",
+ "loss": 1.1361663341522217,
+ "grad_norm": 3.379551887512207,
+ "seconds": 4.024449825286865
+ },
+ {
+ "step": 5,
+ "init": "2026100518",
+ "loss": 1.0784796476364136,
+ "grad_norm": 3.357921600341797,
+ "seconds": 4.032755613327026
+ },
+ {
+ "step": 6,
+ "init": "2026100618",
+ "loss": 1.227004885673523,
+ "grad_norm": 5.038512706756592,
+ "seconds": 4.0339744091033936
+ },
+ {
+ "step": 7,
+ "init": "2026100412",
+ "loss": 1.057813286781311,
+ "grad_norm": 3.085665225982666,
+ "seconds": 4.032515048980713
+ },
+ {
+ "step": 8,
+ "init": "2026100612",
+ "loss": 1.0449199676513672,
+ "grad_norm": 3.1186587810516357,
+ "seconds": 4.036591053009033
+ },
+ {
+ "step": 9,
+ "init": "2026100418",
+ "loss": 1.078609585762024,
+ "grad_norm": 2.8268682956695557,
+ "seconds": 4.032351970672607
+ }
+ ],
+ "peak_gib": 14.064476490020752,
+ "median_step_s": 4.032351970672607,
+ "eval_after": {
+ "2026103012": {
+ "loss": 1.0226144790649414,
+ "per_var": {
+ "100m_u_component_of_wind": 0.1377466470003128,
+ "100m_v_component_of_wind": 0.12740720808506012,
+ "10m_u_component_of_wind": 0.1455896943807602,
+ "10m_v_component_of_wind": 0.1344599574804306,
+ "2m_temperature": 0.09660831838846207,
+ "cyclone_all_wind_disc": 0.0,
+ "cyclone_exists_gaussian_unit_mode": 0.0,
+ "cyclone_usa_pres_disc": 0.0,
+ "cyclone_usa_r34_ne_radius_disc": 0.0,
+ "cyclone_usa_r34_nw_radius_disc": 0.0,
+ "cyclone_usa_r34_se_radius_disc": 0.0,
+ "cyclone_usa_r34_sw_radius_disc": 0.0,
+ "cyclone_usa_r50_ne_radius_disc": 0.0,
+ "cyclone_usa_r50_nw_radius_disc": 0.0,
+ "cyclone_usa_r50_se_radius_disc": 0.0,
+ "cyclone_usa_r50_sw_radius_disc": 0.0,
+ "cyclone_usa_r64_ne_radius_disc": 0.0,
+ "cyclone_usa_r64_nw_radius_disc": 0.0,
+ "cyclone_usa_r64_se_radius_disc": 0.0,
+ "cyclone_usa_r64_sw_radius_disc": 0.0,
+ "cyclone_usa_rmw_disc": 0.0,
+ "cyclone_usa_wind_disc": 0.0,
+ "geopotential": 0.060519445687532425,
+ "mean_sea_level_pressure": 0.06435970962047577,
+ "sea_surface_temperature": 1.6431505680084229,
+ "specific_humidity": 0.21444392204284668,
+ "temperature": 0.11650239676237106,
+ "total_precipitation_6hr": 0.1395885944366455,
+ "u_component_of_wind": 0.12552106380462646,
+ "v_component_of_wind": 0.11156003922224045,
+ "vertical_velocity": 0.28006401658058167
+ }
+ }
+ }
+}
\ No newline at end of file
diff --git a/docs/results/finetune/head_sequential_pool24gb_rtx5090.log b/docs/results/finetune/head_sequential_pool24gb_rtx5090.log
new file mode 100644
index 0000000..141eb09
--- /dev/null
+++ b/docs/results/finetune/head_sequential_pool24gb_rtx5090.log
@@ -0,0 +1,33 @@
+example 2026100412 built in 2.2s
+WeatherNext2 ckpt 1: training 5.57M of 183.8M params; masked targets: 18 (cyclone channels, sea_surface_temperature)
+WARNING:root:No normalization location found for cyclone_exists_gaussian_unit_mode
+WARNING:root:No normalization scale found for cyclone_exists_gaussian_unit_mode
+WARNING:absl:No coordinates found for dim noise_channels of array Size: 128B
+JitTracer(float32[1,1,32])
+Dimensions without coordinates: batch, channels, noise_channels., using index for param name. This may make model surgery more error prone if changing coordinates at fine-tuning time.
+/home/raymxnd/Workspace/WeatherForcast/.venv/lib/python3.12/site-packages/scipy/sparse/_index.py:216: SparseEfficiencyWarning: Changing the sparsity structure of a csr_matrix is expensive. lil and dok are more efficient.
+ self._set_arrayXarray(i, j, x)
+W1101 01:40:55.719289 203927 hlo_rematerialization.cc:3363] Can't reduce memory use below 18.54GiB (19905351642 bytes) by rematerialization; only reduced to 34.18GiB (36699016636 bytes), down from 34.18GiB (36699028940 bytes) originally
+train step compiled in 87s: temp 12.27 GiB, arguments 1.79 GiB
+WARNING:root:No normalization location found for cyclone_exists_gaussian_unit_mode
+WARNING:root:No normalization scale found for cyclone_exists_gaussian_unit_mode
+WARNING:absl:No coordinates found for dim noise_channels of array Size: 128B
+JitTracer(float32[1,1,32])
+Dimensions without coordinates: batch, channels, noise_channels., using index for param name. This may make model surgery more error prone if changing coordinates at fine-tuning time.
+/home/raymxnd/Workspace/WeatherForcast/.venv/lib/python3.12/site-packages/scipy/sparse/_index.py:216: SparseEfficiencyWarning: Changing the sparsity structure of a csr_matrix is expensive. lil and dok are more efficient.
+ self._set_arrayXarray(i, j, x)
+W1101 01:42:01.995738 203927 hlo_rematerialization.cc:3363] Can't reduce memory use below 18.52GiB (19885511591 bytes) by rematerialization; only reduced to 31.11GiB (33406662528 bytes), down from 31.99GiB (34353549912 bytes) originally
+eval before (63s): 2026103012 1.02786
+step 0 2026100512 loss 1.04834 |g| 3.679e+00 4.33s
+step 1 2026100600 loss 0.95126 |g| 3.759e+00 4.01s
+step 2 2026100500 loss 1.07848 |g| 5.345e+00 4.02s
+step 3 2026100606 loss 1.08901 |g| 3.226e+00 4.02s
+step 4 2026100506 loss 1.13617 |g| 3.380e+00 4.02s
+step 5 2026100518 loss 1.07848 |g| 3.358e+00 4.03s
+step 6 2026100618 loss 1.22700 |g| 5.039e+00 4.03s
+step 7 2026100412 loss 1.05781 |g| 3.086e+00 4.03s
+step 8 2026100612 loss 1.04492 |g| 3.119e+00 4.04s
+step 9 2026100418 loss 1.07861 |g| 2.827e+00 4.03s
+median step 4.03s, device peak 14.06 GiB
+eval 2026103012: 1.02786 -> 1.02261 (-0.51 %)
+wrote docs/results/finetune/head_sequential_pool24gb_rtx5090.json
diff --git a/docs/validation.md b/docs/validation.md
index 64a5ac5..e3be86c 100644
--- a/docs/validation.md
+++ b/docs/validation.md
@@ -120,6 +120,17 @@ An earlier build without the fused attention kernel ran at 18.5 s/step on the sa
149 autotuned ops differed across three repeats); the most likely explanation is one pick that
did not honour fp32 in that run. For strict fp32 use `--no-autotune` (5/5 matching A100 runs)
and for reproducible runs `--autotune-cache`. Files: [`results/a100_repeats/`](results/a100_repeats/).
+Update (2026-11-01): the canary below reproduced the deviation almost exactly (max rel RMS 1.32e-2,
+median 6.21e-4, same worst fields, min corr equal to 8 digits), so it is not a one-off pick: on
+A100, `cloud/canary_run.py` lands at TF32 level where the repeat script did not. Not yet explained;
+it predates the attention backward (same numbers on v0.1.0).
+
+**Canary after the attention backward (2026-11-01, branch `head-only-finetune`).** Same runner and
+cards as the table above (`modal run cloud/modal_canary.py --cards a100,a10g,l4,t4`, ~$1–2): every
+card ran and fits the 8 GB pool; steady-state s/step and peaks as in the table (A100 1.10 s / 6.18
+GiB, A10G 2.43 / 6.53, L4 3.57 / 6.53, T4 43 / 6.49 on the XLA path). fp32 outputs vs a fresh
+RTX 5090 run of the same code: A10G 1.67e-5, L4 1.85e-5, T4 9.67e-5 max rel RMS; A100 1.32e-2 (the
+note above). Files: [`results/canary_modal/v0.2/`](results/canary_modal/v0.2/).
## 5. Mini model and batch size > 1 (v0.1.1)
@@ -226,6 +237,184 @@ One WN2 layer's masked mesh attention (40,962 nodes, 6 heads × 128 dim, 32-hop
Files: [`results/intel_b580/`](results/intel_b580/) (`fwn_info.txt`, `attn_bench.txt`). All unit tests pass (`pytest -q`).
+## 8. Head-only fine-tuning (prototype)
+
+`faster_weathernext.finetune` trains the head (mesh→grid GNN + grid decoder) with the official
+training loss: the config's full predictor stack and FGN's fair CRPS over 2 members, i.e. an
+effective batch of 2. Gradients are taken w.r.t. the head only, so the frozen encoder, grid→mesh
+GNN and transformer run as in inference. The new option `remat` rematerialises each blocked
+GNN / encoder / decoder step in the backward pass. Results here: RTX 5090, JAX 0.11.2.
+
+**Inference is unchanged by the patch edit** (`remat` defaults to off; `scripts/verify_equivalence.py`
+re-run as in §2): `highest` max rel RMS 1.98e-5, default 4.87e-3 (v0.1.1: 2.00e-5, 4.93e-3); temp
+memory and step time unchanged (4.14 / 4.02 GiB, 1.49 / 0.77 s).
+Files: [`results/finetune/equivalence_rtx5090_remat_option.json`](results/finetune/equivalence_rtx5090_remat_option.json).
+
+**Head gradients vs the official GPU path** (`scripts/verify_grads.py`; the runs in §8 and §9 were made
+when it and `scripts/finetune.py` were called `verify_head_grads.py` / `finetune_head.py`; WeatherNextCyclones_Mini
+at 1°, Google's sample, same inputs / targets / noise). faster-weathernext runs with remat and
+small blocks (8,192 grid points, 16,384 edges), so the 1° grid passes through 8 blocks in every
+blocked module. The fair CRPS has sign terms (|pred − target|, |pred_i − pred_j|): rounding-level
+forward differences flip individual gradient contributions, so gradient agreement is judged
+against a floor, a second official implementation (dense `mha` instead of `triblockdiag_mha`),
+not against zero. Head: 2.54M parameters in 86 tensors; per-tensor errors are over the 58 tensors
+whose gradient norm exceeds 1e-3 of the largest (the others enter the global numbers).
+
+| precision | comparison | loss rel diff | grad global rel L2 | cosine | max per-tensor rel L2 |
+|---|---|---|---|---|---|
+| highest | fwn vs official | 9.9e-8 | 1.00e-3 | 0.99999950 | 2.3e-3 |
+| highest | fwn (`strict_fp32` attention) vs official | 0 | 8.9e-4 | 0.99999962 | 1.4e-3 |
+| highest | official `mha` vs official (floor) | 0 | 5.7e-4 | 0.99999984 | 2.4e-3 |
+| default | fwn vs official | 1.1e-5 | 1.62e-2 | 0.99987 | 0.21 |
+| default | official `mha` vs official (floor) | 1.4e-5 | 1.58e-2 | 0.99988 | 0.21 |
+
+At fp32 the global error is within 2× of the floor (which itself moved from 4.3e-4 to 5.7e-4
+between two runs: XLA autotuning) and the worst tensor is no worse than the floor's; at TF32 the
+two rows are the same. The worst tensors are output layers of individual variables (10 m wind,
+upper-level u; at TF32 the cyclone channels, whose loss terms are tiny). All gradients finite.
+Temp memory per step (Mini, effective batch 2): official 7.36 GiB (`highest`), faster-weathernext
+1.48 GiB with remat and 9.91 GiB without; 0.22 s vs 0.42 s per step.
+Files: [`results/finetune/head_grads_mini_rtx5090.json`](results/finetune/head_grads_mini_rtx5090.json) (+ `.log`).
+
+**WeatherNext 2 at 0.25°** (`scripts/finetune.py`, IFS open data). Compiled training step:
+20.03 GiB temp + 1.79 GiB arguments (default precision). Rematerialisation is what makes this
+fit: compiled for the CPU backend (whose buffer assignment is larger; the loss forward alone is
+21.1 GiB there vs 8.3 GiB on the GPU), the step needs 34.3 GiB temp with remat and 249.6 GiB
+without ([`results/finetune/cpu_compile_*.json`](results/finetune/)). The step needs one
+contiguous 20 GiB buffer: in a JAX pool that grows on demand, an evaluation pass run first left
+regions too small for it and the step failed with RESOURCE_EXHAUSTED
+([`results/finetune/finetune_smoke_rtx5090_oom_growing_pool.log`](results/finetune/finetune_smoke_rtx5090_oom_growing_pool.log));
+`finetune.py` therefore preallocates 90 % of the card.
+
+**Training runs** (checkpoint 1; 60 steps of one example each, 3.84 s per step, device peak
+21.8 GiB; Adam, gradient-norm clip 1). Training examples: 60 of 68 IFS open-data init times
+2026-10-04 12Z – 10-24 06Z, each used once; held out: 10-28 00Z/12Z and 10-30 00Z/12Z, a later
+week. The held-out
+loss is the training loss, evaluated before and after with the same two noise keys per example,
+so both evaluations see the same ensemble noise. Targets: the IFS analysis at +6 h, and the 0–6 h
+precipitation of the IFS run; the cyclone channels have no targets and are masked.
+
+| run | held-out loss change (4 inits) | per variable (mean of the 4) |
+|---|---|---|
+| lr 3e-5, no warmup | +6.6 … +7.1 % | 2 m T −8.9 %, precip −11.2 %, SST −15.4 %; pressure-level T/u/v/q/w +12 … +17 % |
+| lr 3e-5, no warmup, SST masked | +10.7 … +11.7 % | 2 m T −8.9 %, precip −11.4 %; pressure-level T/u/v/q/w +13 … +18 % |
+| **lr 3e-6, 10-step warmup, SST masked** | **−1.5 … −1.6 %** | 2 m T −8.4 %, precip −4.5 %, T −2.6 %, z −2.1 %, MSLP −1.7 %; 13 of 14 variables lower, v wind ±0.0 % |
+
+At 3e-5 the gradient norm grew from ~5 to ~13 within 50 steps and every pressure-level variable got
+worse, while the variables where IFS open data differs most from the model's training data
+(2 m temperature, precipitation, and SST, which open data replaces with skin temperature) improved
+on held-out days. At 3e-6 with warmup the gradient norm fell (3.7 → 2.9) and the held-out loss
+dropped on all four days. SST is masked in the last two runs because its open-data stand-in has a
+diurnal cycle real SST does not; masking it is not what fixed the first run (the learning rate was).
+The saved checkpoint (first run) loads with the stock `weathernext` loader and differs from the
+base checkpoint only in head modules: 15 of the 18. The other 3 are the mesh→grid GNN's mesh-node
+update (1.2M of the 5.6M head parameters), whose output the model discards, so they get no gradient.
+
+What this shows: the pipeline runs at 0.25° on a 32 GB GPU and lowers a held-out one-step loss.
+What it does not show: forecast skill. This is one checkpoint, 60 steps, one week of held-out
+6-hour steps, and the loss of a 2-member sample; rollouts, ensemble spread and lead times beyond
+6 h were not evaluated.
+Files: [`results/finetune/finetune_smoke_rtx5090.json`](results/finetune/finetune_smoke_rtx5090.json),
+[`finetune_masksst_rtx5090.json`](results/finetune/finetune_masksst_rtx5090.json),
+[`finetune_lr3e-6_rtx5090.json`](results/finetune/finetune_lr3e-6_rtx5090.json) (+ `.log`; per-variable
+losses before and after for each held-out init).
+
+## 9. Attention backward and full fine-tuning
+
+The fused attention kernel is differentiable (`pallas_attn.py`, `jax.custom_vjp`): the forward also
+stores each row's log-sum-exp, and two tiled kernels compute dQ (per 64-query block, over the
+forward's tile list) and dK/dV (per 64-key block, over the transposed mask's tile list, which for
+the symmetric k-hop mask is the forward's). With `remat`, each transformer layer is
+rematerialised, so the backward pass keeps one [nodes, 768] input per layer instead of q/k/v/output.
+Without differentiation the forward kernel is the same as before. Inference re-checked as in §2 on
+the final code: `highest` max rel RMS 1.97e-5; default 6.54e-3 in that run (worst field the sparse
+`cyclone_exists_gaussian_unit_mode`; vertical velocity 4.9e-3 as usual) and 4.74e-3 on a repeat,
+i.e. TF32 run-to-run variation from timing-based kernel selection (§4 note); temp memory and speed
+unchanged ([`results/finetune/equivalence_rtx5090_attention_bwd.json`](results/finetune/equivalence_rtx5090_attention_bwd.json),
+[`equivalence_rtx5090_attention_bwd_default_repeat.json`](results/finetune/equivalence_rtx5090_attention_bwd_default_repeat.json)).
+
+**Kernel vs references.** `tests/test_pallas_kernel.py`: dQ, dK, dV against an analytic float64
+masked-softmax gradient (n = 300, so padding rows are exercised; symmetric and non-symmetric masks;
+batch 1 and 2; both `bwd_split` settings): the kernel test file passes on the RTX 5090 (21 tests,
+forward and backward; 11 of them also run on CPUs in Pallas interpret mode). One layer of
+WeatherNext 2's exact 32-hop mask (40,962 nodes, 6 heads × 128, Hilbert order;
+`scripts/attn_bench.py --grad`) against JAX autodiff of the XLA attention path:
+
+| precision | dQ rel RMS | dK rel RMS | dV rel RMS | output rel RMS | Pallas fwd / fwd+bwd | XLA fwd / fwd+bwd |
+|---|---|---|---|---|---|---|
+| highest | 1.1e-6 | 1.1e-6 | 8.8e-7 | 8.7e-7 | 31.0 / 115.4 ms | 128.4 / 258.5 ms |
+| default (TF32) | 3.0e-4 | 2.9e-4 | 1.8e-5 | 2.8e-4 | 9.3 / 42.2 ms | 98.0 / 228.8 ms |
+
+Gradients agree to the same level as the outputs. All attention kernels use one dot precision (the
+backward recomputes P from the forward's log-sum-exp). Under `highest` that is 3xTF32, which keeps
+hi/lo copies of each operand in shared memory: the backward kernels (3–4 dots per tile) asked for
+147 KB of sm_120's 99 KB in 64-row programs, so they run in 32-row programs (`bwd_split=2`). TF32
+keeps 64-row programs (43.7 ms with 32 rows); IEEE fp32 dots (`strict_fp32=True`) take 797 ms.
+Files: [`results/attn_bwd/`](results/attn_bwd/). The kernel tests also pass on an H200 (sm_90; the
+kernel's first run on Hopper): [`results/finetune/h200/grads01/pytest_h200.txt`](results/finetune/h200/grads01/pytest_h200.txt).
+
+**Whole-model gradients vs the official path (Mini, 1°, all 56.7M parameters;
+`verify_grads.py --modules all`).** The official path cannot be differentiated end to end on
+a 32 GB GPU (compiled backward 82–116 GiB with the fair CRPS on the RTX 5090). On an H200 (141 GB;
+`cloud/modal_grads.py`, Modal, ~$4 for the three runs below) it needs 102 GiB (`triblockdiag_mha`)
+and 123 GiB (dense `mha`) with the CRPS, 51 and 62 GiB with a one-member loss; faster-weathernext
+needs 3.97 and 3.30 GiB (39.1 / 19.7 GiB without remat). Reference: the official GPU path
+(`triblockdiag_mha`); floor: the official dense `mha` path against it. Two losses: the fair CRPS,
+whose sign terms let rounding flip individual contributions (§8), and `--loss linear`, a fixed
+random linear functional of all outputs, which is smooth. Global relative L2 error of all gradients:
+
+| loss | precision | fwn | fwn, `strict_fp32` attention | fwn, XLA attention (autodiff) | floor: official `mha` |
+|---|---|---|---|---|---|
+| linear | highest | 1.4e-4 | 9.0e-5 | 7.6e-5 | 8.1e-5 |
+| linear | default | 4.4e-2 | — | 2.0e-2 | 4.5e-2 |
+| CRPS | highest | 4.7e-3 | 1.2e-3 | 1.2e-3 | 9.5e-4 |
+| CRPS | default | 5.2e-2 | — | 4.9e-2 | 5.5e-2 |
+
+At the training default (TF32), faster-weathernext's whole-model gradients differ from the official
+path's as much as the official path's two attention implementations differ from each other. At
+`highest`, IEEE fp32 attention (`strict_fp32=True`) is at the floor; the default 3xTF32 attention,
+whose forward differs slightly from fp32, is at 1.7× the floor with the smooth loss and 5× with the
+CRPS (cosine 0.99998880), whose sign terms amplify forward rounding. The highest rows are from the
+final code (run `grads03`); the default rows from `grads01` (that path has not changed since).
+Comparisons fluctuate between runs (XLA autotuning): over the three H200 runs the floor ranged
+7.5e-4–9.5e-4 (CRPS) and 7.6e-5–1.3e-4 (linear), and the XLA-attention row once reached 4.3e-3 with
+the CRPS. In `grads01` the backward used IEEE dots after a 3xTF32 forward (fwn 2.4e-4 linear,
+4.8e-3 CRPS); in `grads02` the training forward used IEEE dots too (6.8e-5, 1.4e-3, at the floor
+but 797 instead of 115 ms per layer); the final version uses 3xTF32 throughout. What sets the
+distance to the official path at `highest` is the attention forward's precision, not whether
+forward and backward match. On the RTX 5090 the same comparisons against the XLA-attention path,
+before those changes, gave the same picture at default precision (4.4e-2 CRPS, 2.9e-2 linear).
+Files: [`results/finetune/h200/`](results/finetune/h200/) (`grads01`–`grads03`; JSON + log),
+[`results/finetune/full_grads_mini_*_rtx5090.json`](results/finetune/).
+
+**Ensemble members one at a time.** At 0.25° the full-model step compiled to 45.8 GiB temp with
+FGN's two members as one batch ([`finetune_full_batched_oom_rtx5090.log`](results/finetune/finetune_full_batched_oom_rtx5090.log)).
+One rematerialised call per member did not help (46.1 GiB,
+[`finetune_full_unrolled_oom_rtx5090.log`](results/finetune/finetune_full_unrolled_oom_rtx5090.log)):
+the calls are independent, so XLA ran both backward passes side by side. Running the members in an
+`hk.scan` with a rematerialised body forces one at a time: 25.5 GiB. The noise is drawn for the
+stacked batch as before, so the loss is unchanged; checked on the CPU (Mini, all parameters, fp32)
+with the CRPS's absolute differences replaced by squares to make the loss smooth: loss rel diff
+8.4e-8, gradients 1.6e-4 global rel L2 vs the batched call, temp 4.14 → 2.55 GiB.
+For the head alone at 0.25° it brings the step from 20.03 to 12.27 GiB temp (+1.79 GiB arguments;
+4.03 instead of 3.84 s per step). With the JAX pool on the RTX 5090 capped at what a 24 GB card
+gets at `XLA_PYTHON_CLIENT_MEM_FRACTION=0.9` (fraction 0.678), 10 head-only training steps ran with
+a device peak of 14.06 GiB; capped at a 16 GB card's pool (14.39 GiB) they also ran, with 0.2 GiB
+to spare (14.16 GiB). No real 24 GB or 16 GB card has been tried. `scripts/finetune.py` runs the
+members one at a time by default (`--sequential-samples on`).
+Files: [`results/finetune/head_sequential_*`](results/finetune/).
+
+**WeatherNext 2 at 0.25°, every parameter** (`finetune.py --modules all`; same data and
+held-out days as §8; SST masked; lr 1e-6 with a 5-step warmup; 20 steps). Compiled step: 25.46 GiB
+temp + 3.12 GiB arguments (parameters, Adam state, one example), device peak 28.6 GiB, so it needs
+`XLA_PYTHON_CLIENT_MEM_FRACTION=0.95` on a 32 GB card (at 0.9 the step failed to allocate:
+[`finetune_full_scan_pool90_oom_rtx5090.log`](results/finetune/finetune_full_scan_pool90_oom_rtx5090.log)).
+8.51 s per step. Held-out loss −1.8 … −2.4 % on the four days (same noise before and after); every
+variable lower: 2 m T −9.3 %, z −6.1 %, precipitation −5.3 %, MSLP −4.6 %, T −3.9 %, the rest
+−0.4 … −1.5 % (SST, masked, −0.1 %). The same caveats as §8 apply: one checkpoint, 20 steps, one
+week of held-out one-step losses, no rollout or spread evaluation.
+Files: [`results/finetune/finetune_full_smoke_rtx5090.json`](results/finetune/finetune_full_smoke_rtx5090.json) (+ `.log`).
+
## Limitations
- Single-step equivalence (§1, §2) used one initialisation (2024-10-07) and the
@@ -236,6 +425,10 @@ Files: [`results/intel_b580/`](results/intel_b580/) (`fwn_info.txt`, `attn_bench
- Without a fixed autotune cache, results are not bitwise reproducible between runs (§4 note).
- TPU (`splash_mha`) numerics were not compared; TPU matmuls use bf16 passes by default and
differ from any of the GPU paths above.
+- Whole-model gradients were compared with the official path on the Mini model at 1° (§9, on an
+ H200: the official backward does not fit a 32 GB GPU). At 0.25° the official backward does not
+ fit any single GPU, so the 0.25° training runs rest on the Mini comparison, the kernel checks
+ against a float64 reference and the real-mask comparison with JAX autodiff.
## README figures
diff --git a/pyproject.toml b/pyproject.toml
index e48f228..e305fd2 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -32,6 +32,8 @@ cuda13 = ["jax[cuda13]>=0.11.2,<0.12"]
oneapi = ["jax-oneapi-plugin==0.11.2", "jax-oneapi-pjrt==0.11.2"]
# The official code at the validated commit (2026-09-04, version 0.3.1.dev).
weathernext = ["weathernext @ git+https://github.com/google-deepmind/weathernext.git@f2f2c5117d2f864e2d5e7c2f9f220db5e1049dd3"]
+# Head-only fine-tuning (faster_weathernext.finetune, prototype).
+train = ["optax>=0.2.8"]
test = ["pytest"]
[project.scripts]
diff --git a/scripts/attn_bench.py b/scripts/attn_bench.py
index 8c29c9e..dbdf952 100644
--- a/scripts/attn_bench.py
+++ b/scripts/attn_bench.py
@@ -6,6 +6,8 @@
hilbert fused tiled kernel, Hilbert-ordered nodes (the default on Ampere+)
python scripts/attn_bench.py --variants xla,pallas,hilbert --precision default
+ python scripts/attn_bench.py --variants xla,hilbert --grad --bwd-warps 4,8 # backward: dQ/dK/dV vs
+ JAX autodiff of the XLA path (chunks rematerialised), and forward+backward time
"""
import argparse
import time
@@ -37,6 +39,10 @@ def main():
p.add_argument("--precision", default="default")
p.add_argument("--iters", type=int, default=10)
p.add_argument("--chunk", type=int, default=512)
+ p.add_argument("--grad", action="store_true", help="also time forward+backward and compare dQ, dK, dV")
+ p.add_argument("--bwd-warps", default=str(pallas_attn.BWD_NUM_WARPS),
+ help="hilbert variant: comma-separated backward-kernel warp counts to try")
+ p.add_argument("--bwd-split", default="auto", help="hilbert variant: comma-separated bwd_split values (auto, 1, 2)")
args = p.parse_args()
t0 = time.time()
@@ -66,7 +72,7 @@ def one(a):
w = jax.nn.softmax(jnp.where(m[None, None], logits, -1e30), axis=-1)
return jnp.einsum("bhqk,bkhd->bqhd", w, vc)
- out = jax.lax.map(one, (qc, starts, bits))
+ out = jax.lax.map(jax.checkpoint(one, prevent_cse=False) if args.grad else one, (qc, starts, bits))
return out.transpose(1, 0, 2, 3, 4).reshape(1, nc * c, H, D)[:, :n]
fns["xla"] = xla
if "pallas" in args.variants:
@@ -79,28 +85,48 @@ def one(a):
htiles = pallas_attn.build_tile_mask(mask[perm][:, perm])
print(f"hilbert tiles: {len(htiles[1]) * pallas_attn.BLOCK_K / (len(htiles[0]) - 1):.0f} keys/query-block")
- def hilbert(q, k, v):
+ def hilbert(q, k, v, warps, split):
pq, pk, pv = (x[:, perm] for x in (q, k, v))
- return pallas_attn.tiled_flash_attention(pq, pk, pv, *htiles, scale=scale)[:, inv]
- fns["hilbert"] = hilbert
-
- outs = {}
+ return pallas_attn.tiled_flash_attention(pq, pk, pv, *htiles, scale=scale, bwd_num_warps=warps,
+ bwd_split=split)[:, inv]
+ combos = [(int(w), None if s == "auto" else int(s))
+ for w in args.bwd_warps.split(",") for s in args.bwd_split.split(",")]
+ for i, (w, sp) in enumerate(combos):
+ name = "hilbert" if i == 0 else f"hilbert_w{w}" + ("" if sp is None else f"s{sp}")
+ fns[name] = lambda q, k, v, w=w, sp=sp: hilbert(q, k, v, w, sp)
+
+ def timed(fn, *a):
+ out = jax.block_until_ready(fn(*a))
+ ts = []
+ for _ in range(args.iters):
+ t0 = time.time()
+ jax.block_until_ready(fn(*a))
+ ts.append(time.time() - t0)
+ return out, np.median(ts)
+
+ names = [n for n in fns if n.split("_w")[0] in args.variants.split(",")]
+ do = jax.random.normal(jax.random.PRNGKey(1), q.shape, jnp.float32)
+ outs, grads = {}, {}
with jax.default_matmul_precision(args.precision):
- for name in args.variants.split(","):
- fn = jax.jit(fns[name])
- out = jax.block_until_ready(fn(q, k, v))
- ts = []
- for _ in range(args.iters):
- t0 = time.time()
- jax.block_until_ready(fn(q, k, v))
- ts.append(time.time() - t0)
+ for name in names:
+ out, t = timed(jax.jit(fns[name]), q, k, v)
outs[name] = np.asarray(out, np.float64)
- print(f"{name:8s} {np.median(ts) * 1e3:7.1f} ms/layer (x24 = {24 * np.median(ts):.2f} s/step)", flush=True)
- ref_name = args.variants.split(",")[0]
+ line = f"{name:10s} {t * 1e3:7.1f} ms/layer (x24 = {24 * t:.2f} s/step)"
+ if args.grad:
+ bwd = jax.jit(lambda q, k, v, do, f=fns[name]: jax.vjp(f, q, k, v)[1](do))
+ g, tg = timed(bwd, q, k, v, do)
+ grads[name] = [np.asarray(x, np.float64) for x in g]
+ line += f" forward+backward {tg * 1e3:7.1f} ms/layer (x24 = {24 * tg:.2f} s)"
+ print(line, flush=True)
+ ref_name = names[0]
+ rel = lambda o, r: np.sqrt(np.mean((o - r) ** 2)) / r.std()
for name, o in outs.items():
if name != ref_name:
- r = outs[ref_name]
- print(f"{name} vs {ref_name}: rel RMS {np.sqrt(np.mean((o - r) ** 2)) / r.std():.2e}")
+ print(f"{name} vs {ref_name}: output rel RMS {rel(o, outs[ref_name]):.2e}", end="")
+ if args.grad:
+ print(" " + " ".join(f"d{x} rel RMS {rel(a, b):.2e}" for x, a, b in zip("qkv", grads[name], grads[ref_name])),
+ end="")
+ print(flush=True)
if __name__ == "__main__":
diff --git a/scripts/finetune.py b/scripts/finetune.py
new file mode 100644
index 0000000..871c47e
--- /dev/null
+++ b/scripts/finetune.py
@@ -0,0 +1,183 @@
+"""Fine-tuning of WeatherNext 2 at 0.25° on IFS open data (prototype).
+
+Trains the head (mesh->grid GNN and grid decoder, finetune.HEAD_MODULES; `--modules all` trains
+every parameter, through the attention kernel's backward) of one checkpoint with the official
+training loss (fair CRPS over 2 members) on one-step examples from IFS open-data analyses
+(ifs.training_example: inputs at T-6h and T; targets the analysis at T+6h and the 0-6 h
+precipitation of the T run; the cyclone channels have no targets and are masked). One example per
+optimizer step; the next example is built on the CPU while the GPU trains on the current one.
+
+Before and after training the loss is evaluated on held-out init times with fixed noise keys, so
+both evaluations see the same ensemble noise. This checks that the pipeline learns; it is not a
+skill evaluation.
+
+ python scripts/finetune.py --train 2026101912 2026101918 2026102000 --eval 2026103012 \
+ --steps 30 --mask sea_surface_temperature --out head.npz --log finetune_run.json
+ python scripts/finetune.py --train 2026102200 --compile-only # memory of one training step
+ XLA_PYTHON_CLIENT_MEM_FRACTION=0.95 python scripts/finetune.py --modules all ...
+ # every parameter: a 25.5 GiB temp buffer + 3.1 GiB of arguments, too much for the 0.9 pool
+"""
+import argparse
+import concurrent.futures
+import datetime as dt
+import json
+import os
+import time
+
+# One preallocated pool: the 0.25° training step needs a single ~20 GiB temp buffer. In a pool that
+# grows on demand, the regions taken earlier by the evaluation pass are kept but too small, and the
+# step then fails to allocate (RESOURCE_EXHAUSTED) although enough memory is free in total.
+os.environ.setdefault("XLA_PYTHON_CLIENT_PREALLOCATE", "true")
+os.environ.setdefault("XLA_PYTHON_CLIENT_MEM_FRACTION", "0.9")
+
+import numpy as np # noqa: E402
+
+
+def parse_init(s):
+ return dt.datetime.strptime(s, "%Y%m%d%H")
+
+
+def main():
+ p = argparse.ArgumentParser()
+ p.add_argument("--model", default="WeatherNext2")
+ p.add_argument("--checkpoint", type=int, default=1)
+ p.add_argument("--train", nargs="+", required=True, help="init times YYYYMMDDHH (inputs T-6h, T; target T+6h)")
+ p.add_argument("--eval", nargs="*", default=[], help="held-out init times")
+ p.add_argument("--steps", type=int, default=10)
+ # 3e-5 without warmup degraded held-out loss by 7-11 % within 60 steps (docs/validation.md §8).
+ p.add_argument("--lr", type=float, default=3e-6)
+ p.add_argument("--warmup", type=int, default=10, help="steps of linear learning-rate warmup from 0")
+ p.add_argument("--weight-decay", type=float, default=0.0)
+ p.add_argument("--clip", type=float, default=1.0, help="global gradient-norm clip")
+ p.add_argument("--eval-keys", type=int, default=2, help="noise keys per eval example")
+ p.add_argument("--precision", default="default")
+ p.add_argument("--seed", type=int, default=0)
+ p.add_argument("--no-precip", action="store_true", help="mask precipitation instead of fetching its target")
+ p.add_argument("--modules", nargs="+", default=None,
+ help='module names to train (default: the head, finetune.HEAD_MODULES); "all" trains every '
+ "parameter, through the attention backward")
+ p.add_argument("--sequential-samples", choices=["on", "off"], default="on",
+ help="run the 2 ensemble members one at a time: same loss, about half the activation "
+ "memory (0.25°: head 20.0 -> 12.3 GiB temp, all parameters 45.8 -> 25.5 GiB), ~5 %% slower")
+ p.add_argument("--mask", nargs="*", default=[],
+ help="extra target variables to leave out of the loss, e.g. sea_surface_temperature: open data "
+ "has no SST, and its stand-in (skin temperature) has a diurnal cycle real SST lacks")
+ p.add_argument("--compile-only", action="store_true")
+ p.add_argument("--no-remat", action="store_true", help="diagnostics: compile without block remat")
+ p.add_argument("--weights-dir")
+ p.add_argument("--out", help="write the fine-tuned checkpoint here (official format)")
+ p.add_argument("--log", default="finetune_run.json")
+ args = p.parse_args()
+
+ import jax
+ import optax
+ import faster_weathernext
+ from faster_weathernext import finetune, ifs, model
+
+ faster_weathernext.enable(remat=not args.no_remat, sequential_samples=args.sequential_samples == "on")
+ task, _ = model.build(args.model)
+ load = lambda s: ifs.training_example(parse_init(s), task, precip=not args.no_precip)
+ t0 = time.time()
+ first = load(args.train[0])
+ print(f"example {args.train[0]} built in {time.time() - t0:.1f}s", flush=True)
+ masked = sorted(set(finetune.all_nan_vars(first[1])) | set(args.mask))
+ _, loss_fn = finetune.build_loss(args.model, masked=masked)
+ params = model.load_params(args.model, args.checkpoint, args.weights_dir)
+ modules = finetune.HEAD_MODULES if not args.modules else "all" if args.modules == ["all"] else tuple(args.modules)
+ head, frozen = finetune.split(params, modules)
+ head = jax.device_put(head)
+ frozen = jax.device_put(frozen)
+ lr = (optax.join_schedules([optax.linear_schedule(0.0, args.lr, args.warmup), optax.constant_schedule(args.lr)],
+ [args.warmup]) if args.warmup else args.lr)
+ optimizer = optax.chain(optax.clip_by_global_norm(args.clip),
+ optax.adamw(lr, weight_decay=args.weight_decay))
+ opt_state = optimizer.init(head)
+ dev = jax.devices()[0]
+ log = dict(model=args.model, checkpoint=args.checkpoint, device=dev.device_kind, jax=jax.__version__,
+ faster_weathernext=faster_weathernext.__version__, args=vars(args), masked=masked,
+ modules=modules, head_params=finetune.head_size(params, modules),
+ total_params=int(sum(np.size(v) for v in jax.tree.leaves(params))))
+ print(f"{args.model} ckpt {args.checkpoint}: training {log['head_params'] / 1e6:.2f}M of "
+ f"{log['total_params'] / 1e6:.1f}M params; masked targets: {len(masked)} "
+ f"({', '.join(['cyclone channels'] + [v for v in masked if not v.startswith('cyclone')])})", flush=True)
+
+ key = jax.random.PRNGKey(args.seed)
+ step_fn = finetune.train_step(loss_fn, optimizer)
+ with jax.default_matmul_precision(args.precision):
+ t0 = time.time()
+ compiled = step_fn.lower(head, frozen, opt_state, key, *first).compile()
+ mem = compiled.memory_analysis()
+ log["compile"] = dict(seconds=time.time() - t0, temp_gib=mem.temp_size_in_bytes / 2**30,
+ argument_gib=mem.argument_size_in_bytes / 2**30,
+ output_gib=mem.output_size_in_bytes / 2**30)
+ print(f"train step compiled in {log['compile']['seconds']:.0f}s: temp {log['compile']['temp_gib']:.2f} GiB, "
+ f"arguments {log['compile']['argument_gib']:.2f} GiB", flush=True)
+ if args.compile_only:
+ with open(args.log, "w") as f:
+ json.dump(log, f, indent=1)
+ print("wrote", args.log)
+ return
+
+ # frozen is an argument, not a closure: closed-over arrays can become executable constants.
+ eval_fn = jax.jit(lambda h, fz, rng, i, t, f: loss_fn(finetune.merge(fz, h), rng, i, t, f))
+ eval_examples = {s: load(s) for s in args.eval}
+ eval_keys = [jax.random.PRNGKey(1000 + i) for i in range(args.eval_keys)]
+
+ def evaluate(h):
+ out = {}
+ for s, ex in eval_examples.items():
+ vals = [jax.device_get(eval_fn(h, frozen, k, *ex)) for k in eval_keys]
+ out[s] = dict(loss=float(np.mean([v[0] for v in vals])),
+ per_var={k: float(np.mean([v[1][k] for v in vals])) for k in vals[0][1]})
+ return out
+
+ if eval_examples:
+ t0 = time.time()
+ log["eval_before"] = evaluate(head)
+ print(f"eval before ({time.time() - t0:.0f}s): "
+ + " ".join(f"{s} {v['loss']:.5f}" for s, v in log["eval_before"].items()), flush=True)
+
+ # Epochs over the training inits, each in a fresh random order.
+ rng = np.random.default_rng(args.seed)
+ names = [args.train[i] for _ in range(-(-args.steps // len(args.train)))
+ for i in rng.permutation(len(args.train))][:args.steps]
+ pool = concurrent.futures.ThreadPoolExecutor(1)
+ ex = first if names[0] == args.train[0] else load(names[0])
+ log["steps"] = []
+ for step, name in enumerate(names):
+ upcoming = pool.submit(load, names[step + 1]) if step + 1 < len(names) else None
+ t0 = time.time()
+ head, opt_state, metrics = compiled(head, frozen, opt_state, jax.random.fold_in(key, step), *ex)
+ metrics = jax.device_get(metrics)
+ dt_step = time.time() - t0
+ log["steps"].append(dict(step=step, init=name, loss=float(metrics["loss"]),
+ grad_norm=float(metrics["grad_norm"]), seconds=dt_step))
+ print(f"step {step:4d} {name} loss {float(metrics['loss']):.5f} |g| {float(metrics['grad_norm']):.3e} "
+ f"{dt_step:.2f}s", flush=True)
+ if upcoming is not None:
+ ex = upcoming.result()
+ stats = dev.memory_stats() or {}
+ log["peak_gib"] = stats.get("peak_bytes_in_use", 0) / 2**30
+ times = [s["seconds"] for s in log["steps"][1:]] or [log["steps"][0]["seconds"]]
+ log["median_step_s"] = float(np.median(times))
+ print(f"median step {log['median_step_s']:.2f}s, device peak {log['peak_gib']:.2f} GiB", flush=True)
+
+ if eval_examples:
+ log["eval_after"] = evaluate(head)
+ for s in eval_examples:
+ b, a = log["eval_before"][s]["loss"], log["eval_after"][s]["loss"]
+ print(f"eval {s}: {b:.5f} -> {a:.5f} ({100 * (a - b) / b:+.2f} %)", flush=True)
+
+ if args.out:
+ finetune.save(args.out, finetune.merge(frozen, head), model=args.model, checkpoint=args.checkpoint,
+ weights_dir=args.weights_dir,
+ note=f"Head-only fine-tuned with faster-weathernext {faster_weathernext.__version__}: "
+ f"{args.steps} steps, lr {args.lr}, IFS open data {args.train[0]}..{args.train[-1]}.")
+ print("wrote", args.out)
+ with open(args.log, "w") as f:
+ json.dump(log, f, indent=1)
+ print("wrote", args.log)
+
+
+if __name__ == "__main__":
+ main()
diff --git a/scripts/verify_grads.py b/scripts/verify_grads.py
new file mode 100644
index 0000000..ed00451
--- /dev/null
+++ b/scripts/verify_grads.py
@@ -0,0 +1,204 @@
+"""Fine-tuning gradients through faster-weathernext vs the official GPU path.
+
+Mini model (1°) on Google's sample data; same checkpoint, inputs, targets and noise (same PRNG key).
+The loss is the official training loss (fair CRPS over 2 members, finetune.build_loss) and the
+gradients are taken w.r.t. the head (mesh->grid GNN + grid decoder; default) or, with
+`--modules all`, every parameter (through the fused attention kernel's backward), at both matmul
+precisions. Paths:
+
+ official patch disabled, attention_type="triblockdiag_mha" (the official GPU path)
+ official_mha patch disabled, attention_type="mha" (the official dense masked attention): a second
+ official implementation, so official_mha vs official is the floor for differences that
+ come only from rounding. The fair CRPS has sign terms (|pred - target|, |pred_i - pred_j|),
+ so rounding-level forward differences flip a few gradient contributions; gradient
+ agreement is judged against this floor, not against 0.
+ fwn patch enabled, remat=True, small blocks (grid_chunk 8192, edge_chunk 16384) so the
+ 65,160-point 1° grid runs through several encoder / decoder / GNN blocks (with the
+ default 32,768-point blocks the encoder and decoder would take the stock path)
+ fwn_xla fwn with the XLA attention path, whose backward is JAX autodiff (with --modules all:
+ an independent reference for the Pallas backward kernels)
+ fwn_noremat fwn without remat, compiled only: what rematerialisation saves
+ fwn_strict (highest only) fwn with IEEE fp32 attention dots instead of the kernel's 3xTF32
+
+`--loss linear` replaces the CRPS by a smooth one: a fixed random linear functional of all output
+fields (one member, no sample dimension). With no sign terms, implementation differences show at
+rounding level; use it where no second official implementation fits as a floor (--modules all).
+
+A case that runs out of GPU memory is recorded as such (the official path's full backward may not
+fit); the comparisons then use the official path if it ran, else fwn_xla, as the reference.
+
+ python scripts/verify_grads.py --data ".../source-hres_forecast_init-2024-10-07 00:00:00_res-1.0_levels-13_steps-04.nc"
+ python scripts/verify_grads.py --modules all --data ... --out full_grads_mini.json
+"""
+import argparse
+import dataclasses
+import gc
+import json
+import os
+import time
+
+os.environ.setdefault("XLA_PYTHON_CLIENT_PREALLOCATE", "false")
+
+import numpy as np # noqa: E402
+
+MODEL = "WeatherNextCyclones_Mini"
+BLOCKS = dict(grid_chunk=8192, edge_chunk=16384)
+
+
+def example(path, task):
+ import xarray
+ from weathernext.utils import data_utils
+ ds = xarray.load_dataset(path).isel(time=slice(0, 3))
+ ds = ds.assign_coords(time=ds.time - ds.time[0])
+ parts = data_utils.extract_inputs_targets_forcings(ds, target_lead_times=slice("6h", "6h"),
+ **dataclasses.asdict(task))
+ return tuple(x.drop_vars("datetime", errors="ignore").compute() for x in parts)
+
+
+def linear_loss(attention_type, weights):
+ """loss_fn with the CRPS replaced by sum_v mean(prediction_v * weights_v) (one member)."""
+ import jax.numpy as jnp
+ import xarray_jax
+ from faster_weathernext import model
+ _, forward = model.build(MODEL, attention_type=attention_type)
+
+ def loss_fn(params, rng, inputs, targets, forcings):
+ preds = forward(params, rng, inputs, targets, forcings)
+ terms = {v: jnp.mean(jnp.nan_to_num(xarray_jax.unwrap_data(preds[v])) * w) for v, w in weights.items()}
+ return sum(terms.values()), terms
+
+ return loss_fn
+
+
+def run_case(params, ex, masked, precision, path, modules, compile_only=False, weights=None):
+ import jax
+ import faster_weathernext
+ from faster_weathernext import finetune
+ if path.startswith("official"):
+ faster_weathernext.disable()
+ attention_type = "mha" if path == "official_mha" else "triblockdiag_mha"
+ else:
+ faster_weathernext.enable(remat=(path != "fwn_noremat"), strict_fp32=(path == "fwn_strict"),
+ attention="xla" if path == "fwn_xla" else "auto", **BLOCKS)
+ attention_type = "chunked_mha"
+ if weights is None:
+ _, loss_fn = finetune.build_loss(MODEL, masked=masked, attention_type=attention_type)
+ else:
+ loss_fn = linear_loss(attention_type, weights)
+ head, frozen = finetune.split(params, modules)
+ fn = finetune.loss_and_grads(loss_fn)
+ call = (head, frozen, jax.random.PRNGKey(0), *ex)
+ out, info = None, {}
+ try:
+ with jax.default_matmul_precision(precision):
+ t0 = time.time()
+ compiled = fn.lower(*call).compile()
+ info = dict(temp_gib=compiled.memory_analysis().temp_size_in_bytes / 2**30, compile_s=time.time() - t0)
+ if not compile_only:
+ jax.block_until_ready(compiled(*call))
+ t0 = time.time()
+ out = jax.device_get(jax.block_until_ready(compiled(*call)))
+ info["step_s"] = time.time() - t0
+ del compiled
+ except jax.errors.JaxRuntimeError as e:
+ if "RESOURCE_EXHAUSTED" not in str(e):
+ raise
+ info["oom"] = str(e).splitlines()[0][:300]
+ out = None
+ gc.collect()
+ jax.clear_caches()
+ temp = f"temp {info['temp_gib']:5.2f} GiB" if "temp_gib" in info else "temp n/a"
+ status = (" OUT OF MEMORY" if "oom" in info else " (compile only)" if compile_only
+ else f" step {info['step_s']:.2f}s loss {float(out[0][0]):.6f}")
+ print(f" {path:12s} {precision:8s} {temp}{status}", flush=True)
+ return out, info
+
+
+def compare(got, ref):
+ """Head gradients `got` vs `ref` (same tree): global and per-parameter relative L2 error."""
+ rows, num, den, dot, n_got = [], 0.0, 0.0, 0.0, 0.0
+ for m in sorted(ref):
+ for k in sorted(ref[m]):
+ a = np.asarray(got[m][k], np.float64).ravel()
+ b = np.asarray(ref[m][k], np.float64).ravel()
+ nb, nd = float(np.linalg.norm(b)), float(np.linalg.norm(a - b))
+ num, den, dot, n_got = num + nd**2, den + nb**2, dot + float(a @ b), n_got + float(a @ a)
+ rows.append(dict(param=f"{m}/{k}", size=a.size, ref_norm=nb, got_norm=float(np.linalg.norm(a)),
+ finite=bool(np.isfinite(a).all()),
+ rel_l2=nd / nb if nb > 0 else (0.0 if nd == 0 else float("inf"))))
+ # Per-parameter errors are reported for tensors whose gradient is not negligible; tiny ones
+ # (e.g. output heads of variables with near-zero loss weight) only enter the global numbers.
+ big = [r for r in rows if r["ref_norm"] > 1e-3 * max(r["ref_norm"] for r in rows)]
+ return dict(global_rel_l2=np.sqrt(num / den), global_cos=dot / np.sqrt(den * n_got),
+ max_rel_l2=max(r["rel_l2"] for r in big), n_params=len(rows), n_compared=len(big),
+ n_zero_both=sum(r["ref_norm"] == 0 and r["got_norm"] == 0 for r in rows),
+ n_zero_ref_only=sum(r["ref_norm"] == 0 and r["got_norm"] > 0 for r in rows),
+ all_finite=all(r["finite"] for r in rows),
+ worst=sorted(big, key=lambda r: -r["rel_l2"])[:5])
+
+
+def main():
+ p = argparse.ArgumentParser()
+ p.add_argument("--data", required=True, help="local path of Google's 1° steps-04 sample")
+ p.add_argument("--precisions", default="highest,default")
+ p.add_argument("--modules", nargs="+", default=None, help='modules to differentiate (default: head); "all"')
+ p.add_argument("--loss", choices=["crps", "linear"], default="crps")
+ p.add_argument("--weights-dir")
+ p.add_argument("--out", default="grads_mini.json")
+ args = p.parse_args()
+
+ import jax
+ import faster_weathernext
+ from faster_weathernext import finetune, model
+
+ modules = finetune.HEAD_MODULES if not args.modules else "all" if args.modules == ["all"] else tuple(args.modules)
+ task, _ = finetune.build_loss(MODEL)
+ params = model.load_params(MODEL, 1, args.weights_dir)
+ ex = example(args.data, task)
+ masked = finetune.all_nan_vars(ex[1])
+ weights = None
+ if args.loss == "linear": # fixed random weights per output field (NaN predictions count as 0)
+ rng = np.random.default_rng(0)
+ weights = {v: rng.standard_normal(da.shape[:1] + da.shape[2:]).astype(np.float32)[:, None]
+ for v, da in ex[1].data_vars.items()}
+ result = dict(model=MODEL, device=jax.devices()[0].device_kind, jax=jax.__version__,
+ faster_weathernext=faster_weathernext.__version__, blocks=BLOCKS, masked=masked,
+ modules=modules, loss=args.loss, head_params=finetune.head_size(params, modules), precisions={})
+ print(f"{MODEL}: differentiating {result['head_params'] / 1e6:.2f}M params ({modules}); "
+ f"masked targets: {masked}", flush=True)
+ for precision in args.precisions.split(","):
+ print(f"[{precision}]", flush=True)
+ r, outs = {}, {}
+ cases = ["official", "official_mha", "fwn", "fwn_xla"] + (["fwn_strict"] if precision == "highest" else [])
+ for path in cases: # fwn_strict: the Pallas kernel's default under `highest` is 3xTF32, not IEEE fp32
+ outs[path], r[path] = run_case(params, ex, masked, precision, path, modules, weights=weights)
+ _, r["fwn_noremat"] = run_case(params, ex, masked, precision, "fwn_noremat", modules, compile_only=True,
+ weights=weights)
+ ref_name = "official" if outs["official"] is not None else "fwn_xla"
+ off_loss, off_grads = outs[ref_name]
+ pairs = [(f"{p} vs {ref_name}" + (" (floor)" if p == "official_mha" else ""), outs[p])
+ for p in ("fwn", "fwn_strict", "fwn_xla", "official_mha") if p != ref_name and outs.get(p) is not None]
+ r["reference"] = ref_name
+ r["checks"] = {}
+ for label, (loss, grads) in pairs:
+ lo, lg = float(off_loss[0]), float(loss[0])
+ per_var = {k: abs(float(loss[1][k]) - float(v)) / abs(float(v)) for k, v in off_loss[1].items() if float(v) != 0}
+ per_var = per_var or {"-": 0.0}
+ g = compare(grads, off_grads)
+ r["checks"][label] = dict(loss=dict(official=lo, other=lg, rel_diff=abs(lg - lo) / abs(lo),
+ max_per_var_rel_diff=max(per_var.values())), grads=g)
+ print(f" {label:34s} loss rel diff {abs(lg - lo) / abs(lo):.2e} (per var max {max(per_var.values()):.2e}) "
+ f"grads: global rel L2 {g['global_rel_l2']:.2e} cos {g['global_cos']:.10f} "
+ f"max per-param rel L2 {g['max_rel_l2']:.2e} ({g['n_compared']}/{g['n_params']} params; "
+ f"zero in both {g['n_zero_both']}, zero only in official {g['n_zero_ref_only']}) "
+ f"finite {g['all_finite']}", flush=True)
+ for w in g["worst"][:3]:
+ print(f" {w['rel_l2']:.2e} {w['param']}", flush=True)
+ result["precisions"][precision] = r
+ with open(args.out, "w") as f:
+ json.dump(result, f, indent=1)
+ print("wrote", args.out)
+
+
+if __name__ == "__main__":
+ main()
diff --git a/src/faster_weathernext/__init__.py b/src/faster_weathernext/__init__.py
index dd807d7..d1606a9 100644
--- a/src/faster_weathernext/__init__.py
+++ b/src/faster_weathernext/__init__.py
@@ -16,7 +16,7 @@ def enable(force: bool = False, **options):
Call before the model is first traced (e.g. before the first call of a jitted forward).
Keyword options update `faster_weathernext.patch.Options` (attention, reorder, strict_fp32,
- grid_chunk, edge_chunk, attn_chunk, layer_loop, replace_attention).
+ grid_chunk, edge_chunk, attn_chunk, layer_loop, replace_attention, remat, sequential_samples).
Returns the active Options.
"""
from faster_weathernext import compat
diff --git a/src/faster_weathernext/compat.py b/src/faster_weathernext/compat.py
index 38a71ed..8374215 100644
--- a/src/faster_weathernext/compat.py
+++ b/src/faster_weathernext/compat.py
@@ -52,6 +52,7 @@ def weathernext_version() -> str:
_REQUIRED_METHODS = [
("weathernext.utils.deep_gnn", "DeepGNN", "_networks_builder"),
("weathernext.utils.points_mesh_gnn", "PointsMeshTypedGraphGNN", "_build_typed_graph"),
+ ("weathernext.weathernext2.fgn", "Predictor", "run_model_with_sample_as_batch"),
]
diff --git a/src/faster_weathernext/finetune.py b/src/faster_weathernext/finetune.py
new file mode 100644
index 0000000..dd6e330
--- /dev/null
+++ b/src/faster_weathernext/finetune.py
@@ -0,0 +1,189 @@
+"""Fine-tuning of WeatherNext 2 (prototype).
+
+By default only the "head" is trained: everything after the mesh transformer, i.e. the mesh->grid
+GNN (`mesh_to_grid_gnn`) and the grid decoder (`grid_decoder`), ~5.6M of WeatherNext 2's 184M
+parameters. Gradients are taken with respect to the trained modules only, so JAX never
+differentiates a frozen prefix of the model: with the head alone, the encoder, grid->mesh GNN and
+transformer run as in inference. `split(params, "all")` (or any set of module names, e.g.
+("mesh_transformer", *HEAD_MODULES)) trains more; the backward pass then runs through the fused
+attention kernel's backward and needs `enable(remat=True)` to fit.
+
+The loss is the official training loss: the checkpoint config's full predictor stack (residual
+normalisation, NaN cleaning, autoregressive wrapper, two-member sample dimension) and FGN's fair
+CRPS with the config's per-variable weights. Target variables the data cannot provide (or should
+not be trained on) are masked out of the loss (`predictor_config`).
+
+ import faster_weathernext, optax
+ from faster_weathernext import finetune, ifs, model
+
+ faster_weathernext.enable(remat=True) # remat: blocked GNN/decoder rematerialised in backward
+ task, _ = model.build("WeatherNext2")
+ inputs, targets, forcings = ifs.training_example(init, task)
+ _, loss_fn = finetune.build_loss("WeatherNext2", masked=finetune.all_nan_vars(targets))
+ head, frozen = finetune.split(model.load_params("WeatherNext2", 1))
+ optimizer = optax.adamw(1e-5)
+ opt_state = optimizer.init(head)
+ step = finetune.train_step(loss_fn, optimizer)
+ head, opt_state, metrics = step(head, frozen, opt_state, rng, inputs, targets, forcings)
+ finetune.save("head.npz", finetune.merge(frozen, head), model="WeatherNext2", checkpoint=1)
+
+Fine-tuned weights are a different model: the equivalence claims in docs/validation.md cover how
+faster-weathernext executes a checkpoint, not the skill of a fine-tuned one.
+"""
+
+import copy
+import os
+
+import haiku as hk
+import jax
+import numpy as np
+
+HEAD_MODULES = ("mesh_to_grid_gnn", "grid_decoder")
+
+
+def is_head(module_name: str, modules=HEAD_MODULES) -> bool:
+ """Whether a Haiku module path (e.g. 'multimodality_forward/grid_decoder/...') is trained:
+ one of its path components is in `modules`, or modules == "all"."""
+ return modules == "all" or bool(set(module_name.split("/")) & set(modules))
+
+
+def split(params, modules=HEAD_MODULES):
+ """(trained, frozen) partition of a parameter tree by module path (see is_head)."""
+ return hk.data_structures.partition(lambda m, n, v: is_head(m, modules), params)
+
+
+def merge(*parts):
+ return hk.data_structures.merge(*parts)
+
+
+def all_nan_vars(targets) -> list:
+ """Target variables that are NaN everywhere (no data): pass them as `masked` to build_loss."""
+ return [v for v, da in targets.data_vars.items() if np.isnan(np.asarray(da.values)).all()]
+
+
+def predictor_config(model: str = "WeatherNext2", *, masked=(), num_samples: int = 2,
+ attention_type: str = "chunked_mha", mask_type: str | None = None):
+ """The checkpoint's training PredictorConfig, with `masked` target variables out of the loss.
+
+ A masked variable gets loss weight 0 and is listed in `vars_with_nan_targets`. Both are needed:
+ the official NaN mechanism only zeroes the loss where a target is NaN (it handles the all-NaN
+ cyclone channels, but not a variable with data, such as SST, whose land NaNs NaNCleaner fills
+ anyway), and a zero weight alone would give 0 * NaN = NaN for NaN targets.
+ """
+ from weathernext.utils import ensemble
+ from weathernext.utils import fiddle_config_io
+ from weathernext.weathernext2 import fgn
+
+ config = fiddle_config_io.get_fiddle_config_by_name(f"weathernext2/configs/{model}")
+ pk = copy.deepcopy(config.predictor_kwargs)
+ tk = pk["noisy_function_kwargs"]["mesh_model_ctor"].keywords["transformer_kwargs"]
+ tk["attention_type"] = attention_type
+ if mask_type is not None:
+ tk["mask_type"] = mask_type
+ unknown = set(masked) - set(config.task.target_variables)
+ if unknown:
+ raise ValueError(f"masked variables are not targets of {model}: {sorted(unknown)}")
+ loss_kwargs = pk["fgn_loss_function_kwargs"]
+ loss_kwargs["vars_with_nan_targets"] = sorted(set(loss_kwargs["vars_with_nan_targets"]) | set(masked))
+ loss_kwargs["loss_for_nan_targets"] = 0.0
+ loss_kwargs["per_variable_weights"] = {**loss_kwargs.get("per_variable_weights", {}),
+ **{v: 0.0 for v in masked}}
+ wrappers = []
+ for w in config.predictor_wrappers:
+ if w["constructor"] is ensemble.WithSampleDim:
+ w = dict(w, kwargs=dict(w.get("kwargs", {}), num_samples=num_samples))
+ wrappers.append(w)
+ return fgn.PredictorConfig(task=config.task, predictor_constructor=config.predictor_constructor,
+ predictor_kwargs=pk, predictor_wrappers=wrappers)
+
+
+def build_loss(model: str = "WeatherNext2", *, masked=(), num_samples: int = 2,
+ attention_type: str = "chunked_mha", mask_type: str | None = None):
+ """(task, loss_fn) with loss_fn(params, rng, inputs, targets, forcings) -> (loss, per-variable losses).
+
+ `loss` is the official training loss averaged over the batch, without the `masked` target
+ variables (see predictor_config; their targets may be NaN). The per-variable losses are
+ unweighted and include masked variables. `attention_type`/`mask_type` are the official
+ transformer settings, as in model.build (with the patch enabled every attention type runs through
+ faster-weathernext's kernel; the official GPU path is attention_type="triblockdiag_mha" with the
+ patch disabled).
+ """
+ import xarray_jax
+ from weathernext.weathernext2 import fgn
+
+ cfg = predictor_config(model, masked=masked, num_samples=num_samples,
+ attention_type=attention_type, mask_type=mask_type)
+
+ @hk.transform
+ def run(inputs, targets, forcings):
+ loss, diagnostics = fgn.construct_predictor(cfg).loss(inputs, targets, forcings)
+ per_var = {k: xarray_jax.unwrap_data(v).mean() for k, v in diagnostics.items()}
+ return xarray_jax.unwrap_data(loss).mean(), per_var
+
+ return cfg.task, run.apply
+
+
+def loss_and_grads(loss_fn):
+ """jit(head, frozen, rng, inputs, targets, forcings) -> ((loss, per-variable), head grads)."""
+
+ def fn(head, frozen, rng, inputs, targets, forcings):
+ return jax.value_and_grad(
+ lambda h: loss_fn(merge(frozen, h), rng, inputs, targets, forcings), has_aux=True)(head)
+
+ return jax.jit(fn)
+
+
+def train_step(loss_fn, optimizer):
+ """jit(head, frozen, opt_state, rng, inputs, targets, forcings) -> (head, opt_state, metrics).
+
+ `head` and `opt_state` are donated. metrics: loss, grad_norm and the per-variable losses.
+ """
+ import optax
+
+ def step(head, frozen, opt_state, rng, inputs, targets, forcings):
+ (loss, per_var), grads = jax.value_and_grad(
+ lambda h: loss_fn(merge(frozen, h), rng, inputs, targets, forcings), has_aux=True)(head)
+ updates, opt_state = optimizer.update(grads, opt_state, head)
+ head = optax.apply_updates(head, updates)
+ return head, opt_state, dict(loss=loss, grad_norm=optax.global_norm(grads), per_var=per_var)
+
+ return jax.jit(step, donate_argnums=(0, 2))
+
+
+def save(path: str, params, *, model: str = "WeatherNext2", checkpoint: int = 1,
+ weights_dir: str | None = None, note: str = ""):
+ """Write params as an official-format checkpoint (loads with the stock weathernext code).
+
+ Description and license are taken from the base checkpoint; `note` is appended to the
+ description. Parameter names and shapes must match the base checkpoint.
+ """
+ from weathernext.utils import checkpoint as ckpt_lib
+ from weathernext.weathernext2 import fgn
+ from faster_weathernext.model import weights_path
+
+ with open(weights_path(model, checkpoint, weights_dir), "rb") as f:
+ base = ckpt_lib.load(f, fgn.CheckPoint)
+ host = jax.device_get(params)
+ out = {m: {k: np.asarray(v) for k, v in host[m].items()} for m in host}
+ ref = {(m, k): v.shape for m, p in base.params.items() for k, v in p.items()}
+ got = {(m, k): v.shape for m, p in out.items() for k, v in p.items()}
+ if ref != got:
+ diff = sorted(set(ref.items()) ^ set(got.items()))[:5]
+ raise ValueError(f"params do not match {model} checkpoint {checkpoint}: {diff}")
+ description = base.description + (f"\n\n{note}" if note else "")
+ with open(path + ".tmp", "wb") as f:
+ ckpt_lib.dump(f, fgn.CheckPoint(params=out, description=description, license=base.license))
+ os.replace(path + ".tmp", path)
+
+
+def load(path: str):
+ """Params of a checkpoint written by `save` (or any official-format checkpoint file)."""
+ from weathernext.utils import checkpoint as ckpt_lib
+ from weathernext.weathernext2 import fgn
+ with open(path, "rb") as f:
+ return ckpt_lib.load(f, fgn.CheckPoint).params
+
+
+def head_size(params, modules=HEAD_MODULES) -> int:
+ head, _ = split(params, modules)
+ return int(sum(np.size(v) for v in jax.tree.leaves(head)))
diff --git a/src/faster_weathernext/ifs.py b/src/faster_weathernext/ifs.py
index edc468f..c990ccc 100644
--- a/src/faster_weathernext/ifs.py
+++ b/src/faster_weathernext/ifs.py
@@ -244,6 +244,29 @@ def rollout_inputs(init: dt.datetime, n_steps: int, task):
return inputs.compute(), targets, forcings
+def training_example(init: dt.datetime, task, precip: bool = True):
+ """(inputs, targets, forcings) for one 6 h training step from IFS open data.
+
+ Inputs are the analyses at init-6h and init (as in `rollout_inputs`); targets are the analysis at
+ init+6h and, with precip=True, the 0-6 h precipitation of the run started at init (the same
+ source `truth_fields` verifies against). Target fields open data does not provide (the cyclone
+ channels; precipitation with precip=False) stay NaN: mask them out of the loss
+ (`finetune.all_nan_vars`). The `datetime` coordinate is dropped so that examples from different
+ init times share one compiled training step.
+ """
+ import dataclasses
+ from weathernext.utils import data_utils
+ ds = build_example(init, 1, task.target_variables)
+ truth = analysis_state(init + dt.timedelta(hours=6))
+ for name in list(PL_VARS) + list(SFC_VARS):
+ ds[name].values[0, 2] = truth[name]
+ if precip and "total_precipitation_6hr" in ds:
+ ds["total_precipitation_6hr"].values[0, 2] = fetch_fields(init, 6, {("tp", "sfc", 0)})[("tp", 0)]
+ inputs, targets, forcings = data_utils.extract_inputs_targets_forcings(
+ ds, target_lead_times=slice("6h", "6h"), **dataclasses.asdict(task))
+ return tuple(x.drop_vars("datetime", errors="ignore").compute() for x in (inputs, targets, forcings))
+
+
def truth_fields(t: dt.datetime) -> dict:
"""Verification at valid time t: IFS analysis msl/10u/10v and the 0-6h IFS precip ending at t."""
f = fetch_fields(t, 0, {("msl", "sfc", 0), ("10u", "sfc", 0), ("10v", "sfc", 0)})
diff --git a/src/faster_weathernext/pallas_attn.py b/src/faster_weathernext/pallas_attn.py
index 37cacd6..396e939 100644
--- a/src/faster_weathernext/pallas_attn.py
+++ b/src/faster_weathernext/pallas_attn.py
@@ -9,8 +9,14 @@
Numerics match the masked softmax of the stock paths up to summation order: masked logits
are set to -1e30 before the softmax (exp underflows to exactly 0), matmuls are TF32 under
JAX's default matmul precision and full fp32 under jax.default_matmul_precision("highest").
+
+The kernel is differentiable (jax.custom_vjp, FlashAttention-2 backward): the forward also stores
+each row's log-sum-exp, and two more tiled kernels compute dQ (per query block, over the same tile
+list) and dK/dV (per key block, over the transposed mask's tile list, which for the symmetric k-hop
+mask is the same list). Without differentiation the forward is unchanged.
"""
+import dataclasses
import functools
import jax
@@ -25,6 +31,7 @@
BLOCK_K = 32 # one uint32 mask word per query row per tile
NUM_WARPS = 4
NUM_STAGES = 1
+BWD_NUM_WARPS = 4 # tuned on the GPU (docs/validation.md §9)
MASK_VALUE = -1e30
# Under jax.default_matmul_precision("highest"): Triton's strict-fp32 ("ieee") dot does not use
# tensor cores and is ~10x slower on consumer Blackwell (sm_120), so by default use 3xTF32
@@ -102,7 +109,7 @@ def build_tile_mask(mask: sparse.spmatrix, block_q: int = BLOCK_Q, block_k: int
return tile_ptr.astype(np.int32), tile_kb.astype(np.int32), bits
-def _kernel(q_ref, k_ref, v_ref, ptr_ref, kb_ref, bits_ref, o_ref, *, block_k, scale, precision):
+def _kernel(q_ref, k_ref, v_ref, ptr_ref, kb_ref, bits_ref, o_ref, *lse_ref, block_k, scale, precision):
qb = pl.program_id(0)
q = plgpu.load(q_ref) # [BQ, D]
block_q, head_dim = q.shape
@@ -127,15 +134,198 @@ def body(t, carry):
init = (jnp.zeros((block_q, head_dim), jnp.float32),
jnp.full((block_q,), -jnp.inf, jnp.float32),
jnp.zeros((block_q,), jnp.float32))
- acc, _, l = lax.fori_loop(ptr_ref[qb], ptr_ref[qb + 1], body, init)
+ acc, m, l = lax.fori_loop(ptr_ref[qb], ptr_ref[qb + 1], body, init)
# Every real query attends at least to itself; the guard only affects padding rows.
o_ref[...] = (acc / jnp.maximum(l, 1e-30)[:, None]).astype(o_ref.dtype)
+ if lse_ref: # training: per-row log-sum-exp for the backward pass (padding rows: finite garbage)
+ lse_ref[0][...] = m + jnp.log(jnp.maximum(l, 1e-30))
+
+
+# Backward (FlashAttention-2): dS = P * (dP - delta), delta = rowsum(dO * O). Both kernels recompute
+# P = exp(S - lse) tile by tile from the forward's log-sum-exp, so no attention matrix is stored.
+# Masked entries are zeroed with `where` rather than through exp underflow, so the finite garbage
+# that padding rows carry (padded dO is 0) cannot leak into real rows.
+# With split = 2 each program takes half of a 64-row tile-list block (32 rows): it walks the
+# block's tiles and reads its half of each tile's mask words, which halves the operands a program
+# keeps in shared memory without a second tile list (tiles empty for that half are just masked).
+
+
+def _row_words(bits_ref, t, split, part):
+ """This program's rows of tile t's mask words (`part`: which of the block's `split` parts)."""
+ if split == 1:
+ return plgpu.load(bits_ref.at[t, :])
+ rows = bits_ref.shape[1] // split
+ return plgpu.load(bits_ref.at[t, pl.ds(part * rows, rows)])
+
+
+def _bwd_dq_kernel(q_ref, k_ref, v_ref, do_ref, lse_ref, delta_ref, ptr_ref, kb_ref, bits_ref, dq_ref,
+ *, block_k, split, scale, precision):
+ """dQ for one (query block, head): walks the query block's key tiles, as the forward does."""
+ qb, part = pl.program_id(0) // split, pl.program_id(0) % split
+ q = plgpu.load(q_ref) # [BQ, D]
+ do = plgpu.load(do_ref) # [BQ, D]
+ lse = plgpu.load(lse_ref) # [BQ]
+ delta = plgpu.load(delta_ref) # [BQ]
+ cols = lax.broadcasted_iota(jnp.uint32, (1, block_k), 1)
+
+ def body(t, dq):
+ keys = pl.ds(kb_ref[t] * block_k, block_k)
+ k = plgpu.load(k_ref.at[keys, :]) # [BK, D]
+ v = plgpu.load(v_ref.at[keys, :])
+ s = plgpu.dot(q, k, trans_b=True, precision=precision) * scale
+ words = _row_words(bits_ref, t, split, part)
+ valid = ((words[:, None] >> cols) & jnp.uint32(1)) != 0
+ p = jnp.where(valid, jnp.exp(s - lse[:, None]), 0.0)
+ dp = plgpu.dot(do, v, trans_b=True, precision=precision)
+ ds = p * (dp - delta[:, None])
+ return dq + plgpu.dot(ds, k, precision=precision) * scale
+
+ dq = lax.fori_loop(ptr_ref[qb], ptr_ref[qb + 1], body, jnp.zeros(q.shape, jnp.float32))
+ dq_ref[...] = dq.astype(dq_ref.dtype)
+
+
+def _bwd_dkv_kernel(q_ref, k_ref, v_ref, do_ref, lse_ref, delta_ref, ptr_ref, kb_ref, bits_ref,
+ dk_ref, dv_ref, *, block_q, split, scale, precision):
+ """dK, dV for one (key block, head): walks the key block's query tiles (the transposed tile list)."""
+ kb, part = pl.program_id(0) // split, pl.program_id(0) % split
+ k = plgpu.load(k_ref) # [BKV, D]
+ v = plgpu.load(v_ref)
+ cols = lax.broadcasted_iota(jnp.uint32, (1, block_q), 1)
+
+ def body(t, carry):
+ dk, dv = carry
+ queries = pl.ds(kb_ref[t] * block_q, block_q)
+ q = plgpu.load(q_ref.at[queries, :]) # [BQT, D]
+ do = plgpu.load(do_ref.at[queries, :])
+ lse = plgpu.load(lse_ref.at[queries]) # [BQT]
+ delta = plgpu.load(delta_ref.at[queries])
+ st = plgpu.dot(k, q, trans_b=True, precision=precision) * scale # [BKV, BQT] = S^T
+ words = _row_words(bits_ref, t, split, part) # bit j of row r: mask[query j, key r]
+ valid = ((words[:, None] >> cols) & jnp.uint32(1)) != 0
+ pt = jnp.where(valid, jnp.exp(st - lse[None, :]), 0.0)
+ dv = dv + plgpu.dot(pt, do, precision=precision)
+ dpt = plgpu.dot(v, do, trans_b=True, precision=precision)
+ dst = pt * (dpt - delta[None, :])
+ dk = dk + plgpu.dot(dst, q, precision=precision) * scale
+ return dk, dv
+
+ zeros = jnp.zeros(k.shape, jnp.float32)
+ dk, dv = lax.fori_loop(ptr_ref[kb], ptr_ref[kb + 1], body, (zeros, zeros))
+ dk_ref[...] = dk.astype(dk_ref.dtype)
+ dv_ref[...] = dv.astype(dv_ref.dtype)
+
+
+@dataclasses.dataclass(frozen=True)
+class _Config:
+ scale: float
+ precision: object
+ block_q: int
+ block_k: int
+ num_warps: int
+ num_stages: int
+ bwd_num_warps: int
+ bwd_split: int
+ interpret: bool
+
+
+def _full(a):
+ return pl.BlockSpec(a.shape, lambda i, h, b: (0,) * a.ndim)
+
+
+def _forward(qh, kh, vh, tiles, cfg, with_lse):
+ batch, heads, n_q, head_dim = qh.shape
+ n_k = kh.shape[2]
+ out_shape = jax.ShapeDtypeStruct((batch, heads, n_q, head_dim), qh.dtype)
+ out_specs = pl.BlockSpec((None, None, cfg.block_q, head_dim), lambda i, h, b: (b, h, i, 0))
+ if with_lse:
+ out_shape = (out_shape, jax.ShapeDtypeStruct((batch, heads, n_q), jnp.float32))
+ out_specs = (out_specs, pl.BlockSpec((None, None, cfg.block_q), lambda i, h, b: (b, h, i)))
+ return pl.pallas_call(
+ functools.partial(_kernel, block_k=cfg.block_k, scale=cfg.scale, precision=cfg.precision),
+ grid=(n_q // cfg.block_q, heads, batch),
+ in_specs=[
+ pl.BlockSpec((None, None, cfg.block_q, head_dim), lambda i, h, b: (b, h, i, 0)),
+ pl.BlockSpec((None, None, n_k, head_dim), lambda i, h, b: (b, h, 0, 0)),
+ pl.BlockSpec((None, None, n_k, head_dim), lambda i, h, b: (b, h, 0, 0)),
+ *(_full(t) for t in tiles),
+ ],
+ out_specs=out_specs,
+ out_shape=out_shape,
+ compiler_params=plgpu.CompilerParams(num_warps=cfg.num_warps, num_stages=cfg.num_stages),
+ name="wn2_tiled_flash_attention",
+ interpret=cfg.interpret,
+ )(qh, kh, vh, *tiles)
+
+
+@functools.partial(jax.custom_vjp, nondiff_argnums=(5,))
+def _attention(qh, kh, vh, tiles, tiles_t, cfg):
+ """Head-major padded attention: qh [B, H, n_q, D], kh/vh [B, H, n_k, D] -> [B, H, n_q, D]."""
+ del tiles_t
+ return _forward(qh, kh, vh, tiles, cfg, with_lse=False)
+
+
+def _attention_fwd(qh, kh, vh, tiles, tiles_t, cfg):
+ # The backward recomputes the logits and normalises them with this pass's log-sum-exp; all
+ # kernels use one dot precision, so the recomputed P is the forward's.
+ out, lse = _forward(qh, kh, vh, tiles, cfg, with_lse=True)
+ return out, (qh, kh, vh, out, lse, tiles, tiles_t)
+
+
+def _attention_bwd(cfg, res, dout):
+ qh, kh, vh, out, lse, tiles, tiles_t = res
+ batch, heads, n_q, head_dim = qh.shape
+ n_k = kh.shape[2]
+ delta = jnp.sum(out * dout, axis=-1) # [B, H, n_q]
+ split = cfg.bwd_split
+ rows_per = cfg.block_q // split
+ params = plgpu.CompilerParams(num_warps=cfg.bwd_num_warps, num_stages=cfg.num_stages)
+ rows = lambda size: pl.BlockSpec((None, None, size, head_dim), lambda i, h, b: (b, h, i, 0))
+ vec = lambda size: pl.BlockSpec((None, None, size), lambda i, h, b: (b, h, i))
+ whole = lambda x: pl.BlockSpec((None, None, *x.shape[2:]), lambda i, h, b: (b, h) + (0,) * (x.ndim - 2))
+
+ dq = pl.pallas_call(
+ functools.partial(_bwd_dq_kernel, block_k=cfg.block_k, split=split, scale=cfg.scale,
+ precision=cfg.precision),
+ grid=(n_q // rows_per, heads, batch),
+ in_specs=[rows(rows_per), whole(kh), whole(vh), rows(rows_per), vec(rows_per), vec(rows_per),
+ *(_full(t) for t in tiles)],
+ out_specs=rows(rows_per),
+ out_shape=jax.ShapeDtypeStruct(qh.shape, qh.dtype),
+ compiler_params=params, name="wn2_tiled_flash_attention_dq", interpret=cfg.interpret,
+ )(qh, kh, vh, dout, lse, delta, *tiles)
+
+ # Key blocks of block_q rows walk query tiles of block_k: the transposed mask in the forward's
+ # tile layout. Keys are padded to whole key blocks; queries already span whole query tiles.
+ n_kb = tiles_t[0].shape[0] - 1
+ pad = lambda x: jnp.pad(x, ((0, 0), (0, 0), (0, n_kb * cfg.block_q - n_k), (0, 0)))
+ dk, dv = pl.pallas_call(
+ functools.partial(_bwd_dkv_kernel, block_q=cfg.block_k, split=split, scale=cfg.scale,
+ precision=cfg.precision),
+ grid=(n_kb * split, heads, batch),
+ in_specs=[whole(qh), rows(rows_per), rows(rows_per), whole(dout), whole(lse), whole(delta),
+ *(_full(t) for t in tiles_t)],
+ out_specs=(rows(rows_per), rows(rows_per)),
+ out_shape=(jax.ShapeDtypeStruct((batch, heads, n_kb * cfg.block_q, head_dim), kh.dtype),) * 2,
+ compiler_params=params, name="wn2_tiled_flash_attention_dkv", interpret=cfg.interpret,
+ )(qh, pad(kh), pad(vh), dout, lse, delta, *tiles_t)
+ return dq, dk[:, :, :n_k], dv[:, :, :n_k], None, None
+
+
+_attention.defvjp(_attention_fwd, _attention_bwd)
def tiled_flash_attention(q, k, v, tile_ptr, tile_kb, tile_bits, *, scale, strict_fp32=False,
- block_q=BLOCK_Q, block_k=BLOCK_K, num_warps=NUM_WARPS,
- num_stages=NUM_STAGES, interpret=False):
- """q, k, v: [batch, n, heads, head_dim] float32 -> [batch, n, heads, head_dim]."""
+ tiles_t=None, block_q=BLOCK_Q, block_k=BLOCK_K, num_warps=NUM_WARPS,
+ num_stages=NUM_STAGES, bwd_num_warps=BWD_NUM_WARPS, bwd_split=None,
+ interpret=False):
+ """q, k, v: [batch, n, heads, head_dim] float32 -> [batch, n, heads, head_dim].
+
+ Differentiable: the backward pass runs two more tiled kernels (dQ per query block; dK and dV per
+ key block). The dK/dV kernel needs the tile list of the transposed mask (`tiles_t`, from
+ build_tile_mask(mask.T)); it defaults to the forward's, which is correct for symmetric masks
+ such as the k-hop mesh neighbourhoods. bwd_split (1 or 2; default: 2 under `highest`, else 1) is
+ the number of backward programs per 64-row block.
+ """
batch, n, heads, head_dim = q.shape
n_qblocks = tile_ptr.shape[0] - 1
n_q = n_qblocks * block_q
@@ -150,21 +340,15 @@ def tiled_flash_attention(q, k, v, tile_ptr, tile_kb, tile_bits, *, scale, stric
else lax.DotAlgorithmPreset.TF32_TF32_F32_X3)
else:
precision = lax.DotAlgorithmPreset.TF32_TF32_F32
- kernel = functools.partial(_kernel, block_k=block_k, scale=scale, precision=precision)
- full = lambda a: pl.BlockSpec(a.shape, lambda i, h, b: (0,) * a.ndim)
- out = pl.pallas_call(
- kernel,
- grid=(n_qblocks, heads, batch),
- in_specs=[
- pl.BlockSpec((None, None, block_q, head_dim), lambda i, h, b: (b, h, i, 0)),
- pl.BlockSpec((None, None, n_k, head_dim), lambda i, h, b: (b, h, 0, 0)),
- pl.BlockSpec((None, None, n_k, head_dim), lambda i, h, b: (b, h, 0, 0)),
- full(tile_ptr), full(tile_kb), full(tile_bits),
- ],
- out_specs=pl.BlockSpec((None, None, block_q, head_dim), lambda i, h, b: (b, h, i, 0)),
- out_shape=jax.ShapeDtypeStruct((batch, heads, n_q, head_dim), q.dtype),
- compiler_params=plgpu.CompilerParams(num_warps=num_warps, num_stages=num_stages),
- name="wn2_tiled_flash_attention",
- interpret=interpret,
- )(qh, kh, vh, jnp.asarray(tile_ptr), jnp.asarray(tile_kb), jnp.asarray(tile_bits))
+ # 3xTF32 (and IEEE fp32) dots keep extra operand copies in shared memory: the forward (2 dots
+ # per tile) fits sm_120's 99 KB in 64-row programs, the backward kernels (3-4 dots) asked for
+ # 147 KB, so under `highest` they run in 32-row programs (bwd_split=2).
+ if bwd_split is None:
+ bwd_split = 1 if precision in (None, lax.DotAlgorithmPreset.TF32_TF32_F32) else 2
+ cfg = _Config(scale=scale, precision=precision, block_q=block_q, block_k=block_k,
+ num_warps=num_warps, num_stages=num_stages, bwd_num_warps=bwd_num_warps, bwd_split=bwd_split,
+ interpret=interpret)
+ tiles = tuple(jnp.asarray(t) for t in (tile_ptr, tile_kb, tile_bits))
+ tiles_t = tiles if tiles_t is None else tuple(jnp.asarray(t) for t in tiles_t)
+ out = _attention(qh, kh, vh, tiles, tiles_t, cfg)
return out[:, :, :n].transpose(0, 2, 1, 3)
diff --git a/src/faster_weathernext/patch.py b/src/faster_weathernext/patch.py
index f32beed..cca6131 100644
--- a/src/faster_weathernext/patch.py
+++ b/src/faster_weathernext/patch.py
@@ -40,7 +40,9 @@
from weathernext.utils import sparse_transformer as st
from weathernext.utils import sparse_transformer_utils as st_utils
from weathernext.utils import typed_graph_net
+from weathernext.utils import casting
from weathernext.utils import xarray_dense
+from weathernext.weathernext2 import fgn
from faster_weathernext import pallas_attn
@@ -65,6 +67,16 @@ class Options:
edge_chunk: int = 65536
# Run the transformer layers as a loop (hk.scan + hk.switch) instead of unrolled.
layer_loop: bool = True
+ # Rematerialise in the backward pass: each grid/edge block (GNNs, encoder, decoder), each
+ # transformer layer, and each XLA attention chunk. Without it, differentiating the blocked scans
+ # stores every block's intermediates, i.e. exactly the whole-grid tensors the blocking avoids.
+ # Has no effect on code that is not differentiated.
+ remat: bool = False
+ # Training: run the network once per element of the stacked (batch x ensemble member) batch,
+ # each call rematerialised, instead of once on the whole batch. The backward pass then holds one
+ # member's activations at a time (about half the memory for FGN's 2 members). The noise is
+ # drawn for the stacked batch as before and split, so results match the batched call.
+ sequential_samples: bool = False
OPTIONS = Options()
@@ -139,6 +151,11 @@ def one_chunk(args):
weights = st_utils.wrap_fn_for_upcast_downcast(logits, jax.nn.softmax)
return jnp.einsum("bhqk,bkhd->bqhd", weights, vc)
+ # With remat, the backward pass recomputes each chunk's logits instead of storing them all.
+ # one_chunk is pure JAX inside a raw lax.map, so jax.checkpoint, not hk.remat: hk.remat threads
+ # Haiku's state (the RNG key) and leaks it out of the lax.map trace (UnexpectedTracerError).
+ if OPTIONS.remat:
+ one_chunk = jax.checkpoint(one_chunk, prevent_cse=False)
out = jax.lax.map(one_chunk, (q, starts, bits))
out = out.transpose(1, 0, 2, 3, 4).reshape(batch, n_chunks * chunk, heads * head_dim)[:, :n]
attn_winit_final = hk.initializers.VarianceScaling(cfg.attn_winit_final_mult / cfg.num_layers)
@@ -152,8 +169,9 @@ def tiled_mha(x, tiles, cfg):
k = st.multihead_linear(x, "k", cfg)
v = st.multihead_linear(x, "v", cfg)
_, _, heads, head_dim = q.shape
- out = pallas_attn.tiled_flash_attention(q, k, v, *tiles, scale=cfg.key_size**-0.5,
- strict_fp32=OPTIONS.strict_fp32)
+ # tiles[3:] is the transposed mask's tile list, present only for non-symmetric masks.
+ out = pallas_attn.tiled_flash_attention(q, k, v, *tiles[:3], tiles_t=tiles[3:] or None,
+ scale=cfg.key_size**-0.5, strict_fp32=OPTIONS.strict_fp32)
attn_winit_final = hk.initializers.VarianceScaling(cfg.attn_winit_final_mult / cfg.num_layers)
return hk.Linear(cfg.d_model, name="mha_final", w_init=attn_winit_final)(
out.reshape(batch, n, heads * head_dim))
@@ -168,6 +186,7 @@ def tiled_mha(x, tiles, cfg):
Decoder=xarray_dense.DataArrayDictDenseDecoder,
build_typed_graph=points_mesh_gnn.PointsMeshTypedGraphGNN._build_typed_graph,
get_adjacency_matrix=mesh_transformer._get_adjacency_matrix,
+ run_model_with_sample_as_batch=fgn.Predictor.run_model_with_sample_as_batch,
)
@@ -193,6 +212,10 @@ def _transformer_init(self, adj_mat, attention_k_hop, attention_type, mask_type,
if perm is not None:
mask = mask.tocsr()[perm][:, perm]
self.mask = pallas_attn.build_tile_mask(mask)
+ # The backward pass walks the transposed mask. Powers of a symmetric adjacency (the mesh's)
+ # are symmetric, so the forward tile list serves; checking the adjacency is cheap.
+ if (adj_mat != adj_mat.T).nnz:
+ self.mask += pallas_attn.build_tile_mask(mask.T.tocsr())
self.num_padding_nodes = ("pallas", perm)
else:
starts, bits, window, n_kv = build_chunked_mask(adj_mat**attention_k_hop, OPTIONS.attn_chunk)
@@ -252,14 +275,16 @@ def apply(block, x):
if perm is not None:
node_features = jnp.take(node_features, jnp.asarray(perm, jnp.int32), axis=1)
+ # With remat, the backward pass keeps only each layer's input and recomputes the layer
+ # (attention q/k/v/output alone would otherwise be ~1 GiB per layer at 0.25°, batch 2).
if OPTIONS.layer_loop:
branches = [functools.partial(apply, b) for b in blocks]
- x, _ = hk.scan(lambda x, i: (hk.switch(i, branches, x), None), node_features,
- jnp.arange(len(blocks)))
+ layer = _maybe_remat(lambda x, i: hk.switch(i, branches, x))
+ x, _ = hk.scan(lambda x, i: (layer(x, i), None), node_features, jnp.arange(len(blocks)))
else:
x = node_features
for b in blocks:
- x = apply(b, x)
+ x = _maybe_remat(functools.partial(apply, b))(x)
x = st.layernorm(x, create_scale=False, create_offset=False)
x = dense.LinearNormConditioning(name=self.name + "_final_norm_conditioning")(
x, norm_conditioning=jnp.expand_dims(global_norm_conditioning, 1))
@@ -297,8 +322,15 @@ def _pad_rows(x, rows):
return x if x.shape[0] == rows else jnp.pad(x, [(0, rows - x.shape[0])] + [(0, 0)] * (x.ndim - 1))
+def _maybe_remat(fn):
+ # prevent_cse=False: these functions run as scan/map bodies, where XLA cannot CSE the
+ # recomputation away across iterations.
+ return hk.remat(fn, prevent_cse=False) if OPTIONS.remat else fn
+
+
def _update_blocks_in_place(fn, blocks, xs=()):
"""blocks[i] <- fn(blocks[i], *xs[i]), as a scan whose carry XLA updates in place."""
+ fn = _maybe_remat(fn)
def body(buf, args):
i, *rest = args
@@ -441,14 +473,21 @@ def _points_to_mesh(self, graph, edge_key, f):
mesh_rcv_proj = f["pre_gather_receivers_for_edges_fn"][name](mesh) if gather_rcv else None
acc_dtype = jnp.float32 if self._f32_aggregation else mesh.dtype
- def one_block(acc, xs):
- e, s, r = xs
+ def messages(e, s, r):
# Row-wise matmul commutes with the gather; gathering first avoids a [points, F] buffer.
sent = f["pre_gather_senders_for_edges_fn"][name](gather_scatter_ops.gather_with_fill(pts, s))
recv = gather_scatter_ops.gather_with_fill(mesh_rcv_proj, r) if gather_rcv else None
- msg = f["update_edge_fn"][name](
+ return f["update_edge_fn"][name](
f["pre_gather_edges_for_edges_fn"][name](self._encode_edges(e)), sent, recv,
graph.context.features)
+
+ # Only the messages are rematerialised: the accumulation is linear, so the backward pass
+ # needs no per-block copy of the [num_mesh, B, F] accumulator.
+ messages = _maybe_remat(messages)
+
+ def one_block(acc, xs):
+ e, s, r = xs
+ msg = messages(e, s, r)
acc = acc + jax.ops.segment_sum(msg.astype(acc_dtype), r, num_mesh, indices_are_sorted=True)
return acc, None
@@ -576,7 +615,7 @@ def one_block(i):
# grid afterwards is free at batch 1 (size-1 axes) but a second 6 GiB copy at batch 2.
return jnp.moveaxis(out.reshape(out.shape[0], block, out.shape[-1]), 0, 1) # [block, batch, latent]
- stacked = hk.map(one_block, jnp.arange(nb)) # [nb, block, batch, latent]
+ stacked = hk.map(_maybe_remat(one_block), jnp.arange(nb)) # [nb, block, batch, latent]
points = stacked.reshape((nb * block,) + stacked.shape[2:])[:n]
batch, latent = points.shape[1], points.shape[2]
# Back to the [batch, lat, lon, latent] contract; the model flattens it to points-major again
@@ -615,7 +654,7 @@ def one_block(i):
out = stock_call(jnp.moveaxis(b, 0, 1).reshape(batch, 1, block, feat), block_template)
return tuple(xarray_jax.unwrap_data(out[name]) for name in names)
- outs = hk.map(one_block, jnp.arange(nb))
+ outs = hk.map(_maybe_remat(one_block), jnp.arange(nb))
result = {}
for name, o in zip(names, outs):
da = output_template[name]
@@ -624,6 +663,61 @@ def one_block(i):
return xr.Dataset(result) if isinstance(output_template, xr.Dataset) else result
+# ----------------------------------------------------------------------------------------------
+# Ensemble members one at a time (training)
+
+
+def _take(ds, i):
+ """Element i of a stacked-batch Dataset, as batch 1 without the (MultiIndex) batch coordinates."""
+ if ds is None:
+ return None
+ out = {}
+ for name, da in ds.data_vars.items():
+ if "batch" not in da.dims: # static fields
+ out[name] = da
+ continue
+ x = jax.lax.dynamic_slice_in_dim(xarray_jax.unwrap_data(da), i, 1, axis=da.dims.index("batch"))
+ out[name] = xarray_jax.DataArray(x, dims=da.dims,
+ coords={k: c for k, c in da.coords.items() if "batch" not in c.dims})
+ return xr.Dataset(out)
+
+
+def _run_model_with_sample_as_batch(self, inputs, targets, forcings, is_training, add_noise=True):
+ """fgn.Predictor.run_model_with_sample_as_batch; with sequential_samples the network runs once
+ per stacked element, inside hk.scan with a rematerialised body. Independent per-element calls are
+ not enough: XLA then runs the elements' backward passes side by side (46.1 vs 45.8 GiB at 0.25°);
+ a scan's backward runs its iterations one after another. fgn.Predictor is not a Haiku module and
+ the network module is built in its __init__, so parameter names are unchanged."""
+ if not OPTIONS.sequential_samples:
+ return _ORIG["run_model_with_sample_as_batch"](self, inputs, targets, forcings, is_training,
+ add_noise=add_noise)
+ # As the official method: stack batch and sample, draw the noise for the whole stack.
+ stack = lambda ds: ds.rename({"batch": "_batch"}).stack(batch=("_batch", "sample")).transpose("batch", ...)
+ inputs, targets = stack(inputs), stack(targets)
+ forcings = stack(forcings) if forcings is not None else None
+ if add_noise:
+ inputs, forcings = self._noise_generator(inputs, forcings,
+ default_dtype=casting.infer_floating_dtype(inputs))
+ n = inputs.sizes["batch"]
+ dims = {}
+
+ def element(carry, i):
+ preds = self._noisy_function(inputs=_take(inputs, i), targets_template=_take(targets, i),
+ forcings=_take(forcings, i), is_training=is_training)
+ dims.update({v: da.dims for v, da in preds.data_vars.items()})
+ return carry, {v: xarray_jax.unwrap_data(da) for v, da in preds.data_vars.items()}
+
+ _, stacked = hk.scan(hk.remat(element), None, jnp.arange(n))
+ data = {}
+ for v, x in stacked.items(): # [n, ..., 1 (batch), ...] -> [..., n (batch), ...]
+ ax = dims[v].index("batch")
+ x = jnp.moveaxis(x, 0, ax)
+ data[v] = xarray_jax.DataArray(x.reshape(x.shape[:ax] + (n,) + x.shape[ax + 2:]), dims=dims[v],
+ coords=targets[v].coords)
+ preds = xr.Dataset(data).unstack("batch").rename({"_batch": "batch"}).transpose("sample", "batch", ...)
+ return preds.drop_indexes(["sample", "batch"]).reset_coords(["sample", "batch"], drop=True)
+
+
# ----------------------------------------------------------------------------------------------
# Install / uninstall
@@ -645,6 +739,7 @@ def install():
points_mesh_gnn.PointsMeshTypedGraphGNN._build_typed_graph = _build_typed_graph
xarray_dense.DataArrayDictDenseEncoder = ChunkedDenseEncoder
xarray_dense.DataArrayDictDenseDecoder = ChunkedDenseDecoder
+ fgn.Predictor.run_model_with_sample_as_batch = _run_model_with_sample_as_batch
_enabled = True
@@ -659,6 +754,7 @@ def uninstall():
points_mesh_gnn.PointsMeshTypedGraphGNN._build_typed_graph = _ORIG["build_typed_graph"]
xarray_dense.DataArrayDictDenseEncoder = _ORIG["Encoder"]
xarray_dense.DataArrayDictDenseDecoder = _ORIG["Decoder"]
+ fgn.Predictor.run_model_with_sample_as_batch = _ORIG["run_model_with_sample_as_batch"]
_enabled = False
diff --git a/tests/test_finetune.py b/tests/test_finetune.py
new file mode 100644
index 0000000..5626b0f
--- /dev/null
+++ b/tests/test_finetune.py
@@ -0,0 +1,130 @@
+"""CPU tests: head/frozen parameter split, remat of the blocked scans, checkpoint save/load."""
+import haiku as hk
+import jax
+import jax.numpy as jnp
+import numpy as np
+import pytest
+
+from faster_weathernext import finetune
+from faster_weathernext import patch
+
+PARAMS = {
+ "multimodality_forward/grid_encoder/shared_dense/mlp/~/linear_0": {"w": np.ones((3, 4), np.float32)},
+ "multimodality_forward/mesh_transformer/transformer/block_00/mha_final": {"w": np.ones((4, 4), np.float32)},
+ "multimodality_forward/mesh_to_grid_gnn/deep_gnn/processor_nodes_0_point_nodes": {"b": np.ones(4, np.float32)},
+ "multimodality_forward/grid_decoder/split_output_linear": {"w_2m_temperature": np.ones((4, 1), np.float32)},
+ "multimodality_forward/grid_to_mesh_gnn/deep_gnn/processor_nodes_0_mesh_nodes": {"b": np.ones(4, np.float32)},
+}
+
+
+def test_split_selects_mesh_to_grid_and_decoder_only():
+ head, frozen = finetune.split(PARAMS)
+ assert sorted(m.split("/")[1] for m in head) == ["grid_decoder", "mesh_to_grid_gnn"]
+ assert not any(finetune.is_head(m) for m in frozen)
+ assert finetune.merge(frozen, head).keys() == PARAMS.keys()
+ assert finetune.head_size(PARAMS) == 4 + 4
+ trained, frozen = finetune.split(PARAMS, "all")
+ assert trained.keys() == PARAMS.keys() and not frozen
+ trained, _ = finetune.split(PARAMS, ("mesh_transformer", *finetune.HEAD_MODULES))
+ assert len(trained) == 3
+
+
+def test_masked_variables_get_zero_weight_and_nan_masking():
+ pytest.importorskip("weathernext")
+ cfg = finetune.predictor_config("WeatherNext2", masked=["sea_surface_temperature", "cyclone_usa_wind_disc"])
+ kw = cfg.predictor_kwargs["fgn_loss_function_kwargs"]
+ # A zero weight is what removes a variable that has data (SST); the NaN list alone would not.
+ assert kw["per_variable_weights"]["sea_surface_temperature"] == 0.0
+ assert kw["per_variable_weights"]["2m_temperature"] == 1.0 # others keep the config's weights
+ assert {"sea_surface_temperature", "cyclone_usa_wind_disc"} <= set(kw["vars_with_nan_targets"])
+ with pytest.raises(ValueError, match="not targets"):
+ finetune.predictor_config("WeatherNext2", masked=["not_a_variable"])
+
+
+@pytest.fixture
+def remat_option():
+ saved = patch.OPTIONS.remat
+ yield
+ patch.OPTIONS.remat = saved
+
+
+def test_remat_leaves_values_and_grads_unchanged(remat_option):
+ def f(blocks, xs):
+ lin = hk.Linear(5, name="lin") # built outside the scan, as ChunkedDeepGNN captures its modules
+
+ def fn(b, x):
+ return jnp.tanh(lin(b) + x)
+ return jnp.sum(patch._update_blocks_in_place(fn, blocks, (xs,)) ** 2)
+
+ model = hk.transform(f)
+ blocks = jnp.asarray(np.random.default_rng(0).standard_normal((4, 3, 5)), jnp.float32)
+ xs = jnp.asarray(np.random.default_rng(1).standard_normal((4, 3, 5)), jnp.float32)
+ params = model.init(jax.random.PRNGKey(0), blocks, xs)
+ results = {}
+ for remat in (False, True):
+ patch.OPTIONS.remat = remat
+ results[remat] = jax.value_and_grad(lambda p, b: model.apply(p, None, b, xs), argnums=(0, 1))(params, blocks)
+ (v0, (gp0, gb0)), (v1, (gp1, gb1)) = results[False], results[True]
+ np.testing.assert_allclose(v0, v1, rtol=1e-6)
+ np.testing.assert_allclose(gb0, gb1, rtol=1e-5, atol=1e-6)
+ for a, b in zip(jax.tree.leaves(gp0), jax.tree.leaves(gp1)):
+ np.testing.assert_allclose(a, b, rtol=1e-5, atol=1e-6)
+
+
+def test_remat_is_in_the_jaxpr_only_when_enabled(remat_option):
+ def f(blocks):
+ lin = hk.Linear(5, name="lin")
+ return patch._update_blocks_in_place(lin, blocks)
+
+ model = hk.transform(f)
+ blocks = jnp.ones((2, 3, 5), jnp.float32)
+ params = model.init(jax.random.PRNGKey(0), blocks)
+ for remat in (False, True):
+ patch.OPTIONS.remat = remat
+ jaxpr = str(jax.make_jaxpr(lambda p: jax.grad(lambda q: model.apply(q, None, blocks).sum())(p))(params))
+ assert ("checkpoint" in jaxpr or "remat" in jaxpr) == remat
+
+
+def test_save_roundtrip_and_shape_check(tmp_path, monkeypatch):
+ from weathernext.utils import checkpoint as ckpt_lib
+ from weathernext.weathernext2 import fgn
+ from faster_weathernext import model
+
+ base = tmp_path / "base.npz"
+ with open(base, "wb") as f:
+ ckpt_lib.dump(f, fgn.CheckPoint(params=PARAMS, description="base", license="CC BY 4.0"))
+ monkeypatch.setattr(model, "weights_path", lambda *a, **k: str(base))
+
+ head, frozen = finetune.split(PARAMS)
+ head = jax.tree.map(lambda v: v * 2, head)
+ out = tmp_path / "tuned.npz"
+ finetune.save(str(out), finetune.merge(frozen, head), note="head-only test")
+ loaded = finetune.load(str(out))
+ for m, p in finetune.merge(frozen, head).items():
+ for k, v in p.items():
+ np.testing.assert_array_equal(loaded[m][k], v)
+ with open(out, "rb") as f:
+ ck = ckpt_lib.load(f, fgn.CheckPoint)
+ assert ck.description == "base\n\nhead-only test" and ck.license == "CC BY 4.0"
+
+ bad = finetune.merge(frozen, jax.tree.map(lambda v: v[..., :1], head))
+ with pytest.raises(ValueError, match="do not match"):
+ finetune.save(str(tmp_path / "bad.npz"), bad)
+
+
+def test_take_slices_one_stacked_member_and_keeps_static_fields():
+ import xarray
+ import xarray_jax
+ a = np.arange(2 * 3 * 4, dtype=np.float32).reshape(1, 2, 3, 4) # batch, sample, lat, lon
+ ds = xarray.Dataset({"x": (("batch", "sample", "lat", "lon"), a),
+ "static": (("lat", "lon"), np.ones((3, 4), np.float32))},
+ coords={"lat": np.arange(3.0), "lon": np.arange(4.0)})
+ st = ds.rename({"batch": "_batch"}).stack(batch=("_batch", "sample")).transpose("batch", ...)
+
+ def f(i):
+ one = patch._take(st, i)
+ assert dict(one["x"].sizes) == {"batch": 1, "lat": 3, "lon": 4} and "batch" not in one["static"].dims
+ return xarray_jax.unwrap_data(one["x"])
+
+ np.testing.assert_array_equal(jax.jit(f)(1)[0], a[0, 1])
+ np.testing.assert_array_equal(jax.jit(f)(0)[0], a[0, 0])
diff --git a/tests/test_pallas_kernel.py b/tests/test_pallas_kernel.py
index d1ef659..df644ec 100644
--- a/tests/test_pallas_kernel.py
+++ b/tests/test_pallas_kernel.py
@@ -53,3 +53,66 @@ def test_matches_dense_reference(precision, tol, batch):
ref = _reference(q[b], k[b], v[b], m.toarray(), d**-0.5)
err = np.abs(np.asarray(out[b], np.float64) - ref).max() / np.abs(ref).max()
assert err < tol, (b, err)
+
+
+def _reference_grads(q, k, v, do, mask, scale):
+ """float64 dQ, dK, dV of the masked softmax attention for one batch element ((n, heads, d))."""
+ q, k, v, do = (x.astype(np.float64) for x in (q, k, v, do))
+ s = np.where(mask[None], np.einsum("nhd,mhd->hnm", q, k) * scale, -np.inf)
+ p = np.exp(s - s.max(-1, keepdims=True))
+ p /= p.sum(-1, keepdims=True) # [h, n, m]
+ dp = np.einsum("nhd,mhd->hnm", do, v)
+ ds = p * (dp - (dp * p).sum(-1, keepdims=True))
+ dq = np.einsum("hnm,mhd->nhd", ds, k) * scale
+ dk = np.einsum("hnm,nhd->mhd", ds, q) * scale
+ dv = np.einsum("hnm,nhd->mhd", p, do)
+ return dq, dk, dv
+
+
+@pytest.mark.parametrize("split", [1, 2])
+@pytest.mark.parametrize("symmetric", [True, False])
+@pytest.mark.parametrize("batch", [1, 2])
+@pytest.mark.parametrize("precision,tol", [("highest", 1e-5), ("default", 3e-2)])
+def test_grads_match_dense_reference(precision, tol, batch, symmetric, split):
+ """Backward kernels (dQ per query block, dK/dV per key block over the transposed tile list)."""
+ if precision == "default" and not _on_gpu():
+ pytest.skip("TF32 only exists on GPUs")
+ rng = np.random.default_rng(2)
+ n, h, d = 300, 2, 64 # n is not a multiple of the 64-row blocks: padding rows are exercised
+ q, k, v, do = (rng.standard_normal((batch, n, h, d)).astype(np.float32) for _ in range(4))
+ m = sparse.random(n, n, density=0.08, random_state=3, format="csr").astype(bool)
+ m = ((m + m.T) if symmetric else m) + sparse.identity(n, dtype=bool)
+ m = m.tocsr().astype(bool)
+ tiles = pallas_attn.build_tile_mask(m)
+ tiles_t = None if symmetric else pallas_attn.build_tile_mask(m.T.tocsr())
+ fn = lambda q, k, v: pallas_attn.tiled_flash_attention(q, k, v, *tiles, tiles_t=tiles_t, scale=d**-0.5,
+ bwd_split=split, interpret=not _on_gpu())
+ try:
+ with jax.default_matmul_precision(precision):
+ _, vjp = jax.vjp(fn, *(jnp.asarray(x) for x in (q, k, v)))
+ grads = vjp(jnp.asarray(do))
+ except NotImplementedError as e:
+ pytest.skip(f"Pallas interpret mode unsupported: {e}")
+ for b in range(batch):
+ for name, got, ref in zip("qkv", grads, _reference_grads(q[b], k[b], v[b], do[b], m.toarray(), d**-0.5)):
+ err = np.abs(np.asarray(got[b], np.float64) - ref).max() / np.abs(ref).max()
+ assert err < tol, (name, b, err)
+
+
+def test_nonsymmetric_mask_needs_transposed_tiles():
+ """Reusing the forward tiles for a non-symmetric mask gives wrong dK/dV (guards the default)."""
+ rng = np.random.default_rng(4)
+ n, h, d = 130, 1, 32
+ q, k, v, do = (rng.standard_normal((1, n, h, d)).astype(np.float32) for _ in range(4))
+ m = (sparse.random(n, n, density=0.1, random_state=5, format="csr").astype(bool)
+ + sparse.identity(n, dtype=bool)).tocsr()
+ tiles = pallas_attn.build_tile_mask(m)
+ try:
+ _, vjp = jax.vjp(lambda k: pallas_attn.tiled_flash_attention(jnp.asarray(q), k, jnp.asarray(v), *tiles,
+ scale=d**-0.5, interpret=not _on_gpu()),
+ jnp.asarray(k))
+ (dk,) = vjp(jnp.asarray(do))
+ except NotImplementedError as e:
+ pytest.skip(f"Pallas interpret mode unsupported: {e}")
+ ref = _reference_grads(q[0], k[0], v[0], do[0], m.toarray(), d**-0.5)[1]
+ assert np.abs(np.asarray(dk[0], np.float64) - ref).max() / np.abs(ref).max() > 1e-2