File size: 3,478 Bytes
02720c5 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 | from __future__ import annotations
import torch
import torch.nn.functional as F
from torch import nn
from torch.nn.attention import SDPBackend
try:
from fla.modules.layernorm import rms_norm as fused_rms_norm
except ImportError:
fused_rms_norm = None
class RMSNorm(nn.Module):
"""Zero-centered RMSNorm as used by Qwen3-Next.
On CUDA this runs FLA's single-pass Triton kernel with the ``1 + weight``
scale kept in FP32, which matches the eager FP32 math to rounding.
"""
def __init__(self, hidden_size: int, eps: float) -> None:
super().__init__()
self.weight = nn.Parameter(torch.zeros(hidden_size))
self.eps = eps
def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
scale = 1.0 + self.weight.float()
if fused_rms_norm is not None and hidden_states.is_cuda:
return fused_rms_norm(hidden_states, scale, None, eps=self.eps)
input_dtype = hidden_states.dtype
normalized = hidden_states.float()
normalized = normalized * torch.rsqrt(
normalized.square().mean(dim=-1, keepdim=True) + self.eps
)
return (normalized * scale).to(dtype=input_dtype)
def linear_cross_entropy_reference(
states: torch.Tensor,
weight: torch.Tensor,
targets: torch.Tensor,
*,
z_loss_coefficient: float = 0.0,
ignore_index: int = -100,
) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
"""Eager ``(objective, cross_entropy, z_loss)`` over materialized logits."""
logits = F.linear(states, weight)
cross_entropy = F.cross_entropy(logits, targets, ignore_index=ignore_index)
reduced_dtype = torch.promote_types(logits.dtype, torch.float32)
if z_loss_coefficient > 0.0:
# Keep logits in their training dtype for logsumexp. Casting the full
# [tokens, vocab] tensor to fp32 materializes it again in the autograd
# graph; upcast only the reduced [tokens] partition.
log_partition = torch.logsumexp(logits, dim=-1).to(reduced_dtype)
z_loss = log_partition.square()[targets.ne(ignore_index)].mean()
else:
z_loss = cross_entropy.new_zeros((), dtype=reduced_dtype)
objective = cross_entropy + z_loss_coefficient * z_loss
return objective, cross_entropy.detach(), z_loss.detach()
def linear_token_logprobs_reference(
states: torch.Tensor,
weight: torch.Tensor,
targets: torch.Tensor,
*,
ignore_index: int = -100,
) -> torch.Tensor:
"""Eager fp32 ``log_softmax(states @ weight.T)[row, target]``; ignored rows are 0."""
valid = targets.ne(ignore_index)
logprobs = F.linear(states, weight).float().log_softmax(dim=-1)
picked = logprobs.gather(1, targets.masked_fill(~valid, 0)[:, None]).squeeze(1)
return picked.masked_fill(~valid, 0.0)
def softcap(values: torch.Tensor, cap: float) -> torch.Tensor:
return cap * torch.tanh(values / cap)
def situ_glu_reference(gate_up: torch.Tensor, gate_cap: float, up_cap: float) -> torch.Tensor:
"""``softcap(gate) * sigmoid(gate) * softcap(up)`` over ``[..., gate | up]``."""
gate, up = gate_up.chunk(2, dim=-1)
return softcap(gate, gate_cap) * torch.sigmoid(gate) * softcap(up, up_cap)
def attention_sdpa_backends(device: torch.device) -> list[SDPBackend]:
if device.type != "cuda":
return [SDPBackend.MATH]
return [
SDPBackend.CUDNN_ATTENTION,
SDPBackend.FLASH_ATTENTION,
SDPBackend.EFFICIENT_ATTENTION,
]
|