Integrate Colossal.AI Engine

[siddk]

🚀 Feature

Similar to the DeepSpeed/FairScale integrations, it'd be really cool for PyTorch Lightning to expose an API for integrating Colossal.AI.

Motivation

Having 3D parallelism and other optimizations as a simple plugin for PyTorch Lightning would make scaling large models super easy!

Pitch

Colossal AI exposes a simple engine-based API similar to DeepSpeed (https://www.colossalai.org/docs/basics/engine_trainer#engine)... should be straightforward to integrate?

cc @Borda @akihironitta @justusschock

[justusschock]

Hey @siddk that would be awesome. Would you be interested to contribute? I believe that you are right and the integration should be very similar to DeepSpeed. Maybe it is even possible to factor out some common utilities.

We would really appreciate your contribution and definitely help you to iterate over an initial PR :)

[siddk]

Hey @justusschock - I wish I had the time, but this would be a high-value feature for both some of the work we're doing on training large models at Stanford (https://crfm.stanford.edu/), and a crucial part of several large-scale projects.

If there was 1-2 members of the team (or other contributors), I'd definitely be down to help test, and provide feedback, I just don't think I'd be able to take on all of this alone.

[marsggbo]

A good suggestion! Any progress now?

[tchaton]

Hey @marsggbo @siddk @rohitgr7,

I reached out to Collosal-AI Team hpcaitech/ColossalAI#1330.
Let s see if we can collaborate with their team to address your needs.

We will keep you updated on progresses.

Best,
T.C

[binmakeswell]

Hi, thank you for your attention and advice, Colossal-AI Team is willing to help provide support for this feature as soon as possible.