Description
When exporting the gemma model with huggingface transformers,
import torch
from transformers import AutoTokenizer, AutoModelForCausalLM
tokenizer = AutoTokenizer.from_pretrained("google/gemma-7b")
model = AutoModelForCausalLM.from_pretrained("google/gemma-7b")
print("Model loaded")
kwargs = tokenizer("What is your name?")
input_ids = torch.tensor([kwargs["input_ids"]])
attention_mask = torch.tensor([kwargs["attention_mask"]])
program = torch.export.export(model, (input_ids, attention_mask))
fails with
Traceback (most recent call last):
File "/home/justinchu/dev/torch-onnx/gemma.py", line 31, in <module>
program = torch.export.export(model, (input_ids, attention_mask))
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
File "/home/justinchu/anaconda3/envs/onnx/lib/python3.11/site-packages/torch/export/__init__.py", line 174, in export
return _export(
^^^^^^^^
File "/home/justinchu/anaconda3/envs/onnx/lib/python3.11/site-packages/torch/export/_trace.py", line 945, in wrapper
raise e
File "/home/justinchu/anaconda3/envs/onnx/lib/python3.11/site-packages/torch/export/_trace.py", line 928, in wrapper
ep = fn(*args, **kwargs)
^^^^^^^^^^^^^^^^^^^
File "/home/justinchu/anaconda3/envs/onnx/lib/python3.11/site-packages/torch/export/exported_program.py", line 89, in wrapper
return fn(*args, **kwargs)
^^^^^^^^^^^^^^^^^^^
File "/home/justinchu/anaconda3/envs/onnx/lib/python3.11/site-packages/torch/export/_trace.py", line 1533, in _export
exported_program = ExportedProgram(
^^^^^^^^^^^^^^^^
File "/home/justinchu/anaconda3/envs/onnx/lib/python3.11/site-packages/torch/export/exported_program.py", line 246, in __init__
self.verifier().check(self)
File "/home/justinchu/anaconda3/envs/onnx/lib/python3.11/site-packages/torch/_export/verifier.py", line 154, in check
self._check_graph_module(ep.graph_module)
File "/home/justinchu/anaconda3/envs/onnx/lib/python3.11/site-packages/torch/_export/verifier.py", line 220, in _check_graph_module
_check_val(node)
File "/home/justinchu/anaconda3/envs/onnx/lib/python3.11/site-packages/torch/_export/verifier.py", line 62, in _check_val
raise SpecViolationError(f"Node.meta {node.name} is missing val field.")
torch._export.verifier.SpecViolationError: Node.meta _enter_autocast is missing val field.
Environment
PyTorch version: 2.4.0+cu121
Is debug build: False
CUDA used to build PyTorch: 12.1
ROCM used to build PyTorch: N/A
OS: Microsoft Azure Linux 3.0 (x86_64)
GCC version: (GCC) 13.2.0
Clang version: Could not collect
CMake version: Could not collect
Libc version: glibc-2.38
Python version: 3.11.9 (main, Apr 19 2024, 16:48:06) [GCC 11.2.0] (64-bit runtime)
Python platform: Linux-6.6.22.1-2.azl3-x86_64-with-glibc2.38
Is CUDA available: False
CUDA runtime version: No CUDA
CUDA_MODULE_LOADING set to: N/A
GPU models and configuration: No CUDA
Nvidia driver version: No CUDA
cuDNN version: No CUDA
HIP runtime version: N/A
MIOpen runtime version: N/A
Is XNNPACK available: True
CPU:
Architecture: x86_64
CPU op-mode(s): 32-bit, 64-bit
Address sizes: 48 bits physical, 48 bits virtual
Byte Order: Little Endian
CPU(s): 64
On-line CPU(s) list: 0-63
Vendor ID: AuthenticAMD
Model name: AMD EPYC 7763 64-Core Processor
CPU family: 25
Model: 1
Thread(s) per core: 2
Core(s) per socket: 32
Socket(s): 1
Stepping: 1
BogoMIPS: 4890.86
Flags: fpu vme de pse tsc msr pae mce cx8 apic sep mtrr pge mca cmov pat pse36 clflush mmx fxsr sse sse2 ht syscall nx mmxext fxsr_opt pdpe1gb rdtscp lm rep_good nopl cpuid extd_apicid pni pclmulqdq ssse3 fma cx16 sse4_1 sse4_2 movbe popcnt aes xsave avx f16c rdrand hypervisor lahf_lm cmp_legacy cr8_legacy abm sse4a misalignsse 3dnowprefetch osvw topoext vmmcall fsgsbase bmi1 avx2 smep bmi2 erms rdseed adx smap clflushopt clwb sha_ni xsaveopt xsavec xgetbv1 xsaves xsaveerptr arat umip vaes vpclmulqdq rdpid fsrm
Hypervisor vendor: Microsoft
Virtualization type: full
L1d cache: 1 MiB (32 instances)
L1i cache: 1 MiB (32 instances)
L2 cache: 16 MiB (32 instances)
L3 cache: 32 MiB (1 instance)
NUMA node(s): 1
NUMA node0 CPU(s): 0-63
Vulnerability Gather data sampling: Not affected
Vulnerability Itlb multihit: Not affected
Vulnerability L1tf: Not affected
Vulnerability Mds: Not affected
Vulnerability Meltdown: Not affected
Vulnerability Mmio stale data: Not affected
Vulnerability Reg file data sampling: Not affected
Vulnerability Retbleed: Not affected
Vulnerability Spec rstack overflow: Vulnerable: Safe RET, no microcode
Vulnerability Spec store bypass: Vulnerable
Vulnerability Spectre v1: Mitigation; usercopy/swapgs barriers and __user pointer sanitization
Vulnerability Spectre v2: Mitigation; Retpolines, STIBP disabled, RSB filling, PBRSB-eIBRS Not affected
Vulnerability Srbds: Not affected
Vulnerability Tsx async abort: Not affected
Versions of relevant libraries:
[pip3] model-explorer-onnx==0.2.7
[pip3] mypy==1.10.1
[pip3] mypy-extensions==1.0.0
[pip3] numpy==1.26.4
[pip3] onnx-weekly==1.17.0.dev20240722
[pip3] onnxruntime==1.18.0
[pip3] onnxruntime-training==1.18.0
[pip3] onnxscript==0.1.0.dev20240703
[pip3] torch==2.4.0
[pip3] torch-onnx==0.0.34
[pip3] torchvision==0.19.0
[pip3] triton==3.0.0
[conda] numpy 1.26.4 pypi_0 pypi
[conda] torch 2.4.0 pypi_0 pypi
[conda] torch-onnx 0.0.34 pypi_0 pypi
[conda] torchvision 0.19.0 pypi_0 pypi
[conda] triton 3.0.0 pypi_0 pypi
cc @avikchaudhuri @gmagogsfm @zhxchen17 @tugsbayasgalan @angelayi @suo @ydwu4 @penguinwu