Repository navigation
Fix LoMa BF16 CUDA probing with CPU bookkeeping nodes - #4895
fringehush-gif wants to merge 2 commits into
Conversation
Inspect CUDA graph assignments so LoMa BF16 accepts supported CPU bookkeeping nodes while keeping numerical work on CUDA. ### Changes - Allow selected CPU shape/index operators during CUDA capability probes; reject other CPU operators and require CUDA work. - Keep strict probing on older ONNX Runtime and CoreML, and FP32 fallback on pre-Ampere GPUs. ### Tests - **RTX 3060 / ORT 1.30:** 11 tests pass. Real LoMa B FP32/BF16 self-matches: **64/64 and 128/128**; default FP32 matches the base exactly. - Reproduced the original strict BF16 probe failure. ### Compatibility - **Windows 10 / GTX 1080 / ORT 1.30 CUDA 12:** 11 tests pass without skips. GPU FP32 and BF16-to-FP32 fallback match base and CPU exactly at 64/128 features. - ORT 1.23.2 and 1.30 CPU SDKs: **9 tests each, 20 matching calls**, exact pairs. ### Build - Changed CUDA-enabled source units compile against both SDKs. - Full CLI links; `--help` exits 0. Fixes colmap#4787. Totally vibe coded
Inspect CUDA graph assignments so LoMa BF16 accepts supported CPU bookkeeping nodes while keeping numerical work on CUDA. ### Changes - Allow selected CPU shape/index operators during CUDA capability probes; reject other CPU operators and require CUDA work. - Keep strict probing on older ONNX Runtime and CoreML, and FP32 fallback on pre-Ampere GPUs. ### Tests - **RTX 3060 / ORT 1.30:** 11 tests pass. Real LoMa B FP32/BF16 self-matches: **64/64 and 128/128**; default FP32 matches the base exactly. - Reproduced the original strict BF16 probe failure. ### Compatibility - **Windows 10 / GTX 1080 / ORT 1.30 CUDA 12:** 11 tests pass without skips. GPU FP32 and BF16-to-FP32 fallback match base and CPU exactly at 64/128 features. - ORT 1.23.2 and 1.30 CPU SDKs: **9 tests each, 20 matching calls**, exact pairs. ### Build - Changed CUDA-enabled source units compile against both SDKs. - Full CLI links; `--help` exits 0. Upstream review: colmap#4895. Totally vibe coded
- Fix CMake plain/keyword signature mixing for colmap_util_onnx_test, which broke configure on CUDA builds - Link colmap_feature against colmap_util_cuda (loma.cc needs the symbol) - Extract generic CudaDeviceSupportsBf16() helper to util/cuda and replace the LoMa-internal BF16 device check with a thin wrapper; the helper fails closed (false) on invalid input or query failure - Document that the CUDA probe inspects operator types only, and log allowed CPU nodes at VLOG(2) - Add cuda_test for the helper; document the onnx_test probe models
|
Thanks a lot for tracking this down and for the thorough testing! I opened #4924 as an alternative fix for #4787 and would suggest superseding this PR with it. The reasoning: The root cause is the strict probe, not specific operators. ONNX Runtime's CPU placement (
Dropping The compute capability check is better done before loading. Here it runs after the bf16 session (and, for extraction, the detector) has been downloaded and built, and throws a synthetic #4924 thus:
It is +130/−95 (net +35 lines, vs. +255 here), without hand-encoded ONNX test models. Note that the ClangTidy job also fails here ( Since you can reproduce the issue on both an RTX 3060 and a GTX 1080, it would be great if you could check that #4924 runs the bf16 matcher on the 3060 and falls back to fp32 on the 1080 — I could not run the ORT 1.30 CUDA EP locally. |
) 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 `ONNXModel`s, 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 #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.
|
That explanation went a bit over my head but I think I got the general gist. I appreciate the explanation. I'm a total novice and this is entirely vibe coded like I mentioned, but anyway. I noticed it's been merged into main since but I will test it and post a follow up in #4924. |
Inspect CUDA graph assignments so LoMa BF16 accepts supported CPU bookkeeping nodes while keeping numerical work on CUDA. ### Changes - Allow selected CPU shape/index operators during CUDA capability probes; reject other CPU operators and require CUDA work. - Keep strict probing on older ONNX Runtime and CoreML, and FP32 fallback on pre-Ampere GPUs. ### Tests - **RTX 3060 / ORT 1.30:** 11 tests pass. Real LoMa B FP32/BF16 self-matches: **64/64 and 128/128**; default FP32 matches the base exactly. - Reproduced the original strict BF16 probe failure. ### Compatibility - **Windows 10 / GTX 1080 / ORT 1.30 CUDA 12:** 11 tests pass without skips. GPU FP32 and BF16-to-FP32 fallback match base and CPU exactly at 64/128 features. - ORT 1.23.2 and 1.30 CPU SDKs: **9 tests each, 20 matching calls**, exact pairs. ### Build - Changed CUDA-enabled source units compile against both SDKs. - Full CLI links; `--help` exits 0. Upstream review: colmap#4895. Totally vibe coded
Inspect CUDA graph assignments so LoMa BF16 accepts supported CPU bookkeeping nodes while keeping numerical work on CUDA.
Changes
Tests
Compatibility
Build
--helpexits 0.Fixes #4787.
Totally vibe coded