Skip to content

Disable CUDNN_SOFTMAX_FAST or use a separate math mode variable for softmax #506

Open
@chengchingwen

Description

@chengchingwen

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

Metadata

Metadata

Assignees

No one assigned

    Labels

    Type

    No type

    Projects

    No projects

    Milestone

    No milestone

    Relationships

    None yet

    Development

    No branches or pull requests

    Issue actions