modify rms norm

DeepLink-org · Apr 1, 2024 · 3f894a2 · 3f894a2
1 parent 27da2ec
commit 3f894a2
Show file tree

Hide file tree

Showing 3 changed files with 2 additions and 2 deletions.
diff --git a/deeplink_ext/common/rms_norm/__init__.py b/deeplink_ext/common/rms_norm/__init__.py
@@ -1,4 +1,4 @@
 from .deeplink import rms_norm_out, rms_norm, rms_norm_backward_out, rms_norm_backward
 
 
-all = ["rms_norm_out", "rms_norm", "rms_norm_backward_out", "rms_norm_backward"]
+all = ["rms_norm_out", "rms_norm", "rms_norm_backward_out", "rms_norm_backward"]
diff --git a/deeplink_ext/common/rms_norm/deeplink.py b/deeplink_ext/common/rms_norm/deeplink.py
@@ -76,4 +76,3 @@ def rms_norm_backward(input, grad_output, inv_rms, normalized_shape, weight, bia
     )
 
     return [grad_input, grad_weight, grad_bias]
-
diff --git a/deeplink_ext/patch_lightllm.py b/deeplink_ext/patch_lightllm.py
@@ -52,6 +52,7 @@ def patch_token_softmax_reducev_inference():
 
         def patch_rms_norm_lightllm():
             from .common.rms_norm.deeplink import rms_norm
+
             rms_norm_pack.rmsnorm_forward = rms_norm
 
         def patch_rotary_emb():
Original file line number	Diff line number	Diff line change
Expand Up		@@ -76,4 +76,3 @@ def rms_norm_backward(input, grad_output, inv_rms, normalized_shape, weight, bia
		)

		return [grad_input, grad_weight, grad_bias]