Upload folder using huggingface_hub

config.json

@@ -0,0 +1,29 @@
+{
+  "architectures": [
+    "GPTBertForMaskedLM"
+  ],
+  "attention_probs_dropout_prob": 0.1,
+  "auto_map": {
+    "AutoConfig": "configuration_gpt_bert.GPTBertConfig",
+    "AutoModel": "modeling_gpt_bert.GPTBertModel",
+    "AutoModelForCausalLM": "modeling_gpt_bert.GPTBertForCausalLM",
+    "AutoModelForMaskedLM": "modeling_gpt_bert.GPTBertForMaskedLM"
+  },
+  "bos_token_id": 0,
+  "dtype": "float32",
+  "eos_token_id": 2,
+  "hidden_dropout_prob": 0.1,
+  "hidden_size": 720,
+  "intermediate_size": 2048,
+  "layer_norm_eps": 1e-05,
+  "mask_token_id": 4,
+  "max_position_embeddings": 512,
+  "model_type": "gpt_bert_mntp",
+  "num_attention_heads": 12,
+  "num_hidden_layers": 12,
+  "pad_token_id": 1,
+  "position_bucket_size": 32,
+  "transformers_version": "4.57.6",
+  "unk_token_id": 3,
+  "vocab_size": 16000
+}

configuration_gpt_bert.py

@@ -0,0 +1,51 @@
+from transformers import PretrainedConfig
+
+
+class GPTBertConfig(PretrainedConfig):
+    model_type = "gpt_bert_mntp"
+
+    def __init__(
+        self,
+        vocab_size=16000,
+        hidden_size=720,
+        intermediate_size=2048,
+        num_hidden_layers=12,
+        num_attention_heads=12,
+        max_position_embeddings=512,
+        position_bucket_size=32,
+        hidden_dropout_prob=0.1,
+        attention_probs_dropout_prob=0.1,
+        layer_norm_eps=1e-5,
+        pad_token_id=1,
+        bos_token_id=0,
+        eos_token_id=2,
+        unk_token_id=3,
+        mask_token_id=4,
+        **kwargs,
+    ):
+        super().__init__(
+            pad_token_id=pad_token_id,
+            bos_token_id=bos_token_id,
+            eos_token_id=eos_token_id,
+            **kwargs,
+        )
+        self.vocab_size = vocab_size
+        self.hidden_size = hidden_size
+        self.intermediate_size = intermediate_size
+        self.num_hidden_layers = num_hidden_layers
+        self.num_attention_heads = num_attention_heads
+        self.max_position_embeddings = max_position_embeddings
+        self.position_bucket_size = position_bucket_size
+        self.hidden_dropout_prob = hidden_dropout_prob
+        self.attention_probs_dropout_prob = attention_probs_dropout_prob
+        self.layer_norm_eps = layer_norm_eps
+        self.unk_token_id = unk_token_id
+        self.mask_token_id = mask_token_id
+        self.tie_word_embeddings = True
+        self.architectures = ["GPTBertForMaskedLM"]
+        self.auto_map = {
+            "AutoConfig": "configuration_gpt_bert.GPTBertConfig",
+            "AutoModel": "modeling_gpt_bert.GPTBertModel",
+            "AutoModelForMaskedLM": "modeling_gpt_bert.GPTBertForMaskedLM",
+            "AutoModelForCausalLM": "modeling_gpt_bert.GPTBertForCausalLM",
+        }

model.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:7159b2e4b9f841c6e72ce1c1b1e9e210e34ac8d10a2cd3f72bec62b47b731269
+size 385344608

modeling_gpt_bert.py

