[PROPOSAL] Project Elixir

[1SAA]

Elixir (also called Gemini2.0)

Elixir is new feature used to substitute current Gemini in ColossalAI. Elixir should not be viewed as an incremental update of Gemini. It is a new designed feature which does the same thing as Gemini but they are based on different principles. The difference is that Gemini dynamically allocates chunk footprints and profiles CUDA memory usage while Elixir allocates and profiles statically before running models.

Users can track the development of Elixir here.

cc @FrankLeeeee, @ver217 , @YuliangLiu0306 , @kurisusnowdeng.

Motivation (Drawbacks of Gemini)

• poor compatibility: Gemini fails on a lot of models now, such as ResNet, EfficentNet, Swin Transformer, etc. The reason is that the parameter hook mechanism of Gemini is based on ColoParameter, a derived class from ColoTensor which is used for "automatic" tensor parallelism. The poor compatibility of ColoTensor makes Gemini fail at a large range of models. As we shall deprecate ColoTensor and use DTensor instead for real automatic tensor parallelism, we can design a new parameter hook mechanism for all common models.
• CUDA memory fragments: Since Gemini dynamically allocates and releases chunk footprints, it generates tons of fragments in CUDA memory. The result is that auto policy always fails easily. Though the profiler detects enough memory space for future activations, there has no enough continous memory space and raises an OOM error instead.
• performance drop: Gemini's excellent performance drops in a slow communication environment. In our tests, the cross-node training drops the TFLOPS to 90 from 180. In order to alleviate the issue caused by slow communications, we require overlapping communications with computations which is hard to implement in the current dynamic framework.

What's New in Elixir

• A robust parameter hook mechanism used to upload and offload parameters. Currently, this robust mechanism is tested on a large range of models in torchvision, timm and HuggingFace. It is compatible with enormous popular models such as ResNet, ViT, Swin Transformer, GPT, Albert, OPT, T5, UNet, etc.
• Elixir traces models before running it. Elixir can get the called order of parameters and the maximum activation memory usage. Thus, Elixir is able to allocate memory footprints, decide whether to use a specific parallel pattern like ZeRO2 or ZeRO3, and look for an optimal chunk initialization plan which minimizes the total communication time. This feature is based on a parameter order tracer and a memory profiler which are under development.
• Elixir enables prefetching chunks to overlap communications with computations. After Elixir detects the called order of all parameters, it can prefetch the next used chunk during traning.

High Level Insights

Elixir consists of three main parts SearchAlgorihtm, ElixirModule and ElixirOptimizer. ElixirModule and ElixirOptimizer are designed to be coupled together. One ElixirOptimizer is stricted corresponds to one specific ElixirModule.

• SearchAlgorithm is assigned by users. It will look for a chunk initialization plan and initialize a allocated memory pool based on the plan. There are three types of SearchAlgorithm. They are PrefetchSA, MinimumWasteSA and SimpleSA. PrefetchSA needs users provide a step function for model tracing. MinimumWasteSA doesn't need a step function, it look for a chunk plan which minimizes CUDA memory usage. SimpleSA is used for unit tests, it just simply split parameters into blocks with a pointed number of blocks.
• ElixirModule is very similar with the current GeminiDDP. It supports forward and backward.
• ElixirOptimizer is almost the same as current ZeroOptimizer. But it should supports full fp32 training.

The high level control flow is shown below:

Low Level Memory Management

The bottom memory management of chunks is based on this paper written by me.

[HongLouyemeng]

你好，这个项目会合并到colossal-ai到官方库吗？