Add GPTNeoX AutoTP support (#3778)

Co-authored-by: NOlatunji Ruwase <olruwase@microsoft.com>
Co-authored-by: NMolly Smith <112220543+molly-smith@users.noreply.github.com>

deepspeed/module_inject/auto_tp.py

@@ -108,6 +108,8 @@ class AutoTP():
                     gem_list = gem_list + [layer]
                 elif 'down_proj' in layer:
                     gem_list = gem_list + [layer]
+                elif 'attention.dense' in layer and 'GPTNeoX' in str(model):
+                    gem_list = gem_list + [layer]
                 elif 'self_attention.dense' in layer and 'falcon' in str(
                         type(module)):  # this is a hack to get the right linear layer for this model!
                     gem_list = gem_list + [layer]