Draft+ for SDXL [draft] #222

rohitrango · 2024-06-25T23:35:24Z

What does this PR do ?

Add a one line overview of what this PR aims to accomplish.

Changelog

Please update the CHANGELOG.md under next version with high level changes in this PR.

Usage

You can potentially add a usage example below

# Add a code snippet demonstrating how to use this

Before your PR is "Ready for review"

Pre checks:

Make sure you read and followed Contributor guidelines
Did you write any new necessary tests?
Did you add or update any necessary documentation? Make sure to also update the NeMo Framework User Guide which contains the tutorials

Checklist when contributing a new algorithm

Does the trainer resume and restore model state all states?
Does the trainer support all parallelism techniques(PP, TP, DP)?
Does the trainer support max_steps=-1 and validation?
Does the trainer only call APIs defined in alignable_interface.py?
Does the trainer have proper logging?

Additional Information

Related to # (issue)

terrykong

Thanks @rohitrango for the contribution!

I've left some comments. Let me know if there are any questions

.gitignore

examples/mm/stable_diffusion/launch_draft_oci_multinode_xl.sh

nemo_aligner/algorithms/supervised.py

nemo_aligner/models/mm/stable_diffusion/megatron_sdxl_draftp_model.py

nemo_aligner/utils/train_script_utils.py

nemo_aligner/models/mm/stable_diffusion/megatron_sdxl_draftp_model.py

examples/mm/stable_diffusion/train_sdxl_draftp.py

examples/mm/stable_diffusion/conf/draftp_sdxl.yaml

nemo_aligner/models/mm/stable_diffusion/megatron_sdxl_draftp_model.py

terrykong · 2024-07-08T19:13:04Z

Requires NVIDIA/NeMo#9543

edit: Also requires NVIDIA/NeMo#9654 to have feature in RC branch

Signed-off-by: Rohit Jena <[email protected]>

sampling instead Signed-off-by: Rohit Jena <[email protected]>

Signed-off-by: Rohit Jena <[email protected]>

for more information, see https://pre-commit.ci Signed-off-by: Rohit Jena <[email protected]>

Signed-off-by: Rohit Jena <[email protected]>

for more information, see https://pre-commit.ci Signed-off-by: Rohit Jena <[email protected]>

Signed-off-by: Rohit Jena <[email protected]>

for more information, see https://pre-commit.ci Signed-off-by: Rohit Jena <[email protected]>

Signed-off-by: Rohit Jena <[email protected]>

for more information, see https://pre-commit.ci Signed-off-by: Rohit Jena <[email protected]>

terrykong

Thanks @rohitrango for adding this new feature!

