File size: 3,342 Bytes
2719787 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 | """MLM head for the hybrid MLM/CLM objective.
Ported from temp/gpt-bert-main/pretraining/model.py MaskClassifier, but
configurable to match the trunk's customizability:
- norm_type: "layernorm" or "rmsnorm" (reuses the existing field).
- ffn_activation: "gelu", "swiglu", "relu", or "identity" (reuses the
existing field). The head's nonlinearity follows this choice.
- use_bias_ffn: whether the linear layers in the head use bias.
- weight tying: the final linear is tied to the word embedding, matching
gpt-bert's MaskClassifier and our lm_head.
The head is controlled by model config (mlm_head_enabled) so checkpoint
loading and HF export remain strict and honest.
"""
import torch
import torch.nn as nn
from .transformer_components import RMSNorm, FeedForward
def _create_norm(norm_type: str, hidden_dim: int) -> nn.Module:
if norm_type == "layernorm":
return nn.LayerNorm(hidden_dim)
if norm_type == "rmsnorm":
return RMSNorm(hidden_dim)
raise ValueError(f"Unsupported norm_type: {norm_type}")
class MaskClassifier(nn.Module):
"""MLM prediction head with configurable norm and activation.
Architecture (matching gpt-bert's MaskClassifier):
norm -> FeedForward (activation, bias configurable) -> norm -> dropout -> linear (tied)
The final linear is tied to the word embedding weight, matching
gpt-bert's MaskClassifier and our lm_head.
Args:
hidden_dim: Model hidden dimension.
vocab_size: Vocabulary size for the output projection.
norm_type: "layernorm" or "rmsnorm".
ffn_activation: "gelu", "swiglu", "relu", or "identity".
use_bias_ffn: Whether the intermediate linear layers use bias.
dropout: Dropout probability.
word_embedding: The word embedding weight to tie the final linear to.
"""
def __init__(
self,
hidden_dim: int,
vocab_size: int,
norm_type: str,
ffn_activation: str,
use_bias_ffn: bool,
dropout: float,
word_embedding: nn.Parameter,
):
super().__init__()
self.hidden_dim = hidden_dim
self.vocab_size = vocab_size
self.norm_type = norm_type
self.ffn_activation = ffn_activation
self.use_bias_ffn = use_bias_ffn
self.dropout_p = dropout
self.norm1 = _create_norm(norm_type, hidden_dim)
self.feed_forward = FeedForward(
input_dim=hidden_dim,
hidden_dim=hidden_dim,
dropout=dropout,
activation=ffn_activation,
use_bias=use_bias_ffn,
)
self.norm2 = _create_norm(norm_type, hidden_dim)
self.dropout = nn.Dropout(dropout)
self.linear_out = nn.Linear(hidden_dim, vocab_size, bias=False)
self.linear_out.weight = word_embedding
def forward(self, x: torch.Tensor) -> torch.Tensor:
"""Apply the MLM head to contextualized embeddings.
Args:
x: Contextualized embeddings [batch, seq_len, hidden_dim] or
flattened [num_masked, hidden_dim].
Returns:
Logits [batch, seq_len, vocab_size] or [num_masked, vocab_size].
"""
x = self.norm1(x)
x = self.feed_forward(x)
x = self.norm2(x)
x = self.dropout(x)
x = self.linear_out(x)
return x
|