Skip to content

Fix layer_norm backward to reuse the forward dropout mask and scale dx1 by rowscale - #1023

Open
cyboTiger wants to merge 1 commit into
state-spaces:mainfrom
cyboTiger:main
Open

Fix layer_norm backward to reuse the forward dropout mask and scale dx1 by rowscale#1023
cyboTiger wants to merge 1 commit into
state-spaces:mainfrom
cyboTiger:main

Conversation

@cyboTiger

Copy link
Copy Markdown

Summary

This PR fixes two issues in triton kernel mamba/blob/main/mamba_ssm/ops/triton/layernorm.py:

  • the backward kernel use tl.rand to regenerate dropout mask rather than reuse dropout mask generated in forward pass, which can cause incorrect gradient backpropagation during training
  • row_scale and x1 cannot co-exist in backward kernel

Description

The backward kernel regenerated the dropout mask from the seeds with tl.rand instead 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_mask replaces seeds in _layer_norm_bwd).

dx1 (parallel LayerNorm) was written out without applying the second half of rowscale, unlike dx and unlike the forward kernel, which scales x1 by rowscale[M + row]. The forward host asserts that rejected rowscale together with x1 are removed (the kernels already supported it; rowscale now has 2 * M rows 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.py covering 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-scaled dx1).

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant