Fix/pin bump

.license-info/third_party_licenses.csv

@@ -5,7 +5,7 @@ ImageIO,2.37.3,BSD-2-Clause,https://github.com/imageio/imageio
 Jinja2,3.1.6,BSD License,https://github.com/pallets/jinja/
 Markdown,3.10.2,BSD-3-Clause,https://Python-Markdown.github.io/
 MarkupSafe,3.0.3,BSD-3-Clause,https://github.com/pallets/markupsafe/
-NATTEN,0.21.5,MIT,https://natten.org
+NATTEN,0.21.7,MIT,https://natten.org
 PyYAML,6.0.3,MIT License,https://pyyaml.org/
 Pygments,2.20.0,BSD-2-Clause,https://pygments.org
 Werkzeug,3.1.8,BSD-3-Clause,https://github.com/pallets/werkzeug/
@@ -118,7 +118,7 @@ numba,0.65.0,BSD License,https://numba.pydata.org
 numcodecs,0.16.5,MIT,https://github.com/zarr-developers/numcodecs
 numpy,2.4.4,BSD-3-Clause AND 0BSD AND MIT AND Zlib AND CC0-1.0,https://numpy.org
 nvdlfw_inspect,0.2.2,Apache2,https://github.com/NVIDIA/nvidia-dlfw-inspect
-nvidia-physicsnemo,2.2.0a0,Apache-2.0,https://github.com/NVIDIA/physicsnemo
+nvidia-physicsnemo,2.3.0a0,Apache-2.0,https://github.com/NVIDIA/physicsnemo
 oauthlib,3.3.1,BSD-3-Clause,https://github.com/oauthlib/oauthlib
 omegaconf,2.3.0,BSD License,https://github.com/omry/omegaconf
 onnx,1.21.0,Apache-2.0,https://onnx.ai/

CHANGELOG.md

@@ -4,6 +4,16 @@ All notable changes to this project will be documented in this file.
 
 ## [Unreleased]
 
