Hi,
I have a huggingface model trained with bfp16. I tried to load the model with bfp16 and export it using torch.onnx.export, but got the following error RuntimeError: unexpected tensor scalar type. My code/detailed error is below. Is this because of the bfp16 exporting is not supported?
If I change the model loading code to pegasus_model = PegasusForConditionalGeneration.from_pretrained(model_path), the export below can work. Is there a way to export a model with bfp16 weights?
I was using torch version 1.10.1+cu113 on an a10g gpu.
model_path = "xx"
pegasus_model = PegasusForConditionalGeneration.from_pretrained(model_path, torch_dtype=torch.bfloat16)
tokenizer = PegasusTokenizer.from_pretrained(model_path)
def export_encoder(model, args, exported_model_path):
model.eval()
with torch.no_grad():
_ = torch.onnx._export(model,
args,
exported_model_path,
export_params=True,
opset_version=12,
input_names=['input_ids'],
output_names=['hidden_states'],
dynamic_axes={
'input_ids': {0:'batch', 1: 'sequence'},
'hidden_states': {0:'batch', 1: 'sequence'},
})
export_text = "This is a great one"
export_input = tokenizer(export_text, return_tensors='pt')
export_encoder(model.model.encoder, export_input['input_ids'], output_encoder_path)
Detailed Error:
~/pegasus_exp/lib/python3.9/site-packages/torch/onnx/init.py in _export(*args, **kwargs)
26 def _export(*args, **kwargs):
27 from torch.onnx import utils
—> 28 result = utils._export(*args, **kwargs)
29 return result
30
~/pegasus_exp/lib/python3.9/site-packages/torch/onnx/utils.py in _export(model, args, f, export_params, verbose, training, input_names, output_names, operator_export_type, export_type, example_outputs, opset_version, do_constant_folding, dynamic_axes, keep_initializers_as_inputs, fixed_batch_size, custom_opsets, add_node_names, use_external_data_format, onnx_shape_inference)
722
723 graph, params_dict, torch_out =
→ 724 _model_to_graph(model, args, verbose, input_names,
725 output_names, operator_export_type,
726 example_outputs, val_do_constant_folding,
~/pegasus_exp/lib/python3.9/site-packages/torch/onnx/utils.py in _model_to_graph(model, args, verbose, input_names, output_names, operator_export_type, example_outputs, do_constant_folding, _disable_torch_constant_prop, fixed_batch_size, training, dynamic_axes)
495 params_dict = _get_named_param_dict(graph, params)
496
→ 497 graph = _optimize_graph(graph, operator_export_type,
498 _disable_torch_constant_prop=_disable_torch_constant_prop,
499 fixed_batch_size=fixed_batch_size, params_dict=params_dict,
~/pegasus_exp/lib/python3.9/site-packages/torch/onnx/utils.py in _optimize_graph(graph, operator_export_type, _disable_torch_constant_prop, fixed_batch_size, params_dict, dynamic_axes, input_names, module)
214 dynamic_axes = {} if dynamic_axes is None else dynamic_axes
215 torch._C._jit_pass_onnx_set_dynamic_input_shape(graph, dynamic_axes, input_names)
→ 216 graph = torch._C._jit_pass_onnx(graph, operator_export_type)
217 torch._C._jit_pass_lint(graph)
218
~/pegasus_exp/lib/python3.9/site-packages/torch/onnx/init.py in _run_symbolic_function(*args, **kwargs)
371 def _run_symbolic_function(*args, **kwargs):
372 from torch.onnx import utils
→ 373 return utils._run_symbolic_function(*args, **kwargs)
374
375
~/pegasus_exp/lib/python3.9/site-packages/torch/onnx/utils.py in _run_symbolic_function(g, block, n, inputs, env, operator_export_type)
1030 return None
1031 attrs = {k: n[k] for k in n.attributeNames()}
→ 1032 return symbolic_fn(g, *inputs, **attrs)
1033
1034 elif ns == “prim”:
~/pegasus_exp/lib/python3.9/site-packages/torch/onnx/symbolic_helper.py in wrapper(g, *args, **kwargs)
170 if len(kwargs) == 1:
171 assert “_outputs” in kwargs
→ 172 return fn(g, *args, **kwargs)
173
174 return wrapper
~/pegasus_exp/lib/python3.9/site-packages/torch/onnx/symbolic_opset9.py in embedding(g, weight, indices, padding_idx, scale_grad_by_freq, sparse)
494 “ONNX does not support not updating the embedding vector at padding_idx during training.”)
495
→ 496 return g.op(“Gather”, weight, indices)
497
498
~/pegasus_exp/lib/python3.9/site-packages/torch/onnx/utils.py in _graph_op(g, opname, *raw_args, **kwargs)
926 if _onnx_shape_inference:
927 from torch.onnx.symbolic_helper import _export_onnx_opset_version as opset_version
→ 928 torch._C._jit_pass_onnx_node_shape_type_inference(n, _params_dict, opset_version)
929
930 if outputs == 1:
RuntimeError: unexpected tensor scalar type