Fix LoMa bf16 fallback on CUDA by loading it like any other model by sarlinpe · Pull Request #4924 · colmap/colmap

sarlinpe · 2026-10-29T21:54:57Z

Fixes #4787. Supersedes #4895.

Problem

With --LomaMatching.use_bf16 1 on CUDA, the bf16 LoMa matcher always fails to initialize and falls back to fp32:

Failed to initialize the bf16 LoMa matcher (This session contains graph nodes that are assigned to the default CPU EP, but fallback to CPU EP has been explicitly disabled by the user.); falling back to fp32

The bf16 models are loaded as capability probes, which set session.disable_cpu_ep_fallback. But ONNX Runtime deliberately runs small shape computations on the CPU, whatever the precision of the model: its CUDA Shape/Size kernels write their output to host memory, and GetCpuPreferredNodes then moves their int64 consumers (Gather, Unsqueeze, Concat, Cast, Reshape, ...) to the CPU EP. Tensors in bf16 or fp16 are never moved. The dynamic-shape matcher graphs contain such nodes, so the probe always rejects them. The fp32 models are not loaded as probes, so they run with the same CPU nodes.

Fix

Load the bf16 models like any other model, and decide up front whether to use them:

The bf16 models are only used with the CUDA execution provider, on GPUs with compute capability 8.0 (Ampere) or newer. Otherwise, the fp32 models are loaded directly, without first downloading and loading the bf16 models.
- On older GPUs, creating a bf16 session can succeed while inference fails, which the fallback on initialization failures would not catch.
- On the CPU, use_bf16 is now a no-op instead of trying to load the bf16 models.
The bf16 models are loaded as regular ONNXModels, with a fallback to fp32 if initialization throws an Ort::Exception. This removes ONNXModel's is_capability_probe, whose only user was LoMa.
New CudaDeviceSupportsBf16() queries the device. It returns false if the device cannot be queried, and on HIP builds.
New SelectONNXCudaDevice() resolves gpu_index to the device that ONNXModel runs on, shared with the bf16 check.
Updated the option docs (C++, pycolmap, doc/features.rst).

This is the same pattern as the fp16 FlexGlue matcher in #4908: reduced precision only on CUDA, a regular session, and a fallback to fp32 on initialization failures.

Comparison with #4895

#4895 keeps the strict probe and instead inspects the EP graph assignment (ORT >= 1.24) against an allowlist of CPU operators. This PR avoids:

the allowlist, which is brittle: which nodes ORT moves to the CPU depends on data flow, not on the operator type. For example, FlexGlue's matchers run Reshape on the CPU, which the allowlist rejects. Any rejected node silently falls back to fp32 again.
the ORT_API_VERSION >= 24 gating; the fix also works with older ORT.
downloading and loading the bf16 models on GPUs without bf16 support before rejecting them.
~80 lines of hand-encoded ONNX test models.

Removing the strict probe barely weakens the check: the bf16 tensors themselves are never moved to the CPU, and the CPU EP has almost no bf16 compute kernels, so a bf16 op without a CUDA kernel still fails at initialization and falls back to fp32.

Tests

util/cuda_test: CudaDeviceSupportsBf16 returns false for invalid devices and, with a GPU, leaves no CUDA error.
util/onnx_test: SelectONNXCudaDevice, and an invalid model throws the Ort::Exception that the LoMa fallback catches.
feature/loma_test, util/onnx_test, util/cuda_test (also with CUDA_VISIBLE_DEVICES="") pass locally (CUDA + ONNX build).
Not yet run end-to-end with the CUDA EP: ORT 1.30's CUDA EP does not load on my machine (older system CUDA runtime). The bf16 LoMa-B matcher with CPU shape nodes was validated in Fix LoMa BF16 CUDA probing with CPU bookkeeping nodes #4895 (64/64 and 128/128 self-matches on an RTX 3060).

Stack: the FlexGlue stack builds on this PR: #4907, #4921, #4908, #4915. #4908 selects its fp16 matcher the same way, with a new CudaDeviceHasTensorCores() (compute capability ≥ 7.0) next to CudaDeviceSupportsBf16(). #4915 uses SelectONNXCudaDevice() for the GPU feature cache.

