Conversation
abhinavg4
left a comment
abhinavg4
left a comment
There was a problem hiding this comment.
Left minor comments. Overall looks ok.
|
|
||
|
|
||
| # P2: Look into get_batch_on_this_cp_rank | ||
| def split_inputs_cp(x: Tensor, seq_dim: int, cp_group: ProcessGroup) -> Tensor: |
There was a problem hiding this comment.
What is checkpoint parallelism and where are these functions used? Maybe use their equivalent from Mcore?
There was a problem hiding this comment.
it wasn't used in the new repo, removed it.
| @@ -0,0 +1,343 @@ | |||
| # Copyright (c) 2025, NVIDIA CORPORATION. All rights reserved. | |||
There was a problem hiding this comment.
I think @huvunvidia is also using this file? Can we move this to common if yes?
There was a problem hiding this comment.
Yes, Wan will be using DiffusionDataModule in diffusion_energon_datamodule.py.
There was a problem hiding this comment.
moved it to src/megatron/data/common/ along with some additional classes implemented for sequence packing.
|
|
||
| @dataclass | ||
| class DiTWithAdaLNSubmodules(TransformerLayerSubmodules): | ||
| temporal_self_attention: Union[ModuleSpec, type] = IdentityOp |
There was a problem hiding this comment.
I think this should be cross_attention istead?
There was a problem hiding this comment.
@sajadn Can you take a look at this comment ? I see we use cross_attention in the code and not temporal_Self_Attention ?
There was a problem hiding this comment.
you're right. I'll remove it.
huvunvidia
left a comment
huvunvidia
left a comment
There was a problem hiding this comment.
Left some comments.
Thank you Sajad.
| @@ -27,7 +27,6 @@ | |||
| SelfAttention, | |||
There was a problem hiding this comment.
We should not use any nemo_vfm right?
There was a problem hiding this comment.
hmm, did I use nemo_vfm? I can't find it.
There was a problem hiding this comment.
No, I meant I'm not sure why your PR include code in nemo_vfm.
There was a problem hiding this comment.
Yeah probably this file should not be used (incase it's used, I think it was just added to a commit by mistake)
There was a problem hiding this comment.
hey, ya this one was added by mistake. I removed it but still ruff complaint about its styling issues. Apparently when ruff gets applied to a file it's not possible to remove it from the checks. Let's push the changes as they're minor changes and just removes unused imports.
| @@ -0,0 +1,193 @@ | |||
| # SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved. | |||
There was a problem hiding this comment.
Not sure dfm/src/common is the right place to put these tokenizer files.
dfm/src/common should be for common files used by both AutoModel and Mcore paths.
Since only DiT model use these tokenizer files, maybe dfm/src/megatron/models/dit/tokenizers would be a better place?
There was a problem hiding this comment.
I think this is ok since Automodel might use it someday?
There was a problem hiding this comment.
Can we replace cosmos to something else and I don't think we need cosmos/cosmos1
Let's do
dfm/src/common/tokenizers/causal_video_tokenizer.py
There was a problem hiding this comment.
i still like to keep it in cosmos, what if we want to incorporate additional video tokenizers?
There was a problem hiding this comment.
We should not have any cosmos reference just coz cosmos is a seperate project now and we are not using it anymore. Anyways it's a P1
| @@ -0,0 +1,104 @@ | |||
| # Copyright (c) 2024, NVIDIA CORPORATION. All rights reserved. | |||
There was a problem hiding this comment.
Same comments as above, these files shouldn't be in dfm/src/common because it is not used for AutoModel, and only for Mcore's DiT.
There was a problem hiding this comment.
I think it's ok Automodel might use them down the line ?
There was a problem hiding this comment.
I'm okay both ways, let me know which one is preferred.
There was a problem hiding this comment.
I think let's keep it here, in case Automodel uses it. When in confusion, move things to common IMO
| @@ -0,0 +1,343 @@ | |||
| # Copyright (c) 2025, NVIDIA CORPORATION. All rights reserved. | |||
There was a problem hiding this comment.
Yes, Wan will be using DiffusionDataModule in diffusion_energon_datamodule.py.
abhinavg4
left a comment
abhinavg4
left a comment
There was a problem hiding this comment.
Only major comment is the cross attention comment. Thanks
| @@ -0,0 +1,193 @@ | |||
| # SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved. | |||
There was a problem hiding this comment.
I think this is ok since Automodel might use it someday?
| @@ -0,0 +1,193 @@ | |||
| # SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved. | |||
There was a problem hiding this comment.
Can we replace cosmos to something else and I don't think we need cosmos/cosmos1
Let's do
dfm/src/common/tokenizers/causal_video_tokenizer.py
| @@ -0,0 +1,145 @@ | |||
| # SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved. | |||
There was a problem hiding this comment.
just utils but under tokenizers maybe?
| @@ -0,0 +1,104 @@ | |||
| # Copyright (c) 2024, NVIDIA CORPORATION. All rights reserved. | |||
There was a problem hiding this comment.
I think let's keep it here, in case Automodel uses it. When in confusion, move things to common IMO
|
|
||
| @dataclass | ||
| class DiTWithAdaLNSubmodules(TransformerLayerSubmodules): | ||
| temporal_self_attention: Union[ModuleSpec, type] = IdentityOp |
There was a problem hiding this comment.
@sajadn Can you take a look at this comment ? I see we use cross_attention in the code and not temporal_Self_Attention ?
| @@ -27,7 +27,6 @@ | |||
| SelfAttention, | |||
There was a problem hiding this comment.
Yeah probably this file should not be used (incase it's used, I think it was just added to a commit by mistake)
sajadn
force-pushed
the
dit_debug
branch
2 times, most recently
from
c12a7fa to
a120cce
Compare
Signed-off-by: Sajad Norouzi <snorouzi@nvidia.com>
Signed-off-by: Sajad Norouzi <snorouzi@nvidia.com>
Signed-off-by: Sajad Norouzi <snorouzi@nvidia.com>
Signed-off-by: Sajad Norouzi <snorouzi@nvidia.com>
…ify usage of common elements in the codebase. Signed-off-by: Sajad Norouzi <snorouzi@nvidia.com>
- Introduced `batch_ops.py` with functions for tensor broadcasting and element-wise operations (addition, multiplication, subtraction, division). - Added `dynamic_import.py` for dynamic loading of classes or functions from specified paths. - Created `torch_split_tensor_for_cp.py` for tensor splitting and concatenation in checkpoint parallelism. - Implemented a new data module for handling multimodal datasets in `base.py` and `diffusion_energon_datamodule.py`. - Added `diffusion_taskencoder.py` for encoding image/video samples for diffusion tasks. - Updated `edm_pipeline.py` to utilize new utility functions for tensor operations. - Included a script for preparing the Energon dataset in `prepare_energon_dataset_butterfly.py`. Signed-off-by: Sajad Norouzi <snorouzi@nvidia.com>
Signed-off-by: Sajad Norouzi <snorouzi@nvidia.com>
…ec.py to streamline the codebase. Signed-off-by: Sajad Norouzi <snorouzi@nvidia.com>
….py`, updated import paths in `dit.py`, and adjusted training parameters for `task_encoder_seq_length` and `save_interval` to optimize configuration. Signed-off-by: Sajad Norouzi <snorouzi@nvidia.com>
…sitories, dataset location, Docker setup, installation, and running the model on the butterfly dataset. Signed-off-by: Sajad Norouzi <snorouzi@nvidia.com>
…te files including `decode_cosmos_latent.py`, `attention_wan.py`, `diffusers_vae.py`, `dit_embeddings.py`, `dit_inference.py`, and others. Updated import paths and streamlined the codebase for improved maintainability. Signed-off-by: Sajad Norouzi <snorouzi@nvidia.com>
- Introduced `attention_wan.py` with implementations for WAN self and cross attention layers, including adaptive layer normalization. - Added `time_pos_embeddings.py` for various temporal positional embedding strategies. - Created `dit_data_process.py` for data handling in DiT model training, including batch processing and sequence length encoding. - Implemented `dit_model_provider.py` for model configuration and instantiation of DiT models. - Added utility functions in `edm_utils.py` for EDM sampling and scaling processes. These additions enhance the DiT model's capabilities and streamline data processing and model configuration. Signed-off-by: Sajad Norouzi <snorouzi@nvidia.com>
Signed-off-by: Sajad Norouzi <snorouzi@nvidia.com>
Signed-off-by: Sajad Norouzi <snorouzi@nvidia.com>
Signed-off-by: Sajad Norouzi <snorouzi@nvidia.com>
Signed-off-by: Sajad Norouzi <snorouzi@nvidia.com>
Signed-off-by: Sajad Norouzi <snorouzi@nvidia.com>
Signed-off-by: Sajad Norouzi <snorouzi@nvidia.com>
…imal Tokenizer version. Signed-off-by: Sajad Norouzi <snorouzi@nvidia.com>
Signed-off-by: Sajad Norouzi <snorouzi@nvidia.com>
Signed-off-by: Sajad Norouzi <snorouzi@nvidia.com>
…odules class Signed-off-by: Sajad Norouzi <snorouzi@nvidia.com>
…king logic that is shared across models. Signed-off-by: Sajad Norouzi <snorouzi@nvidia.com>
…f tokens. Signed-off-by: Sajad Norouzi <snorouzi@nvidia.com>
Signed-off-by: Sajad Norouzi <snorouzi@nvidia.com>
Signed-off-by: Sajad Norouzi <snorouzi@nvidia.com>
…f common/models. Make inference compatible with the new data loading. Signed-off-by: Sajad Norouzi <snorouzi@nvidia.com>
Signed-off-by: Sajad Norouzi <snorouzi@nvidia.com>
|
/ok to test 5dd2f8c |
Signed-off-by: Sajad Norouzi <snorouzi@nvidia.com>
Signed-off-by: Sajad Norouzi <snorouzi@nvidia.com>
|
/ok to test 80da225 |
* first dit commit. Signed-off-by: Sajad Norouzi <snorouzi@nvidia.com> * inference working with old checkpoint. Signed-off-by: Sajad Norouzi <snorouzi@nvidia.com> * add on_validation_start step to generate video. Signed-off-by: Sajad Norouzi <snorouzi@nvidia.com> * EDM loss bug fix. Signed-off-by: Sajad Norouzi <snorouzi@nvidia.com> * Refactor DIT components: add comments for potential removals and clarify usage of common elements in the codebase. Signed-off-by: Sajad Norouzi <snorouzi@nvidia.com> * Add utility functions for tensor operations and dynamic imports - Introduced `batch_ops.py` with functions for tensor broadcasting and element-wise operations (addition, multiplication, subtraction, division). - Added `dynamic_import.py` for dynamic loading of classes or functions from specified paths. - Created `torch_split_tensor_for_cp.py` for tensor splitting and concatenation in checkpoint parallelism. - Implemented a new data module for handling multimodal datasets in `base.py` and `diffusion_energon_datamodule.py`. - Added `diffusion_taskencoder.py` for encoding image/video samples for diffusion tasks. - Updated `edm_pipeline.py` to utilize new utility functions for tensor operations. - Included a script for preparing the Energon dataset in `prepare_energon_dataset_butterfly.py`. Signed-off-by: Sajad Norouzi <snorouzi@nvidia.com> * remove unused layers from spec. Signed-off-by: Sajad Norouzi <snorouzi@nvidia.com> * Remove unused AdaLN and STDiTLayerWithAdaLN classes from dit_layer_spec.py to streamline the codebase. Signed-off-by: Sajad Norouzi <snorouzi@nvidia.com> * Refactor DIT model files: removed unused functions from `dit_provider.py`, updated import paths in `dit.py`, and adjusted training parameters for `task_encoder_seq_length` and `save_interval` to optimize configuration. Signed-off-by: Sajad Norouzi <snorouzi@nvidia.com> * Add README for DiT model setup: include instructions for cloning repositories, dataset location, Docker setup, installation, and running the model on the butterfly dataset. Signed-off-by: Sajad Norouzi <snorouzi@nvidia.com> * Remove unused files and refactor DIT model components: deleted obsolete files including `decode_cosmos_latent.py`, `attention_wan.py`, `diffusers_vae.py`, `dit_embeddings.py`, `dit_inference.py`, and others. Updated import paths and streamlined the codebase for improved maintainability. Signed-off-by: Sajad Norouzi <snorouzi@nvidia.com> * Add new components for DiT model: - Introduced `attention_wan.py` with implementations for WAN self and cross attention layers, including adaptive layer normalization. - Added `time_pos_embeddings.py` for various temporal positional embedding strategies. - Created `dit_data_process.py` for data handling in DiT model training, including batch processing and sequence length encoding. - Implemented `dit_model_provider.py` for model configuration and instantiation of DiT models. - Added utility functions in `edm_utils.py` for EDM sampling and scaling processes. These additions enhance the DiT model's capabilities and streamline data processing and model configuration. Signed-off-by: Sajad Norouzi <snorouzi@nvidia.com> * revert THD to SBHD and use remove prints. Signed-off-by: Sajad Norouzi <snorouzi@nvidia.com> * add support sequence for sequence packing. Signed-off-by: Sajad Norouzi <snorouzi@nvidia.com> * move files in proper location, uncomment cp related code, cleaning. Signed-off-by: Sajad Norouzi <snorouzi@nvidia.com> * fix lint errors. Signed-off-by: Sajad Norouzi <snorouzi@nvidia.com> * fix README lint error. Signed-off-by: Sajad Norouzi <snorouzi@nvidia.com> * add the missing copyright. Signed-off-by: Sajad Norouzi <snorouzi@nvidia.com> * removed dependency on nemo CausalVideoTokenizer by implementing a minimal Tokenizer version. Signed-off-by: Sajad Norouzi <snorouzi@nvidia.com> * remove nemo vae modules. Signed-off-by: Sajad Norouzi <snorouzi@nvidia.com> * update readme.md Signed-off-by: Sajad Norouzi <snorouzi@nvidia.com> * remove unused temporal_self_attention attribute from DiTWithAdaLNSubmodules class Signed-off-by: Sajad Norouzi <snorouzi@nvidia.com> * create an abstract DiffusionTaskEncoder to encompass the sequence packing logic that is shared across models. Signed-off-by: Sajad Norouzi <snorouzi@nvidia.com> * fix use seq_length instead of packing_buffer_size to set max number of tokens. Signed-off-by: Sajad Norouzi <snorouzi@nvidia.com> * Add the missing copy write. Minor update to the config. Signed-off-by: Sajad Norouzi <snorouzi@nvidia.com> * lint still complains about the nemo_vfm file. Signed-off-by: Sajad Norouzi <snorouzi@nvidia.com> * move data modules to common, move cosmos to common/tokenizer instad of common/models. Make inference compatible with the new data loading. Signed-off-by: Sajad Norouzi <snorouzi@nvidia.com> * uncomment omega config overwrite. Signed-off-by: Sajad Norouzi <snorouzi@nvidia.com> * change import (remote lint complains). Signed-off-by: Sajad Norouzi <snorouzi@nvidia.com> --------- Signed-off-by: Sajad Norouzi <snorouzi@nvidia.com> Co-authored-by: Abhinav Garg <abhinavg@stanford.edu>
4 participants
the code for training and inference of DiT. Data preparation is also included. The current version uses SBHD + packed_sequence which will be removed with the follow up commits.