|
|
|
|
|
from typing import Optional |
|
|
|
|
|
import torch |
|
|
from torch import nn, Tensor |
|
|
|
|
|
from ..ops.triton.layer_norm import RMSNorm, layer_norm_fn |
|
|
|
|
|
|
|
|
class Block(nn.Module): |
|
|
def __init__( |
|
|
self, |
|
|
dim, |
|
|
mixer_cls, |
|
|
mlp_cls, |
|
|
norm_cls=nn.LayerNorm, |
|
|
fused_add_norm=False, |
|
|
residual_in_fp32=False, |
|
|
): |
|
|
""" |
|
|
Simple block wrapping a mixer class with LayerNorm/RMSNorm and residual connection" |
|
|
|
|
|
This Block has a slightly different structure compared to a regular |
|
|
prenorm Transformer block. |
|
|
The standard block is: LN -> MHA/MLP -> Add. |
|
|
[Ref: https://arxiv.org/abs/2002.04745] |
|
|
Here we have: Add -> LN -> Mixer, returning both |
|
|
the hidden_states (output of the mixer) and the residual. |
|
|
This is purely for performance reasons, as we can fuse add and LayerNorm. |
|
|
The residual needs to be provided (except for the very first block). |
|
|
""" |
|
|
super().__init__() |
|
|
self.residual_in_fp32 = residual_in_fp32 |
|
|
self.fused_add_norm = fused_add_norm |
|
|
self.norm = norm_cls(dim) |
|
|
self.mixer = mixer_cls(dim) |
|
|
if mlp_cls is not nn.Identity: |
|
|
self.norm2 = norm_cls(dim) |
|
|
self.mlp = mlp_cls(dim) |
|
|
else: |
|
|
self.mlp = None |
|
|
if self.fused_add_norm: |
|
|
assert RMSNorm is not None, "RMSNorm import fails" |
|
|
assert isinstance( |
|
|
self.norm, (nn.LayerNorm, RMSNorm) |
|
|
), "Only LayerNorm and RMSNorm are supported for fused_add_norm" |
|
|
|
|
|
def forward( |
|
|
self, |
|
|
hidden_states: Tensor, |
|
|
residual: Optional[Tensor] = None, |
|
|
inference_params=None, |
|
|
**mixer_kwargs |
|
|
): |
|
|
r"""Pass the input through the encoder layer. |
|
|
|
|
|
Args: |
|
|
hidden_states: the sequence to the encoder layer (required). |
|
|
residual: hidden_states = Mixer(LN(residual)) |
|
|
""" |
|
|
if not self.fused_add_norm: |
|
|
residual = ( |
|
|
(hidden_states + residual) if residual is not None else hidden_states |
|
|
) |
|
|
hidden_states = self.norm(residual.to(dtype=self.norm.weight.dtype)) |
|
|
if self.residual_in_fp32: |
|
|
residual = residual.to(torch.float32) |
|
|
else: |
|
|
hidden_states, residual = layer_norm_fn( |
|
|
hidden_states, |
|
|
self.norm.weight, |
|
|
self.norm.bias, |
|
|
residual=residual, |
|
|
prenorm=True, |
|
|
residual_in_fp32=self.residual_in_fp32, |
|
|
eps=self.norm.eps, |
|
|
is_rms_norm=isinstance(self.norm, RMSNorm), |
|
|
) |
|
|
hidden_states = self.mixer( |
|
|
hidden_states, inference_params=inference_params, **mixer_kwargs |
|
|
) |
|
|
|
|
|
if self.mlp is not None: |
|
|
if not self.fused_add_norm: |
|
|
residual = hidden_states + residual |
|
|
hidden_states = self.norm2(residual.to(dtype=self.norm2.weight.dtype)) |
|
|
if self.residual_in_fp32: |
|
|
residual = residual.to(torch.float32) |
|
|
else: |
|
|
hidden_states, residual = layer_norm_fn( |
|
|
hidden_states, |
|
|
self.norm2.weight, |
|
|
self.norm2.bias, |
|
|
residual=residual, |
|
|
prenorm=True, |
|
|
residual_in_fp32=self.residual_in_fp32, |
|
|
eps=self.norm2.eps, |
|
|
is_rms_norm=isinstance(self.norm2, RMSNorm), |
|
|
) |
|
|
hidden_states = self.mlp(hidden_states) |
|
|
|
|
|
return hidden_states, residual |
|
|
|
|
|
def allocate_inference_cache(self, batch_size, max_seqlen, dtype=None, **kwargs): |
|
|
return self.mixer.allocate_inference_cache( |
|
|
batch_size, max_seqlen, dtype=dtype, **kwargs |
|
|
) |
|
|
|
