-
-
Notifications
You must be signed in to change notification settings - Fork 129
Open
Labels
Description
Since #455 is merged, I want to point out that CUDNN_SOFTMAX_FAST would easily cause problem for attention operation. In the masking scenario, we would usually set the masked value to -Inf or some really small value, like -1e9. But if we want to use CUDA.math_mode!(CUDA.FAST_MATH) to accelerate the gemm, softmax would actually introduce many NaNs.
MWE:
julia> using CUDA, Flux
julia> x = CUDA.randn(Float32, 512, 10); fill!(x, -1f3);
julia> CUDA.math_mode!(CUDA.FAST_MATH)
julia> softmax(x)
512×10 CuArray{Float32, 2, CUDA.Mem.DeviceBuffer}:
NaN NaN NaN NaN NaN NaN NaN NaN NaN NaN
NaN NaN NaN NaN NaN NaN NaN NaN NaN NaN
NaN NaN NaN NaN NaN NaN NaN NaN NaN NaN
⋮ ⋮
NaN NaN NaN NaN NaN NaN NaN NaN NaN NaN
NaN NaN NaN NaN NaN NaN NaN NaN NaN NaN
NaN NaN NaN NaN NaN NaN NaN NaN NaN NaN
NaN NaN NaN NaN NaN NaN NaN NaN NaN NaN