Integrate with ColossalAI for memory/speed efficiency and improvements

[0xdevalias]

WHAT

Integrate with ColossalAI for memory/speed efficiency and improvements.

WHERE

• https://github.com/hpcaitech/ColossalAI

  • Colossal-AI: A Unified Deep Learning System for Big Model Era

• https://medium.com/@yangyou_berkeley/diffusion-pretraining-and-hardware-fine-tuning-can-be-almost-7x-cheaper-85e970fe207b

HOW

• https://github.com/hpcaitech/ColossalAI/tree/main/examples/images/diffusion

WHEN

Whenever I get to it. Hold your horses... 👌🏻

SEE ALSO

• Community Integration: Making AIGC cheaper, faster, and more efficient. huggingface/diffusers#1212
• [Feature Request]: colossalai integration AUTOMATIC1111/stable-diffusion-webui#4606
• [FEATURE]: Integrate StableDiffusion Dreambooth with ColossalAI hpcaitech/ColossalAI#1935
• Integrate StableDiffusion Dreambooth with ColossalAI TheLastBen/fast-stable-diffusion#493
• Integrate StableDiffusion Dreambooth with ColossalAI ShivamShrirao/diffusers#129

[0xdevalias]

It looks like the dev team of ColossalAI is looking to work directly with huggingface diffusers to simplify integrations in this issue:

• Community Integration: Making AIGC cheaper, faster, and more efficient. huggingface/diffusers#1212

This is very cool @binmakeswell ! Would be very happy to integrate diffusers in colossalAI examples!

Originally posted by @patil-suraj in huggingface/diffusers#1212 (comment)

Thanks for your enthusiastic response and we look forward to working together!

How do you expect it to be carried out?
Take our collaboration with Lightning for example, we communicated in slack for details, provided documentation and eventually did a joint release.

Originally posted by @binmakeswell in huggingface/diffusers#1212 (comment)

And there seems to be at least some initial/cursory interest on the issue I opened on the ColossalAI repo RE: integrating with dreambooth:

• [FEATURE]: Integrate StableDiffusion Dreambooth with ColossalAI hpcaitech/ColossalAI#1935

Hi, thanks for reaching out! I think it’s pretty awesome. Can I check it with our team and get back to you ASAP?

Originally posted by @Sze-qq in hpcaitech/ColossalAI#1935 (comment)

So hopefully things rapidly move towards the 'easier to integrate'/'better documentation' direction of things!

Originally posted by @0xdevalias in ShivamShrirao/diffusers#129 (comment)

[d8ahazard]

So, while I'm always hesitant to say "no" to a feature request...this is a BIG one.

Like, porting Shivam's repo isn't too tedious, because I'm mostly just exposing the args to the UI, adding some stuff to update the text, and tying in extraction and conversion.

This sounds a bit more complex than that. I'd bascially have to go through the existing "train_dreambooth.py" script, rip out some big components, and then replace them with the optimized stuff from ColossalAI.

Still, I definitely see how it could be very useful to other people, so I'll say "Let's put a pin in this one for now." :P

[sgsdxzy]

From my study, current training scripts can be categorized to two types:
One uses the stable diffusion repo, ldm and torch.lightning to train, including the original training in stable diffusion, pokemon diffusion, old waifu diffusion and colossal ai.
The other uses diffusers+accelerate, including ShivamShrirao's diffuser that this extension is based on, kohya_ss, new waifu diffusion, etc.
I find it difficult to adapt tools used by one to the other. The routines are just too different. One thing to note is that webui uses ldm and not diffusers, and I find integrating webui routines into this extension is already not easy. So the current colossal ai example uses the first training pipeline, unless someone understand all these very well I do not expect it to land very soon.

[78Alpha]

From what I get from it, it's like DeepSpeed, does the same thing more or less, but it doesn't have the same level of documentation (None, last I checked). So, it would be a painful process.

[github-actions]

This issue is stale because it has been open 5 days with no activity. Remove stale label or comment or this will be closed in 5 days

[0xdevalias]

This feels relevant to keep open even if it's not something that would immediately be worked on?

[sgsdxzy]

Update: I just find ColossalAI does NOT support Windows. So I doubt it will be of much use (Linux can already use DeepSpeed).

[0xdevalias]

I guess the big question though is whether this is better/faster than deep speed, which it kind of seemed to be pitching that it was.

[github-actions]

This issue is stale because it has been open 5 days with no activity. Remove stale label or comment or this will be closed in 5 days

[d8ahazard]

Yeah, so, we've had a lot of improvements to VRAM and speed, but I don't see ColossalAI in the cards. As such, I'm gonna close this.

[78Alpha]

Just adding a note, looks like it was added to a few things

huggingface/diffusers#1212 (comment)

[0xdevalias]

Any update @binmakeswell ? It would be very great to fasten training process.

I found train_dreambooth.py and train_dreambooth_colossalai.py that are based on diffusers.

Thank you for your interest and contribution to colossalai, we have merged into the hugging face diffuser library.
https://github.com/huggingface/diffusers/tree/main/examples/research_projects/colossalai

Originally posted by @Fazziekey in huggingface/diffusers#1212 (comment)