"""Modern BERTc v3 — ModernBERT release-aligned + Cramming-style ScaledSinusoidal PE. 主要按 release `modernbert-base-pretrain.yaml` 对齐(除 Alt Attn 和 PE): - 架构: 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, field 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(对齐 ModernBERT release) embed_dropout: float = 0.0 mlp_dropout: float = 0.0 attn_out_dropout: float = 0.1 attn_probs_dropout: float = 0.0 # 架构开关(对齐 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" @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 self.attn_probs_dropout = config.attn_probs_dropout # 无 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) self.out_dropout = nn.Dropout(config.attn_out_dropout) 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, dropout_p=self.attn_probs_dropout if self.training else 0.0, is_causal=False, ) # [B, h, L, d] out = out.transpose(1, 2).reshape(B, L, H) return self.out_dropout(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) self.dropout = nn.Dropout(config.mlp_dropout) def forward(self, x: torch.Tensor) -> torch.Tensor: gate, up = self.w_in(x).chunk(2, dim=-1) return self.dropout(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.embed_dropout = nn.Dropout(config.embed_dropout) 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) x = self.embed_dropout(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) # ============ Quick sanity test ============ if __name__ == "__main__": print("=== Test 1: default config (22L/768H release-aligned) ===") cfg = ModernBertConfig() print(f"Config: V={cfg.vocab_size} H={cfg.hidden_size} L={cfg.num_hidden_layers} " f"head={cfg.num_attention_heads} d={cfg.head_dim} I={cfg.intermediate_size} " f"pe_theta={cfg.pe_theta} ln_eps={cfg.layer_norm_eps}") print(f"embed_norm={cfg.embed_norm} skip_first_prenorm={cfg.skip_first_prenorm} " f"final_norm={cfg.final_norm} init={cfg.init_method}") print(f"dropout: embed={cfg.embed_dropout} mlp={cfg.mlp_dropout} " f"attn_out={cfg.attn_out_dropout} attn_probs={cfg.attn_probs_dropout}") model = ModernBertForMLM(cfg) n = model.num_parameters() print(f"params: {n:,} = {n / 1e6:.1f}M") # forward smoke B, L = 2, 128 ids = torch.randint(0, cfg.vocab_size, (B, L)) mask = torch.ones(B, L, dtype=torch.long) labels = torch.randint(0, cfg.vocab_size, (B, L)) out = model(ids, attention_mask=mask, labels=labels) print(f"logits: {out['logits'].shape} loss: {out['loss'].item():.4f}") # backward smoke out["loss"].backward() # 检查所有 trainable param 都有 grad no_grad = [n for n, p in model.named_parameters() if p.requires_grad and p.grad is None] if no_grad: print(f"ERROR: 缺 grad 的参数: {no_grad}") else: print("backward OK (所有参数都有 grad)") # 验证 skip_first_prenorm 生效 assert model.bert.layers[0].skip_norm1, "layer[0].skip_norm1 应为 True" assert not model.bert.layers[1].skip_norm1, "layer[1].skip_norm1 应为 False" assert isinstance(model.bert.layers[0].norm1, nn.Identity), "layer[0].norm1 应为 Identity" assert isinstance(model.bert.layers[1].norm1, LayerNormNoBias), "layer[1].norm1 应为 LayerNormNoBias" print("skip_first_prenorm 配置正确") # 验证 embed_norm + final_norm assert isinstance(model.bert.embed_norm, LayerNormNoBias), "embed_norm 应为 LayerNormNoBias" assert isinstance(model.bert.final_norm, LayerNormNoBias), "final_norm 应为 LayerNormNoBias" print("embed_norm + final_norm 配置正确") # 验证 Megatron init L_layers = cfg.num_hidden_layers expected_scale = (2.0 * L_layers) ** -0.5 # 一个未缩的初始值 std ≈ 0.02,缩后 std ≈ 0.02 * expected_scale o_std = model.bert.layers[5].attn.o.weight.std().item() expected_std = cfg.initializer_range * expected_scale # 允许 ±50% 容差(单 tensor std 估计有 noise) assert 0.5 * expected_std < o_std < 1.5 * expected_std, \ f"Megatron init: attn.o.weight.std={o_std:.6f}, expected≈{expected_std:.6f}" print(f"Megatron init OK: attn.o.weight.std={o_std:.6f} ≈ {expected_std:.6f}") # 验证 no bias for name, p in model.named_parameters(): if "bias" in name and name != "head_bias": print(f"ERROR: 不该有的 bias: {name}") print("Linear/Norm no-bias 配置正确") print("\n=== Test 1b: seg_ids → flex_attention 路径(cross-doc 隔离)===") if torch.cuda.is_available(): # flex_attention 需要 CUDA model_cuda = ModernBertForMLM(cfg).cuda().to(torch.bfloat16) ids_c = ids.cuda() labels_c = labels.cuda() # 构造 seg_ids:chunk 内 0 0 0 ... 0 | 1 1 ... 1 (前半 doc0, 后半 doc1) seg_ids = torch.zeros(B, L, dtype=torch.int32, device="cuda") seg_ids[:, L // 2:] = 1 out_flex = model_cuda(ids_c, seg_ids=seg_ids, labels=labels_c) print(f" flex_attention forward OK loss={out_flex['loss'].item():.4f} " f"logits={out_flex['logits'].shape} dtype={out_flex['logits'].dtype}") out_flex["loss"].backward() no_grad = [n for n, p in model_cuda.named_parameters() if p.requires_grad and p.grad is None] if no_grad: print(f" ERROR: 缺 grad: {no_grad[:5]}") else: print(f" flex_attention backward OK") # 验证 attention 真正隔离了:用相同 ids 但 seg_ids 不同,前半 token 应该输出不同 # (因为 doc0 token 在 seg=0 时只 attend 同 seg=0;改成 seg=2 后 attend 不同的 kv 集) with torch.no_grad(): seg_v1 = torch.zeros(B, L, dtype=torch.int32, device="cuda") seg_v2 = torch.zeros(B, L, dtype=torch.int32, device="cuda") seg_v2[:, :L // 2] = 0 seg_v2[:, L // 2:] = 1 # 前半 vs 后半 隔离 h1 = model_cuda.bert(ids_c, seg_ids=seg_v1) # 全 attend(单 doc) h2 = model_cuda.bert(ids_c, seg_ids=seg_v2) # 前后半隔离 # 前半 token 在 v1 attend 全 chunk,在 v2 只 attend 前半 → 输出应不同 front_diff = (h1[:, :L // 2] - h2[:, :L // 2]).abs().mean().item() print(f" 前半 token 输出差异(应 > 0,验证隔离生效):{front_diff:.6f}") assert front_diff > 1e-3, f"隔离未生效:front_diff={front_diff}" print(f" 跨 doc 隔离工作正常 ✓") else: print(" 跳过(无 CUDA)") print("\n=== Test 2: legacy v1 config (12L/1024H) — 验证向后兼容 ===") cfg_v1 = ModernBertConfig( hidden_size=1024, num_hidden_layers=12, num_attention_heads=16, intermediate_size=2752, embed_norm=False, skip_first_prenorm=False, init_method="normal", ) model_v1 = ModernBertForMLM(cfg_v1) n_v1 = model_v1.num_parameters() print(f"v1-style params: {n_v1 / 1e6:.1f}M") out_v1 = model_v1(ids, attention_mask=mask, labels=labels) print(f"v1 forward OK loss={out_v1['loss'].item():.4f}") print("\n=== All smoke tests passed ===")