The bf16 LoMa models were loaded as capability probes that disabled ONNX Runtime's CPU fallback. ORT deliberately runs small shape computations of the matchers on the CPU, so the bf16 matcher always failed to initialize and fell back to fp32 (#4787). - Only use the bf16 models with the CUDA execution provider on GPUs with compute capability 8.0 or newer, checked before loading any model. - Drop the disabled CPU fallback of capability probes. They still propagate initialization failures without logging an error. - Add CudaDeviceSupportsBf16().

- Remove ONNXModel's is_capability_probe. LoMa now loads the bf16 models as plain sessions with an Ort::Exception fallback, like the FlexGlue fp16 matcher. - Add SelectONNXCudaDevice(), shared by ONNXModel and the LoMa bf16 check, so both resolve gpu_index to the same device. - Skip the CUDA error state check of cuda_test without a GPU, where the initialization error is sticky. - Return false from CudaDeviceSupportsBf16() on HIP. - Remove the Bf16MatcherCpuFallback test, which no longer reaches the fallback.

…4907) ## Summary Adds a `max_num_features` option to `automatic_reconstructor` and to the reconstruction benchmark, to evaluate a feature type with a different number of keypoints than its default. The first use is evaluating learned matchers at 2k/4k/8k keypoints. ### Changes - **`AutomaticReconstructionController::Options::max_num_features`** (default `-1`, off). - When positive, it overrides `max_num_features` of the SIFT, ALIKED and LoMa extraction options. - It is applied after the quality preset, which otherwise sets the number of SIFT features. - **`colmap automatic_reconstructor --max_num_features N`.** - **Benchmark:** `benchmark/reconstruction/evaluate.py --max_num_features N` forwards the option to every `automatic_reconstructor` call (`evaluation/utils.py`). ### Validation - With `--feature loma --max_num_features {2048, 4096, 8192}`, all 898 images of ETH3D DSLR have exactly 2048 / 4096 / 8192 keypoints. - With the option unset, behaviour is unchanged. - `controllers/` tests pass. Part of the FlexGlue stack: #4924, then this PR, #4921, #4908 and #4915. #4908 uses it for its evaluation.

## Summary Adopt the device check and ordinary ONNX Runtime sessions from upstream COLMAP PR colmap#4924. Check BF16 support before loading the model and fall back to FP32 when initialization fails. Remove the graph-assignment probe and its operator allowlist while retaining the fork's unrelated CUDA synchronization work and focused fallback tests. ## Validation - A fresh Linux CUDA build with ONNX Runtime 1.31.1 passed 11 ONNX and five CUDA tests with no skips. - On one real image pair, the BF16 matcher produced 1,848 matches with byte-identical match indices to the pristine upstream output. Separate 128-feature screens produced 68 matches for both CPU selection and GPU initialization-error fallback to FP32. - Native Windows component checks with ONNX Runtime 1.30 passed for the fork and integrated candidate in FP32 and BF16. All four 128-feature self-match outputs were byte-identical to the pristine BF16 result. Unchanged objects were reused; system-header contents were not sealed before their original compilation. - On the GTX 1080 with ONNX Runtime 1.31.1, the integrated component selected FP32 without loading BF16 and produced byte-identical 64- and 128-feature self-matches compared with the pristine and original-fork controls. These are focused functional checks, not full reconstruction comparisons or performance measurements. AI disclosure: totally vibe coded.

fringehush-gif · 2026-10-30T18:19:56Z

I can confirm this work as intended whatever I threw at it.

ORT 1.30 and 1.31.1 produced matching results on RTX 3060 and 3080 Ti under WSL and Windows. Both version selected FP32 before loading BF16 on GTX 1080 on Windows.

Turns out I have quite a lot of stuff lying around so I tried a semi-random matrix of things, including a few different LoMa models not reflected below, both BF16 and FP32, across hardware. This was more of a curiosity and was not tested equally across all scenarios like the ones below.

An unhealthy mix of compilers, CUDA libraries, operating systems and hardware. I couldn't get it to fail. :)

