Issues tracing GPT-2 models for mobile deployment

cphifer: <p>Hey folks,</p>
<p><strong>Problem Description</strong></p>
<p>I am attempting to prepare models based on GPT-2 for deployment to a mobile application using the PyTorch JIT tracing capability, mobile optimization, and lite-interpreter saving (as recommended in various PyTorch formulas / examples). We make use of <code>past_key_values</code>, so have adjusted those examples to account for that input to the model.</p>
<p>The tracing / export process appears to work as expected, but calls to <code>forward</code> on the resulting model result in the following error on mobile (and similar on desktop):</p>
<pre><code class="lang-auto">com.facebook.jni.CppException: XNNPACK Linear not available! Reason: The provided (weight, bias, output_min, output_max) parameters are either invalid individually or their combination is not supported by XNNPACK.
</code></pre>
<p><strong>Steps to reproduce</strong></p>
<p>This was run on an x86-64 system using PyTorch 2.3.0. I have also reproduced the same error in our mobile environment (using the Android API version 2.1.0, the latest available).  Running the following script should reproduce the error, specifically when calling <code>loaded_model</code> to advance:</p>
<pre data-code-wrap="python"><code class="lang-python">import torch

from transformers import AutoModelForCausalLM

from torch.utils.mobile_optimizer import optimize_for_mobile

from concurrent import futures


def trace_and_save(model, example_input, example_past, filename):
    class ModelWrapper(torch.nn.Module):
        def __init__(self, model):
            super().__init__()
            self.model = model

        def forward(self, x, past):
            return self.model(x, return_dict=False, past_key_values=past)

    wrapped_model = ModelWrapper(model)

    print("Tracing...")
    traced_model = torch.jit.trace(
        wrapped_model,
        (example_input, example_past),
    )

    print("Optimizing...")
    optimized_model = optimize_for_mobile(traced_model)

    print("Saving...")
    optimized_model._save_for_lite_interpreter(filename)


def main():
    with torch.no_grad():
        torch.set_default_dtype(torch.float64)

        model = AutoModelForCausalLM.from_pretrained(
            "datificate/gpt2-small-spanish", torchscript=True
        )

        model.eval()

        N, C = 1, 10
        dummy_input = torch.zeros((N, C), dtype=torch.long)

        def make_tensor_pair():
            return (
                torch.zeros(1, 12, 0, 64, dtype=torch.float),
                torch.zeros(1, 12, 0, 64, dtype=torch.float),
            )

        dummy_past = tuple(make_tensor_pair() for _ in range(12))

        trace_and_save(model, dummy_input, dummy_past, "gpt2.ptl")

        loaded_model = torch.jit.load("gpt2.ptl")
        loaded_model.eval()

        all_encoder_layers, _ = loaded_model(dummy_input, dummy_past)
        print(all_encoder_layers)


if __name__ == "__main__":
    main() 
</code></pre>
<p>Any ideas or advice would be very welcome, thanks!</p>

cphifer: <p>For anyone that comes across this thread: We figured out that it was the combination of <code>optimize_for_mobile</code> and <code>set_default_dtype(torch.float64)</code> that was causing our problems.</p>
<p>We introduced the latter due to warnings emitted during the trace regarding a few ‘mismatched’ elements (i.e. elements that were out of the PyTorch default tolerance range). This is acceptable for our application, so we simply reverted to the default 32-bit float tensor type.</p>
<p>Alternatively, we found that skipping the <code>optimize_for_mobile</code> step worked if that precision <em>is</em> important – though there was a significant drop in performance in the mobile setting, as you may expect.</p>