From c593f97b9f4f03f10a0293f35c91679b233e46e4 Mon Sep 17 00:00:00 2001 From: Liu Jianfeng Date: Tue, 19 Jul 2022 09:34:59 +0000 Subject: [PATCH] fix softmax dim of Residual MoE in moe/layer.py --- deepspeed/moe/layer.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/deepspeed/moe/layer.py b/deepspeed/moe/layer.py index 7dcf4144c0e6..399f512bfead 100644 --- a/deepspeed/moe/layer.py +++ b/deepspeed/moe/layer.py @@ -126,6 +126,6 @@ def forward(self, hidden_states, used_token=None): if type(output_mlp) is tuple: output_mlp = output_mlp[0] # Ignore the bias term for now coef = self.coefficient(hidden_states) - coef = torch.nn.functional.softmax(coef, dim=1) + coef = torch.nn.functional.softmax(coef, dim=-1) output = output * coef[..., 0:1] + output_mlp * coef[..., 1:] return output, self.deepspeed_moe.l_aux, self.deepspeed_moe.exp_counts