lifeofcoding/mastermax-7b
124
1# port of models described in RW2# We use the bloom model as a starting point for these model.3# Please refer to the bloom models for usage instructions.4 5import math6import warnings7from typing import Optional, Tuple, Union8 9import torch10import torch.utils.checkpoint11from torch import nn12from torch.nn import BCEWithLogitsLoss, CrossEntropyLoss, LayerNorm, MSELoss13from torch.nn import functional as F14 15from transformers.modeling_outputs import (16 BaseModelOutputWithPastAndCrossAttentions,17 CausalLMOutputWithCrossAttentions,18 QuestionAnsweringModelOutput,19 SequenceClassifierOutputWithPast,20 TokenClassifierOutput,21)22from transformers.modeling_utils import PreTrainedModel23from transformers.utils import logging24from .configuration_RW import RWConfig25 26logger = logging.get_logger(__name__)27 28# NOTE(Hesslow): Unfortunately we did not fuse matmul and bias during training, this means that there's one additional quantization to bfloat16 between the operations.29# In order not to degrade the quality of our HF-port, we keep these characteristics in the final model.30class Linear(nn.Linear):31 def forward(self, input: torch.Tensor) -> torch.Tensor:32 ret = input @ self.weight.T33 if self.bias is None:34 return ret35 else:36 return ret + self.bias37 38 39from einops import rearrange40 41# rotary pos emb helpers (torch.jit.script does not seem to support staticmethod...)42def rotate_half(x):43 x1, x2 = x[..., : x.shape[-1] // 2], x[..., x.shape[-1] // 2 :]44 return torch.cat((-x2, x1), dim=x1.ndim - 1) # dim=-1 triggers a bug in torch < 1.8.045 46 47class RotaryEmbedding(torch.nn.Module):48 """Implementation of RotaryEmbedding from GPT-NeoX.49 This implementation is design to operate on queries and keys that are compatible with50 [batch_size, n_heads_per_partition, seq_len, head_dim] (e.g. MinGPTAttention format).51 """52 53 def __init__(54 self,55 head_dim: int,56 base=10000,57 ):58 super().__init__()59 inv_freq = 1.0 / (base ** (torch.arange(0, head_dim, 2).float() / head_dim))60 self.register_buffer("inv_freq", inv_freq, persistent=False)61 self.head_dim = head_dim62 self.seq_len_cached = None63 self.batch_size_cached = None64 self.cos_cached: torch.Tensor | None = None65 self.sin_cached: torch.Tensor | None = None66 67 def cos_sin(68 self,69 seq_len: int,70 device="cuda",71 dtype=torch.bfloat16,72 ) -> torch.Tensor:73 if seq_len != self.seq_len_cached:74 self.seq_len_cached = seq_len75 t = torch.arange(seq_len, device=device).type_as(self.inv_freq)76 freqs = torch.einsum("i,j->ij", t, self.inv_freq)77 emb = torch.cat((freqs, freqs), dim=-1).to(device)78 79 if dtype in [torch.float16, torch.bfloat16]:80 emb = emb.float()81 82 self.cos_cached = emb.cos()[None, :, :]83 self.sin_cached = emb.sin()[None, :, :]84 85 self.cos_cached = self.cos_cached.type(dtype)86 self.sin_cached = self.sin_cached.type(dtype)87 88 return self.cos_cached, self.sin_cached89 90 def forward(self, q, k):91 batch, seq_len, head_dim = q.shape92 cos, sin = self.cos_sin(seq_len, q.device, q.dtype)93 return (q * cos) + (rotate_half(q) * sin), (k * cos) + (rotate_half(k) * sin)94 95 96def _make_causal_mask(97 input_ids_shape: torch.Size, device: torch.device, past_key_values_length: int98) -> torch.BoolTensor:99 batch_size, target_length = input_ids_shape100 mask = torch.empty((target_length, target_length + past_key_values_length), dtype=torch.bool, device=device)101 # ONNX doesn't support `torch.Tensor.triu` properly, thus we use this workaround102 seq_ids = torch.arange(target_length, device=device)103 mask[:, past_key_values_length:] = seq_ids[:, None] < seq_ids[None, :]104 105 if past_key_values_length > 0:106 mask[:, :past_key_values_length] = False107 108 expanded_mask = mask[None, None, :, :].expand(batch_size, 1, target_length, target_length + past_key_values_length)109 return expanded_mask110 111 112def _expand_mask(mask: torch.Tensor, tgt_length: int) -> torch.BoolTensor:113 batch_size, src_length = mask.shape114 tgt_length = tgt_length if tgt_length is not None else src_length115 116 expanded_mask = ~(mask[:, None, None, :].to(torch.bool))117 return expanded_mask.expand(batch_size, 1, tgt_length, src_length)118 119 120def build_alibi_tensor(attention_mask: torch.Tensor, num_heads: int, dtype: torch.dtype) -> torch.Tensor:121 batch_size, seq_length = attention_mask.shape122 closest_power_of_2 = 2 ** math.floor(math.log2(num_heads))123 base = torch.tensor(124 2 ** (-(2 ** -(math.log2(closest_power_of_2) - 3))), device=attention_mask.device, dtype=torch.float32125 )126 powers = torch.arange(1, 1 + closest_power_of_2, device=attention_mask.device, dtype=torch.int32)127 slopes = torch.pow(base, powers)128 129 if closest_power_of_2 != num_heads:130 extra_base = torch.tensor(131 2 ** (-(2 ** -(math.log2(2 * closest_power_of_2) - 3))), device=attention_mask.device, dtype=torch.float32132 )133 num_remaining_heads = min(closest_power_of_2, num_heads - closest_power_of_2)134 extra_powers = torch.arange(1, 1 + 2 * num_remaining_heads, 2, device=attention_mask.device, dtype=torch.int32)135 slopes = torch.cat([slopes, torch.pow(extra_base, extra_powers)], dim=0)136 137 # Note: alibi will added to the attention bias that will be applied to the query, key product of attention138 # => therefore alibi will have to be of shape (batch_size, num_heads, query_length, key_length)139 # => here we set (batch_size=1, num_heads=num_heads, query_length=1, key_length=max_length)140 # => the query_length dimension will then be broadcasted correctly141 # This is more or less identical to T5's relative position bias:142 # https://github.com/huggingface/transformers/blob/f681437203baa7671de3174b0fa583c349d9d5e1/src/transformers/models/t5/modeling_t5.py#L527143 arange_tensor = ((attention_mask.cumsum(dim=-1) - 1) * attention_mask)[:, None, :]144 alibi = slopes[..., None].bfloat16() * arange_tensor145 return alibi.reshape(batch_size * num_heads, 1, seq_length).to(dtype)146 147 148def dropout_add(x: torch.Tensor, residual: torch.Tensor, prob: float, training: bool) -> torch.Tensor:149 out = F.dropout(x, p=prob, training=training)150 out = residual + out151 return out152 153 154class Attention(nn.Module):155 def __init__(self, config: RWConfig):156 super().__init__()157 158 self.hidden_size = config.hidden_size159 self.num_heads = config.n_head160 self.head_dim = self.hidden_size // self.num_heads161 self.split_size = self.hidden_size162 self.hidden_dropout = config.hidden_dropout163 164 if self.head_dim * self.num_heads != self.hidden_size:165 raise ValueError(166 f"`hidden_size` must be divisible by num_heads (got `hidden_size`: {self.hidden_size} and `num_heads`:"167 f" {self.num_heads})."168 )169 170 self.maybe_rotary = RotaryEmbedding(config.head_dim) if config.rotary else lambda q, k: (q, k)171 172 # Layer-wise attention scaling173 self.inv_norm_factor = 1.0 / math.sqrt(self.head_dim)174 self.beta = self.inv_norm_factor175 176 self.query_key_value = Linear(177 self.hidden_size,178 3 * self.hidden_size if not config.multi_query else (self.hidden_size + 2 * self.head_dim),179 bias=config.bias,180 )181 self.multi_query = config.multi_query182 self.dense = Linear(self.hidden_size, self.hidden_size, bias=config.bias)183 self.attention_dropout = nn.Dropout(config.attention_dropout)184 self.num_kv = config.n_head if not self.multi_query else 1185 186 def _split_heads(self, fused_qkv: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]:187 """188 Split the last dimension into (num_heads, head_dim) without making any copies, results share same memory189 storage as `fused_qkv`190 191 Args:192 fused_qkv (`torch.tensor`, *required*): [batch_size, seq_length, num_heads * 3 * head_dim]193 194 Returns:195 query: [batch_size, seq_length, num_heads, head_dim] key: [batch_size, seq_length, num_heads, head_dim]196 value: [batch_size, seq_length, num_heads, head_dim]197 """198 if not self.multi_query:199 batch_size, seq_length, three_times_hidden_size = fused_qkv.shape200 fused_qkv = fused_qkv.view(batch_size, seq_length, self.num_heads, 3, self.head_dim)201 return fused_qkv[..., 0, :], fused_qkv[..., 1, :], fused_qkv[..., 2, :]202 else:203 batch_size, seq_length, three_times_hidden_size = fused_qkv.shape204 fused_qkv = fused_qkv.view(batch_size, seq_length, self.num_heads + 2, self.head_dim)205 return fused_qkv[..., :-2, :], fused_qkv[..., [-2], :], fused_qkv[..., [-1], :]206 207 def _merge_heads(self, x: torch.Tensor) -> torch.Tensor:208 """209 Merge heads together over the last dimenstion210 211 Args:212 x: (`torch.tensor`, *required*): [batch_size * num_heads, seq_length, head_dim]213 214 Returns:215 torch.tensor: [batch_size, seq_length, num_heads * head_dim]216 """217 # What we want to achieve is:218 # batch_size * num_heads, seq_length, head_dim -> batch_size, seq_length, num_heads * head_dim219 batch_size_and_num_heads, seq_length, _ = x.shape220 batch_size = batch_size_and_num_heads // self.num_heads221 222 # First view to decompose the batch size223 # batch_size * num_heads, seq_length, head_dim -> batch_size, num_heads, seq_length, head_dim224 x = x.view(batch_size, self.num_heads, seq_length, self.head_dim)225 226 # batch_size, num_heads, seq_length, head_dim -> batch_size, seq_length, num_heads, head_dim227 x = x.permute(0, 2, 1, 3)228 229 # batch_size, seq_length, num_heads, head_dim -> batch_size, seq_length, num_heads * head_dim230 return x.reshape(batch_size, seq_length, self.num_heads * self.head_dim)231 232 def forward(233 self,234 hidden_states: torch.Tensor,235 alibi: torch.Tensor,236 attention_mask: torch.Tensor,237 layer_past: Optional[Tuple[torch.Tensor, torch.Tensor]] = None,238 head_mask: Optional[torch.Tensor] = None,239 use_cache: bool = False,240 output_attentions: bool = False,241 ):242 fused_qkv = self.query_key_value(hidden_states) # [batch_size, seq_length, 3 x hidden_size]243 244 # 3 x [batch_size, seq_length, num_heads, head_dim]245 (query_layer, key_layer, value_layer) = self._split_heads(fused_qkv)246 247 batch_size, q_length, _, _ = query_layer.shape248 249 query_layer = query_layer.transpose(1, 2).reshape(batch_size * self.num_heads, q_length, self.head_dim)250 key_layer = key_layer.transpose(1, 2).reshape(251 batch_size * self.num_kv,252 q_length,253 self.head_dim,254 )255 value_layer = value_layer.transpose(1, 2).reshape(batch_size * self.num_kv, q_length, self.head_dim)256 257 query_layer, key_layer = self.maybe_rotary(query_layer, key_layer)258 259 if layer_past is not None:260 past_key, past_value = layer_past261 # concatenate along seq_length dimension:262 # - key: [batch_size * self.num_heads, head_dim, kv_length]263 # - value: [batch_size * self.num_heads, kv_length, head_dim]264 key_layer = torch.cat((past_key, key_layer), dim=1)265 value_layer = torch.cat((past_value, value_layer), dim=1)266 267 _, kv_length, _ = key_layer.shape268 269 if use_cache is True:270 present = (key_layer, value_layer)271 else:272 present = None273 274 if alibi is None:275 query_layer_ = query_layer.reshape(batch_size, self.num_heads, -1, self.head_dim)276 key_layer_ = key_layer.reshape(batch_size, self.num_kv, -1, self.head_dim)277 value_layer_ = value_layer.reshape(batch_size, self.num_kv, -1, self.head_dim)278 279 attn_output = F.scaled_dot_product_attention(280 query_layer_, key_layer_, value_layer_, None, 0.0, is_causal=True281 )282 283 x = attn_output.view(batch_size, self.num_heads, q_length, self.head_dim)284 x = x.permute(0, 2, 1, 3)285 attn_output = x.reshape(batch_size, q_length, self.num_heads * self.head_dim)286 287 output_tensor = self.dense(attn_output)288 289 outputs = (output_tensor, present)290 assert not output_attentions # not supported.291 return outputs292 else:293 attention_mask_float = (attention_mask * 1.0).masked_fill(attention_mask, -1e9).to(torch.bfloat16)294 matmul_result = query_layer @ key_layer.transpose(-1, -2)295 296 # change view to [batch_size, num_heads, q_length, kv_length]297 attention_scores = matmul_result.view(batch_size, self.num_heads, q_length, kv_length)298 299 # cast attention scores to fp32, compute scaled softmax and cast back to initial dtype - [batch_size, num_heads, q_length, kv_length]300 input_dtype = attention_scores.dtype301 # `float16` has a minimum value of -65504.0, whereas `bfloat16` and `float32` have a minimum value of `-3.4e+38`302 if input_dtype == torch.float16 or input_dtype == torch.bfloat16:303 attention_scores = attention_scores.to(torch.float32)304 # attn_weights = torch.masked_fill(attention_scores, attention_mask, torch.finfo(attention_scores.dtype).min)305 attention_probs = F.softmax(306 (attention_scores + alibi.view(batch_size, self.num_heads, 1, -1)) * self.inv_norm_factor + attention_mask_float,307 dim=-1,308 dtype=hidden_states.dtype,309 )310 # [batch_size, num_heads, q_length, kv_length]311 attention_probs = self.attention_dropout(attention_probs)312 313 if head_mask is not None:314 attention_probs = attention_probs * head_mask315 316 # change view [batch_size x num_heads, q_length, kv_length]317 attention_probs_reshaped = attention_probs.view(batch_size * self.num_heads, q_length, kv_length)318 319 # matmul: [batch_size * num_heads, q_length, head_dim]320 context_layer = attention_probs_reshaped @ value_layer321 322 # change view [batch_size, num_heads, q_length, head_dim]323 context_layer = self._merge_heads(context_layer)324 325 output_tensor = self.dense(context_layer)326 327 outputs = (output_tensor, present)328 if output_attentions:329 outputs += (attention_probs,)330 331 return outputs332 333 334class MLP(nn.Module):335 def __init__(self, config: RWConfig):336 super().__init__()337 hidden_size = config.hidden_size338 339 self.dense_h_to_4h = Linear(hidden_size, 4 * hidden_size, bias=config.bias)340 self.act = nn.GELU()341 self.dense_4h_to_h = Linear(4 * hidden_size, hidden_size, bias=config.bias)342 self.hidden_dropout = config.hidden_dropout343 344 def forward(self, x: torch.Tensor) -> torch.Tensor:345 x = self.act(self.dense_h_to_4h(x))346 x = self.dense_4h_to_h(x)347 return x348 349 350class DecoderLayer(nn.Module):351 def __init__(self, config: RWConfig):352 super().__init__()353 hidden_size = config.hidden_size354 355 self.input_layernorm = LayerNorm(hidden_size, eps=config.layer_norm_epsilon)356 self.num_heads = config.n_head357 self.self_attention = Attention(config)358 359 if not config.parallel_attn:360 # unused if parallel attn361 self.post_attention_layernorm = LayerNorm(hidden_size, eps=config.layer_norm_epsilon)362 363 self.mlp = MLP(config)364 365 self.apply_residual_connection_post_layernorm = config.apply_residual_connection_post_layernorm366 self.hidden_dropout = config.hidden_dropout367 368 self.config = config369 370 def forward(371 self,372 hidden_states: torch.Tensor,373 alibi: torch.Tensor,374 attention_mask: torch.Tensor,375 layer_past: Optional[Tuple[torch.Tensor, torch.Tensor]] = None,376 head_mask: Optional[torch.Tensor] = None,377 use_cache: bool = False,378 output_attentions: bool = False,379 ):380 381 layernorm_output = self.input_layernorm(hidden_states)382 residual = hidden_states383 384 # Self attention.385 attn_outputs = self.self_attention(386 layernorm_output,387 layer_past=layer_past,388 attention_mask=attention_mask,389 alibi=alibi,390 head_mask=head_mask,391 use_cache=use_cache,392 output_attentions=output_attentions,393 )394 395 attention_output = attn_outputs[0]396 397 if not self.config.parallel_attn:398 residual = dropout_add(attention_output, residual, self.config.attention_dropout, training=self.training)399 layernorm_output = self.post_attention_layernorm(residual)400 401 outputs = attn_outputs[1:]402 403 # MLP.404 mlp_output = self.mlp(layernorm_output)405 406 if self.config.parallel_attn:407 mlp_output += attention_output408 409 output = dropout_add(mlp_output, residual, self.config.hidden_dropout, training=self.training)410 411 if use_cache:412 outputs = (output,) + outputs413 else:414 outputs = (output,) + outputs[1:]415 416 return outputs # hidden_states, present, attentions417 418 419class RWPreTrainedModel(PreTrainedModel):420 _keys_to_ignore_on_load_missing = [r"h.*.self_attention.scale_mask_softmax.causal_mask", r"lm_head.weight"]421 """422 An abstract class to handle weights initialization and a simple interface for downloading and loading pretrained423 models.424 """425 426 config_class = RWConfig427 base_model_prefix = "transformer"428 supports_gradient_checkpointing = True429 _no_split_modules = ["DecoderLayer"]430 431 def __init__(self, *inputs, **kwargs):432 super().__init__(*inputs, **kwargs)433 434 def _init_weights(self, module: nn.Module):435 """Initialize the weights."""436 if isinstance(module, nn.Linear) or isinstance(module, Linear):437 # Slightly different from the TF version which uses truncated_normal for initialization438 # cf https://github.com/pytorch/pytorch/pull/5617439 module.weight.data.normal_(mean=0.0, std=self.config.initializer_range)440 if module.bias is not None:441 module.bias.data.zero_()442 elif isinstance(module, nn.Embedding):443 module.weight.data.normal_(mean=0.0, std=self.config.initializer_range)444 if module.padding_idx is not None:445 module.weight.data[module.padding_idx].zero_()446 elif isinstance(module, LayerNorm):447 module.bias.data.zero_()448 module.weight.data.fill_(1.0)449 450 def _set_gradient_checkpointing(self, module: nn.Module, value: bool = False):451 if isinstance(module, RWModel):452 module.gradient_checkpointing = value453 454 @staticmethod455 def _convert_to_standard_cache(456 past_key_value: Tuple[Tuple[torch.Tensor, torch.Tensor]], batch_size: int457 ) -> Tuple[Tuple[torch.Tensor, torch.Tensor]]:458 """459 Standardizes the format of the cache so as to match most implementations, i.e. to tuple(tuple([batch_size,460 num_heads, ...]))461 """462 batch_size_times_num_heads, head_dim, seq_length = past_key_value[0][0].shape463 num_heads = batch_size_times_num_heads // batch_size464 # key: [batch_size * num_heads, head_dim, seq_length] -> [batch_size, num_heads, head_dim, seq_length]465 # value: [batch_size * num_heads, seq_length, head_dim] -> [batch_size, num_heads, seq_length, head_dim]466 return tuple(467 (468 layer_past[0].view(batch_size, num_heads, head_dim, seq_length),469 layer_past[1].view(batch_size, num_heads, seq_length, head_dim),470 )471 for layer_past in past_key_value472 )473 474 @staticmethod475 def _convert_to_rw_cache(476 past_key_value: Tuple[Tuple[torch.Tensor, torch.Tensor]]477 ) -> Tuple[Tuple[torch.Tensor, torch.Tensor]]:478 batch_size, num_heads, head_dim, seq_length = past_key_value[0][0].shape479 batch_size_times_num_heads = batch_size * num_heads480 # key: [batch_size, num_heads, head_dim, seq_length] -> [batch_size * num_heads, head_dim, seq_length]481 # value: [batch_size, num_heads, seq_length, head_dim] -> [batch_size * num_heads, seq_length, head_dim]482 return tuple(483 (484 layer_past[0].view(batch_size_times_num_heads, head_dim, seq_length),485 layer_past[1].view(batch_size_times_num_heads, seq_length, head_dim),486 )487 for layer_past in past_key_value488 )489 490 491class RWModel(RWPreTrainedModel):492 def __init__(self, config: RWConfig):493 super().__init__(config)494 495 self.embed_dim = config.hidden_size496 self.num_heads = config.n_head497 self.alibi = config.alibi498 499 # Embedding + LN Embedding500 self.word_embeddings = nn.Embedding(config.vocab_size, self.embed_dim)501 502 # Transformer blocks503 self.h = nn.ModuleList([DecoderLayer(config) for _ in range(config.num_hidden_layers)])504 505 # Final Layer Norm506 self.ln_f = LayerNorm(self.embed_dim, eps=config.layer_norm_epsilon)507 508 self.gradient_checkpointing = False509 510 # Initialize weights and apply final processing511 self.post_init()512 513 def get_input_embeddings(self):514 return self.word_embeddings515 516 def _prepare_attn_mask(517 self, attention_mask: torch.Tensor, input_shape: Tuple[int, int], past_key_values_length: int518 ) -> torch.BoolTensor:519 # create causal mask520 # [batch_size, seq_length] -> [batch_size, 1, tgt_length, src_length]521 combined_attention_mask = None522 device = attention_mask.device523 _, src_length = input_shape524 525 if src_length > 1:526 combined_attention_mask = _make_causal_mask(527 input_shape, device=device, past_key_values_length=past_key_values_length528 )529 530 # [batch_size, seq_length] -> [batch_size, 1, tgt_length, src_length]531 expanded_attn_mask = _expand_mask(attention_mask, tgt_length=src_length)532 combined_attention_mask = (533 expanded_attn_mask if combined_attention_mask is None else expanded_attn_mask | combined_attention_mask534 )535 536 return combined_attention_mask537 538 def set_input_embeddings(self, new_embeddings: torch.Tensor):539 self.word_embeddings = new_embeddings540 541 def forward(542 self,543 input_ids: Optional[torch.LongTensor] = None,544 past_key_values: Optional[Tuple[Tuple[torch.Tensor, torch.Tensor], ...]] = None,545 attention_mask: Optional[torch.Tensor] = None,546 head_mask: Optional[torch.LongTensor] = None,547 inputs_embeds: Optional[torch.LongTensor] = None,548 use_cache: Optional[bool] = None,549 output_attentions: Optional[bool] = None,550 output_hidden_states: Optional[bool] = None,551 return_dict: Optional[bool] = None,552 **deprecated_arguments,553 ) -> Union[Tuple[torch.Tensor, ...], BaseModelOutputWithPastAndCrossAttentions]:554 if deprecated_arguments.pop("position_ids", False) is not False:555 # `position_ids` could have been `torch.Tensor` or `None` so defaulting pop to `False` allows to detect if users were passing explicitly `None`556 warnings.warn(557 "`position_ids` have no functionality in BLOOM and will be removed in v5.0.0. You can safely ignore"558 " passing `position_ids`.",559 FutureWarning,560 )561 if len(deprecated_arguments) > 0:562 raise ValueError(f"Got unexpected arguments: {deprecated_arguments}")563 564 output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions565 output_hidden_states = (566 output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states567 )568 use_cache = use_cache if use_cache is not None else self.config.use_cache569 return_dict = return_dict if return_dict is not None else self.config.use_return_dict570 571 if input_ids is not None and inputs_embeds is not None:572 raise ValueError("You cannot specify both input_ids and inputs_embeds at the same time")573 elif input_ids is not None:574 batch_size, seq_length = input_ids.shape575 elif inputs_embeds is not None:576 batch_size, seq_length, _ = inputs_embeds.shape577 else:578 raise ValueError("You have to specify either input_ids or inputs_embeds")579 580 if past_key_values is None:581 past_key_values = tuple([None] * len(self.h))582 583 # Prepare head mask if needed584 # 1.0 in head_mask indicate we keep the head585 # attention_probs has shape batch_size x num_heads x N x N586 # head_mask has shape n_layer x batch x num_heads x N x N587 head_mask = self.get_head_mask(head_mask, self.config.n_layer)588 589 if inputs_embeds is None:590 inputs_embeds = self.word_embeddings(input_ids)591 592 hidden_states = inputs_embeds593 594 presents = () if use_cache else None595 all_self_attentions = () if output_attentions else None596 all_hidden_states = () if output_hidden_states else None597 598 # Compute alibi tensor: check build_alibi_tensor documentation599 seq_length_with_past = seq_length600 past_key_values_length = 0601 if past_key_values[0] is not None:602 past_key_values_length = past_key_values[0][0].shape[2]603 seq_length_with_past = seq_length_with_past + past_key_values_length604 if attention_mask is None:605 attention_mask = torch.ones((batch_size, seq_length_with_past), device=hidden_states.device)606 else:607 attention_mask = attention_mask.to(hidden_states.device)608 609 if self.alibi:610 alibi = build_alibi_tensor(attention_mask, self.num_heads, dtype=hidden_states.dtype)611 else:612 alibi = None613 614 causal_mask = self._prepare_attn_mask(615 attention_mask,616 input_shape=(batch_size, seq_length),617 past_key_values_length=past_key_values_length,618 )619 620 for i, (block, layer_past) in enumerate(zip(self.h, past_key_values)):621 622 if output_hidden_states:623 all_hidden_states = all_hidden_states + (hidden_states,)624 625 if self.gradient_checkpointing and self.training:626 627 if use_cache:628 logger.warning(629 "`use_cache=True` is incompatible with gradient checkpointing. Setting `use_cache=False`..."630 )631 use_cache = False632 633 def create_custom_forward(module):634 def custom_forward(*inputs):635 # None for past_key_value636 return module(*inputs, use_cache=use_cache, output_attentions=output_attentions)637 638 return custom_forward639 640 outputs = torch.utils.checkpoint.checkpoint(641 create_custom_forward(block),642 hidden_states,643 alibi,644 causal_mask,645 head_mask[i],646 )647 else:648 outputs = block(649 hidden_states,650 layer_past=layer_past,651 attention_mask=causal_mask,652 head_mask=head_mask[i],653 use_cache=use_cache,654 output_attentions=output_attentions,655 alibi=alibi,656 )657 658 hidden_states = outputs[0]659 if use_cache is True:660 presents = presents + (outputs[1],)661 662 if output_attentions:663 all_self_attentions = all_self_attentions + (outputs[2 if use_cache else 1],)664 665 # Add last hidden state666 hidden_states = self.ln_f(hidden_states)667 668 if output_hidden_states:669 all_hidden_states = all_hidden_states + (hidden_states,)670 671 if not return_dict:672 return tuple(v for v in [hidden_states, presents, all_hidden_states, all_self_attentions] if v is not None)673 674 return BaseModelOutputWithPastAndCrossAttentions(675 last_hidden_state=hidden_states,676 past_key_values=presents,677 hidden_states=all_hidden_states,678 attentions=all_self_attentions,679 )680 681 682class RWForCausalLM(RWPreTrainedModel):683 _keys_to_ignore_on_load_missing = [r"h.*.self_attention.scale_mask_softmax.causal_mask", r"lm_head.weight"]684 685 def __init__(self, config: RWConfig):686 super().__init__(config)687 self.transformer = RWModel(config)688 self.lm_head = nn.Linear(config.hidden_size, config.vocab_size, bias=False)689 690 # Initialize weights and apply final processing691 self.post_init()692 693 def get_output_embeddings(self):694 return self.lm_head695 696 def set_output_embeddings(self, new_embeddings: torch.Tensor):697 self.lm_head = new_embeddings698 699 def prepare_inputs_for_generation(700 self,701 input_ids: torch.LongTensor,702 past: Optional[torch.Tensor] = None,703 attention_mask: Optional[torch.Tensor] = None,704 **kwargs,705 ) -> dict:706 # only last token for input_ids if past is not None707 if past:708 input_ids = input_ids[:, -1].unsqueeze(-1)709 710 # the cache may be in the stardard format (e.g. in contrastive search), convert to our's format if needed711 if past[0][0].shape[0] == input_ids.shape[0]:712 past = self._convert_to_rw_cache(past)713 714 return {715 "input_ids": input_ids,716 "past_key_values": past,717 "use_cache": kwargs.get("use_cache"),718 "attention_mask": attention_mask,719 }720 721 def forward(722 self,723 input_ids: Optional[torch.LongTensor] = None,724 past_key_values: Optional[Tuple[Tuple[torch.Tensor, torch.Tensor], ...]] = None,725 attention_mask: Optional[torch.Tensor] = None,726 head_mask: Optional[torch.Tensor] = None,727 inputs_embeds: Optional[torch.Tensor] = None,728 labels: Optional[torch.Tensor] = None,729 use_cache: Optional[bool] = None,730 output_attentions: Optional[bool] = None,731 output_hidden_states: Optional[bool] = None,732 return_dict: Optional[bool] = None,733 **deprecated_arguments,734 ) -> Union[Tuple[torch.Tensor], CausalLMOutputWithCrossAttentions]:735 r"""736 labels (`torch.LongTensor` of shape `(batch_size, sequence_length)`, *optional*):737 Labels for language modeling. Note that the labels **are shifted** inside the model, i.e. you can set738 `labels = input_ids` Indices are selected in `[-100, 0, ..., config.vocab_size]` All labels set to `-100`739 are ignored (masked), the loss is only computed for labels in `[0, ..., config.vocab_size]`740 """741 if deprecated_arguments.pop("position_ids", False) is not False:742 # `position_ids` could have been `torch.Tensor` or `None` so defaulting pop to `False` allows to detect if users were passing explicitly `None`743 warnings.warn(744 "`position_ids` have no functionality in BLOOM and will be removed in v5.0.0. You can safely ignore"745 " passing `position_ids`.",746 FutureWarning,747 )748 if len(deprecated_arguments) > 0:749 raise ValueError(f"Got unexpected arguments: {deprecated_arguments}")750 751 return_dict = return_dict if return_dict is not None else self.config.use_return_dict752 753 transformer_outputs = self.transformer(754 input_ids,755 past_key_values=past_key_values,756 attention_mask=attention_mask,757 head_mask=head_mask,758 inputs_embeds=inputs_embeds,759 use_cache=use_cache,760 output_attentions=output_attentions,761 output_hidden_states=output_hidden_states,762 return_dict=return_dict,763 )764 hidden_states = transformer_outputs[0]765 766 lm_logits = self.lm_head(hidden_states)767 768 loss = None769 if labels is not None:770 # Shift so that tokens < n predict n771 shift_logits = lm_logits[..., :-1, :].contiguous()772 shift_labels = labels[..., 1:].contiguous()773 batch_size, seq_length, vocab_size = shift_logits.shape774 # Flatten the tokens775 loss_fct = CrossEntropyLoss()776 loss = loss_fct(777 shift_logits.view(batch_size * seq_length, vocab_size), shift_labels.view(batch_size * seq_length)778 )779 780 if not return_dict:781 output = (lm_logits,) + transformer_outputs[1:]782 return ((loss,) + output) if loss is not None else output783 784 return CausalLMOutputWithCrossAttentions(785 loss=loss,786 logits=lm_logits,787 past_key_values=transformer_outputs.past_key_values,788 hidden_states=transformer_outputs.hidden_states,789 attentions=transformer_outputs.attentions,790 )791 792 def _reorder_cache(793 self, past: Tuple[Tuple[torch.Tensor, torch.Tensor], ...], beam_idx: torch.LongTensor794 ) -> Tuple[Tuple[torch.Tensor, torch.Tensor], ...]:795 """796 This function is used to re-order the `past_key_values` cache if [`~PreTrainedModel.beam_search`] or797 [`~PreTrainedModel.beam_sample`] is called. This is required to match `past_key_values` with the correct798 beam_idx at every generation step.799 800 Output shares the same memory storage as `past`.801 """802 standardized_past = self._convert_to_standard_cache(past, batch_size=len(beam_idx))803 804 # Get a copy of `beam_idx` on all the devices where we need those indices.805 device_to_beam_idx = {806 past_state.device: beam_idx.to(past_state.device) for layer_past in past for past_state in layer_past807 }808 reordered_past = tuple(809 (810 layer_past[0].index_select(0, device_to_beam_idx[layer_past[0].device]),811 layer_past[1].index_select(0, device_to_beam_idx[layer_past[0].device]),812 )813 for layer_past in standardized_past814 )815 return self._convert_to_rw_cache(reordered_past)816 817 818class RWForSequenceClassification(RWPreTrainedModel):819 _keys_to_ignore_on_load_missing = [r"h.*.self_attention.scale_mask_softmax.causal_mask", r"lm_head.weight"]820 821 def __init__(self, config: RWConfig):822 super().__init__(config)823 self.num_labels = config.num_labels824 self.transformer = RWModel(config)825 self.score = nn.Linear(config.hidden_size, config.num_labels, bias=False)826 827 # Initialize weights and apply final processing828 self.post_init()829 830 def forward(831 self,832 input_ids: Optional[torch.LongTensor] = None,833 past_key_values: Optional[Tuple[Tuple[torch.Tensor, torch.Tensor], ...]] = None,834 attention_mask: Optional[torch.Tensor] = None,835 head_mask: Optional[torch.Tensor] = None,836 inputs_embeds: Optional[torch.Tensor] = None,837 labels: Optional[torch.Tensor] = None,838 use_cache: Optional[bool] = None,839 output_attentions: Optional[bool] = None,840 output_hidden_states: Optional[bool] = None,841 return_dict: Optional[bool] = None,842 **deprecated_arguments,843 ) -> Union[Tuple[torch.Tensor], SequenceClassifierOutputWithPast]:844 r"""845 labels (`torch.LongTensor` of shape `(batch_size,)`, *optional*):846 Labels for computing the sequence classification/regression loss. Indices should be in `[0, ...,847 config.num_labels - 1]`. If `config.num_labels == 1` a regression loss is computed (Mean-Square loss), If848 `config.num_labels > 1` a classification loss is computed (Cross-Entropy).849 """850 if deprecated_arguments.pop("position_ids", False) is not False:851 # `position_ids` could have been `torch.Tensor` or `None` so defaulting pop to `False` allows to detect if users were passing explicitly `None`852 warnings.warn(853 "`position_ids` have no functionality in BLOOM and will be removed in v5.0.0. You can safely ignore"854 " passing `position_ids`.",855 FutureWarning,856 )857 if len(deprecated_arguments) > 0:858 raise ValueError(f"Got unexpected arguments: {deprecated_arguments}")859 860 return_dict = return_dict if return_dict is not None else self.config.use_return_dict861 862 transformer_outputs = self.transformer(863 input_ids,864 past_key_values=past_key_values,865 attention_mask=attention_mask,866 head_mask=head_mask,867 inputs_embeds=inputs_embeds,868 use_cache=use_cache,869 output_attentions=output_attentions,870 output_hidden_states=output_hidden_states,871 return_dict=return_dict,872 )873 874 hidden_states = transformer_outputs[0]875 logits = self.score(hidden_states)876 877 if input_ids is not None:878 batch_size = input_ids.shape[0]879 else:880 batch_size = inputs_embeds.shape[0]881 882 if self.config.pad_token_id is None and batch_size != 1:883 raise ValueError("Cannot handle batch sizes > 1 if no padding token is defined.")884 if self.config.pad_token_id is None:885 sequence_lengths = -1886 else:887 if input_ids is not None:888 sequence_lengths = torch.ne(input_ids, self.config.pad_token_id).sum(dim=-1) - 1889 else:890 sequence_lengths = -1891 logger.warning(892 f"{self.__class__.__name__} will not detect padding tokens in `inputs_embeds`. Results may be "893 "unexpected if using padding tokens in conjunction with `inputs_embeds.`"894 )895 896 pooled_logits = logits[torch.arange(batch_size, device=logits.device), sequence_lengths]897 898 loss = None899 if labels is not None:900 if self.config.problem_type is None:901 if self.num_labels == 1:902 self.config.problem_type = "regression"903 elif self.num_labels > 1 and (labels.dtype == torch.long or labels.dtype == torch.int):904 self.config.problem_type = "single_label_classification"905 else:906 self.config.problem_type = "multi_label_classification"907 908 if self.config.problem_type == "regression":909 loss_fct = MSELoss()910 if self.num_labels == 1:911 loss = loss_fct(pooled_logits.squeeze(), labels.squeeze())912 else:913 loss = loss_fct(pooled_logits, labels)914 elif self.config.problem_type == "single_label_classification":915 loss_fct = CrossEntropyLoss()916 loss = loss_fct(pooled_logits, labels)917 elif self.config.problem_type == "multi_label_classification":918 loss_fct = BCEWithLogitsLoss()919 loss = loss_fct(pooled_logits, labels)920 if not return_dict:921 output = (pooled_logits,) + transformer_outputs[1:]922 return ((loss,) + output) if loss is not None else output923 924 return SequenceClassifierOutputWithPast(925 loss=loss,926 logits=pooled_logits,927 past_key_values=transformer_outputs.past_key_values,928 hidden_states=transformer_outputs.hidden_states,929 attentions=transformer_outputs.attentions,930 )931 932 933class RWForTokenClassification(RWPreTrainedModel):934 _keys_to_ignore_on_load_missing = [r"h.*.self_attention.scale_mask_softmax.causal_mask", r"lm_head.weight"]935 936 def __init__(self, config: RWConfig):937 super().__init__(config)938 self.num_labels = config.num_labels939 940 self.transformer = RWModel(config)941 if hasattr(config, "classifier_dropout") and config.classifier_dropout is not None:942 classifier_dropout = config.classifier_dropout943 elif hasattr(config, "hidden_dropout") and config.hidden_dropout is not None:944 classifier_dropout = config.hidden_dropout945 else:946 classifier_dropout = 0.1947 self.dropout = nn.Dropout(classifier_dropout)948 self.classifier = nn.Linear(config.hidden_size, config.num_labels)949 950 # Initialize weights and apply final processing951 self.post_init()952 953 def forward(954 self,955 input_ids: Optional[torch.LongTensor] = None,956 past_key_values: Optional[Tuple[Tuple[torch.Tensor, torch.Tensor], ...]] = None,957 attention_mask: Optional[torch.Tensor] = None,958 head_mask: Optional[torch.Tensor] = None,959 inputs_embeds: Optional[torch.Tensor] = None,960 labels: Optional[torch.Tensor] = None,961 use_cache: Optional[bool] = None,962 output_attentions: Optional[bool] = None,963 output_hidden_states: Optional[bool] = None,964 return_dict: Optional[bool] = None,965 **deprecated_arguments,966 ) -> Union[Tuple[torch.Tensor], TokenClassifierOutput]:967 r"""968 labels (`torch.LongTensor` of shape `(batch_size,)`, *optional*):969 Labels for computing the sequence classification/regression loss. Indices should be in `[0, ...,970 config.num_labels - 1]`. If `config.num_labels == 1` a regression loss is computed (Mean-Square loss), If971 `config.num_labels > 1` a classification loss is computed (Cross-Entropy).972 """973 if deprecated_arguments.pop("position_ids", False) is not False:974 # `position_ids` could have been `torch.Tensor` or `None` so defaulting pop to `False` allows to detect if users were passing explicitly `None`975 warnings.warn(976 "`position_ids` have no functionality in BLOOM and will be removed in v5.0.0. You can safely ignore"977 " passing `position_ids`.",978 FutureWarning,979 )980 if len(deprecated_arguments) > 0:981 raise ValueError(f"Got unexpected arguments: {deprecated_arguments}")982 983 return_dict = return_dict if return_dict is not None else self.config.use_return_dict984 985 transformer_outputs = self.transformer(986 input_ids,987 past_key_values=past_key_values,988 attention_mask=attention_mask,989 head_mask=head_mask,990 inputs_embeds=inputs_embeds,991 use_cache=use_cache,992 output_attentions=output_attentions,993 output_hidden_states=output_hidden_states,994 return_dict=return_dict,995 )996 997 hidden_states = transformer_outputs[0]998 hidden_states = self.dropout(hidden_states)999 logits = self.classifier(hidden_states)1000 1001 loss = None1002 if labels is not None:1003 batch_size, seq_length = labels.shape1004 loss_fct = CrossEntropyLoss()1005 loss = loss_fct(logits.view(batch_size * seq_length, self.num_labels), labels.view(batch_size * seq_length))1006 1007 if not return_dict:1008 output = (logits,) + transformer_outputs[2:]1009 return ((loss,) + output) if loss is not None else output1010 1011 return TokenClassifierOutput(1012 loss=loss,1013 logits=logits,1014 hidden_states=transformer_outputs.hidden_states,1015 attentions=transformer_outputs.attentions,1016 )1017 1018 1019class RWForQuestionAnswering(RWPreTrainedModel):1020 _keys_to_ignore_on_load_missing = [r"h.*.self_attention.scale_mask_softmax.causal_mask", r"lm_head.weight"]1021 1022 def __init__(self, config):1023 super().__init__(config)1024 self.transformer = RWModel(config)1025 self.qa_outputs = nn.Linear(config.hidden_size, 2)1026 1027 # Initialize weights and apply final processing1028 self.post_init()1029 1030 def forward(1031 self,1032 input_ids: Optional[torch.LongTensor] = None,1033 attention_mask: Optional[torch.FloatTensor] = None,1034 position_ids: Optional[torch.LongTensor] = None,1035 head_mask: Optional[torch.FloatTensor] = None,1036 inputs_embeds: Optional[torch.FloatTensor] = None,1037 start_positions: Optional[torch.LongTensor] = None,1038 end_positions: Optional[torch.LongTensor] = None,1039 output_attentions: Optional[bool] = None,1040 output_hidden_states: Optional[bool] = None,1041 return_dict: Optional[bool] = None,1042 ) -> Union[Tuple, QuestionAnsweringModelOutput]:1043 r"""1044 start_positions (`torch.LongTensor` of shape `(batch_size,)`, *optional*):1045 Labels for position (index) of the start of the labelled span for computing the token classification loss.1046 Positions are clamped to the length of the sequence (`sequence_length`). Position outside of the sequence1047 are not taken into account for computing the loss.1048 end_positions (`torch.LongTensor` of shape `(batch_size,)`, *optional*):1049 Labels for position (index) of the end of the labelled span for computing the token classification loss.1050 Positions are clamped to the length of the sequence (`sequence_length`). Position outside of the sequence1051 are not taken into account for computing the loss.1052 """1053 return_dict = return_dict if return_dict is not None else self.config.use_return_dict1054 1055 outputs = self.transformer(1056 input_ids,1057 attention_mask=attention_mask,1058 position_ids=position_ids,1059 head_mask=head_mask,1060 inputs_embeds=inputs_embeds,1061 output_attentions=output_attentions,1062 output_hidden_states=output_hidden_states,1063 return_dict=return_dict,1064 )1065 1066 sequence_output = outputs[0]1067 1068 logits = self.qa_outputs(sequence_output)1069 start_logits, end_logits = logits.split(1, dim=-1)1070 start_logits = start_logits.squeeze(-1).contiguous()1071 end_logits = end_logits.squeeze(-1).contiguous()1072 1073 total_loss = None1074 if start_positions is not None and end_positions is not None:1075 # If we are on multi-GPU, split add a dimension1076 if len(start_positions.size()) > 1:1077 start_positions = start_positions.squeeze(-1)1078 if len(end_positions.size()) > 1:1079 end_positions = end_positions.squeeze(-1)1080 # sometimes the start/end positions are outside our model inputs, we ignore these terms1081 ignored_index = start_logits.size(1)1082 start_positions = start_positions.clamp(0, ignored_index)1083 end_positions = end_positions.clamp(0, ignored_index)1084 1085 loss_fct = CrossEntropyLoss(ignore_index=ignored_index)1086 start_loss = loss_fct(start_logits, start_positions)1087 end_loss = loss_fct(end_logits, end_positions)1088 total_loss = (start_loss + end_loss) / 21089 1090 if not return_dict:1091 output = (start_logits, end_logits) + outputs[2:]1092 return ((total_loss,) + output) if total_loss is not None else output1093 1094 return QuestionAnsweringModelOutput(1095 loss=total_loss,1096 start_logits=start_logits,1097 end_logits=end_logits,1098 hidden_states=outputs.hidden_states,1099 attentions=outputs.attentions,1100 )