* leftover commit Signed-off-by: Rohit Jena <[email protected]> * leftover commit Signed-off-by: Rohit Jena <[email protected]> * commits Signed-off-by: Rohit Jena <[email protected]> * update gitignore Signed-off-by: Rohit Jena <[email protected]> * init model for SDXL Signed-off-by: Rohit Jena <[email protected]> * correct path in eos script Signed-off-by: Rohit Jena <[email protected]> * modified generate and log_generate scripts to use diffusion engine sampling instead Signed-off-by: Rohit Jena <[email protected]> * fixed most runtime bugs -- check for logical bugs Signed-off-by: Rohit Jena <[email protected]> * examining mapping between hf and nemo Signed-off-by: Rohit Jena <[email protected]> * writing converter script for unet Signed-off-by: Rohit Jena <[email protected]> * tmp commit (moving to eos) Signed-off-by: Rohit Jena <[email protected]> * more changes to draftp xl Signed-off-by: Rohit Jena <[email protected]> * changed batch scripts Signed-off-by: Rohit Jena <[email protected]> * changed launch scripts Signed-off-by: Rohit Jena <[email protected]> * check adapter control Signed-off-by: Rohit Jena <[email protected]> * clean up lora hotswap debugging Signed-off-by: Rohit Jena <[email protected]> * adding fsdp to draftp training Signed-off-by: Rohit Jena <[email protected]> * FSDP now works for SDXL?! Signed-off-by: Rohit Jena <[email protected]> * added custom rule to enable sharding of decoder Signed-off-by: Rohit Jena <[email protected]> * also shard the clip embeddings Signed-off-by: Rohit Jena <[email protected]> * multinode script created + testing 2048 config Signed-off-by: Rohit Jena <[email protected]> * add activation checkpointing Signed-off-by: Rohit Jena <[email protected]> * added activation checkpointing Signed-off-by: Rohit Jena <[email protected]> * added SFT and PEFT support with Draft+ Signed-off-by: Rohit Jena <[email protected]> * corrected init denoise bug Signed-off-by: Rohit Jena <[email protected]> * added multinode setup script - test it Signed-off-by: Rohit Jena <[email protected]> * enable sharding for sdlora Signed-off-by: Rohit Jena <[email protected]> * added multinode script for OCI Signed-off-by: Rohit Jena <[email protected]> * some more refactoring Signed-off-by: Rohit Jena <[email protected]> * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci Signed-off-by: Rohit Jena <[email protected]> * address Terry's comments Signed-off-by: Rohit Jena <[email protected]> * remove + 0 to clone Signed-off-by: Rohit Jena <[email protected]> * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci Signed-off-by: Rohit Jena <[email protected]> * add version guard Signed-off-by: Rohit Jena <[email protected]> * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci Signed-off-by: Rohit Jena <[email protected]> * added todo to merge these APIs later Signed-off-by: Rohit Jena <[email protected]> * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci Signed-off-by: Rohit Jena <[email protected]> --------- Signed-off-by: Rohit Jena <[email protected]> Co-authored-by: Rohit Jena <[email protected]> Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com> Signed-off-by: Alexander Bukharin <[email protected]>

* leftover commit Signed-off-by: Rohit Jena <[email protected]> * leftover commit Signed-off-by: Rohit Jena <[email protected]> * commits Signed-off-by: Rohit Jena <[email protected]> * update gitignore Signed-off-by: Rohit Jena <[email protected]> * init model for SDXL Signed-off-by: Rohit Jena <[email protected]> * correct path in eos script Signed-off-by: Rohit Jena <[email protected]> * modified generate and log_generate scripts to use diffusion engine sampling instead Signed-off-by: Rohit Jena <[email protected]> * fixed most runtime bugs -- check for logical bugs Signed-off-by: Rohit Jena <[email protected]> * examining mapping between hf and nemo Signed-off-by: Rohit Jena <[email protected]> * writing converter script for unet Signed-off-by: Rohit Jena <[email protected]> * tmp commit (moving to eos) Signed-off-by: Rohit Jena <[email protected]> * more changes to draftp xl Signed-off-by: Rohit Jena <[email protected]> * changed batch scripts Signed-off-by: Rohit Jena <[email protected]> * changed launch scripts Signed-off-by: Rohit Jena <[email protected]> * check adapter control Signed-off-by: Rohit Jena <[email protected]> * clean up lora hotswap debugging Signed-off-by: Rohit Jena <[email protected]> * adding fsdp to draftp training Signed-off-by: Rohit Jena <[email protected]> * FSDP now works for SDXL?! Signed-off-by: Rohit Jena <[email protected]> * added custom rule to enable sharding of decoder Signed-off-by: Rohit Jena <[email protected]> * also shard the clip embeddings Signed-off-by: Rohit Jena <[email protected]> * multinode script created + testing 2048 config Signed-off-by: Rohit Jena <[email protected]> * add activation checkpointing Signed-off-by: Rohit Jena <[email protected]> * added activation checkpointing Signed-off-by: Rohit Jena <[email protected]> * added SFT and PEFT support with Draft+ Signed-off-by: Rohit Jena <[email protected]> * corrected init denoise bug Signed-off-by: Rohit Jena <[email protected]> * added multinode setup script - test it Signed-off-by: Rohit Jena <[email protected]> * enable sharding for sdlora Signed-off-by: Rohit Jena <[email protected]> * added multinode script for OCI Signed-off-by: Rohit Jena <[email protected]> * some more refactoring Signed-off-by: Rohit Jena <[email protected]> * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci Signed-off-by: Rohit Jena <[email protected]> * address Terry's comments Signed-off-by: Rohit Jena <[email protected]> * remove + 0 to clone Signed-off-by: Rohit Jena <[email protected]> * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci Signed-off-by: Rohit Jena <[email protected]> * add version guard Signed-off-by: Rohit Jena <[email protected]> * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci Signed-off-by: Rohit Jena <[email protected]> * added todo to merge these APIs later Signed-off-by: Rohit Jena <[email protected]> * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci Signed-off-by: Rohit Jena <[email protected]> --------- Signed-off-by: Rohit Jena <[email protected]> Co-authored-by: Rohit Jena <[email protected]> Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>

