ONNX of switch transformer fails in a weird place

Julian_Buchel: <p>Hi,</p>
<p>I am trying to convert a SwitchTransformer to an ONNX graph using this code</p>
<pre><code class="lang-auto">from transformers import SwitchTransformersEncoderModel, SwitchTransformersConfig, AutoTokenizer
import torch

if __name__ == "__main__":
    device = torch.device("cuda") if torch.cuda.is_available() else torch.device("cpu")
    vocab_size = 1024
    batch_size = 10
    switch_transformer_config = SwitchTransformersConfig(
        vocab_size=vocab_size,
        d_model=256,
        d_kv=64,
        d_ff=512,
        num_layers=2,
        num_sparse_encoder_layers=1,
        num_experts=3,
    )
    # - Create a small switch transformer
    model = SwitchTransformersEncoderModel(switch_transformer_config).to(device).eval()
    random_inputs = (vocab_size * torch.rand(batch_size,vocab_size)).int().to(device)

    # - Check if the model runs on the input
    model(random_inputs)

    torch.onnx.export(
        model,
        random_inputs,
        "onnx-models/switch_transformer.onnx",
        verbose=True
    )
</code></pre>
<p>I get the following error</p>
<pre><code class="lang-auto">input_shape_value == reshape_value || input_shape_value == 1 || reshape_value == 1 INTERNAL ASSERT FAILED at "../torch/csrc/jit/passes/onnx/shape_type_inference.cpp":554, please report a bug to PyTorch. ONNX Expand input shape constraint not satisfied.
</code></pre>
<p>However, this error disappears when I comment out the dropout call in line 379 of <code>transformers/models/switch_transformers/modeling_switch_transformers.py</code> which says</p>
<pre><code class="lang-auto">output = hidden_states + self.dropout(forwarded_states)
</code></pre>
<p>I don’t understand why this is happening as the model successfully runs over the input.</p>
<p>Any help would be appreciated.</p>

Julian_Buchel: <p>Ok, actually it makes sense that it works when I comment out the dropout call as this basically makes the <code>forwarded_state</code> redundant.</p>
<pre><code class="lang-auto">def forward(self, hidden_states, output_router_logits):
    forwarded_states = self.layer_norm(hidden_states)
    forwarded_states = self.mlp(forwarded_states)

    if isinstance(forwarded_states, tuple):
        forwarded_states, router_tuple = forwarded_states
    else:
        router_tuple = None

    output = hidden_states + self.dropout(forwarded_states)

    if output_router_logits and router_tuple is not None:
        output = (output, router_tuple)

    return output
</code></pre>
<p>So the problem arises in the MLP layer, which is in this case the MoE layer.</p>
<pre><code class="lang-auto">def forward(self, hidden_states):
    r"""
    Hold on, this will be slightly tricky to understand In the correct order, a MoE layer does the following:

    1- Gets the `router_mask` from the router. The shape of the mask is `(batch_size, sequence_length, num_expert)`
    and corresponds to the argmax of the `router_probs`. The probabilities are needed in the computation of the
    hidden states : they are broadcasted to the hidden states values (can be interpreted as a scaling factor).

    2- Dispatch the tokens to its associated experts. We do a classic for loop over the experts and assign for each
    expert the corresponding hidden states.

    """
    # Step 1: Get the router_mask from the router as wel as the probabilities
    router_mask, router_probs, router_logits = self.router(hidden_states)
    expert_index = torch.argmax(router_mask, dim=-1)

    # The routers introduced might not always map all the tokens, to a router, which means that some hidden states
    # can be unchanged from one layer to another. That is why the hidden states are cloned before updating only the seleced ones.

    next_states = hidden_states.clone()
    for idx, expert in enumerate(self.experts.values()):

        token_indices = router_mask[:, :, idx].bool()
        next_states[token_indices] = expert(hidden_states[token_indices])

    hidden_states = router_probs * next_states
    return hidden_states, (router_logits, expert_index)
</code></pre>