Repository navigation
fix(conv): use the fp32 fallback when set_precision requests float32 or float64 - #13
Merged
Merged
Conversation
…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.
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Context
Conv2dignored an explicitset_precision("float32")on a GPU: it stayed on the fp16 WMMApath 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(...)"float16""float32""float64"_compile_for_device._apply_precisionre-runs it on an alreadycompiled layer, so
set_precision()works before or afterto()/finalize(), and"float32"→"float16"returns to WMMA.THIS IS MUCH SLOWERand points toset_precision('float16')or leaving it unset.float64still computes in float32: the fallback is fp32 only.Model.set_precisionand theConv2ddocstrings now describe this.Behavior changes
set_precision("float32")or("float64")on a GPU now gives slower, real fp32 convolutionsinstead of fp16 WMMA ones, plus a
UserWarning. Not callingset_precisionor using"float16"is unchanged."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 theROCm devcontainer.
"float32"and"float64"bind thefallback and warn once;
Noneand"float16"stay on WMMA silently;"float32"→"float16"rebinds WMMA; the
"float32"forward is bit-exact with the fallback. The first, second andlast fail on the old code. The other two pin behavior that was already correct.
test_mismatched_policy_warns_once_naming_both_dtypesnow reads the effective dtype afterapplying
float64, since that changes the binding.examples/cifar10/v1_cifar10.pyI changedset_precision("float16")toset_precision("float32"). The warning appears as intended andtraining continues on the fallback, which is far slower than the kernel path.
Follow-ups (out of scope)
tests/ops/test_conv_cast_kernel.py:55(test_matches_clip_and_astype) emitsRuntimeWarning: overflow encountered in cast. The test intentionally casts inputs above65504 to fp16, which overflow to
inf, and doesn't suppress the expected NumPy warning. Theassertions are correct; the fix is in the next PR.