diff --git a/docs/source/features/onnx-transformations.md b/docs/source/features/onnx-transformations.md index 05e05db500..8fc122a293 100644 --- a/docs/source/features/onnx-transformations.md +++ b/docs/source/features/onnx-transformations.md @@ -2583,6 +2583,33 @@ No parameters required. } ``` +### `UnstackGatheredMatMulWeights` + +#### Description + +Static-unroll Mixture-of-Experts (MoE) blocks often keep the weights of all experts in a single `[E, K, N]` initializer and select one expert per `MatMul` with a constant index, either `Gather(W_all, k, axis=0)` with a scalar index or `Gather(W_all, [k], axis=0)` followed by `Squeeze(axes=[0])`. + +Weight-only quantizers such as `OnnxKQuantQuantization` only quantize `MatMul` nodes whose weight is an initializer, so these per-expert `MatMul` nodes are left unquantized. This surgery replaces each gathered weight with a 2-D initializer holding `W_all[k]` so a following quantization pass picks it up. `Gather`/`Squeeze` nodes and stacked initializers that are no longer used are removed. `MatMul` nodes whose expert index is computed at runtime are left unchanged. + +Run it before `OnnxKQuantQuantization` (or another weight-only `MatMul` quantizer) on models with static-unroll MoE experts. + +#### Configurations + +No parameters required. + +#### Example + +```json +{ + "type": "GraphSurgeries", + "surgeries": [ + { + "surgeon": "UnstackGatheredMatMulWeights" + } + ] +} +``` + ## ORT Performance Tuning diff --git a/olive/passes/onnx/graph_surgeries.py b/olive/passes/onnx/graph_surgeries.py index ce453913d8..1074a3d122 100644 --- a/olive/passes/onnx/graph_surgeries.py +++ b/olive/passes/onnx/graph_surgeries.py @@ -3434,6 +3434,134 @@ def _remove_from_graph(cls, graph: ir.Graph) -> int: return removed +class UnstackGatheredMatMulWeights(Surgeon): + """Materialize MatMul weights gathered from a stacked 3-D initializer as 2-D initializers. + + Static-unroll Mixture-of-Experts blocks often keep the weights of all experts in a + single ``[E, K, N]`` initializer and select one expert per MatMul with a constant index:: + + W_k = Gather(W_all, k, axis=0) # scalar k -> [K, N] + W_k = Squeeze(Gather(W_all, [k], axis=0), axes=[0]) # 1-element k -> [K, N] + y = MatMul(x, W_k) + + Weight-only quantizers such as ``OnnxKQuantQuantization`` only quantize MatMuls whose + weight is an initializer, so these per-expert MatMuls stay unquantized. This surgery + replaces each gathered weight with a 2-D initializer holding ``W_all[k]``, so that a + following quantization pass picks them up. Gather/Squeeze nodes and stacked initializers + left without consumers are removed. + + Only constant indices on axis 0 are rewritten. MatMuls whose expert index is computed at + runtime are left unchanged. + + When to use: + Run **before** ``OnnxKQuantQuantization`` (or another weight-only MatMul quantizer) + on models with static-unroll MoE experts. + """ + + def call_ir(self, model: ir.Model) -> ir.Model: + graph = model.graph + unstacked_weights: dict[tuple[str, int], ir.Value] = {} + # all values share one namespace (see _unique_value_name); collect it once since a model can have + # thousands of experts + taken_names = set(graph.initializers) + taken_names.update(value.name for value in graph.inputs if value.name) + taken_names.update(value.name for value in graph.outputs if value.name) + for graph_node in graph: + taken_names.update(output.name for output in graph_node.outputs if output.name) + rewritten = 0 + for node in list(graph): + if node.op_type != "MatMul" or node.domain not in ("", "ai.onnx") or len(node.inputs) != 2: + continue + weight = node.inputs[1] + if weight is None or weight.is_initializer(): + continue + + gathered = self._get_gathered_weight(weight) + if gathered is None: + continue + stacked, index = gathered + + key = (stacked.name, index) + if key not in unstacked_weights: + name = self._get_unique_name(taken_names, f"{stacked.name}_{index}") + taken_names.add(name) + array = np.ascontiguousarray(stacked.const_value.numpy()[index]) + value = ir.Value(name=name, const_value=ir.tensor(array, name=name)) + graph.register_initializer(value) + unstacked_weights[key] = value + node.replace_input_with(1, unstacked_weights[key]) + rewritten += 1 + + if rewritten: + RemoveUnusedNodesPass()(model) + logger.debug("Unstacked %d gathered MatMul weights into %d initializers", rewritten, len(unstacked_weights)) + return model + + @classmethod + def _get_gathered_weight(cls, weight: ir.Value) -> tuple[ir.Value, int] | None: + """Return the stacked initializer and constant expert index that produce ``weight``, if any.""" + producer = weight.producer() + squeezed = False + if producer is not None and producer.op_type == "Squeeze" and producer.domain in ("", "ai.onnx"): + if not cls._squeezes_only_leading_axis(producer): + return None + squeezed = True + producer = producer.inputs[0].producer() if producer.inputs[0] is not None else None + + if producer is None or producer.op_type != "Gather" or producer.domain not in ("", "ai.onnx"): + return None + if producer.attributes.get_int("axis", 0) != 0: + return None + + stacked, indices = producer.inputs + # an initializer that is also a graph input is an overridable default, not a constant + if stacked is None or not stacked.is_initializer() or stacked.is_graph_input() or stacked.const_value is None: + return None + if len(stacked.const_value.shape) != 3: + return None + + if indices is None or indices.is_graph_input(): + return None + index_tensor = ir.convenience.get_const_tensor(indices) + if index_tensor is None: + return None + index_array = index_tensor.numpy() + # a scalar index drops the expert axis, a one-element index keeps it and needs the Squeeze + if (squeezed and index_array.shape != (1,)) or (not squeezed and index_array.ndim != 0): + return None + + num_experts = stacked.const_value.shape[0] + index = int(index_array.reshape(-1)[0]) + if not -num_experts <= index < num_experts: + return None + return stacked, index % num_experts + + @staticmethod + def _squeezes_only_leading_axis(squeeze: ir.Node) -> bool: + if len(squeeze.inputs) > 1 and squeeze.inputs[1] is not None: + # overriding default axes could change the rank of the weight, so they must be constant + if squeeze.inputs[1].is_graph_input(): + return False + axes_tensor = ir.convenience.get_const_tensor(squeeze.inputs[1]) + if axes_tensor is None: + return False + axes = axes_tensor.numpy().reshape(-1).tolist() + else: + # opset < 13 keeps axes as an attribute + axes = list(squeeze.attributes.get_ints("axes", [])) + # rank of the Gather output is 3, so -3 also refers to the leading axis + return axes in ([0], [-3]) + + @staticmethod + def _get_unique_name(taken_names: set[str], base_name: str) -> str: + if base_name not in taken_names: + return base_name + index = 1 + while f"{base_name}_{index}" in taken_names: + index += 1 + return f"{base_name}_{index}" + + class RenameInputDims(Surgeon): """Rename / promote a dimension in an input tensor's shape to a named symbolic dim. diff --git a/test/passes/onnx/test_graph_surgeries.py b/test/passes/onnx/test_graph_surgeries.py index 6763ef9a0a..c132fa5564 100644 --- a/test/passes/onnx/test_graph_surgeries.py +++ b/test/passes/onnx/test_graph_surgeries.py @@ -3464,3 +3464,226 @@ def test_remove_memcpy_chained(tmp_path): x = np.random.randn(4, 8).astype(np.float32) result = sess.run(["output"], {"input": x})[0] np.testing.assert_allclose(result, np.maximum(x, 0), atol=1e-6) + + +def _make_stacked_expert_matmul_model(model_path: Path, weights: np.ndarray): + """Build x -> 3 MatMuls with expert weights gathered from one stacked [E, K, N] initializer. + + expert0: Gather with a scalar Constant index (no Squeeze) + expert1: Gather with a one-element initializer index followed by Squeeze(axes=[0]) + expert_dyn: Gather with a runtime index from a graph input, which must stay unchanged + """ + num_experts, k, n = weights.shape + nodes = [ + helper.make_node( + "Constant", [], ["idx0"], name="Idx0", value=numpy_helper.from_array(np.array(0, dtype=np.int64)) + ), + helper.make_node("Gather", ["W_all", "idx0"], ["w0"], name="Gather0", axis=0), + helper.make_node("MatMul", ["x", "w0"], ["y0"], name="MatMul0"), + helper.make_node("Gather", ["W_all", "idx1"], ["w1_3d"], name="Gather1", axis=0), + helper.make_node("Squeeze", ["w1_3d", "squeeze_axes"], ["w1"], name="Squeeze1"), + helper.make_node("MatMul", ["x", "w1"], ["y1"], name="MatMul1"), + helper.make_node("Gather", ["W_all", "expert_idx"], ["w_dyn"], name="GatherDyn", axis=0), + helper.make_node("MatMul", ["x", "w_dyn"], ["y_dyn"], name="MatMulDyn"), + ] + graph = helper.make_graph( + nodes=nodes, + name="StackedExperts", + inputs=[ + helper.make_tensor_value_info("x", TensorProto.FLOAT, [2, k]), + helper.make_tensor_value_info("expert_idx", TensorProto.INT64, []), + ], + outputs=[helper.make_tensor_value_info(name, TensorProto.FLOAT, [2, n]) for name in ("y0", "y1", "y_dyn")], + initializer=[ + numpy_helper.from_array(weights, "W_all"), + numpy_helper.from_array(np.array([num_experts - 1], dtype=np.int64), "idx1"), + numpy_helper.from_array(np.array([0], dtype=np.int64), "squeeze_axes"), + ], + ) + model = helper.make_model(graph, opset_imports=[helper.make_opsetid("", 17)]) + model.ir_version = 9 + onnx.save(model, model_path) + + +def test_unstack_gathered_matmul_weights_materializes_constant_expert_weights(tmp_path): + rng = np.random.default_rng(0) + weights = rng.standard_normal((3, 4, 5)).astype(np.float32) + model_path = tmp_path / "model.onnx" + _make_stacked_expert_matmul_model(model_path, weights) + + p = create_pass_from_dict( + GraphSurgeries, {"surgeries": [{"surgeon": "UnstackGatheredMatMulWeights"}]}, disable_search=True + ) + output_model = p.run(ONNXModelHandler(model_path=str(model_path)), str(tmp_path / "onnx")) + output_model_def = output_model.load_model() + + # constant-index MatMuls now take 2-D initializers, the runtime-index one is untouched + initializers = {init.name: numpy_helper.to_array(init) for init in output_model_def.graph.initializer} + matmuls = {node.name: node for node in output_model_def.graph.node if node.op_type == "MatMul"} + np.testing.assert_array_equal(initializers[matmuls["MatMul0"].input[1]], weights[0]) + np.testing.assert_array_equal(initializers[matmuls["MatMul1"].input[1]], weights[2]) + assert matmuls["MatMulDyn"].input[1] == "w_dyn" + # the Gather/Squeeze feeding constant experts are removed, the stacked weight stays for the dynamic Gather + node_names = {node.name for node in output_model_def.graph.node} + assert {"Gather0", "Gather1", "Squeeze1"}.isdisjoint(node_names) + assert "GatherDyn" in node_names + assert "W_all" in initializers + + # outputs are unchanged + x = rng.standard_normal((2, 4)).astype(np.float32) + feeds = {"x": x, "expert_idx": np.array(1, dtype=np.int64)} + expected = InferenceSession(str(model_path), providers=["CPUExecutionProvider"]).run(None, feeds) + actual = InferenceSession(output_model.model_path, providers=["CPUExecutionProvider"]).run(None, feeds) + for expected_output, actual_output in zip(expected, actual): + np.testing.assert_allclose(actual_output, expected_output, rtol=1e-6, atol=1e-6) + + +def test_unstack_gathered_matmul_weights_removes_unused_stacked_initializer(tmp_path): + weights = np.arange(2 * 3 * 4, dtype=np.float32).reshape(2, 3, 4) + nodes = [ + helper.make_node("Gather", ["W_all", "idx"], ["w"], name="Gather", axis=0), + helper.make_node("MatMul", ["x", "w"], ["y"], name="MatMul"), + ] + graph = helper.make_graph( + nodes=nodes, + name="SingleExpert", + inputs=[helper.make_tensor_value_info("x", TensorProto.FLOAT, [1, 3])], + outputs=[helper.make_tensor_value_info("y", TensorProto.FLOAT, [1, 4])], + initializer=[ + numpy_helper.from_array(weights, "W_all"), + numpy_helper.from_array(np.array(-1, dtype=np.int64), "idx"), + ], + ) + model = helper.make_model(graph, opset_imports=[helper.make_opsetid("", 17)]) + model.ir_version = 9 + model_path = tmp_path / "model.onnx" + onnx.save(model, model_path) + + p = create_pass_from_dict( + GraphSurgeries, {"surgeries": [{"surgeon": "UnstackGatheredMatMulWeights"}]}, disable_search=True + ) + output_model_def = p.run(ONNXModelHandler(model_path=str(model_path)), str(tmp_path / "onnx")).load_model() + + initializers = {init.name: numpy_helper.to_array(init) for init in output_model_def.graph.initializer} + assert "W_all" not in initializers + assert [node.op_type for node in output_model_def.graph.node] == ["MatMul"] + # negative index -1 selects the last expert + np.testing.assert_array_equal(initializers[output_model_def.graph.node[0].input[1]], weights[-1]) + + +def test_unstack_gathered_matmul_weights_lets_kquant_quantize_expert_matmuls(tmp_path): + from olive.hardware.accelerator import AcceleratorSpec + from olive.passes.onnx.kquant_quantization import OnnxKQuantQuantization + + rng = np.random.default_rng(0) + weights = rng.standard_normal((3, 64, 32)).astype(np.float32) + model_path = tmp_path / "model.onnx" + _make_stacked_expert_matmul_model(model_path, weights) + accelerator_spec = AcceleratorSpec(accelerator_type="CPU", execution_provider="CPUExecutionProvider") + + def count_matmul_nbits(model_handler): + return sum(node.op_type == str(OpType.MatMulNBits) for node in model_handler.load_model().graph.node) + + kquant_config = {"bits": 4, "block_size": 32} + kquant = create_pass_from_dict( + OnnxKQuantQuantization, kquant_config, disable_search=True, accelerator_spec=accelerator_spec + ) + # without the surgery no expert MatMul has an initializer weight + assert count_matmul_nbits(kquant.run(ONNXModelHandler(model_path=str(model_path)), tmp_path / "q0.onnx")) == 0 + + surgeries = create_pass_from_dict( + GraphSurgeries, {"surgeries": [{"surgeon": "UnstackGatheredMatMulWeights"}]}, disable_search=True + ) + unstacked = surgeries.run(ONNXModelHandler(model_path=str(model_path)), str(tmp_path / "unstacked")) + kquant = create_pass_from_dict( + OnnxKQuantQuantization, kquant_config, disable_search=True, accelerator_spec=accelerator_spec + ) + # the two constant-index experts are quantized, the runtime-index one is not + assert count_matmul_nbits(kquant.run(unstacked, tmp_path / "q1.onnx")) == 2 + + +def _run_unstack_surgery(model: onnx.ModelProto, tmp_path: Path) -> ONNXModelHandler: + model_path = tmp_path / "model.onnx" + onnx.save(model, model_path) + p = create_pass_from_dict( + GraphSurgeries, {"surgeries": [{"surgeon": "UnstackGatheredMatMulWeights"}]}, disable_search=True + ) + return p.run(ONNXModelHandler(model_path=str(model_path)), str(tmp_path / "onnx")) + + +def test_unstack_gathered_matmul_weights_avoids_value_name_collisions(tmp_path): + weights = np.arange(2 * 3 * 4, dtype=np.float32).reshape(2, 3, 4) + nodes = [ + helper.make_node("Gather", ["W_all", "idx"], ["w"], name="Gather", axis=0), + helper.make_node("MatMul", ["x", "w"], ["y"], name="MatMul"), + # a node output already uses the name the expert-0 initializer would get + helper.make_node("Identity", ["x"], ["W_all_0"], name="Identity"), + ] + graph = helper.make_graph( + nodes=nodes, + name="NameCollision", + inputs=[helper.make_tensor_value_info("x", TensorProto.FLOAT, [1, 3])], + outputs=[ + helper.make_tensor_value_info("y", TensorProto.FLOAT, [1, 4]), + helper.make_tensor_value_info("W_all_0", TensorProto.FLOAT, [1, 3]), + ], + initializer=[ + numpy_helper.from_array(weights, "W_all"), + numpy_helper.from_array(np.array(0, dtype=np.int64), "idx"), + ], + ) + model = helper.make_model(graph, opset_imports=[helper.make_opsetid("", 17)]) + model.ir_version = 9 + + output_model = _run_unstack_surgery(model, tmp_path) + output_model_def = output_model.load_model() + onnx.checker.check_model(output_model_def) + + matmul = next(node for node in output_model_def.graph.node if node.op_type == "MatMul") + assert matmul.input[1] != "W_all_0" + initializers = {init.name: numpy_helper.to_array(init) for init in output_model_def.graph.initializer} + np.testing.assert_array_equal(initializers[matmul.input[1]], weights[0]) + x = np.ones((1, 3), dtype=np.float32) + y, identity = InferenceSession(output_model.model_path, providers=["CPUExecutionProvider"]).run(None, {"x": x}) + np.testing.assert_allclose(y, x @ weights[0]) + np.testing.assert_array_equal(identity, x) + + +@pytest.mark.parametrize("overridable", ["W_all", "idx1", "squeeze_axes"]) +def test_unstack_gathered_matmul_weights_skips_overridable_defaults(overridable, tmp_path): + rng = np.random.default_rng(0) + weights = rng.standard_normal((3, 4, 5)).astype(np.float32) + nodes = [ + helper.make_node("Gather", ["W_all", "idx1"], ["w1_3d"], name="Gather1", axis=0), + helper.make_node("Squeeze", ["w1_3d", "squeeze_axes"], ["w1"], name="Squeeze1"), + helper.make_node("MatMul", ["x", "w1"], ["y1"], name="MatMul1"), + ] + initializers = { + "W_all": weights, + "idx1": np.array([2], dtype=np.int64), + "squeeze_axes": np.array([0], dtype=np.int64), + } + # an initializer that is also a graph input is a default the caller can override at inference time + graph = helper.make_graph( + nodes=nodes, + name="OverridableDefault", + inputs=[ + helper.make_tensor_value_info("x", TensorProto.FLOAT, [2, 4]), + helper.make_tensor_value_info( + overridable, + helper.np_dtype_to_tensor_dtype(initializers[overridable].dtype), + list(initializers[overridable].shape), + ), + ], + outputs=[helper.make_tensor_value_info("y1", TensorProto.FLOAT, None)], + initializer=[numpy_helper.from_array(value, name) for name, value in initializers.items()], + ) + model = helper.make_model(graph, opset_imports=[helper.make_opsetid("", 17)]) + model.ir_version = 9 + + output_model_def = _run_unstack_surgery(model, tmp_path).load_model() + + # nothing is materialized, so overriding the default still affects the MatMul + matmul = next(node for node in output_model_def.graph.node if node.op_type == "MatMul") + assert matmul.input[1] == "w1" + assert {node.name for node in output_model_def.graph.node} == {"Gather1", "Squeeze1", "MatMul1"}