Dit unit tests by sajadn · Pull Request #68 · NVIDIA-NeMo/DFM

sajadn · 2025-11-20T19:40:37Z

add unit tests
very minor cleanings
add missing iter to the DiffusionDataModuleConfig

copy-pr-bot · 2025-11-20T19:40:40Z

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.

Signed-off-by: sajadn <snorouzi@nvidia.com>

Signed-off-by: Sajad Norouzi <snorouzi@nvidia.com>

Signed-off-by: sajadn <snorouzi@nvidia.com>

sajadn · 2025-11-20T23:21:21Z

/ok to test fa1b884

Signed-off-by: sajadn <snorouzi@nvidia.com>

sajadn · 2025-11-28T15:39:35Z

/ok to test c376224

Signed-off-by: sajadn <snorouzi@nvidia.com>

abhinavg4 · 2025-12-01T12:27:48Z

/ok to test 3e93c2a

abhinavg4

Looks good. Let;s merge once tets pass

abhinavg4 · 2025-11-21T07:15:07Z

dfm/src/common/utils/save_video.py

        "output_params": ["-f", "mp4"],
    }

-    print("video_save_path", video_save_path)


Good cleanup! Removing debug print statements keeps the output clean in production.

abhinavg4 · 2025-11-21T07:16:20Z

dfm/src/megatron/data/common/diffusion_energon_datamodule.py

    def build_datasets(self, context: DatasetBuildContext):
