Initial gLM2 HF port

.gitattributes

@@ -0,0 +1,35 @@
+*.7z filter=lfs diff=lfs merge=lfs -text
+*.arrow filter=lfs diff=lfs merge=lfs -text
+*.bin filter=lfs diff=lfs merge=lfs -text
+*.bz2 filter=lfs diff=lfs merge=lfs -text
+*.ckpt filter=lfs diff=lfs merge=lfs -text
+*.ftz filter=lfs diff=lfs merge=lfs -text
+*.gz filter=lfs diff=lfs merge=lfs -text
+*.h5 filter=lfs diff=lfs merge=lfs -text
+*.joblib filter=lfs diff=lfs merge=lfs -text
+*.lfs.* filter=lfs diff=lfs merge=lfs -text
+*.mlmodel filter=lfs diff=lfs merge=lfs -text
+*.model filter=lfs diff=lfs merge=lfs -text
+*.msgpack filter=lfs diff=lfs merge=lfs -text
+*.npy filter=lfs diff=lfs merge=lfs -text
+*.npz filter=lfs diff=lfs merge=lfs -text
+*.onnx filter=lfs diff=lfs merge=lfs -text
+*.ot filter=lfs diff=lfs merge=lfs -text
+*.parquet filter=lfs diff=lfs merge=lfs -text
+*.pb filter=lfs diff=lfs merge=lfs -text
+*.pickle filter=lfs diff=lfs merge=lfs -text
+*.pkl filter=lfs diff=lfs merge=lfs -text
+*.pt filter=lfs diff=lfs merge=lfs -text
+*.pth filter=lfs diff=lfs merge=lfs -text
+*.rar filter=lfs diff=lfs merge=lfs -text
+*.safetensors filter=lfs diff=lfs merge=lfs -text
+saved_model/**/* filter=lfs diff=lfs merge=lfs -text
+*.tar.* filter=lfs diff=lfs merge=lfs -text
+*.tar filter=lfs diff=lfs merge=lfs -text
+*.tflite filter=lfs diff=lfs merge=lfs -text
+*.tgz filter=lfs diff=lfs merge=lfs -text
+*.wasm filter=lfs diff=lfs merge=lfs -text
+*.xz filter=lfs diff=lfs merge=lfs -text
+*.zip filter=lfs diff=lfs merge=lfs -text
+*.zst filter=lfs diff=lfs merge=lfs -text
+*tfevents* filter=lfs diff=lfs merge=lfs -text

README.md

