导出 llama3.2-11b-vision 出错

[L-jasmine]

错误日志

(base) tom@peaceful-dawn-1385-86b5647d8d-5965n:~/workspace/LLM-TPU/models/Llama3_2-Vision/compile$ python export_onnx.py -m meta-llama/Llama-3.2-11B-Vision -s 512
Loading checkpoint shards: 100%|████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 5/5 [00:15<00:00,  3.07s/it]

LLM config:
    Layers: 40
    Hidden size: 4096
    Query heads: 32
    KV heads: 8
    CrossAttn layers: [3, 8, 13, 18, 23, 28, 33, 38]

VIT config:
    Image size: 448
    Max tiles: 4
    Patch size: 14
    Num Patches: 1025

Convert vision transformer
Traceback (most recent call last):
  File "/home/tom/fssd/LLM-TPU/models/Llama3_2-Vision/compile/export_onnx.py", line 487, in <module>
    convert_vision_transformer()
  File "/home/tom/fssd/LLM-TPU/models/Llama3_2-Vision/compile/export_onnx.py", line 242, in convert_vision_transformer
    torch.onnx.export(
  File "/opt/conda/lib/python3.11/site-packages/torch/onnx/__init__.py", line 375, in export
    export(
  File "/opt/conda/lib/python3.11/site-packages/torch/onnx/utils.py", line 502, in export
    _export(
  File "/opt/conda/lib/python3.11/site-packages/torch/onnx/utils.py", line 1564, in _export
    graph, params_dict, torch_out = _model_to_graph(
                                    ^^^^^^^^^^^^^^^^
  File "/opt/conda/lib/python3.11/site-packages/torch/onnx/utils.py", line 1113, in _model_to_graph
    graph, params, torch_out, module = _create_jit_graph(model, args)
                                       ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/opt/conda/lib/python3.11/site-packages/torch/onnx/utils.py", line 997, in _create_jit_graph
    graph, torch_out = _trace_and_get_graph_from_model(model, args)
                       ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/opt/conda/lib/python3.11/site-packages/torch/onnx/utils.py", line 904, in _trace_and_get_graph_from_model
    trace_graph, torch_out, inputs_states = torch.jit._get_trace_graph(
                                            ^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/opt/conda/lib/python3.11/site-packages/torch/jit/_trace.py", line 1500, in _get_trace_graph
    outs = ONNXTracedModule(
           ^^^^^^^^^^^^^^^^^
  File "/opt/conda/lib/python3.11/site-packages/torch/nn/modules/module.py", line 1736, in _wrapped_call_impl
    return self._call_impl(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/opt/conda/lib/python3.11/site-packages/torch/nn/modules/module.py", line 1747, in _call_impl
    return forward_call(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/opt/conda/lib/python3.11/site-packages/torch/jit/_trace.py", line 139, in forward
    graph, out = torch._C._create_graph_by_tracing(
                 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/opt/conda/lib/python3.11/site-packages/torch/jit/_trace.py", line 130, in wrapper
    outs.append(self.inner(*trace_inputs))
                ^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/opt/conda/lib/python3.11/site-packages/torch/nn/modules/module.py", line 1736, in _wrapped_call_impl
    return self._call_impl(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/opt/conda/lib/python3.11/site-packages/torch/nn/modules/module.py", line 1747, in _call_impl
    return forward_call(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/opt/conda/lib/python3.11/site-packages/torch/nn/modules/module.py", line 1726, in _slow_forward
    result = self.forward(*input, **kwargs)
             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/tom/fssd/LLM-TPU/models/Llama3_2-Vision/compile/export_onnx.py", line 92, in forward
    out = vision(pixel_values, aspect_ratio_ids, aspect_ratio_mask)[0]
          ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/opt/conda/lib/python3.11/site-packages/torch/nn/modules/module.py", line 1736, in _wrapped_call_impl
    return self._call_impl(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/opt/conda/lib/python3.11/site-packages/torch/nn/modules/module.py", line 1747, in _call_impl
    return forward_call(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/opt/conda/lib/python3.11/site-packages/torch/nn/modules/module.py", line 1726, in _slow_forward
    result = self.forward(*input, **kwargs)
             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/opt/conda/lib/python3.11/site-packages/transformers/models/mllama/modeling_mllama.py", line 1505, in forward
    hidden_state = patch_embeds.reshape(4,1280,1600).transpose(1,2)
                   ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
RuntimeError: shape '[4, 1280, 1600]' is invalid for input of size 5242880

已经替换了 files/Llama-3.2-11B-Vision-Instruct/modeling_mllama.py

[sgwzy22]

检查一下当前环境的torch版本，要求2.2.0以上，已添加至requirements.txt

[L-jasmine]

找到问题了，Llama-3.2-11B-Vision-Instruct 可以 但是 Llama-3.2-11B-Vision 不行

[sgwzy22]

👌