@@ -0,0 +1,492 @@
+import math
+from dataclasses import dataclass
+from typing import Optional, Tuple
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from transformers import PreTrainedModel
+from transformers.modeling_outputs import BaseModelOutput, CausalLMOutput, MaskedLMOutput
+from transformers.utils import ModelOutput
+
+try:
+    from .configuration_gpt_bert import GPTBertConfig
+except ImportError:
+    from configuration_gpt_bert import GPTBertConfig
+
+
+@dataclass
+class GPTBertTrainingOutput(ModelOutput):
+    loss: Optional[torch.Tensor] = None
+    logits: Optional[torch.Tensor] = None
+    ce_loss: Optional[torch.Tensor] = None
+    z_loss: Optional[torch.Tensor] = None
+    accuracy: Optional[torch.Tensor] = None
+    num_tokens: Optional[torch.Tensor] = None
+
+
+class GeGLU(nn.Module):
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        value, gate = x.chunk(2, dim=-1)
+        return value * F.gelu(gate, approximate="tanh")
+
+
+def _relative_position_buckets(
+    relative_position: torch.Tensor,
+    bucket_size: int,
+    max_position: int,
+) -> torch.Tensor:
+    sign = torch.sign(relative_position)
+    mid = bucket_size // 2
+    abs_pos = torch.where(
+        (relative_position < mid) & (relative_position > -mid),
+        torch.full_like(relative_position, mid - 1),
+        torch.abs(relative_position).clamp(max=max_position - 1),
+    )
+    safe = abs_pos.clamp(min=mid)
+    log_pos = (
+        torch.ceil(
+            torch.log(safe.float() / mid)
+            / math.log((max_position - 1) / mid)
+            * (mid - 1)
+        ).long()
+        + mid
+    )
+    bucket_pos = torch.where(abs_pos <= mid, relative_position, log_pos * sign)
+    return bucket_size - 1 + bucket_pos.long()
+
+
+class GPTBertEmbeddings(nn.Module):
+    def __init__(self, config: GPTBertConfig):
+        super().__init__()
+        self.hidden_size = config.hidden_size
+        self.word_embeddings = nn.Embedding(
+            config.vocab_size,
+            config.hidden_size,
+            padding_idx=config.pad_token_id,
+        )
+        self.word_norm = nn.LayerNorm(
+            config.hidden_size,
+            eps=config.layer_norm_eps,
+            elementwise_affine=False,
+        )
+        self.relative_embeddings = nn.Parameter(
+            torch.empty(2 * config.position_bucket_size - 1, config.hidden_size)
+        )
+        self.relative_norm = nn.LayerNorm(
+            config.hidden_size,
+            eps=config.layer_norm_eps,
+        )
+        self.dropout = nn.Dropout(config.hidden_dropout_prob)
+        self.reset_parameters()
+
+    def reset_parameters(self):
+        std = math.sqrt(2.0 / (5.0 * self.hidden_size))
+        nn.init.trunc_normal_(
+            self.word_embeddings.weight, mean=0.0, std=std, a=-2 * std, b=2 * std
+        )
+        nn.init.trunc_normal_(
+            self.relative_embeddings, mean=0.0, std=std, a=-2 * std, b=2 * std
+        )
+        if self.word_embeddings.padding_idx is not None:
+            with torch.no_grad():
+                self.word_embeddings.weight[self.word_embeddings.padding_idx].zero_()
+
+    def forward(self, input_ids: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]:
+        x = self.dropout(self.word_norm(self.word_embeddings(input_ids)))
+        rel = self.relative_norm(self.relative_embeddings)
+        return x, rel
+
+
+class GPTBertAttention(nn.Module):
+    def __init__(self, config: GPTBertConfig):
+        super().__init__()
+        if config.hidden_size % config.num_attention_heads != 0:
+            raise ValueError("hidden_size must be divisible by num_attention_heads")
+        self.config = config
+        self.hidden_size = config.hidden_size
+        self.num_heads = config.num_attention_heads
+        self.head_dim = config.hidden_size // config.num_attention_heads
+
+        self.qk_proj = nn.Linear(config.hidden_size, 2 * config.hidden_size)
+        self.vg_proj = nn.Linear(config.hidden_size, 2 * config.hidden_size)
+        self.out_proj = nn.Linear(config.hidden_size, config.hidden_size)
+        self.pre_norm = nn.LayerNorm(
+            config.hidden_size,
+            eps=config.layer_norm_eps,
+            elementwise_affine=False,
+        )
+        self.post_norm = nn.LayerNorm(
+            config.hidden_size,
+            eps=config.layer_norm_eps,
+            elementwise_affine=False,
+        )
+        self.dropout = nn.Dropout(config.attention_probs_dropout_prob)
+        self.out_dropout = nn.Dropout(config.hidden_dropout_prob)
+        self.scale = 1.0 / math.sqrt(3.0 * self.head_dim)
+
+        positions = (
+            torch.arange(config.max_position_embeddings).unsqueeze(1)
+            - torch.arange(config.max_position_embeddings).unsqueeze(0)
+        )
+        buckets = _relative_position_buckets(
+            positions,
+            config.position_bucket_size,
+            config.max_position_embeddings,
+        )
+        self.register_buffer("position_indices", buckets, persistent=False)
+        self.reset_parameters()
+
+    def reset_parameters(self):
+        std = math.sqrt(2.0 / (5.0 * self.hidden_size))
+        for layer in (self.qk_proj, self.vg_proj, self.out_proj):
+            nn.init.trunc_normal_(
+                layer.weight, mean=0.0, std=std, a=-2 * std, b=2 * std
+            )
+            if layer.bias is not None:
+                nn.init.zeros_(layer.bias)
+
+    def forward(
+        self,
+        hidden_states: torch.Tensor,
+        blocked_mask: torch.Tensor,
+        relative_embeddings: torch.Tensor,
+    ) -> torch.Tensor:
+        batch_size, seq_len, _ = hidden_states.shape
+        x = self.pre_norm(hidden_states)
+
+        query, key = self.qk_proj(x).chunk(2, dim=-1)
+        value, gate = self.vg_proj(x).chunk(2, dim=-1)
+        gate = F.gelu(gate)
+
+        query = query.view(batch_size, seq_len, self.num_heads, self.head_dim)
+        key = key.view(batch_size, seq_len, self.num_heads, self.head_dim)
+        value = value.view(batch_size, seq_len, self.num_heads, self.head_dim)
+        query = query.permute(0, 2, 1, 3)
+        key = key.permute(0, 2, 1, 3)
+        value = value.permute(0, 2, 1, 3)
+
+        scores = torch.matmul(query, key.transpose(-1, -2)) * self.scale
+
+        rel_qk = self.qk_proj(self.dropout(relative_embeddings))
+        rel_q, rel_k = rel_qk.chunk(2, dim=-1)
+        indices = self.position_indices[:seq_len, :seq_len]
+        rel_q = F.embedding(indices, rel_q).view(
+            seq_len, seq_len, self.num_heads, self.head_dim
+        )
+        rel_k = F.embedding(indices, rel_k).view(
+            seq_len, seq_len, self.num_heads, self.head_dim
+        )
+
+        scores = scores + torch.einsum(
+            "bhqd,qkhd->bhqk", query, rel_k * self.scale
+        )
+        scores = scores + torch.einsum(
+            "bhkd,qkhd->bhqk", key * self.scale, rel_q
+        )
+
+        scores = scores.masked_fill(blocked_mask, torch.finfo(scores.dtype).min)
+        probs = torch.softmax(scores.float(), dim=-1).to(scores.dtype)
+        probs = self.dropout(probs)
+
+        context = torch.matmul(probs, value)
+        context = context.permute(0, 2, 1, 3).contiguous().view(
+            batch_size, seq_len, self.hidden_size
+        )
+        context = context * gate
+        context = self.post_norm(context)
+        context = self.out_proj(context)
+        return self.out_dropout(context)
+
+
+class GPTBertFeedForward(nn.Module):
+    def __init__(self, config: GPTBertConfig, layer_index: int):
+        super().__init__()
+        self.norm1 = nn.LayerNorm(
+            config.hidden_size,
+            eps=config.layer_norm_eps,
+            elementwise_affine=False,
+        )
+        self.fc1 = nn.Linear(
+            config.hidden_size,
+            2 * config.intermediate_size,
+            bias=False,
+        )
+        self.act = GeGLU()
+        self.norm2 = nn.LayerNorm(
+            config.intermediate_size,
+            eps=config.layer_norm_eps,
+            elementwise_affine=False,
+        )
+        self.fc2 = nn.Linear(
+            config.intermediate_size,
+            config.hidden_size,
+            bias=False,
+        )
+        self.dropout = nn.Dropout(config.hidden_dropout_prob)
+        self.reset_parameters(layer_index)
+
+    def reset_parameters(self, layer_index: int):
+        std = math.sqrt(2.0 / (5.0 * self.fc2.out_features))
+        nn.init.trunc_normal_(self.fc1.weight, mean=0.0, std=std, a=-2 * std, b=2 * std)
+        nn.init.trunc_normal_(self.fc2.weight, mean=0.0, std=std, a=-2 * std, b=2 * std)
+        scale = math.sqrt(1.0 / (2.0 * (1 + layer_index)))
+        with torch.no_grad():
+            self.fc1.weight.mul_(scale)
+            self.fc2.weight.mul_(scale)
+
+    def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
+        x = self.norm1(hidden_states)
+        x = self.fc1(x)
+        x = self.act(x)
+        x = self.norm2(x)
+        x = self.fc2(x)
+        return self.dropout(x)
+
+
+class GPTBertLayer(nn.Module):
+    def __init__(self, config: GPTBertConfig, layer_index: int):
+        super().__init__()
+        self.attention = GPTBertAttention(config)
+        self.ffn = GPTBertFeedForward(config, layer_index)
+
+    def forward(
+        self,
+        hidden_states: torch.Tensor,
+        blocked_mask: torch.Tensor,
+        relative_embeddings: torch.Tensor,
+    ) -> torch.Tensor:
+        hidden_states = hidden_states + self.attention(
+            hidden_states, blocked_mask, relative_embeddings
+        )
+        hidden_states = hidden_states + self.ffn(hidden_states)
+        return hidden_states
+
+
+class GPTBertPreTrainedModel(PreTrainedModel):
+    config_class = GPTBertConfig
+    base_model_prefix = "gpt_bert"
+    supports_gradient_checkpointing = False
+
+    def _init_weights(self, module):
+        # Components initialize themselves to match the LTG/GPT-BERT recipe.
+        return
+
+
+class GPTBertModel(GPTBertPreTrainedModel):
+    def __init__(self, config: GPTBertConfig):
+        super().__init__(config)
+        self.embeddings = GPTBertEmbeddings(config)
+        self.layers = nn.ModuleList(
+            [GPTBertLayer(config, i) for i in range(config.num_hidden_layers)]
+        )
+        self.post_init()
+
+    def get_input_embeddings(self):
+        return self.embeddings.word_embeddings
+
+    def set_input_embeddings(self, value):
+        self.embeddings.word_embeddings = value
+
+    def _build_blocked_mask(
+        self,
+        input_ids: torch.Tensor,
+        attention_mask: Optional[torch.Tensor],
+        is_causal: bool,
+    ) -> torch.Tensor:
+        batch_size, seq_len = input_ids.shape
+        if attention_mask is None:
+            attention_mask = torch.ones(
+                batch_size, seq_len, device=input_ids.device, dtype=torch.long
+            )
+        key_padding = attention_mask.eq(0)[:, None, None, :]
+        if is_causal:
+            causal = torch.ones(
+                seq_len, seq_len, device=input_ids.device, dtype=torch.bool
+            ).triu(diagonal=1)[None, None, :, :]
+            return key_padding | causal
+        return key_padding.expand(batch_size, 1, seq_len, seq_len)
+
+    def forward(
+        self,
+        input_ids: torch.Tensor,
+        attention_mask: Optional[torch.Tensor] = None,
+        is_causal: bool = False,
+        return_dict: bool = True,
+        **kwargs,
+    ):
+        blocked_mask = self._build_blocked_mask(
+            input_ids, attention_mask, is_causal=is_causal
+        )
+        hidden_states, relative_embeddings = self.embeddings(input_ids)
+        for layer in self.layers:
+            hidden_states = layer(
+                hidden_states, blocked_mask, relative_embeddings
+            )
+        if not return_dict:
+            return (hidden_states,)
+        return BaseModelOutput(last_hidden_state=hidden_states)
+
+
+class GPTBertLMHead(nn.Module):
+    def __init__(self, config: GPTBertConfig, embedding_weight: nn.Parameter):
+        super().__init__()
+        self.dense = nn.Linear(config.hidden_size, config.hidden_size)
+        self.activation = nn.GELU()
+        self.norm = nn.LayerNorm(
+            config.hidden_size,
+            eps=config.layer_norm_eps,
+            elementwise_affine=False,
+        )
+        self.decoder = nn.Linear(config.hidden_size, config.vocab_size, bias=False)
+        self.decoder.weight = embedding_weight
+        self.bias = nn.Parameter(torch.zeros(config.vocab_size))
+        self.reset_parameters(config.hidden_size)
+
+    def reset_parameters(self, hidden_size: int):
+        std = math.sqrt(2.0 / (5.0 * hidden_size))
+        nn.init.trunc_normal_(
+            self.dense.weight, mean=0.0, std=std, a=-2 * std, b=2 * std
+        )
+        nn.init.zeros_(self.dense.bias)
+
+    def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
+        x = self.dense(hidden_states)
+        x = self.activation(x)
+        x = self.norm(x)
+        return self.decoder(x) + self.bias
+
+
+class GPTBertForMaskedLM(GPTBertPreTrainedModel):
+    _tied_weights_keys = ["lm_head.decoder.weight"]
+
+    def __init__(self, config: GPTBertConfig):
+        super().__init__(config)
+        self.gpt_bert = GPTBertModel(config)
+        self.lm_head = GPTBertLMHead(
+            config, self.gpt_bert.embeddings.word_embeddings.weight
+        )
+        self.post_init()
+
+    def get_input_embeddings(self):
+        return self.gpt_bert.get_input_embeddings()
+
+    def set_input_embeddings(self, value):
+        self.gpt_bert.set_input_embeddings(value)
+        self.lm_head.decoder.weight = value.weight
+
+    def get_output_embeddings(self):
+        return self.lm_head.decoder
+
+    def set_output_embeddings(self, new_embeddings):
+        self.lm_head.decoder = new_embeddings
+
+    def _selected_stats(
+        self,
+        logits: torch.Tensor,
+        labels: torch.Tensor,
+    ):
+        ce_loss = F.cross_entropy(logits.float(), labels)
+        z_loss = torch.logsumexp(logits.float(), dim=-1).pow(2).mean()
+        accuracy = (logits.argmax(dim=-1) == labels).float().mean()
+        return ce_loss, z_loss, accuracy
+
+    def forward(
+        self,
+        input_ids: torch.Tensor,
+        attention_mask: Optional[torch.Tensor] = None,
+        labels: Optional[torch.Tensor] = None,
+        mode: str = "mntp",
+        z_loss_weight: float = 0.0,
+        return_dict: bool = True,
+        **kwargs,
+    ):
+        if mode not in {"mntp", "masked", "causal"}:
+            raise ValueError(f"Unsupported mode: {mode}")
+        is_causal = mode == "causal"
+        hidden = self.gpt_bert(
+            input_ids=input_ids,
+            attention_mask=attention_mask,
+            is_causal=is_causal,
+            return_dict=True,
+        ).last_hidden_state
+
+        if labels is not None and mode in {"mntp", "masked"}:
+            # MNTP: the hidden state at position i-1 predicts a masked token at i.
+            valid = labels[:, 1:].ne(-100)
+            selected_hidden = hidden[:, :-1][valid]
+            selected_labels = labels[:, 1:][valid]
+            if selected_labels.numel() == 0:
+                raise RuntimeError("MNTP batch contains no prediction targets")
+            selected_logits = self.lm_head(selected_hidden)
+            ce_loss, z_loss, accuracy = self._selected_stats(
+                selected_logits, selected_labels
+            )
+            loss = ce_loss + z_loss_weight * z_loss
+            return GPTBertTrainingOutput(
+                loss=loss,
+                logits=None,
+                ce_loss=ce_loss.detach(),
+                z_loss=z_loss.detach(),
+                accuracy=accuracy.detach(),
+                num_tokens=torch.tensor(
+                    selected_labels.numel(), device=input_ids.device
+                ),
+            )
+
+        if labels is not None and mode == "causal":
+            # Training input is [BOS] + tokens[:-1]; labels are tokens.
+            logits = self.lm_head(hidden)
+            valid = labels.ne(-100)
+            selected_logits = logits[valid]
+            selected_labels = labels[valid]
+            ce_loss, z_loss, accuracy = self._selected_stats(
+                selected_logits, selected_labels
+            )
+            loss = ce_loss + z_loss_weight * z_loss
+            return GPTBertTrainingOutput(
+                loss=loss,
+                logits=None,
+                ce_loss=ce_loss.detach(),
+                z_loss=z_loss.detach(),
+                accuracy=accuracy.detach(),
+                num_tokens=torch.tensor(
+                    selected_labels.numel(), device=input_ids.device
+                ),
+            )
+
+        raw_logits = self.lm_head(hidden)
+        # The BabyLM MNTP backend selects target_position - 1 itself.
+        if not return_dict:
+            return (raw_logits,)
+        return MaskedLMOutput(logits=raw_logits)
+
+
+class GPTBertForCausalLM(GPTBertForMaskedLM):
+    def forward(
+        self,
+        input_ids: torch.Tensor,
+        attention_mask: Optional[torch.Tensor] = None,
+        labels: Optional[torch.Tensor] = None,
+        return_dict: bool = True,
+        **kwargs,
+    ):
+        hidden = self.gpt_bert(
+            input_ids=input_ids,
+            attention_mask=attention_mask,
+            is_causal=True,
+            return_dict=True,
+        ).last_hidden_state
+        logits = self.lm_head(hidden)
+        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, shift_logits.size(-1)).float(),
+                shift_labels.view(-1),
+                ignore_index=-100,
+            )
+        if not return_dict:
+            return (loss, logits) if loss is not None else (logits,)
+        return CausalLMOutput(loss=loss, logits=logits)

