Aluode/PerceptionLabPortable
0
1# 🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨2# This file was automatically generated from src/transformers/models/modernbert/modular_modernbert.py.3# Do NOT edit this file manually as any edits will be overwritten by the generation of4# the file from the modular. If any change should be done, please apply the change to the5# modular_modernbert.py file directly. One of our CI enforces this.6# 🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨7# Copyright 2024 Answer.AI, LightOn, and contributors, and the HuggingFace Inc. team. All rights reserved.8#9#10# Licensed under the Apache License, Version 2.0 (the "License");11# you may not use this file except in compliance with the License.12# You may obtain a copy of the License at13#14# http://www.apache.org/licenses/LICENSE-2.015#16# Unless required by applicable law or agreed to in writing, software17# distributed under the License is distributed on an "AS IS" BASIS,18# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.19# See the License for the specific language governing permissions and20# limitations under the License.21 22import copy23import math24from contextlib import nullcontext25from typing import Optional, Union26 27import torch28import torch.nn.functional as F29from torch import nn30from torch.nn import BCEWithLogitsLoss, CrossEntropyLoss, MSELoss31 32from ...activations import ACT2FN33from ...modeling_attn_mask_utils import _prepare_4d_attention_mask34from ...modeling_layers import GradientCheckpointingLayer35from ...modeling_outputs import (36 BaseModelOutput,37 MaskedLMOutput,38 MultipleChoiceModelOutput,39 QuestionAnsweringModelOutput,40 SequenceClassifierOutput,41 TokenClassifierOutput,42)43from ...modeling_rope_utils import ROPE_INIT_FUNCTIONS, dynamic_rope_update44from ...modeling_utils import PreTrainedModel45from ...utils import auto_docstring, is_flash_attn_2_available, logging46from ...utils.import_utils import is_triton_available47from .configuration_modernbert import ModernBertConfig48 49 50if is_flash_attn_2_available():51 from flash_attn.flash_attn_interface import flash_attn_varlen_qkvpacked_func52 from flash_attn.layers.rotary import RotaryEmbedding53 from flash_attn.ops.triton.rotary import apply_rotary54else:55 RotaryEmbedding = object56 57 58logger = logging.get_logger(__name__)59 60 61class ApplyRotaryEmbUnpad(torch.autograd.Function):62 @staticmethod63 def forward(64 ctx,65 qkv,66 cos,67 sin,68 cu_seqlens: Optional[torch.Tensor] = None,69 max_seqlen: Optional[int] = None,70 ):71 # (total_nnz, 3, nheads, headdim)72 qkv = qkv.contiguous()73 total_nnz, _three, _nheads, headdim = qkv.shape74 # We need qkv to be contiguous so that when we reshape to combine (3, nheads) dimensions,75 # we get the same tensor76 # qk = rearrange(qkv[:, :2], "b_s t h d -> b_s (t h) d")77 qk = qkv[:, :2].view(total_nnz, -1, headdim)78 apply_rotary(79 qk,80 cos,81 sin,82 seqlen_offsets=0,83 cu_seqlens=cu_seqlens,84 max_seqlen=max_seqlen,85 interleaved=False,86 inplace=True,87 )88 89 ctx.save_for_backward(cos, sin, cu_seqlens)90 ctx.max_seqlen = max_seqlen91 return qkv92 93 @staticmethod94 def backward(ctx, do):95 cos, sin, cu_seqlens = ctx.saved_tensors96 do = do.contiguous()97 total_nnz, _three, _nheads, headdim = do.shape98 # We need dqkv to be contiguous so that when we reshape to combine (3, nheads) dimensions,99 # we get the same tensor100 dqk = do[:, :2].view(total_nnz, -1, headdim)101 apply_rotary(102 dqk,103 cos,104 sin,105 seqlen_offsets=0,106 cu_seqlens=cu_seqlens,107 max_seqlen=ctx.max_seqlen,108 interleaved=False,109 inplace=True,110 conjugate=True,111 )112 113 return do, None, None, None, None, None, None114 115 116def apply_rotary_unpadded(117 qkv,118 cos,119 sin,120 cu_seqlens: Optional[torch.Tensor] = None,121 max_seqlen: Optional[int] = None,122):123 """124 Arguments:125 qkv: (total_nnz, 3, nheads, headdim) - input tensor for packed QKV.126 cos, sin: (seqlen_rotary, rotary_dim / 2)127 interleaved: if True, rotate pairs of even and odd dimensions (GPT-J style) instead128 of 1st half and 2nd half (GPT-NeoX style).129 inplace: if True, apply rotary embedding in-place.130 seqlen_offsets: (batch_size,) or int. Each sequence in x is shifted by this amount.131 Most commonly used in inference when we have KV cache.132 cu_seqlens: (batch + 1,) or None133 max_seqlen: int134 Return:135 out: (total_nnz, dim)136 rotary_dim must be <= headdim137 Apply rotary embedding to the first rotary_dim of x.138 """139 return ApplyRotaryEmbUnpad.apply(qkv, cos, sin, cu_seqlens, max_seqlen)140 141 142class ModernBertUnpaddedRotaryEmbedding(RotaryEmbedding):143 """144 The rotary position embeddings applied directly to unpadded sequences.145 """146 147 def __init__(148 self,149 dim: int,150 base: float = 10000.0,151 max_seqlen: Optional[int] = None,152 device: Optional[torch.device] = None,153 dtype: Optional[torch.dtype] = None,154 ):155 """156 max_seqlen: if max_seqlen, device, and dtype are provided, we precompute the cos_sin_cache157 up to max_seqlen. If the max_seqlen, device, or dtype during training/inference differ,158 the cos_sin_cache will be recomputed during the forward pass.159 """160 super().__init__(dim=dim, base=base, device=device, interleaved=False)161 self.max_seqlen = max_seqlen162 163 if max_seqlen is not None and device is not None and dtype is not None:164 self._update_cos_sin_cache(max_seqlen, device=device, dtype=dtype)165 166 def forward(167 self,168 qkv: torch.Tensor,169 cu_seqlens: torch.Tensor,170 max_seqlen: Optional[int] = None,171 ) -> Union[torch.Tensor, tuple[torch.Tensor, torch.Tensor]]:172 """173 Apply rotary embedding *inplace* to qkv.174 qkv: (total_nnz, 3, nheads, headdim)175 cu_seqlens: (batch + 1,) cumulative sequence lengths176 max_seqlen: int max seq length in the batch177 """178 if max_seqlen is not None:179 self._update_cos_sin_cache(max_seqlen, device=qkv.device, dtype=qkv.dtype)180 181 qkv = apply_rotary_unpadded(182 qkv,183 self._cos_cached,184 self._sin_cached,185 cu_seqlens=cu_seqlens,186 max_seqlen=max_seqlen,187 )188 189 return qkv190 191 def extra_repr(self) -> str:192 return f"dim={self.dim}, base={self.base}, scale_base={self.scale_base}"193 194 195class ModernBertEmbeddings(nn.Module):196 """197 Same as BertEmbeddings with a tiny tweak for positional embeddings indexing.198 """199 200 def __init__(self, config: ModernBertConfig):201 super().__init__()202 self.config = config203 self.tok_embeddings = nn.Embedding(config.vocab_size, config.hidden_size, padding_idx=config.pad_token_id)204 self.norm = nn.LayerNorm(config.hidden_size, eps=config.norm_eps, bias=config.norm_bias)205 self.drop = nn.Dropout(config.embedding_dropout)206 207 @torch.compile(dynamic=True)208 def compiled_embeddings(self, input_ids: torch.LongTensor) -> torch.Tensor:209 return self.drop(self.norm(self.tok_embeddings(input_ids)))210 211 def forward(212 self, input_ids: Optional[torch.LongTensor] = None, inputs_embeds: Optional[torch.Tensor] = None213 ) -> torch.Tensor:214 if inputs_embeds is not None:215 hidden_states = self.drop(self.norm(inputs_embeds))216 else:217 hidden_states = (218 self.compiled_embeddings(input_ids)219 if self.config.reference_compile220 else self.drop(self.norm(self.tok_embeddings(input_ids)))221 )222 return hidden_states223 224 225class ModernBertMLP(nn.Module):226 """Applies the GLU at the end of each ModernBERT layer.227 228 Compared to the default BERT architecture, this block replaces :class:`~transformers.model.bert.modeling_bert.BertIntermediate`229 and :class:`~transformers.model.bert.modeling_bert.SelfOutput` with a single module that has similar functionality.230 """231 232 def __init__(self, config: ModernBertConfig):233 super().__init__()234 self.config = config235 self.Wi = nn.Linear(config.hidden_size, int(config.intermediate_size) * 2, bias=config.mlp_bias)236 self.act = ACT2FN[config.hidden_activation]237 self.drop = nn.Dropout(config.mlp_dropout)238 self.Wo = nn.Linear(config.intermediate_size, config.hidden_size, bias=config.mlp_bias)239 240 def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:241 input, gate = self.Wi(hidden_states).chunk(2, dim=-1)242 return self.Wo(self.drop(self.act(input) * gate))243 244 245class ModernBertRotaryEmbedding(nn.Module):246 inv_freq: torch.Tensor # fix linting for `register_buffer`247 248 def __init__(self, config: ModernBertConfig, device=None):249 super().__init__()250 # BC: "rope_type" was originally "type"251 if hasattr(config, "rope_scaling") and isinstance(config.rope_scaling, dict):252 self.rope_type = config.rope_scaling.get("rope_type", config.rope_scaling.get("type"))253 else:254 self.rope_type = "default"255 self.max_seq_len_cached = config.max_position_embeddings256 self.original_max_seq_len = config.max_position_embeddings257 258 self.config = config259 self.rope_init_fn = ROPE_INIT_FUNCTIONS[self.rope_type]260 261 inv_freq, self.attention_scaling = self.rope_init_fn(self.config, device)262 self.register_buffer("inv_freq", inv_freq, persistent=False)263 self.original_inv_freq = self.inv_freq264 265 @torch.no_grad()266 @dynamic_rope_update # power user: used with advanced RoPE types (e.g. dynamic rope)267 def forward(self, x, position_ids):268 inv_freq_expanded = self.inv_freq[None, :, None].float().expand(position_ids.shape[0], -1, 1).to(x.device)269 position_ids_expanded = position_ids[:, None, :].float()270 271 device_type = x.device.type if isinstance(x.device.type, str) and x.device.type != "mps" else "cpu"272 with torch.autocast(device_type=device_type, enabled=False): # Force float32273 freqs = (inv_freq_expanded.float() @ position_ids_expanded.float()).transpose(1, 2)274 emb = torch.cat((freqs, freqs), dim=-1)275 cos = emb.cos() * self.attention_scaling276 sin = emb.sin() * self.attention_scaling277 278 return cos.to(dtype=x.dtype), sin.to(dtype=x.dtype)279 280 281def rotate_half(x):282 """Rotates half the hidden dims of the input."""283 x1 = x[..., : x.shape[-1] // 2]284 x2 = x[..., x.shape[-1] // 2 :]285 return torch.cat((-x2, x1), dim=-1)286 287 288def apply_rotary_pos_emb(q, k, cos, sin, position_ids=None, unsqueeze_dim=1):289 """Applies Rotary Position Embedding to the query and key tensors.290 291 Args:292 q (`torch.Tensor`): The query tensor.293 k (`torch.Tensor`): The key tensor.294 cos (`torch.Tensor`): The cosine part of the rotary embedding.295 sin (`torch.Tensor`): The sine part of the rotary embedding.296 position_ids (`torch.Tensor`, *optional*):297 Deprecated and unused.298 unsqueeze_dim (`int`, *optional*, defaults to 1):299 The 'unsqueeze_dim' argument specifies the dimension along which to unsqueeze cos[position_ids] and300 sin[position_ids] so that they can be properly broadcasted to the dimensions of q and k. For example, note301 that cos[position_ids] and sin[position_ids] have the shape [batch_size, seq_len, head_dim]. Then, if q and302 k have the shape [batch_size, heads, seq_len, head_dim], then setting unsqueeze_dim=1 makes303 cos[position_ids] and sin[position_ids] broadcastable to the shapes of q and k. Similarly, if q and k have304 the shape [batch_size, seq_len, heads, head_dim], then set unsqueeze_dim=2.305 Returns:306 `tuple(torch.Tensor)` comprising of the query and key tensors rotated using the Rotary Position Embedding.307 """308 cos = cos.unsqueeze(unsqueeze_dim)309 sin = sin.unsqueeze(unsqueeze_dim)310 q_embed = (q * cos) + (rotate_half(q) * sin)311 k_embed = (k * cos) + (rotate_half(k) * sin)312 return q_embed, k_embed313 314 315def eager_attention_forward(316 module: "ModernBertAttention",317 qkv: torch.Tensor,318 attention_mask: torch.Tensor,319 sliding_window_mask: torch.Tensor,320 position_ids: Optional[torch.LongTensor],321 local_attention: tuple[int, int],322 bs: int,323 dim: int,324 output_attentions: Optional[bool] = False,325 **_kwargs,326) -> Union[tuple[torch.Tensor, torch.Tensor], tuple[torch.Tensor]]:327 # qkv: [batch_size, seqlen, 3, nheads, headdim]328 cos, sin = module.rotary_emb(qkv, position_ids=position_ids)329 query, key, value = qkv.transpose(3, 1).unbind(dim=2)330 # query, key, value: [batch_size, heads, seq_len, head_dim]331 query, key = apply_rotary_pos_emb(query, key, cos, sin)332 333 scale = module.head_dim**-0.5334 attn_weights = torch.matmul(query, key.transpose(2, 3)) * scale335 336 if local_attention != (-1, -1):337 attention_mask = sliding_window_mask338 339 attn_weights = attn_weights + attention_mask340 341 # upcast attention to fp32342 attn_weights = nn.functional.softmax(attn_weights, dim=-1, dtype=torch.float32).to(query.dtype)343 attn_weights = nn.functional.dropout(attn_weights, p=module.attention_dropout, training=module.training)344 attn_output = torch.matmul(attn_weights, value)345 attn_output = attn_output.transpose(1, 2).contiguous()346 attn_output = attn_output.view(bs, -1, dim)347 if output_attentions:348 return (attn_output, attn_weights)349 return (attn_output,)350 351 352def flash_attention_forward(353 module: "ModernBertAttention",354 qkv: torch.Tensor,355 rotary_emb: ModernBertUnpaddedRotaryEmbedding,356 cu_seqlens: torch.Tensor,357 max_seqlen: int,358 local_attention: tuple[int, int],359 bs: int,360 dim: int,361 target_dtype: torch.dtype = torch.bfloat16,362 **_kwargs,363) -> tuple[torch.Tensor]:364 # (total_seqlen, 3, nheads, headdim)365 qkv = rotary_emb(qkv, cu_seqlens=cu_seqlens, max_seqlen=max_seqlen)366 367 convert_dtype = qkv.dtype not in (torch.float16, torch.bfloat16)368 if convert_dtype:369 # FA2 implementation only supports fp16 and bf16. If FA2 is supported,370 # bfloat16 must be supported as of FA2 2.5.7. (Turing GPUs not supported)371 orig_dtype = qkv.dtype372 qkv = qkv.to(target_dtype)373 374 attn = flash_attn_varlen_qkvpacked_func(375 qkv,376 cu_seqlens=cu_seqlens,377 max_seqlen=max_seqlen,378 dropout_p=module.attention_dropout if module.training else 0.0,379 deterministic=module.deterministic_flash_attn,380 window_size=local_attention,381 )382 attn = attn.to(orig_dtype) # type: ignore383 else:384 attn = flash_attn_varlen_qkvpacked_func(385 qkv,386 cu_seqlens=cu_seqlens,387 max_seqlen=max_seqlen,388 dropout_p=module.attention_dropout if module.training else 0.0,389 deterministic=module.deterministic_flash_attn,390 window_size=local_attention,391 )392 return (attn.view(bs, dim),)393 394 395def sdpa_attention_forward(396 module: "ModernBertAttention",397 qkv: torch.Tensor,398 attention_mask: torch.Tensor,399 sliding_window_mask: torch.Tensor,400 position_ids: Optional[torch.LongTensor],401 local_attention: tuple[int, int],402 bs: int,403 dim: int,404 **_kwargs,405) -> tuple[torch.Tensor]:406 # qkv: [batch_size, seqlen, 3, nheads, headdim]407 cos, sin = module.rotary_emb(qkv, position_ids=position_ids)408 query, key, value = qkv.transpose(3, 1).unbind(dim=2)409 # query, key, value: [batch_size, heads, seq_len, head_dim]410 query, key = apply_rotary_pos_emb(query, key, cos, sin)411 412 if local_attention != (-1, -1):413 attention_mask = sliding_window_mask414 415 attn_output = (416 F.scaled_dot_product_attention(417 query,418 key,419 value,420 dropout_p=module.attention_dropout if module.training else 0.0,421 attn_mask=attention_mask,422 )423 .transpose(1, 2)424 .contiguous()425 )426 attn_output = attn_output.view(bs, -1, dim)427 return (attn_output,)428 429 430MODERNBERT_ATTENTION_FUNCTION = {431 "flash_attention_2": flash_attention_forward,432 "eager": eager_attention_forward,433 "sdpa": sdpa_attention_forward,434}435 436 437class ModernBertAttention(nn.Module):438 """Performs multi-headed self attention on a batch of unpadded sequences.439 440 If Flash Attention 2 is installed, this module uses Flash Attention to improve throughput.441 If Flash Attention 2 is not installed, the implementation will use PyTorch's SDPA kernel,442 which requires padding and unpadding inputs, adding some overhead.443 444 See `forward` method for additional details.445 """446 447 def __init__(self, config: ModernBertConfig, layer_id: Optional[int] = None):448 super().__init__()449 self.config = config450 self.layer_id = layer_id451 452 if config.hidden_size % config.num_attention_heads != 0:453 raise ValueError(454 f"The hidden size ({config.hidden_size}) is not a multiple of the number of attention heads ({config.num_attention_heads})"455 )456 457 self.attention_dropout = config.attention_dropout458 self.deterministic_flash_attn = config.deterministic_flash_attn459 self.num_heads = config.num_attention_heads460 self.head_dim = config.hidden_size // config.num_attention_heads461 self.all_head_size = self.head_dim * self.num_heads462 self.Wqkv = nn.Linear(config.hidden_size, 3 * self.all_head_size, bias=config.attention_bias)463 464 if layer_id % config.global_attn_every_n_layers != 0:465 self.local_attention = (config.local_attention // 2, config.local_attention // 2)466 rope_theta = config.local_rope_theta if config.local_rope_theta is not None else config.global_rope_theta467 max_position_embeddings = config.local_attention468 else:469 self.local_attention = (-1, -1)470 max_position_embeddings = config.max_position_embeddings471 rope_theta = config.global_rope_theta472 473 if config._attn_implementation == "flash_attention_2":474 self.rotary_emb = ModernBertUnpaddedRotaryEmbedding(475 dim=self.head_dim, max_seqlen=max_position_embeddings, base=rope_theta476 )477 else:478 config_copy = copy.deepcopy(config)479 config_copy.rope_theta = rope_theta480 self.rotary_emb = ModernBertRotaryEmbedding(config=config_copy)481 482 self.Wo = nn.Linear(config.hidden_size, config.hidden_size, bias=config.attention_bias)483 self.out_drop = nn.Dropout(config.attention_dropout) if config.attention_dropout > 0.0 else nn.Identity()484 self.pruned_heads = set()485 486 def forward(487 self,488 hidden_states: torch.Tensor,489 output_attentions: Optional[bool] = False,490 **kwargs,491 ) -> torch.Tensor:492 qkv = self.Wqkv(hidden_states)493 494 bs = hidden_states.shape[0]495 if self.config._attn_implementation == "flash_attention_2":496 qkv = qkv.view(-1, 3, self.num_heads, self.head_dim)497 else:498 qkv = qkv.view(bs, -1, 3, self.num_heads, self.head_dim)499 500 attn_outputs = MODERNBERT_ATTENTION_FUNCTION[self.config._attn_implementation](501 self,502 qkv=qkv,503 rotary_emb=self.rotary_emb,504 local_attention=self.local_attention,505 bs=bs,506 dim=self.all_head_size,507 output_attentions=output_attentions,508 **kwargs,509 )510 hidden_states = attn_outputs[0]511 hidden_states = self.out_drop(self.Wo(hidden_states))512 513 return (hidden_states,) + attn_outputs[1:] # add attentions if outputted514 515 516class ModernBertEncoderLayer(GradientCheckpointingLayer):517 def __init__(self, config: ModernBertConfig, layer_id: Optional[int] = None):518 super().__init__()519 self.config = config520 if layer_id == 0:521 self.attn_norm = nn.Identity()522 else:523 self.attn_norm = nn.LayerNorm(config.hidden_size, eps=config.norm_eps, bias=config.norm_bias)524 self.attn = ModernBertAttention(config=config, layer_id=layer_id)525 self.mlp_norm = nn.LayerNorm(config.hidden_size, eps=config.norm_eps, bias=config.norm_bias)526 self.mlp = ModernBertMLP(config)527 528 @torch.compile(dynamic=True)529 def compiled_mlp(self, hidden_states: torch.Tensor) -> torch.Tensor:530 return self.mlp(self.mlp_norm(hidden_states))531 532 def forward(533 self,534 hidden_states: torch.Tensor,535 attention_mask: Optional[torch.Tensor] = None,536 sliding_window_mask: Optional[torch.Tensor] = None,537 position_ids: Optional[torch.LongTensor] = None,538 cu_seqlens: Optional[torch.Tensor] = None,539 max_seqlen: Optional[int] = None,540 output_attentions: Optional[bool] = False,541 ) -> torch.Tensor:542 attn_outputs = self.attn(543 self.attn_norm(hidden_states),544 attention_mask=attention_mask,545 sliding_window_mask=sliding_window_mask,546 position_ids=position_ids,547 cu_seqlens=cu_seqlens,548 max_seqlen=max_seqlen,549 output_attentions=output_attentions,550 )551 hidden_states = hidden_states + attn_outputs[0]552 mlp_output = (553 self.compiled_mlp(hidden_states)554 if self.config.reference_compile555 else self.mlp(self.mlp_norm(hidden_states))556 )557 hidden_states = hidden_states + mlp_output558 559 return (hidden_states,) + attn_outputs[1:] # add attentions if outputted560 561 562@auto_docstring563class ModernBertPreTrainedModel(PreTrainedModel):564 config: ModernBertConfig565 base_model_prefix = "model"566 supports_gradient_checkpointing = True567 _no_split_modules = ["ModernBertEmbeddings", "ModernBertEncoderLayer"]568 _supports_flash_attn = True569 _supports_sdpa = True570 _supports_flex_attn = False571 572 def _init_weights(self, module: nn.Module):573 cutoff_factor = self.config.initializer_cutoff_factor574 if cutoff_factor is None:575 cutoff_factor = 3576 577 def init_weight(module: nn.Module, std: float):578 nn.init.trunc_normal_(579 module.weight,580 mean=0.0,581 std=std,582 a=-cutoff_factor * std,583 b=cutoff_factor * std,584 )585 586 if isinstance(module, nn.Linear):587 if module.bias is not None:588 nn.init.zeros_(module.bias)589 590 stds = {591 "in": self.config.initializer_range,592 "out": self.config.initializer_range / math.sqrt(2.0 * self.config.num_hidden_layers),593 "embedding": self.config.initializer_range,594 "final_out": self.config.hidden_size**-0.5,595 }596 597 if isinstance(module, ModernBertEmbeddings):598 init_weight(module.tok_embeddings, stds["embedding"])599 elif isinstance(module, ModernBertMLP):600 init_weight(module.Wi, stds["in"])601 init_weight(module.Wo, stds["out"])602 elif isinstance(module, ModernBertAttention):603 init_weight(module.Wqkv, stds["in"])604 init_weight(module.Wo, stds["out"])605 elif isinstance(module, ModernBertPredictionHead):606 init_weight(module.dense, stds["out"])607 elif isinstance(module, ModernBertForMaskedLM):608 init_weight(module.decoder, stds["out"])609 elif isinstance(610 module,611 (612 ModernBertForSequenceClassification,613 ModernBertForMultipleChoice,614 ModernBertForTokenClassification,615 ModernBertForQuestionAnswering,616 ),617 ):618 init_weight(module.classifier, stds["final_out"])619 elif isinstance(module, nn.LayerNorm):620 module.weight.data.fill_(1.0)621 if module.bias is not None:622 module.bias.data.zero_()623 624 def _check_and_adjust_attn_implementation(625 self, attn_implementation: Optional[str], is_init_check: bool = False626 ) -> str:627 """628 Checks and dispatches to hhe requested attention implementation.629 """630 # If the user didn't specify anything, try to use flash_attention_2 if available.631 # Otherwise we fall back to the default SDPA -> Eager from the super() method.632 # ModernBert's FA2 implementation correctly handles non-fp16/bf16 dtypes, we don't633 # need the FA2 warning for non-fp16/bf16 dtypes so we set fp16 for the FA2 check.634 635 try:636 attn_implementation = (637 "flash_attention_2"638 if attn_implementation is None and self._flash_attn_2_can_dispatch()639 else attn_implementation640 )641 except (ValueError, ImportError):642 pass643 return super()._check_and_adjust_attn_implementation(644 attn_implementation=attn_implementation, is_init_check=is_init_check645 )646 647 def _maybe_set_compile(self):648 if self.config.reference_compile is False:649 return650 651 if hasattr(self, "hf_device_map") and len(self.hf_device_map) > 1:652 if self.config.reference_compile:653 logger.warning_once(654 "If `accelerate` split the model across devices, `torch.compile` will not work. "655 "Falling back to non-compiled mode."656 )657 self.config.reference_compile = False658 659 if self.device.type == "mps":660 if self.config.reference_compile:661 logger.warning_once(662 "Compiling the model with `torch.compile` and using a `torch.mps` device is not supported. "663 "Falling back to non-compiled mode."664 )665 self.config.reference_compile = False666 667 if self.device.type == "cpu":668 if self.config.reference_compile:669 logger.warning_once(670 "Compiling the model with `torch.compile` and using a `torch.cpu` device is not supported. "671 "Falling back to non-compiled mode."672 )673 self.config.reference_compile = False674 675 if self.config.reference_compile is None:676 self.config.reference_compile = is_triton_available()677 678 def resize_token_embeddings(self, *args, **kwargs):679 model_embeds = super().resize_token_embeddings(*args, **kwargs)680 681 if self.config.reference_compile in {True, None}:682 if self.config.reference_compile:683 logger.warning_once(684 "Resizing token embeddings with `torch.compile` is not supported. Falling back to non-compiled mode."685 )686 self.config.reference_compile = False687 688 return model_embeds689 690 691def _unpad_modernbert_input(692 inputs: torch.Tensor,693 attention_mask: torch.Tensor,694 position_ids: Optional[torch.Tensor] = None,695 labels: Optional[torch.Tensor] = None,696) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, int, Optional[torch.Tensor], Optional[torch.Tensor]]:697 """698 Remove padding from input sequences.699 700 Args:701 inputs: (batch, seqlen, ...) or (batch, seqlen)702 attention_mask: (batch, seqlen), bool / int, 1 means valid and 0 means not valid.703 position_ids: (batch, seqlen), int, position ids704 labels: (batch, seqlen), int, labels705 706 Returns:707 unpadded_inputs: (total_nnz, ...), where total_nnz = number of tokens selected in attention_mask.708 indices: (total_nnz)709 cu_seqlens: (batch + 1), the cumulative sequence lengths710 max_seqlen_in_batch: int711 unpadded_position_ids: (total_nnz) or None712 unpadded_labels: (total_nnz) or None713 """714 seqlens_in_batch = attention_mask.sum(dim=-1, dtype=torch.int32)715 indices = torch.nonzero(attention_mask.flatten(), as_tuple=False).flatten()716 max_seqlen_in_batch = int(seqlens_in_batch.max().item())717 cu_seqlens = torch.nn.functional.pad(torch.cumsum(seqlens_in_batch, dim=0, dtype=torch.int32), (1, 0))718 719 if inputs.dim() == 2:720 unpadded_inputs = inputs.flatten()[indices]721 else:722 batch, seqlen, *rest = inputs.shape723 shape = batch * seqlen724 unpadded_inputs = inputs.view(shape, *rest)[indices]725 726 unpadded_position_ids = position_ids.flatten()[indices] if position_ids is not None else None727 unpadded_labels = labels.flatten()[indices] if labels is not None else None728 729 return unpadded_inputs, indices, cu_seqlens, max_seqlen_in_batch, unpadded_position_ids, unpadded_labels730 731 732def _pad_modernbert_output(733 inputs: torch.Tensor,734 indices: torch.Tensor,735 batch: int,736 seqlen: int,737) -> torch.Tensor:738 """739 Add padding to sequences.740 741 Args:742 inputs: (total_nnz, ...) or (total_nnz,), where total_nnz = number of tokens selected in attention_mask.743 indices: (total_nnz)744 batch: int, batch size745 seqlen: int, max sequence length746 747 Returns:748 padded_inputs: (batch, seqlen, ...) or (batch, seqlen)749 """750 if inputs.dim() == 1:751 output = torch.zeros(batch * seqlen, dtype=inputs.dtype, device=inputs.device)752 output[indices] = inputs753 padded_inputs = output.view(batch, seqlen)754 else:755 _, *rest = inputs.shape756 output = torch.zeros(batch * seqlen, *rest, dtype=inputs.dtype, device=inputs.device)757 output[indices] = inputs758 padded_inputs = output.view(batch, seqlen, *rest)759 760 return padded_inputs761 762 763@auto_docstring764class ModernBertModel(ModernBertPreTrainedModel):765 def __init__(self, config: ModernBertConfig):766 super().__init__(config)767 self.config = config768 self.embeddings = ModernBertEmbeddings(config)769 self.layers = nn.ModuleList(770 [ModernBertEncoderLayer(config, layer_id) for layer_id in range(config.num_hidden_layers)]771 )772 self.final_norm = nn.LayerNorm(config.hidden_size, eps=config.norm_eps, bias=config.norm_bias)773 self.gradient_checkpointing = False774 self.post_init()775 776 def get_input_embeddings(self):777 return self.embeddings.tok_embeddings778 779 def set_input_embeddings(self, value):780 self.embeddings.tok_embeddings = value781 782 @auto_docstring783 def forward(784 self,785 input_ids: Optional[torch.LongTensor] = None,786 attention_mask: Optional[torch.Tensor] = None,787 sliding_window_mask: Optional[torch.Tensor] = None,788 position_ids: Optional[torch.LongTensor] = None,789 inputs_embeds: Optional[torch.Tensor] = None,790 indices: Optional[torch.Tensor] = None,791 cu_seqlens: Optional[torch.Tensor] = None,792 max_seqlen: Optional[int] = None,793 batch_size: Optional[int] = None,794 seq_len: Optional[int] = None,795 output_attentions: Optional[bool] = None,796 output_hidden_states: Optional[bool] = None,797 return_dict: Optional[bool] = None,798 ) -> Union[tuple[torch.Tensor, ...], BaseModelOutput]:799 r"""800 sliding_window_mask (`torch.Tensor` of shape `(batch_size, sequence_length)`, *optional*):801 Mask to avoid performing attention on padding or far-away tokens. In ModernBert, only every few layers802 perform global attention, while the rest perform local attention. This mask is used to avoid attending to803 far-away tokens in the local attention layers when not using Flash Attention.804 indices (`torch.Tensor` of shape `(total_unpadded_tokens,)`, *optional*):805 Indices of the non-padding tokens in the input sequence. Used for unpadding the output.806 cu_seqlens (`torch.Tensor` of shape `(batch + 1,)`, *optional*):807 Cumulative sequence lengths of the input sequences. Used to index the unpadded tensors.808 max_seqlen (`int`, *optional*):809 Maximum sequence length in the batch excluding padding tokens. Used to unpad input_ids and pad output tensors.810 batch_size (`int`, *optional*):811 Batch size of the input sequences. Used to pad the output tensors.812 seq_len (`int`, *optional*):813 Sequence length of the input sequences including padding tokens. Used to pad the output tensors.814 """815 output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions816 output_hidden_states = (817 output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states818 )819 return_dict = return_dict if return_dict is not None else self.config.use_return_dict820 821 if (input_ids is None) ^ (inputs_embeds is not None):822 raise ValueError("You must specify exactly one of input_ids or inputs_embeds")823 824 all_hidden_states = () if output_hidden_states else None825 all_self_attentions = () if output_attentions else None826 827 self._maybe_set_compile()828 829 if input_ids is not None:830 self.warn_if_padding_and_no_attention_mask(input_ids, attention_mask)831 832 if batch_size is None and seq_len is None:833 if inputs_embeds is not None:834 batch_size, seq_len = inputs_embeds.shape[:2]835 else:836 batch_size, seq_len = input_ids.shape[:2]837 device = input_ids.device if input_ids is not None else inputs_embeds.device838 839 if attention_mask is None:840 attention_mask = torch.ones((batch_size, seq_len), device=device, dtype=torch.bool)841 842 repad = False843 if self.config._attn_implementation == "flash_attention_2":844 if indices is None and cu_seqlens is None and max_seqlen is None:845 repad = True846 if inputs_embeds is None:847 with torch.no_grad():848 input_ids, indices, cu_seqlens, max_seqlen, *_ = _unpad_modernbert_input(849 inputs=input_ids, attention_mask=attention_mask850 )851 else:852 inputs_embeds, indices, cu_seqlens, max_seqlen, *_ = _unpad_modernbert_input(853 inputs=inputs_embeds, attention_mask=attention_mask854 )855 else:856 if position_ids is None:857 position_ids = torch.arange(seq_len, device=device).unsqueeze(0)858 859 attention_mask, sliding_window_mask = self._update_attention_mask(860 attention_mask, output_attentions=output_attentions861 )862 863 hidden_states = self.embeddings(input_ids=input_ids, inputs_embeds=inputs_embeds)864 865 for encoder_layer in self.layers:866 if output_hidden_states:867 all_hidden_states = all_hidden_states + (hidden_states,)868 869 layer_outputs = encoder_layer(870 hidden_states,871 attention_mask=attention_mask,872 sliding_window_mask=sliding_window_mask,873 position_ids=position_ids,874 cu_seqlens=cu_seqlens,875 max_seqlen=max_seqlen,876 output_attentions=output_attentions,877 )878 hidden_states = layer_outputs[0]879 if output_attentions and len(layer_outputs) > 1:880 all_self_attentions = all_self_attentions + (layer_outputs[1],)881 882 if output_hidden_states:883 all_hidden_states = all_hidden_states + (hidden_states,)884 885 hidden_states = self.final_norm(hidden_states)886 887 if repad:888 hidden_states = _pad_modernbert_output(889 inputs=hidden_states, indices=indices, batch=batch_size, seqlen=seq_len890 )891 if all_hidden_states is not None:892 all_hidden_states = tuple(893 _pad_modernbert_output(inputs=hs, indices=indices, batch=batch_size, seqlen=seq_len)894 for hs in all_hidden_states895 )896 # If the attention implementation is FA2 and there is no need for repadding, there might still be the batch897 # dimension missing898 elif (899 self.config._attn_implementation == "flash_attention_2"900 and all_hidden_states is not None901 and all_hidden_states[-1].dim() == 2902 ):903 hidden_states = hidden_states.unsqueeze(0)904 all_hidden_states = tuple(hs.unsqueeze(0) for hs in all_hidden_states)905 906 if not return_dict:907 return tuple(v for v in [hidden_states, all_hidden_states, all_self_attentions] if v is not None)908 return BaseModelOutput(909 last_hidden_state=hidden_states,910 hidden_states=all_hidden_states,911 attentions=all_self_attentions,912 )913 914 def _update_attention_mask(self, attention_mask: torch.Tensor, output_attentions: bool) -> torch.Tensor:915 if output_attentions:916 if self.config._attn_implementation == "sdpa":917 logger.warning_once(918 "Outputting attentions is only supported with the 'eager' attention implementation, "919 'not with "sdpa". Falling back to `attn_implementation="eager"`.'920 )921 self.config._attn_implementation = "eager"922 elif self.config._attn_implementation != "eager":923 logger.warning_once(924 "Outputting attentions is only supported with the eager attention implementation, "925 f'not with {self.config._attn_implementation}. Consider setting `attn_implementation="eager"`.'926 " Setting `output_attentions=False`."927 )928 929 global_attention_mask = _prepare_4d_attention_mask(attention_mask, self.dtype)930 931 # Create position indices932 rows = torch.arange(global_attention_mask.shape[2]).unsqueeze(0)933 # Calculate distance between positions934 distance = torch.abs(rows - rows.T)935 936 # Create sliding window mask (1 for positions within window, 0 outside)937 window_mask = (938 (distance <= self.config.local_attention // 2).unsqueeze(0).unsqueeze(0).to(attention_mask.device)939 )940 # Combine with existing mask941 sliding_window_mask = global_attention_mask.masked_fill(window_mask.logical_not(), torch.finfo(self.dtype).min)942 943 return global_attention_mask, sliding_window_mask944 945 946class ModernBertPredictionHead(nn.Module):947 def __init__(self, config: ModernBertConfig):948 super().__init__()949 self.config = config950 self.dense = nn.Linear(config.hidden_size, config.hidden_size, config.classifier_bias)951 self.act = ACT2FN[config.classifier_activation]952 self.norm = nn.LayerNorm(config.hidden_size, eps=config.norm_eps, bias=config.norm_bias)953 954 def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:955 return self.norm(self.act(self.dense(hidden_states)))956 957 958@auto_docstring(959 custom_intro="""960 The ModernBert Model with a decoder head on top that is used for masked language modeling.961 """962)963class ModernBertForMaskedLM(ModernBertPreTrainedModel):964 _tied_weights_keys = ["decoder.weight"]965 966 def __init__(self, config: ModernBertConfig):967 super().__init__(config)968 self.config = config969 self.model = ModernBertModel(config)970 self.head = ModernBertPredictionHead(config)971 self.decoder = nn.Linear(config.hidden_size, config.vocab_size, bias=config.decoder_bias)972 973 self.sparse_prediction = self.config.sparse_prediction974 self.sparse_pred_ignore_index = self.config.sparse_pred_ignore_index975 976 # Initialize weights and apply final processing977 self.post_init()978 979 def get_output_embeddings(self):980 return self.decoder981 982 def set_output_embeddings(self, new_embeddings: nn.Linear):983 self.decoder = new_embeddings984 985 @torch.compile(dynamic=True)986 def compiled_head(self, output: torch.Tensor) -> torch.Tensor:987 return self.decoder(self.head(output))988 989 @auto_docstring990 def forward(991 self,992 input_ids: Optional[torch.LongTensor] = None,993 attention_mask: Optional[torch.Tensor] = None,994 sliding_window_mask: Optional[torch.Tensor] = None,995 position_ids: Optional[torch.Tensor] = None,996 inputs_embeds: Optional[torch.Tensor] = None,997 labels: Optional[torch.Tensor] = None,998 indices: Optional[torch.Tensor] = None,999 cu_seqlens: Optional[torch.Tensor] = None,1000 max_seqlen: Optional[int] = None,1001 batch_size: Optional[int] = None,1002 seq_len: Optional[int] = None,1003 output_attentions: Optional[bool] = None,1004 output_hidden_states: Optional[bool] = None,1005 return_dict: Optional[bool] = None,1006 **kwargs,1007 ) -> Union[tuple[torch.Tensor], MaskedLMOutput]:1008 r"""1009 sliding_window_mask (`torch.Tensor` of shape `(batch_size, sequence_length)`, *optional*):1010 Mask to avoid performing attention on padding or far-away tokens. In ModernBert, only every few layers1011 perform global attention, while the rest perform local attention. This mask is used to avoid attending to1012 far-away tokens in the local attention layers when not using Flash Attention.1013 indices (`torch.Tensor` of shape `(total_unpadded_tokens,)`, *optional*):1014 Indices of the non-padding tokens in the input sequence. Used for unpadding the output.1015 cu_seqlens (`torch.Tensor` of shape `(batch + 1,)`, *optional*):1016 Cumulative sequence lengths of the input sequences. Used to index the unpadded tensors.1017 max_seqlen (`int`, *optional*):1018 Maximum sequence length in the batch excluding padding tokens. Used to unpad input_ids and pad output tensors.1019 batch_size (`int`, *optional*):1020 Batch size of the input sequences. Used to pad the output tensors.1021 seq_len (`int`, *optional*):1022 Sequence length of the input sequences including padding tokens. Used to pad the output tensors.1023 """1024 return_dict = return_dict if return_dict is not None else self.config.use_return_dict1025 self._maybe_set_compile()1026 1027 if self.config._attn_implementation == "flash_attention_2":1028 if indices is None and cu_seqlens is None and max_seqlen is None:1029 if batch_size is None and seq_len is None:1030 if inputs_embeds is not None:1031 batch_size, seq_len = inputs_embeds.shape[:2]1032 else:1033 batch_size, seq_len = input_ids.shape[:2]1034 device = input_ids.device if input_ids is not None else inputs_embeds.device1035 1036 if attention_mask is None:1037 attention_mask = torch.ones((batch_size, seq_len), device=device, dtype=torch.bool)1038 1039 if inputs_embeds is None:1040 with torch.no_grad():1041 input_ids, indices, cu_seqlens, max_seqlen, position_ids, labels = _unpad_modernbert_input(1042 inputs=input_ids, attention_mask=attention_mask, position_ids=position_ids, labels=labels1043 )1044 else:1045 inputs_embeds, indices, cu_seqlens, max_seqlen, position_ids, labels = _unpad_modernbert_input(1046 inputs=inputs_embeds, attention_mask=attention_mask, position_ids=position_ids, labels=labels1047 )1048 1049 outputs = self.model(1050 input_ids=input_ids,1051 attention_mask=attention_mask,1052 sliding_window_mask=sliding_window_mask,1053 position_ids=position_ids,1054 inputs_embeds=inputs_embeds,1055 indices=indices,1056 cu_seqlens=cu_seqlens,1057 max_seqlen=max_seqlen,1058 batch_size=batch_size,1059 seq_len=seq_len,1060 output_attentions=output_attentions,1061 output_hidden_states=output_hidden_states,1062 return_dict=return_dict,1063 )1064 last_hidden_state = outputs[0]1065 1066 if self.sparse_prediction and labels is not None:1067 # flatten labels and output first1068 labels = labels.view(-1)1069 last_hidden_state = last_hidden_state.view(labels.shape[0], -1)1070 1071 # then filter out the non-masked tokens1072 mask_tokens = labels != self.sparse_pred_ignore_index1073 last_hidden_state = last_hidden_state[mask_tokens]1074 labels = labels[mask_tokens]1075 1076 logits = (1077 self.compiled_head(last_hidden_state)1078 if self.config.reference_compile1079 else self.decoder(self.head(last_hidden_state))1080 )1081 1082 loss = None1083 if labels is not None:1084 loss = self.loss_function(logits, labels, vocab_size=self.config.vocab_size, **kwargs)1085 1086 if self.config._attn_implementation == "flash_attention_2":1087 # Logits padding1088 with nullcontext() if self.config.repad_logits_with_grad or labels is None else torch.no_grad():1089 logits = _pad_modernbert_output(inputs=logits, indices=indices, batch=batch_size, seqlen=seq_len)1090 # Hidden states padding1091 if getattr(outputs, "hidden_states", None) is not None:1092 padded_hidden_states = []1093 for hs in outputs.hidden_states:1094 if hs.dim() == 3 and hs.shape[0] == 1:1095 hs = hs.squeeze(0)1096 padded_hidden_states.append(1097 _pad_modernbert_output(inputs=hs, indices=indices, batch=batch_size, seqlen=seq_len)1098 )1099 outputs.hidden_states = tuple(padded_hidden_states)1100 1101 if not return_dict:1102 output = (logits,)1103 return ((loss,) + output) if loss is not None else output1104 1105 return MaskedLMOutput(1106 loss=loss,1107 logits=logits,1108 hidden_states=outputs.hidden_states,1109 attentions=outputs.attentions,1110 )1111 1112 1113@auto_docstring(1114 custom_intro="""1115 The ModernBert Model with a sequence classification head on top that performs pooling.1116 """1117)1118class ModernBertForSequenceClassification(ModernBertPreTrainedModel):1119 def __init__(self, config: ModernBertConfig):1120 super().__init__(config)1121 self.num_labels = config.num_labels1122 self.config = config1123 1124 self.model = ModernBertModel(config)1125 self.head = ModernBertPredictionHead(config)1126 self.drop = torch.nn.Dropout(config.classifier_dropout)1127 self.classifier = nn.Linear(config.hidden_size, config.num_labels)1128 1129 # Initialize weights and apply final processing1130 self.post_init()1131 1132 @auto_docstring1133 def forward(1134 self,1135 input_ids: Optional[torch.LongTensor] = None,1136 attention_mask: Optional[torch.Tensor] = None,1137 sliding_window_mask: Optional[torch.Tensor] = None,1138 position_ids: Optional[torch.Tensor] = None,1139 inputs_embeds: Optional[torch.Tensor] = None,1140 labels: Optional[torch.Tensor] = None,1141 indices: Optional[torch.Tensor] = None,1142 cu_seqlens: Optional[torch.Tensor] = None,1143 max_seqlen: Optional[int] = None,1144 batch_size: Optional[int] = None,1145 seq_len: Optional[int] = None,1146 output_attentions: Optional[bool] = None,1147 output_hidden_states: Optional[bool] = None,1148 return_dict: Optional[bool] = None,1149 **kwargs,1150 ) -> Union[tuple[torch.Tensor], SequenceClassifierOutput]:1151 r"""1152 sliding_window_mask (`torch.Tensor` of shape `(batch_size, sequence_length)`, *optional*):1153 Mask to avoid performing attention on padding or far-away tokens. In ModernBert, only every few layers1154 perform global attention, while the rest perform local attention. This mask is used to avoid attending to1155 far-away tokens in the local attention layers when not using Flash Attention.1156 labels (`torch.LongTensor` of shape `(batch_size,)`, *optional*):1157 Labels for computing the sequence classification/regression loss. Indices should be in `[0, ...,1158 config.num_labels - 1]`. If `config.num_labels == 1` a regression loss is computed (Mean-Square loss), If1159 `config.num_labels > 1` a classification loss is computed (Cross-Entropy).1160 indices (`torch.Tensor` of shape `(total_unpadded_tokens,)`, *optional*):1161 Indices of the non-padding tokens in the input sequence. Used for unpadding the output.1162 cu_seqlens (`torch.Tensor` of shape `(batch + 1,)`, *optional*):1163 Cumulative sequence lengths of the input sequences. Used to index the unpadded tensors.1164 max_seqlen (`int`, *optional*):1165 Maximum sequence length in the batch excluding padding tokens. Used to unpad input_ids and pad output tensors.1166 batch_size (`int`, *optional*):1167 Batch size of the input sequences. Used to pad the output tensors.1168 seq_len (`int`, *optional*):1169 Sequence length of the input sequences including padding tokens. Used to pad the output tensors.1170 """1171 return_dict = return_dict if return_dict is not None else self.config.use_return_dict1172 self._maybe_set_compile()1173 1174 if input_ids is not None:1175 self.warn_if_padding_and_no_attention_mask(input_ids, attention_mask)1176 1177 if batch_size is None and seq_len is None:1178 if inputs_embeds is not None:1179 batch_size, seq_len = inputs_embeds.shape[:2]1180 else:1181 batch_size, seq_len = input_ids.shape[:2]1182 device = input_ids.device if input_ids is not None else inputs_embeds.device1183 1184 if attention_mask is None:1185 attention_mask = torch.ones((batch_size, seq_len), device=device, dtype=torch.bool)1186 1187 outputs = self.model(1188 input_ids=input_ids,1189 attention_mask=attention_mask,1190 sliding_window_mask=sliding_window_mask,1191 position_ids=position_ids,1192 inputs_embeds=inputs_embeds,1193 indices=indices,1194 cu_seqlens=cu_seqlens,1195 max_seqlen=max_seqlen,1196 batch_size=batch_size,1197 seq_len=seq_len,1198 output_attentions=output_attentions,1199 output_hidden_states=output_hidden_states,1200 return_dict=return_dict,