Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Reduce Memory Cost in Flux Training #9829

Merged
merged 14 commits into from
Nov 1, 2024
+12 −0
Merged

Reduce Memory Cost in Flux Training #9829

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
14 commits
Select commit Hold shift + click to select a range
e49d112
Improve NPU performance
Oct 11, 2024
d25d229
Improve NPU performance
Oct 12, 2024
9555cf3
Improve NPU performance
Oct 12, 2024
2b0febe
Improve NPU performance
Oct 14, 2024
0365ae4
Merge branch 'main' into main_npu
leisuzz Oct 14, 2024
458f68f
Merge branch 'huggingface:main' into main_npu
leisuzz Oct 24, 2024
5a742e6
[bugfix] bugfix for npu free memory
Oct 24, 2024
42271f1
[bugfix] bugfix for npu free memory
Oct 24, 2024
62a718c
[bugfix] bugfix for npu free memory
Oct 24, 2024
4ef65c9
Merge branch 'main' into main_npu
sayakpaul Oct 24, 2024
c4e509e
Merge branch 'main' into main_npu
leisuzz Oct 25, 2024
f9966d7
Merge branch 'huggingface:main' into main_npu
leisuzz Nov 1, 2024
2dbbe34
Reduce memory cost for flux training process
Nov 1, 2024
5e4eb2f
Merge branch 'main' into main_npu
leisuzz Nov 1, 2024
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
6 changes: 6 additions & 0 deletions examples/dreambooth/train_dreambooth_flux.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -1740,6 +1740,9 @@ def get_sigmas(timesteps, n_dim=4, dtype=torch.float32):
torch_npu.npu.empty_cache()
gc.collect()

images = None
del pipeline

# Save the lora layers
accelerator.wait_for_everyone()
if accelerator.is_main_process:
Expand Down Expand Up @@ -1798,6 +1801,9 @@ def get_sigmas(timesteps, n_dim=4, dtype=torch.float32):
ignore_patterns=["step_*", "epoch_*"],
)

images = None
del pipeline

accelerator.end_training()


Expand Down
6 changes: 6 additions & 0 deletions examples/dreambooth/train_dreambooth_lora_flux.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -1844,6 +1844,9 @@ def get_sigmas(timesteps, n_dim=4, dtype=torch.float32):
del text_encoder_one, text_encoder_two
free_memory()

images = None
del pipeline

# Save the lora layers
accelerator.wait_for_everyone()
if accelerator.is_main_process:
Expand Down Expand Up @@ -1908,6 +1911,9 @@ def get_sigmas(timesteps, n_dim=4, dtype=torch.float32):
ignore_patterns=["step_*", "epoch_*"],
)

images = None
Copy link
Collaborator

@linoytsaban linoytsaban Nov 1, 2024 •
edited
Loading

Choose a reason for hiding this comment

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

@leisuzz @sayakpaul great contribution thanks! I want to add these changes to the advanced flux script too, is this last part specifically helpful? since it happens after the weights are already uploaded/pushed to the hub

Copy link
Contributor Author

Choose a reason for hiding this comment

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

It actually depends on the limit of memory and the user's hardware, but it will clean some memory usages for sure

del pipeline

accelerator.end_training()


Expand Down
Loading