github-actions bot added Utils Algorithms labels Jun 25, 2024

rohitrango force-pushed the sdxl_draft branch from 783252b to ccbad85 Compare June 26, 2024 18:48

terrykong self-requested a review June 28, 2024 21:54

terrykong requested changes Jun 28, 2024

View reviewed changes

rohitrango force-pushed the sdxl_draft branch from 3b67453 to 2756c26 Compare July 2, 2024 18:41

terrykong reviewed Jul 2, 2024

View reviewed changes

nemo_aligner/models/mm/stable_diffusion/megatron_sdxl_draftp_model.py Outdated Show resolved Hide resolved

terrykong requested changes Jul 2, 2024

View reviewed changes

nemo_aligner/models/mm/stable_diffusion/megatron_sdxl_draftp_model.py Show resolved Hide resolved

terrykong reviewed Jul 2, 2024

View reviewed changes

examples/mm/stable_diffusion/train_sdxl_draftp.py Outdated Show resolved Hide resolved

terrykong reviewed Jul 2, 2024

View reviewed changes

examples/mm/stable_diffusion/train_sdxl_draftp.py Outdated Show resolved Hide resolved

terrykong reviewed Jul 3, 2024

View reviewed changes

examples/mm/stable_diffusion/train_sdxl_draftp.py Show resolved Hide resolved

terrykong reviewed Jul 3, 2024

View reviewed changes

examples/mm/stable_diffusion/conf/draftp_sdxl.yaml Show resolved Hide resolved

terrykong reviewed Jul 3, 2024

View reviewed changes

nemo_aligner/models/mm/stable_diffusion/megatron_sdxl_draftp_model.py Show resolved Hide resolved

yaoyu-33 approved these changes Jul 8, 2024

View reviewed changes

rohitrango requested a review from terrykong July 8, 2024 18:54

Rohit Jena added 14 commits July 8, 2024 16:01

leftover commit

bf062de

Signed-off-by: Rohit Jena <[email protected]>

leftover commit

996f0d1

Signed-off-by: Rohit Jena <[email protected]>

commits

a278ff5

Signed-off-by: Rohit Jena <[email protected]>

update gitignore

f7ebd21

Signed-off-by: Rohit Jena <[email protected]>

init model for SDXL

6203ea7

Signed-off-by: Rohit Jena <[email protected]>

correct path in eos script

1a232d7

Signed-off-by: Rohit Jena <[email protected]>

modified generate and log_generate scripts to use diffusion engine

9d6bc85

sampling instead Signed-off-by: Rohit Jena <[email protected]>

fixed most runtime bugs -- check for logical bugs

d0af4f8

Signed-off-by: Rohit Jena <[email protected]>

examining mapping between hf and nemo

ecc30f0

Signed-off-by: Rohit Jena <[email protected]>

writing converter script for unet

125e61b

Signed-off-by: Rohit Jena <[email protected]>

tmp commit (moving to eos)

bc1617c

Signed-off-by: Rohit Jena <[email protected]>

more changes to draftp xl

2b46f92

