Refactor GroupNorm and log unmatched state_dict keys #989

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Jump to bottom

Merged

CharlelieLrt merged 63 commits into NVIDIA:main from juliusberner:improve_diffusion

Aug 13, 2025

Contributor

juliusberner commented Jun 24, 2025 •

edited by CharlelieLrt

Loading

PhysicsNeMo Pull Request

Description

Diffusion models: consolidate ApexGroupNorm and GroupNorm in models/diffusion/layers.py with a factory get_group_norm that can be used to instantiate either one of them. this ensures that the state_dict remains consistent with previous versions.
Log missing and unexpected keys when loading checkpoints.
In SongUNet-based diffusion models, add persistent=False for deterministic, non-learnable positional embeddings.

Closes #1001 .

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.


          Refactor GroupNorm and log unmatched state_dict keys

b870c49

Signed-off-by: Julius Berner <jberner@nvidia.com>

CharlelieLrt self-requested a review

June 24, 2025 21:05

CharlelieLrt assigned juliusberner

CharlelieLrt added bug 3 - Ready for Review labels

juliusberner and others added 2 commits

June 24, 2025 14:07


          Merge branch 'main' into improve_diffusion

72416a6


          Merge branch 'main' into improve_diffusion

119c983

CharlelieLrt reviewed

View reviewed changes

physicsnemo/models/util_compatibility.py Outdated Show resolved Hide resolved

physicsnemo/models/module.py Outdated Show resolved Hide resolved

CharlelieLrt mentioned this pull request

🐛[BUG]: GroupNorm creates unused parameters when use_apex_gn=True #1001

Closed

pzharrington reviewed

View reviewed changes

physicsnemo/models/module.py Outdated Show resolved Hide resolved

akshaysubr reviewed

View reviewed changes

physicsnemo/models/diffusion/layers.py Outdated Show resolved Hide resolved

Julius Berner added 2 commits

July 9, 2025 01:13


          Refactor GroupNorm and log unmatched state_dict keys

59a6ff1

Signed-off-by: Julius Berner <jberner@nvidia.com>


          Add changes from MR996

f8e01c7

juliusberner force-pushed the improve_diffusion branch from 8767605 to f8e01c7 Compare

July 9, 2025 01:15

CharlelieLrt added 3 commits

July 22, 2025 20:01


          Merge branch 'improve_diffusion' of https://github.com/juliusberner/p…

08f5368

…hysicsnemo into improve_diffusion


          Merge branch 'main' into improve_diffusion

f07cbf3


          Made load_state_dict method semi-private

133d925

Signed-off-by: Charlelie Laurent <claurent@nvidia.com>

CharlelieLrt reviewed

View reviewed changes

physicsnemo/models/diffusion/layers.py Outdated Show resolved Hide resolved

CharlelieLrt added 4 commits

July 31, 2025 09:22


          Move the attention migration into a load_state_dict pre-hook

ca424a8

Signed-off-by: Charlelie Laurent <claurent@nvidia.com>


          Deleted duplicate line in CHANGELOG.md

bae6d39

Signed-off-by: Charlelie Laurent <claurent@nvidia.com>


          Removed warnings in UNetBlock load_state_dict pre-hook

235b928

Signed-off-by: Charlelie Laurent <claurent@nvidia.com>


          Added test for UNetBlock checkpoint loading from v1.0.1

4fed09c

Signed-off-by: Charlelie Laurent <claurent@nvidia.com>

Collaborator

CharlelieLrt commented Aug 1, 2025

Note: need to add tests (see this thread)

CharlelieLrt and others added 3 commits

August 1, 2025 10:36


          Merge branch 'main' into fix-attention-load-state-dict

647528b


          Changed tol in test + added new test with fused_conv_bias=True

f70e0b7

Signed-off-by: Charlelie Laurent <claurent@nvidia.com>


          Merge branch 'fix-attention-load-state-dict' of https://github.com/Ch…

b105b32

…arlelieLrt/modulus into fix-attention-load-state-dict

CharlelieLrt added ! - Release 4 - In Review and removed 3 - Ready for Review labels


          Merge branch 'main' into improve_diffusion

49d78a5

CharlelieLrt added 4 commits

August 11, 2025 11:11


          Added new tests for UNet block

3640bfe

Signed-off-by: Charlelie Laurent <claurent@nvidia.com>


          Added more param validation in Attention

96bca9a

Signed-off-by: Charlelie Laurent <claurent@nvidia.com>


          Added tests for new Attention layer

05debe1

Signed-off-by: Charlelie Laurent <claurent@nvidia.com>


          Merge branch 'main' into improve_diffusion

db8a738

CharlelieLrt self-requested a review

August 11, 2025 18:49

CharlelieLrt approved these changes

View reviewed changes

Collaborator

CharlelieLrt commented Aug 11, 2025

/blossom-ci


          Pin C++ backend for attention op

f2e850c

Signed-off-by: Charlelie Laurent <claurent@nvidia.com>

Collaborator

CharlelieLrt commented Aug 11, 2025

/blossom-ci


          Added reference input data for attention tests

94d4401

Signed-off-by: Charlelie Laurent <claurent@nvidia.com>

Collaborator

CharlelieLrt commented Aug 11, 2025

/blossom-ci


          Some files renaming

e1c8abb

Signed-off-by: Charlelie Laurent <claurent@nvidia.com>

Collaborator

CharlelieLrt commented Aug 11, 2025

/blossom-ci


          Reverted back attention to previous implementation

60b06a8

Signed-off-by: Charlelie Laurent <claurent@nvidia.com>

Collaborator

CharlelieLrt commented Aug 12, 2025

/blossom-ci

CharlelieLrt added 2 commits

August 11, 2025 19:27


          Updates on new tests

7a2fec7

Signed-off-by: Charlelie Laurent <claurent@nvidia.com>


          Updates on new tests

153990a

Signed-off-by: Charlelie Laurent <claurent@nvidia.com>

Collaborator

CharlelieLrt commented Aug 12, 2025

/blossom-ci

Collaborator

prem-krishnan commented Aug 12, 2025

@CharlelieLrt Can we associate this task with an epic?

CharlelieLrt added 7 commits

August 12, 2025 12:21


          Deleted tests ref data

163d1d0

Signed-off-by: Charlelie Laurent <claurent@nvidia.com>


          Group norm test working

a166fe4

Signed-off-by: Charlelie Laurent <claurent@nvidia.com>


          Group norm test working

07f39d2

Signed-off-by: Charlelie Laurent <claurent@nvidia.com>


          Tests for attention layer passing locally

00f39ef

Signed-off-by: Charlelie Laurent <claurent@nvidia.com>


          Removed backend in attention test

f3cc71e

Signed-off-by: Charlelie Laurent <claurent@nvidia.com>


          Modified UNetBlock tests

94cc4da

Signed-off-by: Charlelie Laurent <claurent@nvidia.com>


          Tests for UNetBlock passing locally

b6291ca

Signed-off-by: Charlelie Laurent <claurent@nvidia.com>

Collaborator

CharlelieLrt commented Aug 13, 2025

/blossom-ci


          Merge branch 'main' into improve_diffusion

5366a76

Collaborator

CharlelieLrt commented Aug 13, 2025

/blossom-ci

CharlelieLrt merged commit fb4140c into NVIDIA:main

1 check passed

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

4 - In Review bug