Prevent half-precision mutual information overflow#9023
Conversation
Signed-off-by: kyinhub <kevinpyin@gmail.com>
📝 WalkthroughWalkthrough
Estimated code review effort: 3 (Moderate) | ~20 minutes 🚥 Pre-merge checks | ✅ 5✅ Passed checks (5 passed)
✨ Finishing Touches🧪 Generate unit tests (beta)
Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out. Comment |
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 1
🧹 Nitpick comments (3)
monai/losses/image_dissimilarity.py (2)
327-336: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winUndocumented float16-only workaround + duplicated dtype-selection logic.
The early
img.float()conversion forfloat16beforeclamp(but not forbfloat16) looks asymmetric/dead at first glance, but it's actually a workaround: PyTorch has historically lacked CPU kernels forclamponHalf, e.g. RuntimeError: "clamp_scalar_cpu" not implemented for 'Half', whilebfloat16has materially better CPU op coverage. Worth a one-line comment so a future refactor doesn't "clean up" this asymmetry into a bug.Separately, the
compute_dtype = torch.float32 if X.dtype in (torch.float16, torch.bfloat16) else X.dtypeexpression is duplicated verbatim (structurally) here and inforward(Line 364). Extracting a tiny helper avoids drift if the reduced-precision dtype policy ever changes.♻️ Suggested cleanup
+ `@staticmethod` + def _reduced_precision_compute_dtype(dtype: torch.dtype) -> torch.dtype: + """Returns float32 for float16/bfloat16 inputs, otherwise the input dtype unchanged.""" + return torch.float32 if dtype in (torch.float16, torch.bfloat16) else dtype + def parzen_windowing_gaussian(self, img: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]: ... output_dtype = img.dtype + # torch.clamp lacks a CPU kernel for float16 (bfloat16 has broader CPU op support), + # so upcast float16 before clamping to avoid a runtime error. if output_dtype == torch.float16: img = img.float() - compute_dtype = torch.float32 if img.dtype in (torch.float16, torch.bfloat16) else img.dtype + compute_dtype = self._reduced_precision_compute_dtype(img.dtype)🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@monai/losses/image_dissimilarity.py` around lines 327 - 336, Add a concise comment before the float16-only conversion in the relevant loss computation, documenting the CPU clamp kernel limitation and preserving the intentional asymmetry with bfloat16. Extract the duplicated compute-dtype selection used here and in forward into a small shared helper, then update both call sites to use it.
358-368: 🎯 Functional Correctness | 🔵 Trivial | ⚡ Quick winUnnecessary float32→float16/bfloat16→float32 round-trip before histogram accumulation.
parzen_windowing_gaussiandowncastsweight/probabilityback to the original reduced-precision dtype, thenforwardimmediately upcasts them back to float32 for thebmm/logcomputation. Sinceparzen_windowingis only invoked internally here, this quantizes the Parzen weights to float16/bfloat16 precision for no benefit before the very computation this PR is trying to make numerically stable.Consider adding an internal flag to skip the final downcast when called from
forward(keeping the public dtype-preservation contract intact for direct callers/tests):♻️ Suggested approach
- def parzen_windowing_gaussian(self, img: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]: + def parzen_windowing_gaussian( + self, img: torch.Tensor, restore_input_dtype: bool = True + ) -> tuple[torch.Tensor, torch.Tensor]: ... - if output_dtype in (torch.float16, torch.bfloat16): + if restore_input_dtype and output_dtype in (torch.float16, torch.bfloat16): weight = weight.to(dtype=output_dtype) probability = probability.to(dtype=output_dtype) return weight, probability
parzen_windowing/forwardwould need to threadrestore_input_dtype=Falsethrough for the internal call path.🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@monai/losses/image_dissimilarity.py` around lines 358 - 368, Add an internal restore_input_dtype control to parzen_windowing_gaussian/parzen_windowing and pass restore_input_dtype=False from forward, so the internal path retains float32 accumulation inputs instead of round-tripping through float16/bfloat16. Preserve the existing dtype-restoration behavior by default for direct callers and tests, then keep forward’s histogram computation unchanged.tests/losses/image_dissimilarity/test_global_mutual_information_loss.py (1)
203-217: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winConsider asserting
result.dtypefor symmetry with the sibling test.
test_half_precision_large_constant_volume_is_finitechecksresult.dtype == pred.dtype; this autocast variant doesn't, leaving the dtype-preservation contract unverified under autocast specifically.✅ Suggested addition
self.assertTrue(torch.isfinite(result)) + self.assertEqual(result.dtype, pred.dtype) result.backward()🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@tests/losses/image_dissimilarity/test_global_mutual_information_loss.py` around lines 203 - 217, Update test_cpu_float16_autocast_large_volume_is_finite to assert that result.dtype matches pred.dtype, mirroring the existing dtype assertion in test_half_precision_large_constant_volume_is_finite while preserving the other finiteness, gradient, and device checks.
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Inline comments:
In `@monai/losses/image_dissimilarity.py`:
- Around line 346-357: Update the docstring for the loss method containing the
target.shape validation to document the ValueError raised when pred and target
shapes differ, alongside the existing reduction ValueError in the Raises
section.
---
Nitpick comments:
In `@monai/losses/image_dissimilarity.py`:
- Around line 327-336: Add a concise comment before the float16-only conversion
in the relevant loss computation, documenting the CPU clamp kernel limitation
and preserving the intentional asymmetry with bfloat16. Extract the duplicated
compute-dtype selection used here and in forward into a small shared helper,
then update both call sites to use it.
- Around line 358-368: Add an internal restore_input_dtype control to
parzen_windowing_gaussian/parzen_windowing and pass restore_input_dtype=False
from forward, so the internal path retains float32 accumulation inputs instead
of round-tripping through float16/bfloat16. Preserve the existing
dtype-restoration behavior by default for direct callers and tests, then keep
forward’s histogram computation unchanged.
In `@tests/losses/image_dissimilarity/test_global_mutual_information_loss.py`:
- Around line 203-217: Update test_cpu_float16_autocast_large_volume_is_finite
to assert that result.dtype matches pred.dtype, mirroring the existing dtype
assertion in test_half_precision_large_constant_volume_is_finite while
preserving the other finiteness, gradient, and device checks.
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Pro Plus
Run ID: 61ddddab-1114-4363-8c6e-caadeba8d25b
📒 Files selected for processing (2)
monai/losses/image_dissimilarity.pytests/losses/image_dissimilarity/test_global_mutual_information_loss.py
1be4b14 to
9c32c55
Compare
|
@coderabbitai review |
✅ Action performedReview finished.
|
There was a problem hiding this comment.
🧹 Nitpick comments (1)
tests/losses/image_dissimilarity/test_global_mutual_information_loss.py (1)
195-196: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winDocument the parameterized argument.
Add Google-style
Args:entries fordtypein both test docstrings.
tests/losses/image_dissimilarity/test_global_mutual_information_loss.py#L195-L196: documentdtype.tests/losses/image_dissimilarity/test_global_mutual_information_loss.py#L233-L234: documentdtype.As per path instructions, “Docstrings should be present for all definition which describe each variable, return value, and raised exception in the appropriate section of the Google-style of docstrings.”
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@tests/losses/image_dissimilarity/test_global_mutual_information_loss.py` around lines 195 - 196, Update the docstrings for test_module_cast_with_many_bins_remains_finite at tests/losses/image_dissimilarity/test_global_mutual_information_loss.py:195-196 and the test at tests/losses/image_dissimilarity/test_global_mutual_information_loss.py:233-234 to include Google-style Args: sections documenting the parameterized dtype argument; make no other changes.Source: Path instructions
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Nitpick comments:
In `@tests/losses/image_dissimilarity/test_global_mutual_information_loss.py`:
- Around line 195-196: Update the docstrings for
test_module_cast_with_many_bins_remains_finite at
tests/losses/image_dissimilarity/test_global_mutual_information_loss.py:195-196
and the test at
tests/losses/image_dissimilarity/test_global_mutual_information_loss.py:233-234
to include Google-style Args: sections documenting the parameterized dtype
argument; make no other changes.
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Pro Plus
Run ID: 8d459577-2b3c-46de-9958-44bd80925fea
📒 Files selected for processing (2)
monai/losses/image_dissimilarity.pytests/losses/image_dissimilarity/test_global_mutual_information_loss.py
🚧 Files skipped from review as they are similar to previous changes (1)
- monai/losses/image_dissimilarity.py
Signed-off-by: kyinhub <kevinpyin@gmail.com>
9c32c55 to
76a50a8
Compare
|
@coderabbitai review |
✅ Action performedReview finished.
|
1 participant
Fixes #7622.
Description
Gaussian global mutual information now avoids reduced-precision overflow in two places:
Public weight, probability, loss, device, and gradient dtype contracts remain unchanged; float64 inputs continue to compute in float64. A finite Python scalar fallback preserves the Gaussian preterm when module-wide float16 conversion would otherwise overflow the registered buffer.
Types of changes
./runtests.sh -f -u --net --coverage../runtests.sh --quick --unittests --disttests.make htmlcommand in thedocs/folder.Validation
upstream/devproduces non-finite many-bin float16 Gaussian weights;upstream/devreturns-infwith non-finite gradients on a48³volume;git diff --checkpass.Scope note
The Parzen helper intentionally preserves its historical reduced-precision return dtype before the histogram stage. Avoiding that internal downcast would be a separate numerical-accuracy change rather than the scoped overflow fix.
Sequencing
This PR remains draft and will be rebased after #9022 because both modify
GlobalMutualInformationLoss.