[ET-VK] Clamp the tanh argument in the gelu shader - #22324
Conversation
🔗 Helpful Links🧪 See artifacts and rendered test results at hud.pytorch.org/pr/pytorch/executorch/22324
Note: Links to docs will display an error until the docs builds have been completed.
|
meta-cla
Bot
added
the
CLA Signed
This PR needs a
|
github-actions
Bot
added
ciflow/trunk
module: arm
|
The following ciflow label(s) have been added but CI has not been triggered yet because the workflows are awaiting approval:
Once a maintainer approves the workflows (scroll to the bottom of the PR page), the corresponding CI jobs will be triggered automatically. Please ping one of the reviewers if you do not have access to approve and run workflows. |
1 similar comment
|
The following ciflow label(s) have been added but CI has not been triggered yet because the workflows are awaiting approval:
Once a maintainer approves the workflows (scroll to the bottom of the PR page), the corresponding CI jobs will be triggered automatically. Please ping one of the reviewers if you do not have access to approve and run workflows. |
The gelu shader evaluates the tanh approximation with an unclamped
argument. For an input of x the argument is
sqrt(2/pi) * (x + 0.044715 * x^3)
which grows cubically, so x = -13.24 already yields -93.4. A driver
that evaluates tanh as (e^y - e^-y) / (e^y + e^-y) overflows fp32 at
|y| > ~88 and returns inf/inf = NaN. On a Mali-G76 this makes the
Whisper encoder emit NaN from conv2 for exactly the two activations
below -13, and the first LayerNorm then propagates them across the
whole tensor.
The tanh op in this same file already clamps to +/-15 for this reason;
gelu now does the same. The clamp is numerically free: 1 - tanh(15) is
1.9e-13, well below fp32 epsilon, so the result is bit-identical for
every input that gets clamped.
fbfdafc to
11bac85
Compare
|
Sorry guys, my fault with the original push. |
2 participants
Fixes #22323.
The
gelushader evaluates the tanh approximation with an unclamped argument. That argument grows cubically in the input, sox = -13.24already yields-93.4, and a driver that computestanh(y)as(e^y - e^-y) / (e^y + e^-y)overflows fp32 at|y| > ~88and returnsinf/inf = NaN.The
tanhop in this same file already clamps to+/-15for exactly this reason. This change makesgeludo the same.The clamp costs no accuracy:
1 - tanh(15) = 1.9e-13, well below fp32 epsilon, so the output is bit-identical for every input that gets clamped.Verification
Whisper-tiny encoder, Vulkan delegate, Samsung Galaxy S10+ (Mali-G76), against a CPU reference:
Before the fix the failure is silent: the runner reports success and normal per-iteration timings while returning an all-NaN tensor.
cc @digantdesai @freddan80 @per @zingo @oscarandersson8218 @mansnils @Sebastian-Larsson @robell @rascani