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

オープン 初心者向け
  • 何か問題がありますか？
#3,155 コメント 3 件 リアクション 1 件 担当者 2 名 GitHub で見る

メンテナーはふだん 4 日以内に返信

@justinchuby がすでに取り組んでいます。

2026年11月5日 から。

  • #3158 @aanelluri による — オープン
  • #3163 @copilot-swe-agent による — オープン

評価

難易度
2/5
見積もり時間
1〜3時間
初心者へのやさしさ
76/100
issue の種類
バグ
明瞭さ
明確に書かれている
活発さ
活発
技術スタック
python
領域
backend, machine-learning

調査の方向性

まず、ONNX Script ORT の fusion ルールにある fuse_skip_rms_normalization と関連テストを読みます。issue には optimize_for_ort を使った再現例があります。このルールが input、skip、scale のデータ型をどのように扱うかを確認し、その後、対象を絞った fusion テストを実行します。データ型が一致しない場合は fusion がスキップされ、提供されたデータ型が一致する float16 のケースでは引き続き fusion が行われ、生成されたグラフが ONNX Runtime に読み込まれることを完了条件とします。

索引モデルが issue の本文から書いたものです。

説明

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.

主要言語
Python
スター
457
フォーク
149
平均マージ
3日 18時間
マージ済み PR（30日）
23

環境構築

はじめの一歩

  1. issue を最後まで読み、次にプロジェクトのコントリビューションガイドを読みます。
  2. 着手することを issue にコメントします — 二人が同じ作業をするのを防げます。
  3. リポジトリをフォークし、ブランチを切って変更します。
  4. issue 番号を参照したプルリクエストを送ります。

microsoft/onnxscript のほかの issue

microsoft/onnxscript の issue をすべて見る

似ている issue

Python の issue をもっと見る

新しい issue をメールで受け取る

初心者向けの GitHub issue を短くまとめたダイジェスト。