special_tokens_map.json

@@ -0,0 +1,37 @@
+{
+  "bos_token": {
+    "content": "<s>",
+    "lstrip": false,
+    "normalized": false,
+    "rstrip": false,
+    "single_word": false
+  },
+  "eos_token": {
+    "content": "</s>",
+    "lstrip": false,
+    "normalized": false,
+    "rstrip": false,
+    "single_word": false
+  },
+  "mask_token": {
+    "content": "<mask>",
+    "lstrip": false,
+    "normalized": false,
+    "rstrip": false,
+    "single_word": false
+  },
+  "pad_token": {
+    "content": "<pad>",
+    "lstrip": false,
+    "normalized": false,
+    "rstrip": false,
+    "single_word": false
+  },
+  "unk_token": {
+    "content": "<unk>",
+    "lstrip": false,
+    "normalized": false,
+    "rstrip": false,
+    "single_word": false
+  }
+}

tokenizer_config.json

@@ -0,0 +1,54 @@
+{
+  "added_tokens_decoder": {
+    "0": {
+      "content": "<s>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "1": {
+      "content": "<pad>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "2": {
+      "content": "</s>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "3": {
+      "content": "<unk>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "4": {
+      "content": "<mask>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    }
+  },
+  "backend": "tokenizers",
+  "bos_token": "<s>",
+  "clean_up_tokenization_spaces": false,
+  "eos_token": "</s>",
+  "extra_special_tokens": {},
+  "mask_token": "<mask>",
+  "model_max_length": 512,
+  "pad_token": "<pad>",
+  "tokenizer_class": "PreTrainedTokenizerFast",
+  "unk_token": "<unk>"
+}