Skip to content

DiT Implementation#1080

Merged
CharlelieLrt merged 13 commits intoNVIDIA:mainfrom
Dibyajyoti-Chakraborty:fea-ext-dit
Aug 20, 2025
Merged

DiT Implementation#1080
CharlelieLrt merged 13 commits intoNVIDIA:mainfrom
Dibyajyoti-Chakraborty:fea-ext-dit

Conversation

@Dibyajyoti-Chakraborty
Copy link
Collaborator

@Dibyajyoti-Chakraborty Dibyajyoti-Chakraborty commented Aug 20, 2025

PhysicsNeMo Pull Request

Description

Diffusion Transformer (DiT) model. The DiT model can be accessed in
physicsnemo.experimental.models.dit.DiT. ⚠️Warning: - Experimental feature
subject to future API changes.

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.

Dependencies

timm, apex, tranformer_engine

Copy link
Collaborator

@coreyjadams coreyjadams left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Overall, I like this version of the DiT better: it is easy to follow, and I appreciate you using existing layers more instead of duplicating, that makes maintenance much easier.

I have two small comments about the try/except blocks in the imports. But once we resolve those, I think it's ready to go.

@coreyjadams
Copy link
Collaborator

/blossom-ci

Copy link
Collaborator

@pzharrington pzharrington left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Approved!

Copy link
Collaborator

@CharlelieLrt CharlelieLrt left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

LGTM overall, greatly appreciate the fact that it's reusing many existing components. Juts a few minor comments about formatting, arguments validation, and documentation

Removed defaults for input_size and in_channels

divided into layers.py and dit.py

Added other validations
@CharlelieLrt CharlelieLrt self-requested a review August 20, 2025 19:31
@CharlelieLrt
Copy link
Collaborator

/blossom-ci

@CharlelieLrt CharlelieLrt merged commit 0f3b8c7 into NVIDIA:main Aug 20, 2025
1 check passed
@Dibyajyoti-Chakraborty Dibyajyoti-Chakraborty deleted the fea-ext-dit branch August 20, 2025 20:32
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

4 participants