PC / OS	GPU	ORT and CUDA provider	COLMAP build	LoMa result
Main PC · WSL Ubuntu 24.04	RTX 3060 + 3080 Ti	ORT 1.30, CUDA 13	GCC 12.4, nvcc 13.2	Pass. B128 BF16, real image pair: 1,848 raw / 1,706 verified matches on both cards.
Main PC · WSL Ubuntu 24.04	RTX 3060 + 3080 Ti	ORT 1.31.1, CUDA 13	GCC 12.4, nvcc 13.2	Pass. Same pair, counts and output bytes as 1.30.
Main PC · Windows 11	RTX 3060 + 3080 Ti	ORT 1.30, CUDA 13	MSVC 19.51, CUDA 13.3	Pass. LoMa B BF16 component self-match: 128/128 on both cards.
Main PC · Windows 11	RTX 3060 + 3080 Ti	ORT 1.31.1, CUDA 13	MSVC 19.51, CUDA 13.3	Pass. Same 128/128 and identical output bytes.
Older PC · Windows 10	GTX 1080	ORT 1.30, CUDA 12	MSVC 19.39, CUDA 12.9	Pass. A BF16 request correctly selected FP32 before loading BF16: 64/64 and 128/128 self-matches.
Older PC · Windows 10	GTX 1080	ORT 1.31.1, CUDA 12	MSVC 19.39, CUDA 12.9	Pass. Same FP32 fallback and match outputs. No BF16 inference on the GTX 1080.

sarlinpe · 2026-10-30T22:20:19Z

Great, thank you for confirming with such an extensive combination of setups! :)

## Summary Adopt the device check and ordinary ONNX Runtime sessions from upstream COLMAP PR colmap#4924. Check BF16 support before loading the model and fall back to FP32 when initialization fails. Remove the graph-assignment probe and its operator allowlist while retaining the fork's unrelated CUDA synchronization work and focused fallback tests. ## Validation - A fresh Linux CUDA build with ONNX Runtime 1.31.1 passed 11 ONNX and five CUDA tests with no skips. - On one real image pair, the BF16 matcher produced 1,848 matches with byte-identical match indices to the pristine upstream output. Separate 128-feature screens produced 68 matches for both CPU selection and GPU initialization-error fallback to FP32. - Native Windows component checks with ONNX Runtime 1.30 passed for the fork and integrated candidate in FP32 and BF16. All four 128-feature self-match outputs were byte-identical to the pristine BF16 result. Unchanged objects were reused; system-header contents were not sealed before their original compilation. - On the GTX 1080 with ONNX Runtime 1.31.1, the integrated component selected FP32 without loading BF16 and produced byte-identical 64- and 128-feature self-matches compared with the pristine and original-fork controls. These are focused functional checks, not full reconstruction comparisons or performance measurements. AI disclosure: totally vibe coded.

## Summary - Integrate all 12 upstream commits after the shared base 9371f7b, including LoMa BF16 loading, matching-worker error propagation, keypoint shape scaling, and project-file CLI overrides. - Preserve the fork's CUDA blocking-sync helper alongside upstream's BF16 device check and register the shared CUDA test once. ## Validation - Resolved all merge conflicts; no conflict markers or unmerged paths remain; git diff --check passes. - Incremental Linux CUDA build linked six focused test executables. The prepared custom Ceres closure required its existing Abseil/CUDA library directories on the linker's search path. - CTest passed controllers/base_option_manager_test, feature_matching_utils_test, incremental_pipeline_test, option_manager_test, feature/loma_test, and util/cuda_test (six programs, one run each). - No full suite, Python binding, end-to-end, or performance claim. ## Upstream commits - `ea470c0c4f4a40a3bdfa33d01eb38375c5df64de` Run Windows and remaining pycolmap CI only with the full-ci label (colmap#4904) - `c0339e95482c404b51f38ff9b86866de495984f0` Fix MSVC warnings C4805, C4996, and C4756 (colmap#4910) - `f00d6063bb0b120cf10ee190f7f06f4a34edf432` Expose Ceres solver thresholds in incremental mapping (colmap#4892) - `978a4467f36bbfc9ae18780f2dca253000b0d008` Keep division-camera derivatives finite on the optical axis (colmap#4888) - `98f31541f483788ab550f5c4a4e0804e8674b9e0` Fix leak and check covariance in CovarianceWeightedCostFunction (colmap#4911) - `a467519441171054c7aaa17dffe86fe06431f05c` Avoid stack overflow when compressing union-find chains (colmap#4906) - `404b83db2c2cc04892e14a578ddfa73e0ac9bfb1` Fix LoMa bf16 fallback on CUDA by loading it like any other model (colmap#4924) - `d02e6701dede22eaeee89f73784a8c79f6f289ac` Add max_num_features option to automatic_reconstructor and benchmark (colmap#4907) - `82c3b5db78d96a7e93ca7b8b2af0634e28aa33f4` Propagate errors of feature matching workers instead of terminating (colmap#4921) - `4afe8d1156a71a2ff33974895f221c8eb3740cb1` Rescale keypoint shapes in image coordinates (colmap#4926) - `ce2b5ab9df3801c2b75edbba338c04ad3c8cc6af` Fix vocabulary tree builder descriptor dimensions (colmap#4923) - `f956ce629c21a88d75a7f7b3107298d1faf3dcf3` Fix command-line overrides for project files (colmap#4902)

