From f01311d7ac2e8b3a8f21a3f3ce6cb60418b52a58 Mon Sep 17 00:00:00 2001 From: Zeyuan Hu Date: Thu, 29 Oct 2026 15:02:45 -0700 Subject: [PATCH 1/3] build: pin physicsnemo fae6bb2b to fix the slow depthwise conv (#6) physicsnemo main @ fae6bb2b contains NVIDIA/physicsnemo#2140: Strata's DepthwiseConv defaults to chunk_size="auto" instead of fixed 2-channel vmap chunks. 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. On H100 (t128 checkpoint, BF16 forward) the old pin is 1.07-1.72x slower than auto and launches ~8x more kernels; auto is within 1% of grouped nn.Conv2d and, unlike it, runs above 2**31 conv elements. Outputs are unchanged (FP32 rel RMS 1.7e-6). Every physicsnemo commit containing #2140 requires torch>=2.13, so: - docker: NGC PyTorch 26.01 -> 26.09 (torch 2.14). 26.09 ships an empty pip constraint file, so write one that pins the torch/CUDA/numpy stack (plus packaging<=26.2 for DALI) and keep it active for every step. - NATTEN 0.21.5 -> 0.21.7, built from the PyPI sdist with csrc patched to C++20 (torch 2.14 headers require it; NATTEN hard-codes C++17). - physicsnemo now imports treelib and termcolor; tensordict<0.16 and timm>=1.0.30 per its pyproject. earth2studio needs nest-asyncio and pandas<3 (26.09 ships pandas 3); marimo needs the pip Pygments. - earth2grid 3415382a -> v2025.11.1 (6e39b069): the old commit returns NaN for lat/lon targets west of the first source longitude, failing test_latlon_to_point_grid_regridder_keeps_seam_and_out_of_bounds_lat_targets_finite. - pyproject (torch>=2.13, natten>=0.21.7, physicsnemo>=2.3.0a0, uv source), Makefile, README base image, license inventory and the smoke test follow. Adds test_factory_default_depthwise_conv_uses_auto_chunking: fails on the old pin (chunk_size == 2), passes on fae6bb2b. Signed-off-by: Zeyuan Hu --- .license-info/third_party_licenses.csv | 4 +- CHANGELOG.md | 10 ++++ Makefile | 8 +-- README.md | 2 +- docker/Dockerfile | 76 ++++++++++++++++++-------- docker/build_natten.sh | 26 ++++++--- pyproject.toml | 21 ++++--- requirements.txt | 2 +- screamcast/strata_wrappers.py | 2 +- tests/test_dit_3d_pixel.py | 33 +++++++++++ 10 files changed, 135 insertions(+), 49 deletions(-) diff --git a/.license-info/third_party_licenses.csv b/.license-info/third_party_licenses.csv index ae01398..8c204c5 100644 --- a/.license-info/third_party_licenses.csv +++ b/.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/ diff --git a/CHANGELOG.md b/CHANGELOG.md index 90bad84..b0dcc2a 100644 --- a/CHANGELOG.md +++ b/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 diff --git a/Makefile b/Makefile index a9778cb..b8c99cf 100644 --- a/Makefile +++ b/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 "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 diff --git a/README.md b/README.md index 3bf7e13..0084a71 100644 --- a/README.md +++ b/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) diff --git a/docker/Dockerfile b/docker/Dockerfile index 954805d..cf9fe61 100644 --- a/docker/Dockerfile +++ b/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,10 +49,22 @@ 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 @@ -66,21 +78,21 @@ ARG CUDA_ARCH=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 +103,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,21 +116,29 @@ 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); 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 \ @@ -124,16 +146,18 @@ RUN python3 -m pip install --no-cache-dir --prefer-binary \ pyvers \ huggingface_hub \ safetensors \ - pyyaml + pyyaml \ + "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 +177,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, diff --git a/docker/build_natten.sh b/docker/build_natten.sh index bd98204..5c55291 100644 --- a/docker/build_natten.sh +++ b/docker/build_natten.sh @@ -18,7 +18,7 @@ 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 NATTEN_VERSION="${NATTEN_VERSION:-0.21.7}" export CUDA_ARCH="${CUDA_ARCH:-9.0}" # 9.0=Hopper(H100), 10.0=Blackwell(B200), 8.0=Ampere(A100) export NATTEN_CUDA_ARCH="${CUDA_ARCH}" export NATTEN_N_WORKERS="${NATTEN_N_WORKERS:-16}" @@ -27,12 +27,22 @@ 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" diff --git a/pyproject.toml b/pyproject.toml index 5a6b914..a9c836b 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -12,7 +12,7 @@ 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", @@ -20,7 +20,7 @@ dependencies = [ "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" } diff --git a/requirements.txt b/requirements.txt index b03cb13..25eff80 100644 --- a/requirements.txt +++ b/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 diff --git a/screamcast/strata_wrappers.py b/screamcast/strata_wrappers.py index e8369b6..089b0b2 100644 --- a/screamcast/strata_wrappers.py +++ b/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 diff --git a/tests/test_dit_3d_pixel.py b/tests/test_dit_3d_pixel.py index a406748..bc21c15 100644 --- a/tests/test_dit_3d_pixel.py +++ b/tests/test_dit_3d_pixel.py @@ -279,6 +279,39 @@ def test_factory_translation_production_shape(): assert model._index_is_latlon +def test_factory_default_depthwise_conv_uses_auto_chunking(): + """The default pixel config builds automatically chunked depthwise convs. + + NVlabs/STRATA#6: checkpoints that predate ``use_chunked_depthwise_conv`` + load with its default (True), i.e. physicsnemo's chunked DepthwiseConv. + That is fast only with physicsnemo's ``chunk_size="auto"`` (PR #2140); + the fixed 2-channel chunks of earlier pins launch ~8x more kernels per + forward. Check the built modules, not the config field. + """ + from physicsnemo.experimental.models.strata.layers import DepthwiseConv + + import train as train_module + from screamcast.config import DiTConfig, PixelDiTConfig + + with torch.device("meta"): + model = train_module.build_strata( + 6, + 6, + 1024, + 32, + dit_cfg=DiTConfig(embed_dim=64, n_layers=1, num_heads=2), + pixel_cfg=PixelDiTConfig(use_bilinear_dw_gelu_project_adaln=True), + do_bf16_mixed=False, + depth_levels=4, + grid_type="healpix", + ) + convs = [block.adaln_bilinear_dw_conv for block in model.strata.pixel_blocks] + assert convs + for conv in convs: + assert isinstance(conv, DepthwiseConv) + assert conv.chunk_size == "auto" + + def test_dit_pixel_freeze_pixel_blocks(): D, H, W = 24, 64, 64 model = StrataModel( From 4e165faa9f8cf07a6c2eaceebc64a45689ac5ab0 Mon Sep 17 00:00:00 2001 From: Zeyuan Hu Date: Thu, 29 Oct 2026 15:37:11 -0700 Subject: [PATCH 2/3] build: default NATTEN to Ada+Hopper; test: run characterization in fp32 docker: default CUDA_ARCH to "8.9;9.0" (was 9.0). A 9.0-only NATTEN build aborts the whole process on an Ada GPU with "no kernel image is available for execution on the device" -- no Python exception, so pytest exits with no summary (test_dit_3d, test_dit_3d_pixel and the characterization tests on an RTX 6000 Ada). Also install GitPython and hydra-core with the other physicsnemo runtime imports, so the README no longer has to tell users to add them before launching. tests: test_characterization_inference disables TF32 for the forward and restores it afterwards. NGC images enable TF32, which moves these outputs by ~2e-4 (above ATOL), so goldens generated in one TF32 mode failed in the other. With goldens generated from main (physicsnemo 07cdcc8b), all three production configs reproduce on this branch (physicsnemo fae6bb2b). Signed-off-by: Zeyuan Hu --- Makefile | 2 +- README.md | 14 +++++++------- docker/Dockerfile | 16 +++++++++++----- docker/build_natten.sh | 2 +- tests/test_characterization_inference.py | 18 +++++++++++++++++- 5 files changed, 37 insertions(+), 15 deletions(-) diff --git a/Makefile b/Makefile index b8c99cf..446b614 100644 --- a/Makefile +++ b/Makefile @@ -10,7 +10,7 @@ install: pip install --no-build-isolation -r requirements.txt 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 "termcolor>=3.3.0" "treelib>=1.8.0" + 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: diff --git a/README.md b/README.md index 0084a71..4b28280 100644 --- a/README.md +++ b/README.md @@ -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 diff --git a/docker/Dockerfile b/docker/Dockerfile index cf9fe61..b2cf0b1 100644 --- a/docker/Dockerfile +++ b/docker/Dockerfile @@ -71,10 +71,13 @@ RUN python3 -m pip install --no-cache-dir --prefer-binary \ # 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= 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= 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 @@ -134,7 +137,8 @@ RUN python3 -m pip install --no-cache-dir --prefer-binary marimo # The import of physicsnemo.experimental.models.strata transitively pulls in # all of physicsnemo.nn, which needs warp, tensordict and timm at import # time (and, since the 2.3 dev line, treelib and termcolor via -# physicsnemo.distributed); tensordict/timm declare their own torch/torchvision floors, so they +# 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 \ @@ -147,6 +151,8 @@ RUN python3 -m pip install --no-cache-dir --prefer-binary \ huggingface_hub \ safetensors \ 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 diff --git a/docker/build_natten.sh b/docker/build_natten.sh index 5c55291..205205e 100644 --- a/docker/build_natten.sh +++ b/docker/build_natten.sh @@ -19,7 +19,7 @@ set -e # 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.7}" -export CUDA_ARCH="${CUDA_ARCH:-9.0}" # 9.0=Hopper(H100), 10.0=Blackwell(B200), 8.0=Ampere(A100) +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" diff --git a/tests/test_characterization_inference.py b/tests/test_characterization_inference.py index e374509..58ce4ba 100644 --- a/tests/test_characterization_inference.py +++ b/tests/test_characterization_inference.py @@ -38,7 +38,8 @@ *by sorted name* after construction, so identical (name, shape) param sets produce identical weights regardless of construction order. Only a genuine change to the forward math can move the output. -* **fp32 / eval / fixed input.** Deterministic forward, no bf16, no dropout. +* **fp32 / eval / fixed input.** Deterministic forward, no bf16, no dropout, + and TF32 disabled for the test (see ``_strict_fp32``). Usage ----- @@ -92,6 +93,21 @@ ) +@pytest.fixture(autouse=True) +def _strict_fp32(): + """Run the forward in true fp32, whatever the environment's TF32 default. + + NGC PyTorch images enable TF32 matmuls/convs, which moves these outputs by + ~2e-4, above ATOL; goldens generated with TF32 off would then fail in an + image with it on, and vice versa. Restores the previous flags afterwards. + """ + matmul, cudnn = torch.backends.cuda.matmul, torch.backends.cudnn + saved = matmul.allow_tf32, cudnn.allow_tf32 + matmul.allow_tf32 = cudnn.allow_tf32 = False + yield + matmul.allow_tf32, cudnn.allow_tf32 = saved + + def _fill_params_deterministically(model: torch.nn.Module) -> torch.nn.Module: """Overwrite every parameter with a deterministic value keyed by sorted name. From d38457614e25dfd4f453c940bb71590cd49ef764 Mon Sep 17 00:00:00 2001 From: Zeyuan Hu Date: Thu, 29 Oct 2026 16:00:30 -0700 Subject: [PATCH 3/3] test: give test_from_checkpoint a timestamp for coszr checkpoints Production checkpoints with coszr forcing compute it from time + lead_time; the test passed an empty coords["time"], so the forward raised "coords['time'] must contain at least one timestamp to compute 'coszr'" before reaching the model. Set one timestamp, as the wrapper tests do. Signed-off-by: Zeyuan Hu --- tests/test_earth2studio.py | 2 ++ 1 file changed, 2 insertions(+) diff --git a/tests/test_earth2studio.py b/tests/test_earth2studio.py index 0d7c6ac..837c76a 100644 --- a/tests/test_earth2studio.py +++ b/tests/test_earth2studio.py @@ -424,6 +424,8 @@ def test_from_checkpoint(): coords = OrderedDict(in_coords) coords["face"] = np.array([0]) coords["batch"] = np.empty(2) + # Checkpoints with coszr forcing compute it from time + lead_time. + coords["time"] = np.array([np.datetime64("2020-10-13T00:00:00")]) out, _ = model(x, coords)