From 22a61a3cf513c4beffb59dbc55a4cbe6ac278835 Mon Sep 17 00:00:00 2001 From: Adam Louly Date: Mon, 8 Apr 2024 13:04:12 -0700 Subject: [PATCH] Fix Mixtral Parity test to keep it consistent with Transformers. (#20210) ### Description I recently opened a PR in hf transformers repo to fix an issue on the indexing part. https://github.com/huggingface/transformers/issues/29857 onnx exporter was failing because of the tolist() conversion so we had to remove it. I found out that the code was also a part of our codebase so this PR is to keep the code consistent. --- .../test/python/transformers/test_parity_mixtral_moe.py | 8 ++------ 1 file changed, 2 insertions(+), 6 deletions(-) diff --git a/onnxruntime/test/python/transformers/test_parity_mixtral_moe.py b/onnxruntime/test/python/transformers/test_parity_mixtral_moe.py index 50292f186d..0070462602 100644 --- a/onnxruntime/test/python/transformers/test_parity_mixtral_moe.py +++ b/onnxruntime/test/python/transformers/test_parity_mixtral_moe.py @@ -293,15 +293,11 @@ class MixtralSparseMoeBlock(nn.Module): if top_x.shape[0] == 0: continue - # in torch it is faster to index using lists than torch tensors - top_x_list = top_x.tolist() - idx_list = idx.tolist() - # Index the correct hidden states and compute the expert hidden state for # the current expert. We need to make sure to multiply the output hidden # states by `routing_weights` on the corresponding tokens (top-1 and top-2) - current_state = hidden_states[None, top_x_list].reshape(-1, hidden_dim) - current_hidden_states = expert_layer(current_state) * routing_weights[top_x_list, idx_list, None] + current_state = hidden_states[None, top_x].reshape(-1, hidden_dim) + current_hidden_states = expert_layer(current_state) * routing_weights[top_x, idx, None] # However `index_add_` only support torch tensors for indexing so we'll use # the `top_x` tensor here.