Cannot export model in bfp16 to ONNX

username: sc21
cooked: <p>Hi,</p>
<p>I have a huggingface model trained with bfp16. I tried to load the model with bfp16 and export it using <code>torch.onnx.export</code>, but got the following error <strong>RuntimeError: unexpected tensor scalar type</strong>. My code/detailed error is below. Is this because of the bfp16 exporting is not supported?</p>
<p>If I change the model loading code to <code>pegasus_model = PegasusForConditionalGeneration.from_pretrained(model_path)</code>, the export below can work. Is there a way to export a model with bfp16 weights?</p>
<p>I was using torch version 1.10.1+cu113 on an a10g gpu.</p>
<pre><code class="lang-auto">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)
</code></pre>
<p>Detailed Error:</p>
<blockquote>
<p>~/pegasus_exp/lib/python3.9/site-packages/torch/onnx/<strong>init</strong>.py in _export(*args, **kwargs)<br>
26 def _export(*args, **kwargs):<br>
27     from torch.onnx import utils<br>
—&gt; 28     result = utils._export(*args, **kwargs)<br>
29     return result<br>
30<br>
~/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)<br>
722<br>
723             graph, params_dict, torch_out = <br>
 → 724                 _model_to_graph(model, args, verbose, input_names,<br>
725                                 output_names, operator_export_type,<br>
726                                 example_outputs, val_do_constant_folding,<br>
~/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)<br>
495     params_dict = _get_named_param_dict(graph, params)<br>
496<br>
 → 497     graph = _optimize_graph(graph, operator_export_type,<br>
498                             _disable_torch_constant_prop=_disable_torch_constant_prop,<br>
499                             fixed_batch_size=fixed_batch_size, params_dict=params_dict,<br>
~/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)<br>
214         dynamic_axes = {} if dynamic_axes is None else dynamic_axes<br>
215         torch._C._jit_pass_onnx_set_dynamic_input_shape(graph, dynamic_axes, input_names)<br>
 → 216     graph = torch._C._jit_pass_onnx(graph, operator_export_type)<br>
217     torch._C._jit_pass_lint(graph)<br>
218<br>
~/pegasus_exp/lib/python3.9/site-packages/torch/onnx/<strong>init</strong>.py in _run_symbolic_function(*args, **kwargs)<br>
371 def _run_symbolic_function(*args, **kwargs):<br>
372     from torch.onnx import utils<br>
 → 373     return utils._run_symbolic_function(*args, **kwargs)<br>
374<br>
375<br>
~/pegasus_exp/lib/python3.9/site-packages/torch/onnx/utils.py in _run_symbolic_function(g, block, n, inputs, env, operator_export_type)<br>
1030                     return None<br>
1031                 attrs = {k: n[k] for k in n.attributeNames()}<br>
 → 1032                 return symbolic_fn(g, *inputs, **attrs)<br>
1033<br>
1034         elif ns == “prim”:<br>
~/pegasus_exp/lib/python3.9/site-packages/torch/onnx/symbolic_helper.py in wrapper(g, *args, **kwargs)<br>
170             if len(kwargs) == 1:<br>
171                 assert “_outputs” in kwargs<br>
 → 172             return fn(g, *args, **kwargs)<br>
173<br>
174         return wrapper<br>
~/pegasus_exp/lib/python3.9/site-packages/torch/onnx/symbolic_opset9.py in embedding(g, weight, indices, padding_idx, scale_grad_by_freq, sparse)<br>
494                       “ONNX does not support not updating the embedding vector at padding_idx during training.”)<br>
495<br>
 → 496     return g.op(“Gather”, weight, indices)<br>
497<br>
498<br>
~/pegasus_exp/lib/python3.9/site-packages/torch/onnx/utils.py in _graph_op(g, opname, *raw_args, **kwargs)<br>
926     if _onnx_shape_inference:<br>
927         from torch.onnx.symbolic_helper import _export_onnx_opset_version as opset_version<br>
 → 928         torch._C._jit_pass_onnx_node_shape_type_inference(n, _params_dict, opset_version)<br>
929<br>
930     if outputs == 1:<br>
RuntimeError: unexpected tensor scalar type</p>
</blockquote>