@@ -0,0 +1,172 @@
+---
+library_name: transformers
+tags:
+- biology
+- genomics
+- protein
+- dna
+- language-model
+license: apache-2.0
+datasets:
+- tattabio/OMG
+---
+
+# gLM-150M
+
+Minimal HuggingFace port of the **150M** parameter variant of
+[gLM2](https://huggingface.co/tattabio/gLM2_150M) -- a mixed-modality genomic
+language model that encodes a genomic scaffold using both amino-acid and DNA
+tokens. Pretrained with masked language modeling on the
+[OMG dataset](https://huggingface.co/datasets/tattabio/OMG).
+
+## Architecture
+
+| Parameter | Value |
+|---|---|
+| Layers | 30 |
+| Attention heads | 10 |
+| Embedding dimension | 640 |
+| FFN hidden dimension | 1792 (SwiGLU, multiple_of=256) |
+| Vocabulary size | 37 |
+| Positional encoding | RoPE (base=10000, non-interleaved) |
+| Normalization | RMSNorm |
+| Architecture | Pre-LN Transformer with SwiGLU FFN |
+| Max sequence length | 4096 |
+
+**Vocabulary:** `<cls>`, `<pad>`, `<eos>`, `<unk>`, the 26 IUPAC amino-acid
+letters (`L A G V S E R T I D P K Q N F Y M H W C X B U Z O`, uppercase),
+the 4 DNA nucleotides (`a t c g`, lowercase), strand markers `<+>` / `<->`,
+and `<mask>` / `<sep>`. Amino-acid and nucleotide tokens share the alphabet
+by case (uppercase = amino acid, lowercase = nucleotide).
+
+## Pretraining
+
+- **Objective:** Masked language modeling (30% mask rate)
+- **Data:** [OMG dataset](https://huggingface.co/datasets/tattabio/OMG) (open
+  metagenomic corpus, semantically-deduplicated)
+- **Pretraining tokens:** 315B (bfloat16, context length 4096)
+- **Source checkpoint:** `tattabio/gLM2_150M`
+
+## Parity Verification
+
+All 31 representation levels (embedding + 30 transformer blocks) verified to
+be bit-exact (max abs diff = 0.00) against the original `tattabio/gLM2_150M`
+weights with `attn_implementation="sdpa"`. The added eager and
+`flash_attention_2` backends agree within fp32 kernel drift (atol = 1e-3) and
+bf16 cosine similarity >= 0.999 respectively. Verified on GPU with PyTorch
+2.7 / CUDA 12.
+
+## Related Models
+
+See the full [gLM2 collection](https://huggingface.co/collections/Taykhoom/glm2-6a2e19be671ba44c163c617f).
+
+| Model | Parameters | Notes |
+|---|---|---|
+| **[gLM-150M](https://huggingface.co/Taykhoom/gLM-150M)** | 150M | This model |
+| [gLM-650M](https://huggingface.co/Taykhoom/gLM-650M) | 650M | Larger variant |
+
+## Usage
+
+### Embedding generation
+
+```python
+import torch
+from transformers import AutoTokenizer, AutoModel
+
+tokenizer = AutoTokenizer.from_pretrained("Taykhoom/gLM-150M", trust_remote_code=True)
+model = AutoModel.from_pretrained("Taykhoom/gLM-150M", trust_remote_code=True)
+model.eval()
+
+# Canonical gLM2 input: amino acids (uppercase) + DNA (lowercase) + strand markers.
+sequence = (
+    "<+>MALTKVEKRNRIKRRVRGKISGTQASPRLSVYKSNK"
+    "<+>aatttaaggaa"
+    "<->MLGIDNIERVKPGGLELVDRLVAVNRVTKVTKGGRAFGFSAIVVVGNED"
+)
+enc = tokenizer([sequence], return_tensors="pt")
+
+with torch.no_grad():
+    out = model(**enc)
+
+cls_emb   = out.last_hidden_state[:, 0, :]   # (batch, 640) -- CLS token
+token_emb = out.last_hidden_state             # (batch, seq_len, 640)
+
+# Intermediate layers
+out_all = model(**enc, output_hidden_states=True)
+layer15_emb = out_all.hidden_states[15]       # after block 15
+```
+
+The tokenizer also accepts plain DNA strings (no strand marker) and
+auto-prepares them by lowercasing, replacing `U`/`u` with `t`, and prepending
+`<+>`. The three calls below produce identical token sequences:
+
+```python
+tokenizer(["ATCGATCG", "atcgatcg", "AUCGAUCG"], return_tensors="pt")
+```
+
+### MLM logits
+
+```python
+from transformers import AutoTokenizer, AutoModelForMaskedLM
+
+tokenizer = AutoTokenizer.from_pretrained("Taykhoom/gLM-150M", trust_remote_code=True)
+model = AutoModelForMaskedLM.from_pretrained("Taykhoom/gLM-150M", trust_remote_code=True)
+model.eval()
+
+enc = tokenizer(["<+>MA<mask>K"], return_tensors="pt")
+with torch.no_grad():
+    logits = model(**enc).logits   # (1, seq_len, 37)
+```
+
+### Faster attention backends
+
+```python
+# SDPA (PyTorch 2.0+, default upstream backend) -- recommended for fp32
+model = AutoModel.from_pretrained("Taykhoom/gLM-150M", trust_remote_code=True,
+                                  attn_implementation="sdpa")
+
+# Flash Attention 2 (requires flash-attn package) -- fastest on long sequences
+model = AutoModel.from_pretrained("Taykhoom/gLM-150M", trust_remote_code=True,
+                                  attn_implementation="flash_attention_2",
+                                  dtype=torch.bfloat16)
+```
+
+### Fine-tuning
+
+Standard HF conventions. For sequence-level tasks, pool over non-padding
+positions or use the CLS token embedding as input to a prediction head.
+
+## Implementation Notes
+
+The original gLM2 implementation uses PyTorch SDPA as the only attention
+backend. This HF port adds eager and `flash_attention_2` as separate
+implementations selectable via `attn_implementation`, with eager falling back
+automatically when `output_attentions=True` is requested.
+
+The eager kernel computes the QK matmul and softmax in fp32 even when the
+model is loaded in bf16, matching the numerical behaviour of SDPA and
+`flash_attention_2` in mixed precision.
+
+## Citation
+
+```bibtex
+@article{cornman2024_glm2,
+  title   = {The {OMG} dataset: An Open MetaGenomic corpus for mixed-modality genomic language modeling},
+  author  = {Cornman, Andre and West-Roberts, Jacob and Camargo, Antonio Pedro and Roux, Simon and Beracochea, Martin and Mirdita, Milot and Ovchinnikov, Sergey and Hwang, Yunha},
+  journal = {bioRxiv},
+  year    = {2024},
+  doi     = {10.1101/2024.08.14.607850}
+}
+```
+
+## Credits
+
+Original model and code by Cornman et al. (Tatta Bio). Source:
+[GitHub](https://github.com/TattaBio/gLM2),
+[`tattabio/gLM2_150M` on the Hub](https://huggingface.co/tattabio/gLM2_150M).
+The HF conversion code was authored primarily by [Claude Code](https://claude.ai/code)
+and reviewed manually by Taykhoom Dalal.
+
+## License
+
+Apache 2.0, following the original repository.

config.json

@@ -0,0 +1,27 @@
+{
+  "architectures": [
+    "gLM2ForMaskedLM"
+  ],
+  "auto_map": {
+    "AutoConfig": "configuration_glm2.gLM2Config",
+    "AutoModel": "modeling_glm2.gLM2Model",
+    "AutoModelForMaskedLM": "modeling_glm2.gLM2ForMaskedLM"
+  },
+  "cls_token_id": 0,
+  "depth": 30,
+  "dim": 640,
+  "dtype": "float32",
+  "eos_token_id": 2,
+  "ffn_dim_multiplier": null,
+  "heads": 10,
+  "initializer_range": 0.02,
+  "mask_token_id": 35,
+  "max_position_embeddings": 4096,
+  "model_type": "gLM2",
+  "norm_eps": 1e-05,
+  "pad_token_id": 1,
+  "swiglu_multiple_of": 256,
+  "tie_word_embeddings": false,
+  "transformers_version": "4.57.6",
+  "vocab_size": 37
+}

configuration_glm2.py

@@ -0,0 +1,60 @@
+"""gLM2 model configuration."""
+
+from typing import Optional
+
+from transformers import PretrainedConfig
+from transformers.utils import logging
+
+logger = logging.get_logger(__name__)
+
+
+class gLM2Config(PretrainedConfig):
+    """Configuration class for gLM2.
+
+    Mirrors the hyperparameters in the official `tattabio/gLM2_*` checkpoints
+    plus the standard HuggingFace `PretrainedConfig` knobs.
+    """
+
+    model_type = "gLM2"
+
+    def __init__(
+        self,
+        dim: int = 640,
+        depth: int = 30,
+        heads: int = 10,
+        vocab_size: int = 37,
+        swiglu_multiple_of: int = 256,
+        ffn_dim_multiplier: Optional[float] = None,
+        norm_eps: float = 1e-5,
+        pad_token_id: int = 1,
+        cls_token_id: int = 0,
+        eos_token_id: int = 2,
+        mask_token_id: int = 35,
+        max_position_embeddings: int = 4096,
+        initializer_range: float = 0.02,
+        tie_word_embeddings: bool = False,
+        **kwargs,
+    ):
+        super().__init__(
+            pad_token_id=pad_token_id,
+            tie_word_embeddings=tie_word_embeddings,
+            **kwargs,
+        )
+        self.dim = dim
+        self.depth = depth
+        self.heads = heads
+        self.vocab_size = vocab_size
+        self.swiglu_multiple_of = swiglu_multiple_of
+        self.ffn_dim_multiplier = ffn_dim_multiplier
+        self.norm_eps = norm_eps
+        self.cls_token_id = cls_token_id
+        self.eos_token_id = eos_token_id
+        self.mask_token_id = mask_token_id
+        self.max_position_embeddings = max_position_embeddings
+        self.initializer_range = initializer_range
+
+        self.auto_map = {
+            "AutoConfig": "configuration_glm2.gLM2Config",
+            "AutoModel": "modeling_glm2.gLM2Model",
+            "AutoModelForMaskedLM": "modeling_glm2.gLM2ForMaskedLM",
+        }

glm_tokenizer.py

@@ -0,0 +1,132 @@
+"""Tokenizer for gLM2.
+
+Wraps a BPE-style fast tokenizer with the upstream `tattabio/gLM2_*` vocabulary
+(amino acids, nucleotides, strand markers). On top of the upstream behaviour
+this version adds *automatic DNA preparation* so callers can pass plain DNA
+sequences (e.g. mRNABench-style input) without manually lowercasing or adding
+strand markers.
+"""
+
+from __future__ import annotations
+
+import re
+from typing import List, Optional, Sequence, Union
+
+from tokenizers import Tokenizer
+from tokenizers.models import BPE
+from transformers import PreTrainedTokenizerFast
+
+
+_NUC_CHARS = set("acgtuACGTU")
+_DNA_PREFIX_RE = re.compile(r"^\s*(<\+>|<->)")
+
+
+def _is_pure_dna(seq: str) -> bool:
+    """True if `seq` only contains DNA/RNA characters (ATCGU, any case)."""
+    return len(seq) > 0 and all(c in _NUC_CHARS for c in seq)
+
+
+def prepare_dna_sequence(seq: str, strand: str = "+") -> str:
+    """Normalize a plain DNA/RNA sequence for gLM2.
+
+    - Lower-cases nucleotides so they are tokenized as DNA (gLM2 uses
+      lower-case `a/t/c/g` for nucleotides; upper-case letters are amino
+      acids).
+    - Replaces `U`/`u` with `t` (gLM2's vocab has no uracil token).
+    - Prepends the strand marker (`<+>` or `<->`) if not already present.
+    """
+    if strand not in ("+", "-"):
+        raise ValueError(f"strand must be '+' or '-', got {strand!r}")
+    if _DNA_PREFIX_RE.match(seq):
+        return seq
+    cleaned = seq.lower().replace("u", "t")
+    return f"<{strand}>{cleaned}"
+
+
+class gLM2Tokenizer(PreTrainedTokenizerFast):
+    """gLM2 fast tokenizer with optional automatic DNA preparation."""
+
+    VOCAB = [
+        "<cls>", "<pad>", "<eos>", "<unk>",
+        "L", "A", "G", "V", "S", "E", "R", "T", "I", "D", "P", "K",
+        "Q", "N", "F", "Y", "M", "H", "W", "C", "X", "B", "U", "Z",
+        "O", "a", "t", "c", "g", "<+>", "<->", "<mask>", "<sep>",
+    ]
+
+    def __init__(
+        self,
+        unk_token: str = "<unk>",
+        cls_token: str = "<cls>",
+        pad_token: str = "<pad>",
+        mask_token: str = "<mask>",
+        eos_token: str = "<eos>",
+        sep_token: str = "<sep>",
+        pos_token: str = "<+>",
+        neg_token: str = "<->",
+        auto_prepare_dna: bool = True,
+        model_max_length: int = 4096,
+        **kwargs,
+    ):
+        all_tokens = self.VOCAB
+        token_to_id = {tok: ind for ind, tok in enumerate(all_tokens)}
+
+        bpe = BPE(token_to_id, merges=[], unk_token=str(unk_token))
+        tokenizer = Tokenizer(bpe)
+        special_tokens = [
+            cls_token, pad_token, mask_token, eos_token, sep_token, pos_token, neg_token,
+        ]
+        tokenizer.add_special_tokens(special_tokens)
+
+        super().__init__(
+            tokenizer_object=tokenizer,
+            unk_token=unk_token,
+            cls_token=cls_token,
+            pad_token=pad_token,
+            mask_token=mask_token,
+            eos_token=eos_token,
+            sep_token=sep_token,
+            model_max_length=model_max_length,
+            **kwargs,
+        )
+        self.auto_prepare_dna = auto_prepare_dna
+
+    def _maybe_prepare(
+        self, text: Union[str, Sequence[str]]
+    ) -> Union[str, List[str]]:
+        if not self.auto_prepare_dna:
+            return text  # type: ignore[return-value]
+        if isinstance(text, str):
+            return prepare_dna_sequence(text) if _is_pure_dna(text) else text
+        out: List[str] = []
+        for s in text:
+            if isinstance(s, str) and _is_pure_dna(s):
+                out.append(prepare_dna_sequence(s))
+            else:
+                out.append(s)
+        return out
+
+    def __call__(
+        self,
+        text: Union[str, Sequence[str], None] = None,
+        text_pair=None,
+        text_target=None,
+        text_pair_target=None,
+        *args,
+        **kwargs,
+    ):
+        if text is not None:
+            text = self._maybe_prepare(text)
+        return super().__call__(
+            text,
+            text_pair=text_pair,
+            text_target=text_target,
+            text_pair_target=text_pair_target,
+            *args,
+            **kwargs,
+        )
+
+    def encode(
+        self, text: Union[str, Sequence[str]], *args, **kwargs
+    ) -> List[int]:
+        text = self._maybe_prepare(text)
+        return super().encode(text, *args, **kwargs)

model.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:047facd39472afe082c985058fa042f43a7b10d4ffb2ab51b9a3e6c63e9f3834
+size 609855088

modeling_glm2.py

@@ -0,0 +1,535 @@
+"""PyTorch gLM2 model.
+
+Minimal HuggingFace port of tattabio/gLM2 with three attention implementations
+(eager, sdpa, flash_attention_2) and standard HF outputs.
+
+Architecture is unchanged from the upstream `tattabio/gLM2_*` checkpoints
+(RMSNorm, rotary position embeddings, fused QKV, SwiGLU MLP). Weight names
+match upstream so the same `model.safetensors` loads cleanly.
+"""
+
+import math
+from typing import Optional, Tuple, Union
+
+import torch
+import torch.nn.functional as F
+from torch import nn
+from torch.nn import CrossEntropyLoss
+
+from transformers.modeling_outputs import BaseModelOutput, MaskedLMOutput
+from transformers.modeling_utils import PreTrainedModel
+from transformers.utils import logging
+
+from .configuration_glm2 import gLM2Config
+
+logger = logging.get_logger(__name__)
+
+
+def rotate_half(x: torch.Tensor) -> torch.Tensor:
+    x1, x2 = x.chunk(2, dim=-1)
+    return torch.cat((-x2, x1), dim=-1)
+
+
+def apply_rotary_emb_torch(
+    x: torch.Tensor, cos: torch.Tensor, sin: torch.Tensor
+) -> torch.Tensor:
+    """Apply rotary embeddings to `x`.
+
+    Args:
+        x: (batch, seqlen, nheads, headdim)
+        cos, sin: (seqlen, rotary_dim / 2) - rotary_dim must equal headdim.
+    """
+    seqlen = x.shape[1]
+    cos = cos[:seqlen]
+    sin = sin[:seqlen]
+    cos = cos.to(x.dtype)
+    sin = sin.to(x.dtype)
+    cos = cos.repeat_interleave(2, dim=-1) if False else torch.cat([cos, cos], dim=-1)
+    sin = torch.cat([sin, sin], dim=-1)
+    cos = cos.unsqueeze(-2)
+    sin = sin.unsqueeze(-2)
+    return x * cos + rotate_half(x) * sin
+
+
+class RotaryEmbedding(nn.Module):
+    """Rotary position embeddings.
+
+    Identical numerics to the upstream `tattabio/gLM2_*` `RotaryEmbedding`
+    (non-interleaved, base 10000, no scaling), simplified to the path actually
+    used by the released checkpoints.
+    """
+
+    def __init__(self, dim: int, base: float = 10000.0):
+        super().__init__()
+        self.dim = dim
+        self.base = float(base)
+        inv_freq = 1.0 / (
+            self.base
+            ** (torch.arange(0, dim, 2, dtype=torch.float32) / dim)
+        )
+        self.register_buffer("inv_freq", inv_freq, persistent=False)
+
+        self._seq_len_cached = 0
+        self._cos_cached: Optional[torch.Tensor] = None
+        self._sin_cached: Optional[torch.Tensor] = None
+
+    def _update_cache(self, seqlen: int, device: torch.device, dtype: torch.dtype) -> None:
+        if (
+            seqlen > self._seq_len_cached
+            or self._cos_cached is None
+            or self._cos_cached.device != device
+            or self._cos_cached.dtype != dtype
+        ):
+            self._seq_len_cached = seqlen
+            inv_freq = self.inv_freq
+            if inv_freq.dtype != torch.float32:
+                inv_freq = 1.0 / (
+                    self.base
+                    ** (
+                        torch.arange(0, self.dim, 2, device=device, dtype=torch.float32)
+                        / self.dim
+                    )
+                )
+            t = torch.arange(seqlen, device=device, dtype=torch.float32)
+            freqs = torch.outer(t, inv_freq.to(device=device, dtype=torch.float32))
+            self._cos_cached = torch.cos(freqs).to(dtype)
+            self._sin_cached = torch.sin(freqs).to(dtype)
+
+    def forward(self, qkv: torch.Tensor) -> torch.Tensor:
+        """Apply rotary embeddings to q and k. v is left untouched.
+
+        Args:
+            qkv: (batch, seqlen, 3, nheads, headdim)
+        """
+        seqlen = qkv.shape[1]
+        self._update_cache(seqlen, device=qkv.device, dtype=qkv.dtype)
+        cos = self._cos_cached
+        sin = self._sin_cached
+        q_rot = apply_rotary_emb_torch(qkv[:, :, 0], cos, sin)
+        k_rot = apply_rotary_emb_torch(qkv[:, :, 1], cos, sin)
+        return torch.stack((q_rot, k_rot, qkv[:, :, 2]), dim=2)
+
+
+def rmsnorm_func(
+    hidden_states: torch.Tensor, weight: torch.Tensor, variance_epsilon: torch.Tensor
+) -> torch.Tensor:
+    input_dtype = hidden_states.dtype
+    hidden_states = hidden_states.to(torch.float32)
+    variance = hidden_states.pow(2).mean(-1, keepdim=True)
+    hidden_states = hidden_states * torch.rsqrt(variance + variance_epsilon)
+    return (weight * hidden_states).to(input_dtype)
+
+
+class RMSNorm(nn.Module):
+    """Root-mean-square layer norm."""
+
+    def __init__(self, dim: int, eps: float = 1e-6):
+        super().__init__()
+        self.weight = nn.Parameter(torch.ones(dim))
+        self.register_buffer(
+            "variance_epsilon", torch.tensor(eps), persistent=False
+        )
+
+    def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
+        return rmsnorm_func(hidden_states, self.weight, self.variance_epsilon)
+
+
+class gLM2Attention(nn.Module):
+    """Eager multi-head attention with rotary embeddings."""
+
+    def __init__(self, config: gLM2Config):
+        super().__init__()
+        self.n_heads = config.heads
+        self.head_dim = config.dim // config.heads
+        self.dim = config.dim
+
+        self.wqkv = nn.Linear(config.dim, self.n_heads * self.head_dim * 3, bias=False)
+        self.wo = nn.Linear(self.n_heads * self.head_dim, config.dim, bias=False)
+        self.rotary_emb = RotaryEmbedding(self.head_dim)
+
+    def _qkv(self, x: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
+        bsz, seqlen, _ = x.shape
+        qkv = self.wqkv(x).view(bsz, seqlen, 3, self.n_heads, self.head_dim)
+        qkv = self.rotary_emb(qkv)
+        # qkv: (B, S, 3, H, D) -> (B, H, S, D) for q,k,v
+        qkv = qkv.permute(0, 3, 2, 1, 4)  # (B, H, 3, S, D)
+        q = qkv[:, :, 0]
+        k = qkv[:, :, 1]
+        v = qkv[:, :, 2]
+        return q, k, v
+
+    def _output(self, attn_out: torch.Tensor) -> torch.Tensor:
+        # attn_out: (B, H, S, D) -> (B, S, H*D)
+        bsz, _, seqlen, _ = attn_out.shape
+        out = attn_out.permute(0, 2, 1, 3).contiguous().view(bsz, seqlen, self.n_heads * self.head_dim)
+        return self.wo(out)
+
+    def forward(
+        self,
+        x: torch.Tensor,
+        attention_mask: Optional[torch.Tensor] = None,
+        output_attentions: bool = False,
+    ) -> Tuple[torch.Tensor, Optional[torch.Tensor]]:
+        q, k, v = self._qkv(x)
+        scale = 1.0 / math.sqrt(self.head_dim)
+        # (B, H, S, S). Compute in fp32 for numerical stability under bf16/fp16
+        # (matches what flash-attn / SDPA do internally).
+        scores = torch.matmul(q.float(), k.float().transpose(-2, -1)) * scale
+        if attention_mask is not None:
+            mask = attention_mask[:, None, None, :]
+            scores = scores.masked_fill(mask == 0, torch.finfo(scores.dtype).min)
+        attn = scores.softmax(dim=-1)
+        attn_for_return = attn.to(q.dtype) if output_attentions else None
+        context = torch.matmul(attn, v.float()).to(q.dtype)
+        return self._output(context), attn_for_return
+
+
+class gLM2SdpaAttention(gLM2Attention):
+    """SDPA-backed attention. Falls back to eager when output_attentions=True."""
+
+    def forward(
+        self,
+        x: torch.Tensor,
+        attention_mask: Optional[torch.Tensor] = None,
+        output_attentions: bool = False,
+    ) -> Tuple[torch.Tensor, Optional[torch.Tensor]]:
+        if output_attentions:
+            return super().forward(x, attention_mask=attention_mask, output_attentions=True)
+
+        q, k, v = self._qkv(x)
+        attn_mask = None
+        if attention_mask is not None:
+            # SDPA wants (B, 1, 1, S) bool mask where True = attend.
+            attn_mask = attention_mask[:, None, None, :].bool()
+        out = F.scaled_dot_product_attention(q, k, v, attn_mask=attn_mask)
+        return self._output(out), None
+
+
+class gLM2FlashAttention2(gLM2Attention):
+    """flash-attn 2 backed attention. Falls back to eager when output_attentions=True."""
+
+    def forward(
+        self,
+        x: torch.Tensor,
+        attention_mask: Optional[torch.Tensor] = None,
+        output_attentions: bool = False,
+    ) -> Tuple[torch.Tensor, Optional[torch.Tensor]]:
+        if output_attentions:
+            return super().forward(x, attention_mask=attention_mask, output_attentions=True)
+
+        try:
+            from flash_attn import flash_attn_func, flash_attn_varlen_func
+            from flash_attn.bert_padding import pad_input, unpad_input
+        except ImportError as e:
+            raise ImportError(
+                "flash_attn is required for attn_implementation='flash_attention_2'. "
+                "Install with: pip install flash-attn --no-build-isolation"
+            ) from e
+
+        bsz, seqlen, _ = x.shape
+        qkv = self.wqkv(x).view(bsz, seqlen, 3, self.n_heads, self.head_dim)
+        qkv = self.rotary_emb(qkv)
+        # flash-attn wants (B, S, H, D) per q/k/v.
+        q = qkv[:, :, 0]
+        k = qkv[:, :, 1]
+        v = qkv[:, :, 2]
+
+        orig_dtype = q.dtype
+        if q.dtype not in (torch.float16, torch.bfloat16):
+            q = q.to(torch.bfloat16)
+            k = k.to(torch.bfloat16)
+            v = v.to(torch.bfloat16)
+
+        if attention_mask is not None and (attention_mask == 0).any():
+            attention_mask_bool = attention_mask.bool()  # True = attend
+            q_unpad, indices_q, cu_q, max_q, _ = unpad_input(q, attention_mask_bool)
+            k_unpad, _, cu_k, max_k, _ = unpad_input(k, attention_mask_bool)
+            v_unpad, _, _, _, _ = unpad_input(v, attention_mask_bool)
+            out_unpad = flash_attn_varlen_func(
+                q_unpad,
+                k_unpad,
+                v_unpad,
+                cu_seqlens_q=cu_q,
+                cu_seqlens_k=cu_k,
+                max_seqlen_q=max_q,
+                max_seqlen_k=max_k,
+                causal=False,
+            )
+            out = pad_input(out_unpad, indices_q, bsz, seqlen)
+        else:
+            out = flash_attn_func(q, k, v, causal=False)
+
+        out = out.to(orig_dtype)
+        out = out.contiguous().view(bsz, seqlen, self.n_heads * self.head_dim)
+        return self.wo(out), None
+
+
+GLM2_ATTENTION_CLASSES = {
+    "eager": gLM2Attention,
+    "sdpa": gLM2SdpaAttention,
+    "flash_attention_2": gLM2FlashAttention2,
+}
+
+
+class FeedForward(nn.Module):
+    """SwiGLU MLP."""
+
+    def __init__(
+        self,
+        dim: int,
+        hidden_dim: int,
+        multiple_of: int,
+        ffn_dim_multiplier: Optional[float],
+    ):
+        super().__init__()
+        hidden_dim = int(2 * hidden_dim / 3)
+        if ffn_dim_multiplier is not None:
+            hidden_dim = int(ffn_dim_multiplier * hidden_dim)
+        hidden_dim = multiple_of * ((hidden_dim + multiple_of - 1) // multiple_of)
+
+        self.w1 = nn.Linear(dim, hidden_dim, bias=False)
+        self.w2 = nn.Linear(hidden_dim, dim, bias=False)
+        self.w3 = nn.Linear(dim, hidden_dim, bias=False)
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        return self.w2(F.silu(self.w1(x)) * self.w3(x))
+
+
+class TransformerBlock(nn.Module):
+    """Pre-norm transformer block."""
+
+    def __init__(self, config: gLM2Config):
+        super().__init__()
+        attn_impl = getattr(config, "_attn_implementation", "eager")
+        attn_cls = GLM2_ATTENTION_CLASSES[attn_impl]
+        self.attention = attn_cls(config)
+        self.feed_forward = FeedForward(
+            dim=config.dim,
+            hidden_dim=4 * config.dim,
+            multiple_of=config.swiglu_multiple_of,
+            ffn_dim_multiplier=config.ffn_dim_multiplier,
+        )
+        self.attention_norm = RMSNorm(config.dim, eps=config.norm_eps)
+        self.ffn_norm = RMSNorm(config.dim, eps=config.norm_eps)
+
+    def forward(
+        self,
+        x: torch.Tensor,
+        attention_mask: Optional[torch.Tensor] = None,
+        output_attentions: bool = False,
+    ) -> Tuple[torch.Tensor, Optional[torch.Tensor]]:
+        attn_out, attn_weights = self.attention(
+            self.attention_norm(x),
+            attention_mask=attention_mask,
+            output_attentions=output_attentions,
+        )
+        h = x + attn_out
+        out = h + self.feed_forward(self.ffn_norm(h))
+        return out, attn_weights
+
+
+class TransformerLayers(nn.Module):
+    def __init__(self, config: gLM2Config):
+        super().__init__()
+        self.config = config
+        self.layers = nn.ModuleList(
+            [TransformerBlock(config) for _ in range(config.depth)]
+        )
+
+    def forward(
+        self,
+        x: torch.Tensor,
+        attention_mask: Optional[torch.Tensor] = None,
+        output_hidden_states: bool = False,
+        output_attentions: bool = False,
+    ) -> Tuple[torch.Tensor, Optional[Tuple[torch.Tensor, ...]], Optional[Tuple[torch.Tensor, ...]]]:
+        if x.shape[-1] != self.config.dim:
+            raise ValueError(
+                f"Input feature dim should be {self.config.dim}, but input has shape {x.shape}"
+            )
+        all_hidden_states: list = []
+        all_attentions: list = []
+        if output_hidden_states:
+            all_hidden_states.append(x)
+        for layer in self.layers:
+            x, attn_weights = layer(
+                x, attention_mask=attention_mask, output_attentions=output_attentions
+            )
+            if output_hidden_states:
+                all_hidden_states.append(x)
+            if output_attentions:
+                all_attentions.append(attn_weights)
+        hidden_tuple = tuple(all_hidden_states) if output_hidden_states else None
+        attn_tuple = tuple(all_attentions) if output_attentions else None
+        return x, hidden_tuple, attn_tuple
+
+
+class gLM2PreTrainedModel(PreTrainedModel):
+    """Base class for gLM2 weight init / from_pretrained dispatch."""
+
+    config_class = gLM2Config
+    base_model_prefix = "glm2"
+    supports_gradient_checkpointing = False
+    _supports_sdpa = True
+    _supports_flash_attn_2 = True
+
+    def _init_weights(self, module):
+        std = getattr(self.config, "initializer_range", 0.02)
+        if isinstance(module, nn.Linear):
+            nn.init.normal_(module.weight, std=std)
+            if module.bias is not None:
+                nn.init.zeros_(module.bias)
+        elif isinstance(module, nn.Embedding):
+            nn.init.normal_(module.weight, std=std)
+            if module.padding_idx is not None:
+                with torch.no_grad():
+                    module.weight[module.padding_idx].zero_()
+        elif isinstance(module, RotaryEmbedding):
+            inv_freq = 1.0 / (
+                module.base
+                ** (
+                    torch.arange(
+                        0, module.dim, 2, device=module.inv_freq.device, dtype=torch.float32
+                    )
+                    / module.dim
+                )
+            )
+            with torch.no_grad():
+                module.inv_freq.copy_(inv_freq)
+        elif isinstance(module, RMSNorm):
+            with torch.no_grad():
+                module.variance_epsilon.fill_(self.config.norm_eps)
+
+
+class gLM2Model(gLM2PreTrainedModel):
+    """gLM2 backbone (token embedding + transformer encoder)."""
+
+    def __init__(self, config: gLM2Config):
+        super().__init__(config)
+        self.config = config
+        self.tok_embeddings = nn.Embedding(config.vocab_size, config.dim)
+        self.encoder = TransformerLayers(config)
+        self.post_init()
+
+    def get_input_embeddings(self) -> nn.Embedding:
+        return self.tok_embeddings
+
+    def set_input_embeddings(self, value: nn.Embedding) -> None:
+        self.tok_embeddings = value
+
+    def forward(
+        self,
+        input_ids: torch.Tensor,
+        attention_mask: Optional[torch.Tensor] = None,
+        output_hidden_states: Optional[bool] = None,
+        output_attentions: Optional[bool] = None,
+        return_dict: Optional[bool] = None,
+    ) -> Union[Tuple[torch.Tensor, ...], BaseModelOutput]:
+        output_hidden_states = (
+            output_hidden_states
+            if output_hidden_states is not None
+            else self.config.output_hidden_states
+        )
+        output_attentions = (
+            output_attentions
+            if output_attentions is not None
+            else self.config.output_attentions
+        )
+        return_dict = (
+            return_dict if return_dict is not None else self.config.use_return_dict
+        )
+
+        h = self.tok_embeddings(input_ids)
+        sequence_output, all_hidden_states, all_attentions = self.encoder(
+            h,
+            attention_mask=attention_mask,
+            output_hidden_states=bool(output_hidden_states),
+            output_attentions=bool(output_attentions),
+        )
+
+        if not return_dict:
+            return tuple(
+                v
+                for v in (sequence_output, all_hidden_states, all_attentions)
+                if v is not None
+            )
+
+        return BaseModelOutput(
+            last_hidden_state=sequence_output,
+            hidden_states=all_hidden_states,
+            attentions=all_attentions,
+        )
+
+
+class gLM2LMHead(nn.Module):
+    def __init__(self, config: gLM2Config):
+        super().__init__()
+        self.norm = RMSNorm(config.dim, eps=config.norm_eps)
+        self.proj_output = nn.Linear(config.dim, config.vocab_size, bias=False)
+
+    def forward(self, features: torch.Tensor) -> torch.Tensor:
+        return self.proj_output(self.norm(features))
+
+
+class gLM2ForMaskedLM(gLM2PreTrainedModel):
+    """gLM2 with the masked-language-modeling head."""
+
+    _tied_weights_keys = []
+
+    def __init__(self, config: gLM2Config):
+        super().__init__(config)
+        self.glm2 = gLM2Model(config)
+        self.lm_head = gLM2LMHead(config)
+        self.post_init()
+
+    def get_output_embeddings(self) -> nn.Linear:
+        return self.lm_head.proj_output
+
+    def set_output_embeddings(self, new_embeddings: nn.Linear) -> None:
+        self.lm_head.proj_output = new_embeddings
+
+    def forward(
+        self,
+        input_ids: torch.Tensor,
+        attention_mask: Optional[torch.Tensor] = None,
+        labels: Optional[torch.LongTensor] = None,
+        output_hidden_states: Optional[bool] = None,
+        output_attentions: Optional[bool] = None,
+        return_dict: Optional[bool] = None,
+    ) -> Union[Tuple[torch.Tensor, ...], MaskedLMOutput]:
+        return_dict = (
+            return_dict if return_dict is not None else self.config.use_return_dict
+        )
+
+        outputs = self.glm2(
+            input_ids,
+            attention_mask=attention_mask,
+            output_hidden_states=output_hidden_states,
+            output_attentions=output_attentions,
+            return_dict=True,
+        )
+        sequence_output = outputs.last_hidden_state
+        prediction_scores = self.lm_head(sequence_output)
+
+        masked_lm_loss = None
+        if labels is not None:
+            loss_fct = CrossEntropyLoss()
+            labels = labels.to(prediction_scores.device)
+            masked_lm_loss = loss_fct(
+                prediction_scores.view(-1, self.config.vocab_size), labels.view(-1)
+            )
+
+        if not return_dict:
+            output = (prediction_scores,)
+            if outputs.hidden_states is not None:
+                output = output + (outputs.hidden_states,)
+            if outputs.attentions is not None:
+                output = output + (outputs.attentions,)
+            return ((masked_lm_loss,) + output) if masked_lm_loss is not None else output
+
+        return MaskedLMOutput(
+            loss=masked_lm_loss,
+            logits=prediction_scores,
+            hidden_states=outputs.hidden_states,
+            attentions=outputs.attentions,
+        )

special_tokens_map.json

@@ -0,0 +1,8 @@
+{
+  "cls_token": "<cls>",
+  "eos_token": "<eos>",
+  "mask_token": "<mask>",
+  "pad_token": "<pad>",
+  "sep_token": "<sep>",
+  "unk_token": "<unk>"
+}

tokenizer.json

@@ -0,0 +1,132 @@
+{
+  "version": "1.0",
+  "truncation": null,
+  "padding": null,
+  "added_tokens": [
+    {
+      "id": 0,
+      "content": "<cls>",
+      "single_word": false,
+      "lstrip": false,
+      "rstrip": false,
+      "normalized": false,
+      "special": true
+    },
+    {
+      "id": 1,
+      "content": "<pad>",
+      "single_word": false,
+      "lstrip": false,
+      "rstrip": false,
+      "normalized": false,
+      "special": true
+    },
+    {
+      "id": 2,
+      "content": "<eos>",
+      "single_word": false,
+      "lstrip": false,
+      "rstrip": false,
+      "normalized": false,
+      "special": true
+    },
+    {
+      "id": 3,
+      "content": "<unk>",
+      "single_word": false,
+      "lstrip": false,
+      "rstrip": false,
+      "normalized": false,
+      "special": true
+    },
+    {
+      "id": 33,
+      "content": "<+>",
+      "single_word": false,
+      "lstrip": false,
+      "rstrip": false,
+      "normalized": false,
+      "special": true
+    },
+    {
+      "id": 34,
+      "content": "<->",
+      "single_word": false,
+      "lstrip": false,
+      "rstrip": false,
+      "normalized": false,
+      "special": true
+    },
+    {
+      "id": 35,
+      "content": "<mask>",
+      "single_word": false,
+      "lstrip": false,
+      "rstrip": false,
+      "normalized": false,
+      "special": true
+    },
+    {
+      "id": 36,
+      "content": "<sep>",
+      "single_word": false,
+      "lstrip": false,
+      "rstrip": false,
+      "normalized": false,
+      "special": true
+    }
+  ],
+  "normalizer": null,
+  "pre_tokenizer": null,
+  "post_processor": null,
+  "decoder": null,
+  "model": {
+    "type": "BPE",
+    "dropout": null,
+    "unk_token": "<unk>",
+    "continuing_subword_prefix": null,
+    "end_of_word_suffix": null,
+    "fuse_unk": false,
+    "byte_fallback": false,
+    "vocab": {
+      "<cls>": 0,
+      "<pad>": 1,
+      "<eos>": 2,
+      "<unk>": 3,
+      "L": 4,
+      "A": 5,
+      "G": 6,
+      "V": 7,
+      "S": 8,
+      "E": 9,
+      "R": 10,
+      "T": 11,
+      "I": 12,
+      "D": 13,
+      "P": 14,
+      "K": 15,
+      "Q": 16,
+      "N": 17,
+      "F": 18,
+      "Y": 19,
+      "M": 20,
+      "H": 21,
+      "W": 22,
+      "C": 23,
+      "X": 24,
+      "B": 25,
+      "U": 26,
+      "Z": 27,
+      "O": 28,
+      "a": 29,
+      "t": 30,
+      "c": 31,
+      "g": 32,
+      "<+>": 33,
+      "<->": 34,
+      "<mask>": 35,
+      "<sep>": 36
+    },
+    "merges": []
+  }
+}

tokenizer_config.json

@@ -0,0 +1,83 @@
+{
+  "added_tokens_decoder": {
+    "0": {
+      "content": "<cls>",
+      "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": "<eos>",
+      "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
+    },
+    "33": {
+      "content": "<+>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "34": {
+      "content": "<->",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "35": {
+      "content": "<mask>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "36": {
+      "content": "<sep>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    }
+  },
+  "auto_map": {
+    "AutoTokenizer": [
+      "glm_tokenizer.gLM2Tokenizer",
+      null
+    ]
+  },
+  "clean_up_tokenization_spaces": true,
+  "cls_token": "<cls>",
+  "eos_token": "<eos>",
+  "mask_token": "<mask>",
+  "model_max_length": 4096,
+  "pad_token": "<pad>",
+  "sep_token": "<sep>",
+  "tokenizer_class": "gLM2Tokenizer",
+  "unk_token": "<unk>"
+}