| 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: |
| |
| |
| |
| 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, |
| ] |
|
|
