`SlicesSplit` rule makes `optimizer.optimize` quadratic in the number of `Slice` nodes

Đang mở
  • Có gì không ổn?
#3,154 0 bình luận 0 reaction 0 người được giao Xem trên GitHub

Maintainer thường phản hồi trong vòng 4 ngày

Chưa có ai nhận issue này.

Đánh giá

Độ khó
4/5
Thời gian dự kiến
3-5 ngày
Mức phù hợp với người mới
68/100
Loại issue
Lỗi
Độ rõ ràng
Đặc tả rõ ràng
Mức độ hoạt động
Sôi nổi
Công nghệ
python
Lĩnh vực
machine-learning, performance

Hướng nghiên cứu

Start by reading SimplePatternMatcher.match and _multi_match in onnxscript.rewriter, then inspect the SlicesSplit rule in _DEFAULT_REWRITE_RULES. Use the provided repro to compare its runtime with the other default rules as the number of Slice nodes grows. Done means SlicesSplit no longer exhibits quadratic scaling while preserving rewrite behavior.

Do mô hình lập chỉ mục viết ra từ nội dung của issue.

Mô tả

module: optimizer

SimplePatternMatcher.match handles a pattern with several output nodes by fixing the first output to the candidate node and trying every combination (itertools.product) of graph nodes with the same op type for the other outputs. It also rebuilds an index of all the graph's nodes for each candidate. The only multi-output rule in _DEFAULT_REWRITE_RULES is SlicesSplit, so on a graph with many Slices that one rule takes far longer than all the others, and its time grows with the square of the number of Slices. optimize() runs it on every iteration, and optimize_for_ort calls optimize() three times.

We hit this exporting a mixture-of-experts model whose experts were unrolled into a Slice per expert (24 layers × 32 experts, about 3,900 Slices). torch.onnx.export(optimize=True) took 11 minutes per graph, and optimize_for_ort 13 minutes, against seconds for a dense model of the same size. A py-spy profile put 97% of the time in SimplePatternMatcher.match / _multi_match.

Repro (onnxscript 0.7.2, onnx_ir 1.0.0): n independent Slices of one input, timing SlicesSplit alone against the other 45 default rules together.

"""onnxscript's `SlicesSplit` rewrite rule takes time quadratic in the number of `Slice` nodes in the graph.

The graph is n independent `Slice`s of one input. `SlicesSplit` alone is timed against every other default rewrite
rule together, on the same graph.
"""

import time

import numpy as np
import onnx_ir as ir
import onnxscript.rewriter


def constant(name, values):
    value = ir.Value(name=name, const_value=ir.tensor(np.array(values, dtype=np.int64)))
    value.dtype, value.shape = ir.DataType.INT64, ir.Shape([len(values)])
    return value


def build(n: int) -> ir.Model:
    x = ir.Value(name="x", type=ir.TensorType(ir.DataType.FLOAT), shape=ir.Shape([4 * n, 8]))
    nodes, initializers, outputs = [], [], []
    for i in range(n):
        arguments = [constant(f"starts{i}", [4 * i]), constant(f"ends{i}", [4 * i + 4]), constant(f"axes{i}", [0])]
        part = ir.node("Slice", [x, *arguments])
        relu = ir.node("Relu", [part.outputs[0]])
        relu.outputs[0].name = f"y{i}"
        relu.outputs[0].dtype, relu.outputs[0].shape = ir.DataType.FLOAT, ir.Shape([4, 8])
        nodes += [part, relu]
        initializers += arguments
        outputs.append(relu.outputs[0])
    graph = ir.Graph([x], outputs, nodes=nodes, initializers=initializers, opset_imports={"": 18})
    return ir.Model(graph, ir_version=10)


rules = list(onnxscript.rewriter._DEFAULT_REWRITE_RULES)
slices_split = [rule for rule in rules if rule.name == "SlicesSplit"]
others = [rule for rule in rules if rule.name != "SlicesSplit"]
for n in (250, 500, 1000, 2000):
    times = []
    for selected in (slices_split, others):
        model = build(n)
        start = time.perf_counter()
        onnxscript.rewriter.RewriteRuleSet(selected).apply_to_model(model)
        times.append(time.perf_counter() - start)
    print(f"{n:5d} Slice nodes: SlicesSplit {times[0]:6.2f} s, the {len(others)} other default rules {times[1]:6.2f} s")

  250 Slice nodes: SlicesSplit   0.24 s, the 45 other default rules   0.09 s
  500 Slice nodes: SlicesSplit   0.91 s, the 45 other default rules   0.18 s
 1000 Slice nodes: SlicesSplit   3.56 s, the 45 other default rules   0.36 s
 2000 Slice nodes: SlicesSplit  14.11 s, the 45 other default rules   0.72 s

Each doubling of n makes SlicesSplit 4x slower and the other rules together 2x slower.

Possible fixes:

  • for the other outputs, only try nodes that share an input with the first one, since that's what SlicesSplit requires anyway;
  • build the op_to_nodes index once per pass, not once per candidate;
  • let optimize() take a rule set, so callers can leave a rule out.
Ngôn ngữ chính
Python
Star
457
Fork
149
Merge trung bình
3 ngày 20 giờ
Pull request đã merge (30 ngày)
22

Chuẩn bị môi trường

Bắt đầu từ đâu

  1. Đọc hết issue, rồi đọc hướng dẫn đóng góp của dự án.
  2. Bình luận trên issue rằng bạn sẽ nhận — tránh hai người làm cùng một việc.
  3. Fork repository và làm thay đổi trên một nhánh.
  4. Mở pull request có tham chiếu số hiệu của issue.

Issue khác của microsoft/onnxscript

Tất cả issue của microsoft/onnxscript

Issue tương tự

Thêm issue về Python

Nhận issue mới trong hộp thư của bạn

Bản tóm tắt ngắn những issue GitHub phù hợp với người mới.