Signed-off-by: Rohit Jena <[email protected]>

changed batch scripts

047df83

Signed-off-by: Rohit Jena <[email protected]>

changed launch scripts

8376132

Signed-off-by: Rohit Jena <[email protected]>

Rohit Jena and others added 23 commits July 8, 2024 16:05

check adapter control

d69566c

Signed-off-by: Rohit Jena <[email protected]>

clean up lora hotswap debugging

149e58e

Signed-off-by: Rohit Jena <[email protected]>

adding fsdp to draftp training

8c7cf40

Signed-off-by: Rohit Jena <[email protected]>

FSDP now works for SDXL?!

e56c20f

Signed-off-by: Rohit Jena <[email protected]>

added custom rule to enable sharding of decoder

d5407a7

Signed-off-by: Rohit Jena <[email protected]>

also shard the clip embeddings

bbb36ea

Signed-off-by: Rohit Jena <[email protected]>

multinode script created + testing 2048 config

9cb21a7

Signed-off-by: Rohit Jena <[email protected]>

add activation checkpointing

1efb73b

Signed-off-by: Rohit Jena <[email protected]>

added activation checkpointing

69faaf3

Signed-off-by: Rohit Jena <[email protected]>

added SFT and PEFT support with Draft+

fa018e3

Signed-off-by: Rohit Jena <[email protected]>

corrected init denoise bug

7ca876f

Signed-off-by: Rohit Jena <[email protected]>

added multinode setup script - test it

3b7ff3d

Signed-off-by: Rohit Jena <[email protected]>

enable sharding for sdlora

5ae5cca

Signed-off-by: Rohit Jena <[email protected]>

added multinode script for OCI

b226823

Signed-off-by: Rohit Jena <[email protected]>

some more refactoring

58cbbfb

Signed-off-by: Rohit Jena <[email protected]>

[pre-commit.ci] auto fixes from pre-commit.com hooks

61118c2

for more information, see https://pre-commit.ci Signed-off-by: Rohit Jena <[email protected]>

address Terry's comments

d67bb14

Signed-off-by: Rohit Jena <[email protected]>

remove + 0 to clone

47d7e21

Signed-off-by: Rohit Jena <[email protected]>

[pre-commit.ci] auto fixes from pre-commit.com hooks

3d626d7

for more information, see https://pre-commit.ci Signed-off-by: Rohit Jena <[email protected]>

add version guard

c9692ca

Signed-off-by: Rohit Jena <[email protected]>

[pre-commit.ci] auto fixes from pre-commit.com hooks

6945884

for more information, see https://pre-commit.ci Signed-off-by: Rohit Jena <[email protected]>

added todo to merge these APIs later

e87c824

Signed-off-by: Rohit Jena <[email protected]>

[pre-commit.ci] auto fixes from pre-commit.com hooks

866df67

for more information, see https://pre-commit.ci Signed-off-by: Rohit Jena <[email protected]>

rohitrango force-pushed the sdxl_draft branch from 5ca1ba7 to 866df67 Compare July 8, 2024 23:08

terrykong approved these changes Jul 12, 2024

View reviewed changes

terrykong merged commit 2a454f9 into NVIDIA:main Jul 12, 2024
3 checks passed

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Draft+ for SDXL [draft] #222

Draft+ for SDXL [draft] #222

rohitrango commented Jun 25, 2024

terrykong left a comment

terrykong commented Jul 8, 2024 •

edited

Loading

terrykong left a comment

Draft+ for SDXL [draft] #222

Draft+ for SDXL [draft] #222

Conversation

rohitrango commented Jun 25, 2024

What does this PR do ?

Changelog

Usage

Before your PR is "Ready for review"

Checklist when contributing a new algorithm

Additional Information

terrykong left a comment

Choose a reason for hiding this comment

terrykong commented Jul 8, 2024 • edited Loading

terrykong left a comment

Choose a reason for hiding this comment

terrykong commented Jul 8, 2024 •

edited

Loading