[Optimizer] DQ + MatMul to MatMulNBits support #21180

fajin-corp · 2024-06-26T06:48:38Z

Description

MatMulNBits is a heavily optimized matmul operation. Currently a MatMul can be converted to MatMulNBits to speed up the model inference. However, MatMulNBits is an ORT only op. To make the graph compatible with ONNX ops and utilize MatMulNBits at the same time, we introduce Q/DQ support for MatMulNBits.

To convert MatMul ops in a model to MatMulNBits:

use matmul_4bits_quantizer.py to convert MatMul to DQ + MatMul using QDQ mode.
In ORT session, DQ + MatMul is fused to MatMulNBits

Note

MatMulNBits assume B weight is uint4. When no zp is provided, zp defaults to 8, which is different from DQ. DQ defaults zp to 0 when no zp provided. And DQ supports int4. Therefore some conversions are introduced during DQ + MatMul --> MatMulNBits step.

Perf

Using QDQ format will increase the model initialization time and memory consumption. With current implement, model init time increased from ~4s to ~9s, and memory consumption increased from ~2.8GB to ~4.8GB.
The memory increase is due to

in optimizer, after transpose the B weight, a in-memory tensor proto is created using protobuf's arena.
in finalize step, when saving initializer and prepacking, ORT arena is used to create buffers for initializers.

The memory allocated by arenas cannot be fully deallocated.
If disable ORT arena memory allocation, the memory consumptions of both QDQ format and original format are ~2.2GB.
The time increase is mainly due to multiple memory copy, but can be further optimized.

Motivation and Context

Please see description for details.

onnxruntime/python/tools/quantization/matmul_4bits_quantizer.py

onnxruntime/core/optimizer/qdq_transformer/selectors_actions/qdq_selector_action_transformer.cc

onnxruntime/core/optimizer/qdq_transformer/selectors_actions/qdq_actions.cc

onnxruntime/test/optimizer/qdq_matmulnbits_transformer_test.cc

fajin-corp · 2024-06-27T21:02:21Z

/azp run Windows_CI / Windows-CUDA-12,Windows_SCA / Onnxruntime-SCA-training-CUDA

azure-pipelines · 2024-06-27T21:02:27Z

No pipelines are associated with this pull request.

onnxruntime/core/optimizer/qdq_transformer/selectors_actions/qdq_actions.cc

edgchen1 · 2024-07-02T00:46:27Z

onnxruntime/core/optimizer/qdq_transformer/selectors_actions/qdq_actions.cc

