diff --git a/families/sam2/README.md b/families/sam2/README.md index 9c263b7cd3..4097107ecf 100644 --- a/families/sam2/README.md +++ b/families/sam2/README.md @@ -4,9 +4,66 @@ This family builds SAM2.1 Hiera-S (`facebook/sam2.1-hiera-small`) for `video_segmentation` as six TensorRT plans: one image plan (Hiera backbone, FPN neck, box head) and five tracker plans (the prompt plan for the first frame and four recurrent plans for 1 to 4 history frames, each with memory attention, -memory encoder and mask decoder). The family precision is `mixed_bf16_fp32`, -batch size is 1 and the image plan input is 1024x1024. The family has no -quantized build. +memory encoder and mask decoder). The default precision is `fp16`; the earlier +default `mixed_bf16_fp32` can still be selected. Batch size is 1 and the image +plan input is 1024x1024. The family has no quantized build. + +## Precision + +| Precision | How to select | GEMM, convolution, attention | Layer norms and the other fp32 parts of the graph | Plan inputs and outputs | +|---|---|---|---|---| +| `fp16` (default) | default of `trtmc build`, or `precision="fp16"` | fp16 | fp32 | same for both precisions | +| `mixed_bf16_fp32` | `precision="mixed_bf16_fp32"` in a `BuildRequest` or E2E manifest | bf16 | fp32 | same for both precisions | + +Both precisions build the same graph and only change the 16-bit compute type. +The plan inputs and outputs do not change: FPN levels 0 and 1, the box maps +and the memory features stay bf16, the rest stays fp32. The fp16 build adds one +cast at each bf16 input or output (8 in the image plan, 2 in the prompt plan, 3 +in each recurrent plan), so the native runtime and the bundle layout are the +same for both precisions. The `mixed_bf16_fp32` network is unchanged. + +Measured with TensorRT 11.4.0.106 (11.1.0.106 as a check), batch 1, one GPU, +both precisions built three times (11.1: once) and timed in the same job, +median of 9 interleaved rounds. "Video" is the real unit of work of the +bundle: 5 frames, each the image plan plus the prompt plan (frame 0) or the +recurrent plan with 1 to 4 history frames, with the memory features and object +pointers passed from frame to frame (plans only, no pre- or post-processing). + +| GPU | TensorRT | Video, `mixed_bf16_fp32` (ms) | Video, `fp16` (ms) | Speedup | Image plan | Prompt plan | Recurrent H=1 / H=4 | +|---|---|---|---|---|---|---|---| +| B200 | 11.4 | 18.676 | 17.565 | 1.063x | 1.954 to 1.830 ms | 0.617 to 0.525 ms | 1.065x / 1.042x | +| B300 | 11.4 | 18.484 | 17.245 | 1.072x | 1.995 to 1.841 ms | 0.596 to 0.506 ms | 1.068x / 1.041x | +| B200 | 11.1 | 19.545 | 18.446 | 1.060x | 2.144 to 2.012 ms | 0.617 to 0.534 ms | 1.060x / 1.036x | +| B300 | 11.1 | 19.172 | 17.911 | 1.070x | 2.129 to 1.964 ms | 0.604 to 0.510 ms | 1.067x / 1.036x | + +The six-plan set shrinks from 378 to 300 MiB (B200) and from 381 to 302 MiB +(B300). + +Quality against the PyTorch fp32 `sam2` 1.1.0 video predictor (same +preprocessed pixels; not ground truth), TensorRT 11.4. COCO: 300 COCO 2017 +validation images, the largest object box as the prompt, mask logits > 0 at +256x256. Video: 18 objects in 6 windows of 5 frames of the `sam2` demo clip +"bedroom". Stress: 90 frames (black, white, uniform noise, 1-pixel and 8-pixel +checkerboards, saturated color; full, mid-size and 4-pixel boxes; 5 frames +each). + +| Set | Metric | B200 `mixed_bf16_fp32` | B200 `fp16` | B300 `mixed_bf16_fp32` | B300 `fp16` | +|---|---|---|---|---|---| +| COCO | mask IoU mean (p5) | 0.9895 (0.975) | 0.9956 (0.996) | 0.9922 (0.973) | 0.9970 (0.996) | +| COCO | mask IoU min | 0.089 | 0.462 | 0.458 | 0.462 | +| COCO | FPN relative L2, level 0 / 1 / 2 | 0.016 / 0.035 / 0.017 | 0.003 / 0.008 / 0.004 | 0.016 / 0.035 / 0.017 | 0.004 / 0.008 / 0.004 | +| Video, frames 1 to 4 | mask IoU mean | 0.959 to 0.969 | 0.981 to 0.997 | 0.976 to 0.993 | 0.995 to 0.998 | +| Stress | mask IoU mean (min) | 0.970 (0.000) | 0.999 (0.990) | 0.954 (0.000) | 0.999 (0.991) | +| All | non-finite outputs | 0 | 0 | 0 | 0 | + +The fp16 build is closer to PyTorch on every set; the bf16 build loses the +mask completely on some stress frames and drops to IoU 0.09 on one COCO image +(B200). The COCO +minimum of 0.46 is the same image for both precisions (logit cosine 0.948 +against PyTorch). The largest activation of the PyTorch model on all inputs is +8.9 (image encoder and memory attention block outputs), far below the fp16 +range. TensorRT 11.1 (50 COCO images): mask IoU 0.9991 / 0.9993 for fp16 and +0.9963 / 0.9956 for bf16 (B200 / B300), no non-finite output. ## FP8 quantization study (not shipped) @@ -17,13 +74,16 @@ fp16 build of the same graph, a whole frame gains 1.01x to 1.06x, and the masks get clearly worse. This section records the recipe, the method and the numbers, so that the study does not have to be repeated. The FP8 code was not merged. All numbers come from a local study; this repository has no FP8 build for SAM2. +The study was run when the default precision was `mixed_bf16_fp32`: "bf16" +below is that precision, and the "fp16 control" is close to the current `fp16` +default (the control also changed the plan inputs and outputs to fp16). ### 1. Summary and decision TensorRT 11.4, batch 1. "fp16 control" is the same graph with bf16 replaced by fp16 (experiment only, not a family build). -| Item | bf16 (family default) | fp16 control | FP8 image encoder | FP8 vs default / vs fp16 control | +| Item | bf16 (default at the time) | fp16 control | FP8 image encoder | FP8 vs default / vs fp16 control | |---|---|---|---|---| | Image plan, B200 | 1.976 ms | 1.877 ms | 1.935 ms | 1.02x / 0.97x | | Image plan, B300 | 2.024 ms | 1.893 ms | 1.758 ms | 1.15x / 1.08x | @@ -46,7 +106,7 @@ Decision: do not ship FP8 for SAM2. faster in FP8 (section 5), so the bundle only shrinks by about 8%. - An fp16 build of the same graph is 5% to 8% faster than the bf16 default on the image plan and has a higher mask IoU (0.9991 vs 0.9954). It is a better - first step than FP8 (not part of this study's code). + first step than FP8 and is now the default (see Precision). ### 2. Model and plans @@ -54,7 +114,7 @@ Decision: do not ship FP8 for SAM2. |---|---| | Checkpoint | `facebook/sam2.1-hiera-small`, revision `6c381d9c16faed5e8a7c4a2cd99918bdca8316e4` (the public snapshot with the family's synthetic box head) | | Plans | image plan; prompt plan; recurrent plans H=1, 2, 3, 4 | -| Family precision | `mixed_bf16_fp32` | +| Family precision | `mixed_bf16_fp32` (the default at the time) | | Batch size | 1 (the family builds batch 1 only) | | Image encoder model config | Hiera, embed dim 96, 16 blocks in 4 stages (1, 2, 11, 2), neck dim 256, image 1024 | | Image plan outputs used by the tracker | FPN levels 0, 1, 2 | @@ -145,7 +205,7 @@ bf16 / fp16 numbers within 0.004. | Image plan | Mask IoU vs PyTorch, mean (min) | Mask logit cosine vs PyTorch, mean (min) | FPN relative L2, level 0 / 1 / 2 | |---|---|---|---| -| bf16 (default) | 0.9954 (0.861) | 1.0000 (0.992) | 0.015 / 0.034 / 0.017 | +| bf16 (default at the time) | 0.9954 (0.861) | 1.0000 (0.992) | 0.015 / 0.034 / 0.017 | | fp16 control (floor) | 0.9991 (0.988) | 1.0000 (1.000) | 0.003 / 0.007 / 0.004 | | FP8 | 0.9788 (0.399) | 0.9968 (0.524) | 0.063 / 0.168 / 0.099 | | FP8, bf16 scale constants | 0.9776 (0.394) | 0.9946 (0.254) | 0.064 / 0.169 / 0.099 | @@ -158,7 +218,7 @@ Mask IoU of FP8 against the fp16 control: 0.9789 mean, 0.399 min. B300, TensorRT 11.4, mean mask IoU vs PyTorch per frame (18 objects; frames 1 to 4 use the recurrent plans): -| Frame | bf16 (default) | fp16 control | FP8 | NVFP4 (experiment) | +| Frame | bf16 (default at the time) | fp16 control | FP8 | NVFP4 (experiment) | |---|---|---|---|---| | 0 | 0.990 | 0.999 | 0.978 | 0.890 | | 1 | 0.991 | 0.997 | 0.975 | 0.909 | @@ -281,7 +341,6 @@ predictor. default (0.995) with no collapsed images. - A TensorRT release that builds `attn.proj` in FP8 after the fused attention, or gives a larger FP8 gain on B200 (retest the speed table on both GPUs). -- A change of the family default to fp16 (measured 5% to 8% faster than bf16 on - the image plan): FP8 must then beat fp16, which it does not on 11.4 B200 or - 11.1. +- The family default is now fp16 (see Precision): FP8 must beat it, which the + fp16 control above shows it does not on 11.4 B200 or 11.1. - Batch support above 1 in the family. diff --git a/families/sam2/graph_ops.py b/families/sam2/graph_ops.py index a97397c5d3..ff4e620ac5 100644 --- a/families/sam2/graph_ops.py +++ b/families/sam2/graph_ops.py @@ -28,10 +28,14 @@ class WindowTensor: class TrtLayers: """Add named, strongly typed layers while retaining all host weights.""" - def __init__(self, trt: Any, network: Any, checkpoint: Checkpoint) -> None: + def __init__(self, trt: Any, network: Any, checkpoint: Checkpoint, half: Any = None) -> None: self.trt = trt self.network = network self.checkpoint = checkpoint + # 16-bit compute type of the GEMM, convolution and attention paths: bf16 or fp16. + self.half = trt.bfloat16 if half is None else half + if self.half not in (trt.bfloat16, trt.float16): + raise NetworkBuildError("SAM2 supports only BF16 and FP16 compute") self._owned: list[np.ndarray] = [] @staticmethod @@ -54,6 +58,16 @@ def _bf16_weights(self, values: np.ndarray) -> Any: self._owned.append(rounded) return self.trt.Weights(self.trt.bfloat16, rounded.ctypes.data, rounded.size) + def _fp16_weights(self, values: np.ndarray) -> Any: + source = np.ascontiguousarray(values, dtype=np.float16) + self._owned.append(source) + return self.trt.Weights(source) + + def _half_weights(self, values: np.ndarray) -> Any: + if self.half == self.trt.float16: + return self._fp16_weights(values) + return self._bf16_weights(values) + def _fp32_weights(self, values: np.ndarray) -> Any: source = np.ascontiguousarray(values, dtype=np.float32) self._owned.append(source) @@ -64,7 +78,9 @@ def _typed_weights(self, values: np.ndarray, dtype: Any) -> Any: return self._fp32_weights(values) if dtype == self.trt.bfloat16: return self._bf16_weights(values) - raise NetworkBuildError("SAM2 supports only FP32 and BF16 weights") + if dtype == self.trt.float16: + return self._fp16_weights(values) + raise NetworkBuildError("SAM2 supports only FP32, BF16 and FP16 weights") def cast(self, tensor: Any, dtype: Any, name: str) -> Any: if tensor.dtype == dtype: @@ -220,7 +236,7 @@ def convolution_batch_norm_silu( epsilon: float, name: str, ) -> Any: - tensor = self.cast(tensor, self.trt.bfloat16, f"{name}.input_bf16") + tensor = self.cast(tensor, self.half, f"{name}.input_bf16") kernel_values, bias_values = self._fold_batch_norm( module_name, input_channels, output_channels, kernel, groups, epsilon ) @@ -229,8 +245,8 @@ def convolution_batch_norm_silu( tensor, output_channels, (kernel, kernel), - self._bf16_weights(kernel_values), - self._bf16_weights(bias_values), + self._half_weights(kernel_values), + self._half_weights(bias_values), ), "convolution", f"{name}.conv_bn_folded", @@ -240,7 +256,7 @@ def convolution_batch_norm_silu( layer.num_groups = groups return self.silu(layer.get_output(0), f"{name}.silu") - def linear_bf16( + def linear_half( self, tensor: Any, module_name: str, @@ -249,14 +265,14 @@ def linear_bf16( name: str, ) -> Any: input_shape = self._shape(tensor) - tensor = self.cast(tensor, self.trt.bfloat16, f"{name}.input_bf16") + tensor = self.cast(tensor, self.half, f"{name}.input_bf16") weight = self.constant( f"{module_name}.weight", (output_features, input_features), (output_features, input_features), f"{name}.weight", ) - weight = self.cast(weight, self.trt.bfloat16, f"{name}.weight_bf16") + weight = self.cast(weight, self.half, f"{name}.weight_bf16") weight_shape = [1] * len(input_shape) weight_shape[-2:] = (output_features, input_features) weight = self.shuffle(weight, weight_shape, f"{name}.weight_broadcast") @@ -275,7 +291,7 @@ def linear_bf16( "constant", f"{name}.bias", ).get_output(0) - bias = self.cast(bias, self.trt.bfloat16, f"{name}.bias_bf16") + bias = self.cast(bias, self.half, f"{name}.bias_bf16") return self.elementwise(output, bias, self.trt.ElementWiseOperation.SUM, f"{name}.bias_add") def layer_norm_fp32( diff --git a/families/sam2/image_builder.py b/families/sam2/image_builder.py index 2e24eafafe..535c602c44 100644 --- a/families/sam2/image_builder.py +++ b/families/sam2/image_builder.py @@ -60,10 +60,10 @@ class _Block: class ImageNetworkBuilder: - def __init__(self, trt: Any, network: Any, checkpoint: Checkpoint) -> None: + def __init__(self, trt: Any, network: Any, checkpoint: Checkpoint, half: Any = None) -> None: self.trt = trt self.network = network - self.layers = TrtLayers(trt, network, checkpoint) + self.layers = TrtLayers(trt, network, checkpoint, half) @staticmethod def _prefix(index: int) -> str: @@ -83,7 +83,7 @@ def _attention( weights = f"{self._prefix(index)}.attn" name = self._name(index, "attention") - qkv = self.layers.linear_bf16( + qkv = self.layers.linear_half( tensor, f"{weights}.qkv", input_channels, 3 * output_channels, f"{name}.qkv" ) grouped = self.layers.shuffle( @@ -122,7 +122,7 @@ def _attention( k = self.layers.transpose(k, (0, 2, 1, 3), f"{name}.k_heads") v = self.layers.transpose(v, (0, 2, 1, 3), f"{name}.v_heads") scale = self.layers.scalar( - reciprocal_sqrtf(head_channels), 4, self.trt.bfloat16, f"{name}.q_scale" + reciprocal_sqrtf(head_channels), 4, self.layers.half, f"{name}.q_scale" ) q = self.layers.elementwise( q, scale, self.trt.ElementWiseOperation.PROD, f"{name}.q_scaled" @@ -148,7 +148,7 @@ def _attention( (batch, query_height, query_width, output_channels), f"{name}.image", ) - return self.layers.linear_bf16( + return self.layers.linear_half( image, f"{weights}.proj", output_channels, @@ -164,7 +164,7 @@ def _block(self, tensor: Any, index: int, contract: _Block) -> Any: ) shortcut = tensor if contract.input_channels != contract.output_channels: - projected = self.layers.linear_bf16( + projected = self.layers.linear_half( normalized, f"{weights}.proj", contract.input_channels, @@ -220,7 +220,7 @@ def _block(self, tensor: Any, index: int, contract: _Block) -> Any: mlp_input = self.layers.layer_norm_fp32( result, f"{weights}.norm2", contract.output_channels, 1.0e-6, f"{name}.norm2" ) - mlp = self.layers.linear_bf16( + mlp = self.layers.linear_half( mlp_input, f"{weights}.mlp.layers.0", contract.output_channels, @@ -228,7 +228,7 @@ def _block(self, tensor: Any, index: int, contract: _Block) -> Any: f"{name}.mlp.fc1", ) mlp = self.layers.gelu(mlp, f"{name}.mlp.gelu") - mlp = self.layers.linear_bf16( + mlp = self.layers.linear_half( mlp, f"{weights}.mlp.layers.1", 4 * contract.output_channels, @@ -241,7 +241,7 @@ def _block(self, tensor: Any, index: int, contract: _Block) -> Any: ) def _hiera(self, pixels: Any) -> list[Any]: - tensor = self.layers.cast(pixels, self.trt.bfloat16, "hiera.input_bf16") + tensor = self.layers.cast(pixels, self.layers.half, "hiera.input_bf16") tensor = self.layers.convolution( tensor, "image_encoder.trunk.patch_embed.proj.weight", @@ -309,7 +309,7 @@ def _fpn(self, trunk: list[Any]) -> list[Any]: for level in range(3, -1, -1): module = f"image_encoder.neck.convs.{3 - level}.conv" name = f"fpn.lateral.{level}" - tensor = self.layers.cast(trunk[level], self.trt.bfloat16, f"{name}.input_bf16") + tensor = self.layers.cast(trunk[level], self.layers.half, f"{name}.input_bf16") fpn[level] = self.layers.convolution( tensor, f"{module}.weight", @@ -402,6 +402,9 @@ def _bbox_head(self, fpn: list[Any]) -> tuple[list[Any], list[Any]]: def _output(self, tensor: Any, contract: tuple[str, str, tuple[int, ...]]) -> None: name, dtype_name, shape = contract dtype = self.trt.bfloat16 if dtype_name == "bf16" else self.trt.float32 + if dtype == self.trt.bfloat16 and tensor.dtype == self.layers.half: + # The runtime contract keeps BF16 interface tensors for both compute types. + tensor = self.layers.cast(tensor, dtype, f"{name}.output_bf16") if tensor.dtype != dtype or self.layers._shape(tensor) != shape: raise NetworkBuildError( f"output contract mismatch for {name}: actual shape {self.layers._shape(tensor)}" @@ -419,7 +422,9 @@ def build(self) -> None: self._output(tensor, contract) -def populate_image_network(trt: Any, network: Any, checkpoint: Checkpoint) -> ImageNetworkBuilder: - builder = ImageNetworkBuilder(trt, network, checkpoint) +def populate_image_network( + trt: Any, network: Any, checkpoint: Checkpoint, half: Any = None +) -> ImageNetworkBuilder: + builder = ImageNetworkBuilder(trt, network, checkpoint, half) builder.build() return builder diff --git a/families/sam2/model.py b/families/sam2/model.py index e9ae72a4fc..beeb653ea2 100644 --- a/families/sam2/model.py +++ b/families/sam2/model.py @@ -25,7 +25,11 @@ ) -_PRECISION = "mixed_bf16_fp32" +_PRECISION = "fp16" +# Compute type of the GEMM, convolution and attention paths per precision. Both keep layer +# norms, softmax-sensitive reductions and the runtime float outputs in FP32, and both keep the +# BF16 runtime interface tensors (FPN levels 0 and 1, box maps, memory features). +_HALF_TYPES = {"fp16": "float16", "mixed_bf16_fp32": "bfloat16"} _PUBLIC_CORE_VARIANT = "public_sam2_1_small_with_synthetic_bbox_v1" _DEFAULT_WORKSPACE_BYTES = 8 << 30 _GENERIC_EMBED_CANDIDATES = ( @@ -123,9 +127,13 @@ class _Sam2Model: default_build_precision = _PRECISION @staticmethod - def _require_precision(precision: str) -> None: - if precision != _PRECISION: - raise ValueError(f"SAM2 supports only {_PRECISION!r} precision, got {precision!r}") + def _require_precision(precision: str) -> Any: + if precision not in _HALF_TYPES: + raise ValueError( + f"SAM2 supports only {', '.join(map(repr, _HALF_TYPES))} precision, " + f"got {precision!r}" + ) + return getattr(trt, _HALF_TYPES[precision]) def load_weights( self, model_dir: str, _config: Any, *, precision: str = _PRECISION @@ -160,12 +168,15 @@ def build_engine( precision: str = _PRECISION, verbose: bool = False, ) -> bytes: - self._require_precision(precision) + half = self._require_precision(precision) from .image_builder import populate_image_network + def populate(trt: Any, network: Any, source: Checkpoint) -> Any: + return populate_image_network(trt, network, source, half) + return _serialize( _checkpoint(weights), - populate_image_network, + populate, expected_inputs=1, expected_outputs=9, section="engine_plan", @@ -181,7 +192,7 @@ def build_extra_engines( precision: str = _PRECISION, verbose: bool = False, ) -> dict[str, bytes]: - self._require_precision(precision) + half = self._require_precision(precision) from .tracker_builder import populate_tracker_network checkpoint = _checkpoint(weights) @@ -191,7 +202,7 @@ def build_extra_engines( def populate( trt: Any, network: Any, source: Checkpoint, h: int = history_frames ) -> Any: - return populate_tracker_network(trt, network, source, h) + return populate_tracker_network(trt, network, source, h, half) plans[section] = _serialize( checkpoint, diff --git a/families/sam2/support.py b/families/sam2/support.py index 0503929379..3990fcadf2 100644 --- a/families/sam2/support.py +++ b/families/sam2/support.py @@ -11,4 +11,5 @@ required_files=("sam2.1_hiera_s.yaml", "sam2.1_hiera_small.pt"), tasks=("video_segmentation",), default_task="video_segmentation", + default_precision="fp16", ) diff --git a/families/sam2/tests/manifests/sam2-public-core-l0-fp16.json b/families/sam2/tests/manifests/sam2-public-core-l0-fp16.json new file mode 100644 index 0000000000..71b33d2e39 --- /dev/null +++ b/families/sam2/tests/manifests/sam2-public-core-l0-fp16.json @@ -0,0 +1,19 @@ +{ + "name": "sam2-public-core-l0-fp16", + "hf_id": "facebook/sam2.1-hiera-small", + "hf_revision": "6c381d9c16faed5e8a7c4a2cd99918bdca8316e4", + "bundle": "sam2-public-core-l0-fp16.bundle", + "family": "sam2", + "task": "video_segmentation", + "testcases": [ + { + "name": "sam2-public-core-l0-fp16", + "premerge": true + } + ], + "tensor_parallel_size": 1, + "precision": "fp16", + "image_height": 1280, + "image_width": 1088, + "video_num_frames": 5 +} diff --git a/families/sam2/tests/test_precision.py b/families/sam2/tests/test_precision.py new file mode 100644 index 0000000000..3ac8251a47 --- /dev/null +++ b/families/sam2/tests/test_precision.py @@ -0,0 +1,87 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +from __future__ import annotations + +import numpy as np +import pytest + +trt = pytest.importorskip("tensorrt") + +from families.sam2.image_builder import populate_image_network # noqa: E402 +from families.sam2.model import _Sam2Model # noqa: E402 +from families.sam2.support import describe # noqa: E402 +from families.sam2.tracker_builder import populate_tracker_network # noqa: E402 +from tensorrt_model_connect.model_support import ModelMetadata # noqa: E402 + + +class _Weights: + """Deterministic stand-in for the checkpoint: graph structure only.""" + + def tensor(self, name: str, shape: tuple[int, ...]) -> np.ndarray: + if name.endswith("running_var"): + return np.ones(shape, dtype=np.float32) + seed = sum(name.encode()) % 997 + return np.random.default_rng(seed).normal(0.0, 0.02, shape).astype(np.float32) + + +def _network(populate, half): + builder = trt.Builder(trt.Logger(trt.Logger.ERROR)) + network = builder.create_network(0) + owner = populate(trt, network, _Weights(), half) + return builder, network, owner + + +def _io(network) -> list[tuple[str, object, tuple[int, ...]]]: + tensors = [network.get_input(i) for i in range(network.num_inputs)] + tensors += [network.get_output(i) for i in range(network.num_outputs)] + return [(tensor.name, tensor.dtype, tuple(tensor.shape)) for tensor in tensors] + + +def _matmul_input_types(network) -> set: + types = set() + for index in range(network.num_layers): + layer = network.get_layer(index) + if layer.type == trt.LayerType.MATRIX_MULTIPLY: + types.update(layer.get_input(i).dtype for i in range(layer.num_inputs)) + return types + + +def test_default_precision_is_fp16() -> None: + assert _Sam2Model.default_build_precision == "fp16" + support = describe(ModelMetadata({"model_type": "sam2"}, {})) + assert support is not None and support.default_precision == "fp16" + + +@pytest.mark.parametrize( + ("precision", "half"), + [("fp16", "float16"), ("mixed_bf16_fp32", "bfloat16")], +) +def test_precision_selects_the_compute_type(precision: str, half: str) -> None: + assert _Sam2Model._require_precision(precision) == getattr(trt, half) + + +@pytest.mark.parametrize("precision", ["bf16", "fp32", "mixed_fp16_fp32"]) +def test_unsupported_precision_is_rejected(precision: str) -> None: + with pytest.raises(ValueError, match="SAM2 supports only"): + _Sam2Model._require_precision(precision) + + +@pytest.mark.parametrize( + "populate", + [ + populate_image_network, + lambda t, n, c, half: populate_tracker_network(t, n, c, 0, half), + lambda t, n, c, half: populate_tracker_network(t, n, c, 4, half), + ], + ids=["image", "prompt", "recurrent4"], +) +def test_fp16_keeps_the_runtime_contract(populate) -> None: + _, bf16, bf16_owner = _network(populate, trt.bfloat16) + _, fp16, fp16_owner = _network(populate, trt.float16) + assert _io(fp16) == _io(bf16) + assert trt.bfloat16 in {dtype for _, dtype, _ in _io(fp16)} + assert _matmul_input_types(bf16) <= {trt.bfloat16, trt.float32} + assert trt.float16 in _matmul_input_types(fp16) + assert trt.bfloat16 not in _matmul_input_types(fp16) + del bf16_owner, fp16_owner diff --git a/families/sam2/tracker_builder.py b/families/sam2/tracker_builder.py index 21af9b4611..ea8322ece4 100644 --- a/families/sam2/tracker_builder.py +++ b/families/sam2/tracker_builder.py @@ -42,16 +42,19 @@ class _DecoderOutput: class _TrackerBuilder: - def __init__(self, trt: Any, network: Any, checkpoint: Any) -> None: + def __init__(self, trt: Any, network: Any, checkpoint: Any, half: Any = None) -> None: self.trt = trt self.network = network self.checkpoint = checkpoint - self.layers = TrtLayers(trt, network, checkpoint) + self.layers = TrtLayers(trt, network, checkpoint, half) def _add_input(self, contract: tuple[str, str, tuple[int, ...]]) -> Any: name, dtype_name, shape = contract dtype = getattr(self.trt, dtype_name) tensor = self.network.add_input(name, dtype, shape) + if dtype == self.trt.bfloat16: + # The runtime contract keeps BF16 interface tensors for both compute types. + tensor = self.layers.cast(tensor, self.layers.half, f"{name}.compute") return tensor def _float_constant( @@ -133,7 +136,7 @@ def _conv( groups: int, name: str, ) -> Any: - tensor = self.layers.cast(tensor, self.trt.bfloat16, f"{name}.input") + tensor = self.layers.cast(tensor, self.layers.half, f"{name}.input") return self.layers.convolution( tensor, f"{module}.weight", @@ -155,7 +158,7 @@ def _deconv( output_channels: int, name: str, ) -> Any: - tensor = self.layers.cast(tensor, self.trt.bfloat16, f"{name}.input") + tensor = self.layers.cast(tensor, self.layers.half, f"{name}.input") source_kernel = self.checkpoint.tensor( f"{module}.weight", (input_channels, output_channels, 2, 2) ) @@ -170,8 +173,8 @@ def _deconv( tensor, output_channels * 4, (1, 1), - self.layers._bf16_weights(projection_kernel), - self.layers._bf16_weights(projection_bias), + self.layers._half_weights(projection_kernel), + self.layers._half_weights(projection_bias), ), "pixel-shuffle projection", f"{name}.projection", @@ -227,7 +230,7 @@ def _layer_norm_2d(self, tensor: Any, module: str, channels: int, name: str) -> def _linear_relu( self, tensor: Any, module: str, input_features: int, output_features: int, name: str ) -> Any: - value = self.layers.linear_bf16( + value = self.layers.linear_half( tensor, module, input_features, output_features, f"{name}.linear" ) return self._activation(value, self.trt.ActivationType.RELU, name) @@ -264,14 +267,14 @@ def _make_dense_random_position(self) -> Any: coordinates[0, spatial, 0] = np.float32(2.0) * normalized_x - np.float32(1.0) coordinates[0, spatial, 1] = np.float32(2.0) * normalized_y - np.float32(1.0) grid = self._float_constant( - coordinates, coordinates.shape, "prompt.dense_position.grid", self.trt.bfloat16 + coordinates, coordinates.shape, "prompt.dense_position.grid", self.layers.half ) gaussian = self._checkpoint_constant( "sam_prompt_encoder.pe_layer.positional_encoding_gaussian_matrix", (2, 128), (1, 2, 128), "prompt.dense_position.gaussian", - self.trt.bfloat16, + self.layers.half, ) phase = self.layers.matrix_multiply( grid, @@ -321,13 +324,13 @@ def _build_sparse_prompt(self, box: Any | None, prompt: bool) -> Any: unit = self._binary( doubled, minus_one, self.trt.ElementWiseOperation.SUM, "prompt.box.unit_square" ) - unit = self.layers.cast(unit, self.trt.bfloat16, "prompt.box.unit_square_bf16") + unit = self.layers.cast(unit, self.layers.half, "prompt.box.unit_square_bf16") gaussian = self._checkpoint_constant( "sam_prompt_encoder.pe_layer.positional_encoding_gaussian_matrix", (2, 128), (1, 2, 128), "prompt.box.gaussian", - self.trt.bfloat16, + self.layers.half, ) phase = self.layers.matrix_multiply( unit, @@ -494,13 +497,13 @@ def _attention( query_tokens = self.layers._shape(query)[1] key_tokens = self.layers._shape(key)[1] head_dim = internal_features // heads - projected_q = self.layers.linear_bf16( + projected_q = self.layers.linear_half( query, f"{module}.q_proj", query_features, internal_features, f"{name}.q_proj" ) - projected_k = self.layers.linear_bf16( + projected_k = self.layers.linear_half( key, f"{module}.k_proj", key_value_features, internal_features, f"{name}.k_proj" ) - projected_v = self.layers.linear_bf16( + projected_v = self.layers.linear_half( value, f"{module}.v_proj", key_value_features, internal_features, f"{name}.v_proj" ) q_grouped = self.layers.shuffle( @@ -562,7 +565,7 @@ def _attention( recombined = self.layers.shuffle( token_major, (1, query_tokens, internal_features), f"{name}.recombine" ) - return self.layers.linear_bf16( + return self.layers.linear_half( recombined, f"{module}.out_proj", internal_features, @@ -675,7 +678,7 @@ def _two_way_transformer( hidden = self._linear_relu( queries, f"{prefix}.mlp.layers.0", 256, 2048, f"{name}.mlp.relu" ) - mlp = self.layers.linear_bf16( + mlp = self.layers.linear_half( hidden, f"{prefix}.mlp.layers.1", 2048, 256, f"{name}.mlp.output" ) queries = self._binary( @@ -921,7 +924,7 @@ def _decode( hidden0, f"{module}1", 256, 256, f"decoder.hyper.{token}.relu1" ) hypernetworks.append( - self.layers.linear_bf16( + self.layers.linear_half( hidden1, f"{module}2", 256, 32, f"decoder.hyper.{token}.output" ) ) @@ -952,7 +955,7 @@ def _decode( 256, "decoder.iou.relu1", ) - iou_rank3 = self.layers.linear_bf16( + iou_rank3 = self.layers.linear_half( iou_hidden1, "sam_mask_decoder.iou_prediction_head.layers.2", 256, @@ -978,7 +981,7 @@ def _decode( 256, "decoder.object.relu1", ) - object_score_rank3 = self.layers.linear_bf16( + object_score_rank3 = self.layers.linear_half( object_hidden1, "sam_mask_decoder.pred_obj_score_head.layers.2", 256, @@ -1040,8 +1043,8 @@ def _decode( (1, 1, 1, 1), "decoder.mask.alternate_best", ) - upper = self.layers.scalar(0.05, 4, self.trt.bfloat16, "decoder.stability.upper") - lower = self.layers.scalar(-0.05, 4, self.trt.bfloat16, "decoder.stability.lower") + upper = self.layers.scalar(0.05, 4, self.layers.half, "decoder.stability.upper") + lower = self.layers.scalar(-0.05, 4, self.layers.half, "decoder.stability.lower") intersection_bool = self.layers.elementwise( stability_mask, upper, @@ -1161,7 +1164,7 @@ def _decode( pointer1 = self._linear_relu( pointer0, "obj_ptr_proj.layers.1", 256, 256, "decoder.pointer.relu1" ) - pointer2 = self.layers.linear_bf16( + pointer2 = self.layers.linear_half( pointer1, "obj_ptr_proj.layers.2", 256, 256, "decoder.pointer.output" ) pointer_fp32 = self.layers.cast(pointer2, self.trt.float32, "decoder.pointer.fp32") @@ -1264,11 +1267,11 @@ def _encode_memory( ) normalized = self._layer_norm_2d(depthwise, module + ".norm", 256, name + ".norm") nhwc = self.layers.transpose(normalized, (0, 2, 3, 1), name + ".nhwc") - expanded = self.layers.linear_bf16( + expanded = self.layers.linear_half( nhwc, module + ".pwconv1", 256, 1024, name + ".expand" ) hidden = self.layers.gelu(expanded, name + ".gelu") - projected = self.layers.linear_bf16( + projected = self.layers.linear_half( hidden, module + ".pwconv2", 1024, 256, name + ".project" ) gamma = self._checkpoint_constant( @@ -1323,7 +1326,7 @@ def _memory_attention( memory_positions = [] spatial_position_fp32 = self._make_sine_position(64, "attention.memory.spatial_fp32") spatial_position = self.layers.cast( - spatial_position_fp32, self.trt.bfloat16, "attention.memory.spatial_bf16" + spatial_position_fp32, self.layers.half, "attention.memory.spatial_bf16" ) temporal_table = self._checkpoint_constant( "maskmem_tpos_enc", @@ -1399,7 +1402,7 @@ def _memory_attention( history_frames if index == 0 else index, f"attention.pointer.raw_position.{index}", ) - projected_position = self.layers.linear_bf16( + projected_position = self.layers.linear_half( raw_position, "obj_ptr_tpos_proj", 256, @@ -1546,7 +1549,7 @@ def _memory_attention( name + ".mlp.norm", ) hidden = self._linear_relu(mlp_norm, module + ".linear1", 256, 2048, name + ".mlp.relu") - projected = self.layers.linear_bf16( + projected = self.layers.linear_half( hidden, module + ".linear2", 2048, 256, name + ".mlp.output" ) output = self._binary( @@ -1627,7 +1630,7 @@ def build(self, history_frames: int) -> None: def populate_tracker_network( - trt: Any, network: Any, checkpoint: Any, history_frames: int + trt: Any, network: Any, checkpoint: Any, history_frames: int, half: Any = None ) -> _TrackerBuilder: """Populate one static prompt (H=0) or recurrent (H=1..4) SAM2 network. @@ -1635,6 +1638,6 @@ def populate_tracker_network( network because TensorRT may retain pointers to NumPy-backed weights. """ - builder = _TrackerBuilder(trt, network, checkpoint) + builder = _TrackerBuilder(trt, network, checkpoint, half) builder.build(history_frames) return builder