+- Fix the slow pixel-stage depthwise conv (#6): pin physicsnemo main @
+  `fae6bb2b`, which contains NVIDIA/physicsnemo#2140. Its Strata
+  `DepthwiseConv` defaults to `chunk_size="auto"` instead of fixed
+  2-channel chunks, so every config that leaves
+  `use_chunked_depthwise_conv` at its default (including checkpoints that
+  predate the field) gets the fast path with no config or checkpoint change.
+  physicsnemo now requires torch >= 2.13, so the Docker base moves from NGC
+  PyTorch 26.01 to 26.09; NATTEN goes to 0.21.7 (built as C++20 for the
+  torch 2.14 headers) and earth2grid to v2025.11.1, which fixes NaNs in
+  lat/lon regridding west of the first source longitude.
 - The model architecture now comes from NVIDIA PhysicsNeMo
   (`physicsnemo.experimental.models.strata`); the local `dit_3d.py` /
   `dit_3d_pixel.py` implementations were replaced by thin wrappers

Makefile

@@ -8,10 +8,10 @@ install-linters:
 
 install:
 	pip install --no-build-isolation -r requirements.txt
-	pip install -f https://whl.natten.org "natten>=0.21.5"
-	pip install --no-deps --no-build-isolation https://github.com/NVlabs/earth2grid/archive/3415382a37e414a867d4c6ae2519e3e8afdbece1.tar.gz
-	pip install "warp-lang>=1.14" jaxtyping
-	pip install --no-deps https://github.com/NVIDIA/physicsnemo/archive/07cdcc8b65f10bee733a32a77b4b76e20c70c54c.tar.gz
+	pip install -f https://whl.natten.org "natten>=0.21.7"
+	pip install --no-deps --no-build-isolation https://github.com/NVlabs/earth2grid/archive/6e39b069f0558cd0eff0ec666ede5e125d79b4ff.tar.gz
+	pip install "warp-lang>=1.17" jaxtyping "GitPython>=3.1.62" "hydra-core>=1.3.7" "termcolor>=3.3.0" "treelib>=1.8.0"
+	pip install --no-deps https://github.com/NVIDIA/physicsnemo/archive/fae6bb2b0e96488a052d9627d41e6dc94653f06a.tar.gz
 
 lint:
 	python3 ci/check_licenses.py

README.md

@@ -29,7 +29,7 @@ The models need a recent NVIDIA GPU software stack that cannot come from PyPI
 alone: a CUDA-tuned PyTorch build, NVIDIA DALI (data pipelines), and a NATTEN
 build that matches the installed torch. The supported
 base environment is the [NGC PyTorch container](https://catalog.ngc.nvidia.com/orgs/nvidia/containers/pytorch)
-`nvcr.io/nvidia/pytorch:26.01-py3` (anonymous pulls generally work; a free
+`nvcr.io/nvidia/pytorch:26.09-py3` (anonymous pulls generally work; a free
 NGC account may be required depending on NGC policy).
 
 ### Docker (recommended)
@@ -39,9 +39,9 @@ Build the training/inference image from the repository root:
     docker build -f docker/Dockerfile -t strata .
 
 The image contains the complete environment, including a from-source NATTEN
-build (set `CUDA_ARCH` in [`docker/build_natten.sh`](docker/build_natten.sh)
-for non-Hopper GPUs) and PhysicsNeMo pinned to a commit that includes the
-Strata models.
+build for Ada and Hopper GPUs by default (pass `--build-arg CUDA_ARCH=...`
+for others, e.g. `8.0` on A100; see [`docker/build_natten.sh`](docker/build_natten.sh))
+and PhysicsNeMo pinned to a commit that includes the Strata models.
 
 ### pip install (plain PyPI environment)
 
@@ -144,10 +144,10 @@ export SCREAM_MAIN_ZARR_PATH=/path/to/sdecadal.ne1024pg2_ne1024pg2.F20TR-SCREAMv
 
 The forecast itself is a 4-GPU, 2-step `torchrun`, on whatever machine can
 launch it. Build `docker/Dockerfile` for your GPU (`--build-arg CUDA_ARCH=8.0`
-on A100, `9.0` on H100). `TORCH_COMPILE_DISABLE=1` keeps the first step from
-sitting in Inductor. PhysicsNeMo is installed with `--no-deps` in that image,
-so `treelib`, `gitpython`, `hydra-core`, and `termcolor` have to be
-importable or the process exits at startup.
+on A100; the default covers Ada and H100). `TORCH_COMPILE_DISABLE=1` keeps
+the first step from sitting in Inductor. PhysicsNeMo is installed with
+`--no-deps` in that image; its runtime imports (`treelib`, `gitpython`,
+`hydra-core`, `termcolor`) are installed explicitly.
 
 ```bash
 export SCREAM_AUX_ZARR_PATH=$SCREAM_MAIN_ZARR_PATH

docker/Dockerfile

@@ -24,7 +24,7 @@
 # Engine. Everything below is layered on top of it. Anonymous pulls of the
 # base image generally work; a free NGC account may be required depending on
 # NGC policy: https://catalog.ngc.nvidia.com
-FROM nvcr.io/nvidia/pytorch:26.01-py3
+FROM nvcr.io/nvidia/pytorch:26.09-py3
 
 ENV PIP_NO_CACHE_DIR=1 \
     PYTHONDONTWRITEBYTECODE=1 \
@@ -49,38 +49,53 @@ RUN apt-get update && apt-get install -y --no-install-recommends \
 
 RUN python3 -m pip install -U pip setuptools wheel
 
+# Pin the NGC torch / CUDA / numpy stack. NGC 26.09 ships an EMPTY
+# /etc/pip/constraint.txt, so without this any package that declares e.g.
+# torch>=2.13 would replace NGC's 2.14.0a0 pre-release torch with a PyPI
+# wheel. NGC's nvidia-dali-cuda130 (the data pipelines) requires
+# packaging<=26.2.
+RUN python3 -m pip list --format=freeze \
+    | grep -i -E '^(torch|torchvision|torchaudio|triton|numpy|transformer.engine|apex|flash.attn|cuda-|nvidia-|cupti)' \
+    > /etc/pip/strata-constraint.txt \
+ && echo "packaging<=26.2" >> /etc/pip/strata-constraint.txt \
+ && cat /etc/pip/strata-constraint.txt
+ENV PIP_CONSTRAINT=/etc/pip/strata-constraint.txt
+
 # Plain-pip requirements. NATTEN, earth2grid, and physicsnemo are installed
 # separately below because each needs flags a requirements file cannot
-# express. This step runs with NGC's PIP_CONSTRAINT still active so pip
-# cannot displace the container's torch stack.
+# express. PIP_CONSTRAINT stays active for every step so pip cannot displace
+# the container's torch stack.
 COPY requirements.txt /tmp/requirements.txt
 RUN python3 -m pip install --no-cache-dir --prefer-binary \
     -r /tmp/requirements.txt
 
 # NATTEN built from source: NGC PyTorch is a custom build, so the prebuilt
 # wheels on https://whl.natten.org do not include a matching libnatten.
-# Pass --build-arg CUDA_ARCH=<compute capability> for your GPU
-# (9.0=Hopper/H100 default; `nvidia-smi --query-gpu=compute_cap
-# --format=csv,noheader` prints it; see build_natten.sh).
-ARG CUDA_ARCH=9.0
+# Pass --build-arg CUDA_ARCH=<compute capabilities> for your GPUs, e.g.
+# "8.0" for A100 (`nvidia-smi --query-gpu=compute_cap --format=csv,noheader`
+# prints it; see build_natten.sh). The default covers Ada (RTX 6000 Ada,
+# L40S) and Hopper (H100). On a GPU the build does not cover, the first NATTEN
+# call aborts the whole process with "no kernel image is available for
+# execution on the device" (no Python exception, so pytest exits silently).
+ARG CUDA_ARCH="8.9;9.0"
 COPY docker/build_natten.sh /tmp/build_natten.sh
 RUN CUDA_ARCH=${CUDA_ARCH} bash /tmp/build_natten.sh
 
-# earth2grid is not on PyPI; install the pinned GitHub archive. --no-deps
-# and --no-build-isolation keep it from touching the torch stack.
+# earth2grid is not on PyPI; install the pinned GitHub archive (v2025.11.1).
+# --no-deps and --no-build-isolation keep it from touching the torch stack.
 RUN python3 -m pip install --no-cache-dir --prefer-binary --no-deps --no-build-isolation \
-    https://github.com/NVlabs/earth2grid/archive/3415382a37e414a867d4c6ae2519e3e8afdbece1.tar.gz
+    https://github.com/NVlabs/earth2grid/archive/6e39b069f0558cd0eff0ec666ede5e125d79b4ff.tar.gz
 
-# NGC pins zarr<3 via PIP_CONSTRAINT; the datasets are zarr v3. Clear the
-# constraint for this and the remaining steps.
-ENV PIP_CONSTRAINT=""
+# The datasets are zarr v3. (NGC 26.01 pinned zarr<3 in its constraint file;
+# the strata-constraint.txt above does not constrain zarr.)
 RUN python3 -m pip install --no-cache-dir --upgrade "zarr>=3"
 
 # fme (AI2 Full Model Emulator) and earth2studio power the rollout/inference
 # wrappers. Their runtime deps are installed first so the --no-deps install
 # below cannot skip them; --no-deps stops pip from re-resolving the
 # torch/CUDA/natten stack already baked into the image.
-# torch-harmonics is pinned to 0.8.0 by fme==2026.1.1.
+# torch-harmonics is pinned to 0.8.0 by fme==2026.1.1. earth2studio 0.13.0
+# imports nest_asyncio and requires pandas<3 (NGC 26.09 ships pandas 3).
 RUN python3 -m pip install --no-cache-dir --prefer-binary \
     dacite \
     dask \
@@ -91,7 +106,9 @@ RUN python3 -m pip install --no-cache-dir --prefer-binary \
     imageio_ffmpeg \
     loguru \
     moviepy \
+    nest-asyncio \
     omegaconf \
+    "pandas<3" \
     plotly \
     pygrib \
     "setuptools<81" \
@@ -102,38 +119,51 @@ RUN python3 -m pip install --no-cache-dir --prefer-binary --no-deps \
     fme==2026.1.1 \
     earth2studio==0.13.0
 
-# marimo runs the notebooks/ examples.
+# marimo runs the notebooks/ examples. It needs a newer Pygments than the
+# apt-installed one, which has no pip RECORD file and cannot be uninstalled;
+# shadow it in site-packages instead.
+RUN python3 -m pip install --no-cache-dir --ignore-installed pygments
 RUN python3 -m pip install --no-cache-dir --prefer-binary marimo
 
-# physicsnemo pinned to main @ 07cdcc8b (includes the Strata models and the
-# continuous/stereographic RoPE helpers from PR #1750; pre-2.2, so PyPI does
-# not have it yet). --no-deps is REQUIRED: physicsnemo declares
-# torch>=2.10.0, and the NGC torch is a PEP 440 PRE-release that pip
-# considers older, so a plain install would replace the container's torch.
+# physicsnemo pinned to main @ fae6bb2b (2.3.0a0). It contains PR #2140
+# (merge b26a8805): Strata's depthwise conditioning conv defaults to
+# chunk_size="auto" instead of fixed 2-channel vmap chunks, which fixes the
+# slow pixel-stage conv of NVlabs/STRATA#6 for every config that leaves
+# use_chunked_depthwise_conv at its default. No physicsnemo release contains
+# it yet, and every commit that does requires torch>=2.13 (hence NGC 26.09).
+# --no-deps is REQUIRED: physicsnemo declares torch>=2.13.0, and the NGC
+# torch is a PEP 440 PRE-release that pip considers older, so a plain
+# install would replace the container's torch.
 # The import of physicsnemo.experimental.models.strata transitively pulls in
 # all of physicsnemo.nn, which needs warp, tensordict and timm at import
-# time; tensordict/timm declare their own torch/torchvision floors, so they
+# time (and, since the 2.3 dev line, treelib and termcolor via
+# physicsnemo.distributed; GitPython and hydra-core are needed by the
+# training entry points); tensordict/timm declare their own torch/torchvision floors, so they
 # get the same --no-deps treatment, with their pure-Python runtime deps
 # installed explicitly first.
 RUN python3 -m pip install --no-cache-dir --prefer-binary \
-    "warp-lang>=1.14" \
+    "warp-lang>=1.17" \
     jaxtyping \
     cloudpickle \
     importlib_metadata \
     packaging \
     pyvers \
     huggingface_hub \
     safetensors \
-    pyyaml
+    pyyaml \
+    "GitPython>=3.1.62" \
+    "hydra-core>=1.3.7" \
+    "termcolor>=3.3.0" \
+    "treelib>=1.8.0"
 # NB: tensordict declares orjson, but only as an optional accelerated JSON
 # backend imported lazily (stdlib json is the registered default) — verified
 # import-clean without it, so it is deliberately NOT installed (MPL-2.0).
 RUN python3 -m pip install --no-cache-dir --prefer-binary --no-deps \
-    tensordict \
-    timm
+    "tensordict>=0.15.3,<0.16.0" \
+    "timm>=1.0.30"
 RUN python3 -m pip uninstall -y nvidia-physicsnemo || true
 RUN python3 -m pip install --no-cache-dir --no-deps \
-    https://github.com/NVIDIA/physicsnemo/archive/07cdcc8b65f10bee733a32a77b4b76e20c70c54c.tar.gz
+    https://github.com/NVIDIA/physicsnemo/archive/fae6bb2b0e96488a052d9627d41e6dc94653f06a.tar.gz
 
 # The NGC image chain has shipped triton with an empty backends/nvidia/bin
 # (no ptxas / cuobjdump / nvdisasm -> "RuntimeError: Cannot find ptxas") and
@@ -153,14 +183,20 @@ import typer  # noqa: F401
 import warp
 from earth2grid import healpix_bare  # noqa: F401
 
+import inspect
+
 import physicsnemo
 
 assert natten.HAS_LIBNATTEN, "natten built without libnatten"
-assert physicsnemo.__version__.startswith("2.2"), physicsnemo.__version__
+assert physicsnemo.__version__.startswith("2.3"), physicsnemo.__version__
 from physicsnemo.experimental.models.strata import (  # noqa: F401
     Strata,
     StrataTransformer3D,
 )
+
+# physicsnemo PR #2140 (NVlabs/STRATA#6): automatic depthwise-conv chunking.
+default = inspect.signature(Strata.__init__).parameters["chunk_size_grouped_conv"].default
+assert default == "auto", default
 from physicsnemo.experimental.nn import (  # noqa: F401
     build_axial_rope_cos_sin_2d_continuous,
     spherical_centroid,

docker/build_natten.sh

@@ -18,21 +18,31 @@ set -e
 # Build NATTEN from source against the torch in the current environment.
 # Required inside NGC PyTorch containers: their torch is a custom build, so
 # the prebuilt wheels on https://whl.natten.org ship no matching libnatten.
-export NATTEN_VERSION="${NATTEN_VERSION:-0.21.5}"
-export CUDA_ARCH="${CUDA_ARCH:-9.0}"  # 9.0=Hopper(H100), 10.0=Blackwell(B200), 8.0=Ampere(A100)
+export NATTEN_VERSION="${NATTEN_VERSION:-0.21.7}"
+export CUDA_ARCH="${CUDA_ARCH:-8.9;9.0}"  # ;-separated: 8.0=Ampere(A100), 8.9=Ada, 9.0=Hopper(H100), 10.0=Blackwell(B200)
 export NATTEN_CUDA_ARCH="${CUDA_ARCH}"
 export NATTEN_N_WORKERS="${NATTEN_N_WORKERS:-16}"
 export NATTEN_WITH_CUDA="1"
 export NATTEN_VERBOSE="1"
 
 echo "Building NATTEN ${NATTEN_VERSION} for CUDA architectures: ${CUDA_ARCH}"
 
-git clone --recursive https://github.com/SHI-Labs/NATTEN /opt/natten
-cd /opt/natten
-git checkout "v${NATTEN_VERSION}"
-make clean uninstall || true
-make
-pip install --no-deps --no-build-isolation -v -e . 2>&1
+# Non-editable install from the PyPI sdist; the source tree is removed after.
+src=$(mktemp -d)
+cd "${src}"
+python -m pip download --no-deps --no-binary natten --no-build-isolation \
+    "natten==${NATTEN_VERSION}" -d .
+tar xzf "natten-${NATTEN_VERSION}.tar.gz"
+cd "natten-${NATTEN_VERSION}"
+# torch >= 2.13 headers require C++20 ("#error C++20 or later compatible
+# compiler is required to use ATen"), but NATTEN hard-codes C++17 in
+# csrc/CMakeLists.txt (the CXX_STD default and the host -std flag) and its
+# setup.py cannot pass -DCXX_STD, so patch both before building.
+sed -i -e 's/set(CXX_STD "17"/set(CXX_STD "20"/' -e 's/-std=c++17/-std=c++20/' csrc/CMakeLists.txt
+grep -q 'CXX_STD "20"' csrc/CMakeLists.txt
+python -m pip install --no-deps --no-build-isolation -v . 2>&1
+cd /
+rm -rf "${src}"
 
 echo "NATTEN build complete. Verifying libnatten..."
-python -c "import natten; print('HAS_LIBNATTEN:', natten.HAS_LIBNATTEN)"
+python -c "import natten; print('HAS_LIBNATTEN:', natten.HAS_LIBNATTEN); assert natten.HAS_LIBNATTEN"

pyproject.toml

@@ -12,15 +12,15 @@ readme = "README.md"
 requires-python = ">=3.11"
 license = {text = "Apache 2.0"}
 dependencies = [
-    "torch>=2.10.0",
+    "torch>=2.13.0",
     "numpy>=2.1",
     "pandas>=2.3",
     "matplotlib>=3.10",
     "xarray>=2026.7",
     "zarr>=3",
     # not on PyPI: direct reference to the commit the release was verified
     # against (same archive docker/Dockerfile installs)
-    "earth2grid @ https://github.com/NVlabs/earth2grid/archive/3415382a37e414a867d4c6ae2519e3e8afdbece1.tar.gz",
+    "earth2grid @ https://github.com/NVlabs/earth2grid/archive/6e39b069f0558cd0eff0ec666ede5e125d79b4ff.tar.gz",
     "earth2studio>=0.13",
     "einops>=0.8",
     "scipy>=1.16",
@@ -30,7 +30,7 @@ dependencies = [
     "python-dotenv>=1.2",
     "pyyaml>=6.0",
     "nvtx>=0.2.14",
-    "natten>=0.21.5",
+    "natten>=0.21.7",
     "typer>=0.26",
     "tqdm>=4.67",
     "torchmetrics>=1.9.0",
@@ -41,9 +41,10 @@ dependencies = [
     "pytorch-optimizer>=3.10.0",
     "tensorboard>=2.20.0",
     "lmdb>=2.2.0",
-    # >=2.2.0a0: the Strata models ship in 2.2 (PyPI 2.1.x would import-fail);
-    # the pinned archive below provides 2.2.0a0.
-    "nvidia-physicsnemo>=2.2.0a0",
+    # >=2.3.0a0: the pinned archive below (2.3.0a0) contains physicsnemo
+    # PR #2140, the automatic depthwise-conv chunking that fixes
+    # NVlabs/STRATA#6; 2.2.x and earlier use the slow fixed 2-channel chunks.
+    "nvidia-physicsnemo>=2.3.0a0",
     "h5netcdf>=1.8",
     "netCDF4>=1.7",
     "h5py>=3.16",
@@ -113,6 +114,8 @@ exclude = ["examples"]
 "tests/*.py" = ["S101"]
 
 [tool.uv.sources]
-# Pinned to physicsnemo main including the Strata models + stereographic RoPE
-# (PR #1750). Switch to `nvidia-physicsnemo>=2.2` on PyPI once 2.2.0 is released.
-nvidia-physicsnemo = { url = "https://github.com/NVIDIA/physicsnemo/archive/07cdcc8b65f10bee733a32a77b4b76e20c70c54c.tar.gz" }
+# Pinned to physicsnemo main @ fae6bb2b: the Strata models, stereographic RoPE
+# (PR #1750) and automatic depthwise-conv chunking (PR #2140, NVlabs/STRATA#6).
+# No release contains #2140 yet (latest is 2.2.2); switch to a PyPI release
+# once one does. Every commit containing #2140 requires torch>=2.13.
+nvidia-physicsnemo = { url = "https://github.com/NVIDIA/physicsnemo/archive/fae6bb2b0e96488a052d9627d41e6dc94653f06a.tar.gz" }

requirements.txt

@@ -5,7 +5,7 @@
 #                 or the source build in docker/build_natten.sh on NGC images)
 #   - earth2grid (not on PyPI; pinned GitHub archive, --no-deps)
 #   - nvidia-physicsnemo (pinned GitHub archive with the Strata models;
-#                 --no-deps so its torch>=2.10.0 floor cannot replace a
+#                 --no-deps so its torch>=2.13.0 floor cannot replace a
 #                 container's pre-release torch build)
 # lightning >=2.6.6 everywhere: earlier releases (through 2.6.5) import and
 # execute arbitrary module paths from a checkpoint's _instantiator

screamcast/strata_wrappers.py

@@ -626,7 +626,7 @@ def set_tile_size(self, height: int, width: int) -> None:
         physicsnemo ``Strata`` has no re-tiling hook of its own (documented
         gap), so this reaches into its pixel RoPE buffers to rebuild them for
         ``rope_mode_pixel="axial"``. Mirrors the construction-time code in
-        physicsnemo strata.py (pinned commit 07cdcc8b); revisit on pin bumps.
+        physicsnemo strata.py (pinned commit fae6bb2b); revisit on pin bumps.
         """
         self.strata.backbone.set_tile_size(height, width)
         self.strata.height = height