From 248d63e9108eb8463b064efad19d10faab396dd4 Mon Sep 17 00:00:00 2001 From: Andrey Chernykh Date: Fri, 7 Oct 2022 07:08:22 +0000 Subject: [PATCH] Fix GPT Neo-X multi-gpu inference --- deepspeed/module_inject/replace_module.py | 5 ++--- 1 file changed, 2 insertions(+), 3 deletions(-) diff --git a/deepspeed/module_inject/replace_module.py b/deepspeed/module_inject/replace_module.py index ea0e13726316..a8afe46360fa 100755 --- a/deepspeed/module_inject/replace_module.py +++ b/deepspeed/module_inject/replace_module.py @@ -454,9 +454,8 @@ def transpose(data): dense_w.data = transpose(dense_w.data) def _transpose(x): - num_attention_heads_per_partition = transformer_config.heads // transformer_config.mp_size - attention_head_size = x.shape[-1] // num_attention_heads_per_partition - new_x_shape = x.size()[:-1] + (num_attention_heads_per_partition, + attention_head_size = x.shape[-1] // transformer_config.heads + new_x_shape = x.size()[:-1] + (transformer_config.heads, attention_head_size) x_1 = x.view(*new_x_shape) (q, k, v) = torch.split(x_1, (x_1.shape[-1] // 3), dim=(x_1.dim() - 1))