Use safetensors to store Loras #74
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
|
This is great, the metadata will come in handy! |
|
Awesome! I will certainly check this out |
|
@hafriedlander Hi, so I've merged your #73, can you check these conflicts? |
|
Weird, that should have happened. 🤷 Sure, I'll rebase in a bit. |
hafriedlander
force-pushed
the
feature/safetensors
branch
from
73d58e7
to
a2c8189
Compare
|
@cloneofsimo done, looks good now. I'll do a bit of testing. If I have time I'll update the CLI with a mode to convert to / from safetensors too, but that'll probably be after christmas. |
|
Thanks! I'll use these features myself! Thank you so much for your contribution @hafriedlander! Btw, do you want to become a maintainer? |
|
this is great... i didn't notice the addition of GEGLU till just now, and also because Flax presents pytrees/dicts in alpha order and torch doesn't save tensors with names, it makes it hard to load torch weights to flax version (and i guess a non-GEGLU one in torch as well). i had to write a mess of functions that would build the torch model just to extract its layer names in the order presumably saved by lora, then convert/rename/map those torch.loaded weights to the flax version would it not be better for all methods and compatibilities to just save to JSON? this way no one needs to have "the other library" installed and instantiated to load weights... |
cloneofsimo
added a commit
that referenced
this pull request
Dec 25, 2022
* Add parameter to control rank of decomposition (#28) * ENH: allow controlling rank of approximation * Training script accepts lora_rank * feat : statefully monkeypatch different loras + example ipynb + readme * Fix lora inject, added weight self apply lora (#39) * Develop (#34) * Add parameter to control rank of decomposition (#28) * ENH: allow controlling rank of approximation * Training script accepts lora_rank * feat : statefully monkeypatch different loras + example ipynb + readme Co-authored-by: brian6091 <brian6091@gmail.com> * release : version 0.0.4, now able to tune rank, now add loras dynamically * readme : add brain6091's discussions * fix:inject lora in to_out module list * feat: added weight self apply lora * chore: add import copy * fix: readded r Co-authored-by: Simo Ryu <35953539+cloneofsimo@users.noreply.github.com> Co-authored-by: brian6091 <brian6091@gmail.com> Co-authored-by: SimoRyu <cloneofsimo@korea.ac.kr> * Revert "Fix lora inject, added weight self apply lora (#39)" (#40) This reverts commit fececf3. * fix : rank bug in monkeypatch * fix cli fix * visualizatio on effect of LR * Fix save_steps, max_train_steps, and logging (#45) * v 0.0.5 (#42) * Add parameter to control rank of decomposition (#28) * ENH: allow controlling rank of approximation * Training script accepts lora_rank * feat : statefully monkeypatch different loras + example ipynb + readme * Fix lora inject, added weight self apply lora (#39) * Develop (#34) * Add parameter to control rank of decomposition (#28) * ENH: allow controlling rank of approximation * Training script accepts lora_rank * feat : statefully monkeypatch different loras + example ipynb + readme Co-authored-by: brian6091 <brian6091@gmail.com> * release : version 0.0.4, now able to tune rank, now add loras dynamically * readme : add brain6091's discussions * fix:inject lora in to_out module list * feat: added weight self apply lora * chore: add import copy * fix: readded r Co-authored-by: Simo Ryu <35953539+cloneofsimo@users.noreply.github.com> Co-authored-by: brian6091 <brian6091@gmail.com> Co-authored-by: SimoRyu <cloneofsimo@korea.ac.kr> * Revert "Fix lora inject, added weight self apply lora (#39)" (#40) This reverts commit fececf3. * fix : rank bug in monkeypatch * fix cli fix * visualizatio on effect of LR Co-authored-by: brian6091 <brian6091@gmail.com> Co-authored-by: Davide Paglieri <paglieridavide@gmail.com> * Fix save_steps, max_train_steps, and logging Corrected indenting so checking save_steps, max_train_steps, and updating logs are performed every step instead at the end of an epoch. Co-authored-by: Simo Ryu <35953539+cloneofsimo@users.noreply.github.com> Co-authored-by: brian6091 <brian6091@gmail.com> Co-authored-by: Davide Paglieri <paglieridavide@gmail.com> * Enable resuming (#52) * v 0.0.5 (#42) * Add parameter to control rank of decomposition (#28) * ENH: allow controlling rank of approximation * Training script accepts lora_rank * feat : statefully monkeypatch different loras + example ipynb + readme * Fix lora inject, added weight self apply lora (#39) * Develop (#34) * Add parameter to control rank of decomposition (#28) * ENH: allow controlling rank of approximation * Training script accepts lora_rank * feat : statefully monkeypatch different loras + example ipynb + readme Co-authored-by: brian6091 <brian6091@gmail.com> * release : version 0.0.4, now able to tune rank, now add loras dynamically * readme : add brain6091's discussions * fix:inject lora in to_out module list * feat: added weight self apply lora * chore: add import copy * fix: readded r Co-authored-by: Simo Ryu <35953539+cloneofsimo@users.noreply.github.com> Co-authored-by: brian6091 <brian6091@gmail.com> Co-authored-by: SimoRyu <cloneofsimo@korea.ac.kr> * Revert "Fix lora inject, added weight self apply lora (#39)" (#40) This reverts commit fececf3. * fix : rank bug in monkeypatch * fix cli fix * visualizatio on effect of LR Co-authored-by: brian6091 <brian6091@gmail.com> Co-authored-by: Davide Paglieri <paglieridavide@gmail.com> * Enable resume training unet/text encoder (#48) * Enable resume training unet/text encoder New flags --resume_text_encoder --resume_unet accept the paths to .pt files to resume. Make sure to change the output directory from the previous training session, or else .pt files will be overwritten since training does not resume from previous global step. * Load weights from .pt with inject_trainable_lora Adds new loras argument to inject_trainable_lora function which accepts path to a .pt file containing previously trained weights. Co-authored-by: Simo Ryu <35953539+cloneofsimo@users.noreply.github.com> Co-authored-by: brian6091 <brian6091@gmail.com> Co-authored-by: Davide Paglieri <paglieridavide@gmail.com> * feat : low-rank pivotal tuning * feat : pivotal tuning * v 0.0.6 * Learning rate switching & fix indent (#57) * Learning rate switching & fix indent Make learning rates switch from training textual inversion to unet/text encoder after unfreeze_lora_step. I think this is how it was explained in the paper linked(?) Either way, it might be useful to add another parameter to activate unet/text encoder training at a certain step instead of at unfreeze_lora_step. This would let the user have more control. Also fix indenting to make save_steps and logging work properly. * Fix indent fix accelerator.wait_for_everyone() indent according to original dreambooth training * Re:Fix indent (#58) Fix indenting of accelerator.wait_for_everyone() according to original dreambooth training * ff now training default * feat : dataset * feat : utils to back training * readme : more contents. citations, etc. * fix : weight init * Feature/monkeypatch improvements (#73) * Refactor module replacement to work with nested Linears * Make monkeypatch_remove_lora remove all LoraInjectedLinear instances * Turn off resizing images with --resize=False (#71) * Make image resize optional with --resize Toggle off image resizing using --resize=False. Default is true for to maintain consistent operation. * Make image resize optional with --resize Toggle off image resizing using --resize=False. Default is true for to maintain consistent operation. * Make image resize optional with --resize Toggle off image resizing using --resize=False. Default is true for to maintain consistent operation. * Revert "Turn off resizing images with --resize=False (#71)" (#77) This reverts commit 39affb7. * Use safetensors to store Loras (#74) * Add safetensors supports * Add some documentation for the safetensors load and save methods * Fix typing-related syntax errors in Python < 3.10 introduced in recent refactor (#79) * Fix the --resize=False option (#81) * Make image resize optional with --resize Toggle off image resizing using --resize=False. Default is true for to maintain consistent operation. * Make image resize optional with --resize Toggle off image resizing using --resize=False. Default is true for to maintain consistent operation. * Make image resize optional with --resize Toggle off image resizing using --resize=False. Default is true for to maintain consistent operation. * Fix resize==False functionality * Update train_lora_pt_caption.py * Update train_lora_w_ti.py * Pivotal Tuning with hackable training code for CLI (#83) * feat : save utils on lora * fix : stochastic attribute * feat : cleaner training code * fix : bit of bugs on inspect and trainer * fix : moved pti training to cli * feat : patch now accepts target arg * fix : gelu in target * fix : gradient being way too large : autocast was the problem * fix : hflip * fix : example running well! * merge master Co-authored-by: brian6091 <brian6091@gmail.com> Co-authored-by: Davide Paglieri <paglieridavide@gmail.com> Co-authored-by: hdeezy <82070413+hdeezy@users.noreply.github.com> Co-authored-by: Hamish Friedlander <hafriedlander@gmail.com>
|
@krahnikblis PyTorch pickles can contain arbitrary code. They're not very safe to load - you need to trust the source completely. There's various ways to try and improve safety, but a dedicated attacker can probably work around them. So avoiding using pytorch pickles at all is best. You could just save all the lora weights directly into JSON files, but that's quite slow (and increases size by about 33%). There is a pure python implementation of just the loader side of safetensors, so if you were mostly concerned about being able to use Loras then I could look at adding that? (It wouldn't help when training though). |
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.
See just commit e622ee3, this uses code from a PR that's still open.
Safetensor allows saving:
In a way that is safe to load, even if you don't trust the source
(This is not true for .pt files, which can contain arbitrary executable code).
This PR uses safetensor to save the Lora for multiple modules (unet, text_encoder) in a single file, along with the rank per Lora, and the target classes.