Fine-tuning (experimental): differentiable attention kernel, head-only and whole-model training

AGENTS.md

@@ -92,13 +92,34 @@ 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.
 - **In-place block updates** need the new block wrapped in `optimization_barrier` before the
   `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

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;

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 |
+
+<sub>¹ 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.</sub>
+
+```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
 
 <details>
@@ -177,7 +198,7 @@ Yes. Install `faster-weathernext` (without the `[weathernext]` extra) into the e
 <details>
 <summary><b>Can I train or fine-tune with this?</b></summary>
 
-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).
 </details>
 
 ## Citation and licenses

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]'")
 )
 
 

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 <app id>      # 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)

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

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

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

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
+}

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
+}

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
+}

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
+}

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

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
+}