| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | # SPDX-FileCopyrightText: Tri Dao |
| 4 | # |
| 5 | # This file contains code copied from the flash-linear-attention project. |
| 6 | # The original source code was licensed under the MIT license and included |
| 7 | # the following copyright notice: |
| 8 | # Copyright (c) 2024, Tri Dao. |
| 9 | |
| 10 | # ruff: noqa: E501 |
| 11 | # Based on the Triton LayerNorm tutorial: https://triton-lang.org/main/getting-started/tutorials/05-layer-norm.html |
| 12 | # For the backward pass, we keep weight_grad and bias_grad in registers and accumulate. |
| 13 | # This backward pass is faster for dimensions up to 8k, but after that it's much slower due to register spilling. |
| 14 | # The models we train have hidden dim up to 8k anyway (e.g. Llama 70B), so this is fine. |
| 15 | |
| 16 | from dataclasses import dataclass |
| 17 | |
| 18 | import torch |
| 19 | import torch.nn as nn |
| 20 | import torch.nn.functional as F |
| 21 | from einops import rearrange |
| 22 | |
| 23 | from vllm.model_executor.warmup.jit_warmup import VllmJitKernel |
| 24 | from vllm.model_executor.warmup.jit_warmup_triton_helper import ( |
| 25 | TritonPointerInputVariant, |
| 26 | TritonWarmupTensor, |
| 27 | triton_scalar_specialization_rep, |
| 28 | ) |
| 29 | from vllm.triton_utils import tl, triton |
| 30 | from vllm.utils.math_utils import cdiv, next_power_of_2 |
| 31 | from vllm.utils.platform_utils import num_compute_units |
| 32 | |
| 33 | from .utils import input_guard |
| 34 | |
| 35 | |
| 36 | def rms_norm_ref( |
| 37 | x, |
| 38 | weight, |
| 39 | bias, |
| 40 | z=None, |
| 41 | eps=1e-6, |
| 42 | group_size=None, |
| 43 | norm_before_gate=True, |
| 44 | upcast=True, |
| 45 | ): |
| 46 | dtype = x.dtype |
| 47 | weight = weight.float() |
| 48 | bias = bias.float() if bias is not None else None |
| 49 | if upcast: |
| 50 | x = x.float() |
| 51 | z = z.float() if z is not None else z |
| 52 | if z is not None and not norm_before_gate: |
| 53 | x = x * F.silu(z) |
| 54 | if group_size is None: |
| 55 | rstd = 1 / torch.sqrt((x.square()).mean(dim=-1, keepdim=True) + eps) |
| 56 | out = (x * rstd * weight) + bias if bias is not None else (x * rstd * weight) |
| 57 | else: |
| 58 | x_group = rearrange(x, "... (g d) -> ... g d", d=group_size) |
| 59 | rstd = 1 / torch.sqrt((x_group.square()).mean(dim=-1, keepdim=True) + eps) |
| 60 | out = rearrange(x_group * rstd, "... g d -> ... (g d)") * weight |
| 61 | if bias is not None: |
| 62 | out = out + bias |
| 63 | if z is not None and norm_before_gate: |
| 64 | out *= F.silu(z) |
| 65 | return out.to(dtype) |
| 66 | |
| 67 | |
| 68 | @triton.heuristics( |
| 69 | { |
| 70 | "HAS_BIAS": lambda args: args["B"] is not None, |
| 71 | "HAS_Z": lambda args: args["Z"] is not None, |
| 72 | } |
| 73 | ) |
| 74 | @triton.jit |
| 75 | def layer_norm_fwd_kernel( |
| 76 | X, # pointer to the input |
| 77 | Y, # pointer to the output |
| 78 | W, # pointer to the weights |
| 79 | B, # pointer to the biases |
| 80 | Z, # pointer to the other branch |
| 81 | Mean, # pointer to the mean |
| 82 | Rstd, # pointer to the 1/std |
| 83 | stride_x_row, # how much to increase the pointer when moving by 1 row |
| 84 | stride_y_row, |
| 85 | stride_z_row, |
| 86 | M, # number of rows in X |
| 87 | N: tl.constexpr, # number of columns in X |
| 88 | eps, # epsilon to avoid division by zero |
| 89 | BLOCK_N: tl.constexpr, |
| 90 | ROWS_PER_BLOCK: tl.constexpr, |
| 91 | HAS_BIAS: tl.constexpr, |
| 92 | HAS_Z: tl.constexpr, |
| 93 | NORM_BEFORE_GATE: tl.constexpr, |
| 94 | IS_RMS_NORM: tl.constexpr, |
| 95 | ACTIVATION: tl.constexpr, |
| 96 | ): |
| 97 | # Map the program id to the starting row of X and Y it should compute. |
| 98 | row_start = tl.program_id(0) * ROWS_PER_BLOCK |
| 99 | group = tl.program_id(1) |
| 100 | |
| 101 | # Create 2D tile: [ROWS_PER_BLOCK, BLOCK_N] |
| 102 | rows = row_start + tl.arange(0, ROWS_PER_BLOCK) |
| 103 | cols = tl.arange(0, BLOCK_N) |
| 104 | |
| 105 | # Compute offsets for 2D tile |
| 106 | row_offsets = rows[:, None] * stride_x_row |
| 107 | col_offsets = cols[None, :] + group * N |
| 108 | |
| 109 | # Base pointers |
| 110 | X_base = X + row_offsets + col_offsets |
| 111 | Y_base = Y + rows[:, None] * stride_y_row + col_offsets |
| 112 | |
| 113 | # Create mask for valid rows and columns |
| 114 | row_mask = rows[:, None] < M |
| 115 | col_mask = cols[None, :] < N |
| 116 | mask = row_mask & col_mask |
| 117 | |
| 118 | # Load input data with 2D tile |
| 119 | x = tl.load(X_base, mask=mask, other=0.0).to(tl.float32) |
| 120 | |
| 121 | if HAS_Z and not NORM_BEFORE_GATE: |
| 122 | Z_base = Z + rows[:, None] * stride_z_row + col_offsets |
| 123 | z = tl.load(Z_base, mask=mask, other=0.0).to(tl.float32) |
| 124 | if ACTIVATION == "swish" or ACTIVATION == "silu": |
| 125 | x *= z * tl.sigmoid(z) |
| 126 | elif ACTIVATION == "sigmoid": |
| 127 | x *= tl.sigmoid(z) |
| 128 | |
| 129 | # Compute mean and variance per row (reduce along axis 1) |
| 130 | if not IS_RMS_NORM: |
| 131 | mean = tl.sum(x, axis=1) / N # Shape: [ROWS_PER_BLOCK] |
| 132 | # Store mean for each row |
| 133 | mean_offsets = group * M + rows |
| 134 | mean_mask = rows < M |
| 135 | tl.store(Mean + mean_offsets, mean, mask=mean_mask) |
| 136 | # Broadcast mean back to 2D for subtraction |
| 137 | xbar = tl.where(mask, x - mean[:, None], 0.0) |
| 138 | var = tl.sum(xbar * xbar, axis=1) / N # Shape: [ROWS_PER_BLOCK] |
| 139 | else: |
| 140 | xbar = tl.where(mask, x, 0.0) |
| 141 | var = tl.sum(xbar * xbar, axis=1) / N # Shape: [ROWS_PER_BLOCK] |
| 142 | mean = 0.0 # Placeholder for RMS norm |
| 143 | |
| 144 | rstd = tl.rsqrt(var + eps) # Shape: [ROWS_PER_BLOCK] |
| 145 | |
| 146 | # Store rstd for each row |
| 147 | rstd_offsets = group * M + rows |
| 148 | rstd_mask = rows < M |
| 149 | tl.store(Rstd + rstd_offsets, rstd, mask=rstd_mask) |
| 150 | |
| 151 | # Load weights and biases (broadcast across rows) |
| 152 | w_offsets = cols + group * N |
| 153 | w_mask = cols < N |
| 154 | w = tl.load(W + w_offsets, mask=w_mask, other=0.0).to(tl.float32) |
| 155 | |
| 156 | if HAS_BIAS: |
| 157 | b = tl.load(B + w_offsets, mask=w_mask, other=0.0).to(tl.float32) |
| 158 | |
| 159 | # Normalize and apply linear transformation |
| 160 | if not IS_RMS_NORM: |
| 161 | x_hat = (x - mean[:, None]) * rstd[:, None] |
| 162 | else: |
| 163 | x_hat = x * rstd[:, None] |
| 164 | |
| 165 | y = x_hat * w[None, :] + b[None, :] if HAS_BIAS else x_hat * w[None, :] |
| 166 | |
| 167 | if HAS_Z and NORM_BEFORE_GATE: |
| 168 | Z_base = Z + rows[:, None] * stride_z_row + col_offsets |
| 169 | z = tl.load(Z_base, mask=mask, other=0.0).to(tl.float32) |
| 170 | if ACTIVATION == "swish" or ACTIVATION == "silu": |
| 171 | y *= z * tl.sigmoid(z) |
| 172 | elif ACTIVATION == "sigmoid": |
| 173 | y *= tl.sigmoid(z) |
| 174 | |
| 175 | # Write output |
| 176 | tl.store(Y_base, y, mask=mask) |
| 177 | |
| 178 | |
| 179 | def calc_rows_per_block(M: int, device: torch.device) -> int: |
| 180 | sm_count = num_compute_units(device.index) |
| 181 | rows_per_block = next_power_of_2(cdiv(M, 2 * sm_count)) |
| 182 | rows_per_block = min(rows_per_block, 4) |
| 183 | return rows_per_block |
| 184 | |
| 185 | |
| 186 | class LayerNormFwdKernel(VllmJitKernel["LayerNormFwdKernel.CompileKey"]): |
| 187 | """Own the runtime and compile-only dispatch for FLA layer norm.""" |
| 188 | |
| 189 | @dataclass(frozen=True) |
| 190 | class CompileKey: |
| 191 | x_dtype: torch.dtype |
| 192 | y_dtype: torch.dtype |
| 193 | weight_dtype: torch.dtype |
| 194 | bias_dtype: torch.dtype | None |
| 195 | z_dtype: torch.dtype | None |
| 196 | mean_dtype: torch.dtype | None |
| 197 | rstd_dtype: torch.dtype |
| 198 | input_variant: TritonPointerInputVariant |
| 199 | stride_x_row: int |
| 200 | stride_y_row: int |
| 201 | stride_z_row: int |
| 202 | M: int |
| 203 | N: int |
| 204 | BLOCK_N: int |
| 205 | ROWS_PER_BLOCK: int |
| 206 | NORM_BEFORE_GATE: bool |
| 207 | IS_RMS_NORM: bool |
| 208 | ACTIVATION: str |
| 209 | |
| 210 | kernel = staticmethod(layer_norm_fwd_kernel) |
| 211 | |
| 212 | def dispatch( # type: ignore[override] |
| 213 | self, |
| 214 | *, |
| 215 | x_dtype: torch.dtype, |
| 216 | y_dtype: torch.dtype, |
| 217 | weight_dtype: torch.dtype, |
| 218 | bias_dtype: torch.dtype | None, |
| 219 | z_dtype: torch.dtype | None, |
| 220 | mean_dtype: torch.dtype | None, |
| 221 | rstd_dtype: torch.dtype, |
| 222 | x_aligned: bool, |
| 223 | y_aligned: bool, |
| 224 | weight_aligned: bool, |
| 225 | bias_aligned: bool, |
| 226 | z_aligned: bool, |
| 227 | mean_aligned: bool, |
| 228 | rstd_aligned: bool, |
| 229 | stride_x_row: int, |
| 230 | stride_y_row: int, |
| 231 | stride_z_row: int, |
| 232 | M: int, |
| 233 | N: int, |
| 234 | BLOCK_N: int, |
| 235 | ROWS_PER_BLOCK: int, |
| 236 | norm_before_gate: bool, |
| 237 | is_rms_norm: bool, |
| 238 | activation: str, |
| 239 | ) -> CompileKey: |
| 240 | input_variant = TritonPointerInputVariant.from_alignment( |
| 241 | x=x_aligned, |
| 242 | y=y_aligned, |
| 243 | weight=weight_aligned, |
| 244 | bias=bias_aligned, |
| 245 | z=z_aligned, |
| 246 | mean=mean_aligned, |
| 247 | rstd=rstd_aligned, |
| 248 | ) |
| 249 | return self.CompileKey( |
| 250 | x_dtype=x_dtype, |
| 251 | y_dtype=y_dtype, |
| 252 | weight_dtype=weight_dtype, |
| 253 | bias_dtype=bias_dtype, |
| 254 | z_dtype=z_dtype, |
| 255 | mean_dtype=mean_dtype, |
| 256 | rstd_dtype=rstd_dtype, |
| 257 | input_variant=input_variant, |
| 258 | stride_x_row=triton_scalar_specialization_rep(stride_x_row), |
| 259 | stride_y_row=triton_scalar_specialization_rep(stride_y_row), |
| 260 | stride_z_row=triton_scalar_specialization_rep(stride_z_row), |
| 261 | M=triton_scalar_specialization_rep(M), |
| 262 | N=N, |
| 263 | BLOCK_N=BLOCK_N, |
| 264 | ROWS_PER_BLOCK=ROWS_PER_BLOCK, |
| 265 | NORM_BEFORE_GATE=norm_before_gate, |
| 266 | IS_RMS_NORM=is_rms_norm, |
| 267 | ACTIVATION=activation, |
| 268 | ) |
| 269 | |
| 270 | def get_warmup_keys( |
| 271 | self, |
| 272 | *, |
| 273 | max_num_tokens: int, |
| 274 | rows_per_token: int, |
| 275 | group_size: int, |
| 276 | x_dtype: torch.dtype, |
| 277 | weight_dtype: torch.dtype, |
| 278 | device: torch.device, |
| 279 | norm_before_gate: bool, |
| 280 | is_rms_norm: bool, |
| 281 | activation: str, |
| 282 | ) -> list[CompileKey]: |
| 283 | if max_num_tokens < 1 or rows_per_token < 1: |
| 284 | return [] |
| 285 | |
| 286 | max_fused_size = 65536 // torch.empty((), dtype=x_dtype).element_size() |
| 287 | block_n = min(max_fused_size, triton.next_power_of_2(group_size)) |
| 288 | if group_size > block_n: |
| 289 | return [] |
| 290 | |
| 291 | # Triton's runtime key only distinguishes three integer specialization |
| 292 | # classes. Keep one real M for each (M specialization, rows/block) pair. |
| 293 | m_values: dict[tuple[int, int], int] = {} |
| 294 | for num_tokens in range(1, max_num_tokens + 1): |
| 295 | M = num_tokens * rows_per_token |
| 296 | signature = ( |
| 297 | triton_scalar_specialization_rep(M), |
| 298 | calc_rows_per_block(M, device), |
| 299 | ) |
| 300 | m_values.setdefault(signature, M) |
| 301 | |
| 302 | keys: dict[LayerNormFwdKernel.CompileKey, None] = {} |
| 303 | for M in m_values.values(): |
| 304 | key = self.dispatch( |
| 305 | x_dtype=x_dtype, |
| 306 | y_dtype=x_dtype, |
| 307 | weight_dtype=weight_dtype, |
| 308 | bias_dtype=None, |
| 309 | z_dtype=x_dtype, |
| 310 | mean_dtype=None if is_rms_norm else torch.float32, |
| 311 | rstd_dtype=torch.float32, |
| 312 | x_aligned=True, |
| 313 | y_aligned=True, |
| 314 | weight_aligned=True, |
| 315 | bias_aligned=True, |
| 316 | z_aligned=True, |
| 317 | mean_aligned=True, |
| 318 | rstd_aligned=True, |
| 319 | stride_x_row=group_size, |
| 320 | stride_y_row=group_size, |
| 321 | stride_z_row=group_size, |
| 322 | M=M, |
| 323 | N=group_size, |
| 324 | BLOCK_N=block_n, |
| 325 | ROWS_PER_BLOCK=calc_rows_per_block(M, device), |
| 326 | norm_before_gate=norm_before_gate, |
| 327 | is_rms_norm=is_rms_norm, |
| 328 | activation=activation, |
| 329 | ) |
| 330 | keys[key] = None |
| 331 | return list(keys) |
| 332 | |
| 333 | def compile(self, compile_key: CompileKey) -> None: |
| 334 | warmup = getattr(self.kernel, "warmup", None) |
| 335 | assert warmup is not None |
| 336 | variant = compile_key.input_variant |
| 337 | |
| 338 | def pointer(name: str, dtype: torch.dtype | None) -> TritonWarmupTensor | None: |
| 339 | return None if dtype is None else variant.pointer(name, dtype) |
| 340 | |
| 341 | warmup( |
| 342 | pointer("x", compile_key.x_dtype), |
| 343 | pointer("y", compile_key.y_dtype), |
| 344 | pointer("weight", compile_key.weight_dtype), |
| 345 | pointer("bias", compile_key.bias_dtype), |
| 346 | pointer("z", compile_key.z_dtype), |
| 347 | pointer("mean", compile_key.mean_dtype), |
| 348 | pointer("rstd", compile_key.rstd_dtype), |
| 349 | compile_key.stride_x_row, |
| 350 | compile_key.stride_y_row, |
| 351 | compile_key.stride_z_row, |
| 352 | compile_key.M, |
| 353 | compile_key.N, |
| 354 | 1e-5, |
| 355 | BLOCK_N=compile_key.BLOCK_N, |
| 356 | ROWS_PER_BLOCK=compile_key.ROWS_PER_BLOCK, |
| 357 | HAS_BIAS=compile_key.bias_dtype is not None, |
| 358 | HAS_Z=compile_key.z_dtype is not None, |
| 359 | NORM_BEFORE_GATE=compile_key.NORM_BEFORE_GATE, |
| 360 | IS_RMS_NORM=compile_key.IS_RMS_NORM, |
| 361 | ACTIVATION=compile_key.ACTIVATION, |
| 362 | num_warps=min(max(compile_key.BLOCK_N // 256, 1), 8), |
| 363 | grid=(1, 1), |
| 364 | ) |
| 365 | |
| 366 | def __call__( |
| 367 | self, |
| 368 | x: torch.Tensor, |
| 369 | out: torch.Tensor, |
| 370 | weight: torch.Tensor, |
| 371 | bias: torch.Tensor | None, |
| 372 | z: torch.Tensor | None, |
| 373 | mean: torch.Tensor | None, |
| 374 | rstd: torch.Tensor, |
| 375 | eps: float, |
| 376 | group_size: int, |
| 377 | norm_before_gate: bool, |
| 378 | is_rms_norm: bool, |
| 379 | activation: str, |
| 380 | ) -> None: |
| 381 | M = x.shape[0] |
| 382 | max_fused_size = 65536 // x.element_size() |
| 383 | block_n = min(max_fused_size, triton.next_power_of_2(group_size)) |
| 384 | if group_size > block_n: |
| 385 | raise RuntimeError("This layer norm doesn't support feature dim >= 64KB.") |
| 386 | rows_per_block = calc_rows_per_block(M, x.device) |
| 387 | grid = (cdiv(M, rows_per_block), x.shape[1] // group_size) |
| 388 | self.kernel[grid]( |
| 389 | x, |
| 390 | out, |
| 391 | weight, |
| 392 | bias, |
| 393 | z, |
| 394 | mean, |
| 395 | rstd, |
| 396 | x.stride(0), |
| 397 | out.stride(0), |
| 398 | z.stride(0) if z is not None else 0, |
| 399 | M, |
| 400 | group_size, |
| 401 | eps, |
| 402 | BLOCK_N=block_n, |
| 403 | ROWS_PER_BLOCK=rows_per_block, |
| 404 | HAS_BIAS=bias is not None, |
| 405 | HAS_Z=z is not None, |
| 406 | NORM_BEFORE_GATE=norm_before_gate, |
| 407 | IS_RMS_NORM=is_rms_norm, |
| 408 | ACTIVATION=activation, |
| 409 | num_warps=min(max(block_n // 256, 1), 8), |
| 410 | ) |
| 411 | |
| 412 | |
| 413 | _LAYER_NORM_FWD_KERNEL = LayerNormFwdKernel() |
| 414 | |
| 415 | |
| 416 | def warmup_layer_norm_fwd( |
| 417 | *, |
| 418 | max_num_tokens: int, |
| 419 | rows_per_token: int, |
| 420 | group_size: int, |
| 421 | x_dtype: torch.dtype, |
| 422 | weight_dtype: torch.dtype, |
| 423 | device: torch.device, |
| 424 | norm_before_gate: bool, |
| 425 | is_rms_norm: bool, |
| 426 | activation: str, |
| 427 | ) -> None: |
| 428 | _LAYER_NORM_FWD_KERNEL.warmup( |
| 429 | max_num_tokens=max_num_tokens, |
| 430 | rows_per_token=rows_per_token, |
| 431 | group_size=group_size, |
| 432 | x_dtype=x_dtype, |
| 433 | weight_dtype=weight_dtype, |
| 434 | device=device, |
| 435 | norm_before_gate=norm_before_gate, |
| 436 | is_rms_norm=is_rms_norm, |
| 437 | activation=activation, |
| 438 | ) |
| 439 | |
| 440 | |
| 441 | def layer_norm_fwd( |
| 442 | x: torch.Tensor, |
| 443 | weight: torch.Tensor, |
| 444 | bias: torch.Tensor, |
| 445 | eps: float, |
| 446 | z: torch.Tensor = None, |
| 447 | out: torch.Tensor = None, |
| 448 | group_size: int = None, |
| 449 | norm_before_gate: bool = True, |
| 450 | is_rms_norm: bool = False, |
| 451 | activation: str = "swish", |
| 452 | ): |
| 453 | M, N = x.shape |
| 454 | if group_size is None: |
| 455 | group_size = N |
| 456 | assert N % group_size == 0 |
| 457 | ngroups = N // group_size |
| 458 | assert x.stride(-1) == 1 |
| 459 | if z is not None: |
| 460 | assert z.stride(-1) == 1 |
| 461 | assert z.shape == (M, N) |
| 462 | assert weight.shape == (N,) |
| 463 | assert weight.stride(-1) == 1 |
| 464 | if bias is not None: |
| 465 | assert bias.stride(-1) == 1 |
| 466 | assert bias.shape == (N,) |
| 467 | # allocate output |
| 468 | if out is not None: |
| 469 | assert out.shape == x.shape |
| 470 | else: |
| 471 | out = torch.empty_like(x) |
| 472 | assert out.stride(-1) == 1 |
| 473 | mean = ( |
| 474 | torch.empty((ngroups * M,), dtype=torch.float32, device=x.device) |
| 475 | if not is_rms_norm |
| 476 | else None |
| 477 | ) |
| 478 | rstd = torch.empty((ngroups * M,), dtype=torch.float32, device=x.device) |
| 479 | _LAYER_NORM_FWD_KERNEL( |
| 480 | x, |
| 481 | out, |
| 482 | weight, |
| 483 | bias, |
| 484 | z, |
| 485 | mean, |
| 486 | rstd, |
| 487 | eps, |
| 488 | group_size, |
| 489 | norm_before_gate, |
| 490 | is_rms_norm, |
| 491 | activation, |
| 492 | ) |
| 493 | return out, mean, rstd |
| 494 | |
| 495 | |
| 496 | def _layer_norm_fn_impl( |
| 497 | x, |
| 498 | weight, |
| 499 | bias, |
| 500 | z=None, |
| 501 | eps=1e-6, |
| 502 | group_size=None, |
| 503 | norm_before_gate=True, |
| 504 | is_rms_norm=False, |
| 505 | activation: str = "swish", |
| 506 | ): |
| 507 | """Triton layer/RMS norm with optional gating. |
| 508 | |
| 509 | If z is not None, computes norm(x) * silu(z) when norm_before_gate, |
| 510 | else norm(x * silu(z)). |
| 511 | |
| 512 | This calls the triton kernel directly. The original code wrapped this |
| 513 | in a torch.autograd.Function (LayerNormFn) to save tensors for a |
| 514 | backward pass, but vLLM is inference-only so there is no backward pass. |
| 515 | The autograd wrapper also prevented torch.compile/dynamo from tracing |
| 516 | through the function due to its @staticmethod forward. |
| 517 | """ |
| 518 | x_shape_og = x.shape |
| 519 | x = x.reshape(-1, x.shape[-1]) |
| 520 | if x.stride(-1) != 1: |
| 521 | x = x.contiguous() |
| 522 | if z is not None: |
| 523 | assert z.shape == x_shape_og |
| 524 | z = z.reshape(-1, z.shape[-1]) |
| 525 | if z.stride(-1) != 1: |
| 526 | z = z.contiguous() |
| 527 | weight = weight.contiguous() |
| 528 | if bias is not None: |
| 529 | bias = bias.contiguous() |
| 530 | y, _, _ = layer_norm_fwd( |
| 531 | x, |
| 532 | weight, |
| 533 | bias, |
| 534 | eps, |
| 535 | z=z, |
| 536 | group_size=group_size, |
| 537 | norm_before_gate=norm_before_gate, |
| 538 | is_rms_norm=is_rms_norm, |
| 539 | activation=activation, |
| 540 | ) |
| 541 | return y.reshape(x_shape_og) |
| 542 | |
| 543 | |
| 544 | @input_guard |
| 545 | def layernorm_fn( |
| 546 | x, |
| 547 | weight, |
| 548 | bias, |
| 549 | z=None, |
| 550 | eps=1e-6, |
| 551 | group_size=None, |
| 552 | norm_before_gate=True, |
| 553 | is_rms_norm=False, |
| 554 | activation: str = "swish", |
| 555 | ): |
| 556 | return _layer_norm_fn_impl( |
| 557 | x, weight, bias, z, eps, group_size, norm_before_gate, is_rms_norm, activation |
| 558 | ) |
| 559 | |
| 560 | |
| 561 | @input_guard |
| 562 | def rmsnorm_fn( |
| 563 | x, |
| 564 | weight, |
| 565 | bias, |
| 566 | z=None, |
| 567 | eps=1e-6, |
| 568 | group_size=None, |
| 569 | norm_before_gate=True, |
| 570 | activation: str = "swish", |
| 571 | ): |
| 572 | return _layer_norm_fn_impl( |
| 573 | x, weight, bias, z, eps, group_size, norm_before_gate, True, activation |
| 574 | ) |
| 575 | |
| 576 | |
| 577 | class RMSNormGated(nn.Module): |
| 578 | def __init__( |
| 579 | self, |
| 580 | hidden_size, |
| 581 | eps: float = 1e-5, |
| 582 | group_size: int | None = None, |
| 583 | norm_before_gate: bool = False, |
| 584 | device: torch.device | None = None, |
| 585 | dtype: torch.dtype | None = None, |
| 586 | activation: str = "swish", |
| 587 | ): |
| 588 | """If group_size is not None, we do GroupNorm with each group having group_size elements. |
| 589 | group_size=None is equivalent to group_size=hidden_size (i.e. there's only 1 group). |
| 590 | """ |
| 591 | factory_kwargs = {"device": device, "dtype": dtype} |
| 592 | super().__init__() |
| 593 | self.eps = eps |
| 594 | self.activation = activation |
| 595 | self.weight = nn.Parameter(torch.empty(hidden_size, **factory_kwargs)) |
| 596 | self.register_parameter("bias", None) |
| 597 | self.group_size = group_size |
| 598 | self.norm_before_gate = norm_before_gate |
| 599 | self.reset_parameters() |
| 600 | |
| 601 | def reset_parameters(self): |
| 602 | torch.nn.init.ones_(self.weight) |
| 603 | |
| 604 | def forward(self, x, z=None): |
| 605 | """If z is not None, we do norm(x) * silu(z) if norm_before_gate, else norm(x * silu(z))""" |
| 606 | return rmsnorm_fn( |
| 607 | x, |
| 608 | self.weight, |
| 609 | self.bias, |
| 610 | z=z, |
| 611 | eps=self.eps, |
| 612 | group_size=self.group_size, |
| 613 | norm_before_gate=self.norm_before_gate, |
| 614 | activation=self.activation, |
| 615 | ) |