Repository navigation
Add a gelu_tanh op with f32-internal forward and backward kernels to candle-nn - #4060
Open
DrJesseGlass wants to merge 2 commits into
Open
DrJesseGlass wants to merge 2 commits into
DrJesseGlass wants to merge 2 commits into
Conversation
Tensor::gelu evaluates the tanh approximation in the tensor dtype through the unary macro, so for bf16 it rounds at every intermediate, and its composed backward rounds again at each step. candle_nn::ops::gelu_tanh is a dedicated op whose CUDA and CPU forward and backward compute in f32 (f64 for f64) and round once at the store. Backends without the kernel fall back to composed f32 ops, which autograd differentiates itself. Tests check forward and gradient parity against Tensor::gelu on CPU and CUDA, and the single-rounding property for bf16.
The CPU evaluators computed every dtype in f64 while the CUDA kernels and the composed fallback use f32 for the half types and f32, so results depended on the backend (an input where f32 tanh saturates to -1 gave zero on CUDA and a small nonzero value on CPU). The evaluators are now generic over the accumulator: f32 for bf16, f16 and f32, f64 for f64, matching CUDA.
This branch has not been deployed
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.
Tensor::gelu (the tanh approximation) runs its bf16 and f16 kernels in half arithmetic and its backward as composed half ops, rounding at every intermediate. When training in bf16, the gradient error compounds and the composed backward is slow relative to the forward. #1982 raised the half-precision gelu cost on CPU in 2024.
This PR adds candle_nn::ops::gelu_tanh, a CustomOp1 with a backward, on CPU and CUDA. Forward and backward evaluate in f32 for the half types and f32, in f64 for f64, with one rounding at the end. On other backends, it falls back to a composed form in f32 that autograd differentiates. Tensor::gelu is untouched, so no existing model changes numerics; the op is opt-in for training code. The CPU evaluators are generic over the accumulator so the half types and f32 take the same path the CUDA kernels do.
The test checks forward and dx against Tensor::gelu in f32 and the bf16 forward against the f32 evaluation of the same bf16 inputs to within one bf16 ulp on CPU, CUDA, and Metal.
On an arch64 CPU, forward plus backward is roughly 11x faster than Tensor::gelu in bf16 and 8x in f32 at (4, 256, 8192). This is the trainer-side counterpart of #3722, which vectorizes the inference-side Tensor::gelu on aarch64.