+  }
+
+  OrtThreadPoolParams to;
+  auto tp = concurrency::CreateThreadPool(&onnxruntime::Env::Default(), to,


not sure if we should create a new threadpool here. if we really need this to run with a threadpool, consider adjusting the APIs to pass in an existing threadpool.

fix in next iter

onnxruntime/core/optimizer/qdq_transformer/selectors_actions/qdq_selectors.cc

onnxruntime/core/optimizer/graph_transformer_utils.cc

onnxruntime/core/optimizer/qdq_transformer/selectors_actions/qdq_selectors.cc

onnxruntime/core/optimizer/qdq_transformer/selectors_actions/qdq_actions.cc

skottmckay · 2024-07-02T09:21:36Z

onnxruntime/core/optimizer/qdq_transformer/selectors_actions/qdq_actions.h

+// used together with DQMatMulNodeGroupSelector, which does the sanity check
+struct DQMatMulReplaceWithMatMulNBits : public Action {
+  explicit DQMatMulReplaceWithMatMulNBits(int64_t accuracy_level);
+  Status Run(Graph&, const NodesToOptimize& selected_nodes) const override;


ReplaceWithNew would definitely be preferred.

Not sure why we seem to have a few things using Action directly (SplitReplaceWithQuant and MatMulReplaceWithQLinear), but the intention of the setup was to use something like ReplaceWithNew whenever possible so we're not duplication logic (and unnecessarily increasing binary size) for the common functionality.

onnxruntime/core/optimizer/qdq_transformer/selectors_actions/qdq_selector_action_transformer.cc

onnxruntime/core/optimizer/qdq_transformer/selectors_actions/qdq_selectors.h

onnxruntime/core/optimizer/qdq_transformer/selectors_actions/qdq_selectors.cc

onnxruntime/core/optimizer/qdq_transformer/selectors_actions/qdq_actions.cc

guschmue · 2024-07-11T15:14:41Z

the CI error in ort-web is this test:
https://github.com/microsoft/onnxruntime/blob/main/js/web/test/e2e/browser-test-webgpu-external-data.js
super simple test that is checking if we load external data correctly.
On main the test is passing.

Fails with:
ERROR: 'worker sent an error! http://localhost:9876/base/node_modules/onnxruntime-web/dist/ort-wasm-simd-threaded.jsep.mjs:112: Uncaught TypeError: Sa is not a function'

Unclear how we get there.

### Description This is a partial change from [fajin/qdqmatmulnbitstoolchain](#21180). The original PR is blocked by Web CI failures. MatMulNBits is a heavily optimized matmul operation. Currently a MatMul can be converted to MatMulNBits to speed up the model inference. However, MatMulNBits is an ORT only op. To make the graph compatible with ONNX ops and utilize MatMulNBits at the same time, we introduce Q/DQ support for MatMulNBits. To convert MatMul ops in a model to MatMulNBits: 1. use matmul_4bits_quantizer.py to convert MatMul to DQ + MatMul using QDQ mode. 2. In ORT session, DQ + MatMul is fused to MatMulNBits #### Note MatMulNBits assume B weight is uint4. When no zp is provided, zp defaults to 8, which is different from DQ. DQ defaults zp to 0 when no zp provided. And DQ supports int4. Therefore some conversions are introduced during DQ + MatMul --> MatMulNBits step. #### Perf Using QDQ format will increase the model initialization time and memory consumption. With current implement, model init time increased from ~4s to ~9s, and memory consumption increased from ~2.8GB to ~4.8GB. The memory increase is due to 1. in optimizer, after transpose the B weight, a in-memory tensor proto is created using protobuf's arena. 2. in finalize step, when saving initializer and prepacking, ORT arena is used to create buffers for initializers. The memory allocated by arenas cannot be fully deallocated. If disable ORT arena memory allocation, the memory consumptions of both QDQ format and original format are ~2.2GB. The time increase is mainly due to multiple memory copy, but can be further optimized. ### Motivation and Context Please see description for details.

fajin-corp requested a review from a team as a code owner June 26, 2024 06:48

github-advanced-security bot found potential problems Jun 26, 2024

View reviewed changes

onnxruntime/python/tools/quantization/matmul_4bits_quantizer.py Fixed Show fixed Hide fixed

fajin-corp commented Jun 26, 2024

View reviewed changes

onnxruntime/core/optimizer/qdq_transformer/selectors_actions/qdq_selector_action_transformer.cc Outdated Show resolved Hide resolved

fajin-corp commented Jun 26, 2024

View reviewed changes

onnxruntime/core/optimizer/qdq_transformer/selectors_actions/qdq_actions.cc Outdated Show resolved Hide resolved

fajin-corp commented Jun 26, 2024

View reviewed changes

onnxruntime/test/optimizer/qdq_matmulnbits_transformer_test.cc Show resolved Hide resolved

edgchen1 reviewed Jul 2, 2024

View reviewed changes

skottmckay reviewed Jul 2, 2024

View reviewed changes

fajin-corp force-pushed the fajin/qdqmatmulnbitstoolchain branch 2 times, most recently from b26e0ea to d2a0008 Compare July 8, 2024 23:33

fajin-corp added 18 commits July 12, 2024 12:13

adding python interface

d7d08f7

restarting

1c687d1

finished default quantizer

7612401

only enable quant_format in default quantizer

c5b6175

added UT and fixed qtype in DQ

60c3cf8

added dq matmul selectors

77b7dd9

added selector checks, and init action

1685070

finished attribute insertion

82419a3

finished initializer transpose and append to replacement node

d5d9e61

change target name generation

b2548da

added selector and action to qdq selector transformer

3aa704c

fixed linting

2a10834

fixed building UT

b516487

fixed non-convert ut

3fe19cb

fixed action calling transpose

67542a5

finished changing quantize

95de135

finished modifying transpose

0ad7fe4

updated mlas kernel calling

f81d7bb

fajin-corp added 17 commits July 12, 2024 12:14

fixed mlas scale calc bug

8e8b314

passed UT

5729762

fixed python build

c4805c8

fixing ci

53635fa

fixing ci

425f61b

fixing minimal build

b188f55

fixing ci

1d34c27

change dq matmul tool chain interface for genAI

75bc4c5

pass accuracy from session.config

f8773d9

added UT for accuracy level in session config options

d5b032e

fixed missing transformer path

2d01a38

resolved comments

c17a9cf

fix build

8c2a121

fix ut

baa9389

fixing arm ut

f33dbf9

fix linting

18e0000

corrected UT semantics

1ecf5c5

fajin-corp force-pushed the fajin/qdqmatmulnbitstoolchain branch from 6596e54 to 1ecf5c5 Compare July 12, 2024 19:15

try to fix web ci failure

f95c3d6

fajin-corp mentioned this pull request Jul 16, 2024

add transform part of the dq matmul tool chain #21374

Merged

fajin-corp closed this Jul 23, 2024

fajin-corp deleted the fajin/qdqmatmulnbitstoolchain branch July 23, 2024 20:42

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

[Optimizer] DQ + MatMul to MatMulNBits support #21180

[Optimizer] DQ + MatMul to MatMulNBits support #21180

fajin-corp commented Jun 26, 2024 •

edited

Loading

fajin-corp commented Jun 27, 2024

azure-pipelines bot commented Jun 27, 2024

edgchen1 Jul 2, 2024

fajin-corp Jul 6, 2024

skottmckay Jul 2, 2024

guschmue commented Jul 11, 2024 •

edited

Loading

[Optimizer] DQ + MatMul to MatMulNBits support #21180

[Optimizer] DQ + MatMul to MatMulNBits support #21180

Conversation

fajin-corp commented Jun 26, 2024 • edited Loading

Description

Note

Perf

Motivation and Context

fajin-corp commented Jun 27, 2024

azure-pipelines bot commented Jun 27, 2024

edgchen1 Jul 2, 2024

Choose a reason for hiding this comment

fajin-corp Jul 6, 2024

Choose a reason for hiding this comment

skottmckay Jul 2, 2024

Choose a reason for hiding this comment

guschmue commented Jul 11, 2024 • edited Loading

fajin-corp commented Jun 26, 2024 •

edited

Loading

guschmue commented Jul 11, 2024 •

edited

Loading