AutoTP for training?

[KaiLv69]

Hi. I've noticed that DeepSpeed supports AutoTP for many models in HF. That's really useful. And I'm wondering whether it could be used to training models with TP. That's to say, can we do backpropagation with AutoTP?

[delock]

AutoTP relies on LinarAllreduce, LmHeadLinearAllreduce, Linear Layer in the files below. They don't have backward function yet.
https://github.com/microsoft/DeepSpeed/blob/master/deepspeed/module_inject/layers.py

[RezaYazdaniAminabadi]

I think it should be doable, and as @delock suggested we need some backward functions added for several operations that parallelize the model-execution. In addition, we need some high-level apis to manage the training environment such as data-loader, sharding optimizer states, ckpt load/saving, etc. Most of these are already handled through the Megatron-DeepSpeed repo, however, it only support megatron-type of models. If we can borrow some of those logic and build the core to handle tensor-parallelism then we can support more models from HF or other libraries.

[KaiLv69]

Great! Is there any plan to support that?

[delock]

I see this PR addressing this discussion #5445