V2: 12.97M params, HW optimizer, parallel BBPE, DPO, reasoning, inference, 10 bugs fixed
8594de8 verified | """ | |
| Xavante - embedding_reconfig.py | |
| Responsabilidade: Embedding reconfiguravel (Teorema 12.1/12.2). | |
| Permite expansao dinamica do vocabulario e reconfiguracao de dim do modelo. | |
| """ | |
| from __future__ import annotations | |
| import logging | |
| import math | |
| from typing import Optional | |
| import torch | |
| import torch.nn as nn | |
| logger = logging.getLogger(__name__) | |
| class ReconfigurableEmbedding(nn.Module): | |
| """ | |
| Embedding que suporta: | |
| - Expansao de vocabulario sem reset | |
| - Reconfiguracao de dim (projecao) | |
| - Quantizacao para memoria | |
| """ | |
| def __init__(self, vocab_size: int, d_model: int, padding_idx: Optional[int] = None): | |
| super().__init__() | |
| self.vocab_size = vocab_size | |
| self.d_model = d_model | |
| self.padding_idx = padding_idx | |
| self.weight = nn.Parameter(torch.empty(vocab_size, d_model)) | |
| nn.init.normal_(self.weight, mean=0.0, std=1.0 / math.sqrt(d_model) if False else 0.02) | |
| if padding_idx is not None: | |
| with torch.no_grad(): | |
| self.weight[padding_idx].fill_(0) | |
| def forward(self, idx: torch.Tensor) -> torch.Tensor: | |
| return torch.nn.functional.embedding(idx, self.weight, padding_idx=self.padding_idx) | |
| def expand_vocab(self, new_size: int) -> None: | |
| """Expande o vocabulario preservando pesos antigos.""" | |
| if new_size <= self.vocab_size: | |
| return | |
| old = self.weight.data | |
| new = torch.empty(new_size - self.vocab_size, self.d_model, device=old.device, dtype=old.dtype) | |
| nn.init.normal_(new, mean=0.0, std=0.02) | |
| self.weight = nn.Parameter(torch.cat([old, new], dim=0)) | |
| self.vocab_size = new_size | |
| logger.info("Vocab expandido para %d", new_size) | |
| def reconfigure_dim(self, new_d: int) -> None: | |
| """Reconfigura a dimensao via projecao linear (random init para a parte nova).""" | |
| if new_d == self.d_model: | |
| return | |
| old = self.weight.data # [V, d_model] | |
| if new_d > self.d_model: | |
| pad = torch.empty(old.shape[0], new_d - self.d_model, device=old.device, dtype=old.dtype) | |
| nn.init.normal_(pad, mean=0.0, std=0.02) | |
| new_w = torch.cat([old, pad], dim=1) | |
| else: | |
| new_w = old[:, :new_d] | |
| self.weight = nn.Parameter(new_w) | |
| self.d_model = new_d | |
| logger.info("Dim reconfigurada para %d", new_d) | |
| __all__ = ["ReconfigurableEmbedding"] | |