From 9a1688a6e6fa02be33d812dab857637156df6664 Mon Sep 17 00:00:00 2001 From: mgohil Date: Thu, 15 Oct 2026 12:46:11 +0000 Subject: [PATCH 1/6] [DLA_transforms_9] [9/17] DLA LSTM decomposition Signed-off-by: mgohil --- .../dla_transforms/dla_decompose_lstm.py | 1130 +++++++++++++++++ .../onnx/graph_surgery/make_dla_compatible.py | 2 + .../dla_transforms/test_dla_decompose_lstm.py | 340 +++++ 3 files changed, 1472 insertions(+) create mode 100644 modelopt/onnx/graph_surgery/dla_transforms/dla_decompose_lstm.py create mode 100644 tests/unit/onnx/graph_surgery/dla_transforms/test_dla_decompose_lstm.py diff --git a/modelopt/onnx/graph_surgery/dla_transforms/dla_decompose_lstm.py b/modelopt/onnx/graph_surgery/dla_transforms/dla_decompose_lstm.py new file mode 100644 index 00000000000..07e38d8b0b9 --- /dev/null +++ b/modelopt/onnx/graph_surgery/dla_transforms/dla_decompose_lstm.py @@ -0,0 +1,1130 @@ +# SPDX-FileCopyrightText: Copyright (c) 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Expand ONNX ``LSTM`` nodes into a Conv-based subgraph for DLA-style pipelines. + +Input/output ranks follow the ONNX ``LSTM`` operator (``layout`` attribute default ``0``). +Full IO semantics and optional inputs (``B``, ``sequence_lens``, ``initial_h``, ``initial_c``, +``P``) are defined in the ONNX spec: + +https://onnx.ai/onnx/operators/onnx__LSTM.html + +This pass promotes selected inputs to 4D before the internal transpose/conv path: + +- **Index 0 (X):** ``[seq_length, batch_size, input_size]`` (3-D). +- **Index 1 (W):** ``[num_directions, 4*hidden_size, input_size]`` (3-D; often an initializer). +- **Index 2 (R):** ``[num_directions, 4*hidden_size, hidden_size]`` (3-D; often an initializer). +- **Index 3 (B), optional:** ``[num_directions, 8*hidden_size]`` (2-D if present). +- **Index 4 (sequence_lens):** not rewritten here (1-D); omitted from promotion. +- **Index 5-6 (initial_h / initial_c), optional:** ``[num_directions, batch_size, hidden_size]`` (3-D). +- **Index 7 (P), optional:** not promoted by this pass. + +``layout=1`` swaps batch/sequence layout in the spec; models using non-default layout may need +preprocessing outside this transform. +""" + +import numpy as np +import onnx +from onnx import helper, numpy_helper + +from ...logging_config import logger +from ..utils.graph_utils import ( + GraphCache, + add_unique_initializers, + batch_replace_nodes, + get_constant_by_name, + is_const_dq_input, +) + +# Promote ONNX LSTM inputs listed at https://onnx.ai/onnx/operators/onnx__LSTM.html (see module docstring). +# Initializers / DQ weights: numpy expand-to-4D. Graph tensors: Unsqueeze (axis 0 for typical 3-D X / states). +_LSTM_PROMOTE_INPUT_INDICES = (0, 1, 2, 3, 5, 6) + + +def _expand_array_to_rank4(arr: np.ndarray) -> np.ndarray: + rank = len(arr.shape) + if rank >= 4: + return arr + return np.expand_dims(arr, axis=list(range(4 - rank))) + + +def _validate_lstm_supported(node: onnx.NodeProto) -> str: + """Validate that ``node`` uses only LSTM semantics supported by this pass. + + The decomposition implements ONNX LSTM layout 0 with default activations, + no peepholes, no sequence_lens, no clipping, and input_forget disabled. + Return the normalized direction when validation succeeds. + """ + direction = "forward" + unsupported: list[str] = [] + + for attr in node.attribute: + if attr.name == "direction": + direction = attr.s.decode("utf-8").strip().lower() + elif attr.name == "layout" and int(attr.i) != 0: + unsupported.append(f"layout={int(attr.i)}") + elif attr.name == "clip" and float(attr.f) != 0.0: + unsupported.append(f"clip={float(attr.f)}") + elif attr.name == "input_forget" and int(attr.i) != 0: + unsupported.append("input_forget=1") + elif attr.name == "activations": + activations = [a.decode("utf-8").strip().lower() for a in attr.strings] + default_activations = ["sigmoid", "tanh", "tanh"] + if activations not in (default_activations, default_activations * 2): + unsupported.append(f"activations={activations}") + elif attr.name in ("activation_alpha", "activation_beta") and attr.floats: + unsupported.append(attr.name) + + if direction not in ("forward", "reverse", "bidirectional"): + unsupported.append(f"direction={direction!r}") + if len(node.input) > 4 and node.input[4]: + unsupported.append("sequence_lens") + if len(node.input) > 7 and node.input[7]: + unsupported.append("peephole P") + + if unsupported: + raise ValueError( + f"LSTM {node.name!r}: unsupported feature(s) for DLA decomposition: " + f"{', '.join(unsupported)}" + ) + return direction + + +def _promote_lstm_input_to_rank4( + model, graph, lstm_node, input_index, prep_nodes, unique_initializer, cache: GraphCache +) -> None: + """Promote one LSTM input slot to rank 4 for the decomposition path. + + Shapes follow ONNX ``LSTM`` (layout 0): https://onnx.ai/onnx/operators/onnx__LSTM.html + Initializers and DQ-backed weights are expanded in numpy; runtime tensors get ``Unsqueeze``. + """ + if input_index >= len(lstm_node.input): + return + tensor_name = lstm_node.input[input_index] + if not tensor_name: + return + + init_list = list(cache.init_names) + dq_init_name = is_const_dq_input(init_list, tensor_name, graph) + + if dq_init_name is not None: + deq_node = next( + ( + n + for n in graph.node + if n.op_type == "DequantizeLinear" + and n.output[0] == tensor_name + and n.input[0] == dq_init_name + ), + None, + ) + if deq_node is None: + return + initializer = cache.get_init(dq_init_name) + if initializer is None: + return + init_arr = numpy_helper.to_array(initializer) + init_arr = _expand_array_to_rank4(init_arr) + new_init = numpy_helper.from_array( + init_arr, name=initializer.name + f"_lstm4d_{lstm_node.name}_{input_index}" + ) + add_unique_initializers(graph, unique_initializer, [new_init]) + graph.initializer.remove(initializer) + deq_node.input[0] = new_init.name + return + + if cache.is_init(tensor_name): + initializer = cache.get_init(tensor_name) + if initializer is None: + return + init_arr = numpy_helper.to_array(initializer) + new_arr = _expand_array_to_rank4(init_arr) + new_init = numpy_helper.from_array( + new_arr, name=initializer.name + f"_lstm4d_{lstm_node.name}_{input_index}" + ) + add_unique_initializers(graph, unique_initializer, [new_init]) + refs = sum(1 for n in graph.node for inp in n.input if inp == tensor_name) + if refs <= 1: + graph.initializer.remove(initializer) + lstm_node.input[input_index] = new_init.name + return + + # Graph tensor: always prepend singleton dims to reach 4D. LSTM activations are expected 3D — one Unsqueeze(0) + # is enough when rank is unknown or 3; ranks 1-2 still use multiple leading axes; skip if already 4D+. + shape = cache.get_shape(tensor_name) + if shape is not None and len(shape) >= 4: + return + axes_list = [0] if shape is None or len(shape) == 3 else list(range(4 - len(shape))) + + axes_name = f"{tensor_name}_lstm4d_axes_{lstm_node.name}_{input_index}" + out_name = f"{tensor_name}_lstm4d_{lstm_node.name}_{input_index}" + axes_init = numpy_helper.from_array(np.array(axes_list, dtype=np.int64), name=axes_name) + add_unique_initializers(graph, unique_initializer, [axes_init]) + unsq = helper.make_node( + "Unsqueeze", + [tensor_name, axes_name], + [out_name], + name=f"lstm_prep_unsqueeze_{lstm_node.name}_{input_index}", + ) + prep_nodes.append(unsq) + lstm_node.input[input_index] = out_name + elem_t = cache.get_dtype(tensor_name) + if elem_t is not None and shape is not None: + new_shape = [int(x) for x in shape] + for ax in sorted(axes_list): + new_shape.insert(int(ax), 1) + graph.value_info.append(helper.make_tensor_value_info(out_name, int(elem_t), new_shape)) + + +def _fold_lstm_constants(graph, nodes, unique_initializer): + """Fold the constant Split, bias Add and initial-state Conv nodes created by LSTM.""" + initializers = {tensor.name: tensor for tensor in graph.initializer} + remaining = [] + for node in nodes: + if node.op_type not in {"Split", "Add", "Conv"} or not all( + name in initializers for name in node.input + ): + remaining.append(node) + continue + arrays = [numpy_helper.to_array(initializers[name]) for name in node.input] + if node.op_type == "Split": + axis = next( + (attribute.i for attribute in node.attribute if attribute.name == "axis"), 0 + ) + outputs = np.split(arrays[0], np.cumsum(arrays[1])[:-1], axis=axis) + elif node.op_type == "Add" and arrays[0].shape == arrays[1].shape: + outputs = [(arrays[0].astype(np.float64) + arrays[1]).astype(arrays[0].dtype)] + elif ( + node.op_type == "Conv" + and len(arrays) == 2 + and arrays[0].ndim == 4 + and arrays[1].ndim == 4 + and arrays[0].shape[0] == 1 + and arrays[0].shape[2:] == (1, 1) + and arrays[1].shape[1:] == (arrays[0].shape[1], 1, 1) + ): + # The reference accumulates these dot products sequentially in double precision. + products = arrays[1].reshape(arrays[1].shape[0], -1).astype(np.float64) + products = products * arrays[0].reshape(-1).astype(np.float64) + terms = np.concatenate([np.zeros((products.shape[0], 1)), products], axis=1) + values = np.cumsum(terms, axis=1)[:, -1].astype(arrays[0].dtype) + outputs = [values.reshape(1, arrays[1].shape[0], 1, 1)] + else: + remaining.append(node) + continue + for name, array in zip(node.output, outputs): + tensor = numpy_helper.from_array(array, name) + add_unique_initializers(graph, unique_initializer, [tensor]) + initializers[name] = tensor + return remaining + + +def _apply_decompose_lstm(model, *, defer_rank_conversion=False): + """Decompose ``LSTM`` nodes into basic ONNX ops (IO per ONNX LSTM layout 0).""" + graph = model.graph + cache = GraphCache(graph) + lstm_nodes = [n for n in list(graph.node) if n.op_type == "LSTM"] + node_idx = {id(n): i for i, n in enumerate(graph.node)} + replacements_batch: dict[int, tuple[list, list]] = {} + lstm_counter = 0 + unique_initializer = set() + + def add_nodes_for_time_step( + time_step, + h_t, + c_t, + z_t_iofc, + r_t_iofc, + nodes_to_add, + node, + prefix="forward", + shape_info_map=None, + ): + # add it = sigmoid(Ht * R_i + Z_i) + matmul_h_r_i_node = helper.make_node( + "Conv", + [h_t, r_t_iofc["i"]], + [f"{prefix}_h_r_i_{node.name}_seq_{time_step}"], + name=f"{prefix}_conv_h_r_i_{node.name}_seq_{time_step}", + ) + nodes_to_add.append(matmul_h_r_i_node) + add_h_r_i_z_i_node = helper.make_node( + "Add", + [matmul_h_r_i_node.output[0], z_t_iofc["i"]], + [f"{prefix}_h_r_i_z_i_{node.name}_seq_{time_step}"], + name=f"{prefix}_add_h_r_i_z_i_{node.name}_seq_{time_step}", + ) + nodes_to_add.append(add_h_r_i_z_i_node) + i_t_node = helper.make_node( + "Sigmoid", + [add_h_r_i_z_i_node.output[0]], + [f"{prefix}_i_t_{node.name}_seq_{time_step}"], + name=f"{prefix}_sigmoid_i_t_{node.name}_seq_{time_step}", + ) + nodes_to_add.append(i_t_node) + + # add ft = sigmoid(Ht * R_f + Z_f) + matmul_h_r_f_node = helper.make_node( + "Conv", + [h_t, r_t_iofc["f"]], + [f"{prefix}_h_r_f_{node.name}_seq_{time_step}"], + name=f"{prefix}_conv_h_r_f_{node.name}_seq_{time_step}", + ) + nodes_to_add.append(matmul_h_r_f_node) + add_h_r_f_z_f_node = helper.make_node( + "Add", + [matmul_h_r_f_node.output[0], z_t_iofc["f"]], + [f"{prefix}_h_r_f_z_f_{node.name}_seq_{time_step}"], + name=f"{prefix}_add_h_r_f_z_f_{node.name}_seq_{time_step}", + ) + nodes_to_add.append(add_h_r_f_z_f_node) + ft_node = helper.make_node( + "Sigmoid", + [add_h_r_f_z_f_node.output[0]], + [f"{prefix}_ft_{node.name}_seq_{time_step}"], + name=f"{prefix}_sigmoid_ft_{node.name}_seq_{time_step}", + ) + nodes_to_add.append(ft_node) + + # add ot = sigmoid(Ht * R_o + Z_o) + matmul_h_r_o_node = helper.make_node( + "Conv", + [h_t, r_t_iofc["o"]], + [f"{prefix}_h_r_o_{node.name}_seq_{time_step}"], + name=f"{prefix}_conv_h_r_o_{node.name}_seq_{time_step}", + ) + nodes_to_add.append(matmul_h_r_o_node) + add_h_r_o_z_o_node = helper.make_node( + "Add", + [matmul_h_r_o_node.output[0], z_t_iofc["o"]], + [f"{prefix}_h_r_o_z_o_{node.name}_seq_{time_step}"], + name=f"{prefix}_add_h_r_o_z_o_{node.name}_seq_{time_step}", + ) + nodes_to_add.append(add_h_r_o_z_o_node) + ot_node = helper.make_node( + "Sigmoid", + [add_h_r_o_z_o_node.output[0]], + [f"{prefix}_ot_{node.name}_seq_{time_step}"], + name=f"{prefix}_sigmoid_ot_{node.name}_seq_{time_step}", + ) + nodes_to_add.append(ot_node) + + # add ct = tanh(Ht * R_c + Z_c) + matmul_h_r_c_node = helper.make_node( + "Conv", + [h_t, r_t_iofc["c"]], + [f"{prefix}_h_r_c_{node.name}_seq_{time_step}"], + name=f"{prefix}_conv_h_r_c_{node.name}_seq_{time_step}", + ) + nodes_to_add.append(matmul_h_r_c_node) + add_h_r_c_z_c_node = helper.make_node( + "Add", + [matmul_h_r_c_node.output[0], z_t_iofc["c"]], + [f"{prefix}_h_r_c_z_c_{node.name}_seq_{time_step}"], + name=f"{prefix}_add_h_r_c_z_c_{node.name}_seq_{time_step}", + ) + nodes_to_add.append(add_h_r_c_z_c_node) + ct_node = helper.make_node( + "Tanh", + [add_h_r_c_z_c_node.output[0]], + [f"{prefix}_ct_{node.name}_seq_{time_step}"], + name=f"{prefix}_tanh_ct_{node.name}_seq_{time_step}", + ) + nodes_to_add.append(ct_node) + + # add Ct = ft * Ct-1 + it * ct + mul_ft_ct_1_node = helper.make_node( + "Mul", + [ft_node.output[0], c_t], + [f"{prefix}_ft_ct_1_{node.name}_seq_{time_step}"], + name=f"{prefix}_mul_ft_ct_1_{node.name}_seq_{time_step}", + ) + nodes_to_add.append(mul_ft_ct_1_node) + mul_it_ct_node = helper.make_node( + "Mul", + [i_t_node.output[0], ct_node.output[0]], + [f"{prefix}_it_ct_{node.name}_seq_{time_step}"], + name=f"{prefix}_mul_it_ct_{node.name}_seq_{time_step}", + ) + nodes_to_add.append(mul_it_ct_node) + add_ft_ct_1_it_ct_node = helper.make_node( + "Add", + [mul_ft_ct_1_node.output[0], mul_it_ct_node.output[0]], + [f"{prefix}_C_t_{node.name}_seq_{time_step}"], + name=f"{prefix}_add_ft_ct_1_it_ct_{node.name}_seq_{time_step}", + ) + nodes_to_add.append(add_ft_ct_1_it_ct_node) + + # add Ht = ot * tanh(Ct) + tanh_ct_node = helper.make_node( + "Tanh", + [add_ft_ct_1_it_ct_node.output[0]], + [f"{prefix}_tanh_Ct_{node.name}_seq_{time_step}"], + name=f"{prefix}_tanh_Ct_{node.name}_seq_{time_step}", + ) + nodes_to_add.append(tanh_ct_node) + mul_ot_tanh_ct_node = helper.make_node( + "Mul", + [ot_node.output[0], tanh_ct_node.output[0]], + [f"{prefix}_H_t_{node.name}_seq_{time_step}"], + name=f"{prefix}_mul_ot_tanh_ct_{node.name}_seq_{time_step}", + ) + nodes_to_add.append(mul_ot_tanh_ct_node) + + return mul_ot_tanh_ct_node.output[0], add_ft_ct_1_it_ct_node.output[0] + + def transpose_tensor(graph, tensor_name, perm, nodes_to_add, unique_initializer): + init_list = [init.name for init in graph.initializer] + if is_const_dq_input(init_list, tensor_name, graph) is not None: + init_name = is_const_dq_input(init_list, tensor_name, graph) + deq_node = next( + ( + n + for n in graph.node + if n.op_type == "DequantizeLinear" + and n.output[0] == tensor_name + and n.input[0] == init_name + ), + None, + ) + if deq_node is None: + raise ValueError( + f"LSTM transpose_tensor: DequantizeLinear for DQ input {tensor_name!r} (init {init_name!r}) " + "not found in graph." + ) + initializer = cache.get_init(init_name) + if initializer is None: + raise ValueError( + f"LSTM transpose_tensor: initializer {init_name!r} missing for DQ tensor {tensor_name!r}." + ) + init_arr = numpy_helper.to_array(initializer) + init_arr = _expand_array_to_rank4(init_arr) + + init_arr = np.transpose(init_arr, axes=perm) + new_init = numpy_helper.from_array(init_arr, name=initializer.name + "_transposed") + add_unique_initializers(graph, unique_initializer, [new_init]) + graph.initializer.remove(initializer) + deq_node.input[0] = new_init.name + return tensor_name + elif tensor_name in init_list: + initializer = cache.get_init(tensor_name) + if initializer is None: + raise ValueError( + f"LSTM transpose_tensor: initializer {tensor_name!r} missing from graph." + ) + init_arr = numpy_helper.to_array(initializer) + init_arr = _expand_array_to_rank4(init_arr) + init_arr = np.transpose(init_arr, axes=perm) + new_init = numpy_helper.from_array(init_arr, name=initializer.name + "_transposed") + add_unique_initializers(graph, unique_initializer, [new_init]) + graph.initializer.remove(initializer) + return new_init.name + else: + transpose_node = helper.make_node( + "Transpose", + inputs=[tensor_name], + outputs=[tensor_name + "_transposed"], + name=tensor_name + "_transpose", + perm=perm, + ) + nodes_to_add.append(transpose_node) + return transpose_node.output[0] + + def transpose_inputs(graph, node, unique_initializer, shape_info_map, nodes_to_add): + # return list of names of [X, W, R, B, ih, ic] + + # When batch_size==1 the batch dim is already a leading unary dim after + # 4D promotion. Use [0,3,2,1] to keep it in place instead of swapping + # it to the back with [2,3,0,1]. + transpose_perm1 = [0, 3, 2, 1] if shape_info_map.get("N", 0) == 1 else [2, 3, 0, 1] + transpose_perm2 = [2, 3, 1, 0] + transpose_perm3 = [0, 3, 2, 1] + + # transpose X, ih, ic with perm1 + x_transposed = transpose_tensor( + graph, node.input[0], transpose_perm1, nodes_to_add, unique_initializer + ) + # transpose W, R with perm2 + w_transposed = transpose_tensor( + graph, node.input[1], transpose_perm2, nodes_to_add, unique_initializer + ) + r_transposed = transpose_tensor( + graph, node.input[2], transpose_perm2, nodes_to_add, unique_initializer + ) + # transpose using perm3 and split into bw and br + bw_transposed = None + br_transposed = None + if len(node.input) > 3 and node.input[3]: + b_transposed = transpose_tensor( + graph, node.input[3], transpose_perm3, nodes_to_add, unique_initializer + ) + + bw_br_split_init = numpy_helper.from_array( + np.array([4 * shape_info_map["H"], 4 * shape_info_map["H"]], dtype=np.int64), + name=f"bw_br_split_init_{node.name}", + ) + add_unique_initializers(graph, unique_initializer, [bw_br_split_init]) + bw_br_split_node = helper.make_node( + "Split", + [b_transposed, bw_br_split_init.name], + [f"bw_{node.name}", f"br_{node.name}"], + axis=1, + name=f"split_B_{node.name}_axis_updated", + ) + nodes_to_add.append(bw_br_split_node) + bw_transposed = bw_br_split_node.output[0] + br_transposed = bw_br_split_node.output[1] + + ih_transposed = None + if len(node.input) > 5 and node.input[5]: + ih_transposed = transpose_tensor( + graph, node.input[5], transpose_perm2, nodes_to_add, unique_initializer + ) + else: + ih_init = numpy_helper.from_array( + np.zeros( + [shape_info_map["N"], shape_info_map["H"], shape_info_map["d"], 1], + dtype=np.float32, + ), + name=node.name + "/ih", + ) + add_unique_initializers(graph, unique_initializer, [ih_init]) + ih_transposed = ih_init.name + + if len(node.input) > 6 and node.input[6]: + ic_transposed = transpose_tensor( + graph, node.input[6], transpose_perm2, nodes_to_add, unique_initializer + ) + else: + ic_init = numpy_helper.from_array( + np.zeros( + [shape_info_map["N"], shape_info_map["H"], shape_info_map["d"], 1], + dtype=np.float32, + ), + name=node.name + "/ic", + ) + add_unique_initializers(graph, unique_initializer, [ic_init]) + ic_transposed = ic_init.name + if bw_transposed is None and br_transposed is None: + # Optional B (ONNX LSTM): zero biases broadcast with Add(bw, br) and Add(..., conv_out). + zscalar = np.zeros((1, 1, 1, 1), dtype=np.float32) + bw_init = numpy_helper.from_array( + zscalar, + name=(node.name + "/zero_bw").replace("/", "_"), + ) + br_init = numpy_helper.from_array( + np.array(zscalar, copy=True), + name=(node.name + "/zero_br").replace("/", "_"), + ) + add_unique_initializers(graph, unique_initializer, [bw_init, br_init]) + bw_transposed = bw_init.name + br_transposed = br_init.name + elif bw_transposed is None or br_transposed is None: + raise ValueError( + f"LSTM {node.name!r}: bias B produced only one of bw/br; unsupported partial bias." + ) + return { + "x": x_transposed, + "w": w_transposed, + "r": r_transposed, + "bw": bw_transposed, + "br": br_transposed, + "ih": ih_transposed, + "ic": ic_transposed, + } + + def resolve_lstm_tensor_shape(tensor_name, _seen=None): + """Static shape for LSTM X/W when value_info is missing (DQ, Unsqueeze prep, initializers).""" + if _seen is None: + _seen = set() + if not tensor_name or tensor_name in _seen: + return None + _seen.add(tensor_name) + + s = cache.get_shape(tensor_name) + if s is not None: + return list(s) + + if cache.is_init(tensor_name): + init = cache.get_init(tensor_name) + if init is None: + return None + return list(numpy_helper.to_array(init).shape) + + init_names = list(cache.init_names) + dq_src = is_const_dq_input(init_names, tensor_name, graph) + if dq_src: + return resolve_lstm_tensor_shape(dq_src, _seen) + + prod = cache.get_producer(tensor_name) + if prod is None: + return None + if prod.op_type == "Unsqueeze" and len(prod.input) >= 2: + inner = prod.input[0] + axes_name = prod.input[1] + inner_shape = resolve_lstm_tensor_shape(inner, _seen) + if inner_shape is None: + return None + ax = cache.get_init_array(axes_name) + if ax is None: + ax = get_constant_by_name(model, axes_name) + if ax is None: + return None + axes_flat = np.asarray(ax, dtype=np.int64).flatten().tolist() + out = list(inner_shape) + r = len(out) + norm = [int(a) + r + 1 if int(a) < 0 else int(a) for a in axes_flat] + for a in sorted(norm): + if a < 0 or a > len(out): + return None + out.insert(a, 1) + return out + if prod.op_type in ("Identity", "Cast") and prod.input and prod.input[0]: + return resolve_lstm_tensor_shape(prod.input[0], _seen) + return None + + def create_shape_info_map(node): + x, w = node.input[:2] + x_shape = resolve_lstm_tensor_shape(x) + w_shape = resolve_lstm_tensor_shape(w) + if x_shape is None or w_shape is None: + parts = [] + if x_shape is None: + parts.append(f"X ({x!r})") + if w_shape is None: + parts.append(f"W ({w!r})") + raise ValueError( + f"LSTM {node.name!r}: missing static shapes for {', '.join(parts)}. " + "Could not derive ranks from value_info, initializers, get_shape_from_graph, " + "or Unsqueeze/Identity/Cast chains." + ) + hidden_size = w_shape[-2] // 4 + num_directions = w_shape[-3] + input_size = w_shape[-1] + seq_length = x_shape[-3] + batch_size = x_shape[-2] + return { + "D": input_size, + "H": hidden_size, + "d": num_directions, + "T": seq_length, + "N": batch_size, + } + + def split_direction_wise( + graph, node, unique_initializer, shape_info_map, transposed_inputs, nodes_to_add + ): + if shape_info_map["d"] == 1: + # Unidirectional: reverse and forward paths use the same tensors (only timestep order differs). + return [transposed_inputs, transposed_inputs] + + direction_split_init = numpy_helper.from_array( + np.array([1, 1], dtype=np.int64), name=f"direction_split_init_{node.name}" + ) + add_unique_initializers(graph, unique_initializer, [direction_split_init]) + + forward_list = {} + reverse_list = {} + forward_list["x"] = transposed_inputs["x"] + reverse_list["x"] = transposed_inputs["x"] + to_be_split = ["w", "r", "bw", "br", "ih", "ic"] + for input_name in to_be_split: + input_value = transposed_inputs[input_name] + if not input_value: + raise ValueError( + f"LSTM {node.name!r}: missing tensor for direction split ({input_name!r}); " + "required inputs may be absent or empty." + ) + direction_split_node = helper.make_node( + "Split", + [str(input_value), direction_split_init.name], + [ + f"forward_direction_split_{input_name}_{node.name}", + f"reverse_direction_split_{input_name}_{node.name}", + ], + axis=2, + name=f"split_direction_wise_{input_name}_{node.name}_axis_updated", + ) + nodes_to_add.append(direction_split_node) + forward_list[input_name] = direction_split_node.output[0] + reverse_list[input_name] = direction_split_node.output[1] + return [forward_list, reverse_list] + + def precompute_and_split_seq_len( + graph, node, unique_initializer, shape_info_map, input_list, nodes_to_add, direction, prefix + ): + x_input = input_list["x"] + w_input = input_list["w"] + bw_input = input_list["bw"] + br_input = input_list["br"] + if not x_input or not w_input: + raise ValueError( + f"LSTM {node.name!r} ({prefix}): missing X or W tensor name after transpose " + f"(x={x_input!r}, w={w_input!r})." + ) + if bw_input is None or br_input is None: + raise ValueError( + f"LSTM {node.name!r} ({prefix}): missing bias tensors bw/br (bw={bw_input!r}, br={br_input!r})." + ) + + add_bw_br_node = helper.make_node( + "Add", + [str(bw_input), str(br_input)], + [f"{prefix}_add_bw_br_{node.name}"], + name=f"{prefix}_add_bw_br_{node.name}", + ) + nodes_to_add.append(add_bw_br_node) + + conv_x_w_node = helper.make_node( + "Conv", + [str(x_input), str(w_input)], + [f"{prefix}_conv_x_w_{node.name}"], + name=f"{prefix}_conv_x_w_{node.name}", + ) + nodes_to_add.append(conv_x_w_node) + + z_node = helper.make_node( + "Add", + [add_bw_br_node.output[0], conv_x_w_node.output[0]], + [f"{prefix}_add_bw_br_conv_x_w_{node.name}"], + name=f"{prefix}_add_bw_br_conv_x_w_{node.name}", + ) + nodes_to_add.append(z_node) + + split_z_iofc_init = numpy_helper.from_array( + np.array( + [ + shape_info_map["H"], + shape_info_map["H"], + shape_info_map["H"], + shape_info_map["H"], + ], + dtype=np.int64, + ), + name=f"split_z_iofc_init_{node.name}", + ) + add_unique_initializers(graph, unique_initializer, [split_z_iofc_init]) + split_z_iofc_node = helper.make_node( + "Split", + [z_node.output[0], split_z_iofc_init.name], + [ + f"{prefix}_split_z_i_{node.name}", + f"{prefix}_split_z_o_{node.name}", + f"{prefix}_split_z_f_{node.name}", + f"{prefix}_split_z_c_{node.name}", + ], + axis=1, + name=f"{prefix}_split_z_iofc_{node.name}_axis_updated", + ) + nodes_to_add.append(split_z_iofc_node) + + split_seq_init = numpy_helper.from_array( + np.array([1] * shape_info_map["T"], dtype=np.int64), + name=f"split_z_iofc_seq_init_{node.name}", + ) + add_unique_initializers(graph, unique_initializer, [split_seq_init]) + + split_z_iofc_seq = {} + for split_z_iofc_output, split_z_iofc_name in zip( + split_z_iofc_node.output, ["i", "o", "f", "c"] + ): + split_seq_node = helper.make_node( + "Split", + [split_z_iofc_output, split_seq_init.name], + [ + f"{prefix}_split_z_{split_z_iofc_name}_{time_step}_{node.name}" + for time_step in range(shape_info_map["T"]) + ], + axis=3, + name=f"{prefix}_split_z_iofc_seq_{split_z_iofc_name}_{node.name}_axis_updated", + ) + nodes_to_add.append(split_seq_node) + split_z_iofc_seq[split_z_iofc_name] = (split_seq_node.output)[::direction] + + return split_z_iofc_seq + + def split_r_iofc( + graph, node, unique_initializer, shape_info_map, input_list, nodes_to_add, prefix + ): + r_input = input_list["r"] + if not r_input: + raise ValueError( + f"LSTM {node.name!r} ({prefix}): missing R tensor after transpose (r={r_input!r})." + ) + split_r_iofc_init = numpy_helper.from_array( + np.array( + [ + shape_info_map["H"], + shape_info_map["H"], + shape_info_map["H"], + shape_info_map["H"], + ], + dtype=np.int64, + ), + name=f"split_r_iofc_init_{node.name}", + ) + add_unique_initializers(graph, unique_initializer, [split_r_iofc_init]) + split_r_iofc_node = helper.make_node( + "Split", + [str(r_input), split_r_iofc_init.name], + [ + f"{prefix}_split_r_i_{node.name}", + f"{prefix}_split_r_o_{node.name}", + f"{prefix}_split_r_f_{node.name}", + f"{prefix}_split_r_c_{node.name}", + ], + axis=0, + name=f"{prefix}_split_r_iofc_{node.name}_axis_updated", + ) + nodes_to_add.append(split_r_iofc_node) + return { + "i": split_r_iofc_node.output[0], + "o": split_r_iofc_node.output[1], + "f": split_r_iofc_node.output[2], + "c": split_r_iofc_node.output[3], + } + + def transpose_outputs( + graph, node, unique_initializer, input_name, output_name, nodes_to_add, squeeze_axes=None + ): + """Apply output Transpose then optional Squeeze for hidden/cell states only. + + Y stays 4D [T, d, N, H] after perm [3,2,0,1]. Y_h / Y_c use squeeze axis 0: + [1, d, N, H] -> [d, N, H] (ONNX 3D states). + """ + if not output_name: + return None + transpose_perm = [3, 2, 0, 1] + if squeeze_axes: + transposed_tmp = f"{output_name}_lstm_tr_{node.name}".replace("/", "_") + transpose_out = transposed_tmp + else: + transpose_out = output_name + + if not input_name: + raise ValueError( + f"LSTM {node.name!r}: transpose_outputs got empty input for output {output_name!r}." + ) + transpose_node = helper.make_node( + "Transpose", + [str(input_name)], + [str(transpose_out)], + name=f"{output_name}_transpose_{node.name}", + perm=transpose_perm, + ) + nodes_to_add.append(transpose_node) + + if squeeze_axes: + axes_name = f"{output_name}_lstm_squeeze_axes_{node.name}".replace("/", "_") + axes_init = numpy_helper.from_array( + np.array(squeeze_axes, dtype=np.int64), name=axes_name + ) + add_unique_initializers(graph, unique_initializer, [axes_init]) + squeeze_node = helper.make_node( + "Squeeze", + [transpose_out, axes_name], + [output_name], + name=f"{output_name}_lstm_squeeze_{node.name}", + ) + nodes_to_add.append(squeeze_node) + return squeeze_node.output[0] + + new_perm = [0, 1, 2, 3] + consumer_nodes_y_output = cache.get_consumers(transpose_node.output[0]) + if ( + len(consumer_nodes_y_output) == 1 + and consumer_nodes_y_output[0][0].op_type == "Transpose" + ): + consumer = consumer_nodes_y_output[0][0] + shape = cache.get_shape(output_name) + perm = next( + (list(attr.ints) for attr in consumer.attribute if attr.name == "perm"), None + ) + followers = cache.get_consumers(consumer.output[0]) + # The later rank pass resolves 0 dimensions from original output metadata. + restored_by_rank_pass = ( + defer_rank_conversion + and consumer.output[0] not in cache.graph_output_names + and bool(followers) + and all( + next_node.op_type == "Reshape" + and len(next_node.input) > 1 + and (target := cache.get_init_array(next_node.input[1])) is not None + and 0 < target.size < 4 + and cache.get_shape(next_node.output[0]) is not None + for next_node, _ in followers + ) + ) + noop = ( + shape is not None + and perm is not None + and len(shape) == len(perm) + and ([shape[index] for index in perm] == shape or restored_by_rank_pass) + and [index for index in perm if shape[index] != 1] + == [index for index, dim in enumerate(shape) if dim != 1] + ) + for attr in consumer.attribute if noop else []: + if attr.name == "perm": + del attr.ints[:] + attr.ints.extend(int(x) for x in new_perm) + break + + return transpose_node.output[0] + + if ( + lstm_nodes + and next((opset.version for opset in model.opset_import if not opset.domain), 0) < 13 + ): + raise ValueError("LSTM decomposition requires ONNX opset >= 13") + for node in lstm_nodes: + _validate_lstm_supported(node) + shape_info = create_shape_info_map(node) + if any(value <= 0 for value in shape_info.values()): + raise ValueError(f"LSTM {node.name!r}: all dimensions must be static and positive") + dtype = cache.get_dtype(node.input[0]) + if dtype != onnx.TensorProto.FLOAT: + raise ValueError(f"LSTM {node.name!r}: decomposition supports FP32 inputs") + names = set() + for index, node in enumerate(lstm_nodes): + if not node.name or node.name in names: + node.name = cache.unique_name(node.name or f"lstm_{index}") + names.add(node.name) + direction = _validate_lstm_supported(node) + prep_nodes = [] + for idx in _LSTM_PROMOTE_INPUT_INDICES: + _promote_lstm_input_to_rank4( + model, graph, node, idx, prep_nodes, unique_initializer, cache + ) + # Rebuild cache after promotion added new initializers/value_info/nodes + cache.rebuild() + nodes_to_add = list(prep_nodes) + shape_info_map = create_shape_info_map(node) + if direction == "bidirectional" and int(shape_info_map["d"]) != 2: + raise ValueError( + f"LSTM {node.name!r}: direction is bidirectional but num_directions (from W) is " + f"{shape_info_map['d']!r}; expected 2." + ) + transposed_inputs = transpose_inputs( + graph, node, unique_initializer, shape_info_map, nodes_to_add + ) + forward_list, reverse_list = split_direction_wise( + graph, node, unique_initializer, shape_info_map, transposed_inputs, nodes_to_add + ) + y_forward = None + y_h_forward = None + y_c_forward = None + y_reverse = None + y_h_reverse = None + y_c_reverse = None + if direction in ("forward", "bidirectional"): + prefix = "forward" + split_z_iofc_seq = precompute_and_split_seq_len( + graph, + node, + unique_initializer, + shape_info_map, + forward_list, + nodes_to_add, + direction=1, + prefix=prefix, + ) + split_r_iofc_map = split_r_iofc( + graph, + node, + unique_initializer, + shape_info_map, + forward_list, + nodes_to_add, + prefix=prefix, + ) + h_t, c_t = forward_list["ih"], forward_list["ic"] + h_t_list = [] + c_t_list = [] + for time_step in range(shape_info_map["T"]): + split_z_iofc_seq_time_step = {k: v[time_step] for k, v in split_z_iofc_seq.items()} + h_t, c_t = add_nodes_for_time_step( + time_step, + h_t, + c_t, + split_z_iofc_seq_time_step, + split_r_iofc_map, + nodes_to_add, + node, + prefix=prefix, + shape_info_map=shape_info_map, + ) + h_t_list.append(h_t) + c_t_list.append(c_t) + + # concatenate along seq_len dim here axis is 3 + h_t_concat_node = helper.make_node( + "Concat", + h_t_list, + [f"{prefix}_h_t_concat_{node.name}"], + name=f"{prefix}_concat_h_t_concat_{node.name}_axis_updated", + axis=3, + ) + nodes_to_add.append(h_t_concat_node) + y_forward = h_t_concat_node.output[0] + y_h_forward = h_t_list[-1] + y_c_forward = c_t_list[-1] + if direction in ("reverse", "bidirectional"): + prefix = "reverse" + split_z_iofc_seq = precompute_and_split_seq_len( + graph, + node, + unique_initializer, + shape_info_map, + reverse_list, + nodes_to_add, + direction=-1, + prefix=prefix, + ) + split_r_iofc_map = split_r_iofc( + graph, + node, + unique_initializer, + shape_info_map, + reverse_list, + nodes_to_add, + prefix=prefix, + ) + h_t, c_t = reverse_list["ih"], reverse_list["ic"] + h_t_list = [] + c_t_list = [] + for time_step in range(shape_info_map["T"]): + split_z_iofc_seq_time_step = {k: v[time_step] for k, v in split_z_iofc_seq.items()} + h_t, c_t = add_nodes_for_time_step( + time_step, + h_t, + c_t, + split_z_iofc_seq_time_step, + split_r_iofc_map, + nodes_to_add, + node, + prefix=prefix, + shape_info_map=shape_info_map, + ) + h_t_list.append(h_t) + c_t_list.append(c_t) + + # concatenate along seq_len dim here axis is 3 + h_t_concat_node = helper.make_node( + "Concat", + h_t_list[::-1], + [f"{prefix}_h_t_concat_{node.name}"], + name=f"{prefix}_concat_h_t_concat_{node.name}_axis_updated", + axis=3, + ) + nodes_to_add.append(h_t_concat_node) + y_reverse = h_t_concat_node.output[0] + y_h_reverse = h_t_list[-1] + y_c_reverse = c_t_list[-1] + y_before_transpose = None + y_h_before_transpose = None + y_c_before_transpose = None + if direction == "bidirectional": + if y_forward is None or y_reverse is None: + raise ValueError( + f"LSTM {node.name!r}: bidirectional concat needs both forward and reverse " + f"tensors; got y_forward={y_forward!r}, y_reverse={y_reverse!r}. " + "If you still see this after updating, the direction attribute may be " + "unusual (expected 'bidirectional')." + ) + y_concat_node = helper.make_node( + "Concat", + [str(y_forward), str(y_reverse)], + [f"y_concat_{node.name}"], + name=f"concat_y_concat_{node.name}_axis_updated", + axis=2, + ) + nodes_to_add.append(y_concat_node) + y_before_transpose = y_concat_node.output[0] + + if len(node.output) > 1 and node.output[1]: + if y_h_forward is None or y_h_reverse is None: + raise ValueError( + f"LSTM {node.name!r}: bidirectional Y_h concat missing " + f"y_h_forward={y_h_forward!r}, y_h_reverse={y_h_reverse!r}" + ) + y_h_concat_node = helper.make_node( + "Concat", + [str(y_h_forward), str(y_h_reverse)], + [f"y_h_concat_{node.name}"], + name=f"concat_y_h_concat_{node.name}_axis_updated", + axis=2, + ) + nodes_to_add.append(y_h_concat_node) + y_h_before_transpose = y_h_concat_node.output[0] + + if len(node.output) > 2 and node.output[2]: + if y_c_forward is None or y_c_reverse is None: + raise ValueError( + f"LSTM {node.name!r}: bidirectional Y_c concat missing " + f"y_c_forward={y_c_forward!r}, y_c_reverse={y_c_reverse!r}" + ) + y_c_concat_node = helper.make_node( + "Concat", + [str(y_c_forward), str(y_c_reverse)], + [f"y_c_concat_{node.name}"], + name=f"concat_y_c_concat_{node.name}_axis_updated", + axis=2, + ) + nodes_to_add.append(y_c_concat_node) + y_c_before_transpose = y_c_concat_node.output[0] + + elif direction == "forward": + y_before_transpose = y_forward + y_h_before_transpose = y_h_forward + y_c_before_transpose = y_c_forward + elif direction == "reverse": + y_before_transpose = y_reverse + y_h_before_transpose = y_h_reverse + y_c_before_transpose = y_c_reverse + else: + raise ValueError(f"Invalid direction: {direction}") + + transpose_outputs( + graph, node, unique_initializer, y_before_transpose, node.output[0], nodes_to_add + ) + if len(node.output) > 1 and node.output[1]: + transpose_outputs( + graph, + node, + unique_initializer, + y_h_before_transpose, + node.output[1], + nodes_to_add, + squeeze_axes=[0], + ) + if len(node.output) > 2 and node.output[2]: + transpose_outputs( + graph, + node, + unique_initializer, + y_c_before_transpose, + node.output[2], + nodes_to_add, + squeeze_axes=[0], + ) + + nodes_to_add = _fold_lstm_constants(graph, nodes_to_add, unique_initializer) + replacements_batch[node_idx[id(node)]] = ([node], nodes_to_add) + lstm_counter += 1 + + batch_replace_nodes(graph, replacements_batch) + logger.debug("Updated %d LSTM nodes with manual decomposition", lstm_counter) + + return model diff --git a/modelopt/onnx/graph_surgery/make_dla_compatible.py b/modelopt/onnx/graph_surgery/make_dla_compatible.py index 787a17767d3..bb6168d140d 100644 --- a/modelopt/onnx/graph_surgery/make_dla_compatible.py +++ b/modelopt/onnx/graph_surgery/make_dla_compatible.py @@ -32,6 +32,7 @@ from ..logging_config import logger from .dla_transforms.dla_cast_to_fp32 import _apply_cast_to_fp32 from .dla_transforms.dla_constants_to_initializers import _constants_to_initializers_from_model +from .dla_transforms.dla_decompose_lstm import _apply_decompose_lstm from .dla_transforms.dla_graph_cleanup import _remove_unused_initializers from .dla_transforms.dla_matmul_to_transpose_conv_transpose import ( _apply_matmul_to_transpose_conv_transpose, @@ -81,6 +82,7 @@ def _apply_constants_to_initializers(model, **options): ("ReshapeTransposeReshape", _apply_reshape_transpose_reshape), ("ReshapeQDQReduceSum", _apply_reshape_qdq_reduce_sum), ("ReshapeReduceSum", _apply_reshape_reduce_sum), + ("DecomposeLstm", partial(_apply_decompose_lstm, defer_rank_conversion=True)), ( "MatMulToTransposeConvTranspose", partial(_apply_matmul_to_transpose_conv_transpose, defer_rank_conversion=True), diff --git a/tests/unit/onnx/graph_surgery/dla_transforms/test_dla_decompose_lstm.py b/tests/unit/onnx/graph_surgery/dla_transforms/test_dla_decompose_lstm.py new file mode 100644 index 00000000000..5fa65172d0e --- /dev/null +++ b/tests/unit/onnx/graph_surgery/dla_transforms/test_dla_decompose_lstm.py @@ -0,0 +1,340 @@ +# SPDX-FileCopyrightText: Copyright (c) 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +# SPDX-FileCopyrightText: Copyright (c) 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Unit tests for :mod:`modelopt.onnx.graph_surgery.dla_transforms.dla_decompose_lstm`. + +"""ORT CPU parity tests for LSTM decomposition (forward / reverse / bidirectional). + +Each parametrized run writes two ONNX files under +``scratch_space/lstm_decompose_test_models/`` (base + decomposed for that direction). +The decomposed file is passed through :func:`onnx.shape_inference.infer_shapes` before save +(when inference succeeds). +""" + +from __future__ import annotations + +import copy +from pathlib import Path + +import numpy as np +import onnx +import pytest +from onnx import TensorProto, helper, numpy_helper + +from modelopt.onnx.graph_surgery.dla_transforms.dla_decompose_lstm import ( + _apply_decompose_lstm, + _fold_lstm_constants, +) + +pytest.importorskip("onnxruntime", reason="onnxruntime required for ORT parity checks") + +_REPO_ROOT = Path(__file__).resolve().parents[5] +# Written on each parametrized run (six files total: three base + three decomposed). +_SAVED_MODELS_DIR = _REPO_ROOT / "scratch_space" / "lstm_decompose_test_models" + + +def _save_lstm_test_pair( + base: onnx.ModelProto, transformed: onnx.ModelProto, direction: str +) -> None: + _SAVED_MODELS_DIR.mkdir(parents=True, exist_ok=True) + base_path = _SAVED_MODELS_DIR / f"lstm_base_{direction}.onnx" + dec_path = _SAVED_MODELS_DIR / f"lstm_decomposed_{direction}.onnx" + onnx.save(base, str(base_path)) + try: + from onnxruntime.tools.symbolic_shape_infer import ( + SymbolicShapeInference, # type: ignore[import-untyped] + ) + + inferred = SymbolicShapeInference.infer_shapes(transformed, auto_merge_symbolic_dims=True) + transformed_to_save = inferred if inferred is not None else transformed + except Exception: + transformed_to_save = transformed + onnx.save(transformed_to_save, str(dec_path)) + print( + f"[test_dla_decompose_lstm] saved {base_path.name} and {dec_path.name} under {_SAVED_MODELS_DIR}" + ) + + +@pytest.mark.parametrize( + ("state", "kernel", "expected"), + [([16777216, 1, -16777216], [1, 1, 1], 1), ([0, 0, 0], [-1, -1, -1], 0)], +) +def test_constant_initial_state_convolution_matches_cpp_double_accumulation( + state, kernel, expected +): + values = np.array(state, dtype=np.float32).reshape(1, 3, 1, 1) + weights = np.array(kernel, dtype=np.float32).reshape(1, 3, 1, 1) + graph = helper.make_graph( + [], + "constant_state", + [], + [], + initializer=[ + numpy_helper.from_array(values, "state"), + numpy_helper.from_array(weights, "weights"), + ], + ) + node = helper.make_node("Conv", ["state", "weights"], ["folded"]) + assert _fold_lstm_constants(graph, [node], {"state", "weights"}) == [] + tensors = {tensor.name: numpy_helper.to_array(tensor) for tensor in graph.initializer} + assert tensors["folded"].tobytes() == np.array([expected], dtype=np.float32).tobytes() + np.testing.assert_array_equal(tensors["state"], values) + + +# Some onnxruntime builds reject very new ONNX IR versions from recent onnx packages. +_MAX_IR_VERSION_FOR_ORT = 9 +_OPSET_VERSION = 17 + +_T = 4 +_N = 5 +_D_IN = 3 +_H = 16 + + +def _init(name: str, arr: np.ndarray) -> onnx.TensorProto: + return numpy_helper.from_array(arr, name=name) + + +def _vi(name: str, elem_type: int, shape: list[int | str]) -> onnx.ValueInfoProto: + return helper.make_tensor_value_info(name, elem_type, shape) + + +def _clamp_ir_version(model: onnx.ModelProto) -> None: + model.ir_version = min(model.ir_version, _MAX_IR_VERSION_FOR_ORT) + + +def _lstm_input_shape_summary(model: onnx.ModelProto, input_name: str) -> str | tuple: + """Return a printable shape for an LSTM input (initializer, graph input, or value_info).""" + graph = model.graph + for init in graph.initializer: + if init.name == input_name: + return tuple(numpy_helper.to_array(init).shape) + for vi in list(graph.input) + list(graph.value_info): + if vi.name != input_name: + continue + dims: list[int | str] = [] + for d in vi.type.tensor_type.shape.dim: + if d.dim_value: + dims.append(int(d.dim_value)) + elif d.dim_param: + dims.append(d.dim_param) + else: + dims.append("?") + return tuple(dims) + return f"" + + +def _print_lstm_input_shapes(model: onnx.ModelProto, *, direction: str) -> None: + graph = model.graph + lstm_nodes = [n for n in graph.node if n.op_type == "LSTM"] + if not lstm_nodes: + print(f"[test_dla_decompose_lstm] direction={direction}: no LSTM node in graph") + return + node = lstm_nodes[0] + print(f"[test_dla_decompose_lstm] direction={direction} LSTM name={node.name!r} input slots:") + for idx, inp in enumerate(node.input): + if not inp: + print(f" [{idx}] (empty optional)") + continue + sh = _lstm_input_shape_summary(model, inp) + print(f" [{idx}] {inp!r} shape={sh}") + + +def _make_lstm_model(direction: str) -> onnx.ModelProto: + """Single-LSTM graph: X input, W/R/B initializers, Y + Y_h + Y_c outputs.""" + assert direction in ("forward", "reverse", "bidirectional") + num_dir = 2 if direction == "bidirectional" else 1 + + seeds = {"forward": 42, "reverse": 43, "bidirectional": 44} + rng = np.random.default_rng(seeds[direction]) + + w = rng.standard_normal((num_dir, 4 * _H, _D_IN)).astype(np.float32) + r = rng.standard_normal((num_dir, 4 * _H, _H)).astype(np.float32) + b = rng.standard_normal((num_dir, 8 * _H)).astype(np.float32) + + inits = [_init("W", w), _init("R", r), _init("B", b)] + nodes = [ + helper.make_node( + "LSTM", + ["X", "W", "R", "B"], + ["Y", "Y_h", "Y_c"], + name="lstm0", + hidden_size=_H, + direction=direction, + ) + ] + inputs = [_vi("X", TensorProto.FLOAT, [_T, _N, _D_IN])] + outputs = [ + _vi("Y", TensorProto.FLOAT, [_T, num_dir, _N, _H]), + _vi("Y_h", TensorProto.FLOAT, [num_dir, _N, _H]), + _vi("Y_c", TensorProto.FLOAT, [num_dir, _N, _H]), + ] + graph = helper.make_graph(nodes, f"lstm_{direction}", inputs, outputs, initializer=inits) + model = helper.make_model(graph, opset_imports=[helper.make_opsetid("", _OPSET_VERSION)]) + _clamp_ir_version(model) + onnx.checker.check_model(model) + return model + + +def _single_lstm_node(model: onnx.ModelProto) -> onnx.NodeProto: + return next(n for n in model.graph.node if n.op_type == "LSTM") + + +def test_downstream_transpose_keeps_its_nontrivial_permutation(): + model = _make_lstm_model("bidirectional") + model.graph.node.append( + helper.make_node("Transpose", ["Y"], ["permuted"], perm=[0, 2, 1, 3], name="downstream") + ) + model.graph.output[0].CopyFrom(_vi("permuted", TensorProto.FLOAT, [_T, _N, 2, _H])) + feeds = {"X": np.random.default_rng(8).standard_normal((_T, _N, _D_IN)).astype(np.float32)} + expected = _run_ort(model, feeds) + candidate = onnx.shape_inference.infer_shapes(copy.deepcopy(model)) + _apply_decompose_lstm(candidate) + actual = _run_ort(candidate, feeds) + for name in expected: + np.testing.assert_allclose(actual[name], expected[name], rtol=1e-5, atol=1e-5) + + +def _expect_unsupported_lstm(model: onnx.ModelProto, expected: str) -> None: + from modelopt.onnx.graph_surgery.utils.graph_utils import infer_shapes + + model = infer_shapes(copy.deepcopy(model)) + with pytest.raises(ValueError, match=expected): + _apply_decompose_lstm(model) + + +def _run_ort(model: onnx.ModelProto, feeds: dict[str, np.ndarray]) -> dict[str, np.ndarray]: + import onnxruntime as ort # type: ignore[import-not-found] + + _clamp_ir_version(model) + sess = ort.InferenceSession( + model.SerializeToString(), + providers=["CPUExecutionProvider"], + ) + out: dict[str, np.ndarray] = {} + for o in model.graph.output: + out[o.name] = sess.run([o.name], feeds)[0] + return out + + +def _cosine_similarity(a: np.ndarray, b: np.ndarray) -> float: + a64 = np.asarray(a, dtype=np.float64).ravel() + b64 = np.asarray(b, dtype=np.float64).ravel() + na = float(np.linalg.norm(a64)) + nb = float(np.linalg.norm(b64)) + if na == 0.0 and nb == 0.0: + return 1.0 + if na == 0.0 or nb == 0.0: + return 0.0 + return float(np.dot(a64, b64) / (na * nb)) + + +def _assert_outputs_match(ref: dict[str, np.ndarray], got: dict[str, np.ndarray]) -> None: + assert ref.keys() == got.keys() + for name in ref: + a, b = ref[name], got[name] + assert a.shape == b.shape, f"{name}: shape {a.shape} vs {b.shape}" + cos = _cosine_similarity(a, b) + assert cos >= 1.0 - 1e-5, f"{name}: cosine similarity {cos} expected ~1.0" + np.testing.assert_allclose(a, b, rtol=1e-4, atol=1e-4, err_msg=f"{name}: value mismatch") + + +@pytest.mark.parametrize("direction", ["forward", "reverse", "bidirectional"]) +def test_lstm_decompose_ort_cpu_parity(direction: str) -> None: + """Decomposed graph matches reference LSTM on CPU EP (cosine ~1, tight allclose).""" + model = _make_lstm_model(direction) + _print_lstm_input_shapes(model, direction=direction) + rng = np.random.default_rng(7) + x = rng.standard_normal((_T, _N, _D_IN)).astype(np.float32) + feeds = {"X": x} + + ref = _run_ort(model, feeds) + + dec = copy.deepcopy(model) + from modelopt.onnx.graph_surgery.utils.graph_utils import infer_shapes + + dec = infer_shapes(dec) + dec = _apply_decompose_lstm(dec) + _clamp_ir_version(dec) + onnx.checker.check_model(dec) + + got = _run_ort(dec, feeds) + _save_lstm_test_pair(model, dec, direction) + _assert_outputs_match(ref, got) + + +@pytest.mark.parametrize( + ("attr_name", "attr_value", "expected"), + [ + ("layout", 1, "layout=1"), + ("clip", 0.5, "clip=0.5"), + ("input_forget", 1, "input_forget=1"), + ], +) +def test_lstm_decompose_rejects_unsupported_scalar_attrs(attr_name, attr_value, expected) -> None: + """Unsupported scalar LSTM attributes must fail before decomposition.""" + model = _make_lstm_model("forward") + node = _single_lstm_node(model) + if isinstance(attr_value, float): + node.attribute.append(helper.make_attribute(attr_name, attr_value)) + else: + node.attribute.append(helper.make_attribute(attr_name, int(attr_value))) + + _expect_unsupported_lstm(model, expected) + + +def test_lstm_decompose_rejects_peephole_input() -> None: + """Peephole input P is not represented by this decomposition.""" + model = _make_lstm_model("forward") + node = _single_lstm_node(model) + p = np.zeros((1, 3 * _H), dtype=np.float32) + model.graph.initializer.append(_init("P", p)) + del node.input[:] + node.input.extend(["X", "W", "R", "B", "", "", "", "P"]) + + _expect_unsupported_lstm(model, "peephole P") + + +def test_lstm_decompose_rejects_sequence_lens_input() -> None: + """sequence_lens changes sequence semantics and is not handled here.""" + model = _make_lstm_model("forward") + node = _single_lstm_node(model) + seq_lens = np.array([_T] * _N, dtype=np.int32) + model.graph.initializer.append(_init("seq_lens", seq_lens)) + del node.input[:] + node.input.extend(["X", "W", "R", "B", "seq_lens"]) + + _expect_unsupported_lstm(model, "sequence_lens") + + +def test_lstm_decompose_rejects_custom_activations() -> None: + """Only the default Sigmoid/Tanh/Tanh LSTM activation set is supported.""" + model = _make_lstm_model("forward") + node = _single_lstm_node(model) + node.attribute.append(helper.make_attribute("activations", ["Relu", "Tanh", "Tanh"])) + + _expect_unsupported_lstm(model, "activations") + + +def test_lstm_decompose_rejects_activation_params() -> None: + """Custom activation parameters imply non-default activation semantics.""" + model = _make_lstm_model("forward") + node = _single_lstm_node(model) + node.attribute.append(helper.make_attribute("activation_alpha", [0.1])) + + _expect_unsupported_lstm(model, "activation_alpha") From 3e5d0d63eb6414cf774510e249dadddadf7ee277 Mon Sep 17 00:00:00 2001 From: mgohil Date: Mon, 26 Oct 2026 21:40:31 +0000 Subject: [PATCH 2/6] Preserve LSTM optional inputs and shared tensor contracts Signed-off-by: mgohil --- .../dla_transforms/dla_decompose_lstm.py | 309 ++++++++---------- .../dla_transforms/test_dla_decompose_lstm.py | 284 +++++++++++----- 2 files changed, 346 insertions(+), 247 deletions(-) diff --git a/modelopt/onnx/graph_surgery/dla_transforms/dla_decompose_lstm.py b/modelopt/onnx/graph_surgery/dla_transforms/dla_decompose_lstm.py index 07e38d8b0b9..35eee9f5dcc 100644 --- a/modelopt/onnx/graph_surgery/dla_transforms/dla_decompose_lstm.py +++ b/modelopt/onnx/graph_surgery/dla_transforms/dla_decompose_lstm.py @@ -47,13 +47,16 @@ get_constant_by_name, is_const_dq_input, ) +from .dla_constants_to_initializers import _add_legacy_initializer_inputs +from .dla_graph_cleanup import _subgraph_captures # Promote ONNX LSTM inputs listed at https://onnx.ai/onnx/operators/onnx__LSTM.html (see module docstring). -# Initializers / DQ weights: numpy expand-to-4D. Graph tensors: Unsqueeze (axis 0 for typical 3-D X / states). +# Preserve DQ layouts by promoting their outputs; fixed initializer views can use metadata alone. _LSTM_PROMOTE_INPUT_INDICES = (0, 1, 2, 3, 5, 6) def _expand_array_to_rank4(arr: np.ndarray) -> np.ndarray: + """Prepend singleton dimensions for the LSTM convolution layout.""" rank = len(arr.shape) if rank >= 4: return arr @@ -75,7 +78,7 @@ def _validate_lstm_supported(node: onnx.NodeProto) -> str: direction = attr.s.decode("utf-8").strip().lower() elif attr.name == "layout" and int(attr.i) != 0: unsupported.append(f"layout={int(attr.i)}") - elif attr.name == "clip" and float(attr.f) != 0.0: + elif attr.name == "clip": unsupported.append(f"clip={float(attr.f)}") elif attr.name == "input_forget" and int(attr.i) != 0: unsupported.append("input_forget=1") @@ -108,7 +111,7 @@ def _promote_lstm_input_to_rank4( """Promote one LSTM input slot to rank 4 for the decomposition path. Shapes follow ONNX ``LSTM`` (layout 0): https://onnx.ai/onnx/operators/onnx__LSTM.html - Initializers and DQ-backed weights are expanded in numpy; runtime tensors get ``Unsqueeze``. + Fixed private views use initializer metadata; DQ and runtime values use ``Unsqueeze``. """ if input_index >= len(lstm_node.input): return @@ -116,47 +119,21 @@ def _promote_lstm_input_to_rank4( if not tensor_name: return - init_list = list(cache.init_names) - dq_init_name = is_const_dq_input(init_list, tensor_name, graph) - - if dq_init_name is not None: - deq_node = next( - ( - n - for n in graph.node - if n.op_type == "DequantizeLinear" - and n.output[0] == tensor_name - and n.input[0] == dq_init_name - ), - None, - ) - if deq_node is None: - return - initializer = cache.get_init(dq_init_name) - if initializer is None: - return - init_arr = numpy_helper.to_array(initializer) - init_arr = _expand_array_to_rank4(init_arr) - new_init = numpy_helper.from_array( - init_arr, name=initializer.name + f"_lstm4d_{lstm_node.name}_{input_index}" - ) - add_unique_initializers(graph, unique_initializer, [new_init]) - graph.initializer.remove(initializer) - deq_node.input[0] = new_init.name - return - - if cache.is_init(tensor_name): + if cache.is_init(tensor_name) and tensor_name not in cache.graph_input_names: initializer = cache.get_init(tensor_name) if initializer is None: return - init_arr = numpy_helper.to_array(initializer) - new_arr = _expand_array_to_rank4(init_arr) - new_init = numpy_helper.from_array( - new_arr, name=initializer.name + f"_lstm4d_{lstm_node.name}_{input_index}" + new_init = graph.initializer.add() + new_init.CopyFrom(initializer) + new_init.name = cache.unique_name( + initializer.name + f"_lstm4d_{lstm_node.name}_{input_index}" ) - add_unique_initializers(graph, unique_initializer, [new_init]) + original_dims = list(new_init.dims) + del new_init.dims[:] + new_init.dims.extend([1] * max(0, 4 - len(original_dims)) + original_dims) + unique_initializer.add(new_init.name) refs = sum(1 for n in graph.node for inp in n.input if inp == tensor_name) - if refs <= 1: + if refs <= 1 and tensor_name not in cache.graph_output_names | _subgraph_captures(graph): graph.initializer.remove(initializer) lstm_node.input[input_index] = new_init.name return @@ -168,15 +145,15 @@ def _promote_lstm_input_to_rank4( return axes_list = [0] if shape is None or len(shape) == 3 else list(range(4 - len(shape))) - axes_name = f"{tensor_name}_lstm4d_axes_{lstm_node.name}_{input_index}" - out_name = f"{tensor_name}_lstm4d_{lstm_node.name}_{input_index}" + axes_name = cache.unique_name(f"{tensor_name}_lstm4d_axes_{lstm_node.name}_{input_index}") + out_name = cache.unique_name(f"{tensor_name}_lstm4d_{lstm_node.name}_{input_index}") axes_init = numpy_helper.from_array(np.array(axes_list, dtype=np.int64), name=axes_name) add_unique_initializers(graph, unique_initializer, [axes_init]) unsq = helper.make_node( "Unsqueeze", [tensor_name, axes_name], [out_name], - name=f"lstm_prep_unsqueeze_{lstm_node.name}_{input_index}", + name=cache.unique_name(f"lstm_prep_unsqueeze_{lstm_node.name}_{input_index}"), ) prep_nodes.append(unsq) lstm_node.input[input_index] = out_name @@ -216,10 +193,16 @@ def _fold_lstm_constants(graph, nodes, unique_initializer): and arrays[1].shape[1:] == (arrays[0].shape[1], 1, 1) ): # The reference accumulates these dot products sequentially in double precision. - products = arrays[1].reshape(arrays[1].shape[0], -1).astype(np.float64) - products = products * arrays[0].reshape(-1).astype(np.float64) - terms = np.concatenate([np.zeros((products.shape[0], 1)), products], axis=1) - values = np.cumsum(terms, axis=1)[:, -1].astype(arrays[0].dtype) + weights = arrays[1].reshape(arrays[1].shape[0], -1) + state = arrays[0].reshape(-1).astype(np.float64) + values = np.empty(weights.shape[0], dtype=arrays[0].dtype) + rows = max(1, 2**18 // (weights.shape[1] + 1)) + for start in range(0, len(values), rows): + stop = min(start + rows, len(values)) + terms = np.zeros((stop - start, weights.shape[1] + 1), dtype=np.float64) + np.multiply(weights[start:stop], state, out=terms[:, 1:]) + np.cumsum(terms, axis=1, out=terms) + values[start:stop] = terms[:, -1] outputs = [values.reshape(1, arrays[1].shape[0], 1, 1)] else: remaining.append(node) @@ -235,11 +218,28 @@ def _apply_decompose_lstm(model, *, defer_rank_conversion=False): """Decompose ``LSTM`` nodes into basic ONNX ops (IO per ONNX LSTM layout 0).""" graph = model.graph cache = GraphCache(graph) - lstm_nodes = [n for n in list(graph.node) if n.op_type == "LSTM"] + local_functions = { + (function.domain, function.name, function.overload) for function in model.functions + } + lstm_nodes = [ + n + for n in graph.node + if n.op_type == "LSTM" + and n.domain in {"", "ai.onnx"} + and not n.overload + and (n.domain, n.op_type, n.overload) not in local_functions + ] node_idx = {id(n): i for i, n in enumerate(graph.node)} replacements_batch: dict[int, tuple[list, list]] = {} lstm_counter = 0 unique_initializer = set() + allocated_names = {} + + def new_name(proposed): + """Reuse internal controls while allocating names outside the original graph namespace.""" + if proposed not in allocated_names: + allocated_names[proposed] = cache.unique_name(proposed) + return allocated_names[proposed] def add_nodes_for_time_step( time_step, @@ -253,25 +253,26 @@ def add_nodes_for_time_step( shape_info_map=None, ): # add it = sigmoid(Ht * R_i + Z_i) + """Build the gates and state update for one sequence position.""" matmul_h_r_i_node = helper.make_node( "Conv", [h_t, r_t_iofc["i"]], - [f"{prefix}_h_r_i_{node.name}_seq_{time_step}"], - name=f"{prefix}_conv_h_r_i_{node.name}_seq_{time_step}", + [new_name(f"{prefix}_h_r_i_{node.name}_seq_{time_step}")], + name=new_name(f"{prefix}_conv_h_r_i_{node.name}_seq_{time_step}"), ) nodes_to_add.append(matmul_h_r_i_node) add_h_r_i_z_i_node = helper.make_node( "Add", [matmul_h_r_i_node.output[0], z_t_iofc["i"]], - [f"{prefix}_h_r_i_z_i_{node.name}_seq_{time_step}"], - name=f"{prefix}_add_h_r_i_z_i_{node.name}_seq_{time_step}", + [new_name(f"{prefix}_h_r_i_z_i_{node.name}_seq_{time_step}")], + name=new_name(f"{prefix}_add_h_r_i_z_i_{node.name}_seq_{time_step}"), ) nodes_to_add.append(add_h_r_i_z_i_node) i_t_node = helper.make_node( "Sigmoid", [add_h_r_i_z_i_node.output[0]], - [f"{prefix}_i_t_{node.name}_seq_{time_step}"], - name=f"{prefix}_sigmoid_i_t_{node.name}_seq_{time_step}", + [new_name(f"{prefix}_i_t_{node.name}_seq_{time_step}")], + name=new_name(f"{prefix}_sigmoid_i_t_{node.name}_seq_{time_step}"), ) nodes_to_add.append(i_t_node) @@ -279,22 +280,22 @@ def add_nodes_for_time_step( matmul_h_r_f_node = helper.make_node( "Conv", [h_t, r_t_iofc["f"]], - [f"{prefix}_h_r_f_{node.name}_seq_{time_step}"], - name=f"{prefix}_conv_h_r_f_{node.name}_seq_{time_step}", + [new_name(f"{prefix}_h_r_f_{node.name}_seq_{time_step}")], + name=new_name(f"{prefix}_conv_h_r_f_{node.name}_seq_{time_step}"), ) nodes_to_add.append(matmul_h_r_f_node) add_h_r_f_z_f_node = helper.make_node( "Add", [matmul_h_r_f_node.output[0], z_t_iofc["f"]], - [f"{prefix}_h_r_f_z_f_{node.name}_seq_{time_step}"], - name=f"{prefix}_add_h_r_f_z_f_{node.name}_seq_{time_step}", + [new_name(f"{prefix}_h_r_f_z_f_{node.name}_seq_{time_step}")], + name=new_name(f"{prefix}_add_h_r_f_z_f_{node.name}_seq_{time_step}"), ) nodes_to_add.append(add_h_r_f_z_f_node) ft_node = helper.make_node( "Sigmoid", [add_h_r_f_z_f_node.output[0]], - [f"{prefix}_ft_{node.name}_seq_{time_step}"], - name=f"{prefix}_sigmoid_ft_{node.name}_seq_{time_step}", + [new_name(f"{prefix}_ft_{node.name}_seq_{time_step}")], + name=new_name(f"{prefix}_sigmoid_ft_{node.name}_seq_{time_step}"), ) nodes_to_add.append(ft_node) @@ -302,22 +303,22 @@ def add_nodes_for_time_step( matmul_h_r_o_node = helper.make_node( "Conv", [h_t, r_t_iofc["o"]], - [f"{prefix}_h_r_o_{node.name}_seq_{time_step}"], - name=f"{prefix}_conv_h_r_o_{node.name}_seq_{time_step}", + [new_name(f"{prefix}_h_r_o_{node.name}_seq_{time_step}")], + name=new_name(f"{prefix}_conv_h_r_o_{node.name}_seq_{time_step}"), ) nodes_to_add.append(matmul_h_r_o_node) add_h_r_o_z_o_node = helper.make_node( "Add", [matmul_h_r_o_node.output[0], z_t_iofc["o"]], - [f"{prefix}_h_r_o_z_o_{node.name}_seq_{time_step}"], - name=f"{prefix}_add_h_r_o_z_o_{node.name}_seq_{time_step}", + [new_name(f"{prefix}_h_r_o_z_o_{node.name}_seq_{time_step}")], + name=new_name(f"{prefix}_add_h_r_o_z_o_{node.name}_seq_{time_step}"), ) nodes_to_add.append(add_h_r_o_z_o_node) ot_node = helper.make_node( "Sigmoid", [add_h_r_o_z_o_node.output[0]], - [f"{prefix}_ot_{node.name}_seq_{time_step}"], - name=f"{prefix}_sigmoid_ot_{node.name}_seq_{time_step}", + [new_name(f"{prefix}_ot_{node.name}_seq_{time_step}")], + name=new_name(f"{prefix}_sigmoid_ot_{node.name}_seq_{time_step}"), ) nodes_to_add.append(ot_node) @@ -325,22 +326,22 @@ def add_nodes_for_time_step( matmul_h_r_c_node = helper.make_node( "Conv", [h_t, r_t_iofc["c"]], - [f"{prefix}_h_r_c_{node.name}_seq_{time_step}"], - name=f"{prefix}_conv_h_r_c_{node.name}_seq_{time_step}", + [new_name(f"{prefix}_h_r_c_{node.name}_seq_{time_step}")], + name=new_name(f"{prefix}_conv_h_r_c_{node.name}_seq_{time_step}"), ) nodes_to_add.append(matmul_h_r_c_node) add_h_r_c_z_c_node = helper.make_node( "Add", [matmul_h_r_c_node.output[0], z_t_iofc["c"]], - [f"{prefix}_h_r_c_z_c_{node.name}_seq_{time_step}"], - name=f"{prefix}_add_h_r_c_z_c_{node.name}_seq_{time_step}", + [new_name(f"{prefix}_h_r_c_z_c_{node.name}_seq_{time_step}")], + name=new_name(f"{prefix}_add_h_r_c_z_c_{node.name}_seq_{time_step}"), ) nodes_to_add.append(add_h_r_c_z_c_node) ct_node = helper.make_node( "Tanh", [add_h_r_c_z_c_node.output[0]], - [f"{prefix}_ct_{node.name}_seq_{time_step}"], - name=f"{prefix}_tanh_ct_{node.name}_seq_{time_step}", + [new_name(f"{prefix}_ct_{node.name}_seq_{time_step}")], + name=new_name(f"{prefix}_tanh_ct_{node.name}_seq_{time_step}"), ) nodes_to_add.append(ct_node) @@ -348,22 +349,22 @@ def add_nodes_for_time_step( mul_ft_ct_1_node = helper.make_node( "Mul", [ft_node.output[0], c_t], - [f"{prefix}_ft_ct_1_{node.name}_seq_{time_step}"], - name=f"{prefix}_mul_ft_ct_1_{node.name}_seq_{time_step}", + [new_name(f"{prefix}_ft_ct_1_{node.name}_seq_{time_step}")], + name=new_name(f"{prefix}_mul_ft_ct_1_{node.name}_seq_{time_step}"), ) nodes_to_add.append(mul_ft_ct_1_node) mul_it_ct_node = helper.make_node( "Mul", [i_t_node.output[0], ct_node.output[0]], - [f"{prefix}_it_ct_{node.name}_seq_{time_step}"], - name=f"{prefix}_mul_it_ct_{node.name}_seq_{time_step}", + [new_name(f"{prefix}_it_ct_{node.name}_seq_{time_step}")], + name=new_name(f"{prefix}_mul_it_ct_{node.name}_seq_{time_step}"), ) nodes_to_add.append(mul_it_ct_node) add_ft_ct_1_it_ct_node = helper.make_node( "Add", [mul_ft_ct_1_node.output[0], mul_it_ct_node.output[0]], - [f"{prefix}_C_t_{node.name}_seq_{time_step}"], - name=f"{prefix}_add_ft_ct_1_it_ct_{node.name}_seq_{time_step}", + [new_name(f"{prefix}_C_t_{node.name}_seq_{time_step}")], + name=new_name(f"{prefix}_add_ft_ct_1_it_ct_{node.name}_seq_{time_step}"), ) nodes_to_add.append(add_ft_ct_1_it_ct_node) @@ -371,54 +372,24 @@ def add_nodes_for_time_step( tanh_ct_node = helper.make_node( "Tanh", [add_ft_ct_1_it_ct_node.output[0]], - [f"{prefix}_tanh_Ct_{node.name}_seq_{time_step}"], - name=f"{prefix}_tanh_Ct_{node.name}_seq_{time_step}", + [new_name(f"{prefix}_tanh_Ct_{node.name}_seq_{time_step}")], + name=new_name(f"{prefix}_tanh_Ct_{node.name}_seq_{time_step}"), ) nodes_to_add.append(tanh_ct_node) mul_ot_tanh_ct_node = helper.make_node( "Mul", [ot_node.output[0], tanh_ct_node.output[0]], - [f"{prefix}_H_t_{node.name}_seq_{time_step}"], - name=f"{prefix}_mul_ot_tanh_ct_{node.name}_seq_{time_step}", + [new_name(f"{prefix}_H_t_{node.name}_seq_{time_step}")], + name=new_name(f"{prefix}_mul_ot_tanh_ct_{node.name}_seq_{time_step}"), ) nodes_to_add.append(mul_ot_tanh_ct_node) return mul_ot_tanh_ct_node.output[0], add_ft_ct_1_it_ct_node.output[0] def transpose_tensor(graph, tensor_name, perm, nodes_to_add, unique_initializer): + """Transpose a private initializer or preserve runtime data through a node.""" init_list = [init.name for init in graph.initializer] - if is_const_dq_input(init_list, tensor_name, graph) is not None: - init_name = is_const_dq_input(init_list, tensor_name, graph) - deq_node = next( - ( - n - for n in graph.node - if n.op_type == "DequantizeLinear" - and n.output[0] == tensor_name - and n.input[0] == init_name - ), - None, - ) - if deq_node is None: - raise ValueError( - f"LSTM transpose_tensor: DequantizeLinear for DQ input {tensor_name!r} (init {init_name!r}) " - "not found in graph." - ) - initializer = cache.get_init(init_name) - if initializer is None: - raise ValueError( - f"LSTM transpose_tensor: initializer {init_name!r} missing for DQ tensor {tensor_name!r}." - ) - init_arr = numpy_helper.to_array(initializer) - init_arr = _expand_array_to_rank4(init_arr) - - init_arr = np.transpose(init_arr, axes=perm) - new_init = numpy_helper.from_array(init_arr, name=initializer.name + "_transposed") - add_unique_initializers(graph, unique_initializer, [new_init]) - graph.initializer.remove(initializer) - deq_node.input[0] = new_init.name - return tensor_name - elif tensor_name in init_list: + if tensor_name in init_list: initializer = cache.get_init(tensor_name) if initializer is None: raise ValueError( @@ -427,7 +398,9 @@ def transpose_tensor(graph, tensor_name, perm, nodes_to_add, unique_initializer) init_arr = numpy_helper.to_array(initializer) init_arr = _expand_array_to_rank4(init_arr) init_arr = np.transpose(init_arr, axes=perm) - new_init = numpy_helper.from_array(init_arr, name=initializer.name + "_transposed") + new_init = numpy_helper.from_array( + init_arr, name=new_name(initializer.name + "_transposed") + ) add_unique_initializers(graph, unique_initializer, [new_init]) graph.initializer.remove(initializer) return new_init.name @@ -436,7 +409,7 @@ def transpose_tensor(graph, tensor_name, perm, nodes_to_add, unique_initializer) "Transpose", inputs=[tensor_name], outputs=[tensor_name + "_transposed"], - name=tensor_name + "_transpose", + name=new_name(tensor_name + "_transpose"), perm=perm, ) nodes_to_add.append(transpose_node) @@ -448,6 +421,7 @@ def transpose_inputs(graph, node, unique_initializer, shape_info_map, nodes_to_a # When batch_size==1 the batch dim is already a leading unary dim after # 4D promotion. Use [0,3,2,1] to keep it in place instead of swapping # it to the back with [2,3,0,1]. + """Prepare activation, weight, bias and initial-state convolution layouts.""" transpose_perm1 = [0, 3, 2, 1] if shape_info_map.get("N", 0) == 1 else [2, 3, 0, 1] transpose_perm2 = [2, 3, 1, 0] transpose_perm3 = [0, 3, 2, 1] @@ -473,15 +447,15 @@ def transpose_inputs(graph, node, unique_initializer, shape_info_map, nodes_to_a bw_br_split_init = numpy_helper.from_array( np.array([4 * shape_info_map["H"], 4 * shape_info_map["H"]], dtype=np.int64), - name=f"bw_br_split_init_{node.name}", + name=new_name(f"bw_br_split_init_{node.name}"), ) add_unique_initializers(graph, unique_initializer, [bw_br_split_init]) bw_br_split_node = helper.make_node( "Split", [b_transposed, bw_br_split_init.name], - [f"bw_{node.name}", f"br_{node.name}"], + [new_name(f"bw_{node.name}"), new_name(f"br_{node.name}")], axis=1, - name=f"split_B_{node.name}_axis_updated", + name=new_name(f"split_B_{node.name}_axis_updated"), ) nodes_to_add.append(bw_br_split_node) bw_transposed = bw_br_split_node.output[0] @@ -498,7 +472,7 @@ def transpose_inputs(graph, node, unique_initializer, shape_info_map, nodes_to_a [shape_info_map["N"], shape_info_map["H"], shape_info_map["d"], 1], dtype=np.float32, ), - name=node.name + "/ih", + name=new_name(node.name + "/ih"), ) add_unique_initializers(graph, unique_initializer, [ih_init]) ih_transposed = ih_init.name @@ -513,20 +487,20 @@ def transpose_inputs(graph, node, unique_initializer, shape_info_map, nodes_to_a [shape_info_map["N"], shape_info_map["H"], shape_info_map["d"], 1], dtype=np.float32, ), - name=node.name + "/ic", + name=new_name(node.name + "/ic"), ) add_unique_initializers(graph, unique_initializer, [ic_init]) ic_transposed = ic_init.name if bw_transposed is None and br_transposed is None: # Optional B (ONNX LSTM): zero biases broadcast with Add(bw, br) and Add(..., conv_out). - zscalar = np.zeros((1, 1, 1, 1), dtype=np.float32) + zscalar = np.zeros((1, 1, shape_info_map["d"], 1), dtype=np.float32) bw_init = numpy_helper.from_array( zscalar, - name=(node.name + "/zero_bw").replace("/", "_"), + name=new_name((node.name + "/zero_bw").replace("/", "_")), ) br_init = numpy_helper.from_array( np.array(zscalar, copy=True), - name=(node.name + "/zero_br").replace("/", "_"), + name=new_name((node.name + "/zero_br").replace("/", "_")), ) add_unique_initializers(graph, unique_initializer, [bw_init, br_init]) bw_transposed = bw_init.name @@ -561,7 +535,7 @@ def resolve_lstm_tensor_shape(tensor_name, _seen=None): init = cache.get_init(tensor_name) if init is None: return None - return list(numpy_helper.to_array(init).shape) + return list(init.dims) init_names = list(cache.init_names) dq_src = is_const_dq_input(init_names, tensor_name, graph) @@ -573,7 +547,7 @@ def resolve_lstm_tensor_shape(tensor_name, _seen=None): return None if prod.op_type == "Unsqueeze" and len(prod.input) >= 2: inner = prod.input[0] - axes_name = prod.input[1] + axes_name = new_name(prod.input[1]) inner_shape = resolve_lstm_tensor_shape(inner, _seen) if inner_shape is None: return None @@ -596,6 +570,7 @@ def resolve_lstm_tensor_shape(tensor_name, _seen=None): return None def create_shape_info_map(node): + """Resolve the fixed dimensions required for sequence unrolling.""" x, w = node.input[:2] x_shape = resolve_lstm_tensor_shape(x) w_shape = resolve_lstm_tensor_shape(w) @@ -626,12 +601,13 @@ def create_shape_info_map(node): def split_direction_wise( graph, node, unique_initializer, shape_info_map, transposed_inputs, nodes_to_add ): + """Separate the forward and reverse convolution inputs.""" if shape_info_map["d"] == 1: # Unidirectional: reverse and forward paths use the same tensors (only timestep order differs). return [transposed_inputs, transposed_inputs] direction_split_init = numpy_helper.from_array( - np.array([1, 1], dtype=np.int64), name=f"direction_split_init_{node.name}" + np.array([1, 1], dtype=np.int64), name=new_name(f"direction_split_init_{node.name}") ) add_unique_initializers(graph, unique_initializer, [direction_split_init]) @@ -651,11 +627,11 @@ def split_direction_wise( "Split", [str(input_value), direction_split_init.name], [ - f"forward_direction_split_{input_name}_{node.name}", - f"reverse_direction_split_{input_name}_{node.name}", + new_name(f"forward_direction_split_{input_name}_{node.name}"), + new_name(f"reverse_direction_split_{input_name}_{node.name}"), ], axis=2, - name=f"split_direction_wise_{input_name}_{node.name}_axis_updated", + name=new_name(f"split_direction_wise_{input_name}_{node.name}_axis_updated"), ) nodes_to_add.append(direction_split_node) forward_list[input_name] = direction_split_node.output[0] @@ -665,6 +641,7 @@ def split_direction_wise( def precompute_and_split_seq_len( graph, node, unique_initializer, shape_info_map, input_list, nodes_to_add, direction, prefix ): + """Compute input contributions and split them by sequence position.""" x_input = input_list["x"] w_input = input_list["w"] bw_input = input_list["bw"] @@ -682,24 +659,24 @@ def precompute_and_split_seq_len( add_bw_br_node = helper.make_node( "Add", [str(bw_input), str(br_input)], - [f"{prefix}_add_bw_br_{node.name}"], - name=f"{prefix}_add_bw_br_{node.name}", + [new_name(f"{prefix}_add_bw_br_{node.name}")], + name=new_name(f"{prefix}_add_bw_br_{node.name}"), ) nodes_to_add.append(add_bw_br_node) conv_x_w_node = helper.make_node( "Conv", [str(x_input), str(w_input)], - [f"{prefix}_conv_x_w_{node.name}"], - name=f"{prefix}_conv_x_w_{node.name}", + [new_name(f"{prefix}_conv_x_w_{node.name}")], + name=new_name(f"{prefix}_conv_x_w_{node.name}"), ) nodes_to_add.append(conv_x_w_node) z_node = helper.make_node( "Add", [add_bw_br_node.output[0], conv_x_w_node.output[0]], - [f"{prefix}_add_bw_br_conv_x_w_{node.name}"], - name=f"{prefix}_add_bw_br_conv_x_w_{node.name}", + [new_name(f"{prefix}_add_bw_br_conv_x_w_{node.name}")], + name=new_name(f"{prefix}_add_bw_br_conv_x_w_{node.name}"), ) nodes_to_add.append(z_node) @@ -713,26 +690,26 @@ def precompute_and_split_seq_len( ], dtype=np.int64, ), - name=f"split_z_iofc_init_{node.name}", + name=new_name(f"split_z_iofc_init_{node.name}"), ) add_unique_initializers(graph, unique_initializer, [split_z_iofc_init]) split_z_iofc_node = helper.make_node( "Split", [z_node.output[0], split_z_iofc_init.name], [ - f"{prefix}_split_z_i_{node.name}", - f"{prefix}_split_z_o_{node.name}", - f"{prefix}_split_z_f_{node.name}", - f"{prefix}_split_z_c_{node.name}", + new_name(f"{prefix}_split_z_i_{node.name}"), + new_name(f"{prefix}_split_z_o_{node.name}"), + new_name(f"{prefix}_split_z_f_{node.name}"), + new_name(f"{prefix}_split_z_c_{node.name}"), ], axis=1, - name=f"{prefix}_split_z_iofc_{node.name}_axis_updated", + name=new_name(f"{prefix}_split_z_iofc_{node.name}_axis_updated"), ) nodes_to_add.append(split_z_iofc_node) split_seq_init = numpy_helper.from_array( np.array([1] * shape_info_map["T"], dtype=np.int64), - name=f"split_z_iofc_seq_init_{node.name}", + name=new_name(f"split_z_iofc_seq_init_{node.name}"), ) add_unique_initializers(graph, unique_initializer, [split_seq_init]) @@ -744,11 +721,13 @@ def precompute_and_split_seq_len( "Split", [split_z_iofc_output, split_seq_init.name], [ - f"{prefix}_split_z_{split_z_iofc_name}_{time_step}_{node.name}" + new_name(f"{prefix}_split_z_{split_z_iofc_name}_{time_step}_{node.name}") for time_step in range(shape_info_map["T"]) ], axis=3, - name=f"{prefix}_split_z_iofc_seq_{split_z_iofc_name}_{node.name}_axis_updated", + name=new_name( + f"{prefix}_split_z_iofc_seq_{split_z_iofc_name}_{node.name}_axis_updated" + ), ) nodes_to_add.append(split_seq_node) split_z_iofc_seq[split_z_iofc_name] = (split_seq_node.output)[::direction] @@ -758,6 +737,7 @@ def precompute_and_split_seq_len( def split_r_iofc( graph, node, unique_initializer, shape_info_map, input_list, nodes_to_add, prefix ): + """Split the recurrent kernels into the four gate weights.""" r_input = input_list["r"] if not r_input: raise ValueError( @@ -773,20 +753,20 @@ def split_r_iofc( ], dtype=np.int64, ), - name=f"split_r_iofc_init_{node.name}", + name=new_name(f"split_r_iofc_init_{node.name}"), ) add_unique_initializers(graph, unique_initializer, [split_r_iofc_init]) split_r_iofc_node = helper.make_node( "Split", [str(r_input), split_r_iofc_init.name], [ - f"{prefix}_split_r_i_{node.name}", - f"{prefix}_split_r_o_{node.name}", - f"{prefix}_split_r_f_{node.name}", - f"{prefix}_split_r_c_{node.name}", + new_name(f"{prefix}_split_r_i_{node.name}"), + new_name(f"{prefix}_split_r_o_{node.name}"), + new_name(f"{prefix}_split_r_f_{node.name}"), + new_name(f"{prefix}_split_r_c_{node.name}"), ], axis=0, - name=f"{prefix}_split_r_iofc_{node.name}_axis_updated", + name=new_name(f"{prefix}_split_r_iofc_{node.name}_axis_updated"), ) nodes_to_add.append(split_r_iofc_node) return { @@ -808,7 +788,7 @@ def transpose_outputs( return None transpose_perm = [3, 2, 0, 1] if squeeze_axes: - transposed_tmp = f"{output_name}_lstm_tr_{node.name}".replace("/", "_") + transposed_tmp = new_name(f"{output_name}_lstm_tr_{node.name}".replace("/", "_")) transpose_out = transposed_tmp else: transpose_out = output_name @@ -821,13 +801,13 @@ def transpose_outputs( "Transpose", [str(input_name)], [str(transpose_out)], - name=f"{output_name}_transpose_{node.name}", + name=new_name(f"{output_name}_transpose_{node.name}"), perm=transpose_perm, ) nodes_to_add.append(transpose_node) if squeeze_axes: - axes_name = f"{output_name}_lstm_squeeze_axes_{node.name}".replace("/", "_") + axes_name = new_name(f"{output_name}_lstm_squeeze_axes_{node.name}".replace("/", "_")) axes_init = numpy_helper.from_array( np.array(squeeze_axes, dtype=np.int64), name=axes_name ) @@ -836,7 +816,7 @@ def transpose_outputs( "Squeeze", [transpose_out, axes_name], [output_name], - name=f"{output_name}_lstm_squeeze_{node.name}", + name=new_name(f"{output_name}_lstm_squeeze_{node.name}"), ) nodes_to_add.append(squeeze_node) return squeeze_node.output[0] @@ -972,8 +952,8 @@ def transpose_outputs( h_t_concat_node = helper.make_node( "Concat", h_t_list, - [f"{prefix}_h_t_concat_{node.name}"], - name=f"{prefix}_concat_h_t_concat_{node.name}_axis_updated", + [new_name(f"{prefix}_h_t_concat_{node.name}")], + name=new_name(f"{prefix}_concat_h_t_concat_{node.name}_axis_updated"), axis=3, ) nodes_to_add.append(h_t_concat_node) @@ -1024,8 +1004,8 @@ def transpose_outputs( h_t_concat_node = helper.make_node( "Concat", h_t_list[::-1], - [f"{prefix}_h_t_concat_{node.name}"], - name=f"{prefix}_concat_h_t_concat_{node.name}_axis_updated", + [new_name(f"{prefix}_h_t_concat_{node.name}")], + name=new_name(f"{prefix}_concat_h_t_concat_{node.name}_axis_updated"), axis=3, ) nodes_to_add.append(h_t_concat_node) @@ -1046,8 +1026,8 @@ def transpose_outputs( y_concat_node = helper.make_node( "Concat", [str(y_forward), str(y_reverse)], - [f"y_concat_{node.name}"], - name=f"concat_y_concat_{node.name}_axis_updated", + [new_name(f"y_concat_{node.name}")], + name=new_name(f"concat_y_concat_{node.name}_axis_updated"), axis=2, ) nodes_to_add.append(y_concat_node) @@ -1062,8 +1042,8 @@ def transpose_outputs( y_h_concat_node = helper.make_node( "Concat", [str(y_h_forward), str(y_h_reverse)], - [f"y_h_concat_{node.name}"], - name=f"concat_y_h_concat_{node.name}_axis_updated", + [new_name(f"y_h_concat_{node.name}")], + name=new_name(f"concat_y_h_concat_{node.name}_axis_updated"), axis=2, ) nodes_to_add.append(y_h_concat_node) @@ -1078,8 +1058,8 @@ def transpose_outputs( y_c_concat_node = helper.make_node( "Concat", [str(y_c_forward), str(y_c_reverse)], - [f"y_c_concat_{node.name}"], - name=f"concat_y_c_concat_{node.name}_axis_updated", + [new_name(f"y_c_concat_{node.name}")], + name=new_name(f"concat_y_c_concat_{node.name}_axis_updated"), axis=2, ) nodes_to_add.append(y_c_concat_node) @@ -1125,6 +1105,7 @@ def transpose_outputs( lstm_counter += 1 batch_replace_nodes(graph, replacements_batch) + _add_legacy_initializer_inputs(model) logger.debug("Updated %d LSTM nodes with manual decomposition", lstm_counter) return model diff --git a/tests/unit/onnx/graph_surgery/dla_transforms/test_dla_decompose_lstm.py b/tests/unit/onnx/graph_surgery/dla_transforms/test_dla_decompose_lstm.py index 5fa65172d0e..f2a9abb6dbe 100644 --- a/tests/unit/onnx/graph_surgery/dla_transforms/test_dla_decompose_lstm.py +++ b/tests/unit/onnx/graph_surgery/dla_transforms/test_dla_decompose_lstm.py @@ -13,23 +13,13 @@ # See the License for the specific language governing permissions and # limitations under the License. -# SPDX-FileCopyrightText: Copyright (c) 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# SPDX-License-Identifier: Apache-2.0 -# # Unit tests for :mod:`modelopt.onnx.graph_surgery.dla_transforms.dla_decompose_lstm`. -"""ORT CPU parity tests for LSTM decomposition (forward / reverse / bidirectional). - -Each parametrized run writes two ONNX files under -``scratch_space/lstm_decompose_test_models/`` (base + decomposed for that direction). -The decomposed file is passed through :func:`onnx.shape_inference.infer_shapes` before save -(when inference succeeds). -""" +"""ORT CPU parity tests for LSTM decomposition (forward / reverse / bidirectional).""" from __future__ import annotations import copy -from pathlib import Path import numpy as np import onnx @@ -40,34 +30,9 @@ _apply_decompose_lstm, _fold_lstm_constants, ) +from modelopt.onnx.graph_surgery.utils.graph_utils import infer_shapes -pytest.importorskip("onnxruntime", reason="onnxruntime required for ORT parity checks") - -_REPO_ROOT = Path(__file__).resolve().parents[5] -# Written on each parametrized run (six files total: three base + three decomposed). -_SAVED_MODELS_DIR = _REPO_ROOT / "scratch_space" / "lstm_decompose_test_models" - - -def _save_lstm_test_pair( - base: onnx.ModelProto, transformed: onnx.ModelProto, direction: str -) -> None: - _SAVED_MODELS_DIR.mkdir(parents=True, exist_ok=True) - base_path = _SAVED_MODELS_DIR / f"lstm_base_{direction}.onnx" - dec_path = _SAVED_MODELS_DIR / f"lstm_decomposed_{direction}.onnx" - onnx.save(base, str(base_path)) - try: - from onnxruntime.tools.symbolic_shape_infer import ( - SymbolicShapeInference, # type: ignore[import-untyped] - ) - - inferred = SymbolicShapeInference.infer_shapes(transformed, auto_merge_symbolic_dims=True) - transformed_to_save = inferred if inferred is not None else transformed - except Exception: - transformed_to_save = transformed - onnx.save(transformed_to_save, str(dec_path)) - print( - f"[test_dla_decompose_lstm] saved {base_path.name} and {dec_path.name} under {_SAVED_MODELS_DIR}" - ) +ort = pytest.importorskip("onnxruntime", reason="onnxruntime required for ORT parity checks") @pytest.mark.parametrize( @@ -77,6 +42,7 @@ def _save_lstm_test_pair( def test_constant_initial_state_convolution_matches_cpp_double_accumulation( state, kernel, expected ): + """Constant initial state convolution matches cpp double accumulation.""" values = np.array(state, dtype=np.float32).reshape(1, 3, 1, 1) weights = np.array(kernel, dtype=np.float32).reshape(1, 3, 1, 1) graph = helper.make_graph( @@ -107,54 +73,20 @@ def test_constant_initial_state_convolution_matches_cpp_double_accumulation( def _init(name: str, arr: np.ndarray) -> onnx.TensorProto: + """Init.""" return numpy_helper.from_array(arr, name=name) def _vi(name: str, elem_type: int, shape: list[int | str]) -> onnx.ValueInfoProto: + """Vi.""" return helper.make_tensor_value_info(name, elem_type, shape) def _clamp_ir_version(model: onnx.ModelProto) -> None: + """Clamp ir version.""" model.ir_version = min(model.ir_version, _MAX_IR_VERSION_FOR_ORT) -def _lstm_input_shape_summary(model: onnx.ModelProto, input_name: str) -> str | tuple: - """Return a printable shape for an LSTM input (initializer, graph input, or value_info).""" - graph = model.graph - for init in graph.initializer: - if init.name == input_name: - return tuple(numpy_helper.to_array(init).shape) - for vi in list(graph.input) + list(graph.value_info): - if vi.name != input_name: - continue - dims: list[int | str] = [] - for d in vi.type.tensor_type.shape.dim: - if d.dim_value: - dims.append(int(d.dim_value)) - elif d.dim_param: - dims.append(d.dim_param) - else: - dims.append("?") - return tuple(dims) - return f"" - - -def _print_lstm_input_shapes(model: onnx.ModelProto, *, direction: str) -> None: - graph = model.graph - lstm_nodes = [n for n in graph.node if n.op_type == "LSTM"] - if not lstm_nodes: - print(f"[test_dla_decompose_lstm] direction={direction}: no LSTM node in graph") - return - node = lstm_nodes[0] - print(f"[test_dla_decompose_lstm] direction={direction} LSTM name={node.name!r} input slots:") - for idx, inp in enumerate(node.input): - if not inp: - print(f" [{idx}] (empty optional)") - continue - sh = _lstm_input_shape_summary(model, inp) - print(f" [{idx}] {inp!r} shape={sh}") - - def _make_lstm_model(direction: str) -> onnx.ModelProto: """Single-LSTM graph: X input, W/R/B initializers, Y + Y_h + Y_c outputs.""" assert direction in ("forward", "reverse", "bidirectional") @@ -192,10 +124,12 @@ def _make_lstm_model(direction: str) -> onnx.ModelProto: def _single_lstm_node(model: onnx.ModelProto) -> onnx.NodeProto: + """Single lstm node.""" return next(n for n in model.graph.node if n.op_type == "LSTM") def test_downstream_transpose_keeps_its_nontrivial_permutation(): + """Downstream transpose keeps its nontrivial permutation.""" model = _make_lstm_model("bidirectional") model.graph.node.append( helper.make_node("Transpose", ["Y"], ["permuted"], perm=[0, 2, 1, 3], name="downstream") @@ -211,16 +145,14 @@ def test_downstream_transpose_keeps_its_nontrivial_permutation(): def _expect_unsupported_lstm(model: onnx.ModelProto, expected: str) -> None: - from modelopt.onnx.graph_surgery.utils.graph_utils import infer_shapes - + """Expect unsupported lstm.""" model = infer_shapes(copy.deepcopy(model)) with pytest.raises(ValueError, match=expected): _apply_decompose_lstm(model) def _run_ort(model: onnx.ModelProto, feeds: dict[str, np.ndarray]) -> dict[str, np.ndarray]: - import onnxruntime as ort # type: ignore[import-not-found] - + """Run ort.""" _clamp_ir_version(model) sess = ort.InferenceSession( model.SerializeToString(), @@ -233,6 +165,7 @@ def _run_ort(model: onnx.ModelProto, feeds: dict[str, np.ndarray]) -> dict[str, def _cosine_similarity(a: np.ndarray, b: np.ndarray) -> float: + """Cosine similarity.""" a64 = np.asarray(a, dtype=np.float64).ravel() b64 = np.asarray(b, dtype=np.float64).ravel() na = float(np.linalg.norm(a64)) @@ -245,6 +178,7 @@ def _cosine_similarity(a: np.ndarray, b: np.ndarray) -> float: def _assert_outputs_match(ref: dict[str, np.ndarray], got: dict[str, np.ndarray]) -> None: + """Assert outputs match.""" assert ref.keys() == got.keys() for name in ref: a, b = ref[name], got[name] @@ -258,7 +192,6 @@ def _assert_outputs_match(ref: dict[str, np.ndarray], got: dict[str, np.ndarray] def test_lstm_decompose_ort_cpu_parity(direction: str) -> None: """Decomposed graph matches reference LSTM on CPU EP (cosine ~1, tight allclose).""" model = _make_lstm_model(direction) - _print_lstm_input_shapes(model, direction=direction) rng = np.random.default_rng(7) x = rng.standard_normal((_T, _N, _D_IN)).astype(np.float32) feeds = {"X": x} @@ -266,15 +199,12 @@ def test_lstm_decompose_ort_cpu_parity(direction: str) -> None: ref = _run_ort(model, feeds) dec = copy.deepcopy(model) - from modelopt.onnx.graph_surgery.utils.graph_utils import infer_shapes - dec = infer_shapes(dec) dec = _apply_decompose_lstm(dec) _clamp_ir_version(dec) onnx.checker.check_model(dec) got = _run_ort(dec, feeds) - _save_lstm_test_pair(model, dec, direction) _assert_outputs_match(ref, got) @@ -338,3 +268,191 @@ def test_lstm_decompose_rejects_activation_params() -> None: node.attribute.append(helper.make_attribute("activation_alpha", [0.1])) _expect_unsupported_lstm(model, "activation_alpha") + + +@pytest.mark.parametrize("direction", ["forward", "reverse", "bidirectional"]) +def test_optional_bias_omission_preserves_all_outputs(direction): + """Execute all three directions without an optional bias tensor.""" + model = _make_lstm_model(direction) + del model.graph.node[0].input[3:] + model.graph.initializer.remove(next(t for t in model.graph.initializer if t.name == "B")) + feeds = {"X": np.random.default_rng(2).standard_normal((_T, _N, _D_IN)).astype(np.float32)} + expected = _run_ort(model, feeds) + transformed = _apply_decompose_lstm(onnx.shape_inference.infer_shapes(copy.deepcopy(model))) + onnx.checker.check_model(transformed, full_check=True) + _assert_outputs_match(expected, _run_ort(transformed, feeds)) + + +@pytest.mark.parametrize("name", ["W", "R", "B"]) +@pytest.mark.parametrize("axis", [None, 1, -1]) +def test_shared_dq_weight_preserves_source_layout_and_axes(name, axis): + """Preserve a public quantized source and DQ view while decomposing its LSTM use.""" + model = _make_lstm_model("forward") + weight = next(t for t in model.graph.initializer if t.name == name) + values = numpy_helper.to_array(weight) + quantized = np.clip(np.rint(values * 16), -128, 127).astype(np.int8) + dims = list(quantized.shape) + weight.CopyFrom(_init(name + "_q", quantized)) + controls = np.float32(0.0625) if axis is None else np.full(dims[axis], 0.0625, np.float32) + zeros = np.int8(0) if axis is None else np.zeros(dims[axis], np.int8) + model.graph.initializer.extend([_init(name + "_scale", controls), _init(name + "_zero", zeros)]) + dq = helper.make_node( + "DequantizeLinear", [name + "_q", name + "_scale", name + "_zero"], [name], name="weight_dq" + ) + if axis is not None: + dq.attribute.append(helper.make_attribute("axis", axis)) + model.graph.node.insert(0, dq) + model.graph.output.extend( + [_vi(name + "_q", TensorProto.INT8, dims), _vi(name, TensorProto.FLOAT, dims)] + ) + feeds = {"X": np.random.default_rng(3).standard_normal((_T, _N, _D_IN)).astype(np.float32)} + expected = _run_ort(model, feeds) + transformed = _apply_decompose_lstm(onnx.shape_inference.infer_shapes(copy.deepcopy(model))) + onnx.checker.check_model(transformed, full_check=True) + _assert_outputs_match(expected, _run_ort(transformed, feeds)) + + +@pytest.mark.parametrize("name", ["W", "R", "B"]) +def test_public_fixed_lstm_inputs_keep_their_original_view(name): + """Retain fixed weights or bias that are also graph outputs.""" + model = _make_lstm_model("forward") + weight = next(t for t in model.graph.initializer if t.name == name) + model.graph.output.append(_vi(name, TensorProto.FLOAT, list(weight.dims))) + feeds = {"X": np.random.default_rng(4).standard_normal((_T, _N, _D_IN)).astype(np.float32)} + expected = _run_ort(model, feeds) + transformed = _apply_decompose_lstm(onnx.shape_inference.infer_shapes(copy.deepcopy(model))) + onnx.checker.check_model(transformed, full_check=True) + _assert_outputs_match(expected, _run_ort(transformed, feeds)) + + +@pytest.mark.parametrize("name", ["W", "R", "B"]) +def test_overridable_lstm_inputs_use_runtime_values(name): + """Preserve graph-input initializer defaults as runtime LSTM controls.""" + model = _make_lstm_model("forward") + weight = next(t for t in model.graph.initializer if t.name == name) + model.graph.input.append(_vi(name, TensorProto.FLOAT, list(weight.dims))) + feeds = { + "X": np.random.default_rng(5).standard_normal((_T, _N, _D_IN)).astype(np.float32), + name: -numpy_helper.to_array(weight), + } + expected = _run_ort(model, feeds) + transformed = _apply_decompose_lstm(onnx.shape_inference.infer_shapes(copy.deepcopy(model))) + onnx.checker.check_model(transformed, full_check=True) + _assert_outputs_match(expected, _run_ort(transformed, feeds)) + + +@pytest.mark.parametrize("outputs", [["", "Y_h", "Y_c"], ["Y"], ["", "Y_h"]]) +def test_optional_lstm_outputs_remain_valid(outputs): + """Support omitted sequence and state output slots without introducing empty names.""" + model = _make_lstm_model("forward") + del model.graph.node[0].output[:] + model.graph.node[0].output.extend(outputs) + retained = [value for value in model.graph.output if value.name in outputs] + del model.graph.output[:] + model.graph.output.extend(retained) + feeds = {"X": np.random.default_rng(6).standard_normal((_T, _N, _D_IN)).astype(np.float32)} + expected = _run_ort(model, feeds) + transformed = _apply_decompose_lstm(onnx.shape_inference.infer_shapes(copy.deepcopy(model))) + onnx.checker.check_model(transformed, full_check=True) + _assert_outputs_match(expected, _run_ort(transformed, feeds)) + + +def test_explicit_zero_clip_refuses_before_mutation(): + """Reject a zero clipping threshold rather than interpreting it as omitted clipping.""" + model = _make_lstm_model("forward") + model.graph.node[0].attribute.append(helper.make_attribute("clip", 0.0)) + onnx.checker.check_model(model, full_check=True) + before = model.SerializeToString() + with pytest.raises(ValueError, match="clip"): + _apply_decompose_lstm(model) + assert model.SerializeToString() == before + + +def test_custom_lstm_function_is_preserved(): + """Leave a model-local operation whose name happens to be LSTM unchanged.""" + function = helper.make_function( + "custom", + "LSTM", + ["X"], + ["Y"], + [helper.make_node("Neg", ["X"], ["Y"])], + [helper.make_opsetid("", 17)], + ) + model = helper.make_model( + helper.make_graph( + [helper.make_node("LSTM", ["X"], ["Y"], domain="custom")], + "local_lstm", + [_vi("X", TensorProto.FLOAT, [2, 3])], + [_vi("Y", TensorProto.FLOAT, [2, 3])], + ), + functions=[function], + opset_imports=[helper.make_opsetid("", 17), helper.make_opsetid("custom", 1)], + ) + _clamp_ir_version(model) + onnx.checker.check_model(model, full_check=True) + before = model.SerializeToString() + _apply_decompose_lstm(model) + assert model.SerializeToString() == before + + +@pytest.mark.parametrize("kind", ["node", "initializer"]) +def test_generated_names_preserve_existing_public_values(kind): + """Allocate private names without overwriting an unrelated public graph value.""" + model = _make_lstm_model("forward") + name = "forward_conv_x_w_lstm0" if kind == "node" else "split_z_iofc_init_lstm0" + if kind == "node": + model.graph.node.insert(0, helper.make_node("Identity", ["X"], [name], name="existing")) + model.graph.output.append(_vi(name, TensorProto.FLOAT, [_T, _N, _D_IN])) + else: + model.graph.initializer.append(_init(name, np.asarray([97], np.int64))) + model.graph.output.append(_vi(name, TensorProto.INT64, [1])) + feeds = {"X": np.random.default_rng(9).standard_normal((_T, _N, _D_IN)).astype(np.float32)} + expected = _run_ort(model, feeds) + transformed = _apply_decompose_lstm(onnx.shape_inference.infer_shapes(copy.deepcopy(model))) + onnx.checker.check_model(transformed, full_check=True) + _assert_outputs_match(expected, _run_ort(transformed, feeds)) + + +def test_legacy_ir_registers_all_new_lstm_initializers(): + """Keep the legacy initializer-input contract after standalone decomposition.""" + model = _make_lstm_model("forward") + model.ir_version = 3 + for tensor in model.graph.initializer: + model.graph.input.append(_vi(tensor.name, tensor.data_type, list(tensor.dims))) + onnx.checker.check_model(model, full_check=True) + feeds = {"X": np.random.default_rng(10).standard_normal((_T, _N, _D_IN)).astype(np.float32)} + expected = _run_ort(model, feeds) + transformed = _apply_decompose_lstm(onnx.shape_inference.infer_shapes(copy.deepcopy(model))) + onnx.checker.check_model(transformed, full_check=True) + _assert_outputs_match(expected, _run_ort(transformed, feeds)) + + +def test_buffered_constant_state_convolution_keeps_sequential_rounding(): + """Retain established double accumulation across multiple bounded row buffers.""" + rng = np.random.default_rng(11) + weights = rng.standard_normal((513, 1024)).astype(np.float32) + state = rng.standard_normal(1024).astype(np.float32) + terms = np.concatenate( + [np.zeros((513, 1)), weights.astype(np.float64) * state.astype(np.float64)], axis=1 + ) + expected = np.cumsum(terms, axis=1)[:, -1].astype(np.float32) + graph = helper.make_graph( + [], + "buffered_state", + [], + [], + initializer=[ + _init("state", state.reshape(1, 1024, 1, 1)), + _init("weights", weights.reshape(513, 1024, 1, 1)), + ], + ) + assert ( + _fold_lstm_constants( + graph, + [helper.make_node("Conv", ["state", "weights"], ["folded"])], + {"state", "weights"}, + ) + == [] + ) + actual = numpy_helper.to_array(next(t for t in graph.initializer if t.name == "folded")) + assert actual.tobytes() == expected.tobytes() From 9f309ac713cf6d63b6fc5c517c26469d89039639 Mon Sep 17 00:00:00 2001 From: mgohil Date: Mon, 26 Oct 2026 22:00:09 +0000 Subject: [PATCH 3/6] Resolve native LSTM input contracts before rank preparation Signed-off-by: mgohil --- .../dla_transforms/dla_decompose_lstm.py | 87 +++++++++++++++---- .../dla_transforms/test_dla_decompose_lstm.py | 74 ++++++++++++++++ 2 files changed, 146 insertions(+), 15 deletions(-) diff --git a/modelopt/onnx/graph_surgery/dla_transforms/dla_decompose_lstm.py b/modelopt/onnx/graph_surgery/dla_transforms/dla_decompose_lstm.py index 35eee9f5dcc..12432f42434 100644 --- a/modelopt/onnx/graph_surgery/dla_transforms/dla_decompose_lstm.py +++ b/modelopt/onnx/graph_surgery/dla_transforms/dla_decompose_lstm.py @@ -45,7 +45,6 @@ add_unique_initializers, batch_replace_nodes, get_constant_by_name, - is_const_dq_input, ) from .dla_constants_to_initializers import _add_legacy_initializer_inputs from .dla_graph_cleanup import _subgraph_captures @@ -143,7 +142,8 @@ def _promote_lstm_input_to_rank4( shape = cache.get_shape(tensor_name) if shape is not None and len(shape) >= 4: return - axes_list = [0] if shape is None or len(shape) == 3 else list(range(4 - len(shape))) + rank = len(shape) if shape is not None else (2 if input_index == 3 else 3) + axes_list = list(range(max(0, 4 - rank))) axes_name = cache.unique_name(f"{tensor_name}_lstm4d_axes_{lstm_node.name}_{input_index}") out_name = cache.unique_name(f"{tensor_name}_lstm4d_{lstm_node.name}_{input_index}") @@ -537,17 +537,28 @@ def resolve_lstm_tensor_shape(tensor_name, _seen=None): return None return list(init.dims) - init_names = list(cache.init_names) - dq_src = is_const_dq_input(init_names, tensor_name, graph) - if dq_src: - return resolve_lstm_tensor_shape(dq_src, _seen) - prod = cache.get_producer(tensor_name) - if prod is None: + if ( + prod is None + or prod.overload + or (prod.domain, prod.op_type, prod.overload) in local_functions + ): return None - if prod.op_type == "Unsqueeze" and len(prod.input) >= 2: + standard = prod.domain in {"", "ai.onnx"} + if prod.op_type == "DequantizeLinear" and ( + standard + or ( + prod.domain == "com.microsoft" + and any( + entry.domain == prod.domain and entry.version == 1 + for entry in model.opset_import + ) + ) + ): + return resolve_lstm_tensor_shape(prod.input[0], _seen) + if standard and prod.op_type == "Unsqueeze" and len(prod.input) >= 2: inner = prod.input[0] - axes_name = new_name(prod.input[1]) + axes_name = prod.input[1] inner_shape = resolve_lstm_tensor_shape(inner, _seen) if inner_shape is None: return None @@ -558,17 +569,61 @@ def resolve_lstm_tensor_shape(tensor_name, _seen=None): return None axes_flat = np.asarray(ax, dtype=np.int64).flatten().tolist() out = list(inner_shape) - r = len(out) - norm = [int(a) + r + 1 if int(a) < 0 else int(a) for a in axes_flat] + rank = len(out) + len(axes_flat) + norm = [int(a) + rank if int(a) < 0 else int(a) for a in axes_flat] for a in sorted(norm): if a < 0 or a > len(out): return None out.insert(a, 1) return out - if prod.op_type in ("Identity", "Cast") and prod.input and prod.input[0]: + if standard and prod.op_type in ("Identity", "Cast") and prod.input and prod.input[0]: return resolve_lstm_tensor_shape(prod.input[0], _seen) return None + def resolve_lstm_tensor_dtype(name, seen=None): + """Resolve supported dtype-preserving producers without a prior inference pass.""" + dtype = cache.get_dtype(name) + if dtype is not None: + return dtype + seen = set() if seen is None else seen + if name in seen: + return None + seen.add(name) + producer = cache.get_producer(name) + if ( + producer is None + or producer.overload + or (producer.domain, producer.op_type, producer.overload) in local_functions + ): + return None + standard = producer.domain in {"", "ai.onnx"} + if producer.op_type == "DequantizeLinear" and ( + standard + or ( + producer.domain == "com.microsoft" + and any( + entry.domain == producer.domain and entry.version == 1 + for entry in model.opset_import + ) + ) + ): + output_dtype = next( + ( + attribute.i + for attribute in producer.attribute + if attribute.name == "output_dtype" + ), + 0, + ) + return output_dtype or cache.get_dtype(producer.input[1]) + if standard and producer.op_type == "Cast": + return next( + (attribute.i for attribute in producer.attribute if attribute.name == "to"), None + ) + if standard and producer.op_type in {"Identity", "Unsqueeze"}: + return resolve_lstm_tensor_dtype(producer.input[0], seen) + return None + def create_shape_info_map(node): """Resolve the fixed dimensions required for sequence unrolling.""" x, w = node.input[:2] @@ -868,12 +923,14 @@ def transpose_outputs( and next((opset.version for opset in model.opset_import if not opset.domain), 0) < 13 ): raise ValueError("LSTM decomposition requires ONNX opset >= 13") + validated_shapes = {} for node in lstm_nodes: _validate_lstm_supported(node) shape_info = create_shape_info_map(node) if any(value <= 0 for value in shape_info.values()): raise ValueError(f"LSTM {node.name!r}: all dimensions must be static and positive") - dtype = cache.get_dtype(node.input[0]) + validated_shapes[id(node)] = shape_info + dtype = resolve_lstm_tensor_dtype(node.input[0]) if dtype != onnx.TensorProto.FLOAT: raise ValueError(f"LSTM {node.name!r}: decomposition supports FP32 inputs") names = set() @@ -890,7 +947,7 @@ def transpose_outputs( # Rebuild cache after promotion added new initializers/value_info/nodes cache.rebuild() nodes_to_add = list(prep_nodes) - shape_info_map = create_shape_info_map(node) + shape_info_map = validated_shapes[id(node)] if direction == "bidirectional" and int(shape_info_map["d"]) != 2: raise ValueError( f"LSTM {node.name!r}: direction is bidirectional but num_directions (from W) is " diff --git a/tests/unit/onnx/graph_surgery/dla_transforms/test_dla_decompose_lstm.py b/tests/unit/onnx/graph_surgery/dla_transforms/test_dla_decompose_lstm.py index f2a9abb6dbe..c1c35d9ae91 100644 --- a/tests/unit/onnx/graph_surgery/dla_transforms/test_dla_decompose_lstm.py +++ b/tests/unit/onnx/graph_surgery/dla_transforms/test_dla_decompose_lstm.py @@ -456,3 +456,77 @@ def test_buffered_constant_state_convolution_keeps_sequential_rounding(): ) actual = numpy_helper.to_array(next(t for t in graph.initializer if t.name == "folded")) assert actual.tobytes() == expected.tobytes() + + +@pytest.mark.parametrize("role", ["X", "W", "R", "B"]) +def test_dequantized_inputs_without_value_info_preserve_numerics(role): + """Resolve the native LSTM contract without relying on a prior inference pass.""" + model = _make_lstm_model("forward") + rng = np.random.default_rng(12) + values = ( + rng.standard_normal((_T, _N, _D_IN)).astype(np.float32) + if role == "X" + else numpy_helper.to_array(next(t for t in model.graph.initializer if t.name == role)) + ) + quantized = np.clip(np.rint(values * 16), -128, 127).astype(np.int8) + feeds = {"X": rng.standard_normal((_T, _N, _D_IN)).astype(np.float32)} + if role == "X": + model.graph.input[0].CopyFrom(_vi("X_q", TensorProto.INT8, list(quantized.shape))) + feeds = {"X_q": quantized} + else: + weight = next(t for t in model.graph.initializer if t.name == role) + weight.CopyFrom(_init(role + "_q", quantized)) + model.graph.initializer.extend( + [_init(role + "_scale", np.float32(0.0625)), _init(role + "_zero", np.int8(0))] + ) + model.graph.node.insert( + 0, + helper.make_node( + "DequantizeLinear", [role + "_q", role + "_scale", role + "_zero"], [role] + ), + ) + onnx.checker.check_model(model, full_check=True) + expected = _run_ort(model, feeds) + transformed = _apply_decompose_lstm(copy.deepcopy(model)) + onnx.checker.check_model(transformed, full_check=True) + _assert_outputs_match(expected, _run_ort(transformed, feeds)) + + +@pytest.mark.parametrize("constant_axes", [False, True]) +def test_unsqueezed_weight_without_value_info_reads_existing_axes(constant_axes): + """Look up existing axes by their actual initializer or Constant tensor name.""" + model = _make_lstm_model("forward") + weight = next(t for t in model.graph.initializer if t.name == "W") + weight.CopyFrom(_init("W_flat", numpy_helper.to_array(weight).reshape(4 * _H, _D_IN))) + axes = _init("weight_axes", np.asarray([0], np.int64)) + if constant_axes: + model.graph.node.insert(0, helper.make_node("Constant", [], ["weight_axes"], value=axes)) + else: + model.graph.initializer.append(axes) + model.graph.node.insert( + int(constant_axes), helper.make_node("Unsqueeze", ["W_flat", "weight_axes"], ["W"]) + ) + feeds = {"X": np.random.default_rng(13).standard_normal((_T, _N, _D_IN)).astype(np.float32)} + onnx.checker.check_model(model, full_check=True) + expected = _run_ort(model, feeds) + transformed = _apply_decompose_lstm(copy.deepcopy(model)) + onnx.checker.check_model(transformed, full_check=True) + _assert_outputs_match(expected, _run_ort(transformed, feeds)) + + +@pytest.mark.parametrize("axes", [[0, 1], [-3, -2]]) +def test_unsqueezed_activation_without_value_info_handles_multiple_axes(axes): + """Resolve positive and negative multiple axes against the Unsqueeze output rank.""" + model = _make_lstm_model("forward") + model.graph.input[0].CopyFrom(_vi("X_flat", TensorProto.FLOAT, [_D_IN])) + model.graph.initializer.append(_init("activation_axes", np.asarray(axes, np.int64))) + model.graph.node.insert(0, helper.make_node("Unsqueeze", ["X_flat", "activation_axes"], ["X"])) + model.graph.output[0].CopyFrom(_vi("Y", TensorProto.FLOAT, [1, 1, 1, _H])) + for output in model.graph.output[1:]: + output.CopyFrom(_vi(output.name, TensorProto.FLOAT, [1, 1, _H])) + feeds = {"X_flat": np.random.default_rng(14).standard_normal(_D_IN).astype(np.float32)} + onnx.checker.check_model(model, full_check=True) + expected = _run_ort(model, feeds) + transformed = _apply_decompose_lstm(copy.deepcopy(model)) + onnx.checker.check_model(transformed, full_check=True) + _assert_outputs_match(expected, _run_ort(transformed, feeds)) From d5a12808c2129e844c8be19d5ea21f5a2e7f972a Mon Sep 17 00:00:00 2001 From: mgohil Date: Mon, 26 Oct 2026 22:19:03 +0000 Subject: [PATCH 4/6] Allocate private runtime LSTM transpose outputs Signed-off-by: mgohil --- .../dla_transforms/dla_decompose_lstm.py | 2 +- .../dla_transforms/test_dla_decompose_lstm.py | 52 ++++++++++++++++--- 2 files changed, 45 insertions(+), 9 deletions(-) diff --git a/modelopt/onnx/graph_surgery/dla_transforms/dla_decompose_lstm.py b/modelopt/onnx/graph_surgery/dla_transforms/dla_decompose_lstm.py index 12432f42434..c2b6e0a1127 100644 --- a/modelopt/onnx/graph_surgery/dla_transforms/dla_decompose_lstm.py +++ b/modelopt/onnx/graph_surgery/dla_transforms/dla_decompose_lstm.py @@ -408,7 +408,7 @@ def transpose_tensor(graph, tensor_name, perm, nodes_to_add, unique_initializer) transpose_node = helper.make_node( "Transpose", inputs=[tensor_name], - outputs=[tensor_name + "_transposed"], + outputs=[cache.unique_name(tensor_name + "_transposed")], name=new_name(tensor_name + "_transpose"), perm=perm, ) diff --git a/tests/unit/onnx/graph_surgery/dla_transforms/test_dla_decompose_lstm.py b/tests/unit/onnx/graph_surgery/dla_transforms/test_dla_decompose_lstm.py index c1c35d9ae91..b631575a81c 100644 --- a/tests/unit/onnx/graph_surgery/dla_transforms/test_dla_decompose_lstm.py +++ b/tests/unit/onnx/graph_surgery/dla_transforms/test_dla_decompose_lstm.py @@ -87,7 +87,7 @@ def _clamp_ir_version(model: onnx.ModelProto) -> None: model.ir_version = min(model.ir_version, _MAX_IR_VERSION_FOR_ORT) -def _make_lstm_model(direction: str) -> onnx.ModelProto: +def _make_lstm_model(direction: str, batch: int = _N) -> onnx.ModelProto: """Single-LSTM graph: X input, W/R/B initializers, Y + Y_h + Y_c outputs.""" assert direction in ("forward", "reverse", "bidirectional") num_dir = 2 if direction == "bidirectional" else 1 @@ -110,11 +110,11 @@ def _make_lstm_model(direction: str) -> onnx.ModelProto: direction=direction, ) ] - inputs = [_vi("X", TensorProto.FLOAT, [_T, _N, _D_IN])] + inputs = [_vi("X", TensorProto.FLOAT, [_T, batch, _D_IN])] outputs = [ - _vi("Y", TensorProto.FLOAT, [_T, num_dir, _N, _H]), - _vi("Y_h", TensorProto.FLOAT, [num_dir, _N, _H]), - _vi("Y_c", TensorProto.FLOAT, [num_dir, _N, _H]), + _vi("Y", TensorProto.FLOAT, [_T, num_dir, batch, _H]), + _vi("Y_h", TensorProto.FLOAT, [num_dir, batch, _H]), + _vi("Y_c", TensorProto.FLOAT, [num_dir, batch, _H]), ] graph = helper.make_graph(nodes, f"lstm_{direction}", inputs, outputs, initializer=inits) model = helper.make_model(graph, opset_imports=[helper.make_opsetid("", _OPSET_VERSION)]) @@ -188,12 +188,13 @@ def _assert_outputs_match(ref: dict[str, np.ndarray], got: dict[str, np.ndarray] np.testing.assert_allclose(a, b, rtol=1e-4, atol=1e-4, err_msg=f"{name}: value mismatch") +@pytest.mark.parametrize("batch", [1, _N]) @pytest.mark.parametrize("direction", ["forward", "reverse", "bidirectional"]) -def test_lstm_decompose_ort_cpu_parity(direction: str) -> None: +def test_lstm_decompose_ort_cpu_parity(direction: str, batch: int) -> None: """Decomposed graph matches reference LSTM on CPU EP (cosine ~1, tight allclose).""" - model = _make_lstm_model(direction) + model = _make_lstm_model(direction, batch) rng = np.random.default_rng(7) - x = rng.standard_normal((_T, _N, _D_IN)).astype(np.float32) + x = rng.standard_normal((_T, batch, _D_IN)).astype(np.float32) feeds = {"X": x} ref = _run_ort(model, feeds) @@ -530,3 +531,38 @@ def test_unsqueezed_activation_without_value_info_handles_multiple_axes(axes): transformed = _apply_decompose_lstm(copy.deepcopy(model)) onnx.checker.check_model(transformed, full_check=True) _assert_outputs_match(expected, _run_ort(transformed, feeds)) + + +@pytest.mark.parametrize( + ("role", "index"), [("X", 0), ("W", 1), ("R", 2), ("B", 3), ("initial_h", 5), ("initial_c", 6)] +) +def test_runtime_input_transpose_preserves_existing_tensor_names(role, index): + """Allocate runtime transpose outputs without redefining an unrelated public tensor.""" + model = _make_lstm_model("forward") + rng = np.random.default_rng(16) + feeds = {"X": rng.standard_normal((_T, _N, _D_IN)).astype(np.float32)} + if role in {"W", "R", "B"}: + tensor = next(tensor for tensor in model.graph.initializer if tensor.name == role) + shape = list(tensor.dims) + model.graph.input.append(_vi(role, TensorProto.FLOAT, shape)) + feeds[role] = -numpy_helper.to_array(tensor) + elif role in {"initial_h", "initial_c"}: + shape = [1, _N, _H] + model.graph.input.append(_vi(role, TensorProto.FLOAT, shape)) + node = model.graph.node[0] + while len(node.input) <= index: + node.input.append("") + node.input[index] = role + feeds[role] = rng.standard_normal(shape).astype(np.float32) + else: + shape = [_T, _N, _D_IN] + existing = f"{role}_lstm4d_lstm0_{index}_transposed" + model.graph.node.insert( + 0, helper.make_node("Identity", [role], [existing], name="existing_view") + ) + model.graph.output.append(_vi(existing, TensorProto.FLOAT, shape)) + onnx.checker.check_model(model, full_check=True) + expected = _run_ort(model, feeds) + transformed = _apply_decompose_lstm(onnx.shape_inference.infer_shapes(copy.deepcopy(model))) + onnx.checker.check_model(transformed, full_check=True) + _assert_outputs_match(expected, _run_ort(transformed, feeds)) From a04348413b749959d0c8f2304e62ea48ce7dec18 Mon Sep 17 00:00:00 2001 From: mgohil Date: Mon, 26 Oct 2026 22:48:34 +0000 Subject: [PATCH 5/6] Preserve LSTM follower layouts and validate direction counts before mutation Signed-off-by: mgohil --- .../dla_transforms/dla_decompose_lstm.py | 49 +++++-- .../dla_transforms/test_dla_decompose_lstm.py | 136 +++++++++++++++++- 2 files changed, 176 insertions(+), 9 deletions(-) diff --git a/modelopt/onnx/graph_surgery/dla_transforms/dla_decompose_lstm.py b/modelopt/onnx/graph_surgery/dla_transforms/dla_decompose_lstm.py index c2b6e0a1127..ca789588571 100644 --- a/modelopt/onnx/graph_surgery/dla_transforms/dla_decompose_lstm.py +++ b/modelopt/onnx/graph_surgery/dla_transforms/dla_decompose_lstm.py @@ -234,6 +234,7 @@ def _apply_decompose_lstm(model, *, defer_rank_conversion=False): lstm_counter = 0 unique_initializer = set() allocated_names = {} + captured_names = _subgraph_captures(graph) def new_name(proposed): """Reuse internal controls while allocating names outside the original graph namespace.""" @@ -881,6 +882,9 @@ def transpose_outputs( if ( len(consumer_nodes_y_output) == 1 and consumer_nodes_y_output[0][0].op_type == "Transpose" + and consumer_nodes_y_output[0][0].domain in {"", "ai.onnx"} + and not consumer_nodes_y_output[0][0].overload + and (consumer_nodes_y_output[0][0].domain, "Transpose", "") not in local_functions ): consumer = consumer_nodes_y_output[0][0] shape = cache.get_shape(output_name) @@ -888,14 +892,17 @@ def transpose_outputs( (list(attr.ints) for attr in consumer.attribute if attr.name == "perm"), None ) followers = cache.get_consumers(consumer.output[0]) - # The later rank pass resolves 0 dimensions from original output metadata. restored_by_rank_pass = ( defer_rank_conversion - and consumer.output[0] not in cache.graph_output_names + and consumer.output[0] not in cache.graph_output_names | captured_names and bool(followers) and all( next_node.op_type == "Reshape" + and next_node.domain in {"", "ai.onnx"} + and not next_node.overload + and (next_node.domain, "Reshape", "") not in local_functions and len(next_node.input) > 1 + and next_node.input[1] not in cache.graph_input_names and (target := cache.get_init_array(next_node.input[1])) is not None and 0 < target.size < 4 and cache.get_shape(next_node.output[0]) is not None @@ -910,6 +917,31 @@ def transpose_outputs( and [index for index in perm if shape[index] != 1] == [index for index, dim in enumerate(shape) if dim != 1] ) + if noop and restored_by_rank_pass: + for follower, _ in followers: + target = cache.get_init_array(follower.input[1]) + if np.any(target == 0): + # Copied dimensions refer to the original transpose layout. + original_shape = cache.get_shape(follower.output[0]) + assert original_shape is not None + resolved = np.array(target, copy=True) + for index in np.flatnonzero(target == 0): + resolved[index] = original_shape[index] + target_name = cache.unique_name(follower.input[1] + "_lstm_layout") + add_unique_initializers( + graph, + unique_initializer, + [numpy_helper.from_array(resolved, target_name)], + ) + follower.input[1] = target_name + if noop: + for value in graph.value_info: + if value.name == consumer.output[0]: + value.type.tensor_type.shape.CopyFrom( + helper.make_tensor_value_info( + "", onnx.TensorProto.FLOAT, shape + ).type.tensor_type.shape + ) for attr in consumer.attribute if noop else []: if attr.name == "perm": del attr.ints[:] @@ -925,10 +957,16 @@ def transpose_outputs( raise ValueError("LSTM decomposition requires ONNX opset >= 13") validated_shapes = {} for node in lstm_nodes: - _validate_lstm_supported(node) + direction = _validate_lstm_supported(node) shape_info = create_shape_info_map(node) if any(value <= 0 for value in shape_info.values()): raise ValueError(f"LSTM {node.name!r}: all dimensions must be static and positive") + expected_directions = 2 if direction == "bidirectional" else 1 + if int(shape_info["d"]) != expected_directions: + raise ValueError( + f"LSTM {node.name!r}: direction is {direction} but num_directions (from W) is " + f"{shape_info['d']!r}; expected {expected_directions}." + ) validated_shapes[id(node)] = shape_info dtype = resolve_lstm_tensor_dtype(node.input[0]) if dtype != onnx.TensorProto.FLOAT: @@ -948,11 +986,6 @@ def transpose_outputs( cache.rebuild() nodes_to_add = list(prep_nodes) shape_info_map = validated_shapes[id(node)] - if direction == "bidirectional" and int(shape_info_map["d"]) != 2: - raise ValueError( - f"LSTM {node.name!r}: direction is bidirectional but num_directions (from W) is " - f"{shape_info_map['d']!r}; expected 2." - ) transposed_inputs = transpose_inputs( graph, node, unique_initializer, shape_info_map, nodes_to_add ) diff --git a/tests/unit/onnx/graph_surgery/dla_transforms/test_dla_decompose_lstm.py b/tests/unit/onnx/graph_surgery/dla_transforms/test_dla_decompose_lstm.py index b631575a81c..d022ccf8d98 100644 --- a/tests/unit/onnx/graph_surgery/dla_transforms/test_dla_decompose_lstm.py +++ b/tests/unit/onnx/graph_surgery/dla_transforms/test_dla_decompose_lstm.py @@ -24,7 +24,8 @@ import numpy as np import onnx import pytest -from onnx import TensorProto, helper, numpy_helper +from onnx import AttributeProto, TensorProto, helper, numpy_helper +from onnx.compose import add_prefix from modelopt.onnx.graph_surgery.dla_transforms.dla_decompose_lstm import ( _apply_decompose_lstm, @@ -566,3 +567,136 @@ def test_runtime_input_transpose_preserves_existing_tensor_names(role, index): transformed = _apply_decompose_lstm(onnx.shape_inference.infer_shapes(copy.deepcopy(model))) onnx.checker.check_model(transformed, full_check=True) _assert_outputs_match(expected, _run_ort(transformed, feeds)) + + +def _make_lstm_followed_by_copied_reshape(target, view="private"): + """Create a valid unary-dimension transpose followed by a potentially shared target.""" + model = _make_lstm_model("forward", 2) + model.graph.initializer.append(_init("shape_target", np.asarray(target, np.int64))) + model.graph.node.extend( + [ + helper.make_node( + "Transpose", ["Y"], ["layout"], perm=[0, 2, 1, 3], name="layout_change" + ), + helper.make_node( + "Reshape", ["layout", "shape_target"], ["reshaped"], name="restore_shape" + ), + ] + ) + model.graph.output[0].CopyFrom(_vi("reshaped", TensorProto.FLOAT, None)) + feeds = {"X": np.random.default_rng(17).standard_normal((_T, 2, _D_IN)).astype(np.float32)} + if view == "public": + model.graph.output.append(_vi("shape_target", TensorProto.INT64, [3])) + elif view == "overridable": + model.graph.input.append(_vi("shape_target", TensorProto.INT64, [3])) + elif view == "captured": + model.graph.input.append(_vi("condition", TensorProto.BOOL, [])) + feeds["condition"] = np.array(True) + branch = helper.make_graph( + [helper.make_node("Identity", ["layout"], ["branch_output"])], + "layout_reader", + [], + [_vi("branch_output", TensorProto.FLOAT, [_T, 2, 1, _H])], + ) + model.graph.node.append( + helper.make_node( + "If", ["condition"], ["captured_output"], then_branch=branch, else_branch=branch + ) + ) + model.graph.output.append(_vi("captured_output", TensorProto.FLOAT, [_T, 2, 1, _H])) + elif view == "shared": + model.graph.input.append(_vi("other", TensorProto.FLOAT, [_T, 3, _H])) + feeds["other"] = np.arange(_T * 3 * _H, dtype=np.float32).reshape(_T, 3, _H) + model.graph.node.append( + helper.make_node("Reshape", ["other", "shape_target"], ["other_out"]) + ) + model.graph.output.append(_vi("other_out", TensorProto.FLOAT, [_T, 3, _H])) + model = onnx.shape_inference.infer_shapes(model) + onnx.checker.check_model(model, full_check=True) + return model, feeds + + +@pytest.mark.parametrize("target", [[0, 0, _H], [0, -1, 0], [-1, 0, 0], [_T, 2, _H]]) +@pytest.mark.parametrize("defer", [False, True]) +def test_deferred_transpose_keeps_copied_reshape_dimensions(target, defer): + """Preserve copied dimensions when moving unary axes through a deferred reshape.""" + model, feeds = _make_lstm_followed_by_copied_reshape(target) + expected = _run_ort(model, feeds) + transformed = _apply_decompose_lstm(copy.deepcopy(model), defer_rank_conversion=defer) + onnx.checker.check_model(transformed, full_check=True) + _assert_outputs_match(expected, _run_ort(transformed, feeds)) + + +@pytest.mark.parametrize("view", ["public", "shared", "overridable", "captured"]) +def test_deferred_copied_reshape_preserves_target_owners(view): + """Preserve public, shared and runtime structural controls while changing a private layout.""" + model, feeds = _make_lstm_followed_by_copied_reshape([0, 0, _H], view) + expected = _run_ort(model, feeds) + transformed = _apply_decompose_lstm(copy.deepcopy(model), defer_rank_conversion=True) + onnx.checker.check_model(transformed, full_check=True) + _assert_outputs_match(expected, _run_ort(transformed, feeds)) + + +@pytest.mark.parametrize("operator", ["Transpose", "Reshape"]) +def test_deferred_layout_gate_preserves_model_local_followers(operator): + """Preserve opaque follower functions whose names match standard layout operators.""" + model, feeds = _make_lstm_followed_by_copied_reshape([_T, 2, _H]) + node = next(node for node in model.graph.node if node.op_type == operator) + node.domain = "custom" + if operator == "Transpose": + inner = helper.make_node("Transpose", ["X"], ["moved"]) + inner.attribute.append( + AttributeProto(name="perm", type=AttributeProto.INTS, ref_attr_name="perm") + ) + function = helper.make_function( + "custom", + operator, + ["X"], + ["Y"], + [inner, helper.make_node("Softmax", ["moved"], ["Y"], axis=1)], + [helper.make_opsetid("", 17)], + attributes=["perm"], + ) + else: + function = helper.make_function( + "custom", + operator, + ["X", "shape"], + ["Y"], + [ + helper.make_node("Softmax", ["X"], ["probabilities"], axis=1), + helper.make_node("Reshape", ["probabilities", "shape"], ["Y"]), + ], + [helper.make_opsetid("", 17)], + ) + model.functions.append(function) + model.opset_import.append(helper.make_opsetid("custom", 1)) + onnx.checker.check_model(model, full_check=True) + expected = _run_ort(model, feeds) + transformed = _apply_decompose_lstm(copy.deepcopy(model), defer_rank_conversion=True) + onnx.checker.check_model(transformed, full_check=True) + _assert_outputs_match(expected, _run_ort(transformed, feeds)) + + +@pytest.mark.parametrize( + ("direction", "actual_directions"), [("forward", 2), ("reverse", 2), ("bidirectional", 1)] +) +@pytest.mark.parametrize("late", [False, True]) +def test_direction_count_refusal_precedes_all_lstm_mutation(direction, actual_directions, late): + """Refuse a direction-count mismatch before rewriting this or an earlier LSTM.""" + model = _make_lstm_model("bidirectional" if actual_directions == 2 else "forward") + node = model.graph.node[0] + next( + attribute for attribute in node.attribute if attribute.name == "direction" + ).s = direction.encode() + if late: + first = add_prefix(_make_lstm_model("forward"), "first/") + model.graph.node.insert(0, first.graph.node[0]) + model.graph.initializer.extend(first.graph.initializer) + model.graph.input.extend(first.graph.input) + model.graph.output.extend(first.graph.output) + onnx.checker.check_model(model) + before = model.SerializeToString() + with pytest.raises(ValueError, match="num_directions"): + _apply_decompose_lstm(model) + assert model.SerializeToString() == before From f30b877b0dbd4474dbca14ddf73171bbba57eee8 Mon Sep 17 00:00:00 2001 From: mgohil Date: Mon, 26 Oct 2026 22:52:10 +0000 Subject: [PATCH 6/6] Reuse compatible LSTM reshape controls without duplicating constants Signed-off-by: mgohil --- .../dla_transforms/dla_decompose_lstm.py | 32 +++++++++++++++---- 1 file changed, 26 insertions(+), 6 deletions(-) diff --git a/modelopt/onnx/graph_surgery/dla_transforms/dla_decompose_lstm.py b/modelopt/onnx/graph_surgery/dla_transforms/dla_decompose_lstm.py index ca789588571..775bd1f7909 100644 --- a/modelopt/onnx/graph_surgery/dla_transforms/dla_decompose_lstm.py +++ b/modelopt/onnx/graph_surgery/dla_transforms/dla_decompose_lstm.py @@ -927,13 +927,33 @@ def transpose_outputs( resolved = np.array(target, copy=True) for index in np.flatnonzero(target == 0): resolved[index] = original_shape[index] - target_name = cache.unique_name(follower.input[1] + "_lstm_layout") - add_unique_initializers( - graph, - unique_initializer, - [numpy_helper.from_array(resolved, target_name)], + target_name = follower.input[1] + copied_axes = np.flatnonzero(target == 0) + shared_resolution = ( + target_name not in cache.graph_output_names | captured_names + and all( + user.op_type == "Reshape" + and user.domain in {"", "ai.onnx"} + and not user.overload + and (user.domain, "Reshape", "") not in local_functions + and input_index == 1 + and (user_shape := cache.get_shape(user.output[0])) is not None + and all(user_shape[axis] == resolved[axis] for axis in copied_axes) + for user, input_index in cache.get_consumers(target_name) + ) ) - follower.input[1] = target_name + if shared_resolution: + initializer = cache.get_init(target_name) + assert initializer is not None + initializer.CopyFrom(numpy_helper.from_array(resolved, target_name)) + else: + target_name = cache.unique_name(target_name + "_lstm_layout") + add_unique_initializers( + graph, + unique_initializer, + [numpy_helper.from_array(resolved, target_name)], + ) + follower.input[1] = target_name if noop: for value in graph.value_info: if value.name == consumer.output[0]: