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

Abierto Apto para principiantes
  • ¿Algo no está bien?
#3,155 0 comentarios 0 reacciones 0 asignados Ver en GitHub

Los mantenedores suelen responder en 4 días

Nadie ha tomado este issue todavía.

Evaluación

Dificultad
2/5
Tiempo estimado
1-3 horas
Aptitud para principiantes
76/100
Tipo de issue
Error
Claridad
Bien especificado
Estado de actividad
Activo
Stack tecnológico
python
Área
backend, machine-learning

Línea de trabajo

Empieza leyendo fuse_skip_rms_normalization en las reglas de fusión de ONNX Script ORT y las pruebas relacionadas; el issue incluye un caso reproducible que usa optimize_for_ort. Comprueba cómo trata la regla los tipos de datos de input, skip y scale, y luego ejecuta las pruebas de fusión específicas. El objetivo se cumple cuando se omite la fusión si los tipos de datos no coinciden, mientras que el caso float16 proporcionado, con tipos coincidentes, sigue fusionándose y el grafo resultante se carga en ONNX Runtime.

Escrito por el modelo de indexación a partir del texto del issue.

Descripción

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.

Lenguaje dominante
Python
Estrellas
457
Forks
149
Merge medio
3 d 20 h
PR fusionados (30 d)
22

Preparar el entorno

Primeros pasos

  1. Lee el issue completo y luego la guía de contribución del proyecto.
  2. Comenta en el issue que vas a ocuparte — evita que dos personas hagan lo mismo.
  3. Haz un fork del repositorio y trabaja en una rama.
  4. Abre un pull request que haga referencia al número del issue.

Más de microsoft/onnxscript

Todos los issues de microsoft/onnxscript

Issues similares

Más issues de Python

Recibe los nuevos issues en tu correo

Un resumen breve de issues de GitHub para principiantes.