filipstrand/mflux

Feature Request: Train lora

Opened this issue · 3 comments

kirel commented

Docs say

actual training support is coming

— is there any work visible somewhere?

anthonywu commented

@kirel follow this adjacent work: https://github.com/ml-explore/mlx-examples/tree/main/flux#training-example - this project is upstream from the MLX team, and I think mflux could/should just accept the trained Loras from that project

azrahello commented

After training a LoRA with MLX-examples, I tried using the LoRA with MFLUX but it’s not working. Every time I try to train and use a LoRA with MFLUX, it also downloads the model again, and I don’t know why, but it’s very annoying. I’m not sure if it’s a kind of LoRA that’s not compatible with MFLUX or if the issue comes from a mistake I made during the training

filipstrand commented

As Anthony suggested, right now I would also recommend the mlx-examples project if you want to run fine-tuning locally. I quickly tried loading a LoRA from mlx-examples, and can confirm that it does not work with mflux currently. Inspecting the weight matrices, there seem to be some other naming conventions. But I think a renaming and reshaping can fix this. For example,

state_dict["lora_unet_single_blocks_37_linear2.lora_down.weight"] = state_dict["single_blocks.37.linear2.lora_a"].transpose(0, 1)

I have not actually tried this yet as my one fine-tuning attempt failed (probably my fault for not training with enough iterations). But I can at least confirm that the shapes seem to match and the LoRA does load in mflux after the renaming and reshaping. If anyone could point me to a successful LoRA that was fine-tuned from mlx-examples, it would be really nice as a test case to confirm that the translation is working as expected.

Regarding fine-tuning support in mflux - I have personally not found the time lately to really dig into this feature as I first planned (mflux is still a nights and weekends project for me, but I would have loved to be able to spend more time on it :) ). My hope is to eventually find the time to implement this, and, of course, help with reviewing PRs if anyone else contributes a working implementation.

  • 👍1