Is ZeRO-3 compatible to tensor parallel?

[skyshine102]

In recent blogs posts,

AWS:

• https://aws.amazon.com/tw/blogs/machine-learning/amazon-sagemaker-model-parallel-library-now-accelerates-pytorch-fsdp-workloads-by-up-to-20/

Nvidia NeMO:

• FSDP + Tensor Parallelism NVIDIA/NeMo#7897

AWS makes transformer engine's tensor parallel into FSDP, which is "similar" algorithm to ds ZeRO-3.

• I'm curious if we can do something similary with ds ZeRO-3? (I know Megatron uses ZeRO-1+TP)
  • If so, how can we wrap sub-modules for TP and the rest for ZeRO-3? While FSDP provides auto_wrap function, I struggled to find a correspondence in deepspeed.

(Really looking forwards to seeing these hybrid sharding technology being easily accessible)

[tjruwase]

@skyshine102, zero-3 can be combined with tensor parallelism and we validated this combination in our megatron-deespeed repo a while ago. This tutorial might be helpful: https://www.deepspeed.ai/tutorials/zero/

[skyshine102]

Cool. Thanks for the pointer. I will take a look