[CUDNN] Support BFloat16 #2987
Open
[CUDNN] Support BFloat16 #2987
+27
−1
Conversation
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
Contributor
|
Your PR requires formatting changes to meet the project's style guidelines. Click here to view the suggested changes.diff --git a/lib/cudnn/src/util.jl b/lib/cudnn/src/util.jl
index 8923ff9b5..c7ec0c2bd 100644
--- a/lib/cudnn/src/util.jl
+++ b/lib/cudnn/src/util.jl
@@ -4,13 +4,13 @@ using BFloat16s: BFloat16
cptr(x,a::DenseCuArray{Float64})=Float64[x]
cptr(x,a::DenseCuArray{Float32})=Float32[x]
cptr(x,a::DenseCuArray{Float16})=Float32[x]
-cptr(x,a::DenseCuArray{BFloat16})=Float32[x]
+cptr(x, a::DenseCuArray{BFloat16}) = Float32[x]
# Conversion between Julia and cuDNN datatypes
cudnnDataType(::Type{Float16})=CUDNN_DATA_HALF
cudnnDataType(::Type{Float32})=CUDNN_DATA_FLOAT
cudnnDataType(::Type{Float64})=CUDNN_DATA_DOUBLE
-cudnnDataType(::Type{BFloat16})=CUDNN_DATA_BFLOAT16
+cudnnDataType(::Type{BFloat16}) = CUDNN_DATA_BFLOAT16
cudnnDataType(::Type{Int8}) = CUDNN_DATA_INT8
cudnnDataType(::Type{UInt8}) = CUDNN_DATA_UINT8
cudnnDataType(::Type{Int32}) = CUDNN_DATA_INT32
@@ -21,7 +21,7 @@ cudnnDataType(::Type{Int32}) = CUDNN_DATA_INT32
juliaDataType(a)=(a==CUDNN_DATA_HALF ? Float16 :
a==CUDNN_DATA_FLOAT ? Float32 :
a==CUDNN_DATA_DOUBLE ? Float64 :
- a==CUDNN_DATA_BFLOAT16 ? BFloat16 :
+ a == CUDNN_DATA_BFLOAT16 ? BFloat16 :
a==CUDNN_DATA_INT8 ? Int8 :
a==CUDNN_DATA_UINT8 ? UInt8 :
a==CUDNN_DATA_INT32 ? Int32 : error())
diff --git a/lib/cudnn/test/activation.jl b/lib/cudnn/test/activation.jl
index 7b7f2f01a..4164e0231 100644
--- a/lib/cudnn/test/activation.jl
+++ b/lib/cudnn/test/activation.jl
@@ -62,8 +62,8 @@ activationtest(alpha=2)
activationtest(beta=2)
# BFloat16 tests
-(ax,ay) = randn.(BFloat16, (10,10))
-(cx,cy) = CuArray.((ax,ay))
-activationtest(mode=CUDNN_ACTIVATION_SIGMOID)
-activationtest(mode=CUDNN_ACTIVATION_RELU)
-activationtest(mode=CUDNN_ACTIVATION_TANH)
+(ax, ay) = randn.(BFloat16, (10, 10))
+(cx, cy) = CuArray.((ax, ay))
+activationtest(mode = CUDNN_ACTIVATION_SIGMOID)
+activationtest(mode = CUDNN_ACTIVATION_RELU)
+activationtest(mode = CUDNN_ACTIVATION_TANH)
diff --git a/lib/cudnn/test/softmax.jl b/lib/cudnn/test/softmax.jl
index 68967bc1d..ab446813c 100644
--- a/lib/cudnn/test/softmax.jl
+++ b/lib/cudnn/test/softmax.jl
@@ -46,7 +46,7 @@ softmaxtest(algo=CUDNN_SOFTMAX_ACCURATE)
softmaxtest(algo=CUDNN_SOFTMAX_LOG)
# BFloat16 tests
-ax,ay = randn(BFloat16,10,10),randn(BFloat16,10,10)
-cx,cy = CuArray.((ax,ay))
+ax, ay = randn(BFloat16, 10, 10), randn(BFloat16, 10, 10)
+cx, cy = CuArray.((ax, ay))
softmaxtest()
-softmaxtest(algo=CUDNN_SOFTMAX_LOG)
+softmaxtest(algo = CUDNN_SOFTMAX_LOG) |
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
1 participant
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.
This PR defines methods for making cuDNN work with
BFloat16s.BFloat16.In the following example, I show how the new methods fixes the
BFloat16backward pass ofFlux.logitcrossentropy:Before
Note: Core.BFloat16 === BFloat16s.BFloat16, but I didn't explicitly import in this REPL session.
After defining cudnnDataType(::Type{BFloat16})
After defining scalingParameter(::Type{BFloat16}, val)
I also define a
cptrmethod for consistency, but it appears the function isn't used anywhere.Tests are added for softmax, activations, and pooling. I initially also tested convolutions, normalization, RNNs, and MHA but they don't appear to support BFloat16.
Adding BFloat16s.jl as a dependency does not affect compilation since it's already a dependency of CUDA.jl.
Along with my proposed fix in FluxML/Optimisers.jl#215, this has allowed me to train LLMs in BFloat16 with Flux.jl in Julia v1.12. I am still working on an Optimisers.jl, but these together would be a significant unlock for my lab.