trt engine build

gistfile1.py

# pyright: basic, reportMissingImports=false
# Container-only build tool (tensorrt/torch/onnx are present only inside the
# TRT build image, not bazel deps); strict typing buys nothing here.
"""Build a multi-profile TensorRT plan engine for an encoder cross-encoder
reranker (Ettin 68m/400m today; historically the Nemotron 500m).

Takes an already-exported ONNX and compiles a serialized TRT engine for
the local GPU/TRT version — engines are not portable across either, so we
always build on the target box (GH200/aarch64). The ONNX source is either:
  - `source.onnx` from `export_hf_to_onnx.py` (the Ettin chain — the path
    `build_ettin_reranker_trt.sh` drives), or
  - a vendor-shipped ONNX (e.g. the one NVIDIA NIM ships alongside its
    x86-only Nemotron plan, which we can't use on aarch64).

Inputs:
- An ONNX with NO QDQ ops (weakly-typed bf16/fp16/fp32 build) — or one
  PTQ'd with QDQ nodes for a real fp8 build (see the fp8 note below).
- Inputs:  `input_ids` [B, S] (int64), `attention_mask` [B, S] (int64).
- Output:  `logits` [B, 1].

Output layout — one subdir per precision, consumed by the build scripts
that stage the engine into a Triton model repo (e.g.
`build_ettin_reranker_trt.sh` copies `engines/bf16/model.engine` →
`<repo>/<model>/1/model.plan`):

    <output_dir>/
        <precision>/
            model.engine
            profiles.json     # {"max_batch": N, "profiles": [...]}

Usage (typical — runs inside the TRT container; invoked for you by
build_ettin_reranker_trt.sh, shown here standalone):

    docker run --rm --gpus all \\
        -v $HOME/ettin-trt-build:/work \\
        -v $(pwd):/scripts:ro \\
        nvcr.io/nvidia/tensorrt:26.03-py3 \\
        python3 /scripts/build_trt_engine.py \\
            --onnx /work/onnx/source.onnx \\
            --output /work/engines \\
            --precisions bf16 \\
            --seq-lens 256 \\
            --max-batch 512 --opt-batches 128

Notes on fp8: an ONNX exported without QDQ nodes (our default) makes the
`BuilderFlag.FP8` path yield no actual fp8 layers. To get a real fp8
engine you must first PTQ the ONNX with NVIDIA ModelOpt (TensorRT-ModelOpt
toolkit) using a small calibration set (~100-500 representative
query+passage pairs). That's a separate step; this script supports
fp8 as a builder flag but is honest in the log when no QDQ is present.
"""

import argparse
import json
import os
import time
from pathlib import Path

import tensorrt as trt  # type: ignore[import-not-found]  # only inside TRT container


def list_qdq_ops(onnx_path: str) -> int:
    """Returns count of QuantizeLinear+DequantizeLinear nodes in the graph.

    Used purely to log a warning when fp8 is requested but the graph has
    no QDQ — TRT's FP8 builder flag without QDQ produces no fp8 layers.
    """
    try:
        import onnx  # type: ignore[import-not-found]
    except ImportError:
        return -1  # not enough info; let the build proceed
    m = onnx.load(onnx_path, load_external_data=False)
    return sum(
        1 for n in m.graph.node if n.op_type in ("QuantizeLinear", "DequantizeLinear")
    )


