[poc][v1.x] ONNX Add an option to un-interleave BERT

[Zha0q1]

Add model (BERT) specific logic to un-interleave the self attention mat mul. This can potentially speed up inference with trt 8.0 whose compiler can recognize the new pattern.

default usage model_specific_logics='gluonnlp_bert'

converted_model_path = mx.onnx.export_model(sym_file, params_file, input_shapes,
                                            input_types, onnx_file,
                                            model_specific_logics='gluonnlp_bert_uninterleaved')

When the model is not bert base (meaning hidden != 768 or num_heads != 12), e.g. bert large
the usage is:

cheat_sheet = {
    'qkv_hidden': 1024,
    'num_heads': 16,
    'head_dim': 64
}
converted_model_path = mx.onnx.export_model(sym_file, params_file, input_shapes,
                                            input_types, onnx_file,
                                            model_specific_logics='gluonnlp_bert_uninterleaved',
                                            cheat_sheet=cheat_sheet)

This option to un-interleave self-attention would also work with bert-variants such as roberta, distilbert, and 'ernie'

The first screenshot is the old graph, the second is the new graph. Note that use of onnx-sim is required

@TristonC @MoisesHer @waytrue17 @josephevans

[mxnet-bot]

Hey @Zha0q1 , Thanks for submitting the PR
All tests are already queued to run once. If tests fail, you can trigger one or more tests again with the following commands:

• To trigger all jobs: @mxnet-bot run ci [all]
• To trigger specific jobs: @mxnet-bot run ci [job1, job2]

---

CI supported jobs: [clang, miscellaneous, unix-gpu, edge, centos-cpu, unix-cpu, centos-gpu, windows-cpu, website, windows-gpu, sanity]

---

Note:
Only following 3 categories can trigger CI :PR Author, MXNet Committer, Jenkins Admin.
All CI tests must pass before the PR can be merged.

[waytrue17]

Would this logic benefits specifically for TRT? If so, shall we consider to name it like '_gluonnlp_bert_trt'?

[szha]

It probably makes more sense to have gluonnlp export a BERT graph that doesn't involve interleaving than doing it in mxnet. framework shouldn't have knowledge about or rely on the implementation of its ecosystem packages.

[szha]

these options are semantically unclear and hard to maintain

[Zha0q1]

It probably makes more sense to have gluonnlp export a BERT graph that doesn't involve interleaving than doing it in mxnet. framework shouldn't have knowledge about or rely on the implementation of its ecosystem packages.

I agree. Alternatively we might be able to make mx2onnx support loading custom conversion functions dynamically and put these functions to gluonnlp? Or if even that is too hacky this pr can serve as a poc to evaluate the performance benefit with trt 8.0 by un-interleaving the matrix multiplication

[Zha0q1]

@szha @ptrendx for insights

[ptrendx]

@Zha0q1 What is the difficulty in doing this transformation after the export? You should be able to modify the ONNX graph itself, right?