`fuse_skip_rms_normalization` fuses a float16 residual into a norm whose scale is float32, which ONNX Runtime rejects

Aperta Adatta ai principianti
  • Qualcosa non va?
#3,155 3 commenti 1 reazione 2 assegnatari Vedi su GitHub

I maintainer di solito rispondono entro 4 giorni

@justinchuby ci sta già lavorando.

Dal 5/11/2026.

  • #3158 di @aanelluri — aperta
  • #3163 di @copilot-swe-agent — aperta

Valutazione

Difficoltà
2/5
Tempo stimato
1-3 ore
Idoneità per principianti
76/100
Tipo di issue
Bug
Chiarezza
Specificata chiaramente
Stato di attività
Attiva
Stack tecnologico
python
Ambito
backend, machine-learning

Direzione di ricerca

Inizia leggendo fuse_skip_rms_normalization nelle regole di fusione ONNX Script ORT e i test correlati; l’issue fornisce un riproduttore che usa optimize_for_ort. Verifica come la regola gestisce i tipi di dati di input, skip e scale, quindi esegui i test di fusione mirati. Il risultato è raggiunto quando la fusione viene saltata in caso di tipi di dati non corrispondenti, mentre il caso float16 fornito, con tipi corrispondenti, continua a fondersi e il grafo risultante viene caricato in ONNX Runtime.

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

Descrizione

module: optimizer

In a float16 model whose RMS norm computes in float32 and scales by a float32 1 + weight (Gemma's), optimize_for_ort fuses the residual add and the norm into a SkipSimplifiedLayerNormalization with float16 input and skip but a float32 scale. ONNX Runtime refuses to load the session. The rule should check that the input, skip and scale dtypes match before fusing.

Repro (torch 2.13.0, onnxscript 0.7.2, onnxruntime 1.31.1):

"""onnxscript's skip RMS norm fusion fuses a float16 residual add into a norm computed in float32, which gives a
`SkipSimplifiedLayerNormalization` whose inputs have different dtypes, and ONNX Runtime rejects the graph.

The module is a Gemma-style residual block in float16: `x + y`, then an RMS norm computed in float32, scaled by
`1 + weight` in float32, and cast back.
"""

import collections

import onnx_ir as ir
import onnxruntime as ort
import torch
from onnxscript.rewriter.ort_fusions import optimize_for_ort


class Block(torch.nn.Module):
    def __init__(self, hidden=64):
        super().__init__()
        self.weight = torch.nn.Parameter(torch.zeros(hidden, dtype=torch.float16))

    def forward(self, x, y):
        residual = x + y
        states = residual.to(torch.float32)
        states = states * torch.rsqrt(states.pow(2).mean(-1, keepdim=True) + 1e-6)
        return (states * (1.0 + self.weight.float())).to(torch.float16), residual


inputs = (torch.randn(2, 8, 64, dtype=torch.float16), torch.randn(2, 8, 64, dtype=torch.float16))
program = torch.onnx.export(Block(), inputs, dynamo=True)
model, counts = optimize_for_ort(program.model)
for node in model.graph.all_nodes():
    if node.op_type == "SkipSimplifiedLayerNormalization":
        print("fused:", node.op_type, "input dtypes:", [str(value.dtype) for value in node.inputs[:3]])
print("ops:", dict(collections.Counter(node.op_type for node in model.graph.all_nodes())))
try:
    ort.InferenceSession(ir.to_proto(model).SerializeToString(), providers=["CPUExecutionProvider"])
    print("ONNX Runtime loaded the fused graph")
except Exception as error:
    print("ONNX Runtime:", str(error).splitlines()[0][:300])

fused: SkipSimplifiedLayerNormalization input dtypes: ['FLOAT16', 'FLOAT16', 'FLOAT']
ONNX Runtime: [ONNXRuntimeError] : 1 : FAIL : Type Error: Type parameter (T) of Optype (SkipSimplifiedLayerNormalization) bound to different types (tensor(float16) and tensor(float) in node (node_SkipSimplifiedLayerNormalization_8).

With a Llama-style norm (float16 weight applied after the cast back) the fusion is valid, and all three inputs are float16.

Lingua principale
Python
Stelle
457
Fork
149
Merge medio
3g 16h
PR unite (30g)
24

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.