## Summary Adopt the device check and ordinary ONNX Runtime sessions from upstream COLMAP PR colmap#4924. Check BF16 support before loading the model and fall back to FP32 when initialization fails. Remove the graph-assignment probe and its operator allowlist while retaining the fork's unrelated CUDA synchronization work and focused fallback tests. ## Validation - A fresh Linux CUDA build with ONNX Runtime 1.31.1 passed 11 ONNX and five CUDA tests with no skips. - On one real image pair, the BF16 matcher produced 1,848 matches with byte-identical match indices to the pristine upstream output. Separate 128-feature screens produced 68 matches for both CPU selection and GPU initialization-error fallback to FP32. - Native Windows component checks with ONNX Runtime 1.30 passed for the fork and integrated candidate in FP32 and BF16. All four 128-feature self-match outputs were byte-identical to the pristine BF16 result. Unchanged objects were reused; system-header contents were not sealed before their original compilation. - On the GTX 1080 with ONNX Runtime 1.31.1, the integrated component selected FP32 without loading BF16 and produced byte-identical 64- and 128-feature self-matches compared with the pristine and original-fork controls. These are focused functional checks, not full reconstruction comparisons or performance measurements. AI disclosure: totally vibe coded.

sarlinpe mentioned this pull request Oct 29, 2026

Fix LoMa BF16 CUDA probing with CPU bookkeeping nodes #4895

Closed

sarlinpe added 2 commits October 29, 2026 22:38

Address review nits: name the GPU in the bf16 warning, tighten tests

3ecaa71

sarlinpe added the full-ci Run Windows and remaining pycolmap CI label Oct 29, 2026

sarlinpe requested a review from ahojnnes October 29, 2026 22:49

sarlinpe added this pull request to stack #4925 October 29, 2026 23:28

This was referenced Oct 29, 2026

Add FlexGlue feature matcher #4908

Merged

FlexGlue: cache the dense features on the GPU #4915

Merged

Add max_num_features option to automatic_reconstructor and benchmark #4907

Merged

ahojnnes approved these changes Oct 30, 2026

View reviewed changes

ahojnnes merged commit 404b83d into main Oct 30, 2026
18 of 21 checks passed

ahojnnes deleted the sarlinpe/loma-bf16-probe branch October 30, 2026 08:37

Navigation Menu

Repository navigation

Fix LoMa bf16 fallback on CUDA by loading it like any other model - #4924

Fix LoMa bf16 fallback on CUDA by loading it like any other model#4924
ahojnnes merged 3 commits into
mainfrom
sarlinpe/loma-bf16-probe

sarlinpe commented Oct 29, 2026 •

edited

Loading

Uh oh!

Uh oh!

fringehush-gif commented Oct 30, 2026

Uh oh!

sarlinpe commented Oct 30, 2026

Uh oh!

Reviewers

Uh oh!

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants

Conversation

sarlinpe commented Oct 29, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Problem

Fix

Comparison with #4895

Tests

Uh oh!

Uh oh!

fringehush-gif commented Oct 30, 2026

Uh oh!

sarlinpe commented Oct 30, 2026

Uh oh!

Reviewers

Uh oh!

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants

sarlinpe commented Oct 29, 2026 •

edited

Loading