Reduce Memory Cost in Flux Training #9829
Conversation
|
The docs for this PR live here. All of your documentation changes will be reflected on that endpoint. The docs are available until 30 days after the last update. |
|
Thanks for your PR. Can you also maybe report the memory it helps in saving? |
|
With the some configs, without these modifications, I faced OOM with batch size = 1. After modification, I can run the test with batch size = 16. |
* Improve NPU performance * Improve NPU performance * Improve NPU performance * Improve NPU performance * [bugfix] bugfix for npu free memory * [bugfix] bugfix for npu free memory * [bugfix] bugfix for npu free memory * Reduce memory cost for flux training process --------- Co-authored-by: 蒋硕 <[email protected]> Co-authored-by: Sayak Paul <[email protected]>
| @@ -1908,6 +1911,9 @@ def get_sigmas(timesteps, n_dim=4, dtype=torch.float32): | |||
| ignore_patterns=["step_*", "epoch_*"], | |||
| ) | |||
|
|
|||
| images = None | |||
There was a problem hiding this comment.
@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
There was a problem hiding this comment.
It actually depends on the limit of memory and the user's hardware, but it will clean some memory usages for sure
) * memory improvement as done here: #9829 * fix bug * fix bug * style --------- Co-authored-by: Sayak Paul <[email protected]>
* Improve NPU performance * Improve NPU performance * Improve NPU performance * Improve NPU performance * [bugfix] bugfix for npu free memory * [bugfix] bugfix for npu free memory * [bugfix] bugfix for npu free memory * Reduce memory cost for flux training process --------- Co-authored-by: 蒋硕 <[email protected]> Co-authored-by: Sayak Paul <[email protected]>
) * memory improvement as done here: #9829 * fix bug * fix bug * style --------- Co-authored-by: Sayak Paul <[email protected]>
What does this PR do?
Reduce memory cost while training flux model, and reduce OOM issues.
Helps to fix issues: #9790, #9732, #9501, #9500, and #9156.
Before submitting
documentation guidelines, and
here are tips on formatting docstrings.
Who can review?
Anyone in the community is free to review the PR once the tests have passed. Feel free to tag
members/contributors who may be interested in your PR.