-
Notifications
You must be signed in to change notification settings - Fork 2.5k
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.
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
FLUX memory management improvements #6791
Merged
Merged
Conversation
This file contains 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
brandonrising
approved these changes
Aug 29, 2024
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Lgtm
…efaults were misleading, because the config defaults take precedence over them.
…efore loading it.
…to persist after loading from a state dict. This manifested as state tensors being left on the GPU even when a model had been offloaded to the CPU cache.
… that all model references are locally-scoped so that the two models don't have to be help in memory at the same time.
… clear the cache based on the on-disk model size.
42992cc
to
4e4b6c6
Compare
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Labels
backend
PRs that change backend files
invocations
PRs that change invocations
python
PRs that change python files
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.
Summary
This PR contains several improvements to memory management for FLUX workflows.
It is now possible to achieve better FLUX model caching performance, but this still requires users to manually configure their
ram
/vram
settings. E.g. avram
setting of 16.0 should allow for all quantized FLUX models to be kept in memory on the GPU.Changes:
InvokeLinear8bitLt
that was causing some tensors to be left on the GPU when the model was offloaded to the CPU. (This class is getting very messy due to the non-standard state_dict handling inbnb.nn.Linear8bitLt
. )Future:
We should revisit our default ram/vram configs. The current defaults are very conservative, and users could see major performance improvements from tuning these values.
QA Instructions
I tested the FLUX workflow with the following configurations and verified that the cache hit rates and memory usage matched the expected behaviour:
ram = 16
andvram = 16
ram = 16
andvram = 1
ram = 1
andvram = 1
Note that the changes in this PR are not isolated to FLUX. Since we now check the size of models on disk, we may see slight changes in model cache offload patterns for other models as well.
Checklist