[BUG]TensorRT engine export fails with UNSUPPORTED_NODE errors after TE FP8 training
bug
After successfully completing FP8 training and saving the checkpoint (e.g., latest.pth), I am trying to deploy the model using TensorRT. While I can successfully export the model to ONNX, the subsequent conversion from ONNX to a TensorRT engine fails. The TensorRT parser throws several UNSUPPORTED_NODE and UNSUPPORTED_NODE_ATTR errors.
Steps/Code to reproduce bug
Here is the script I use for exporting:
code
```Python
import argparse
import os
import torch
import transformer_engine.pytorch as te
from hydra.utils import instantiate
from omegaconf import OmegaConf
from transformer_engine.pytorch.export import te_translation_table
from flow_planner.data.utils.collect import collect_batch
from flow_planner.export_qat import _apply_trt_onnx_rewrites
from flow_planner.trainer import _build_te_fp8_recipe
from onnx_export.src.export_onnx import OnnxFlowInferenceWrapper, build_export_inputs
OUTPUT_ONNX_NAME = "te_model.onnx"
OUTPUT_ENGINE_NAME = "te_model.engine"
def _apply_trt_rewrites_inplace(onnx_path: str) -> int:
import onnx
onnx_model = onnx.load(onnx_path)
fixed = _apply_trt_onnx_rewrites(onnx_model)
if fixed > 0:
onnx.save(onnx_model, onnx_path)
return fixed
def build_trt_engine(onnx_path: str, engine_path: str, workspace_mb: int, trt_log_level: str) -> None:
import tensorrt as trt
levels = {
"internal_error": trt.Logger.INTERNAL_ERROR,
"error": trt.Logger.ERROR,
"warning": trt.Logger.WARNING,
"info": trt.Logger.INFO,
"verbose": trt.Logger.VERBOSE,
}
logger = trt.Logger(levels[str(trt_log_level).strip().lower()])
builder = trt.Builder(logger)
rewritten_nodes = _apply_trt_rewrites_inplace(onnx_path)
if rewritten_nodes > 0:
print(f"[INFO] Applied {rewritten_nodes} ONNX TensorRT compatibility rewrites before engine build.")
network_flags = 1 << int(trt.NetworkDefinitionCreationFlag.EXPLICIT_BATCH)
if hasattr(trt.NetworkDefinitionCreationFlag, "STRONGLY_TYPED"):
network_flags |= 1 << int(trt.NetworkDefinitionCreationFlag.STRONGLY_TYPED)
network = builder.create_network(network_flags)
parser = trt.OnnxParser(network, logger)
if hasattr(parser, "parse_from_file"):
ok = parser.parse_from_file(onnx_path)
else:
with open(onnx_path, "rb") as f:
ok = parser.parse(f.read())
if not ok:
errors =[str(parser.get_error(i)) for i in range(parser.num_errors)]
raise RuntimeError("Failed to parse ONNX for TensorRT:\n" + "\n".join(errors))
config = builder.create_builder_config()
config.set_memory_pool_limit(trt.MemoryPoolType.WORKSPACE, int(workspace_mb) * 1024 * 1024)
profile = builder.create_optimization_profile()
for i in range(network.num_inputs):
inp = network.get_input(i)
shape = tuple(int(dim) for dim in inp.shape)
profile.set_shape(inp.name, shape, shape, shape)
print(f"[TRT] input shape {inp.name} = {shape}")
config.add_optimization_profile(profile)
serialized = builder.build_serialized_network(network, config)
if serialized is None:
raise RuntimeError("TensorRT build failed: build_serialized_network returned None")
with open(engine_path, "wb") as f:
f.write(serialized)
print(f"TensorRT engine built: {engine_path}")
def parse_args() -> argparse.Namespace:
parser = argparse.ArgumentParser(description="Export TE FP8 checkpoint to ONNX and optional TensorRT engine")
parser.add_argument("--config", type=str, required=True, help="Path to the config file")
parser.add_argument("--checkpoint", type=str, required=True, help="Path to the checkpoint file")
parser.add_argument("--build-engine", action="store_true", help="Build a TensorRT engine after ONNX export")
parser.add_argument("--engine", type=str, default=None, help="Optional TensorRT engine output path")
parser.add_argument("--trt-workspace-mb", type=int, default=4096, help="TensorRT workspace size in MB")
parser.add_argument(
"--trt-log-level",
type=str,
default="warning",
choices=["internal_error", "error", "warning", "info", "verbose"],
help="TensorRT logger verbosity",
)
return parser.parse_args()
def _build_real_export_inputs(
cfg,
model: torch.nn.Module,
device: torch.device,
input_dtype: torch.dtype,
) -> tuple[tuple[torch.Tensor, ...], str] | None:
try:
dataset_cfg = None
dataset_split = ""
for split in ("test", "train"):
dataset_cfg = OmegaConf.select(cfg, f"data.dataset.{split}", default=None)
if dataset_cfg is not None:
dataset_split = split
break
if dataset_cfg is None:
return None
dataset = instantiate(dataset_cfg)
if len(dataset) == 0:
return None
sample = collect_batch([dataset[0]]).to(device)
model_inputs, _ = model.data_processor.sample_to_model_input(
sample,
device=device,
kinematic=model.kinematic,
is_training=False,
)
p = model.planner_params
sample_steps = int(model.flow_ode.sample_params["sample_steps"])
export_inputs = (
model_inputs["neighbor_past"].to(dtype=input_dtype).contiguous(),
model_inputs["lanes"].to(dtype=input_dtype).contiguous(),
model_inputs["lanes_speedlimit"].to(dtype=input_dtype).contiguous(),
model_inputs["lanes_has_speedlimit"].to(dtype=torch.bool).contiguous(),
model_inputs["map_objects"].to(dtype=input_dtype).contiguous(),
model_inputs["routes"].to(dtype=input_dtype).contiguous(),
torch.randn(
1,
int(model.action_num),
int(p["action_len"]),
int(p["state_dim"]),
device=device,
dtype=input_dtype,
),
torch.linspace(0.0, 1.0, sample_steps + 1, device=device, dtype=input_dtype),
torch.tensor([float(model.cfg_weight)], device=device, dtype=input_dtype),
)
return export_inputs, dataset_split
except Exception as exc:
print(f"[WARN] Failed to build real warmup/export inputs, fallback to synthetic sample: {exc}")
return None
def main() -> None:
args = parse_args()
config_path = os.path.abspath(args.config)
checkpoint_path = os.path.abspath(args.checkpoint)
out_dir = os.path.dirname(checkpoint_path)
output_onnx = os.path.join(out_dir, OUTPUT_ONNX_NAME)
output_engine = os.path.abspath(args.engine) if args.engine else os.path.join(out_dir, OUTPUT_ENGINE_NAME)
cfg = OmegaConf.load(config_path)
base_model = instantiate(cfg.model)
ckpt = torch.load(checkpoint_path, map_location="cpu")
base_model.load_state_dict(ckpt["ema_state_dict"])
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
model = base_model.to(device).eval()
wrapper = OnnxFlowInferenceWrapper(model, use_trt_friendly_assemble=True).to(device).eval()
real_export = _build_real_export_inputs(cfg, model, device, input_dtype=torch.float32)
if real_export is None:
export_inputs, _ = build_export_inputs(model, device, input_dtype=torch.float32)
print("[INFO] Warmup/export inputs: synthetic fixed-shape sample.")
else:
export_inputs, dataset_split = real_export
print(f"[INFO] Warmup/export inputs: real preprocessed dataset sample from split={dataset_split}.")
fp8_recipe = _build_te_fp8_recipe(cfg)
with torch.inference_mode(), te.autocast(enabled=True, recipe=fp8_recipe):
# TE requires a forward warmup before ONNX export, using the same FP8 recipe.
_ = wrapper(*export_inputs)
with te.onnx_export(enabled=True):
torch.onnx.export(
wrapper,
export_inputs,
output_onnx,
opset_version=19,
dynamo=True,
custom_translation_table=te_translation_table,
external_data=False,
input_names=[
"neighbor_past",
"lanes",
"lanes_speedlimit",
"lanes_has_speedlimit",
"map_objects",
"routes",
"x_init",
"t_schedule",
"cfg_weight",
],
output_names=["pred"],
)
print(f"ONNX export finished: {output_onnx}")
if args.build_engine:
build_trt_engine(
output_onnx,
output_engine,
workspace_mb=args.trt_workspace_mb,
trt_log_level=args.trt_log_level,
)
print(f"TRT engine export finished: {output_engine}")
else:
print("TRT engine export skipped. Pass --build-engine to build a TensorRT engine.")
if __name__ == "__main__":
main()
```
Exporting to ONNX is successful, but exporting to the TRT Engine fails with the following errors:
```python
04/22/2026-07:14:53] [TRT] [E] ModelImporter.cpp:950: --- Begin node ---
input: "split_303"
input: "val_20"
output: "getitem_743"
name: "node_aten_getitem_20407"
op_type: "aten_getitem"
domain: "pkg.onnxscript.torch_lib"
metadata_props {
key: "namespace"
value: ": onnx_export.src.export_onnx.OnnxFlowInferenceWrapper/getitem_743: <built-in function getitem>"
}
metadata_props {
key: "pkg.torch.onnx.class_hierarchy"
value: "[\'onnx_export.src.export_onnx.OnnxFlowInferenceWrapper\', \'<built-in function getitem>\']"
}
metadata_props {
key: "pkg.torch.onnx.fx_node"
value: "%getitem_743 : [num_users=1] = call_function[target=operator.getitem](args = (%split_303, 1), kwargs = {})"
}
metadata_props {
key: "pkg.torch.onnx.name_scopes"
value: "[\'\', \'getitem_743\']"
}
metadata_props {
key: "pkg.torch.onnx.stack_trace"
value: "File \"/kargobot-vepfs-zone-c/common_rw/yanchaowei1_v/Flow-Planner/onnx_export/src/export_onnx.py\", line 540, in forward\n x = self._step_update(x, t_schedule[3:4], t_schedule[4:5], cfg_weight, enc0_2, enc1_2, mask0_2, mask1_2, routes_cond0_2, routes_mask_2, token_dist_2)"
}
[04/22/2026-07:14:53] [TRT] [E] ModelImporter.cpp:951: --- End node ---
[04/22/2026-07:14:53] [TRT] [E] ModelImporter.cpp:954: ERROR: onnxOpCheckers.cpp:986 In function checkSequenceAt:
[8] false
Traceback (most recent call last):
File "/kargobot-vepfs-zone-c/common_rw/yanchaowei1_v/Flow-Planner/flow_planner/core/TE_export/export_onnx.py", line 202, in <module>
main()
File "/kargobot-vepfs-zone-c/common_rw/yanchaowei1_v/Flow-Planner/flow_planner/core/TE_export/export_onnx.py", line 192, in main
build_trt_engine(
File "/kargobot-vepfs-zone-c/common_rw/yanchaowei1_v/Flow-Planner/flow_planner/core/TE_export/export_onnx.py", line 44, in build_trt_engine
raise RuntimeError("Failed to parse ONNX for TensorRT:\n" + "\n".join(errors))
RuntimeError: Failed to parse ONNX for TensorRT:
In node 100 with name: node_ScatterND_100 and operator: ScatterND (importScatterND): UNSUPPORTED_NODE_ATTR: Assertion failed: !attrs.count("reduction"): Attribute reduction is not supported.
In node 2389 with name: n0 and operator: SplitToSequence (checkSplitToSequence): UNSUPPORTED_NODE: false
In node 2390 with name: n0 and operator: SequenceAt (checkSequenceAt): UNSUPPORTED_NODE: false
In node 2391 with name: n0 and operator: SequenceAt (checkSequenceAt): UNSUPPORTED_NODE: false
In node 2397 with name: n0 and operator: SplitToSequence (checkSplitToSequence): UNSUPPORTED_NODE: false
In node 2398 with name: n0 and operator: SequenceAt (checkSequenceAt): UNSUPPORTED_NODE: false
In node 2399 with name: n0 and operator: SequenceAt (checkSequenceAt): UNSUPPORTED_NODE: false
In node 2406 with name: n0 and operator: SplitToSequence (checkSplitToSequence): UNSUPPORTED_NODE: false
In node 2407 with name: n0 and operator: SequenceAt (checkSequenceAt): UNSUPPORTED_NODE: false
In node 2408 with name: n0 and operator: SequenceAt (checkSequenceAt): UNSUPPORTED_NODE: false
```
Expected behavior
I would like to confirm whether Transformer Engine (TE) supports TensorRT engine export directly after FP8 training.
Specifically, I'd like to clarify whether the TE FP8 training workflow inherently serves as an FP8 quantizer for TensorRT deployment.
Environment overview
```
Environment location: local machine
Method of Transformer Engine install: pip install
OS: Ubuntu 22.04
Python: 3.11
PyTorch: 2.6.0+cu124
Transformer Engine: 2.13.0
Device details
GPU model:RTX4090
TensorRT Version: 10.4.0
```
Additional context
The error trace indicates that the TensorRT parser is failing due to unsupported ONNX operators such as ScatterND (with the reduction attribute), SplitToSequence, and SequenceAt. I am using ONNX opset 19.
Thank you so much if you could solve this.
关闭于 2026-05-06 2 条评论