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

Aperta
  • Qualcosa non va?
#3,154 0 commenti 0 reazioni 0 assegnatari Vedi su GitHub

I maintainer di solito rispondono entro 4 giorni

Nessuno ha ancora preso questa issue.

Valutazione

Difficoltà
4/5
Tempo stimato
3-5 giorni
Idoneità per principianti
68/100
Tipo di issue
Bug
Chiarezza
Specificata chiaramente
Stato di attività
Attiva
Stack tecnologico
python
Ambito
machine-learning, performance

Direzione di ricerca

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.

Scritto dal modello di indicizzazione a partire dal testo della issue.

Descrizione

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.
Lingua principale
Python
Stelle
457
Fork
149
Merge medio
3g 20h
PR unite (30g)
22

Preparare l'ambiente

Come iniziare

  1. Leggi tutta la issue e poi la guida ai contributi del progetto.
  2. Commenta sulla issue per dire che te ne occupi tu — evita che due persone facciano lo stesso lavoro.
  3. Fai un fork del repository e lavora su un branch.
  4. Apri una pull request che faccia riferimento al numero della issue.

Altre issue di microsoft/onnxscript

Tutte le issue di microsoft/onnxscript

Issue simili

Altre issue su Python

Ricevi le nuove issue nella tua casella

Un breve riepilogo di issue GitHub adatte ai principianti.