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..446b614 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 "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 diff --git a/README.md b/README.md index 3bf7e13..4b28280 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) @@ -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 954805d..b2cf0b1 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,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= 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 -# 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,21 +119,30 @@ 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 \ @@ -124,16 +150,20 @@ RUN python3 -m pip install --no-cache-dir --prefer-binary \ 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, diff --git a/docker/build_natten.sh b/docker/build_natten.sh index bd98204..205205e 100644 --- a/docker/build_natten.sh +++ b/docker/build_natten.sh @@ -18,8 +18,8 @@ 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" @@ -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_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. 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( 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)