def build_one(
    onnx_path: str,
    out_dir: str,
    precision: str,
    seq_lens: list[int],
    max_batch: int,
    workspace_gb: int,
    opt_batches: list[int] | None = None,
) -> bool:
    os.makedirs(out_dir, exist_ok=True)
    engine_path = os.path.join(out_dir, "model.engine")
    meta_path = os.path.join(out_dir, "profiles.json")

    if os.path.exists(engine_path):
        print(f"  {precision}: {engine_path} exists, skipping")
        return True

    logger = trt.Logger(trt.Logger.WARNING)
    builder = trt.Builder(logger)
    # FP8 ONNX has explicit QDQ nodes, so the precision is fully encoded
    # in the graph. NVIDIA recommends STRONGLY_TYPED for fp8 networks —
    # TRT 10 emits "A strongly typed network is recommended..." warnings
    # otherwise, and may fall back to fp16 for matmuls it can't infer.
    # For bf16/fp16 (no QDQ in the graph) we still use weakly-typed
    # EXPLICIT_BATCH + a precision flag, since the ONNX is generic.
    if precision == "fp8":
        network = builder.create_network(
            1 << int(trt.NetworkDefinitionCreationFlag.STRONGLY_TYPED)
        )
    else:
        network = builder.create_network(
            1 << int(trt.NetworkDefinitionCreationFlag.EXPLICIT_BATCH)
        )
    parser = trt.OnnxParser(network, logger)
    if not parser.parse_from_file(onnx_path):
        for i in range(parser.num_errors):
            print(f"    parser error: {parser.get_error(i)}")
        return False

    cfg = builder.create_builder_config()
    cfg.set_memory_pool_limit(trt.MemoryPoolType.WORKSPACE, workspace_gb << 30)
    # Only set precision flags on weakly-typed networks. STRONGLY_TYPED
    # networks pick their precision from QDQ ops in the graph and reject
    # builder-level precision flags.
    if precision == "fp16":
        cfg.set_flag(trt.BuilderFlag.FP16)
    elif precision == "bf16":
        cfg.set_flag(trt.BuilderFlag.BF16)
    elif precision == "fp8":
        pass  # encoded by QDQ in the strongly-typed graph
    elif precision == "fp32":
        # No reduced-precision flag → TRT builds a full fp32 engine. Slower
        # than fp16/bf16 (no tensor-core reduced-precision path) but
        # numerically faithful to the PyTorch fp32 reference. The fp16/bf16
        # mixed-precision builds drift ~0.38 (systematic) + ~0.2 (per-doc
        # std) in the output logits vs fp32 — enough to reshuffle dense
        # candidate pools and cost head-of-ranking recall; vLLM (fp32
        # accumulation) does not. Use fp32 when ranking fidelity matters.
        pass
    else:
        raise ValueError(f"unsupported precision: {precision}")

    # If opt_batches is empty/None, fall back to a single opt at bs=64
    # (one profile per seq bucket). Otherwise create
    # len(seq_lens) × len(opt_batches) profiles, each tuned for a
    # specific (seq_bucket, bs) cell. Profile 0 has lowest seq + lowest
    # opt_batch and indices increase by seq first, then opt_batch.
    # Clamp each requested opt batch to max_batch, THEN dedupe — clamping first
    # then de-duping avoids redundant identical profiles when several requested
    # values clamp to the same max_batch (e.g. opt_batches=[128,256], max_batch=64
    # -> a single opt at 64, not two).
    raw_opts = opt_batches if opt_batches else [min(64, max_batch)]
    effective_opts = sorted({min(b, max_batch) for b in raw_opts})
    profiles = []
    prev = 0
    profile_idx = 0
    for max_seq in sorted(seq_lens):
        min_seq = 1 if prev == 0 else prev + 1
        for opt_batch in effective_opts:
            prof = builder.create_optimization_profile()
            prof.set_shape(
                "input_ids",
                (1, min_seq),
                (opt_batch, max_seq),
                (max_batch, max_seq),
            )
            prof.set_shape(
                "attention_mask",
                (1, min_seq),
                (opt_batch, max_seq),
                (max_batch, max_seq),
            )
            cfg.add_optimization_profile(prof)
            profiles.append(
                {
                    "index": profile_idx,
                    "min_seq": min_seq,
                    "max_seq": max_seq,
                    "opt_batch": opt_batch,
                }
            )
            profile_idx += 1
        prev = max_seq

    print(
        f"  {precision}: building with {len(profiles)} profiles, "
        f"max_batch={max_batch}, ws={workspace_gb}GB..."
    )
    t0 = time.time()
    blob = builder.build_serialized_network(network, cfg)
    if blob is None:
        print(f"  {precision}: BUILD FAILED")
        return False
    with open(engine_path, "wb") as f:
        f.write(blob)
    with open(meta_path, "w") as f:
        json.dump({"max_batch": max_batch, "profiles": profiles}, f, indent=2)
    sz_mb = os.path.getsize(engine_path) / 1e6
    print(f"  {precision}: wrote {sz_mb:.0f} MB in {time.time() - t0:.0f}s")
    return True


