Fail to get onnx with torch.onnx.export

vfdff: <p>I’m try to export a onnx model base on <strong>torch.onnx.export</strong> for llama v2, using the script testLlame_om.py can be executed, but after it finish, there is no llama_opt.onnx generated without error message.</p>
<pre><code class="lang-auto">from transformers import LlamaForCausalLM, AutoTokenizer
import torch
import torch_npu
import time

torch.set_printoptions(profile="full")
# torch.npu.set_compile_mode(jit_compile=True)
device="npu:1"

#下载好的hf模型地址
hf_model_path = 'models--daryl149--llama-2-7b-hf'
Model = LlamaForCausalLM.from_pretrained(hf_model_path, device_map=device)
tokenizer = AutoTokenizer.from_pretrained(hf_model_path)

# print(Model)

class SumModule(torch.nn.Module):
   def forward(self, x):
       time_start = time.time()
       generate_ids = Model.generate(x, max_length=512)
       time_end = time.time()

def export():
    models = SumModule()
    input_ids = torch.tensor([1, 3923, 596, 5655, 6975, 30, 3639, 596, 5655, 6975, 30, 3639, 596, 5655, 6975, 30, 3639, 596, 5655, 6975, 30, 3639, 596, 5655, 6975, 30, 3639, 596, 5655, 6975, 30, 3639, 596, 5655, 6975, 30, 3639, 596, 5655, 6975, 30, 3639, 596, 5655, 6975, 30, 3639, 596, 5655, 6975, 30, 3639, 596, 5655, 6975, 30, 3639, 596, 5655, 6975, 30, 3639, 596, 5655, 6975, 30, 3639, 596, 5655, 6975, 30, 3639, 596, 5655, 6975, 30, 3639, 596, 5655, 6975, 30, 3639, 596, 5655, 6975, 30, 3639, 596, 5655, 6975, 30, 3639, 596, 5655, 6975, 30, 3639, 596, 5655, 6975, 30, 3639, 596, 5655, 6975, 30, 3639, 596, 5655, 6975, 30, 3639, 596, 5655, 6975, 30, 3639, 596, 5655, 6975, 30, 3639, 596, 5655, 6975, 30, 3639, 596, 5655, 6975, 30, 3639, 596, 5655, 6975, 30, 3639, 596, 5655, 6975, 30, 3639, 596, 5655, 6975, 30, 3639, 596, 5655, 6975, 30, 3639, 596, 5655, 6975, 30, 3639, 596, 5655, 6975, 30, 3639, 596, 5655, 6975, 30, 3639, 596, 5655, 6975, 30, 3639, 596, 5655, 6975, 30, 3639, 596, 5655, 6975, 30, 3639, 596, 5655, 6975, 30, 3639, 596, 5655, 6975, 30, 3639, 596, 5655, 6975, 30, 3639, 596, 5655, 6975, 30, 3639, 596, 5655, 6975, 30, 3639, 596, 5655, 6975, 30, 3639, 596, 5655, 6975, 30, 3639, 596, 5655, 6975, 30, 3639, 596, 5655, 6975, 30, 3639, 596, 5655, 6975, 30, 3639, 596, 5655, 6975, 30, 3639, 596, 5655, 6975, 30, 3639, 596, 5655, 6975, 30, 3639, 596, 5655, 6975, 30, 3639, 596, 5655, 6975, 30])
    input_ids_npu = input_ids.to(device).unsqueeze(0)


    models.eval()  # set infer mode
    torch.onnx.export(
        models,
        (input_ids_npu),
        "llama_opt.onnx",
        opset_version=11,
        export_params=False,
        verbose=False,
        input_names=["x"],
        output_names=["z"],
    )


if __name__ == "__main__":
    export()
</code></pre>

ptrblck: <p>I could not reproduce any issues on CUDA. Do you see this issue only when <code>npu</code> is used?</p>

vfdff: <p>can you get the file <strong>llama_opt.onnx</strong> base on your CUDA after it finish running?</p>
<p>as I don’t have a cuda machine, I don’t sure the issue only exist on <strong>npu</strong>.</p>
<p>PS：<br>
<strong>optimum-cli export onnx --model models–daryl149–llama-2-7b-hf onnx_model2 --task text-generation</strong> can successful generate a onnx file, so the source code of  models–daryl149–llama-2-7b-hf are complete.</p>
<blockquote>
<p>refer to <a href="https://stackoverflow.com/questions/75962909/convert-llama-to-onnx" class="inline-onebox" rel="noopener nofollow ugc">python - Convert LLaMA to ONNX - Stack Overflow</a></p>
</blockquote>