liger-kernel-nightly 0.6.2.dev20251013144132__py3-none-any.whl → 0.6.2.dev20251014205028__py3-none-any.whl
This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
- liger_kernel/ops/poly_norm.py +386 -0
- liger_kernel/transformers/__init__.py +2 -0
- liger_kernel/transformers/functional.py +5 -0
- liger_kernel/transformers/poly_norm.py +42 -0
- {liger_kernel_nightly-0.6.2.dev20251013144132.dist-info → liger_kernel_nightly-0.6.2.dev20251014205028.dist-info}/METADATA +1 -1
- {liger_kernel_nightly-0.6.2.dev20251013144132.dist-info → liger_kernel_nightly-0.6.2.dev20251014205028.dist-info}/RECORD +10 -8
- {liger_kernel_nightly-0.6.2.dev20251013144132.dist-info → liger_kernel_nightly-0.6.2.dev20251014205028.dist-info}/LICENSE +0 -0
- {liger_kernel_nightly-0.6.2.dev20251013144132.dist-info → liger_kernel_nightly-0.6.2.dev20251014205028.dist-info}/NOTICE +0 -0
- {liger_kernel_nightly-0.6.2.dev20251013144132.dist-info → liger_kernel_nightly-0.6.2.dev20251014205028.dist-info}/WHEEL +0 -0
- {liger_kernel_nightly-0.6.2.dev20251013144132.dist-info → liger_kernel_nightly-0.6.2.dev20251014205028.dist-info}/top_level.txt +0 -0
@@ -0,0 +1,386 @@
|
|
1
|
+
import operator
|
2
|
+
|
3
|
+
import torch
|
4
|
+
import triton
|
5
|
+
import triton.language as tl
|
6
|
+
|
7
|
+
from liger_kernel.ops.utils import calculate_settings
|
8
|
+
from liger_kernel.ops.utils import compare_version
|
9
|
+
from liger_kernel.ops.utils import ensure_contiguous
|
10
|
+
|
11
|
+
if compare_version("triton", operator.ge, "3.0.0"):
|
12
|
+
try:
|
13
|
+
from triton.language.extra.libdevice import rsqrt
|
14
|
+
except ModuleNotFoundError:
|
15
|
+
from triton.language.extra.cuda.libdevice import rsqrt
|
16
|
+
else:
|
17
|
+
from triton.language.math import rsqrt
|
18
|
+
|
19
|
+
|
20
|
+
@triton.jit
|
21
|
+
def _poly_norm_forward_kernel(
|
22
|
+
Y_ptr,
|
23
|
+
Y_row_stride,
|
24
|
+
X_ptr,
|
25
|
+
X_row_stride,
|
26
|
+
W_ptr, # weight: [3] for [w0, w1, w2]
|
27
|
+
B_ptr, # bias: scalar
|
28
|
+
RSTD_ptr, # cache rstd for backward: shape (n_rows, 3)
|
29
|
+
RSTD_row_stride,
|
30
|
+
n_cols,
|
31
|
+
eps,
|
32
|
+
BLOCK_SIZE: tl.constexpr,
|
33
|
+
):
|
34
|
+
"""
|
35
|
+
PolyNorm formula:
|
36
|
+
y = w₀·norm(x³) + w₁·norm(x²) + w₂·norm(x) + b
|
37
|
+
where norm(u) = u / sqrt(mean(u²) + ε)
|
38
|
+
|
39
|
+
Reference:
|
40
|
+
1. https://github.com/BryceZhuo/PolyCom/
|
41
|
+
2. https://arxiv.org/pdf/2411.03884
|
42
|
+
|
43
|
+
Cache rstd values for backward pass
|
44
|
+
"""
|
45
|
+
row_idx = tl.program_id(0).to(tl.int64)
|
46
|
+
col_offsets = tl.arange(0, BLOCK_SIZE)
|
47
|
+
mask = col_offsets < n_cols
|
48
|
+
|
49
|
+
# Load pointers
|
50
|
+
Y_ptr += row_idx * Y_row_stride
|
51
|
+
X_ptr += row_idx * X_row_stride
|
52
|
+
RSTD_ptr += row_idx * RSTD_row_stride
|
53
|
+
|
54
|
+
# Load input row
|
55
|
+
X_row = tl.load(X_ptr + col_offsets, mask=mask, other=0.0)
|
56
|
+
|
57
|
+
# Load weights and bias
|
58
|
+
w0 = tl.load(W_ptr + 0)
|
59
|
+
w1 = tl.load(W_ptr + 1)
|
60
|
+
w2 = tl.load(W_ptr + 2)
|
61
|
+
b = tl.load(B_ptr)
|
62
|
+
|
63
|
+
# Compute x³, x², x
|
64
|
+
X_pow3 = X_row * X_row * X_row
|
65
|
+
X_pow2 = X_row * X_row
|
66
|
+
X_pow1 = X_row
|
67
|
+
|
68
|
+
# Compute norm(x³): norm(u) = u * rsqrt(mean(u²) + eps)
|
69
|
+
mean_square_3 = tl.sum(X_pow3 * X_pow3, axis=0) / n_cols
|
70
|
+
rstd_3 = rsqrt(mean_square_3 + eps)
|
71
|
+
norm_x3 = X_pow3 * rstd_3
|
72
|
+
|
73
|
+
# Compute norm(x²)
|
74
|
+
mean_square_2 = tl.sum(X_pow2 * X_pow2, axis=0) / n_cols
|
75
|
+
rstd_2 = rsqrt(mean_square_2 + eps)
|
76
|
+
norm_x2 = X_pow2 * rstd_2
|
77
|
+
|
78
|
+
# Compute norm(x)
|
79
|
+
mean_square_1 = tl.sum(X_pow1 * X_pow1, axis=0) / n_cols
|
80
|
+
rstd_1 = rsqrt(mean_square_1 + eps)
|
81
|
+
norm_x1 = X_pow1 * rstd_1
|
82
|
+
|
83
|
+
# Cache rstd values for backward
|
84
|
+
tl.store(RSTD_ptr + 0, rstd_3)
|
85
|
+
tl.store(RSTD_ptr + 1, rstd_2)
|
86
|
+
tl.store(RSTD_ptr + 2, rstd_1)
|
87
|
+
|
88
|
+
# Compute output: y = w₀·norm(x³) + w₁·norm(x²) + w₂·norm(x) + b
|
89
|
+
Y_row = w0 * norm_x3 + w1 * norm_x2 + w2 * norm_x1 + b
|
90
|
+
|
91
|
+
# Store output
|
92
|
+
tl.store(Y_ptr + col_offsets, Y_row, mask=mask)
|
93
|
+
|
94
|
+
|
95
|
+
@triton.jit
|
96
|
+
def _poly_norm_backward_kernel(
|
97
|
+
dY_ptr,
|
98
|
+
dY_row_stride,
|
99
|
+
dX_ptr,
|
100
|
+
dX_row_stride,
|
101
|
+
X_ptr,
|
102
|
+
X_row_stride,
|
103
|
+
W_ptr,
|
104
|
+
RSTD_ptr,
|
105
|
+
RSTD_row_stride,
|
106
|
+
dW_ptr, # shape: (n_programs, 3)
|
107
|
+
dW_row_stride,
|
108
|
+
dB_ptr, # shape: (n_programs,)
|
109
|
+
n_rows,
|
110
|
+
n_cols,
|
111
|
+
rows_per_program: tl.constexpr,
|
112
|
+
BLOCK_SIZE: tl.constexpr,
|
113
|
+
):
|
114
|
+
"""
|
115
|
+
PolyNorm Backward Kernel Gradient:
|
116
|
+
∂L/∂x_i = Σ_p w_p * [p*x_i^(p-1) * grad_i/D_p - (p/d)*x_i^(2p-1) * S_p/(D_p³)]
|
117
|
+
|
118
|
+
where:
|
119
|
+
- D_p = RMS(x^p) = 1/rstd_p
|
120
|
+
- S_p = sum(grad * x^p) over the row
|
121
|
+
- d = n_cols
|
122
|
+
- p ∈ {3, 2, 1}
|
123
|
+
"""
|
124
|
+
row_block_id = tl.program_id(0).to(tl.int64)
|
125
|
+
row_start = row_block_id * rows_per_program
|
126
|
+
row_end = min((row_block_id + 1) * rows_per_program, n_rows)
|
127
|
+
col_offsets = tl.arange(0, BLOCK_SIZE)
|
128
|
+
mask = col_offsets < n_cols
|
129
|
+
|
130
|
+
# Initialize accumulators for weight and bias gradients (scalars)
|
131
|
+
dW0_acc = 0.0
|
132
|
+
dW1_acc = 0.0
|
133
|
+
dW2_acc = 0.0
|
134
|
+
dB_acc = 0.0
|
135
|
+
|
136
|
+
# Load weights
|
137
|
+
w0 = tl.load(W_ptr + 0).to(tl.float32)
|
138
|
+
w1 = tl.load(W_ptr + 1).to(tl.float32)
|
139
|
+
w2 = tl.load(W_ptr + 2).to(tl.float32)
|
140
|
+
|
141
|
+
dY_ptr += row_start * dY_row_stride
|
142
|
+
dX_ptr += row_start * dX_row_stride
|
143
|
+
X_ptr += row_start * X_row_stride
|
144
|
+
RSTD_ptr += row_start * RSTD_row_stride
|
145
|
+
|
146
|
+
for _ in range(row_start, row_end):
|
147
|
+
# Load input and gradient
|
148
|
+
dY_row = tl.load(dY_ptr + col_offsets, mask=mask, other=0.0).to(tl.float32)
|
149
|
+
X_row = tl.load(X_ptr + col_offsets, mask=mask, other=0.0).to(tl.float32)
|
150
|
+
|
151
|
+
# Load cached rstd values
|
152
|
+
rstd_3 = tl.load(RSTD_ptr + 0).to(tl.float32)
|
153
|
+
rstd_2 = tl.load(RSTD_ptr + 1).to(tl.float32)
|
154
|
+
rstd_1 = tl.load(RSTD_ptr + 2).to(tl.float32)
|
155
|
+
|
156
|
+
# Compute powers
|
157
|
+
X_pow3 = X_row * X_row * X_row
|
158
|
+
X_pow2 = X_row * X_row
|
159
|
+
X_pow1 = X_row
|
160
|
+
|
161
|
+
# Accumulate bias gradient: dB = sum(dY)
|
162
|
+
dB_acc += tl.sum(dY_row, axis=0)
|
163
|
+
|
164
|
+
# Compute gradient w.r.t. input using closed-form formula
|
165
|
+
# For p=3: ∂L/∂x from w0 * norm(x³)
|
166
|
+
S_3 = tl.sum(dY_row * X_pow3, axis=0) # scalar
|
167
|
+
grad_x_3 = w0 * (
|
168
|
+
3.0 * X_pow2 * rstd_3 * dY_row
|
169
|
+
- (3.0 / n_cols) * X_row * X_row * X_row * X_row * X_row * (rstd_3 * rstd_3 * rstd_3) * S_3
|
170
|
+
)
|
171
|
+
|
172
|
+
# For p=2: ∂L/∂x from w1 * norm(x²)
|
173
|
+
S_2 = tl.sum(dY_row * X_pow2, axis=0) # scalar
|
174
|
+
grad_x_2 = w1 * (
|
175
|
+
2.0 * X_row * rstd_2 * dY_row - (2.0 / n_cols) * X_row * X_row * X_row * (rstd_2 * rstd_2 * rstd_2) * S_2
|
176
|
+
)
|
177
|
+
|
178
|
+
# For p=1: ∂L/∂x from w2 * norm(x)
|
179
|
+
S_1 = tl.sum(dY_row * X_pow1, axis=0) # scalar
|
180
|
+
grad_x_1 = w2 * (1.0 * rstd_1 * dY_row - (1.0 / n_cols) * X_row * (rstd_1 * rstd_1 * rstd_1) * S_1)
|
181
|
+
|
182
|
+
# Accumulate weight gradients using closed-form: dW_p = rstd_p * S_p
|
183
|
+
dW0_acc += rstd_3 * S_3
|
184
|
+
dW1_acc += rstd_2 * S_2
|
185
|
+
dW2_acc += rstd_1 * S_1
|
186
|
+
|
187
|
+
# Total gradient
|
188
|
+
dX_row = grad_x_3 + grad_x_2 + grad_x_1
|
189
|
+
|
190
|
+
# Store gradient
|
191
|
+
tl.store(dX_ptr + col_offsets, dX_row, mask=mask)
|
192
|
+
|
193
|
+
# Update pointers
|
194
|
+
dY_ptr += dY_row_stride
|
195
|
+
dX_ptr += dX_row_stride
|
196
|
+
X_ptr += X_row_stride
|
197
|
+
RSTD_ptr += RSTD_row_stride
|
198
|
+
|
199
|
+
# Store accumulated gradients (scalars)
|
200
|
+
tl.store(dW_ptr + row_block_id * dW_row_stride + 0, dW0_acc)
|
201
|
+
tl.store(dW_ptr + row_block_id * dW_row_stride + 1, dW1_acc)
|
202
|
+
tl.store(dW_ptr + row_block_id * dW_row_stride + 2, dW2_acc)
|
203
|
+
tl.store(dB_ptr + row_block_id, dB_acc)
|
204
|
+
|
205
|
+
|
206
|
+
def poly_norm_forward(X, W, B, eps=1e-6):
|
207
|
+
"""
|
208
|
+
PolyNorm Forward Pass
|
209
|
+
|
210
|
+
Args:
|
211
|
+
X: input tensor of shape (*, H) where H is hidden dimension
|
212
|
+
W: weight tensor of shape (3,) for [w0, w1, w2]
|
213
|
+
B: bias scalar tensor
|
214
|
+
eps: epsilon for numerical stability
|
215
|
+
|
216
|
+
Returns:
|
217
|
+
Y: output tensor of same shape as X
|
218
|
+
X: reshaped input (for backward)
|
219
|
+
RSTD: cached rstd values (for backward)
|
220
|
+
BLOCK_SIZE: block size used
|
221
|
+
num_warps: number of warps used
|
222
|
+
"""
|
223
|
+
shape = X.shape
|
224
|
+
dim = shape[-1]
|
225
|
+
X = X.view(-1, dim)
|
226
|
+
n_rows, n_cols = X.shape
|
227
|
+
BLOCK_SIZE, num_warps = calculate_settings(n_cols)
|
228
|
+
|
229
|
+
# RSTD is to cache rstd for each row
|
230
|
+
Y = torch.empty((n_rows, n_cols), dtype=X.dtype, device=X.device)
|
231
|
+
RSTD = torch.empty((n_rows, 3), dtype=torch.float32, device=X.device)
|
232
|
+
|
233
|
+
# Check constraints
|
234
|
+
assert W.shape[0] == 3, "Weight tensor must have shape (3,)"
|
235
|
+
assert B.numel() == 1, "Bias must be a scalar"
|
236
|
+
|
237
|
+
# XPU-specific optimization
|
238
|
+
kernel_args = {}
|
239
|
+
if X.device.type == "xpu":
|
240
|
+
kernel_args["grf_mode"] = "large"
|
241
|
+
|
242
|
+
# Launch kernel
|
243
|
+
_poly_norm_forward_kernel[(n_rows,)](
|
244
|
+
Y,
|
245
|
+
Y.stride(0),
|
246
|
+
X,
|
247
|
+
X.stride(0),
|
248
|
+
W,
|
249
|
+
B,
|
250
|
+
RSTD,
|
251
|
+
RSTD.stride(0),
|
252
|
+
n_cols,
|
253
|
+
eps,
|
254
|
+
BLOCK_SIZE=BLOCK_SIZE,
|
255
|
+
num_warps=num_warps,
|
256
|
+
**kernel_args,
|
257
|
+
)
|
258
|
+
|
259
|
+
return Y.view(*shape), X, RSTD, BLOCK_SIZE, num_warps
|
260
|
+
|
261
|
+
|
262
|
+
def poly_norm_backward(dY, X, W, RSTD, BLOCK_SIZE, num_warps, in_place):
|
263
|
+
"""
|
264
|
+
PolyNorm Backward Pass
|
265
|
+
|
266
|
+
Args:
|
267
|
+
dY: gradient of output
|
268
|
+
X: input tensor (already reshaped to 2D)
|
269
|
+
W: weight tensor
|
270
|
+
RSTD: cached rstd values from forward
|
271
|
+
BLOCK_SIZE: block size from forward
|
272
|
+
num_warps: number of warps from forward
|
273
|
+
in_place: whether to in-place modify dY to store dX (saves memory)
|
274
|
+
|
275
|
+
Returns:
|
276
|
+
dX: gradient w.r.t. input
|
277
|
+
dW: gradient w.r.t. weight
|
278
|
+
dB: gradient w.r.t. bias
|
279
|
+
"""
|
280
|
+
shape = dY.shape
|
281
|
+
dim = shape[-1]
|
282
|
+
dY = dY.view(-1, dim)
|
283
|
+
n_rows, n_cols = dY.shape
|
284
|
+
|
285
|
+
# Get number of SMs for parallelization
|
286
|
+
import math
|
287
|
+
|
288
|
+
sm_count = 1
|
289
|
+
if X.device.type == "cuda":
|
290
|
+
sm_count = torch.cuda.get_device_properties(X.device).multi_processor_count
|
291
|
+
elif X.device.type == "xpu":
|
292
|
+
sm_count = torch.xpu.get_device_properties(X.device).gpu_eu_count
|
293
|
+
|
294
|
+
# Allocate or reuse gradients
|
295
|
+
if in_place is True:
|
296
|
+
dX = dY
|
297
|
+
else:
|
298
|
+
dX = torch.zeros_like(dY)
|
299
|
+
|
300
|
+
_dW = torch.empty((sm_count, 3), dtype=torch.float32, device=W.device)
|
301
|
+
_dB = torch.empty((sm_count,), dtype=torch.float32, device=W.device)
|
302
|
+
|
303
|
+
rows_per_program = math.ceil(n_rows / sm_count)
|
304
|
+
grid = (sm_count,)
|
305
|
+
|
306
|
+
# XPU-specific optimization
|
307
|
+
kernel_args = {}
|
308
|
+
if X.device.type == "xpu":
|
309
|
+
kernel_args["grf_mode"] = "large"
|
310
|
+
|
311
|
+
# Launch backward kernel
|
312
|
+
_poly_norm_backward_kernel[grid](
|
313
|
+
dY,
|
314
|
+
dY.stride(0),
|
315
|
+
dX,
|
316
|
+
dX.stride(0),
|
317
|
+
X,
|
318
|
+
X.stride(0),
|
319
|
+
W,
|
320
|
+
RSTD,
|
321
|
+
RSTD.stride(0),
|
322
|
+
_dW,
|
323
|
+
_dW.stride(0),
|
324
|
+
_dB,
|
325
|
+
n_rows,
|
326
|
+
n_cols,
|
327
|
+
rows_per_program,
|
328
|
+
BLOCK_SIZE=BLOCK_SIZE,
|
329
|
+
num_warps=num_warps,
|
330
|
+
**kernel_args,
|
331
|
+
)
|
332
|
+
|
333
|
+
# Reduce gradients across SMs
|
334
|
+
dX = dX.view(*shape)
|
335
|
+
dW = _dW.sum(dim=0).to(W.dtype)
|
336
|
+
dB = _dB.sum().to(W.dtype)
|
337
|
+
|
338
|
+
return dX, dW, dB
|
339
|
+
|
340
|
+
|
341
|
+
class LigerPolyNormFunction(torch.autograd.Function):
|
342
|
+
"""
|
343
|
+
PolyNorm Function with forward and backward pass
|
344
|
+
|
345
|
+
PolyNorm formula:
|
346
|
+
y = w₀·norm(x³) + w₁·norm(x²) + w₂·norm(x) + b
|
347
|
+
where norm(u) = u / sqrt(mean(u²) + ε)
|
348
|
+
|
349
|
+
Backward uses closed-form gradient:
|
350
|
+
∂L/∂x_i = Σ_p w_p * [p*x_i^(p-1) * grad_i/D_p - (p/d)*x_i^(2p-1) * S_p/(D_p³)]
|
351
|
+
"""
|
352
|
+
|
353
|
+
@staticmethod
|
354
|
+
@ensure_contiguous
|
355
|
+
def forward(ctx, X, W, B, eps=1e-6, in_place=True):
|
356
|
+
"""
|
357
|
+
Args:
|
358
|
+
X: input tensor of shape (B, T, H) or (BxT, H)
|
359
|
+
W: weight tensor of shape (3,) for [w0, w1, w2]
|
360
|
+
B: bias scalar
|
361
|
+
eps: epsilon for numerical stability
|
362
|
+
in_place: whether to in-place modify grad_output in backward (saves memory)
|
363
|
+
|
364
|
+
Returns:
|
365
|
+
Y: output tensor of same shape as X
|
366
|
+
"""
|
367
|
+
Y, X, RSTD, BLOCK_SIZE, num_warps = poly_norm_forward(X, W, B, eps)
|
368
|
+
ctx.BLOCK_SIZE = BLOCK_SIZE
|
369
|
+
ctx.num_warps = num_warps
|
370
|
+
ctx.in_place = in_place
|
371
|
+
ctx.save_for_backward(X, W, RSTD)
|
372
|
+
return Y
|
373
|
+
|
374
|
+
@staticmethod
|
375
|
+
@ensure_contiguous
|
376
|
+
def backward(ctx, grad_output):
|
377
|
+
"""
|
378
|
+
Args:
|
379
|
+
grad_output: gradient of output
|
380
|
+
|
381
|
+
Returns:
|
382
|
+
dX, dW, dB: gradients w.r.t. X, W, B
|
383
|
+
"""
|
384
|
+
X, W, RSTD = ctx.saved_tensors
|
385
|
+
dX, dW, dB = poly_norm_backward(grad_output, X, W, RSTD, ctx.BLOCK_SIZE, ctx.num_warps, ctx.in_place)
|
386
|
+
return dX, dW, dB, None, None
|
@@ -15,6 +15,7 @@ from liger_kernel.transformers.layer_norm import LigerLayerNorm # noqa: F401
|
|
15
15
|
from liger_kernel.transformers.llama4_rope import liger_llama4_text_rotary_pos_emb # noqa: F401
|
16
16
|
from liger_kernel.transformers.llama4_rope import liger_llama4_vision_rotary_pos_emb # noqa: F401
|
17
17
|
from liger_kernel.transformers.multi_token_attention import LigerMultiTokenAttention # noqa: F401
|
18
|
+
from liger_kernel.transformers.poly_norm import LigerPolyNorm # noqa: F401
|
18
19
|
from liger_kernel.transformers.rms_norm import LigerRMSNorm # noqa: F401
|
19
20
|
from liger_kernel.transformers.rope import liger_rotary_pos_emb # noqa: F401
|
20
21
|
from liger_kernel.transformers.softmax import LigerSoftmax # noqa: F401
|
@@ -137,6 +138,7 @@ __all__ = [
|
|
137
138
|
"LigerJSD",
|
138
139
|
"LigerLayerNorm",
|
139
140
|
"LigerFusedAddRMSNorm",
|
141
|
+
"LigerPolyNorm",
|
140
142
|
"LigerRMSNorm",
|
141
143
|
"liger_rotary_pos_emb",
|
142
144
|
"liger_llama4_text_rotary_pos_emb",
|
@@ -12,6 +12,7 @@ from liger_kernel.ops.jsd import LigerJSDFunction
|
|
12
12
|
from liger_kernel.ops.kl_div import LigerKLDivLossFunction
|
13
13
|
from liger_kernel.ops.layer_norm import LigerLayerNormFunction
|
14
14
|
from liger_kernel.ops.multi_token_attention import LigerMultiTokenAttentionFunction
|
15
|
+
from liger_kernel.ops.poly_norm import LigerPolyNormFunction
|
15
16
|
from liger_kernel.ops.qwen2vl_mrope import LigerQwen2VLMRopeFunction
|
16
17
|
from liger_kernel.ops.rms_norm import LigerRMSNormFunction
|
17
18
|
from liger_kernel.ops.rope import LigerRopeFunction
|
@@ -258,6 +259,10 @@ def liger_rms_norm(X, W, eps, offset: float = 0.0, casting_mode: str = "llama",
|
|
258
259
|
return LigerRMSNormFunction.apply(X, W, eps, offset, casting_mode, in_place)
|
259
260
|
|
260
261
|
|
262
|
+
def liger_poly_norm(X, W, B, eps=1e-6, in_place=True):
|
263
|
+
return LigerPolyNormFunction.apply(X, W, B, eps, in_place)
|
264
|
+
|
265
|
+
|
261
266
|
def liger_fused_add_rms_norm(X, R, W, eps, offset: float = 0.0, casting_mode: str = "llama", in_place: bool = True):
|
262
267
|
return LigerFusedAddRMSNormFunction.apply(X, R, W, eps, offset, casting_mode, in_place)
|
263
268
|
|
@@ -0,0 +1,42 @@
|
|
1
|
+
import torch
|
2
|
+
import torch.nn as nn
|
3
|
+
|
4
|
+
from liger_kernel.ops.poly_norm import LigerPolyNormFunction
|
5
|
+
|
6
|
+
|
7
|
+
class LigerPolyNorm(nn.Module):
|
8
|
+
"""
|
9
|
+
PolyNorm layer wrapper for Liger kernel.
|
10
|
+
|
11
|
+
PolyNorm formula:
|
12
|
+
y = w₀·norm(x³) + w₁·norm(x²) + w₂·norm(x) + b
|
13
|
+
where norm(u) = u / sqrt(mean(u²) + ε)
|
14
|
+
|
15
|
+
Reference:
|
16
|
+
https://github.com/BryceZhuo/PolyCom/
|
17
|
+
|
18
|
+
Args:
|
19
|
+
eps: epsilon for numerical stability (default: 1e-6)
|
20
|
+
in_place: whether to in-place modify grad_output in backward to save memory (default: False).
|
21
|
+
Set to True to save memory if grad_output is not needed elsewhere.
|
22
|
+
"""
|
23
|
+
|
24
|
+
def __init__(self, eps=1e-6, in_place=True):
|
25
|
+
super().__init__()
|
26
|
+
# Align with PolyCom reference: initialize weights to (1/3, 1/3, 1/3) and bias to 1.0
|
27
|
+
self.weight = nn.Parameter(torch.full((3,), 1.0 / 3.0))
|
28
|
+
self.bias = nn.Parameter(torch.tensor(1.0))
|
29
|
+
self.variance_epsilon = eps
|
30
|
+
self.in_place = in_place
|
31
|
+
|
32
|
+
def forward(self, hidden_states):
|
33
|
+
return LigerPolyNormFunction.apply(
|
34
|
+
hidden_states,
|
35
|
+
self.weight,
|
36
|
+
self.bias,
|
37
|
+
self.variance_epsilon,
|
38
|
+
self.in_place,
|
39
|
+
)
|
40
|
+
|
41
|
+
def extra_repr(self):
|
42
|
+
return f"weight_shape={tuple(self.weight.shape)}, eps={self.variance_epsilon}, in_place={self.in_place}"
|
@@ -31,6 +31,7 @@ liger_kernel/ops/kl_div.py,sha256=ZjGdDLKWksHT9dZ0xF_TDgAkj5cuMTwwT5tr9E-_24o,87
|
|
31
31
|
liger_kernel/ops/layer_norm.py,sha256=WmiORsIyufOhazmYZTPjeSc5Z-xTAYwXAKqUcCv_dlY,9807
|
32
32
|
liger_kernel/ops/llama4_rope.py,sha256=-aqdZzllklTN8b9--e-TsWY_ntGCN8-tyseT4x0bd8s,8223
|
33
33
|
liger_kernel/ops/multi_token_attention.py,sha256=Oz_RXDp-OSS_R_HuGmaETHdAJ7Toda_70OfE7TXMUlY,7645
|
34
|
+
liger_kernel/ops/poly_norm.py,sha256=MLgI8Ea93fugKibHCUauQ2ASYVXCvpPZe5v3kQZU6po,11152
|
34
35
|
liger_kernel/ops/qwen2vl_mrope.py,sha256=3GExhYpLgB4VUtyZyjRk8XjEur3W4EWF6HQ67ML5vBU,8481
|
35
36
|
liger_kernel/ops/rms_norm.py,sha256=DtvsWN5YktFAoc0JYSAwVeoZfryBFJlX-ipU7ooP01A,18891
|
36
37
|
liger_kernel/ops/rope.py,sha256=v-7JHRrv-5ImoROkpKfl30WwWI4qTa2tAl7zQeB4ml4,8956
|
@@ -41,12 +42,12 @@ liger_kernel/ops/tvd.py,sha256=FHJtLQI95ijqgg9UtaHpMAjSCiPxB6CduPwPMcGxelc,6405
|
|
41
42
|
liger_kernel/ops/utils.py,sha256=uoFKQqo-34N2TWQNvXMFywqGiOMMXNEVBxVojzlUAa0,3836
|
42
43
|
liger_kernel/ops/experimental/embedding.py,sha256=tolj3tItkzpSb30zWqDN2_yX4ectflaQ8HMyKyFIQc8,4172
|
43
44
|
liger_kernel/ops/experimental/mm_int8int2.py,sha256=TrS9lpwekrik_w5qE7AhMJD1bcq-OidjtbsW80oZ6IM,13314
|
44
|
-
liger_kernel/transformers/__init__.py,sha256=
|
45
|
+
liger_kernel/transformers/__init__.py,sha256=d0H4knUp93iR3OPR3lpYriZYCvC-w_i2cDTYtcgfhzo,9107
|
45
46
|
liger_kernel/transformers/auto_model.py,sha256=0qCTRZt280Bj_LcFdzo9hlaR-BWNazawXOGgoCZjgEg,1545
|
46
47
|
liger_kernel/transformers/cross_entropy.py,sha256=z3KTWQnFxr_IZaVjtYt0ZNEWQdDdYThN35xWkHlDGH0,1683
|
47
48
|
liger_kernel/transformers/dyt.py,sha256=i-4GPaMrl-jab9TVI5qN0-H9qycn_mCbV82ozU4nbmU,723
|
48
49
|
liger_kernel/transformers/fsdp.py,sha256=CUiyjTmjkjY7pLXQv8ly9rnzgXw6529csd9pvtJNMYc,3096
|
49
|
-
liger_kernel/transformers/functional.py,sha256
|
50
|
+
liger_kernel/transformers/functional.py,sha256=a8EGYjHDg34rhnaD4JpU8I20XJ7xiqJvqqjoh4NcwYk,8022
|
50
51
|
liger_kernel/transformers/fused_add_rms_norm.py,sha256=7_Bzg-x6lLe6W1qG2DtjDALhEpNZlC6N5GppEs9cTYY,1199
|
51
52
|
liger_kernel/transformers/fused_linear_cross_entropy.py,sha256=toa54dpmJduoZLhU3lJA-HPZ03MYcMKekDWPcdYjvYA,2020
|
52
53
|
liger_kernel/transformers/fused_linear_jsd.py,sha256=bZ4otCvWBuOnA5XdQL-FzZVItJlDt-ht9e_pG7PG93E,3999
|
@@ -60,6 +61,7 @@ liger_kernel/transformers/layer_norm.py,sha256=c9pk3PEasOKYR0rhe5e5nNrnYKVCEW4VC
|
|
60
61
|
liger_kernel/transformers/llama4_rope.py,sha256=kS6PSHEwf3dS7hD7C7p8S0geugx2EMCiP0h0F7LsUoY,3639
|
61
62
|
liger_kernel/transformers/monkey_patch.py,sha256=TUmx8aY0lonyThcATirRBdSs7uItVvnBggohjBItBuQ,106060
|
62
63
|
liger_kernel/transformers/multi_token_attention.py,sha256=K3NIY9_5TPgZ4_Rahn0xnkMXxD_fmlJHK4CWGYvGQp0,1752
|
64
|
+
liger_kernel/transformers/poly_norm.py,sha256=g5tC75i3qy1_N26ZUP-jfpct7ivQAEdJfIfx8IXzeyE,1377
|
63
65
|
liger_kernel/transformers/qwen2vl_mrope.py,sha256=5EwSqrMdsL9MYspeBMXBsNJKvH0MOmRrtJXAJlnnlOI,1047
|
64
66
|
liger_kernel/transformers/rms_norm.py,sha256=vkekcvTeWY8vL4H6hg3t0XeY0Ew_3OFMPHuzqlxPPVw,2719
|
65
67
|
liger_kernel/transformers/rope.py,sha256=ZTrTORSAyfcFIKjk6XEeYmk4ROH7xXED9L4g2NFntlE,999
|
@@ -99,9 +101,9 @@ liger_kernel/transformers/trainer/__init__.py,sha256=p7yQfklV8-467qSz_ZMimkbDF7H
|
|
99
101
|
liger_kernel/transformers/trainer/orpo_trainer.py,sha256=tX0h63aOFe3rNqTmk6JpMf75UPo981yzEa6TghnjS0Q,5370
|
100
102
|
liger_kernel/triton/__init__.py,sha256=qCiCamzCRv6lpV8IqpAc9YMdNKC7GKurClWceQPnlis,92
|
101
103
|
liger_kernel/triton/monkey_patch.py,sha256=Rd0hUHAzDkFfHvnX7-PBaNK5EKnZhtfM_h-fgQH9HPY,1568
|
102
|
-
liger_kernel_nightly-0.6.2.
|
103
|
-
liger_kernel_nightly-0.6.2.
|
104
|
-
liger_kernel_nightly-0.6.2.
|
105
|
-
liger_kernel_nightly-0.6.2.
|
106
|
-
liger_kernel_nightly-0.6.2.
|
107
|
-
liger_kernel_nightly-0.6.2.
|
104
|
+
liger_kernel_nightly-0.6.2.dev20251014205028.dist-info/LICENSE,sha256=OhzLDHJ0to4a8sodVLELZiCFylZ1NAAYLs-HrjPy0ag,1312
|
105
|
+
liger_kernel_nightly-0.6.2.dev20251014205028.dist-info/METADATA,sha256=6VDasn5yo1wPa73CAIS4iRzr6TJ_cWpSjF_QbD5r1sM,24777
|
106
|
+
liger_kernel_nightly-0.6.2.dev20251014205028.dist-info/NOTICE,sha256=njwnoPZLh9AN8SJQzxvCGLHi-8X__AvWRze6joNXIY8,2066
|
107
|
+
liger_kernel_nightly-0.6.2.dev20251014205028.dist-info/WHEEL,sha256=iAkIy5fosb7FzIOwONchHf19Qu7_1wCWyFNR5gu9nU0,91
|
108
|
+
liger_kernel_nightly-0.6.2.dev20251014205028.dist-info/top_level.txt,sha256=2eghu4hA3LnkM7ElW92tQ8zegWKgSbeo-k-aGe1YnvY,13
|
109
|
+
liger_kernel_nightly-0.6.2.dev20251014205028.dist-info/RECORD,,
|
File without changes
|
File without changes
|
File without changes
|