liger-kernel-nightly 0.6.0.dev20250718050840__py3-none-any.whl → 0.6.0.dev20250718080702__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/fused_add_rms_norm.py +412 -0
- liger_kernel/transformers/__init__.py +2 -0
- liger_kernel/transformers/functional.py +5 -0
- liger_kernel/transformers/fused_add_rms_norm.py +39 -0
- {liger_kernel_nightly-0.6.0.dev20250718050840.dist-info → liger_kernel_nightly-0.6.0.dev20250718080702.dist-info}/METADATA +1 -1
- {liger_kernel_nightly-0.6.0.dev20250718050840.dist-info → liger_kernel_nightly-0.6.0.dev20250718080702.dist-info}/RECORD +10 -8
- {liger_kernel_nightly-0.6.0.dev20250718050840.dist-info → liger_kernel_nightly-0.6.0.dev20250718080702.dist-info}/LICENSE +0 -0
- {liger_kernel_nightly-0.6.0.dev20250718050840.dist-info → liger_kernel_nightly-0.6.0.dev20250718080702.dist-info}/NOTICE +0 -0
- {liger_kernel_nightly-0.6.0.dev20250718050840.dist-info → liger_kernel_nightly-0.6.0.dev20250718080702.dist-info}/WHEEL +0 -0
- {liger_kernel_nightly-0.6.0.dev20250718050840.dist-info → liger_kernel_nightly-0.6.0.dev20250718080702.dist-info}/top_level.txt +0 -0
@@ -0,0 +1,412 @@
|
|
1
|
+
import math
|
2
|
+
import operator
|
3
|
+
|
4
|
+
import torch
|
5
|
+
import triton
|
6
|
+
import triton.language as tl
|
7
|
+
|
8
|
+
from liger_kernel.ops.utils import calculate_settings
|
9
|
+
from liger_kernel.ops.utils import compare_version
|
10
|
+
from liger_kernel.ops.utils import ensure_contiguous
|
11
|
+
from liger_kernel.ops.utils import torch_to_triton_dtype
|
12
|
+
|
13
|
+
if compare_version("triton", operator.ge, "3.0.0"):
|
14
|
+
try:
|
15
|
+
# typical import path with dispatch available
|
16
|
+
from triton.language.extra.libdevice import rsqrt
|
17
|
+
except ModuleNotFoundError:
|
18
|
+
# for working with NGC containers
|
19
|
+
from triton.language.extra.cuda.libdevice import rsqrt
|
20
|
+
else:
|
21
|
+
from triton.language.math import rsqrt
|
22
|
+
|
23
|
+
|
24
|
+
_CASTING_MODE_NONE: tl.constexpr = tl.constexpr(-1)
|
25
|
+
_CASTING_MODE_LLAMA: tl.constexpr = tl.constexpr(0)
|
26
|
+
_CASTING_MODE_GEMMA: tl.constexpr = tl.constexpr(1)
|
27
|
+
|
28
|
+
|
29
|
+
@triton.jit
|
30
|
+
def _fused_add_rms_norm_forward_kernel(
|
31
|
+
Y_ptr,
|
32
|
+
Y_row_stride,
|
33
|
+
S_ptr, # output residual
|
34
|
+
S_row_stride,
|
35
|
+
X_ptr,
|
36
|
+
X_row_stride,
|
37
|
+
R_ptr, # input residual
|
38
|
+
R_row_stride,
|
39
|
+
W_ptr,
|
40
|
+
W_row_stride,
|
41
|
+
RSTD_ptr,
|
42
|
+
RSTD_row_stride,
|
43
|
+
n_cols,
|
44
|
+
eps,
|
45
|
+
offset,
|
46
|
+
casting_mode: tl.constexpr, # constexpr so the `if` blocks can be optimized out
|
47
|
+
BLOCK_SIZE: tl.constexpr,
|
48
|
+
):
|
49
|
+
"""
|
50
|
+
This kernel computes the following:
|
51
|
+
1. hidden_states = residual + hidden_states
|
52
|
+
2. residual = hidden_states
|
53
|
+
3. hidden_states = rmsnorm(hidden_states)
|
54
|
+
|
55
|
+
This is a commonly used pattern in the decoder layers of LLMs.
|
56
|
+
Some examples:
|
57
|
+
1. https://github.com/huggingface/transformers/blob/0dc2df5ddafe3cb5824ad24e85beba13e0aa6726/src/transformers/models/qwen3/modeling_qwen3.py#L271
|
58
|
+
2. https://github.com/huggingface/transformers/blob/0dc2df5ddafe3cb5824ad24e85beba13e0aa6726/src/transformers/models/llama4/modeling_llama4.py#L393
|
59
|
+
|
60
|
+
This kernel is inspired by the rms_norm forward kernel, and is adapted to support the residual addition in the forward pass.
|
61
|
+
The backward pass is also adapted to support the residual addition in the backward pass.
|
62
|
+
"""
|
63
|
+
|
64
|
+
row_idx = tl.program_id(0).to(tl.int64)
|
65
|
+
col_offsets = tl.arange(0, BLOCK_SIZE)
|
66
|
+
mask = col_offsets < n_cols
|
67
|
+
|
68
|
+
Y_ptr += row_idx * Y_row_stride
|
69
|
+
S_ptr += row_idx * S_row_stride
|
70
|
+
X_ptr += row_idx * X_row_stride
|
71
|
+
R_ptr += row_idx * R_row_stride
|
72
|
+
RSTD_ptr += row_idx * RSTD_row_stride
|
73
|
+
|
74
|
+
X_row = tl.load(X_ptr + col_offsets, mask=mask, other=0)
|
75
|
+
R_row = tl.load(R_ptr + col_offsets, mask=mask, other=0)
|
76
|
+
S_row = X_row + R_row
|
77
|
+
tl.store(S_ptr + col_offsets, S_row, mask=mask)
|
78
|
+
S_row_dtype = S_row.dtype
|
79
|
+
W_row = tl.load(W_ptr + col_offsets, mask=mask, other=0)
|
80
|
+
|
81
|
+
# On Llama, only rstd is computed on fp32
|
82
|
+
if casting_mode == _CASTING_MODE_LLAMA:
|
83
|
+
S_row = S_row.to(tl.float32)
|
84
|
+
|
85
|
+
# Gemma computes everything on fp32, and then casts back the output to the original dtype
|
86
|
+
if casting_mode == _CASTING_MODE_GEMMA:
|
87
|
+
W_row = W_row.to(tl.float32)
|
88
|
+
S_row = S_row.to(tl.float32)
|
89
|
+
|
90
|
+
if casting_mode == _CASTING_MODE_NONE:
|
91
|
+
eps = eps.to(S_row_dtype)
|
92
|
+
offset = offset.to(S_row_dtype)
|
93
|
+
|
94
|
+
mean_square = tl.sum(S_row * S_row, axis=0) / n_cols
|
95
|
+
rstd = rsqrt(mean_square + eps)
|
96
|
+
|
97
|
+
# We can save time by caching rms with minimal memory overhead
|
98
|
+
# because rms is much smaller compared to X_row, as rms is for each row.
|
99
|
+
# However, on the computation side, it can save 4 operations (*, sum, /, sqrt).
|
100
|
+
tl.store(RSTD_ptr, rstd)
|
101
|
+
|
102
|
+
S_row = S_row * rstd
|
103
|
+
|
104
|
+
# On Llama, the multiplication with the weight is done on the original dtype
|
105
|
+
if casting_mode == _CASTING_MODE_LLAMA:
|
106
|
+
S_row = S_row.to(S_row_dtype)
|
107
|
+
|
108
|
+
Y_row = S_row * (offset + W_row)
|
109
|
+
|
110
|
+
if casting_mode == _CASTING_MODE_GEMMA:
|
111
|
+
Y_row = Y_row.to(S_row_dtype)
|
112
|
+
|
113
|
+
tl.store(Y_ptr + col_offsets, Y_row, mask=mask)
|
114
|
+
|
115
|
+
|
116
|
+
@triton.jit
|
117
|
+
def _fused_add_rms_norm_backward_kernel(
|
118
|
+
dY_ptr,
|
119
|
+
dY_row_stride,
|
120
|
+
dS_out_ptr,
|
121
|
+
dS_out_row_stride,
|
122
|
+
dX_ptr,
|
123
|
+
dX_row_stride,
|
124
|
+
X_ptr,
|
125
|
+
X_row_stride,
|
126
|
+
X_dtype: tl.constexpr,
|
127
|
+
W_ptr,
|
128
|
+
W_row_stride,
|
129
|
+
RSTD_ptr,
|
130
|
+
RSTD_row_stride,
|
131
|
+
dW_ptr,
|
132
|
+
dW_row_stride,
|
133
|
+
n_rows,
|
134
|
+
n_cols,
|
135
|
+
offset,
|
136
|
+
rows_per_program: tl.constexpr,
|
137
|
+
casting_mode: tl.constexpr,
|
138
|
+
BLOCK_SIZE: tl.constexpr,
|
139
|
+
has_dS_out: tl.constexpr,
|
140
|
+
):
|
141
|
+
"""
|
142
|
+
This kernel is adapted from the rms_norm backward kernel, and is adapted to support the residual
|
143
|
+
addition in the backward pass. For the following code pattern:
|
144
|
+
1. hidden_states = residual + hidden_states
|
145
|
+
2. residual = hidden_states
|
146
|
+
3. hidden_states = rmsnorm(hidden_states)
|
147
|
+
|
148
|
+
The gradient of hidden_states and residual comes out be exactly same. The value of this gradient is
|
149
|
+
the sum of the gradient of the hidden_states in step 3 and the gradient of the residual in step 2.
|
150
|
+
|
151
|
+
The backward pass computation logic is same as the rms_norm backward kernel, except that the gradient
|
152
|
+
of the hidden_states in step 3 and the gradient of the residual in step 2 are summed up.
|
153
|
+
"""
|
154
|
+
|
155
|
+
row_block_id = tl.program_id(0).to(tl.int64)
|
156
|
+
row_start = row_block_id * rows_per_program
|
157
|
+
row_end = min((row_block_id + 1) * rows_per_program, n_rows)
|
158
|
+
col_offsets = tl.arange(0, BLOCK_SIZE)
|
159
|
+
mask = col_offsets < n_cols
|
160
|
+
|
161
|
+
dW_row = tl.zeros((BLOCK_SIZE,), dtype=tl.float32)
|
162
|
+
|
163
|
+
dY_ptr += row_start * dY_row_stride
|
164
|
+
dX_ptr += row_start * dX_row_stride
|
165
|
+
if has_dS_out:
|
166
|
+
dS_out_ptr += row_start * dS_out_row_stride
|
167
|
+
|
168
|
+
X_ptr += row_start * X_row_stride
|
169
|
+
RSTD_ptr += row_start
|
170
|
+
|
171
|
+
W_row = tl.load(W_ptr + col_offsets, mask=mask, other=0.0)
|
172
|
+
W_row = W_row + offset
|
173
|
+
|
174
|
+
for _ in range(row_start, row_end):
|
175
|
+
dY_row = tl.load(dY_ptr + col_offsets, mask=mask, other=0.0)
|
176
|
+
X_row = tl.load(X_ptr + col_offsets, mask=mask, other=0.0)
|
177
|
+
|
178
|
+
# Get cached rms
|
179
|
+
rstd_row = tl.load(RSTD_ptr)
|
180
|
+
|
181
|
+
X_row = X_row.to(tl.float32)
|
182
|
+
|
183
|
+
# Different bacward graphs for different casting modes
|
184
|
+
if casting_mode == _CASTING_MODE_LLAMA:
|
185
|
+
m = (dY_row * W_row).to(tl.float32)
|
186
|
+
|
187
|
+
elif casting_mode == _CASTING_MODE_GEMMA:
|
188
|
+
dY_row = dY_row.to(tl.float32)
|
189
|
+
m = dY_row * W_row
|
190
|
+
else:
|
191
|
+
m = dY_row * W_row
|
192
|
+
|
193
|
+
dX_row = rstd_row * m
|
194
|
+
|
195
|
+
if has_dS_out:
|
196
|
+
dS_out_row = tl.load(dS_out_ptr + col_offsets, mask=mask, other=0.0)
|
197
|
+
dX_row += (rstd_row) * (
|
198
|
+
-(1 / n_cols) * rstd_row * rstd_row * tl.sum(m * X_row, axis=0) * X_row
|
199
|
+
) + dS_out_row
|
200
|
+
dS_out_ptr += dS_out_row_stride
|
201
|
+
else:
|
202
|
+
dX_row += (rstd_row) * (-(1 / n_cols) * rstd_row * rstd_row * tl.sum(m * X_row, axis=0) * X_row)
|
203
|
+
|
204
|
+
# calculate the gradient of W
|
205
|
+
if casting_mode == _CASTING_MODE_LLAMA:
|
206
|
+
dW_row += dY_row * (X_row * rstd_row).to(X_dtype)
|
207
|
+
else:
|
208
|
+
# here X_row is already in fp32 (see previous if block)
|
209
|
+
dW_row += dY_row * (X_row * rstd_row)
|
210
|
+
|
211
|
+
tl.store(dX_ptr + col_offsets, dX_row.to(X_dtype), mask=mask)
|
212
|
+
|
213
|
+
dY_ptr += dY_row_stride
|
214
|
+
dX_ptr += dX_row_stride
|
215
|
+
X_ptr += X_row_stride
|
216
|
+
RSTD_ptr += RSTD_row_stride
|
217
|
+
|
218
|
+
tl.store(dW_ptr + row_block_id * dW_row_stride + col_offsets, dW_row, mask=mask)
|
219
|
+
|
220
|
+
|
221
|
+
_str_to_casting_mode = {
|
222
|
+
"llama": _CASTING_MODE_LLAMA.value,
|
223
|
+
"gemma": _CASTING_MODE_GEMMA.value,
|
224
|
+
"none": _CASTING_MODE_NONE.value,
|
225
|
+
}
|
226
|
+
|
227
|
+
|
228
|
+
def fused_add_rms_norm_forward(X, R, W, eps, offset, casting_mode):
|
229
|
+
if not isinstance(casting_mode, int):
|
230
|
+
assert casting_mode in _str_to_casting_mode, f"Invalid casting mode: {casting_mode}"
|
231
|
+
casting_mode = _str_to_casting_mode[casting_mode]
|
232
|
+
else:
|
233
|
+
assert casting_mode in _str_to_casting_mode.values(), f"Invalid casting mode: {casting_mode}"
|
234
|
+
|
235
|
+
shape = X.shape
|
236
|
+
dim = shape[-1]
|
237
|
+
X = X.view(-1, dim)
|
238
|
+
R = R.view(-1, dim)
|
239
|
+
n_rows, n_cols = X.shape
|
240
|
+
BLOCK_SIZE, num_warps = calculate_settings(n_cols)
|
241
|
+
|
242
|
+
Y = torch.empty((n_rows, n_cols), dtype=X.dtype, device=X.device)
|
243
|
+
S = torch.empty((n_rows, n_cols), dtype=X.dtype, device=X.device)
|
244
|
+
# RSTD is to cache rstd for each row
|
245
|
+
# RSTD is always computed/stored in fp32 if we are using Llama or Gemma casting mode
|
246
|
+
rstd_dtype = torch.float32 if casting_mode in (_CASTING_MODE_LLAMA.value, _CASTING_MODE_GEMMA.value) else X.dtype
|
247
|
+
RSTD = torch.empty(n_rows, dtype=rstd_dtype, device=X.device)
|
248
|
+
|
249
|
+
# Check constraints.
|
250
|
+
assert X.shape[1] == W.shape[0], "Incompatible hidden size dimension between tensor1.shape[1] and tensor2.shape[0]"
|
251
|
+
|
252
|
+
# XPU-specific optimization
|
253
|
+
kernel_args = {}
|
254
|
+
if X.device.type == "xpu":
|
255
|
+
kernel_args["grf_mode"] = "large"
|
256
|
+
|
257
|
+
# TODO: add _block_fused_add_rms_norm_forward_kernel
|
258
|
+
_fused_add_rms_norm_forward_kernel[(n_rows,)](
|
259
|
+
Y,
|
260
|
+
Y.stride(0),
|
261
|
+
S,
|
262
|
+
S.stride(0),
|
263
|
+
X,
|
264
|
+
X.stride(0),
|
265
|
+
R,
|
266
|
+
R.stride(0),
|
267
|
+
W,
|
268
|
+
W.stride(0),
|
269
|
+
RSTD,
|
270
|
+
RSTD.stride(0),
|
271
|
+
n_cols,
|
272
|
+
eps,
|
273
|
+
offset,
|
274
|
+
casting_mode,
|
275
|
+
BLOCK_SIZE=BLOCK_SIZE,
|
276
|
+
num_warps=num_warps,
|
277
|
+
**kernel_args, # XPU-specific optimization
|
278
|
+
)
|
279
|
+
|
280
|
+
return Y.view(*shape), S.view(*shape), RSTD, BLOCK_SIZE, num_warps, casting_mode
|
281
|
+
|
282
|
+
|
283
|
+
def fused_add_rms_norm_backward(dY, dS_out, S, W, RSTD, offset, casting_mode, BLOCK_SIZE, num_warps, in_place):
|
284
|
+
shape = dY.shape
|
285
|
+
dim = shape[-1]
|
286
|
+
dY = dY.view(-1, dim)
|
287
|
+
dS_out = dS_out.view(-1, dim)
|
288
|
+
S = S.view(-1, dim)
|
289
|
+
n_rows, n_cols = dY.shape
|
290
|
+
|
291
|
+
sm_count = 1
|
292
|
+
if S.device.type == "cuda":
|
293
|
+
sm_count = torch.cuda.get_device_properties(S.device).multi_processor_count
|
294
|
+
elif S.device.type == "xpu":
|
295
|
+
sm_count = torch.xpu.get_device_properties(S.device).gpu_eu_count
|
296
|
+
|
297
|
+
# fp32 for numerical stability especially.
|
298
|
+
_dW = torch.empty((sm_count, n_cols), dtype=torch.float32, device=W.device)
|
299
|
+
|
300
|
+
if n_cols > BLOCK_SIZE:
|
301
|
+
raise RuntimeError("This layer norm doesn't support feature dim >= 64KB.")
|
302
|
+
rows_per_program = math.ceil(n_rows / sm_count)
|
303
|
+
grid = (sm_count,)
|
304
|
+
|
305
|
+
if in_place is True:
|
306
|
+
dX = dY
|
307
|
+
else:
|
308
|
+
dX = torch.empty_like(dY)
|
309
|
+
|
310
|
+
# XPU-specific optimization
|
311
|
+
kernel_args = {}
|
312
|
+
if S.device.type == "xpu":
|
313
|
+
kernel_args["grf_mode"] = "large"
|
314
|
+
|
315
|
+
# TODO: add _block_fused_add_rms_norm_backward_kernel
|
316
|
+
_fused_add_rms_norm_backward_kernel[grid](
|
317
|
+
dY,
|
318
|
+
dY.stride(0),
|
319
|
+
dS_out,
|
320
|
+
dS_out.stride(0),
|
321
|
+
dX,
|
322
|
+
dX.stride(0),
|
323
|
+
S,
|
324
|
+
S.stride(0),
|
325
|
+
torch_to_triton_dtype[S.dtype],
|
326
|
+
W,
|
327
|
+
W.stride(0),
|
328
|
+
RSTD,
|
329
|
+
RSTD.stride(0),
|
330
|
+
_dW,
|
331
|
+
_dW.stride(0),
|
332
|
+
n_rows,
|
333
|
+
n_cols,
|
334
|
+
offset,
|
335
|
+
rows_per_program,
|
336
|
+
casting_mode,
|
337
|
+
BLOCK_SIZE=BLOCK_SIZE,
|
338
|
+
num_warps=num_warps,
|
339
|
+
has_dS_out=dS_out is not None,
|
340
|
+
**kernel_args, # XPU-specific optimization
|
341
|
+
)
|
342
|
+
|
343
|
+
dX = dX.view(*shape)
|
344
|
+
dW = _dW.sum(dim=0).to(W.dtype)
|
345
|
+
|
346
|
+
return dX, dX, dW # dR is equal to dX
|
347
|
+
|
348
|
+
|
349
|
+
class LigerFusedAddRMSNormFunction(torch.autograd.Function):
|
350
|
+
"""
|
351
|
+
Performs a fused operation that first adds a residual tensor to the hidden_states tensor (`X`), then applies RMSNorm (Root Mean Square Normalization) to the result using the weight tensor `W`, with optional offset and casting mode.
|
352
|
+
|
353
|
+
This class implements the following sequence, commonly used in transformer decoder layers:
|
354
|
+
1. hidden_states = residual + hidden_states
|
355
|
+
2. residual = hidden_states (after addition)
|
356
|
+
3. hidden_states = rmsnorm(hidden_states)
|
357
|
+
|
358
|
+
Both the normalized hidden_states and the updated residual are returned as outputs.
|
359
|
+
|
360
|
+
Some models use an 'offset' to shift the weight tensor `W` by a constant value. For example, Gemma
|
361
|
+
uses an offset of 1.0, so the computation becomes `(X / RMS(X)) * (W + 1.0)` instead of the usual
|
362
|
+
`(X / RMS(X)) * W`. You can pass the offset value as an argument to the forward function.
|
363
|
+
|
364
|
+
In addition, different models cast their inputs at different places during RMSNorm computation. For
|
365
|
+
example, Gemma casts everything to fp32 before starting the computation, while Llama casts only the
|
366
|
+
inverse RMS to fp32. You can specify the casting mode using the `casting_mode` argument. We currently
|
367
|
+
support the following casting modes (they match HuggingFace Transformers' implementations):
|
368
|
+
- 'llama': matches the Llama implementation, where only the inverse RMS is computed on fp32.
|
369
|
+
- 'gemma': matches the Gemma implementation, where everything is cast to fp32, then computed, then cast back to the original dtype.
|
370
|
+
- 'none': no casting is done. The computation is done in the original dtype. This saves memory and is slightly faster, but has more error w.r.t. the original implementation.
|
371
|
+
|
372
|
+
The `in_place` option determines whether to modify dY in-place to store dX. This defaults to `True` to save memory.
|
373
|
+
"""
|
374
|
+
|
375
|
+
@staticmethod
|
376
|
+
@ensure_contiguous
|
377
|
+
def forward(ctx, X, R, W, eps, offset=0.0, casting_mode="llama", in_place=False):
|
378
|
+
"""
|
379
|
+
X: (B, T, H) or (BxT, H)
|
380
|
+
W: (H,)
|
381
|
+
"""
|
382
|
+
# TODO: add row_mode
|
383
|
+
Y, S, RSTD, BLOCK_SIZE, num_warps, casting_mode = fused_add_rms_norm_forward(X, R, W, eps, offset, casting_mode)
|
384
|
+
ctx.offset = offset
|
385
|
+
ctx.casting_mode = casting_mode
|
386
|
+
ctx.in_place = in_place
|
387
|
+
ctx.BLOCK_SIZE = BLOCK_SIZE
|
388
|
+
ctx.num_warps = num_warps
|
389
|
+
ctx.save_for_backward(S, W, RSTD)
|
390
|
+
return Y, S
|
391
|
+
|
392
|
+
@staticmethod
|
393
|
+
@ensure_contiguous
|
394
|
+
def backward(ctx, dY, dS_out):
|
395
|
+
"""
|
396
|
+
Y: (B, T, H) or (BxT, H)
|
397
|
+
"""
|
398
|
+
S, W, RSTD = ctx.saved_tensors
|
399
|
+
dX, dR, dW = fused_add_rms_norm_backward(
|
400
|
+
dY,
|
401
|
+
dS_out,
|
402
|
+
S,
|
403
|
+
W,
|
404
|
+
RSTD,
|
405
|
+
ctx.offset,
|
406
|
+
ctx.casting_mode,
|
407
|
+
ctx.BLOCK_SIZE,
|
408
|
+
ctx.num_warps,
|
409
|
+
ctx.in_place,
|
410
|
+
)
|
411
|
+
|
412
|
+
return dX, dR, dW, None, None, None, None, None
|
@@ -5,6 +5,7 @@ from typing import TYPE_CHECKING
|
|
5
5
|
# Always-safe imports (independent of 'transformers')
|
6
6
|
from liger_kernel.transformers.cross_entropy import LigerCrossEntropyLoss # noqa: F401
|
7
7
|
from liger_kernel.transformers.dyt import LigerDyT # noqa: F401
|
8
|
+
from liger_kernel.transformers.fused_add_rms_norm import LigerFusedAddRMSNorm # noqa: F401
|
8
9
|
from liger_kernel.transformers.fused_linear_cross_entropy import LigerFusedLinearCrossEntropyLoss # noqa: F401
|
9
10
|
from liger_kernel.transformers.fused_linear_jsd import LigerFusedLinearJSD # noqa: F401
|
10
11
|
from liger_kernel.transformers.geglu import LigerGEGLUMLP # noqa: F401
|
@@ -121,6 +122,7 @@ __all__ = [
|
|
121
122
|
"LigerGEGLUMLP",
|
122
123
|
"LigerJSD",
|
123
124
|
"LigerLayerNorm",
|
125
|
+
"LigerFusedAddRMSNorm",
|
124
126
|
"LigerRMSNorm",
|
125
127
|
"liger_rotary_pos_emb",
|
126
128
|
"LigerBlockSparseTop2MLP",
|
@@ -2,6 +2,7 @@ from typing import Optional
|
|
2
2
|
|
3
3
|
from liger_kernel.ops.cross_entropy import LigerCrossEntropyFunction
|
4
4
|
from liger_kernel.ops.dyt import LigerDyTFunction
|
5
|
+
from liger_kernel.ops.fused_add_rms_norm import LigerFusedAddRMSNormFunction
|
5
6
|
from liger_kernel.ops.fused_linear_cross_entropy import LigerFusedLinearCrossEntropyFunction
|
6
7
|
from liger_kernel.ops.fused_linear_jsd import LigerFusedLinearJSDFunction
|
7
8
|
from liger_kernel.ops.fused_neighborhood_attention import LigerFusedNeighborhoodAttentionFunction
|
@@ -253,6 +254,10 @@ def liger_rms_norm(X, W, eps, offset: float = 0.0, casting_mode: str = "llama",
|
|
253
254
|
return LigerRMSNormFunction.apply(X, W, eps, offset, casting_mode, in_place)
|
254
255
|
|
255
256
|
|
257
|
+
def liger_fused_add_rms_norm(X, R, W, eps, offset: float = 0.0, casting_mode: str = "llama", in_place: bool = True):
|
258
|
+
return LigerFusedAddRMSNormFunction.apply(X, R, W, eps, offset, casting_mode, in_place)
|
259
|
+
|
260
|
+
|
256
261
|
def liger_rope(q, k, cos, sin, position_ids=None, unsqueeze_dim=1):
|
257
262
|
return LigerRopeFunction.apply(q, k, cos, sin, position_ids, unsqueeze_dim)
|
258
263
|
|
@@ -0,0 +1,39 @@
|
|
1
|
+
import torch
|
2
|
+
import torch.nn as nn
|
3
|
+
|
4
|
+
from liger_kernel.ops.fused_add_rms_norm import LigerFusedAddRMSNormFunction
|
5
|
+
|
6
|
+
|
7
|
+
class LigerFusedAddRMSNorm(nn.Module):
|
8
|
+
def __init__(
|
9
|
+
self,
|
10
|
+
hidden_size,
|
11
|
+
eps=1e-6,
|
12
|
+
offset=0.0,
|
13
|
+
casting_mode="llama",
|
14
|
+
init_fn="ones",
|
15
|
+
in_place=False,
|
16
|
+
):
|
17
|
+
super().__init__()
|
18
|
+
assert init_fn in [
|
19
|
+
"ones",
|
20
|
+
"zeros",
|
21
|
+
], f"init_fn must be either 'ones' or 'zeros', got {init_fn}"
|
22
|
+
self.weight = nn.Parameter(torch.ones(hidden_size) if init_fn == "ones" else torch.zeros(hidden_size))
|
23
|
+
self.variance_epsilon, self.offset, self.casting_mode, self.in_place = (eps, offset, casting_mode, in_place)
|
24
|
+
|
25
|
+
def forward(self, hidden_states, residual):
|
26
|
+
return LigerFusedAddRMSNormFunction.apply(
|
27
|
+
hidden_states,
|
28
|
+
residual,
|
29
|
+
self.weight,
|
30
|
+
self.variance_epsilon,
|
31
|
+
self.offset,
|
32
|
+
self.casting_mode,
|
33
|
+
self.in_place,
|
34
|
+
)
|
35
|
+
|
36
|
+
def extra_repr(self):
|
37
|
+
return (
|
38
|
+
f"{tuple(self.weight.shape)}, eps={self.variance_epsilon}, offset={self.offset}, in_place={self.in_place}"
|
39
|
+
)
|
@@ -19,6 +19,7 @@ liger_kernel/chunked_loss/simpo_loss.py,sha256=fy2w8KbhMrBv7b1jdIeH3bBFxY52bPQPZ
|
|
19
19
|
liger_kernel/ops/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
|
20
20
|
liger_kernel/ops/cross_entropy.py,sha256=e8THGnhOcy_0SbOLABx67HEM7-B8a8pG7nDKbCRpQKM,19123
|
21
21
|
liger_kernel/ops/dyt.py,sha256=gCLz4S8aul8SY9nvIGaoK67aGb7U9MJRQdo3ONqmQYs,5417
|
22
|
+
liger_kernel/ops/fused_add_rms_norm.py,sha256=UBqmlqFCmhSAIpkNKd8rrfXatX7Z4J9bp2dX9A0lrJQ,14017
|
22
23
|
liger_kernel/ops/fused_linear_cross_entropy.py,sha256=5fbGhN85n3zf0uIdJ7PYHWIRzTf0VTFiS0ARtOmqIP0,11020
|
23
24
|
liger_kernel/ops/fused_linear_jsd.py,sha256=CSoprxb-YcJy-YUKiTcYkxN8sb9h2kdk_iHuncvSV5c,9683
|
24
25
|
liger_kernel/ops/fused_neighborhood_attention.py,sha256=vPi5xbnh6wxyZehaqo6Tuilqo2fN5SGDiONjnNmIKqs,35556
|
@@ -39,12 +40,13 @@ liger_kernel/ops/tvd.py,sha256=FHJtLQI95ijqgg9UtaHpMAjSCiPxB6CduPwPMcGxelc,6405
|
|
39
40
|
liger_kernel/ops/utils.py,sha256=uoFKQqo-34N2TWQNvXMFywqGiOMMXNEVBxVojzlUAa0,3836
|
40
41
|
liger_kernel/ops/experimental/embedding.py,sha256=tolj3tItkzpSb30zWqDN2_yX4ectflaQ8HMyKyFIQc8,4172
|
41
42
|
liger_kernel/ops/experimental/mm_int8int2.py,sha256=TrS9lpwekrik_w5qE7AhMJD1bcq-OidjtbsW80oZ6IM,13314
|
42
|
-
liger_kernel/transformers/__init__.py,sha256=
|
43
|
+
liger_kernel/transformers/__init__.py,sha256=VoHQp5emsAJAouql37RuvtGFeZCoMIHgoIxfsyYMTc8,7564
|
43
44
|
liger_kernel/transformers/auto_model.py,sha256=0qCTRZt280Bj_LcFdzo9hlaR-BWNazawXOGgoCZjgEg,1545
|
44
45
|
liger_kernel/transformers/cross_entropy.py,sha256=z3KTWQnFxr_IZaVjtYt0ZNEWQdDdYThN35xWkHlDGH0,1683
|
45
46
|
liger_kernel/transformers/dyt.py,sha256=i-4GPaMrl-jab9TVI5qN0-H9qycn_mCbV82ozU4nbmU,723
|
46
47
|
liger_kernel/transformers/fsdp.py,sha256=CUiyjTmjkjY7pLXQv8ly9rnzgXw6529csd9pvtJNMYc,3096
|
47
|
-
liger_kernel/transformers/functional.py,sha256=
|
48
|
+
liger_kernel/transformers/functional.py,sha256=PXnACWD7kzgge50RdOUuvtmOTS7DVkkrL7mm0cX5bOc,7734
|
49
|
+
liger_kernel/transformers/fused_add_rms_norm.py,sha256=7_Bzg-x6lLe6W1qG2DtjDALhEpNZlC6N5GppEs9cTYY,1199
|
48
50
|
liger_kernel/transformers/fused_linear_cross_entropy.py,sha256=O8Sg5BT81nTaY9fSGoOY9dOD9ekibwwiuXhdUHaxntQ,1742
|
49
51
|
liger_kernel/transformers/fused_linear_jsd.py,sha256=bZ4otCvWBuOnA5XdQL-FzZVItJlDt-ht9e_pG7PG93E,3999
|
50
52
|
liger_kernel/transformers/fused_neighborhood_attention.py,sha256=TxYDUAt9B6WSP14aJP66C_2Mbds2sSIPGnamhUSTrC8,7957
|
@@ -90,9 +92,9 @@ liger_kernel/transformers/trainer/__init__.py,sha256=p7yQfklV8-467qSz_ZMimkbDF7H
|
|
90
92
|
liger_kernel/transformers/trainer/orpo_trainer.py,sha256=tX0h63aOFe3rNqTmk6JpMf75UPo981yzEa6TghnjS0Q,5370
|
91
93
|
liger_kernel/triton/__init__.py,sha256=qCiCamzCRv6lpV8IqpAc9YMdNKC7GKurClWceQPnlis,92
|
92
94
|
liger_kernel/triton/monkey_patch.py,sha256=Rd0hUHAzDkFfHvnX7-PBaNK5EKnZhtfM_h-fgQH9HPY,1568
|
93
|
-
liger_kernel_nightly-0.6.0.
|
94
|
-
liger_kernel_nightly-0.6.0.
|
95
|
-
liger_kernel_nightly-0.6.0.
|
96
|
-
liger_kernel_nightly-0.6.0.
|
97
|
-
liger_kernel_nightly-0.6.0.
|
98
|
-
liger_kernel_nightly-0.6.0.
|
95
|
+
liger_kernel_nightly-0.6.0.dev20250718080702.dist-info/LICENSE,sha256=OhzLDHJ0to4a8sodVLELZiCFylZ1NAAYLs-HrjPy0ag,1312
|
96
|
+
liger_kernel_nightly-0.6.0.dev20250718080702.dist-info/METADATA,sha256=mNkIMGPTMPdmmjDsW54kDe0WhPi8Ep0Cpt4koWQQuaE,24672
|
97
|
+
liger_kernel_nightly-0.6.0.dev20250718080702.dist-info/NOTICE,sha256=njwnoPZLh9AN8SJQzxvCGLHi-8X__AvWRze6joNXIY8,2066
|
98
|
+
liger_kernel_nightly-0.6.0.dev20250718080702.dist-info/WHEEL,sha256=iAkIy5fosb7FzIOwONchHf19Qu7_1wCWyFNR5gu9nU0,91
|
99
|
+
liger_kernel_nightly-0.6.0.dev20250718080702.dist-info/top_level.txt,sha256=2eghu4hA3LnkM7ElW92tQ8zegWKgSbeo-k-aGe1YnvY,13
|
100
|
+
liger_kernel_nightly-0.6.0.dev20250718080702.dist-info/RECORD,,
|
File without changes
|
File without changes
|
File without changes
|