A100 fp32 deviation explained: the host CPU builds a different graph

AGENTS.md

@@ -106,9 +106,17 @@ below the ensemble member-to-member spread (docs/validation.md).
   (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.
+- **XLA autotuning is timing-based**: kernel picks (and last bits) differ between runs. Use
+  `--autotune-cache` for reproducible runs.
+- **The grid↔mesh graphs depend on the host CPU** (official code): they are built with NumPy from
+  float32 lat/lon at trace time, and NumPy's math code path, picked from the host CPU (AVX-512,
+  AVX2, baseline), changes last bits and with them which grid points near the ball-query radius
+  get an edge. An AVX2 host
+  (Modal's A100-40GB) and an AVX-512 host differ in 13 of 1.6M grid→mesh edges and 517 mesh→grid
+  triangles: 1.3e-2 max rel RMS at `highest`, which looks like TF32 and is not. Before reading a
+  cross-machine fp32 difference as GPU numerics, compare the graph hashes in `cloud/canary_run.py`'s
+  JSON (or `scripts/graph_edges.py`); `--edges` pins a graph. Modal may serve "A100-40GB" requests
+  on 80GB cards: check `device`.
 - **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

CHANGELOG.md

@@ -1,5 +1,18 @@
 # Changelog
 
+## Unreleased
+
+- The A100 fp32 deviation in the canaries (1.3e-2 max rel RMS vs the RTX 5090) is explained, and
+  it is not GPU precision. The official code builds its grid↔mesh graphs with NumPy on the host
+  CPU, and Modal's A100-40GB hosts (AVX2 only) build a slightly different graph than AVX-512
+  hosts. `cloud/canary_run.py` records the graph hashes and NumPy's SIMD targets, runs on
+  another machine's graph with `--edges`, and skips the rollout with `--steps 0`. New
+  `scripts/graph_edges.py` counts and saves the graphs on the CPU; new `cloud/modal_a100_tf32.py`
+  and `cloud/modal_graph_hosts.py` are the probes. Validation: docs/validation.md §4, A100 note.
+- README FAQ and `--no-autotune` help: no longer recommend `--no-autotune` for strict fp32
+  (`--precision highest` gives fp32-level results; autotuning was not the cause). The FAQ now
+  says why results can differ between machines.
+
 ## 0.2.0 — 2026-11-01
 
 Fine-tuning (experimental), IFS open-data mirror fallback, Intel Arc. Inference is unchanged:

README.md

@@ -168,7 +168,9 @@ No. Weights are the official fp32 checkpoints, attention is exact, and every lay
 <details>
 <summary><b>Why aren't results bit-identical between runs?</b></summary>
 
-XLA picks GPU kernels by timing them, so kernel choices — and the last bits of the result — vary between runs, with or without this project. Use `fwn forecast --autotune-cache kernels.txt` to save the choices on the first run and reuse them. For strict fp32, use `--precision highest --no-autotune`.
+XLA picks GPU kernels by timing them, so kernel choices — and the last bits of the result — vary between runs, with or without this project. Use `fwn forecast --autotune-cache kernels.txt` to save the choices on the first run and reuse them (`--no-autotune` also fixes them; neither is needed for fp32-level results, which `--precision highest` gives).
+
+Between machines, results can differ by more than that: the official code builds its grid↔mesh graph on the CPU with NumPy, and CPUs with and without AVX-512 get a few different edges. In one fp32 step that is up to 1.3 × 10⁻² relative RMS in the most sensitive field (median 6 × 10⁻⁴), around those edges ([details](docs/validation.md#4-other-gpus-modal-vastai)).
 </details>
 
 <details>

cloud/canary_run.py

@@ -6,12 +6,20 @@
 
   1. one step at precision "highest", saved on every 4th grid point (--npz) for cross-GPU
      comparison (scripts/compare_npz.py);
-  2. an N-step rollout at default precision: per-step wall time, finite check, pool peak.
+  2. an N-step rollout at default precision: per-step wall time, finite check, pool peak
+     (--steps 0 skips it, e.g. for precision probes).
+
+The JSON also records the grid<->mesh graphs (edge count and hash) and NumPy's SIMD targets:
+the official code builds the graphs on the host CPU, and the grid->mesh edge set depends on
+NumPy's code path there (docs/validation.md §4, A100 note). Compare fp32 outputs across machines
+only when the graph hashes match, or pass --edges (indices saved by scripts/graph_edges.py
+--save-edges on another machine) to run on that machine's graphs.
 
   python cloud/canary_run.py --data SAMPLE.nc --json out.json [--npz out.npz] [--steps 4]
 """
 import argparse
 import dataclasses
+import hashlib
 import json
 import os
 import time
@@ -24,6 +32,7 @@
 p.add_argument("--npz", default="")
 p.add_argument("--steps", type=int, default=4)
 p.add_argument("--label", default="")
+p.add_argument("--edges", default="", help="grid<->mesh edge indices from scripts/graph_edges.py --save-edges")
 args = p.parse_args()
 res = dict(label=args.label, env={k: v for k, v in os.environ.items() if k.startswith(("XLA_", "JAX_"))}, ok=False)
 
@@ -34,16 +43,33 @@
   import xarray
   import faster_weathernext
   from faster_weathernext import model, patch
+  from numpy._core._multiarray_umath import __cpu_dispatch__, __cpu_features__
   from weathernext.utils import data_utils
+  from weathernext.utils import points_mesh_gnn
+
+  res.update(cpu=next((l.split(":", 1)[1].strip() for l in open("/proc/cpuinfo") if l.startswith("model name")), ""),
+             numpy=np.__version__, numpy_simd=[f for f in __cpu_dispatch__ if __cpu_features__.get(f)], graphs={})
+  get_edge_indices = points_mesh_gnn.PointsMeshTypedGraphGNN._get_edge_indices
+
+  saved = dict(np.load(args.edges)) if args.edges else None
+
+  def record_edges(self, *a):
+    name = self._connectivity_type.name
+    edges = ((saved[f"{name}_points"].astype(int), saved[f"{name}_mesh"].astype(int)) if saved
+             else get_edge_indices(self, *a))
+    res["graphs"][name] = dict(n_edges=int(len(edges[0])), sha1=hashlib.sha1(
+        np.stack(edges).astype(np.int64).tobytes()).hexdigest()[:16], loaded=bool(saved))
+    return edges
 
+  points_mesh_gnn.PointsMeshTypedGraphGNN._get_edge_indices = record_edges
   faster_weathernext.enable()
   G = 2**30
   dev = jax.devices()[0]
   res.update(device=dev.device_kind, jax=jax.__version__, faster_weathernext=faster_weathernext.__version__,
              compute_capability=getattr(dev, "compute_capability", ""), attention=patch.attention_impl(),
              pool_gib=(dev.memory_stats() or {}).get("bytes_limit", 0) / G)
   print(f"{dev.device_kind} cc={res['compute_capability']} attention={res['attention']} "
-        f"pool {res['pool_gib']:.2f} GiB", flush=True)
+        f"pool {res['pool_gib']:.2f} GiB; host {res['cpu']}, NumPy SIMD {res['numpy_simd']}", flush=True)
   task, forward = model.build("WeatherNextCyclones")
   params = model.load_params("WeatherNextCyclones", 1, args.weights_dir)
   ds = xarray.load_dataset(args.data)
@@ -67,27 +93,29 @@
     print(f"highest-precision step saved ({len(fields)} fields) in {res['highest_first_call_s']:.0f}s", flush=True)
     del pred
 
-  lead = pd.to_timedelta(np.arange(1, args.steps + 1) * 6, unit="h")
-  tvars = {}
-  for v, da in targets1.data_vars.items():
-    shape = list(da.shape)
-    shape[da.dims.index("time")] = args.steps
-    tvars[v] = (da.dims, np.broadcast_to(np.float32(np.nan), shape))
-  targets = xarray.Dataset(tvars, coords={**{k: c for k, c in targets1.coords.items() if k != "time"}, "time": lead})
-  init = pd.Timestamp(ds["datetime"].values[0, 1]).to_pydatetime()
-  fds = xarray.Dataset(coords=dict(time=lead, lon=ds.lon, datetime=(("batch", "time"), np.array(
-      [[init + td for td in lead.to_pytimedelta()]], dtype="datetime64[ns]"))))
-  data_utils.add_derived_vars(fds)
-  forcings = fds[list(task.forcing_variables)].drop_vars("datetime")
-  forcings = forcings.astype({v: forcings1[v].dtype for v in forcings1.data_vars})
   res["step_s"], res["finite"] = [], []
-  t0 = time.time()
-  for pred in model.rollout(forward, params, jax.random.PRNGKey(0), inputs, targets, forcings):
-    msl = np.asarray(pred["mean_sea_level_pressure"])
-    res["step_s"].append(time.time() - t0)
-    res["finite"].append(bool(np.isfinite(msl).all()))
-    print(f"  step {len(res['step_s'])}: {res['step_s'][-1]:.1f}s  msl min {np.nanmin(msl) / 100:.1f} hPa", flush=True)
+  if args.steps > 0:
+    lead = pd.to_timedelta(np.arange(1, args.steps + 1) * 6, unit="h")
+    tvars = {}
+    for v, da in targets1.data_vars.items():
+      shape = list(da.shape)
+      shape[da.dims.index("time")] = args.steps
+      tvars[v] = (da.dims, np.broadcast_to(np.float32(np.nan), shape))
+    targets = xarray.Dataset(tvars, coords={**{k: c for k, c in targets1.coords.items() if k != "time"}, "time": lead})
+    init = pd.Timestamp(ds["datetime"].values[0, 1]).to_pydatetime()
+    fds = xarray.Dataset(coords=dict(time=lead, lon=ds.lon, datetime=(("batch", "time"), np.array(
+        [[init + td for td in lead.to_pytimedelta()]], dtype="datetime64[ns]"))))
+    data_utils.add_derived_vars(fds)
+    forcings = fds[list(task.forcing_variables)].drop_vars("datetime")
+    forcings = forcings.astype({v: forcings1[v].dtype for v in forcings1.data_vars})
     t0 = time.time()
+    for pred in model.rollout(forward, params, jax.random.PRNGKey(0), inputs, targets, forcings):
+      msl = np.asarray(pred["mean_sea_level_pressure"])
+      res["step_s"].append(time.time() - t0)
+      res["finite"].append(bool(np.isfinite(msl).all()))
+      print(f"  step {len(res['step_s'])}: {res['step_s'][-1]:.1f}s  msl min {np.nanmin(msl) / 100:.1f} hPa", flush=True)
+      t0 = time.time()
+  print(f"graphs: {res['graphs']}", flush=True)
   res["peak_gib"] = (dev.memory_stats() or {}).get("peak_bytes_in_use", 0) / G
   res["ok"] = all(res["finite"])
   print(f"peak {res['peak_gib']:.2f} GiB of pool {res['pool_gib']:.2f} GiB; ok={res['ok']}", flush=True)

cloud/modal_a100_tf32.py

@@ -0,0 +1,123 @@
+"""A100 fp32 probe: why does the canary's "highest" step land at TF32 level on A100-SXM4-40GB?
+
+Runs the canary's fp32 step (cloud/canary_run.py --steps 0) in fresh processes with different
+XLA settings on an A100-40GB and an A100-80GB. The base run on each card also dumps the autotune
+picks, the HLO, the thunk sequence and XLA's device description. Compare each .npz with an
+RTX 5090 run of the same code (scripts/compare_npz.py). The "graph" step saves the host's
+grid<->mesh graphs and "edges5090" runs the step on the RTX 5090's graphs instead (the cause:
+see docs/validation.md §4, A100 note).
+
+  modal run --detach cloud/modal_a100_tf32.py --tag tf32probe01
+  modal run --detach cloud/modal_a100_tf32.py --tag tf32probe02 --v40 graph,edges5090 --v80 ""
+  modal run --detach cloud/modal_a100_tf32.py --tag tf32probe03 --v40 "" --v80 numpy_avx2
+  modal volume get fwn-canary results/tf32probe01 ./tf32probe01
+  modal app stop -y <app id>      # detached apps can linger after finishing
+
+Cost (2026 prices): about $1 (A100-40GB ~18 min at ~$2.10/h, A100-80GB ~6 min at ~$2.50/h).
+"""
+import os
+import subprocess
+import time
+
+import modal
+
+REPO = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
+DATA = "source-hres_forecast_init-2024-10-07 00:00:00_res-0.25_levels-13_steps-01.nc"
+
+app = modal.App("fwn-a100-tf32")
+vol = modal.Volume.from_name("fwn-canary", 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"])
+    .run_commands("pip install '/opt/faster-weathernext[cuda12,weathernext]'")
+)
+
+# name -> (extra environment, dump autotune picks + HLO, extra canary_run.py arguments)
+VARIANTS = {
+    "base": ({}, True, []),
+    "at0": ({"XLA_FLAGS": "--xla_gpu_autotune_level=0"}, False, []),
+    "tf32off": ({"NVIDIA_TF32_OVERRIDE": "0"}, False, []),
+    "notritongemm": ({"XLA_FLAGS": "--xla_gpu_enable_triton_gemm=false"}, False, []),
+    # The RTX 5090's grid<->mesh graphs (scripts/graph_edges.py --save-edges, uploaded with
+    # `modal volume put fwn-canary edges_rtx5090.npz edges/edges_rtx5090.npz`).
+    "edges5090": ({}, False, ["--edges", "/data/edges/edges_rtx5090.npz"]),
+    # NumPy's AVX2 code path on an AVX-512 host: same host-side graph and features as the
+    # A100-40GB hosts (AVX2 only), so the GPUs are compared on identical inputs.
+    "numpy_avx2": ({"NPY_DISABLE_CPU_FEATURES": "X86_V4 AVX512_ICL AVX512_SPR"}, False, []),
+}
+
+
+def _run(card: str, tag: str, names: list[str]) -> dict:
+  import json
+  out = f"/data/results/{tag}"
+  os.makedirs(out, exist_ok=True)
+  smi = subprocess.run(["nvidia-smi", "--query-gpu=name,driver_version,memory.total,clocks.max.sm,ecc.mode.current",
+                        "--format=csv"], capture_output=True, text=True).stdout
+  summary = {"card": card, "nvidia_smi": smi,
+             "env": {k: v for k, v in os.environ.items() if any(s in k for s in ("TF32", "XLA", "JAX", "CUBLAS", "CUDNN"))}}
+  print(f"[{card}] {smi}", flush=True)
+  for name in names:
+    if name == "graph":  # this host's graphs, saved for reuse elsewhere (CPU only)
+      subprocess.run(["python", "/opt/faster-weathernext/scripts/graph_edges.py", "--weights-dir", "/data",
+                      "--data", f"/data/{DATA}", "--json", f"{out}/{card}_graph.json",
+                      "--save-edges", f"{out}/{card}_edges.npz"], env={**os.environ, "JAX_PLATFORMS": "cpu"})
+      summary[name] = json.load(open(f"{out}/{card}_graph.json"))
+      print(f"[{card}/{name}] {summary[name]['cpu']} {summary[name]['dispatch']} "
+            + "; ".join(f"{g['connectivity']} {g['n_edges']} {g['sha1']}" for g in summary[name]["graphs"]), flush=True)
+      vol.commit()
+      continue
+    extra, dump, extra_args = VARIANTS[name]
+    env = {**os.environ, "XLA_PYTHON_CLIENT_PREALLOCATE": "false", **extra}
+    if dump:
+      d = f"/tmp/dump_{card}_{name}"
+      env["XLA_FLAGS"] = (env.get("XLA_FLAGS", "") + f" --xla_dump_to={d} --xla_dump_hlo_as_text"
+                          " --xla_dump_hlo_module_re=.*lambda.*"
+                          f" --xla_gpu_dump_autotune_results_to={out}/{card}_{name}_autotune.txt").strip()
+    cmd = ["python", "/opt/faster-weathernext/cloud/canary_run.py", "--weights-dir", "/data",
+           "--data", f"/data/{DATA}", "--json", f"{out}/{card}_{name}.json", "--npz", f"{out}/{card}_{name}.npz",
+           "--steps", "0", "--label", f"{card}_{name}", *extra_args]
+    t0 = time.time()
+    proc = subprocess.run(cmd, env=env, capture_output=True, text=True)
+    with open(f"{out}/{card}_{name}.log", "w") as f:
+      f.write(proc.stdout + proc.stderr)
+    if dump:
+      # HLO before/after optimization, thunk sequence (which op runs as which kernel or library
+      # call), device description and autotune picks; not the large buffer-assignment files.
+      keep = r"optimizations.txt$|thunk_sequence.txt$|gpu_target_config.pbtxt$|autotune_results.pbtxt$|memory-usage-report.txt$"
+      subprocess.run(f"cd {d} && tar czf {out}/{card}_{name}_dump.tgz $(ls | grep -E '{keep}')", shell=True, check=False)
+    path = f"{out}/{card}_{name}.json"
+    summary[name] = json.load(open(path)) if os.path.exists(path) else {"error": f"exit {proc.returncode}"}
+    summary[name].pop("traceback", None)
+    summary[name]["wall_s"] = time.time() - t0
+    print(f"[{card}/{name}] exit {proc.returncode} device {summary[name].get('device')} "
+          f"{summary[name]['wall_s']:.0f}s {summary[name].get('error', '')}", flush=True)
+    vol.commit()
+  with open(f"{out}/{card}_summary.json", "w") as f:
+    json.dump(summary, f, indent=1)
+  vol.commit()
+  return summary
+
+
+@app.function(image=image, gpu="A100-40GB", volumes={"/data": vol}, timeout=3600, cpu=8, memory=32768)
+def a100_40(tag: str, names: list[str]):
+  return _run("a100_40", tag, names)
+
+
+@app.function(image=image, gpu="A100-80GB", volumes={"/data": vol}, timeout=3600, cpu=8, memory=32768)
+def a100_80(tag: str, names: list[str]):
+  return _run("a100_80", tag, names)
+
+
+@app.local_entrypoint()
+def main(tag: str = "tf32probe01", v40: str = "base,at0,tf32off,notritongemm", v80: str = "base"):
+  calls = []
+  if v40:
+    calls.append(a100_40.spawn(tag, v40.split(",")))
+  if v80:
+    calls.append(a100_80.spawn(tag, v80.split(",")))
+  for call in calls:
+    r = call.get()
+    print(f"== {r['card']}: " + ", ".join(f"{k}: {v.get('device')} ok={v.get('ok')} {v.get('error', '')}"
+                                         for k, v in r.items() if isinstance(v, dict) and k != "env"), flush=True)