Skip to content

fix(conv): use the fp32 fallback when set_precision requests float32 or float64 - #13

Merged
AlexanderSoftCode merged 1 commit into
mainfrom
fix/conv-honor-fp32-precision
Oct 30, 2026
Merged

AlexanderSoftCode merged 1 commit into
mainfrom
fix/conv-honor-fp32-precision

Conversation

@AlexanderSoftCode

Copy link
Copy Markdown
Owner

Context

Conv2d ignored an explicit set_precision("float32") on a GPU: it stayed on the fp16 WMMA
path and only warned that the policy was ignored. It now honors the request by binding the fp32
fallback, and warns that this is much slower. One commit.

Conv precision

set_precision(...) GPU with WMMA before
not called / "float16" WMMA, no warning same
"float32" fp32 fallback, slow warning WMMA, "policy is ignored"
"float64" fp32 fallback, slow warning that also says float64 isn't honored WMMA, "policy is ignored"
  • The binding happens in _compile_for_device. _apply_precision re-runs it on an already
    compiled layer, so set_precision() works before or after to()/finalize(), and
    "float32" → "float16" returns to WMMA.
  • The warning has no speed ratio (it would go stale as the paths change). It says
    THIS IS MUCH SLOWER and points to set_precision('float16') or leaving it unset.
  • float64 still computes in float32: the fallback is fp32 only.
  • Model.set_precision and the Conv2d docstrings now describe this.

Behavior changes

  • set_precision("float32") or ("float64") on a GPU now gives slower, real fp32 convolutions
    instead of fp16 WMMA ones, plus a UserWarning. Not calling set_precision or using
    "float16" is unchanged.
  • A model saved under "float32" loads onto the fallback path and warns once.

Testing

  • python3 -m unittest discover tests: 1125 tests pass on NumPy and CuPy (59 skipped), in the
    ROCm devcontainer.
  • 5 new GPU-only tests (skipped on NumPy or without WMMA): "float32" and "float64" bind the
    fallback and warn once; None and "float16" stay on WMMA silently; "float32" → "float16"
    rebinds WMMA; the "float32" forward is bit-exact with the fallback. The first, second and
    last fail on the old code. The other two pin behavior that was already correct.
  • test_mismatched_policy_warns_once_naming_both_dtypes now reads the effective dtype after
    applying float64, since that changes the binding.
  • Checked by hand on the RX 9070 XT: in examples/cifar10/v1_cifar10.py I changed
    set_precision("float16") to set_precision("float32"). The warning appears as intended and
    training continues on the fallback, which is far slower than the kernel path.
  • Not verified: CUDA, the RTX 3050, and a GPU without tensor cores.

Follow-ups (out of scope)

  • tests/ops/test_conv_cast_kernel.py:55 (test_matches_clip_and_astype) emits
    RuntimeWarning: overflow encountered in cast. The test intentionally casts inputs above
    65504 to fp16, which overflow to inf, and doesn't suppress the expected NumPy warning. The
    assertions are correct; the fix is in the next PR.

…or float64

Conv2d stayed on the fp16 WMMA path under set_precision("float32") on a
GPU and only warned that the policy was ignored. An explicit float32 or
float64 request now binds the strided fp32 fallback instead, and warns
that this is much slower, pointing at set_precision("float16") or
leaving it unset. The warning gives no speed ratio, since that would go
stale as the paths change.

_compile_for_device picks the path, and _apply_precision re-runs it on a
compiled layer, so set_precision() works before or after to()/finalize()
and float32 -> float16 returns to WMMA. float64 still computes in
float32, because the fallback is fp32 only; the warning says so. Not
calling set_precision, or float16, is unchanged.

Model.set_precision and the Conv2d docstrings describe the new behavior.
Adds GPU-only tests for each case, including a bit-exact check against
the fallback, and updates the one test whose expected dtype changed.
@AlexanderSoftCode
AlexanderSoftCode merged commit 886466b into main Oct 30, 2026
5 checks passed
@AlexanderSoftCode
AlexanderSoftCode deleted the fix/conv-honor-fp32-precision branch October 30, 2026 01:47
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant