[zyhu-hu]

Fix the slow pixel-stage depthwise conv (#6): pin physicsnemo fae6bb2b

Fixes #6.

What changes

physicsnemo main @ fae6bb2b contains NVIDIA/physicsnemo#2140: Strata's
DepthwiseConv now defaults to chunk_size="auto" instead of fixed
2-channel vmap chunks. STRATA only passes use_chunked_depthwise_conv
(never a chunk size), so every config that leaves it at the default,
including checkpoints that predate the field, gets the fast path. No
config, model code or checkpoint change is needed.

Every physicsnemo commit that contains #2140 requires torch >= 2.13, so the
environment moves too:

	before	after	why
NGC PyTorch base	26.01 (torch 2.10)	26.09 (torch 2.14)	physicsnemo needs torch >= 2.13
physicsnemo	07cdcc8b (2.2.0a0)	fae6bb2b (2.3.0a0)	contains #2140
NATTEN	0.21.5 (git, editable)	0.21.7 (PyPI sdist, C++20 patch)	torch 2.14 headers require C++20; NATTEN hard-codes C++17
earth2grid	3415382a (2025.7.1)	6e39b069 (v2025.11.1)	the old commit returns NaN for lat/lon targets west of the first source longitude (test_latlon_to_point_grid_regridder_keeps_seam_and_out_of_bounds_lat_targets_finite fails on main)

Dockerfile details for NGC 26.09:

• 26.09 ships an empty /etc/pip/constraint.txt. The Dockerfile writes
  /etc/pip/strata-constraint.txt (the NGC torch/CUDA/numpy stack plus
  packaging<=26.2 for DALI) and keeps it active for every step. 26.01's
  zarr<3 constraint is gone, so it no longer needs clearing.
• physicsnemo 2.3 imports treelib and termcolor (via
  physicsnemo.distributed); tensordict>=0.15.3,<0.16 and timm>=1.0.30
  follow its pyproject.
• earth2studio 0.13 imports nest_asyncio and requires pandas<3 (26.09
  ships pandas 3).
• marimo needs the pip Pygments (the apt one has no RECORD file).
• The smoke test asserts physicsnemo 2.3 and chunk_size_grouped_conv
  defaulting to "auto".

pyproject.toml (torch, natten, physicsnemo floors and the uv source),
Makefile, the README base image, the license inventory and a code comment
pinned to 07cdcc8b are updated to match.

Test added

tests/test_dit_3d_pixel.py::test_factory_default_depthwise_conv_uses_auto_chunking
builds the model from the default PixelDiTConfig (on the meta device) and
checks that every pixel-block conv is a DepthwiseConv with
chunk_size == "auto". It fails on main's pin (chunk_size == 2) and
passes on this branch.

Performance and numerics (H100 80GB HBM3)

Full network forward, t128 production checkpoint, BF16 autocast. Time is
the median ms per forward; kernels is the number of GPU kernel launches in
one forward:

tile x batch	ms: old pin (chunk 2)	ms: auto	ms: grouped nn.Conv2d	chunk2 / auto	kernels: chunk 2	kernels: auto
128 x 1	405.8	235.4	233.2	1.72x	14,071	1,787
128 x 4	1016.3	930.8	923.0	1.09x	20,312	1,912
160 x 1	506.1	365.3	362.0	1.39x	16,131	1,807
256 x 1	1006.9	919.9	910.8	1.09x	20,275	1,871
320 x 1	1586.8	1478.2	fails (32-bit index limit)	1.07x	28,503	1,955

The kernel ratio (8-15x) is much larger than the time ratio because the
extra kernels are small: each adds about 4-14 us, while attention and MLPs,
which are the same in all three columns, dominate the forward. At 128 x 1
the ~12,300 extra kernels add 170 ms to a 235 ms forward (1.72x); at larger
tiles the rest of the network grows faster than the overhead (1.07-1.09x).
On faster GPUs the same overhead is a larger share (GB300: up to 4.5x).

• Isolated conv, BF16 training step at 128 x 1: 765 ms (chunk 2) vs 47 ms
  (auto) vs 45 ms (grouped).
• Outputs: FP32 (TF32 off) relative RMS 1.7e-6 between implementations at
  every size; auto and chunk 2 give identical outputs.
• GB300 shows the same picture: auto within 2% of grouped, old pin
  1.6-4.5x slower.

Keeping use_chunked_depthwise_conv=True as the default is deliberate:
grouped nn.Conv2d fails at or above 2**31 conv elements (tile 160 x
batch 4, tile 320 x batch 1).

Also in this PR

• Default CUDA_ARCH is now "8.9;9.0" (was 9.0). With a 9.0-only
  NATTEN build, the first NATTEN call on an Ada GPU aborts the whole
  process with "no kernel image is available for execution on the device".
  There is no Python exception, so pytest exits with no summary. On an RTX
  6000 Ada this kills test_dit_3d, test_dit_3d_pixel and the
  characterization tests in main's image.
• GitPython and hydra-core are installed with the other physicsnemo
  runtime imports. The README no longer tells users to add them.
• test_characterization_inference disables TF32 for its forward pass. NGC
  images enable TF32, which moves these outputs by about 2e-4, above the
  1e-4 tolerance, so goldens generated in one TF32 mode failed in the other.

Verification

• make lint passes.
• docker build -f docker/Dockerfile . builds with the new default; NATTEN
  has sm_89 and sm_90a kernels, and the smoke test passes.
• Test suite in the new image on an RTX 6000 Ada, with the SCREAM data
  (S3), the aux data and SCREAMCAST_CHECKPOINT set to the t128 production
  checkpoint, --ignore=tests/distributed: 229 passed, 2 skipped. The
  2 skips are the multi-GPU tests in test_distributed.py.
  tests/distributed (8-GPU torchrun) was not run.
• test_from_checkpoint needed a fix to run at all: it passed an empty
  coords["time"], which coszr checkpoints reject. It now sets one
  timestamp, as the wrapper tests do.
• The earth2grid seam test fails on main and passes here.
• Characterization goldens were generated from main's code with
  physicsnemo 07cdcc8b, in the same image with TF32 off. All three
  production configs (sweep1 DiT3D and both pixeldit configs) reproduce
  within the test's 1e-4 tolerance with physicsnemo fae6bb2b. So the pin
  bump does not change production inference.
• The production-checkpoint numbers above come from screamcast's
  scripts/debug/bench_depthwise_conv.py and from_checkpoint tests on the
  equivalent screamcast image, not from this repo's test suite.

[nbren12]

Thanks for doing this. Merging