| import math |
|
|
| import torch |
| import torch.nn as nn |
| import torch.nn.functional as F |
|
|
| from src.model.config import BoccaccioConfig |
| from src.model.layers import RMSNorm, SwiGLU |
| from src.model.attention import GroupedQueryAttention |
|
|
|
|
| class BoccaccioBlock(nn.Module): |
|
|
| def __init__(self, config: BoccaccioConfig) -> None: |
| super().__init__() |
| self.input_layernorm = RMSNorm(config.hidden_size, config.norm_eps) |
| self.self_attn = GroupedQueryAttention(config) |
| self.post_attention_layernorm = RMSNorm(config.hidden_size, config.norm_eps) |
| self.mlp = SwiGLU(config.hidden_size, config.intermediate_size) |
|
|
| def forward( |
| self, |
| hidden_states: torch.Tensor, |
| position_ids: torch.Tensor, |
| attention_mask: torch.Tensor | None = None, |
| ) -> torch.Tensor: |
| residual = hidden_states |
| hidden_states = self.input_layernorm(hidden_states) |
| hidden_states = self.self_attn(hidden_states, position_ids, attention_mask) |
| hidden_states = residual + hidden_states |
|
|
| residual = hidden_states |
| hidden_states = self.post_attention_layernorm(hidden_states) |
| hidden_states = self.mlp(hidden_states) |
| hidden_states = residual + hidden_states |
|
|
| return hidden_states |
|
|
|
|
| class BoccaccioModel(nn.Module): |
|
|
| def __init__(self, config: BoccaccioConfig) -> None: |
| super().__init__() |
| self.config = config |
| self.embed_tokens = nn.Embedding(config.vocab_size, config.hidden_size) |
| self.layers = nn.ModuleList( |
| [BoccaccioBlock(config) for _ in range(config.num_layers)] |
| ) |
| self.norm = RMSNorm(config.hidden_size, config.norm_eps) |
| self._init_weights() |
|
|
| def forward( |
| self, |
| input_ids: torch.Tensor, |
| position_ids: torch.Tensor | None = None, |
| attention_mask: torch.Tensor | None = None, |
| ) -> torch.Tensor: |
| batch_size, seq_len = input_ids.shape |
| hidden_states = self.embed_tokens(input_ids) |
|
|
| if position_ids is None: |
| position_ids = ( |
| torch.arange(seq_len, device=input_ids.device) |
| .unsqueeze(0) |
| .expand(batch_size, -1) |
| ) |
|
|
| for layer in self.layers: |
| hidden_states = layer(hidden_states, position_ids, attention_mask) |
|
|
| hidden_states = self.norm(hidden_states) |
| return hidden_states |
|
|
| def _init_weights(self) -> None: |
| std = self.config.initializer_range |
| residual_std = std / math.sqrt(2 * self.config.num_layers) |
|
|
| for module in self.modules(): |
| if isinstance(module, nn.Linear): |
| name = None |
| for n, m in self.named_modules(): |
| if m is module: |
| name = n |
| break |
| if name is not None and ( |
| name.endswith(".o_proj") or name.endswith(".down_proj") |
| ): |
| nn.init.normal_(module.weight, mean=0.0, std=residual_std) |
| else: |
| nn.init.normal_(module.weight, mean=0.0, std=std) |
| if module.bias is not None: |
| nn.init.zeros_(module.bias) |
| elif isinstance(module, nn.Embedding): |
| nn.init.normal_(module.weight, mean=0.0, std=std) |
|
|
|
|
| class BoccaccioForCausalLM(nn.Module): |
|
|
| def __init__(self, config: BoccaccioConfig) -> None: |
| super().__init__() |
| self.config = config |
| self.model = BoccaccioModel(config) |
| self.lm_head = nn.Linear(config.hidden_size, config.vocab_size, bias=False) |
|
|
| if config.tie_word_embeddings: |
| self.lm_head.weight = self.model.embed_tokens.weight |
|
|
| def forward( |
| self, |
| input_ids: torch.Tensor, |
| position_ids: torch.Tensor | None = None, |
| attention_mask: torch.Tensor | None = None, |
| labels: torch.Tensor | None = None, |
| ) -> dict[str, torch.Tensor | None]: |
| hidden_states = self.model(input_ids, position_ids, attention_mask) |
| logits = self.lm_head(hidden_states) |
|
|
| loss = None |
| if labels is not None: |
| shift_logits = logits[..., :-1, :].contiguous() |
| shift_labels = labels[..., 1:].contiguous() |
| loss = F.cross_entropy( |
| shift_logits.view(-1, self.config.vocab_size), |
| shift_labels.view(-1), |
| ignore_index=-100, |
| ) |
|
|
| return {"loss": loss, "logits": logits} |
|
|
| def count_parameters(self) -> int: |
| return sum(p.numel() for p in self.parameters() if p.requires_grad) |
|
|
