[Bug] Illegal Memory Access in Attention Layer when Using fp16

[DanielSHKao]

Thanks for your impactful work of deepspeed. I am experiencing difficulties with model parallelism when using fp16 with transformer-based model. Specifically, I received the following bug message:

, where the self.model has the following architecture...:

I am using deepspeed 0.10.3, transformers 4.31.0, torch 2.1.2+cu118 and CUDA 11.8 on 2 RTX4090 GPUs. Do you have any idea how to solve this bug?

Thanks in advance.