def main() -> int:
    ap = argparse.ArgumentParser()
    ap.add_argument(
        "--onnx",
        required=True,
        help="Path to the source ONNX file (e.g. source.onnx from export_hf_to_onnx.py).",
    )
    ap.add_argument(
        "--output",
        required=True,
        help=(
            "Output dir. One subdir per precision is written: "
            "<output>/<precision>/{model.engine,profiles.json}."
        ),
    )
    ap.add_argument(
        "--precisions",
        default="bf16",
        help="Comma-separated, any of: fp16, bf16, fp8.",
    )
    ap.add_argument(
        "--seq-lens",
        default="128,256,512,1024,2048",
        help="Comma-separated optimization-profile seq buckets.",
    )
    ap.add_argument("--max-batch", type=int, default=128)
    ap.add_argument("--workspace-gb", type=int, default=24)
    ap.add_argument(
        "--output-subdir",
        default="",
        help=(
            "Override the per-precision output sub-dir name. Default "
            "(empty) writes to <output>/<precision>/. Set to e.g. "
            "'fp8-mo-32-128-512' to keep variant engines side-by-side "
            "without overwriting earlier builds. Only valid when "
            "--precisions has a single entry."
        ),
    )
    ap.add_argument(
        "--opt-batches",
        default="",
        help=(
            "Comma-separated batch sizes to optimize each seq profile for. "
            "Empty (default) → single profile per seq with opt_batch=64. "
            "Set e.g. '32,128,512' to create 3 × len(seq_lens) "
            "profiles, each tuned for a specific (bs, seq) cell — Triton "
            "sub-models become <prefix>_seq_<L>_bopt_<B>."
        ),
    )
    args = ap.parse_args()

    onnx_path = os.path.abspath(args.onnx)
    if not os.path.exists(onnx_path):
        print(f"ERROR: ONNX not found at {onnx_path}")
        return 1
    out_root = os.path.abspath(args.output)
    Path(out_root).mkdir(parents=True, exist_ok=True)

    qdq_count = list_qdq_ops(onnx_path)
    print(f"==> ONNX: {onnx_path} (QDQ nodes: {qdq_count})")

    precisions = [p.strip() for p in args.precisions.split(",") if p.strip()]
    seq_lens = sorted(int(x) for x in args.seq_lens.split(",") if x.strip())
    opt_batches = [int(x) for x in args.opt_batches.split(",") if x.strip()]

    if "fp8" in precisions and qdq_count == 0:
        print(
            "WARNING: fp8 requested but ONNX has 0 QDQ nodes. TRT will set "
            "the FP8 BuilderFlag, but no layers will actually run in fp8 "
            "since explicit-quantization is required. To get a real fp8 "
            "engine, PTQ the ONNX with NVIDIA ModelOpt first."
        )

    if args.output_subdir and len(precisions) != 1:
        print(
            "ERROR: --output-subdir requires exactly one --precisions entry "
            f"(got {len(precisions)}: {precisions})"
        )
        return 1

    ok = True
    for p in precisions:
        sub = args.output_subdir or p
        out_dir = os.path.join(out_root, sub)
        if not build_one(
            onnx_path,
            out_dir,
            p,
            seq_lens,
            args.max_batch,
            args.workspace_gb,
            opt_batches=opt_batches or None,
        ):
            ok = False

    return 0 if ok else 2


if __name__ == "__main__":
    raise SystemExit(main())