#!/usr/bin/env python3 # SPDX-FileCopyrightText: Copyright (c) 2026 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. """ Build TensorRT engines from exported ONNX models. Supports two modes: - single: Build engine for a single ONNX model - full_pipeline: Build engines for all pipeline components (ViT, LLM, State Encoder, Action Encoder, DiT, Action Decoder) Shape profiles are automatically derived from the ONNX models. Usage: # Full pipeline: python scripts/deployment/build_tensorrt_engine.py \ --mode full_pipeline \ --onnx-dir ./gr00t_trt_deployment/onnx \ --engine-dir ./gr00t_trt_deployment/engines \ --precision bf16 """ from dataclasses import dataclass import logging import os import time from typing import Literal from _trt_contract import load_export_metadata, validate_export_metadata from gr00t.deployment.modes import FULL_PIPELINE_COMPONENTS, BuildEngineMode import onnx import tensorrt as trt import tyro # Set up logging logging.basicConfig(level=logging.INFO, format="%(asctime)s - %(levelname)s - %(message)s") logger = logging.getLogger(__name__) # STRONGLY_TYPED precision sanity check: TRT 10+ STRONGLY_TYPED reads # precision from the ONNX tensor types and ignores --precision builder # flags. Catch the silent mismatch (user asks fp16, ONNX is bf16, engine # silently builds bf16) before burning build time. Indirected through # dtype *names* so the helper can be unit-tested without TensorRT. _PRECISION_TO_TRT_DTYPE_NAME: dict[str, str] = { "bf16": "BF16", "fp16": "HALF", "fp32": "FLOAT", "fp8": "FP8", } def _check_strongly_typed_precision_match( network_dtype_names: set[str], requested_precision: str ) -> None: """Raise if --precision cannot be honored by this STRONGLY_TYPED network.""" expected = _PRECISION_TO_TRT_DTYPE_NAME.get(requested_precision) if expected is None: raise ValueError( f"Unknown precision: {requested_precision!r}. " f"Expected one of {sorted(_PRECISION_TO_TRT_DTYPE_NAME)}." ) if expected not in network_dtype_names: raise ValueError( f"--precision={requested_precision} cannot be honored by this ONNX. " f"STRONGLY_TYPED (TRT 10+) reads precision from ONNX tensor types " f"and ignores builder flags. Network has tensor dtypes " f"{sorted(network_dtype_names)}; none of them are {expected}. " f"Either re-export the ONNX with the requested precision, or " f"pass --precision matching the existing ONNX dtypes." ) # When fp32 is requested, the network must not contain any reduced-precision # tensors. STRONGLY_TYPED won't promote BF16/FP16/FP8 to FLOAT, so a mixed # BF16+FLOAT network silently runs at BF16 for those tensors despite the # caller asking for fp32. if requested_precision == "fp32": reduced = {"BF16", "HALF", "FP8"} & network_dtype_names if reduced: raise ValueError( f"--precision=fp32 cannot be honored: network also contains " f"reduced-precision tensors {sorted(reduced)}. STRONGLY_TYPED " f"won't promote them to FLOAT, so the engine would silently " f"run mixed precision. Re-export the ONNX as pure FP32, or " f"pass --precision matching the dominant reduced dtype." ) def _precision_from_onnx_path(onnx_path: str, default: str) -> str: """Return the precision tag suffixed in the ONNX filename (e.g. ``vit_fp32.onnx`` → ``"fp32"``), else ``default``. Used so the full-pipeline build mirrors the export's per-component dtype instead of forwarding the pipeline-wide ``--precision`` to a mismatched ONNX. """ stem = os.path.splitext(os.path.basename(onnx_path))[0] for tag in _PRECISION_TO_TRT_DTYPE_NAME: if stem.endswith(f"_{tag}"): return tag return default # ============================================================ # Auto Shape Profile from ONNX # ============================================================ def derive_shapes_from_onnx(onnx_path, max_batch=8): """Read an ONNX model and derive min/opt/max shape profiles. For each input: - Fixed dimensions (concrete values) are kept as-is across min/opt/max. - Dynamic batch dimension: min=1, opt=1, max=max_batch. - Dynamic sequence dimensions: min=1, opt=concrete_value, max=2*concrete_value. (concrete_value comes from the ONNX model's shape hints) Returns (min_shapes, opt_shapes, max_shapes) dicts. """ model = onnx.load(onnx_path, load_external_data=False) min_shapes, opt_shapes, max_shapes = {}, {}, {} for inp in model.graph.input: name = inp.name dims = inp.type.tensor_type.shape.dim min_shape, opt_shape, max_shape = [], [], [] for i, d in enumerate(dims): if d.dim_value > 0: # Fixed dimension — use as-is min_shape.append(d.dim_value) opt_shape.append(d.dim_value) max_shape.append(d.dim_value) else: # Dynamic dimension if i == 0: # Batch dimension min_shape.append(1) opt_shape.append(1) max_shape.append(max_batch) else: # Sequence/spatial dimension — use generous range # We don't know the "typical" value from ONNX alone, # so use 1 / 1 / large_max. The builder will optimize for opt. min_shape.append(1) opt_shape.append(1) max_shape.append(512) min_shapes[name] = tuple(min_shape) opt_shapes[name] = tuple(opt_shape) max_shapes[name] = tuple(max_shape) return min_shapes, opt_shapes, max_shapes def derive_shapes_with_hint(onnx_path, opt_seq_lens=None, max_batch=8): """Derive shapes from ONNX, with optional sequence length hints. Args: onnx_path: Path to ONNX model opt_seq_lens: Dict mapping dynamic dim names to optimal sequence lengths. e.g. {"sa_seq_len": 51, "vl_seq_len": 280, "sequence_length": 280} max_batch: Maximum batch size """ model = onnx.load(onnx_path, load_external_data=False) opt_seq_lens = opt_seq_lens or {} min_shapes, opt_shapes, max_shapes = {}, {}, {} for inp in model.graph.input: name = inp.name dims = inp.type.tensor_type.shape.dim min_shape, opt_shape, max_shape = [], [], [] for i, d in enumerate(dims): if d.dim_value > 0: # Fixed dimension min_shape.append(d.dim_value) opt_shape.append(d.dim_value) max_shape.append(d.dim_value) else: dim_name = d.dim_param if d.dim_param else f"dim_{i}" if dim_name == "batch_size": # Batch dimension (at any index) min_shape.append(1) opt_shape.append(1) max_shape.append(max_batch) elif dim_name in opt_seq_lens: # Named dynamic dim with a hint opt_val = opt_seq_lens[dim_name] min_shape.append(1) opt_shape.append(opt_val) max_shape.append(max(opt_val * 2, opt_val + 64)) else: # Unknown dynamic dim — use wide range min_shape.append(1) opt_shape.append(256) max_shape.append(512) min_shapes[name] = tuple(min_shape) opt_shapes[name] = tuple(opt_shape) max_shapes[name] = tuple(max_shape) return min_shapes, opt_shapes, max_shapes # ============================================================ # Engine Builder # ============================================================ def build_engine( onnx_path: str, engine_path: str, precision: str = "bf16", workspace_mb: int = 8192, min_shapes: dict = None, opt_shapes: dict = None, max_shapes: dict = None, trt_severity=None, ): """Build TensorRT engine from ONNX model. Args: onnx_path: Path to ONNX model engine_path: Path to save TensorRT engine precision: Precision mode ('fp32', 'fp16', 'bf16', 'fp8') workspace_mb: Workspace size in MB min_shapes: Minimum input shapes (dict: name -> shape tuple) opt_shapes: Optimal input shapes (dict: name -> shape tuple) max_shapes: Maximum input shapes (dict: name -> shape tuple) """ logger.info("=" * 80) logger.info("TensorRT Engine Builder") logger.info("=" * 80) logger.info(f"ONNX model: {onnx_path}") logger.info(f"Engine output: {engine_path}") logger.info(f"Precision: {precision.upper()}") logger.info(f"Workspace: {workspace_mb} MB") logger.info("=" * 80) TRT_LOGGER = trt.Logger(trt.Logger.VERBOSE if trt_severity is None else trt_severity) # Create builder and network logger.info("\n[Step 1/5] Creating TensorRT builder...") builder = trt.Builder(TRT_LOGGER) # TRT 10.x prefers STRONGLY_TYPED; EXPLICIT_BATCH is the 9.x fallback. use_strongly_typed = hasattr(trt.NetworkDefinitionCreationFlag, "STRONGLY_TYPED") if use_strongly_typed: network_flags = 1 << int(trt.NetworkDefinitionCreationFlag.STRONGLY_TYPED) logger.info("Using STRONGLY_TYPED network (TRT 10.x+)") else: network_flags = 1 << int(trt.NetworkDefinitionCreationFlag.EXPLICIT_BATCH) logger.info("Using EXPLICIT_BATCH network (TRT 9.x fallback)") network = builder.create_network(network_flags) parser = trt.OnnxParser(network, TRT_LOGGER) # Parse ONNX model logger.info("\n[Step 2/5] Parsing ONNX model...") if not parser.parse_from_file(onnx_path): logger.error("Failed to parse ONNX file") for error in range(parser.num_errors): logger.error(parser.get_error(error)) raise RuntimeError("ONNX parsing failed") logger.info(f"Network inputs: {network.num_inputs}") for i in range(network.num_inputs): inp = network.get_input(i) logger.info(f" Input {i}: {inp.name} {inp.shape}") logger.info(f"Network outputs: {network.num_outputs}") for i in range(network.num_outputs): out = network.get_output(i) logger.info(f" Output {i}: {out.name} {out.shape}") # Create builder config logger.info("\n[Step 3/5] Configuring builder...") config = builder.create_builder_config() config.profiling_verbosity = trt.ProfilingVerbosity.DETAILED logger.info("Enabled DETAILED profiling verbosity for engine inspection") config.set_memory_pool_limit(trt.MemoryPoolType.WORKSPACE, workspace_mb * (1024**2)) if use_strongly_typed: network_dtype_names: set[str] = set() for i in range(network.num_inputs): network_dtype_names.add(network.get_input(i).dtype.name) for i in range(network.num_outputs): network_dtype_names.add(network.get_output(i).dtype.name) _check_strongly_typed_precision_match(network_dtype_names, precision) logger.info( f"Precision '{precision}' matches ONNX tensor dtypes (STRONGLY_TYPED, " f"network has {sorted(network_dtype_names)})" ) else: # Weak-typed fallback: explicitly set precision flags if precision == "fp16": config.set_flag(trt.BuilderFlag.FP16) logger.info("Enabled FP16 mode") elif precision == "bf16": config.set_flag(trt.BuilderFlag.BF16) logger.info("Enabled BF16 mode") elif precision == "fp8": config.set_flag(trt.BuilderFlag.FP8) config.set_flag(trt.BuilderFlag.BF16) logger.info("Enabled FP8 + BF16 mode") elif precision == "fp32": logger.info("Using FP32 (default precision)") else: raise ValueError(f"Unknown precision: {precision}") # Set optimization profiles for dynamic shapes if min_shapes and opt_shapes and max_shapes: logger.info("\n[Step 4/5] Setting optimization profiles...") profile = builder.create_optimization_profile() for i in range(network.num_inputs): inp = network.get_input(i) input_name = inp.name if input_name in min_shapes: min_shape = min_shapes[input_name] opt_shape = opt_shapes[input_name] max_shape = max_shapes[input_name] profile.set_shape(input_name, min_shape, opt_shape, max_shape) logger.info(f" {input_name}:") logger.info(f" min: {min_shape}") logger.info(f" opt: {opt_shape}") logger.info(f" max: {max_shape}") config.add_optimization_profile(profile) else: raise RuntimeError("Provide min/max and opt shapes for dynamic axes") # Build engine logger.info("\n[Step 5/5] Building TensorRT engine...") start_time = time.time() serialized_engine = builder.build_serialized_network(network, config) build_time = time.time() - start_time if serialized_engine is None: raise RuntimeError("Failed to build TensorRT engine") logger.info(f"Engine built in {build_time:.1f} seconds ({build_time / 60:.1f} minutes)") # Save engine logger.info(f"\nSaving engine to {engine_path}...") os.makedirs(os.path.dirname(engine_path) or ".", exist_ok=True) with open(engine_path, "wb") as f: f.write(serialized_engine) engine_size_mb = os.path.getsize(engine_path) / (1024**2) logger.info(f"Engine saved! Size: {engine_size_mb:.2f} MB") logger.info("\n" + "=" * 80) logger.info("ENGINE BUILD COMPLETE!") logger.info("=" * 80) logger.info(f"Engine file: {engine_path}") logger.info(f"Size: {engine_size_mb:.2f} MB") logger.info(f"Build time: {build_time:.1f}s") logger.info(f"Precision: {precision.upper()}") logger.info("=" * 80) return engine_path # ============================================================ # Full Pipeline Builder # ============================================================ def build_full_pipeline( onnx_dir, engine_dir, precision="bf16", workspace_mb=8192, trt_severity=None, only: frozenset[str] | None = None, allow_default_hints: bool = False, ): """Build all TRT engines for the full pipeline. Shape profiles are automatically derived from the ONNX models. Dynamic sequence dimensions use hints based on typical inference shapes. Args: onnx_dir: Directory containing exported ONNX models engine_dir: Directory to save TRT engines precision: Precision mode workspace_mb: Workspace size in MB only: Restrict the build to this subset of component names (from ``FULL_PIPELINE_COMPONENTS``). ``None`` builds the full 7. A partial export (e.g. ``action_head``, which keeps ViT/LLM in PyTorch) must pass its produced subset so the completeness check requires exactly those, not the full pipeline. """ os.makedirs(engine_dir, exist_ok=True) # Sequence/patch hints for the TRT shape profiles come from export_metadata.json # (single source of truth). A missing, stale, or incomplete bundle is rejected so # the build can't silently bake wrong shapes; --allow-default-hints opts into the # hardcoded GR1 single-view fallbacks. metadata = load_export_metadata(onnx_dir) try: if metadata is None: raise ValueError(f"no export_metadata.json found in {onnx_dir}") validate_export_metadata(metadata, source="build_full_pipeline", engine_path=onnx_dir) except ValueError as e: if not allow_default_hints: raise ValueError( f"{e}. Re-export with the current exporter, or pass --allow-default-hints " "to build with hardcoded GR1 single-view shape hints (the engine may get " "wrong sequence/patch shapes)." ) from e logger.warning("%s; using default shape hints (--allow-default-hints).", e) metadata = None if metadata is not None: # The engine must be built at the precision it was exported for; a drift here # produces a valid-but-wrong engine. Per-component precision is still taken from # each ONNX filename below — this guards the pipeline-wide default. if metadata["precision"] != precision: raise ValueError( f"build_full_pipeline: --precision={precision} but export_metadata.json in " f"{onnx_dir} records precision={metadata['precision']!r}. Build at the " f"exported precision (--precision {metadata['precision']}) or re-export." ) seq_hints = { "sa_seq_len": metadata["sa_seq_len"], "vl_seq_len": metadata["vl_seq_len"], "sequence_length": metadata["llm_seq_len"], "seq_len": metadata["llm_seq_len"], # N1.7 LLM dynamic dim name "num_patches": metadata["num_patches"], "num_merged_patches": metadata["num_merged_patches"], "num_vis_tokens": metadata["num_vis_tokens"], # N1.7 deepstack } logger.info(f"Loaded shape hints from export_metadata.json in {onnx_dir}: {seq_hints}") else: seq_hints = { "sa_seq_len": 51, # 1 state + action_horizon "vl_seq_len": 280, # typical backbone output seq_len "sequence_length": 280, # LLM seq_len } logger.warning(f"Using default shape hints (no usable metadata): {seq_hints}") # Build order, ONNX candidates, and engine filenames come from the shared # component table (single source of truth). ``only`` restricts to the subset # a partial export produced. FP32 ViT is preferred for accuracy and falls # back to BF16; the engine filename stays precision-neutral (vit.engine) # because the input ONNX may be either FP32 or BF16 — the real precision is # recorded in export_metadata.json and inspectable via TRT tooling. if only is not None: valid_names = {c.name for c in FULL_PIPELINE_COMPONENTS} unknown = set(only) - valid_names if unknown: raise ValueError( f"Unknown pipeline component(s) {sorted(unknown)}; " f"valid components: {sorted(valid_names)}" ) components: list[tuple[str, str, str]] = [] for component in FULL_PIPELINE_COMPONENTS: if only is not None and component.name not in only: continue onnx_file = next( (c for c in component.onnx_candidates if os.path.exists(os.path.join(onnx_dir, c))), component.onnx_candidates[0], ) components.append((component.name, onnx_file, component.engine)) results: list[tuple[str, str, str]] = [] skipped: list[tuple[str, str]] = [] # (name, onnx_path) for components with no ONNX input for name, onnx_file, engine_file in components: onnx_path = os.path.join(onnx_dir, onnx_file) if not os.path.exists(onnx_path): logger.warning(f"Skipping {name}: ONNX file not found at {onnx_path}") skipped.append((name, onnx_path)) continue logger.info(f"\n{'#' * 80}") logger.info(f"# Building {name} engine") logger.info(f"{'#' * 80}") engine_path = os.path.join(engine_dir, engine_file) # Pick the precision that actually matches this ONNX's tensor types. # The full_pipeline export is mixed-precision (ViT FP32, rest BF16), # so the pipeline-wide ``precision`` argument is the default but each # component uses what it was actually exported with. component_precision = _precision_from_onnx_path(onnx_path, default=precision) if component_precision != precision: logger.info( f" Using precision={component_precision} for {name} (from ONNX filename); " f"pipeline default is {precision}" ) try: # Derive shapes from the ONNX model itself min_shapes, opt_shapes, max_shapes = derive_shapes_with_hint( onnx_path, opt_seq_lens=seq_hints ) logger.info(f" Auto-derived shape profiles for {name}:") for input_name in opt_shapes: logger.info( f" {input_name}: min={min_shapes[input_name]} " f"opt={opt_shapes[input_name]} max={max_shapes[input_name]}" ) build_engine( onnx_path=onnx_path, engine_path=engine_path, precision=component_precision, workspace_mb=workspace_mb, min_shapes=min_shapes, opt_shapes=opt_shapes, max_shapes=max_shapes, trt_severity=trt_severity, ) results.append((name, engine_path, "SUCCESS")) except Exception as e: logger.error(f"Failed to build {name} engine: {e}") results.append((name, engine_path, f"FAILED: {e}")) # Print summary logger.info("\n" + "=" * 80) logger.info("FULL PIPELINE BUILD SUMMARY") logger.info("=" * 80) for name, path, status in results: logger.info(f" {name:20s} -> {status}") logger.info("=" * 80) # Every component must build; missing ONNX inputs and failed builds are # equally fatal, otherwise an empty/half-built engine dir exits 0. failures = [(name, status) for name, _, status in results if status.startswith("FAILED")] if failures or skipped: parts = [] if failures: parts.append( f"{len(failures)}/{len(components)} engine(s) failed: " + "; ".join(f"{name} ({status})" for name, status in failures) ) if skipped: parts.append( f"{len(skipped)}/{len(components)} component(s) had no ONNX input: " + ", ".join(f"{name} ({path})" for name, path in skipped) ) raise RuntimeError("Pipeline build incomplete — " + " | ".join(parts)) # ============================================================ # Main # ============================================================ @dataclass class BuildConfig: """Configuration for building TensorRT engines from ONNX models.""" mode: BuildEngineMode = BuildEngineMode.single """Build mode: 'single' (one engine) or 'full_pipeline' (all engines).""" onnx: str | None = None """Path to ONNX model (single mode).""" engine: str | None = None """Path to save TensorRT engine (single mode).""" onnx_dir: str = "./gr00t_trt_deployment/onnx" """Directory with ONNX models (full_pipeline mode).""" engine_dir: str = "./gr00t_trt_deployment/engines" """Directory to save engines (full_pipeline mode).""" precision: Literal["fp32", "fp16", "bf16", "fp8"] = "bf16" """Precision mode (default: bf16).""" workspace: int = 8192 """Workspace size in MB (default: 8192).""" allow_default_hints: bool = False """full_pipeline: build with hardcoded GR1 single-view shape hints when export_metadata.json is missing/stale/incomplete, instead of failing. The engine may get wrong sequence/patch shapes — use only for legacy bundles.""" def main(args: BuildConfig | None = None, trt_severity=None): if args is None: args = tyro.cli(BuildConfig) if args.mode == "full_pipeline": build_full_pipeline( onnx_dir=args.onnx_dir, engine_dir=args.engine_dir, precision=args.precision, workspace_mb=args.workspace, trt_severity=trt_severity, allow_default_hints=args.allow_default_hints, ) else: if not args.onnx or not args.engine: raise ValueError("--onnx and --engine are required in single mode") # Auto-derive shapes from the ONNX model min_shapes, opt_shapes, max_shapes = derive_shapes_with_hint(args.onnx) build_engine( onnx_path=args.onnx, engine_path=args.engine, precision=args.precision, workspace_mb=args.workspace, min_shapes=min_shapes, opt_shapes=opt_shapes, max_shapes=max_shapes, trt_severity=trt_severity, ) if __name__ == "__main__": config = tyro.cli(BuildConfig) main(config)