ColossalAI vs. Fairscale

[nostalgicimp]

Hi, Happy Mid-Autumn Day~

I saw this impressive figure for Github star history among ColossalAI, Megatron and Fairscale. There must be great reasons why ColossalAI is gaining popularity with strong momentum to have surpassed Fairscale in such a sharp slope. How do you compare ColossalAI with Fairscale, since Fairscale (FSDP) is native/friendly to pytorch, and also has efficient memory management, data parallelism, 1-D model/pipeline parallelism with FSDP, and model offloading to CPU.

Other than NVMe offloading and ZeRO Redundancy Optimizer with chuck-based memory management, does Colossal AI has more distinctive features that beat Fairscale? How much does the ZeRO Redundancy Optimizer influence the performance compared to Fairscale which does not have this feature?

Also in terms of fp16 mix precision support, how do you compare ZeRO, ColossalAI AMP, or Fairscale? since fp16 and zero configuration are not compatible, and Pipeline only support ColossalAI naive AMP currently.

[FrankLeeeee]

Colossal-AI is native to PyTorch as well, just that it is not maintained by Meta. Compared to FairScale, our ZeRO and Offloading modules surely outperform those in FairScale. This is because that we build our system from the Tensor layer, we created a new Tensor called ColoTensor (compatible with torch.Tensor) and design in such a way that it supports distributed training natively. Besides that, we also support tensor and pipeline parallelisms, one exciting feature coming is auto parallelism which automates all parallelisms.

As for mixed precision training, The ZeRO module in ColossalAI does use FP16 and has some modifications. Thus, it does not need colossalai.amp, or in other words, not compatible.

[FrankLeeeee]

Yes, unfortunately, zero and naive amp will cast all ops to fp16, including batch norm. This is similar to Apex O3. Currently we find this works fine.

[nostalgicimp]

What do you mean by "we find this works fine"? Apex O3?
Does mix-precision conflict with colossalAI's tensor/pipeline parallelism such that if I want to use mix-precision, I can only use data-parallelism with no "ZeRO-like" capability?

[nostalgicimp]

I got confused of naive.amp by the following 3 statements. Could you also help explain a little bit?

• "The last method is similar to Apex O2 level." I assume this is saying naive.amp is similar to Apex O2 option.
• "O2: “Almost FP16” Mixed Precision: O2 casts the model weights to FP16, patches the model’s forward method to cast input data to FP16, keeps batchnorms in FP32, maintains FP32 master weights, updates the optimizer’s param_groups so that the optimizer.step() acts directly on the FP32 weights (followed by FP32 master weight->FP16 model weight copies if necessary), and implements dynamic loss scaling (unless overridden). " This is saying O2 almost cast all to fp16, but still keeps batchnorms in fp32, and maintains fp32 master weights.
• "naive.amp achieved mixed precision training while maintaining compatibility with complex tensor and pipeline parallelism. This AMP mode will cast all operations into fp16" Is it saying naive.amp downcast all to fp16?

[FrankLeeeee]

Hi, sorry for the confusion. Just to clarify here. Apex O2 keeps BN in fp32 and cast other OPs to fp16 while maintaining a copy of master weights. In ColossalAI, naive AMP and zero cast all OPs to fp16 and keep a copy of master weight, but do not keep BN in fp32, thus, this is similar to O2, but not the same as O2. We have tested the model convergence for ResNet with BN in fp16 and the result is fine.

[FrankLeeeee]

What do you mean by "we find this works fine"? Apex O3? Does mix-precision conflict with colossalAI's tensor/pipeline parallelism such that if I want to use mix-precision, I can only use data-parallelism with no "ZeRO-like" capability?

Naive amp is compatible with tensor/pipeline parallelism.