flux lora train support

[AmazDeng]

@FrankLeeeee @gothicx @tiansiyuan @jeffra

Does ColossalAI support training Flux model? For example, if I'm using a LoRA paradigm and need to redefine the processor within Flux, is this training method supported by ColossalAI?Additionally, does ColossalAI support custom models?

[flybird11111]

Hello, ColossalAI's TP and PP strategies do not support Flux. You may try using the Gemini strategy, but please note that the Gemini strategy does not support LoRA training.No, ColossalAI does not support custom models.

[AmazDeng]

@flybird11111
Thanks for your answer.

1.So, what specific models does ColossalAI support? I haven't seen a list of supported models.
2.The Open-Sora project (https://github.com/hpcaitech/Open-Sora) uses the ColossalAI engine. I've noticed that this project contains a lot of custom components(for example: https://github.com/hpcaitech/Open-Sora/blob/main/opensora/models/mmdit/distributed.py). Why is it able to run with ColossalAI? Can I consider Open-Sora as a custom model [in the context of ColossalAI's usage]?

[flybird11111]

In OpenSora, there are customized policies that need to be implemented by themself. The models that have already been implemented can be found in the README of https://github.com/hpcaitech/ColossalAI.

[AmazDeng]

@flybird11111

I noticed that Open-Sora and ColossalAI are from the same team. Does this mean that if I were to modify Flux myself into a training paradigm of LoRA + custom processor, and then use ColossalAI's parallel training techniques, the difficulty would be extremely high? Since Open-Sora and ColossalAI are from the same team, it's relatively easier for them to make such modifications. Can I understand it this way? Do I need to understand ColossalAI's various internals/kernels to perform such modifications?