Adds DiffusionUNet3D and its modules to physicsnemo.experimental

[CharlelieLrt]

PhysicsNeMo Pull Request

Description

Replaces #1220 from @AbVishwas, which had diverged too much from main.

This PR enables training and sampling diffusion models on 3D structured domains (volumetric or grid-based scientific data) with the addition of a DiffusionUNet3D backbone, which satisfies the DiffusionModel protocol and plugs straight into the existing preconditioners, losses, and samplers in physicsnemo.diffusion. Conditioning is exposed through a single TensorDict argument with optional "vector" and "volume" keys. The PR also adds reusable 3D building blocks (Conv3D, GroupNorm3D, UNetAttention3D, UNetBlock3D) under physicsnemo.experimental.nn that can be composed for other 3D architectures.

Checklist

• I am familiar with the Contributing Guidelines.
• New or existing tests cover these changes.
• The documentation is up to date with these changes.
• The CHANGELOG.md is up to date with these changes.
• An issue is linked to this pull request.
• If I am implementing a new model or modifying any existing model, I have followed the Models Implementation Coding Standards.

Dependencies

Review Process

All PRs are reviewed by the PhysicsNeMo team before merging.

Depending on which files are changed, GitHub may automatically assign a maintainer for review.

We are also testing AI-based code review tools (e.g., Greptile), which may add automated comments with a confidence score.
This score reflects the AI’s assessment of merge readiness and is not a qualitative judgment of your work, nor is
it an indication that the PR will be accepted / rejected.

AI-generated feedback should be reviewed critically for usefulness.
You are not required to respond to every AI comment, but they are intended to help both authors and reviewers.
Please react to Greptile comments with 👍 or 👎 to provide feedback on their accuracy.

[copy-pr-bot]

This pull request requires additional validation before any workflows can run on NVIDIA's runners.

Pull request vetters can view their responsibilities here.

Contributors can view more details about this message here.

[greptile-apps]

Greptile Summary

This PR adds DiffusionUNet3D and its building blocks (Conv3d, GroupNorm, UNetBlock3D) to physicsnemo.experimental, extending the DDPM++/NCSN++ architecture to 3D volumetric data. There are three P1 defects to address before merging:

• The Examples doctest calls SongUNet3D(...) instead of DiffusionUNet3D(...), so CI doctests will fail.
• When label_dim > 0 and class_labels=None is passed (the default), the forward method crashes with a TypeError because class_labels is used without a None guard.
• Passing a 2-element list for img_resolution silently sets the depth equal to the height instead of raising a ValueError.

Important Files Changed

Filename	Overview
physicsnemo/experimental/models/diffusion_unets/diffusion_unet_3d.py	New 3D U-Net diffusion model. Three P1 bugs: doctest references nonexistent SongUNet3D, class_labels=None silently crashes when map_label is set, and a 2-element img_resolution list silently assigns the wrong depth. Multiple MOD-003b/c, MOD-005, MOD-006, MOD-010 violations.
physicsnemo/experimental/nn/diffusion_unet_3d_blocks.py	New 3D Conv, GroupNorm, and UNetBlock3D layers. Logic looks correct. Missing r""" docstring prefix (MOD-003b) and Forward/Outputs docstring sections (MOD-003c) on all three classes; missing jaxtyping forward annotations (MOD-006).
physicsnemo/experimental/models/diffusion_unets/init.py	New package init exporting DiffusionUNet3D. Missing newline at end of file (cosmetic).
physicsnemo/experimental/nn/init.py	Adds UNetBlock3D, Conv3d, GroupNorm to the experimental nn public API. Change is straightforward.

Comments Outside Diff (1)

1. physicsnemo/experimental/models/diffusion_unets/diffusion_unet_3d.py, line 446-515 (link)

   MOD-005: No input shape validation in forward

   The forward method performs no shape checks on x, noise_labels, class_labels, or augment_labels. Per the coding standards, all tensor arguments must be validated at the top of the method, guarded by if not torch.compiler.is_compiling():, with standardised error messages that include the actual shape. Missing validation leads to cryptic errors deep in the computation graph instead of at the API boundary.

   File Used: CODING_STANDARDS/MODELS_IMPLEMENTATION.md (source)

_{Reviews (1): Last reviewed commit: "Adds DiffusionUNet3D and its modules to ..." | Re-trigger Greptile}

[CharlelieLrt]

/ok to test 194214d

[pzharrington]

Largely looks good to me. Regarding ShardTensor, I don't see any issues that would present problems, the patterns seem to pretty closely follow the 2D SongUNet patterns and we know that works. For maximum sanity, you could add some domain parallelism tests explicitly targeting this architecture.

[CharlelieLrt]

Regarding ShardTensor, I don't see any issues that would present problems, the patterns seem to pretty closely follow the 2D SongUNet patterns and we know that works. For maximum sanity, you could add some domain parallelism tests explicitly targeting this architecture.

Good idea, added two sanity tests in 477acd9 under test/domain_parallel/models/test_diffusion_unet_3d.py (unconditional and conditional, both sharded along H), gated by @pytest.mark.multigpu_static.

[CharlelieLrt]

/blossom-ci

[coreyjadams]

Approving for the domain parallelism tests, thanks for adding those. Recommend you increase the size along the domain parallel dimension a bit but otherwise looks good to me.

[CharlelieLrt]

/blossom-ci