ZeRO stage 3 is not a good choice either for the same reason - more inter-node communications required.
And since we have ZeRO, the other benefit is ZeRO-Offload. Since this is stage 1 optimizer states can be offloaded to CPU.
Implementations:
- Megatron-DeepSpeed
- OSLO
🤗 Transformers status: not yet implemented, since we have no PP and TP.
FlexFlow
FlexFlow also solves the parallelization problem in a slightly different approach.
Paper: "Beyond Data and Model Parallelism for Deep Neural Networks" by Zhihao Jia, Matei Zaharia, Alex Aiken
It performs a sort of 4D Parallelism over Sample-Operator-Attribute-Parameter.
Sample = Data Parallelism (sample-wise parallel)
Operator = Parallelize a single operation into several sub-operations
Attribute = Data Parallelism (length-wise parallel)
Parameter = Model Parallelism (regardless of dimension - horizontal or vertical)
