Harry MellorandClaude Opus 5
8c1557a79c
Signed-off-by: Harry Mellor <19981378+hmellor@users.noreply.github.com> Signed-off-by: Harry Mellor <hej.mellor@gmail.com> Co-authored-by: Claude Opus 5 (1M context) <noreply@anthropic.com>
81 lines
2.6 KiB
Python
81 lines
2.6 KiB
Python
# SPDX-License-Identifier: Apache-2.0
|
|
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
|
import torch
|
|
from torch import Tensor
|
|
|
|
from ..op import register_op
|
|
|
|
|
|
@register_op
|
|
def rms_norm(
|
|
x: Tensor, weight: Tensor | None, epsilon: float, variance_size: int | None = None
|
|
) -> Tensor:
|
|
"""Weighted root-mean-square layer normalization."""
|
|
orig_dtype = x.dtype
|
|
x = x.to(torch.float32)
|
|
x_var = x if variance_size is None else x[..., :variance_size]
|
|
variance = x_var.pow(2).mean(dim=-1, keepdim=True)
|
|
x = x * torch.rsqrt(variance + epsilon)
|
|
if weight is not None:
|
|
x = x.to(weight.dtype) * weight
|
|
return x.to(orig_dtype)
|
|
|
|
|
|
@rms_norm.register_input_generator
|
|
def _rms_norm_input_generator(
|
|
num_tokens: int,
|
|
hidden_size: int,
|
|
dtype: torch.dtype,
|
|
epsilon: float = 1e-5,
|
|
device: torch.device | str | None = None,
|
|
) -> tuple:
|
|
x = torch.randn(num_tokens, hidden_size, dtype=dtype, device=device)
|
|
weight = torch.randn(hidden_size, dtype=dtype, device=device)
|
|
return x, weight, epsilon
|
|
|
|
|
|
# Reductions in rms_norm accumulate rounding error at large shapes
|
|
# (e.g. 32768x16384), causing a few elements out of millions to exceed
|
|
# the default float16 tolerance.
|
|
rms_norm.override_tolerance(torch.float16, atol=1e-2, rtol=2e-3)
|
|
|
|
|
|
@register_op(allow_inplace=True)
|
|
def fused_add_rms_norm(
|
|
x: Tensor,
|
|
x_residual: Tensor,
|
|
weight: Tensor | None,
|
|
epsilon: float,
|
|
variance_size: int | None = None,
|
|
) -> tuple[Tensor, Tensor]:
|
|
"""Fused add and weighted root-mean-square layer normalization."""
|
|
orig_dtype = x.dtype
|
|
x = x.to(torch.float32)
|
|
x = x + x_residual.to(torch.float32)
|
|
x_residual = x.to(orig_dtype)
|
|
|
|
x_var = x if variance_size is None else x[..., :variance_size]
|
|
variance = x_var.pow(2).mean(dim=-1, keepdim=True)
|
|
x = x * torch.rsqrt(variance + epsilon)
|
|
if weight is not None:
|
|
x = x.to(weight.dtype) * weight
|
|
return x.to(orig_dtype), x_residual
|
|
|
|
|
|
# fused_add_rms_norm has similar rounding error accumulation as rms_norm
|
|
fused_add_rms_norm.override_tolerance(torch.float16, atol=1e-2, rtol=2e-3)
|
|
|
|
|
|
@fused_add_rms_norm.register_input_generator
|
|
def _fused_add_rms_norm_input_generator(
|
|
num_tokens: int,
|
|
hidden_size: int,
|
|
dtype: torch.dtype,
|
|
epsilon: float = 1e-5,
|
|
device: torch.device | str | None = None,
|
|
) -> tuple:
|
|
x = torch.randn(num_tokens, hidden_size, dtype=dtype, device=device)
|
|
x_residual = torch.randn(num_tokens, hidden_size, dtype=dtype, device=device)
|
|
weight = torch.randn(hidden_size, dtype=dtype, device=device)
|
|
return x, x_residual, weight, epsilon
|