Fix layer_norm backward to reuse the forward dropout mask and scale dx1 by rowscale - #1023
Open
cyboTiger wants to merge 1 commit into
Open
Fix layer_norm backward to reuse the forward dropout mask and scale dx1 by rowscale#1023cyboTiger wants to merge 1 commit into
cyboTiger wants to merge 1 commit into
Conversation
…k and support scale dx1 by rowscale
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.
Summary
This PR fixes two issues in triton kernel
mamba/blob/main/mamba_ssm/ops/triton/layernorm.py:tl.randto regenerate dropout mask rather than reuse dropout mask generated in forward pass, which can cause incorrect gradient backpropagation during trainingrow_scaleandx1cannot co-exist in backward kernelDescription
The backward kernel regenerated the dropout mask from the seeds with
tl.randinstead of reusing the mask stored in the forward pass. Correctness therefore depended on the Triton PRNG being stable across two separately-compiled kernels (which is not guaranteed across Triton versions). The forward now always stores the mask when dropout is enabled and the backward consumes it (dropout_maskreplacesseedsin_layer_norm_bwd).dx1(parallel LayerNorm) was written out without applying the second half ofrowscale, unlikedxand unlike the forward kernel, which scalesx1byrowscale[M + row]. The forward host asserts that rejectedrowscaletogether withx1are removed (the kernels already supported it;rowscalenow has2 * Mrows in that case), and the backward applies the scale.Reference functions fixed for 3-D inputs with
rowscale.Tests
new
tests/ops/triton/test_layernorm.pycovering forward/backward correctness across dropout/x1/residual/bias/rms/rowscale/dtype combinations, plus regression tests that fail on the old implementation (doctored-mask reuse, rowscale-scaleddx1).