saymrwulf/onnxruntime
1
0
Fork 0
mirror of https://github.com/saymrwulf/onnxruntime.git synced 2026-07-11 17:48:34 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions

Fix Mixtral Parity test to keep it consistent with Transformers. (#20210)

Browse source 
### 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.
This commit is contained in:
Adam Louly 2024-04-08 13:04:12 -07:00 committed by GitHub
parent 908a76d675
commit 22a61a3cf5
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
1 changed files with 2 additions and 6 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

8
onnxruntime/test/python/transformers/test_parity_mixtral_moe.py
View file

@ -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.