| """Bidirectional SSM block with explicit pad and segment reset semantics.""" |
|
|
| from __future__ import annotations |
|
|
| import torch |
| from torch import nn |
|
|
| from .attention import FeedForward, RMSNorm |
| from .padding import masked_hidden |
| from .ssm import MaskedScanSSM |
|
|
|
|
| class BidirectionalSSMLayer(nn.Module): |
| def __init__(self, config): |
| super().__init__() |
| self.reset_on_pad = config.ssm_reset_on_pad |
| self.reset_on_segment = config.ssm_reset_on_segment |
| self.fuse = config.ssm_fuse |
| self.norm = RMSNorm(config.hidden_size, config.norm_eps) |
| self.fwd = MaskedScanSSM(config.hidden_size, config.ssm_state_size) |
| self.bwd = MaskedScanSSM(config.hidden_size, config.ssm_state_size) |
| if self.fuse == "concat": |
| self.out = nn.Linear(config.hidden_size * 2, config.hidden_size) |
| elif self.fuse == "concat_product": |
| self.out = nn.Linear(config.hidden_size * 3, config.hidden_size) |
| elif self.fuse == "gated_sum": |
| self.gate = nn.Linear(config.hidden_size, config.hidden_size) |
| self.out = nn.Linear(config.hidden_size, config.hidden_size) |
| else: |
| raise ValueError(f"unknown ssm_fuse: {self.fuse}") |
| self.ffn_norm = RMSNorm(config.hidden_size, config.norm_eps) |
| self.ffn = FeedForward( |
| config.hidden_size, |
| config.intermediate_size, |
| config.hidden_dropout, |
| config.hidden_activation, |
| ) |
|
|
| def forward( |
| self, |
| hidden_states: torch.Tensor, |
| attention_mask: torch.Tensor, |
| segment_ids: torch.Tensor | None = None, |
| ) -> torch.Tensor: |
| residual = hidden_states |
| x = self.norm(hidden_states) |
| y_fwd = self.fwd(x, attention_mask, segment_ids, self.reset_on_pad, self.reset_on_segment) |
| y_bwd_rev = self.bwd( |
| torch.flip(x, dims=[1]), |
| torch.flip(attention_mask, dims=[1]), |
| torch.flip(segment_ids, dims=[1]) if segment_ids is not None else None, |
| self.reset_on_pad, |
| self.reset_on_segment, |
| ) |
| y_bwd = torch.flip(y_bwd_rev, dims=[1]) |
| if self.fuse == "concat": |
| y = self.out(torch.cat([y_fwd, y_bwd], dim=-1)) |
| elif self.fuse == "concat_product": |
| y = self.out(torch.cat([y_fwd, y_bwd, y_fwd * y_bwd], dim=-1)) |
| else: |
| gate = torch.sigmoid(self.gate(x)) |
| y = self.out(gate * y_fwd + (1.0 - gate) * y_bwd) |
| hidden_states = masked_hidden(residual + y, attention_mask) |
| hidden_states = hidden_states + self.ffn(self.ffn_norm(hidden_states)) |
| return masked_hidden(hidden_states, attention_mask) |
|
|
|
|
