"""Modern BERTc 骨干网络。ModernBERT release 对齐 + Cramming 式 ScaledSinusoidal PE。 只依赖 torch。**state_dict 的 key 不能动** —— 改任何模块名或嵌套层级都会让 HF 上已发布的六个模型权重全部失配,而模型照样能随机初始化跑起来、不报错。 改动后跑 test/test_reproduce_sota.py 验证。 两个已发布规格(都用同一份代码,只是 config 不同): BERTc-165M (v4-Mid) 12L / 1024H / 2752I / 16 heads BERTc-315M (v4-Large) 24L / 1024H / 2752I / 16 heads 主要按 release `modernbert-base-pretrain.yaml` 对齐(除 Alt Attn 和 PE): - 默认 config: 22L / 768H / 1152I (GLU) / 12 heads,head_dim=64 - **ScaledSinusoidal 位置编码**(Hua et al. 2022 FLASH;Cramming 实测短 seq 比 RoPE 更值:计算几乎免费,RoPE 收益被 5-10% 速度损失抵消) - GeGLU FFN(glu + gelu) - LayerNorm 无 bias(eps=1e-5),非 RMSNorm - pre-norm 布局 + skip_first_prenorm - embed_norm + final_norm - Megatron-style init:残差层 W 缩 1/sqrt(2L) - 全无 Linear bias - Dropout: 全 0(Cramming 论据:short single-epoch 无 overfit risk) - tied word embedding - flex_attention compiled(支持 cross-doc 隔离 via seg_ids) 不上的 ModernBERT 特性: - Alternating Attention(我们走全局 attention) - Unpadded packing + cu_seqlens(我们定长 pack) - RoPE(换 ScaledSinusoidal,见 Cramming Section 4.2) 参数量(默认 22L/768H/1152I,V=12536): emb (tied) : 12536 × 768 ≈ 9.6M embed_norm : 768 × 1 ≈ 1K(no bias) per layer: norm1/2 : 768 × 2 ≈ 1.5K Q K V O : 4 × 768² ≈ 2.36M GeGLU (W_in=2I, W_out): 768×2304 + 1152×768 ≈ 2.66M total per layer : ≈ 5.0M × 22 layers : ≈ 110M final_norm : 768 × 1 head: dense + norm + gelu: 768×768 + 768 ≈ 0.59M head_bias (V,) : 12.5K total ≈ 130M """ from dataclasses import dataclass from typing import Optional import math import torch import torch.nn as nn import torch.nn.functional as F from torch.nn.attention.flex_attention import ( flex_attention as _flex_attention_raw, create_block_mask, ) # torch.compile 是 lazy 的 — module import 不触发 trace,first call 时才编译。 # ModernBERT 源码用 mode="max-autotune-no-cudagraphs",我们用 default mode 平衡 # (first-call 编译几秒,vs max-autotune 可能几十秒)。 # 训练时默认调 _flex_attention(compiled);smoke 也走这条路径,保证统一。 _flex_attention = torch.compile(_flex_attention_raw) # ============ Config ============ @dataclass class ModernBertConfig: vocab_size: int = 12536 hidden_size: int = 768 num_hidden_layers: int = 22 num_attention_heads: int = 12 intermediate_size: int = 1152 max_position_embeddings: int = 1024 pad_token_id: int = 12531 mask_token_id: int = 12535 pe_theta: float = 10000.0 # ScaledSinusoidal 频率 base(Vaswani 2017 默认) layer_norm_eps: float = 1e-5 initializer_range: float = 0.02 tie_word_embeddings: bool = True # 全程无 dropout:预训练数据量远大于参数量,不存在过拟合,dropout 只拖慢 # 收敛。所以没有 dropout 配置项 —— 已发布的 config.json 里那几个 dropout # 字段都是 0,load_backbone 按字段名过滤,多出来的会被忽略。 # 架构开关(对齐 release) embed_norm: bool = True # embedding 后立刻 LayerNorm skip_first_prenorm: bool = True # 第 1 层不做 pre-norm final_norm: bool = True # 最后一层后 LayerNorm # init init_method: str = "megatron" # "megatron"(残差层 ×1/sqrt(2L))或 "normal" @classmethod def from_dict(cls, d: dict) -> "ModernBertConfig": """按字段名过滤后构造。config.json 可能带这个版本不认识的键 —— 已发布的六个模型就带着四个 dropout 字段(现在全程无 dropout,字段去掉了), 直接 `ModernBertConfig(**d)` 会 TypeError。""" return cls(**{k: v for k, v in d.items() if k in cls.__dataclass_fields__}) @property def head_dim(self) -> int: assert self.hidden_size % self.num_attention_heads == 0 return self.hidden_size // self.num_attention_heads # ============ LayerNorm(no bias)============ class LayerNormNoBias(nn.Module): """LayerNorm with weight only (no bias). 对齐 ModernBERT release。""" def __init__(self, hidden_size: int, eps: float = 1e-5): super().__init__() self.weight = nn.Parameter(torch.ones(hidden_size)) self.eps = eps self.normalized_shape = (hidden_size,) def forward(self, x: torch.Tensor) -> torch.Tensor: return F.layer_norm(x, self.normalized_shape, self.weight, None, self.eps) # ============ ScaledSinusoidal Position Embedding(Hua et al. 2022 / Cramming)============ class ScaledSinusoidalPE(nn.Module): """Scaled sinusoidal positional embedding(Hua 2022 FLASH paper)。 标准 sinusoidal:PE[pos, 2i]=sin(pos/θ^(2i/d)), PE[pos, 2i+1]=cos(...)。 `scale_factor` 是一个 learnable 标量,初始 1/sqrt(d)。 用法:embedding 之后直接 `x = embed + pos_emb(input_ids)`,跟所有层共享。 比 RoPE 便宜:只在 embedding 层 fire 一次,attention 里 0 开销。 """ def __init__(self, embedding_dim: int, max_seq_length: int, theta: float = 10000.0): super().__init__() pe = torch.zeros(max_seq_length, embedding_dim) position = torch.arange(0, max_seq_length, dtype=torch.float).unsqueeze(1) div_term = torch.exp( torch.arange(0, embedding_dim, 2).float() * (-math.log(theta) / embedding_dim) ) pe[:, 0::2] = torch.sin(position * div_term) pe[:, 1::2] = torch.cos(position * div_term) pe = pe.unsqueeze(0) # [1, L, d] self.register_buffer("pe", pe, persistent=False) self.scale_factor = nn.Parameter(torch.tensor([1.0 / embedding_dim ** 0.5])) def forward(self, seq_len: int) -> torch.Tensor: return self.scale_factor * self.pe[:, :seq_len, :] # ============ Attention(bidirectional,无 RoPE,位置走 ScaledSinusoidal)============ class ModernBertAttention(nn.Module): def __init__(self, config: ModernBertConfig): super().__init__() self.num_heads = config.num_attention_heads self.head_dim = config.head_dim self.scale = self.head_dim ** -0.5 # 无 bias self.qkv = nn.Linear(config.hidden_size, 3 * config.hidden_size, bias=False) self.o = nn.Linear(config.hidden_size, config.hidden_size, bias=False) def forward(self, x: torch.Tensor, block_mask=None, attention_mask: Optional[torch.Tensor] = None) -> torch.Tensor: """三种 attention 模式: - block_mask 非空 → flex_attention(block-diag,跨 doc 隔离,训练时) - block_mask 空,attention_mask 非空 → SDPA + pad mask(fine-tune) - 都空 → SDPA 全可见 位置信息走 ScaledSinusoidal,已在 embedding 层加,attention 里无 cos/sin 计算。 """ B, L, H = x.shape qkv = self.qkv(x).reshape(B, L, 3, self.num_heads, self.head_dim) q, k, v = qkv.unbind(dim=2) # 各 [B, L, h, d] q = q.transpose(1, 2) k = k.transpose(1, 2) v = v.transpose(1, 2) if block_mask is not None: # flex_attention(compiled):任意 mask + flash 速度;不支持 dropout_p out = _flex_attention(q, k, v, block_mask=block_mask) else: sdpa_mask = None if attention_mask is not None: sdpa_mask = attention_mask[:, None, None, :].to(torch.bool) out = F.scaled_dot_product_attention( q, k, v, attn_mask=sdpa_mask, is_causal=False, ) # [B, h, L, d] out = out.transpose(1, 2).reshape(B, L, H) return self.o(out) # ============ GeGLU MLP ============ class GeGLU(nn.Module): """Linear(H, 2*I) → split → GELU(gate) * up → Linear(I, H). 全无 bias。 """ def __init__(self, config: ModernBertConfig): super().__init__() I = config.intermediate_size self.w_in = nn.Linear(config.hidden_size, 2 * I, bias=False) self.w_out = nn.Linear(I, config.hidden_size, bias=False) def forward(self, x: torch.Tensor) -> torch.Tensor: gate, up = self.w_in(x).chunk(2, dim=-1) return self.w_out(F.gelu(gate) * up) # ============ Layer(pre-norm,支持 skip_first_prenorm)============ class ModernBertLayer(nn.Module): def __init__(self, config: ModernBertConfig, is_first: bool = False): super().__init__() # is_first + skip_first_prenorm:第 1 层 attention 前不做 pre-norm # (因为 embed_norm 已经 norm 过一次了) self.skip_norm1 = is_first and config.skip_first_prenorm self.norm1 = nn.Identity() if self.skip_norm1 else LayerNormNoBias(config.hidden_size, eps=config.layer_norm_eps) self.attn = ModernBertAttention(config) self.norm2 = LayerNormNoBias(config.hidden_size, eps=config.layer_norm_eps) self.mlp = GeGLU(config) def forward(self, x, block_mask=None, attention_mask=None): x = x + self.attn(self.norm1(x), block_mask, attention_mask) x = x + self.mlp(self.norm2(x)) return x # ============ Backbone ============ class ModernBertModel(nn.Module): def __init__(self, config: ModernBertConfig): super().__init__() self.config = config self.embed = nn.Embedding(config.vocab_size, config.hidden_size, padding_idx=config.pad_token_id) # ScaledSinusoidal PE(Cramming-style),加在 embedding 后 self.pos_emb = ScaledSinusoidalPE( embedding_dim=config.hidden_size, max_seq_length=config.max_position_embeddings, theta=config.pe_theta, ) self.embed_norm = (LayerNormNoBias(config.hidden_size, eps=config.layer_norm_eps) if config.embed_norm else nn.Identity()) self.layers = nn.ModuleList( [ModernBertLayer(config, is_first=(i == 0)) for i in range(config.num_hidden_layers)] ) self.final_norm = (LayerNormNoBias(config.hidden_size, eps=config.layer_norm_eps) if config.final_norm else nn.Identity()) # init: Megatron-style 残差缩放 self.apply(self._init_weights) if config.init_method == "megatron": self._megatron_residual_init() def _init_weights(self, m): if isinstance(m, nn.Linear): nn.init.normal_(m.weight, std=self.config.initializer_range) if m.bias is not None: nn.init.zeros_(m.bias) elif isinstance(m, nn.Embedding): nn.init.normal_(m.weight, std=self.config.initializer_range) if m.padding_idx is not None: with torch.no_grad(): m.weight[m.padding_idx].zero_() elif isinstance(m, LayerNormNoBias): nn.init.ones_(m.weight) def _megatron_residual_init(self): """对每个 residual 路径的输出 W 缩 1/sqrt(2*L)。 防止深层网络早期 forward variance 爆炸。 residual outputs: attn.o, mlp.w_out。 """ L = self.config.num_hidden_layers scale = (2.0 * L) ** -0.5 for layer in self.layers: with torch.no_grad(): layer.attn.o.weight.mul_(scale) layer.mlp.w_out.weight.mul_(scale) def forward(self, input_ids: torch.Tensor, seg_ids: Optional[torch.Tensor] = None, attention_mask: Optional[torch.Tensor] = None) -> torch.Tensor: """seg_ids: [B, L] int32/uint8,同 doc 同 id;非空时走 flex_attention 跨 doc 隔离。 attention_mask: [B, L] 0/1,只在 seg_ids=None 时使用(fine-tune 路径)。 位置信息:ScaledSinusoidal 加在 embedding 后,attention 内部无位置计算。 """ B, L = input_ids.shape x = self.embed(input_ids) x = x + self.pos_emb(L).to(x.dtype) # 加 scaled sinusoidal PE x = self.embed_norm(x) block_mask = self._build_block_mask(seg_ids, B, L) if seg_ids is not None else None for layer in self.layers: x = layer(x, block_mask, attention_mask) x = self.final_norm(x) return x def _build_block_mask(self, seg_ids: torch.Tensor, B: int, L: int): """seg_ids: [B, L] 用 flex_attention 构造 doc-internal mask。 mask_mod 闭包捕获 seg_ids,在 batch/query/kv 索引下查 doc 是否一致。 """ seg_ids_long = seg_ids.to(torch.int32) def mask_mod(b, h, q_idx, kv_idx): return seg_ids_long[b, q_idx] == seg_ids_long[b, kv_idx] # H=None 让 mask 跨 head 共享(同 doc 隔离与 head 无关) return create_block_mask(mask_mod, B=B, H=None, Q_LEN=L, KV_LEN=L, device=seg_ids.device) # ============ MLM head(tied embedding)============ class ModernBertForMLM(nn.Module): """MLM head 简化版(Cramming Section 4.2 推荐): - 无 nonlinear head(去 Dense + LN + GeLU)— "without ill effect" - 无 decoder bias(去 head_bias)— "drop the decoder bias" - 仅 tied embedding projection:logits = h @ embed.weight.T - final LayerNorm 已经在 bert.final_norm 提供,这里不需重复 省参数 ~0.6M,forward 略快。""" def __init__(self, config: ModernBertConfig): super().__init__() self.config = config self.bert = ModernBertModel(config) def get_input_embeddings(self): return self.bert.embed def forward(self, input_ids, seg_ids=None, attention_mask=None, labels=None): h = self.bert(input_ids, seg_ids=seg_ids, attention_mask=attention_mask) # [B, L, H],bert 内已 final_norm # 直接 tied embedding projection,无 nonlinear head 也无 bias logits = F.linear(h, self.bert.embed.weight) # [B, L, V] loss = None if labels is not None: loss = F.cross_entropy( logits.view(-1, self.config.vocab_size), labels.view(-1), ignore_index=-100, ) return {"logits": logits, "loss": loss} def num_parameters(self): return sum(p.numel() for p in self.parameters() if p.requires_grad)