fix(export): keep multiplier depthwise convs FP32 so the MobileCLIP image encoder builds as FP16

CHANGELOG.md

@@ -33,6 +33,12 @@ history of this codebase and was never published. This release is `[0.5.0]`.
 
 ### Fixed
 
+- The MobileCLIP image encoder now builds as an FP16 engine on TensorRT 11.1 instead of falling
+  back to FP32. The build stopped at `Could not find any implementation for node
+  .../reparam_conv/Conv + PWN(...)`: TensorRT 11.1 has no FP16 kernel for a depthwise conv with a
+  channel multiplier, so `bake_fp16_onnx` keeps those four convs in FP32. The engine is 75 MB
+  against 147 MB, 1.5x (batch 8) to 2.1x (batch 32) faster per image, cosine similarity to the FP32
+  engine at least 0.9998 on 200 COCO images (#224).
 - The `curation-vlm-worker` no longer polls every project each cycle and takes a 409
   `vlm_not_configured` from the ones with no VLM (#207, about 2.7 idle cores at 18 projects). Each
   cycle it asks the registry and each project's activation which projects have an active VLM,

docs/PERFORMANCE.md

@@ -812,10 +812,13 @@ The earlier private FP16 measurement of about 170 embeddings/s (5.9 ms) matches
   labelled data (plan section 7) are open.
 - The fix bakes FP16 for every model that calls `bake_fp16_onnx`. On TensorRT 11.1 the MobileCLIP
   image encoder's v0.5.0 bake also failed at parse (`ElementWiseOperation DIV must have same input
-  types`); with the reconcile pass it parses, but `trtexec` then stops at
-  `Could not find any implementation for node .../reparam_conv/Conv + PWN(...)`, so its exporter
-  still falls back to FP32 (its existing, printed fallback). That is a separate TensorRT-side
-  failure, tracked in #224; no other baked model has a mixed-type op.
+  types`); with the reconcile pass it parses, but the build stopped at
+  `Could not find any implementation for node .../reparam_conv/Conv + PWN(...)` (#224). The failing
+  layers are the depthwise 7x7 stride-2 convs with a channel multiplier of 2 (the three
+  `downsample/proj/proj.0/reparam_conv` convs) and the final 3x3 conv (640 groups, 1280 outputs):
+  TensorRT 11.1 has no FP16 kernel for a grouped conv with one input channel per group and more
+  outputs than groups. `bake_fp16_onnx` now leaves those four convs in FP32 and the engine builds as
+  FP16 (see "MobileCLIP image encoder FP16" below). No other baked model has such a conv.
 - The installer's FP32 fallback is kept (an FP32 engine beats none) but is no longer silent: it
   prints a warning naming the model, the FP16 failure reason and the consequence, writes the model
   to `.install/precision.tsv`, records the group as `degraded` in `groups.tsv` (a re-run retries FP16
@@ -971,6 +974,23 @@ cd benchmarks
 
 ---
 
+### MobileCLIP image encoder FP16
+
+One RTX 3080 Ti (12 GB, no foreign load during the runs), TensorRT 11.1, dynamic batch 1/8/32 at
+256 x 256, engine-only `trtexec` timings (no Triton, no preprocessing):
+
+| Engine | Size | Build | Batch 8 GPU compute | Batch 8 images/s | Batch 32 GPU compute | Batch 32 images/s |
+|---|---:|---:|---:|---:|---:|---:|
+| FP32 | 146.6 MB | 62 s | 12.7 ms | 627 | 41.5 ms | 771 |
+| FP16 (4 convs FP32) | 75.1 MB | 65 s | 8.3 ms | 966 | 19.9 ms | 1604 |
+
+Parity of the FP16 engine against the FP32 engine on the first 200 images of the pinned COCO
+benchmark set (hashes verified): cosine similarity min 0.99986, mean 0.99998; top-1 nearest
+neighbour agrees for 98.5 percent of images and the top-5 neighbour overlap is 99.1 percent. The
+engine built by `export_mobileclip_image_encoder.py` itself (71.5 MB, 72 s, no fallback) measures
+cosine min 0.99984, mean 0.99998, top-1 99.5 percent, top-5 99.3 percent. Per-model Triton
+`perf_analyzer` runs were not made; the engine-only figures above are the speedup bound.
+
 ## Tuning Parameters
 
 ### Worker Count Optimization

docs/design/triton_pipeline_optimization_plan.md

@@ -821,9 +821,9 @@ changed. Raw numbers: `docs/benchmarks/v060_fp16_pe.json` and
   limit: it is 74 percent busy, and the host-side item 2 of section 11.1 is the next to measure.
 - [ ] Accuracy gates of section 7 (retrieval, clustering) on labelled sets: not run; only a 200-image
   engine parity check is recorded.
-- [ ] MobileCLIP image encoder: its FP16 bake also failed at parse on TensorRT 11.1; after this fix it
-  parses but the build fails in TensorRT (`Could not find any implementation` for a reparam conv), so
-  its exporter still falls back to FP32 (#224).
+- [x] MobileCLIP image encoder FP16 (#224): the build failed on four depthwise convs with a channel
+  multiplier (no FP16 kernel in TensorRT 11.1); `bake_fp16_onnx` keeps them FP32 and the engine is
+  FP16 (cosine min 0.99984 against FP32 on 200 COCO images).
 
 ## 12. References
 

export/export_mobileclip_image_encoder.py

@@ -291,15 +291,13 @@ def _build_engine(onnx_path, plan_path, fp16, max_batch_size):
     if fp16:
         from trt_utils import bake_fp16_onnx
 
-        # NOTE: the ModelOpt AutoCast path (used by export_yolo26.py) fails
-        # to build this graph's MobileOne reparam-conv + pointwise fusion
-        # ("Could not find any implementation for node
-        # …stages.1/downsample/proj/proj.0/reparam_conv/Conv + PWN…").
-        # trt_utils.bake_fp16_onnx (onnxconverter-common's plain
-        # weight/activation cast, keep_io_types=True) sidesteps that fusion
-        # entirely -- it doesn't rewrite the graph the way ModelOpt does --
-        # and is the same helper used successfully by the ArcFace, SCRFD and
-        # PaddleOCR-det exporters.
+        # trt_utils.bake_fp16_onnx (onnxconverter-common's plain weight/
+        # activation cast, keep_io_types=True) is the helper the ArcFace,
+        # SCRFD and PaddleOCR-det exporters use. It keeps this graph's
+        # depthwise convs with a channel multiplier (the stride-2 downsample
+        # projections and the final conv) in FP32: TensorRT 11.1 has no FP16
+        # kernel for them and stops with "Could not find any implementation
+        # for node .../reparam_conv/Conv + PWN(...)" otherwise (#224).
         print('  Baking FP16 (onnxconverter-common, TRT 11 typed builds)...')
         onnx_path = bake_fp16_onnx(onnx_path)
 

export/trt_utils.py

@@ -197,6 +197,27 @@ def _reconcile_float_types(model, fp32_ops: set[str]) -> int:
     raise ValueError(f'float types did not settle after {_MAX_RECONCILE_PASSES} passes')
 
 
+def _multiplier_depthwise_convs(model) -> list[str]:
+    """Names of depthwise Convs whose output has more channels than their group.
+
+    TensorRT 11.1 (strongly typed) finds no FP16 kernel for a grouped Conv with
+    one input channel per group and a channel multiplier above 1 (MobileCLIP's
+    stride-2 downsample projections and its final conv) and the build stops
+    with "Could not find any implementation for node .../Conv + PWN(...)".
+    Only Convs with initializer weights are inspected.
+    """
+    weights = {init.name: init for init in model.graph.initializer}
+    names = []
+    for node in model.graph.node:
+        if node.op_type != 'Conv' or len(node.input) < 2 or node.input[1] not in weights:
+            continue
+        group = next((a.i for a in node.attribute if a.name == 'group'), 1)
+        out_channels, channels_per_group = weights[node.input[1]].dims[:2]
+        if group > 1 and channels_per_group == 1 and out_channels > group:
+            names.append(node.name)
+    return names
+
+
 def bake_fp16_onnx(onnx_path: str | Path, output_path: str | Path | None = None) -> Path:
     """Bake FP16 mixed precision into an ONNX for TRT >= 11 typed builds.
 
@@ -212,6 +233,8 @@ def bake_fp16_onnx(onnx_path: str | Path, output_path: str | Path | None = None)
     output can mix FP32 and FP16 operands at one op, which TensorRT 11
     rejects at parse. :func:`_reconcile_float_types` repairs that, and
     raises (callers fall back to FP32) if the graph does not settle.
+    Depthwise Convs with a channel multiplier stay FP32
+    (:func:`_multiplier_depthwise_convs`): TensorRT 11.1 has no FP16 kernel for them.
 
     On older TRT (classic FP16 flag still present) this is a no-op and
     returns the input path — :func:`enable_fp16` handles precision there.
@@ -237,7 +260,10 @@ def bake_fp16_onnx(onnx_path: str | Path, output_path: str | Path | None = None)
     }
     model = onnx.load(str(onnx_path))
     fp16_model = float16.convert_float_to_float16(
-        model, keep_io_types=True, op_block_list=sorted(fp32_ops)
+        model,
+        keep_io_types=True,
+        op_block_list=sorted(fp32_ops),
+        node_block_list=_multiplier_depthwise_convs(model),
     )
     _reconcile_float_types(fp16_model, fp32_ops)
     out = Path(output_path) if output_path else Path(onnx_path).with_suffix('.fp16.onnx')

tests/test_trt_utils.py

@@ -435,3 +435,88 @@ def run(path: Path) -> np.ndarray:
             return session.run(None, {'x': x})[0]
 
         np.testing.assert_allclose(run(baked), run(src), rtol=2e-2, atol=2e-2)
+
+
+def _channel_multiplier_depthwise_onnx(path: Path) -> Path:
+    """Depthwise 3x3 conv with a channel multiplier of 2, then a pointwise conv.
+
+    The first is MobileOne's stride-2 downsample projection in miniature: the
+    shape TensorRT 11.1 has no FP16 tactic for.
+    """
+    import numpy as np
+    import onnx
+    from onnx import TensorProto, helper, numpy_helper
+
+    rng = np.random.default_rng(0)
+    depthwise_w = numpy_helper.from_array(rng.normal(size=(8, 1, 3, 3)).astype(np.float32), 'dw_w')
+    pointwise_w = numpy_helper.from_array(rng.normal(size=(4, 8, 1, 1)).astype(np.float32), 'pw_w')
+    images = helper.make_tensor_value_info('images', TensorProto.FLOAT, [1, 4, 8, 8])
+    out = helper.make_tensor_value_info('out', TensorProto.FLOAT, [1, 4, 4, 4])
+    graph = helper.make_graph(
+        [
+            helper.make_node(
+                'Conv',
+                ['images', 'dw_w'],
+                ['mid'],
+                name='dw',
+                group=4,
+                kernel_shape=[3, 3],
+                strides=[2, 2],
+                pads=[1, 1, 1, 1],
+            ),
+            helper.make_node('Conv', ['mid', 'pw_w'], ['out'], name='pw', kernel_shape=[1, 1]),
+        ],
+        'dw',
+        [images],
+        [out],
+        initializer=[depthwise_w, pointwise_w],
+    )
+    model = helper.make_model(graph, opset_imports=[helper.make_opsetid('', 17)], ir_version=8)
+    onnx.checker.check_model(model)
+    onnx.save(model, str(path))
+    return path
+
+
+class TestBakeFp16OnnxKeepsMultiplierDepthwiseConvsFp32:
+    """TensorRT 11.1 finds no FP16 kernel for a depthwise Conv whose output has
+    more channels than its group (MobileCLIP's downsample projections and final
+    conv): the build stops with "Could not find any implementation for node
+    .../reparam_conv/Conv + PWN(...)". Those convs stay FP32, the rest FP16."""
+
+    def test_multiplier_depthwise_conv_stays_fp32_and_others_are_fp16(
+        self, tmp_path: Path, monkeypatch: pytest.MonkeyPatch
+    ) -> None:
+        pytest.importorskip('onnxconverter_common')
+        _install_trt11(monkeypatch)
+        import onnx
+        import trt_utils
+
+        src = _channel_multiplier_depthwise_onnx(tmp_path / 'dw.onnx')
+
+        baked = trt_utils.bake_fp16_onnx(src, tmp_path / 'dw.fp16.onnx')
+
+        model = onnx.load(str(baked))
+        onnx.checker.check_model(model)
+        types = {i.name: i.data_type for i in model.graph.initializer}
+        assert types['dw_w'] == onnx.TensorProto.FLOAT
+        assert types['pw_w'] == onnx.TensorProto.FLOAT16
+        assert model.graph.output[0].type.tensor_type.elem_type == onnx.TensorProto.FLOAT
+
+    def test_the_baked_graph_computes_the_same_values(
+        self, tmp_path: Path, monkeypatch: pytest.MonkeyPatch
+    ) -> None:
+        ort = pytest.importorskip('onnxruntime')
+        pytest.importorskip('onnxconverter_common')
+        import numpy as np
+        import trt_utils
+
+        _install_trt11(monkeypatch)
+        src = _channel_multiplier_depthwise_onnx(tmp_path / 'dw.onnx')
+        baked = trt_utils.bake_fp16_onnx(src, tmp_path / 'dw.fp16.onnx')
+        x = np.random.default_rng(1).normal(size=(1, 4, 8, 8)).astype(np.float32)
+
+        def run(path: Path) -> np.ndarray:
+            session = ort.InferenceSession(str(path), providers=['CPUExecutionProvider'])
+            return session.run(None, {'images': x})[0]
+
+        np.testing.assert_allclose(run(baked), run(src), rtol=2e-2, atol=2e-2)