wehub-resource-sync
219 lines
8.1 KiB
Python
219 lines
8.1 KiB
Python
#
|
|
# SPDX-FileCopyrightText: Copyright (c) 1993-2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
|
|
# SPDX-License-Identifier: Apache-2.0
|
|
#
|
|
# Licensed under the Apache License, Version 2.0 (the "License");
|
|
# you may not use this file except in compliance with the License.
|
|
# You may obtain a copy of the License at
|
|
#
|
|
# http://www.apache.org/licenses/LICENSE-2.0
|
|
#
|
|
# Unless required by applicable law or agreed to in writing, software
|
|
# distributed under the License is distributed on an "AS IS" BASIS,
|
|
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
|
# See the License for the specific language governing permissions and
|
|
# limitations under the License.
|
|
#
|
|
|
|
import os
|
|
import re
|
|
import tempfile
|
|
|
|
import onnx
|
|
import onnx_graphsurgeon as gs
|
|
import torch
|
|
from onnx import shape_inference
|
|
from polygraphy.backend.onnx.loader import fold_constants
|
|
|
|
from demo_diffusion.model import load
|
|
from demo_diffusion.utils_modelopt import (
|
|
cast_convtranspose_io,
|
|
cast_fp8_mha_io,
|
|
cast_layernorm_io,
|
|
cast_resize_io,
|
|
convert_fp16_io,
|
|
convert_zp_fp8,
|
|
)
|
|
|
|
# FIXME update callsites after serialization support for torch.compile is added
|
|
TORCH_INFERENCE_MODELS = ["default", "reduce-overhead", "max-autotune"]
|
|
|
|
|
|
def optimize_checkpoint(model, torch_inference: str):
|
|
"""Optimize a torch model checkpoint using torch.compile."""
|
|
if not torch_inference or torch_inference == "eager":
|
|
return model
|
|
assert torch_inference in TORCH_INFERENCE_MODELS
|
|
return torch.compile(model, mode=torch_inference, dynamic=False, fullgraph=False)
|
|
|
|
|
|
class Optimizer:
|
|
|
|
def __init__(self, onnx_graph, verbose=False, version=None):
|
|
self.graph = gs.import_onnx(onnx_graph)
|
|
self.verbose = verbose
|
|
self.version = version
|
|
|
|
def info(self, prefix):
|
|
if self.verbose:
|
|
print(
|
|
f"{prefix} .. {len(self.graph.nodes)} nodes, {len(self.graph.tensors().keys())} tensors, {len(self.graph.inputs)} inputs, {len(self.graph.outputs)} outputs"
|
|
)
|
|
|
|
def cleanup(self, return_onnx=False):
|
|
self.graph.cleanup().toposort()
|
|
return gs.export_onnx(self.graph) if return_onnx else self.graph
|
|
|
|
def select_outputs(self, keep, names=None):
|
|
self.graph.outputs = [self.graph.outputs[o] for o in keep]
|
|
if names:
|
|
for i, name in enumerate(names):
|
|
self.graph.outputs[i].name = name
|
|
|
|
def fold_constants(self, return_onnx=False):
|
|
onnx_graph = fold_constants(gs.export_onnx(self.graph), allow_onnxruntime_shape_inference=True)
|
|
self.graph = gs.import_onnx(onnx_graph)
|
|
if return_onnx:
|
|
return onnx_graph
|
|
|
|
def infer_shapes(self, return_onnx=False):
|
|
onnx_graph = gs.export_onnx(self.graph)
|
|
if load.onnx_graph_needs_external_data(onnx_graph):
|
|
temp_dir = tempfile.TemporaryDirectory().name
|
|
os.makedirs(temp_dir, exist_ok=True)
|
|
onnx_orig_path = os.path.join(temp_dir, "model.onnx")
|
|
onnx_inferred_path = os.path.join(temp_dir, "inferred.onnx")
|
|
onnx.save_model(
|
|
onnx_graph,
|
|
onnx_orig_path,
|
|
save_as_external_data=True,
|
|
all_tensors_to_one_file=True,
|
|
convert_attribute=False,
|
|
)
|
|
onnx.shape_inference.infer_shapes_path(onnx_orig_path, onnx_inferred_path)
|
|
onnx_graph = onnx.load(onnx_inferred_path)
|
|
else:
|
|
onnx_graph = shape_inference.infer_shapes(onnx_graph)
|
|
|
|
self.graph = gs.import_onnx(onnx_graph)
|
|
if return_onnx:
|
|
return onnx_graph
|
|
|
|
def clip_add_hidden_states(self, hidden_layer_offset, return_onnx=False):
|
|
hidden_layers = -1
|
|
onnx_graph = gs.export_onnx(self.graph)
|
|
for i in range(len(onnx_graph.graph.node)):
|
|
for j in range(len(onnx_graph.graph.node[i].output)):
|
|
name = onnx_graph.graph.node[i].output[j]
|
|
if "layers" in name:
|
|
hidden_layers = max(int(name.split(".")[1].split("/")[0]), hidden_layers)
|
|
for i in range(len(onnx_graph.graph.node)):
|
|
for j in range(len(onnx_graph.graph.node[i].output)):
|
|
if onnx_graph.graph.node[i].output[j] == "/text_model/encoder/layers.{}/Add_1_output_0".format(
|
|
hidden_layers + hidden_layer_offset
|
|
):
|
|
onnx_graph.graph.node[i].output[j] = "hidden_states"
|
|
for j in range(len(onnx_graph.graph.node[i].input)):
|
|
if onnx_graph.graph.node[i].input[j] == "/text_model/encoder/layers.{}/Add_1_output_0".format(
|
|
hidden_layers + hidden_layer_offset
|
|
):
|
|
onnx_graph.graph.node[i].input[j] = "hidden_states"
|
|
if return_onnx:
|
|
return onnx_graph
|
|
|
|
def fuse_mha_qkv_int8_sq(self):
|
|
tensors = self.graph.tensors()
|
|
keys = tensors.keys()
|
|
|
|
# mha : fuse QKV QDQ nodes
|
|
# mhca : fuse KV QDQ nodes
|
|
q_pat = (
|
|
"/down_blocks.\\d+/attentions.\\d+/transformer_blocks"
|
|
".\\d+/attn\\d+/to_q/input_quantizer/DequantizeLinear_output_0"
|
|
)
|
|
k_pat = (
|
|
"/down_blocks.\\d+/attentions.\\d+/transformer_blocks"
|
|
".\\d+/attn\\d+/to_k/input_quantizer/DequantizeLinear_output_0"
|
|
)
|
|
v_pat = (
|
|
"/down_blocks.\\d+/attentions.\\d+/transformer_blocks"
|
|
".\\d+/attn\\d+/to_v/input_quantizer/DequantizeLinear_output_0"
|
|
)
|
|
|
|
qs = list(
|
|
sorted(
|
|
map(
|
|
lambda x: x.group(0), # type: ignore
|
|
filter(lambda x: x is not None, [re.match(q_pat, key) for key in keys]),
|
|
)
|
|
)
|
|
)
|
|
ks = list(
|
|
sorted(
|
|
map(
|
|
lambda x: x.group(0), # type: ignore
|
|
filter(lambda x: x is not None, [re.match(k_pat, key) for key in keys]),
|
|
)
|
|
)
|
|
)
|
|
vs = list(
|
|
sorted(
|
|
map(
|
|
lambda x: x.group(0), # type: ignore
|
|
filter(lambda x: x is not None, [re.match(v_pat, key) for key in keys]),
|
|
)
|
|
)
|
|
)
|
|
|
|
removed = 0
|
|
assert len(qs) == len(ks) == len(vs), "Failed to collect tensors"
|
|
for q, k, v in zip(qs, ks, vs):
|
|
is_mha = all(["attn1" in tensor for tensor in [q, k, v]])
|
|
is_mhca = all(["attn2" in tensor for tensor in [q, k, v]])
|
|
assert (is_mha or is_mhca) and (not (is_mha and is_mhca))
|
|
|
|
if is_mha:
|
|
tensors[k].outputs[0].inputs[0] = tensors[q]
|
|
tensors[v].outputs[0].inputs[0] = tensors[q]
|
|
del tensors[k]
|
|
del tensors[v]
|
|
removed += 2
|
|
else: # is_mhca
|
|
tensors[k].outputs[0].inputs[0] = tensors[v]
|
|
del tensors[k]
|
|
removed += 1
|
|
print(f"Removed {removed} QDQ nodes")
|
|
return removed # expected 72 for L2.5
|
|
|
|
def modify_int8_graph(self):
|
|
# Cast LayerNorm scale/bias from FP16 to FP32 to match INT8 DQ activations.
|
|
cast_layernorm_io(self.graph)
|
|
# Fuse QKV QDQ nodes for INT8 SmoothQuant.
|
|
self.fuse_mha_qkv_int8_sq()
|
|
|
|
def cast_convtranspose_io(self):
|
|
cast_convtranspose_io(self.graph)
|
|
|
|
def cast_resize_io(self, output_dtype):
|
|
cast_resize_io(self.graph, output_dtype=output_dtype)
|
|
|
|
def modify_fp8_graph(self, is_fp16_io=True):
|
|
onnx_graph = gs.export_onnx(self.graph)
|
|
# Convert INT8 Zero to FP8.
|
|
onnx_graph = convert_zp_fp8(onnx_graph)
|
|
|
|
self.graph = gs.import_onnx(onnx_graph)
|
|
# Add cast nodes to Resize I/O.
|
|
cast_resize_io(self.graph)
|
|
# Convert model inputs and outputs to fp16 I/O.
|
|
if is_fp16_io:
|
|
convert_fp16_io(self.graph)
|
|
# Add cast nodes to MHA's BMM1 and BMM2's I/O.
|
|
cast_fp8_mha_io(self.graph)
|
|
|
|
def flux_convert_rope_weight_type(self):
|
|
for node in self.graph.nodes:
|
|
if node.op == "Einsum":
|
|
print(f"Fixed RoPE (Rotary Position Embedding) weight type: {node.name}")
|
|
return gs.export_onnx(self.graph)
|