ColossalAI icon indicating copy to clipboard operation
ColossalAI copied to clipboard

flux lora train support

Open AmazDeng opened this issue 4 months ago • 4 comments

@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?

AmazDeng avatar Aug 15 '25 08:08 AmazDeng

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.

flybird11111 avatar Aug 18 '25 08:08 flybird11111

@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]?

AmazDeng avatar Aug 18 '25 12:08 AmazDeng

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.

flybird11111 avatar Aug 19 '25 02:08 flybird11111

@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?

AmazDeng avatar Aug 20 '25 03:08 AmazDeng