-        return self.dataset.train_dataloader(), self.dataset.val_dataloader(), self.dataset.test_dataloader()
+        return (
+            iter(self.dataset.train_dataloader()),


Nice improvement! Wrapping dataloaders with iter() makes the interface more explicit and reduces potential confusion in downstream usage.

huvunvidia

Left a few comments.

huvunvidia · 2025-12-02T19:33:39Z

dfm/src/megatron/recipes/dit/dit.py

        DiTModelProvider: Configuration for the DiT-S model.
    """
-    return DiTModelProvider(
+    return DiTXLModelProvider(


Nit: In docstring it notes "DiT-S" but we provide DiTXLModelProvider.

huvunvidia · 2025-12-02T19:37:23Z

pyproject.toml

+    "mediapy>=1.2.4",
+    "megatron-bridge",
+    "wandb[media]>=0.23.0",
+]


Just to make sure, are mediapy and wandb included in OSRB?
Any dependency needs to be approved for open-source.

oops, this is a great question. I'm not sure about that.

sajadn · 2025-12-03T01:42:21Z

/ok to test 99ada83

abhinavg4

LOkks good

pablo-garay · 2025-12-03T02:19:14Z

Approved from Automation perspective

abhinavg4 · 2025-12-03T08:39:08Z

/ok to test f007a7f

* edm and data preprocess tests. Signed-off-by: sajadn <snorouzi@nvidia.com> * Minor cleanings for DiT. Signed-off-by: Sajad Norouzi <snorouzi@nvidia.com> * add dit unit test. Signed-off-by: Sajad Norouzi <snorouzi@nvidia.com> * add iter to the DiffusionDataModule. Signed-off-by: sajadn <snorouzi@nvidia.com> * add missing copyright. Signed-off-by: sajadn <snorouzi@nvidia.com> * use 'no caption' if caption is not present. Signed-off-by: sajadn <snorouzi@nvidia.com> * fix dit inference bug. Add wanbd to inference code. Signed-off-by: sajadn <snorouzi@nvidia.com> * update the DiT configs to be aligned with the original paper. Signed-off-by: sajadn <snorouzi@nvidia.com> * add wandb[video] and mediapy to uv. Signed-off-by: sajadn <snorouzi@nvidia.com> * adjust pos_ids in mock_dataset to have batch dimension, fuse adaLN layers, use DiTSelfAttention. Signed-off-by: sajadn <snorouzi@nvidia.com> * fix the diffusion sample size bug. Signed-off-by: sajadn <snorouzi@nvidia.com> * fix broken tests. Signed-off-by: sajadn <snorouzi@nvidia.com> --------- Signed-off-by: sajadn <snorouzi@nvidia.com> Signed-off-by: Sajad Norouzi <snorouzi@nvidia.com> Co-authored-by: Abhinav Garg <abhinavg@stanford.edu> Signed-off-by: Lawrence Lane <llane@nvidia.com>

* edm and data preprocess tests. Signed-off-by: sajadn <snorouzi@nvidia.com> * Minor cleanings for DiT. Signed-off-by: Sajad Norouzi <snorouzi@nvidia.com> * add dit unit test. Signed-off-by: Sajad Norouzi <snorouzi@nvidia.com> * add iter to the DiffusionDataModule. Signed-off-by: sajadn <snorouzi@nvidia.com> * add missing copyright. Signed-off-by: sajadn <snorouzi@nvidia.com> * use 'no caption' if caption is not present. Signed-off-by: sajadn <snorouzi@nvidia.com> * fix dit inference bug. Add wanbd to inference code. Signed-off-by: sajadn <snorouzi@nvidia.com> * update the DiT configs to be aligned with the original paper. Signed-off-by: sajadn <snorouzi@nvidia.com> * add wandb[video] and mediapy to uv. Signed-off-by: sajadn <snorouzi@nvidia.com> * adjust pos_ids in mock_dataset to have batch dimension, fuse adaLN layers, use DiTSelfAttention. Signed-off-by: sajadn <snorouzi@nvidia.com> * fix the diffusion sample size bug. Signed-off-by: sajadn <snorouzi@nvidia.com> * fix broken tests. Signed-off-by: sajadn <snorouzi@nvidia.com> --------- Signed-off-by: sajadn <snorouzi@nvidia.com> Signed-off-by: Sajad Norouzi <snorouzi@nvidia.com> Co-authored-by: Abhinav Garg <abhinavg@stanford.edu>

sajadn added 4 commits November 20, 2025 11:44

edm and data preprocess tests.

f976ccf

Signed-off-by: sajadn <snorouzi@nvidia.com>

Minor cleanings for DiT.

8f324b5

Signed-off-by: Sajad Norouzi <snorouzi@nvidia.com>

add dit unit test.

d271bd6

Signed-off-by: Sajad Norouzi <snorouzi@nvidia.com>

add iter to the DiffusionDataModule.

2f46b43

Signed-off-by: sajadn <snorouzi@nvidia.com>

sajadn force-pushed the dit_unit_tests branch from 4b75aa7 to 2f46b43 Compare November 20, 2025 19:44

sajadn added 2 commits November 20, 2025 11:49

add missing copyright.

728250e

Signed-off-by: sajadn <snorouzi@nvidia.com>

use 'no caption' if caption is not present.

fa1b884

Signed-off-by: sajadn <snorouzi@nvidia.com>

copy-pr-bot bot temporarily deployed to test November 20, 2025 23:21 Inactive

copy-pr-bot bot temporarily deployed to nemo-ci November 20, 2025 23:21 Inactive

copy-pr-bot bot temporarily deployed to nemo-ci November 20, 2025 23:57 Inactive

copy-pr-bot bot temporarily deployed to nemo-ci November 21, 2025 00:15 Inactive

fix dit inference bug. Add wanbd to inference code.

dee1153

Signed-off-by: sajadn <snorouzi@nvidia.com>

copy-pr-bot bot temporarily deployed to test November 28, 2025 15:39 Inactive

copy-pr-bot bot temporarily deployed to nemo-ci November 28, 2025 15:39 Inactive

update the DiT configs to be aligned with the original paper.

644970d

Signed-off-by: sajadn <snorouzi@nvidia.com>

sajadn force-pushed the dit_unit_tests branch from c376224 to 644970d Compare November 28, 2025 16:01

add wandb[video] and mediapy to uv.

a87268b

Signed-off-by: sajadn <snorouzi@nvidia.com>

sajadn requested a review from a team as a code owner November 28, 2025 16:12

Merge branch 'main' into dit_unit_tests

3e93c2a

copy-pr-bot bot temporarily deployed to test December 1, 2025 12:27 Inactive

copy-pr-bot bot temporarily deployed to nemo-ci December 1, 2025 12:28 Inactive

abhinavg4 previously approved these changes Dec 1, 2025

View reviewed changes

copy-pr-bot bot had a problem deploying to nemo-ci December 1, 2025 13:12 Failure

copy-pr-bot bot temporarily deployed to nemo-ci December 1, 2025 13:12 Inactive

copy-pr-bot bot temporarily deployed to nemo-ci December 1, 2025 18:23 Inactive

copy-pr-bot bot temporarily deployed to nemo-ci December 1, 2025 19:51 Inactive

copy-pr-bot bot temporarily deployed to nemo-ci December 1, 2025 20:17 Inactive

huvunvidia reviewed Dec 2, 2025

View reviewed changes

Merge branch 'main' into dit_unit_tests

99ada83

copy-pr-bot bot temporarily deployed to test December 3, 2025 01:42 Inactive

copy-pr-bot bot temporarily deployed to nemo-ci December 3, 2025 01:42 Inactive

abhinavg4 approved these changes Dec 3, 2025

View reviewed changes

sajadn enabled auto-merge (squash) December 3, 2025 02:02

pablo-garay approved these changes Dec 3, 2025

View reviewed changes

copy-pr-bot bot temporarily deployed to nemo-ci December 3, 2025 03:28 Inactive

copy-pr-bot bot temporarily deployed to nemo-ci December 3, 2025 06:14 Inactive

Merge branch 'main' into dit_unit_tests

f007a7f

copy-pr-bot bot temporarily deployed to test December 3, 2025 08:39 Inactive

copy-pr-bot bot temporarily deployed to nemo-ci December 3, 2025 08:39 Inactive

copy-pr-bot bot temporarily deployed to nemo-ci December 3, 2025 08:42 Inactive

copy-pr-bot bot temporarily deployed to nemo-ci December 3, 2025 09:06 Inactive

sajadn merged commit 1061749 into main Dec 3, 2025
15 checks passed

chtruong814 deleted the dit_unit_tests branch January 29, 2026 20:26

Conversation

sajadn commented Nov 20, 2025

Uh oh!

copy-pr-bot bot commented Nov 20, 2025

Uh oh!

sajadn commented Nov 20, 2025

Uh oh!

sajadn commented Nov 28, 2025

Uh oh!

abhinavg4 commented Dec 1, 2025

Uh oh!

abhinavg4 left a comment

Choose a reason for hiding this comment

Uh oh!

abhinavg4 Nov 21, 2025

Choose a reason for hiding this comment

Uh oh!

abhinavg4 Nov 21, 2025

Choose a reason for hiding this comment

Uh oh!

huvunvidia left a comment

Choose a reason for hiding this comment

Uh oh!

huvunvidia Dec 2, 2025

Choose a reason for hiding this comment

Uh oh!

huvunvidia Dec 2, 2025

Choose a reason for hiding this comment

Uh oh!

sajadn Dec 2, 2025

Choose a reason for hiding this comment

Uh oh!

sajadn commented Dec 3, 2025

Uh oh!

abhinavg4 left a comment

Choose a reason for hiding this comment

Uh oh!

pablo-garay commented Dec 3, 2025

Uh oh!

abhinavg4 commented Dec 3, 2025

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

4 participants

Comments