Add DiT Readme. #61
Merged
Add DiT Readme. #61
Conversation
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Contributor
sajadn
commented
Nov 18, 2025
sajadn
commented
- add dit readme.
Signed-off-by: sajadn <[email protected]>
abhinavg4
requested changes
Nov 19, 2025
Contributor
abhinavg4
left a comment
abhinavg4
left a comment
There was a problem hiding this comment.
Left a bunch of comments also please check these comments: https://github.com/NVIDIA-NeMo/DFM/pull/38/files
Signed-off-by: Sajad Norouzi <[email protected]>
Signed-off-by: Sajad Norouzi <[email protected]>
abhinavg4
previously approved these changes
Nov 21, 2025
Contributor
abhinavg4
left a comment
abhinavg4
left a comment
There was a problem hiding this comment.
Great documentation. Loved it. Left a few minor comments.
| The following repositories need to be cloned with specific commit hashes: | ||
|
|
||
| #### Megatron-LM | ||
| The script below prepares the dataset to be compatible with Energon. |
Contributor
There was a problem hiding this comment.
Nit: packs the HF dataset into Webdataset format. Format compatible with Energon sounds something mystic
| git checkout dit_debug | ||
| cd .. | ||
| energon prepare ./ | ||
| /usr/local/lib/python3.12/dist-packages/torch/cuda/__init__.py:63: FutureWarning: The pynvml package is deprecated. Please install nvidia-ml-py instead. If you did not install pynvml directly, please report this to the maintainers of the package that installed pynvml for you. |
Contributor
There was a problem hiding this comment.
Remove this print. Not needed
Contributor
Author
There was a problem hiding this comment.
do you want me to get rid of the whole prints of energon prepare ./?
| Once you have the dataset and container ready, you can start training the DiT model on your own dataset. This repository leverages [sequence packing](https://docs.nvidia.com/nemo-framework/user-guide/24.09/nemotoolkit/features/optimizations/sequence_packing.html) to maximize training efficiency. Sequence packing stacks multiple samples into a single sequence instead of padding individual samples to a fixed length; therefore, `micro_batch_size` must be set to 1. Additionally, `qkv_format` should be set to `thd` to signal to Transformer Engine that sequence packing is enabled. | ||
| For data loading, Energon provides two key hyperparameters related to sequence packing: `task_encoder_seq_length` and `packing_buffer_size`. The `task_encoder_seq_length` parameter controls the maximum sequence length passed to the model, while `packing_buffer_size` determines the number of samples processed to create different buckets. You can look at `select_samples_to_pack` and `pack_selected_samples` methods of [DiffusionTaskEncoderWithSequencePacking](https://github.com/NVIDIA-NeMo/DFM/blob/main/dfm/src/megatron/data/common/diffusion_task_encoder_with_sp.py#L50) to get a better sense of these parameters. |
Contributor
There was a problem hiding this comment.
Great section. Also provide link to energon documentation for these params?
Signed-off-by: Sajad Norouzi <[email protected]>
Signed-off-by: Sajad Norouzi <[email protected]>
Contributor
|
/ok to test 045f421 |
abhinavg4
reviewed
Nov 21, 2025
| @@ -1,77 +1,184 @@ | |||
| # DiT (Diffusion Transformer) Model Setup | |||
Contributor
There was a problem hiding this comment.
Please put this under docs/megatron/models/DiT/dit.md
Signed-off-by: Sajad Norouzi <[email protected]>
Contributor
|
/ok to test 7a832f6 |
lbliii
pushed a commit
that referenced
this pull request
Dec 3, 2025
* Add DiT Readme. Signed-off-by: sajadn <[email protected]> * Update DiT readme. Signed-off-by: Sajad Norouzi <[email protected]> * Minor wording update. Signed-off-by: Sajad Norouzi <[email protected]> --------- Signed-off-by: sajadn <[email protected]> Signed-off-by: Sajad Norouzi <[email protected]> Signed-off-by: Lawrence Lane <[email protected]>
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.
This suggestion is invalid because no changes were made to the code.
Suggestions cannot be applied while the pull request is closed.
Suggestions cannot be applied while viewing a subset of changes.
Only one suggestion per line can be applied in a batch.
Add this suggestion to a batch that can be applied as a single commit.
Applying suggestions on deleted lines is not supported.
You must change the existing code in this line in order to create a valid suggestion.
Outdated suggestions cannot be applied.
This suggestion has been applied or marked resolved.
Suggestions cannot be applied from pending reviews.
Suggestions cannot be applied on multi-line comments.
Suggestions cannot be applied while the pull request is queued to merge.
Suggestion cannot be applied right now. Please check back later.