patching lora load for faster loading #39
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
fofr
force-pushed
the
fast-lora-load
branch
from
126730e
to
34b4947
Compare
andreasjansson
approved these changes
Sep 27, 2024
| @@ -101,6 +102,9 @@ def setup(self) -> None: # pyright: ignore | |||
| "FLUX.1-dev", | |||
| torch_dtype=torch.bfloat16, | |||
| ).to("cuda") | |||
| dev_pipe.__class__.load_lora_into_transformer = classmethod( | |||
There was a problem hiding this comment.
This is fine, but could probably also done a similar pattern to https://github.com/replicate/flux-fine-tuner/blob/main/submodule_patches.py to avoid duplicating this line.
| @@ -430,6 +440,7 @@ def load_single_lora(self, lora_url: str, model: str): | |||
| lora_path = self.weights_cache.ensure(lora_url) | |||
| pipe.load_lora_weights(lora_path, adapter_name="main") | |||
| self.loaded_lora_urls[model] = LoadedLoRAs(main=lora_url, extra=None) | |||
| pipe = pipe.to("cuda") | |||
There was a problem hiding this comment.
Is it not already on gpu at this point? How did it work before? Might be worth adding a comment here.
andreasjansson
force-pushed
the
fast-lora-load
branch
from
0861945
to
f53e34c
Compare
|
Did a test-only cog-safe-push, it passed https://github.com/replicate/flux-fine-tuner/actions/runs/11072824989 |
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.
Profiled lora loading in prod to see what was taking so long - it looks like our old nemesis
kaiming_uniform, and other such functions that are used to randomly initialize empty tensors.conveniently, someone just pushed functionality to peft to disable this. So we don't have to get too wild to turn it off; we just need to patch
load_lora_into_transformerto take advantage of that functionality until it's integrated into diffusers.That's what this PR does. With this I saw lora load times (after download) drop from about 10 seconds to about 1.2 seconds. Tested locally w/dev and schnell.
Important
Patch
load_lora_into_transformerto uselow_cpu_mem_usage=True, significantly reducing LoRA load times, and update dependencies accordingly.load_lora_into_transformerinlora_loading_patch.pyto uselow_cpu_mem_usage=True, reducing LoRA load times from ~10s to ~1.2s.predict.pyto use the patchedload_lora_into_transformerforFluxPipeline,FluxImg2ImgPipeline, andFluxInpaintPipeline.peftversion from0.12.0to0.13.0incog.yaml./weights-cacheto.dockerignore.This description was created by
for f53e34c. It will automatically update as commits are pushed.