papahawk/falcon-40b
116
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 (config.n_head_kv * 2 + config.n_head) * self.head_dim,179 bias=config.bias,180 )181 self.dense = Linear(self.hidden_size, self.hidden_size, bias=config.bias)182 self.attention_dropout = nn.Dropout(config.attention_dropout)183 self.num_kv = config.n_head_kv184 185 def _split_heads(self, fused_qkv: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]:186 """187 Split the last dimension into (num_heads, head_dim), results share same memory188 storage as `fused_qkv`189 190 Args:191 fused_qkv (`torch.tensor`, *required*): [batch_size, seq_length, num_heads * 3 * head_dim]192 193 Returns:194 query: [batch_size, seq_length, num_heads, head_dim]195 key: [batch_size, seq_length, num_heads, head_dim]196 value: [batch_size, seq_length, num_heads, head_dim]197 """198 batch, seq_len, _ = fused_qkv.shape199 qkv = fused_qkv.view(batch, seq_len, -1, self.num_heads // self.num_kv + 2, 64)200 q = qkv[:, :, :, :-2]201 k = qkv[:, :, :, [-2]]202 v = qkv[:, :, :, [-1]]203 k = torch.broadcast_to(k, q.shape)204 v = torch.broadcast_to(v, q.shape)205 206 q, k, v = [207 rearrange(208 x,209 "batch seq_len group num_heads head_dim ->\210 batch seq_len (group num_heads) head_dim",211 head_dim=self.head_dim,212 )213 for x in [q, k, v]214 ]215 return q, k, v216 217 def _merge_heads(self, x: torch.Tensor) -> torch.Tensor:218 """219 Merge heads together over the last dimenstion220 221 Args:222 x: (`torch.tensor`, *required*): [batch_size * num_heads, seq_length, head_dim]223 224 Returns:225 torch.tensor: [batch_size, seq_length, num_heads * head_dim]226 """227 # What we want to achieve is:228 # batch_size * num_heads, seq_length, head_dim -> batch_size, seq_length, num_heads * head_dim229 batch_size_and_num_heads, seq_length, _ = x.shape230 batch_size = batch_size_and_num_heads // self.num_heads231 232 # First view to decompose the batch size233 # batch_size * num_heads, seq_length, head_dim -> batch_size, num_heads, seq_length, head_dim234 x = x.view(batch_size, self.num_heads, seq_length, self.head_dim)235 236 # batch_size, num_heads, seq_length, head_dim -> batch_size, seq_length, num_heads, head_dim237 x = x.permute(0, 2, 1, 3)238 239 # batch_size, seq_length, num_heads, head_dim -> batch_size, seq_length, num_heads * head_dim240 return x.reshape(batch_size, seq_length, self.num_heads * self.head_dim)241 242 def forward(243 self,244 hidden_states: torch.Tensor,245 alibi: torch.Tensor,246 attention_mask: torch.Tensor,247 layer_past: Optional[Tuple[torch.Tensor, torch.Tensor]] = None,248 head_mask: Optional[torch.Tensor] = None,249 use_cache: bool = False,250 output_attentions: bool = False,251 ):252 fused_qkv = self.query_key_value(hidden_states) # [batch_size, seq_length, 3 x hidden_size]253 254 # 3 x [batch_size, seq_length, num_heads, head_dim]255 (query_layer, key_layer, value_layer) = self._split_heads(fused_qkv)256 257 batch_size, q_length, _, _ = query_layer.shape258 259 query_layer = query_layer.transpose(1, 2).reshape(batch_size * self.num_heads, q_length, self.head_dim)260 key_layer = key_layer.transpose(1, 2).reshape(261 batch_size * self.num_heads,262 q_length,263 self.head_dim,264 )265 value_layer = value_layer.transpose(1, 2).reshape(batch_size * self.num_heads, q_length, self.head_dim)266 267 query_layer, key_layer = self.maybe_rotary(query_layer, key_layer)268 269 if layer_past is not None:270 past_key, past_value = layer_past271 # concatenate along seq_length dimension:272 # - key: [batch_size * self.num_heads, head_dim, kv_length]273 # - value: [batch_size * self.num_heads, kv_length, head_dim]274 key_layer = torch.cat((past_key, key_layer), dim=1)275 value_layer = torch.cat((past_value, value_layer), dim=1)276 277 _, kv_length, _ = key_layer.shape278 279 if use_cache is True:280 present = (key_layer, value_layer)281 else:282 present = None283 284 if alibi is None:285 query_layer_ = query_layer.reshape(batch_size, self.num_heads, -1, self.head_dim)286 key_layer_ = key_layer.reshape(batch_size, self.num_heads, -1, self.head_dim)287 value_layer_ = value_layer.reshape(batch_size, self.num_heads, -1, self.head_dim)288 289 attn_output = F.scaled_dot_product_attention(290 query_layer_, key_layer_, value_layer_, None, 0.0, is_causal=True291 )292 293 x = attn_output.view(batch_size, self.num_heads, q_length, self.head_dim)294 x = x.permute(0, 2, 1, 3)295 attn_output = x.reshape(batch_size, q_length, self.num_heads * self.head_dim)296 297 output_tensor = self.dense(attn_output)298 299 outputs = (output_tensor, present)300 assert not output_attentions # not supported.301 return outputs302 else:303 attention_mask_float = (attention_mask * 1.0).masked_fill(attention_mask, -1e9).to(torch.bfloat16)304 matmul_result = query_layer @ key_layer.transpose(-1, -2)305 306 # change view to [batch_size, num_heads, q_length, kv_length]307 attention_scores = matmul_result.view(batch_size, self.num_heads, q_length, kv_length)308 309 # cast attention scores to fp32, compute scaled softmax and cast back to initial dtype - [batch_size, num_heads, q_length, kv_length]310 input_dtype = attention_scores.dtype311 # `float16` has a minimum value of -65504.0, whereas `bfloat16` and `float32` have a minimum value of `-3.4e+38`312 if input_dtype == torch.float16 or input_dtype == torch.bfloat16:313 attention_scores = attention_scores.to(torch.float32)314 # attn_weights = torch.masked_fill(attention_scores, attention_mask, torch.finfo(attention_scores.dtype).min)315 attention_probs = F.softmax(316 (attention_scores + alibi.view(batch_size, self.num_heads, 1, -1)) * self.inv_norm_factor317 + attention_mask_float,318 dim=-1,319 dtype=hidden_states.dtype,320 )321 # [batch_size, num_heads, q_length, kv_length]322 attention_probs = self.attention_dropout(attention_probs)323 324 if head_mask is not None:325 attention_probs = attention_probs * head_mask326 327 # change view [batch_size x num_heads, q_length, kv_length]328 attention_probs_reshaped = attention_probs.view(batch_size * self.num_heads, q_length, kv_length)329 330 # matmul: [batch_size * num_heads, q_length, head_dim]331 context_layer = attention_probs_reshaped @ value_layer332 333 # change view [batch_size, num_heads, q_length, head_dim]334 context_layer = self._merge_heads(context_layer)335 336 output_tensor = self.dense(context_layer)337 338 outputs = (output_tensor, present)339 if output_attentions:340 outputs += (attention_probs,)341 342 return outputs343 344 345class MLP(nn.Module):346 def __init__(self, config: RWConfig):347 super().__init__()348 hidden_size = config.hidden_size349 350 self.dense_h_to_4h = Linear(hidden_size, 4 * hidden_size, bias=config.bias)351 self.act = nn.GELU()352 self.dense_4h_to_h = Linear(4 * hidden_size, hidden_size, bias=config.bias)353 self.hidden_dropout = config.hidden_dropout354 355 def forward(self, x: torch.Tensor) -> torch.Tensor:356 x = self.act(self.dense_h_to_4h(x))357 x = self.dense_4h_to_h(x)358 return x359 360 361class DecoderLayer(nn.Module):362 def __init__(self, config: RWConfig):363 super().__init__()364 hidden_size = config.hidden_size365 366 self.ln_attn = LayerNorm(hidden_size, eps=config.layer_norm_epsilon)367 self.ln_mlp = LayerNorm(hidden_size, eps=config.layer_norm_epsilon)368 369 self.num_heads = config.n_head370 self.self_attention = Attention(config)371 372 self.mlp = MLP(config)373 374 self.apply_residual_connection_post_layernorm = config.apply_residual_connection_post_layernorm375 self.hidden_dropout = config.hidden_dropout376 377 self.config = config378 379 def forward(380 self,381 hidden_states: torch.Tensor,382 alibi: torch.Tensor,383 attention_mask: torch.Tensor,384 layer_past: Optional[Tuple[torch.Tensor, torch.Tensor]] = None,385 head_mask: Optional[torch.Tensor] = None,386 use_cache: bool = False,387 output_attentions: bool = False,388 ):389 390 ln_attn = self.ln_attn(hidden_states)391 ln_mlp = self.ln_mlp(hidden_states)392 393 residual = hidden_states394 395 # Self attention.396 attn_outputs = self.self_attention(397 ln_attn,398 layer_past=layer_past,399 attention_mask=attention_mask,400 alibi=alibi,401 head_mask=head_mask,402 use_cache=use_cache,403 output_attentions=output_attentions,404 )405 406 attention_output = attn_outputs[0]407 408 outputs = attn_outputs[1:]409 410 # MLP.411 mlp_output = self.mlp(ln_mlp)412 413 output = dropout_add(414 mlp_output + attention_output, residual, self.config.hidden_dropout, training=self.training415 )416 417 if use_cache:418 outputs = (output,) + outputs419 else:420 outputs = (output,) + outputs[1:]421 422 return outputs # hidden_states, present, attentions423 424 425class RWPreTrainedModel(PreTrainedModel):426 _keys_to_ignore_on_load_missing = [r"h.*.self_attention.scale_mask_softmax.causal_mask", r"lm_head.weight"]427 """428 An abstract class to handle weights initialization and a simple interface for downloading and loading pretrained429 models.430 """431 432 config_class = RWConfig433 base_model_prefix = "transformer"434 supports_gradient_checkpointing = True435 _no_split_modules = ["DecoderLayer"]436 437 def __init__(self, *inputs, **kwargs):438 super().__init__(*inputs, **kwargs)439 440 def _init_weights(self, module: nn.Module):441 """Initialize the weights."""442 if isinstance(module, nn.Linear) or isinstance(module, Linear):443 # Slightly different from the TF version which uses truncated_normal for initialization444 # cf https://github.com/pytorch/pytorch/pull/5617445 module.weight.data.normal_(mean=0.0, std=self.config.initializer_range)446 if module.bias is not None:447 module.bias.data.zero_()448 elif isinstance(module, nn.Embedding):449 module.weight.data.normal_(mean=0.0, std=self.config.initializer_range)450 if module.padding_idx is not None:451 module.weight.data[module.padding_idx].zero_()452 elif isinstance(module, LayerNorm):453 module.bias.data.zero_()454 module.weight.data.fill_(1.0)455 456 def _set_gradient_checkpointing(self, module: nn.Module, value: bool = False):457 if isinstance(module, RWModel):458 module.gradient_checkpointing = value459 460 @staticmethod461 def _convert_to_standard_cache(462 past_key_value: Tuple[Tuple[torch.Tensor, torch.Tensor]], batch_size: int463 ) -> Tuple[Tuple[torch.Tensor, torch.Tensor]]:464 """465 Standardizes the format of the cache so as to match most implementations, i.e. to tuple(tuple([batch_size,466 num_heads, ...]))467 """468 batch_size_times_num_heads, head_dim, seq_length = past_key_value[0][0].shape469 num_heads = batch_size_times_num_heads // batch_size470 # key: [batch_size * num_heads, head_dim, seq_length] -> [batch_size, num_heads, head_dim, seq_length]471 # value: [batch_size * num_heads, seq_length, head_dim] -> [batch_size, num_heads, seq_length, head_dim]472 return tuple(473 (474 layer_past[0].view(batch_size, num_heads, head_dim, seq_length),475 layer_past[1].view(batch_size, num_heads, seq_length, head_dim),476 )477 for layer_past in past_key_value478 )479 480 @staticmethod481 def _convert_to_rw_cache(482 past_key_value: Tuple[Tuple[torch.Tensor, torch.Tensor]]483 ) -> Tuple[Tuple[torch.Tensor, torch.Tensor]]:484 batch_size, num_heads, head_dim, seq_length = past_key_value[0][0].shape485 batch_size_times_num_heads = batch_size * num_heads486 # key: [batch_size, num_heads, head_dim, seq_length] -> [batch_size * num_heads, head_dim, seq_length]487 # value: [batch_size, num_heads, seq_length, head_dim] -> [batch_size * num_heads, seq_length, head_dim]488 return tuple(489 (490 layer_past[0].view(batch_size_times_num_heads, head_dim, seq_length),491 layer_past[1].view(batch_size_times_num_heads, seq_length, head_dim),492 )493 for layer_past in past_key_value494 )495 496 497class RWModel(RWPreTrainedModel):498 def __init__(self, config: RWConfig):499 super().__init__(config)500 501 self.embed_dim = config.hidden_size502 self.num_heads = config.n_head503 self.alibi = config.alibi504 505 # Embedding + LN Embedding506 self.word_embeddings = nn.Embedding(config.vocab_size, self.embed_dim)507 508 # Transformer blocks509 self.h = nn.ModuleList([DecoderLayer(config) for _ in range(config.num_hidden_layers)])510 511 # Final Layer Norm512 self.ln_f = LayerNorm(self.embed_dim, eps=config.layer_norm_epsilon)513 514 self.gradient_checkpointing = False515 516 # Initialize weights and apply final processing517 self.post_init()518 519 def get_input_embeddings(self):520 return self.word_embeddings521 522 def _prepare_attn_mask(523 self, attention_mask: torch.Tensor, input_shape: Tuple[int, int], past_key_values_length: int524 ) -> torch.BoolTensor:525 # create causal mask526 # [batch_size, seq_length] -> [batch_size, 1, tgt_length, src_length]527 combined_attention_mask = None528 device = attention_mask.device529 _, src_length = input_shape530 531 if src_length > 1:532 combined_attention_mask = _make_causal_mask(533 input_shape, device=device, past_key_values_length=past_key_values_length534 )535 536 # [batch_size, seq_length] -> [batch_size, 1, tgt_length, src_length]537 expanded_attn_mask = _expand_mask(attention_mask, tgt_length=src_length)538 combined_attention_mask = (539 expanded_attn_mask if combined_attention_mask is None else expanded_attn_mask | combined_attention_mask540 )541 542 return combined_attention_mask543 544 def set_input_embeddings(self, new_embeddings: torch.Tensor):545 self.word_embeddings = new_embeddings546 547 def forward(548 self,549 input_ids: Optional[torch.LongTensor] = None,550 past_key_values: Optional[Tuple[Tuple[torch.Tensor, torch.Tensor], ...]] = None,551 attention_mask: Optional[torch.Tensor] = None,552 head_mask: Optional[torch.LongTensor] = None,553 inputs_embeds: Optional[torch.LongTensor] = None,554 use_cache: Optional[bool] = None,555 output_attentions: Optional[bool] = None,556 output_hidden_states: Optional[bool] = None,557 return_dict: Optional[bool] = None,558 **deprecated_arguments,559 ) -> Union[Tuple[torch.Tensor, ...], BaseModelOutputWithPastAndCrossAttentions]:560 if deprecated_arguments.pop("position_ids", False) is not False:561 # `position_ids` could have been `torch.Tensor` or `None` so defaulting pop to `False` allows to detect if users were passing explicitly `None`562 warnings.warn(563 "`position_ids` have no functionality in BLOOM and will be removed in v5.0.0. You can safely ignore"564 " passing `position_ids`.",565 FutureWarning,566 )567 if len(deprecated_arguments) > 0:568 raise ValueError(f"Got unexpected arguments: {deprecated_arguments}")569 570 output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions571 output_hidden_states = (572 output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states573 )574 use_cache = use_cache if use_cache is not None else self.config.use_cache575 return_dict = return_dict if return_dict is not None else self.config.use_return_dict576 577 if input_ids is not None and inputs_embeds is not None:578 raise ValueError("You cannot specify both input_ids and inputs_embeds at the same time")579 elif input_ids is not None:580 batch_size, seq_length = input_ids.shape581 elif inputs_embeds is not None:582 batch_size, seq_length, _ = inputs_embeds.shape583 else:584 raise ValueError("You have to specify either input_ids or inputs_embeds")585 586 if past_key_values is None:587 past_key_values = tuple([None] * len(self.h))588 589 # Prepare head mask if needed590 # 1.0 in head_mask indicate we keep the head591 # attention_probs has shape batch_size x num_heads x N x N592 # head_mask has shape n_layer x batch x num_heads x N x N593 head_mask = self.get_head_mask(head_mask, self.config.n_layer)594 595 if inputs_embeds is None:596 inputs_embeds = self.word_embeddings(input_ids)597 598 hidden_states = inputs_embeds599 600 presents = () if use_cache else None601 all_self_attentions = () if output_attentions else None602 all_hidden_states = () if output_hidden_states else None603 604 # Compute alibi tensor: check build_alibi_tensor documentation605 seq_length_with_past = seq_length606 past_key_values_length = 0607 if past_key_values[0] is not None:608 past_key_values_length = past_key_values[0][0].shape[2]609 seq_length_with_past = seq_length_with_past + past_key_values_length610 if attention_mask is None:611 attention_mask = torch.ones((batch_size, seq_length_with_past), device=hidden_states.device)612 else:613 attention_mask = attention_mask.to(hidden_states.device)614 615 if self.alibi:616 alibi = build_alibi_tensor(attention_mask, self.num_heads, dtype=hidden_states.dtype)617 else:618 alibi = None619 620 causal_mask = self._prepare_attn_mask(621 attention_mask,622 input_shape=(batch_size, seq_length),623 past_key_values_length=past_key_values_length,624 )625 626 for i, (block, layer_past) in enumerate(zip(self.h, past_key_values)):627 628 if output_hidden_states:629 all_hidden_states = all_hidden_states + (hidden_states,)630 631 if self.gradient_checkpointing and self.training:632 633 if use_cache:634 logger.warning(635 "`use_cache=True` is incompatible with gradient checkpointing. Setting `use_cache=False`..."636 )637 use_cache = False638 639 def create_custom_forward(module):640 def custom_forward(*inputs):641 # None for past_key_value642 return module(*inputs, use_cache=use_cache, output_attentions=output_attentions)643 644 return custom_forward645 646 outputs = torch.utils.checkpoint.checkpoint(647 create_custom_forward(block),648 hidden_states,649 alibi,650 causal_mask,651 head_mask[i],652 )653 else:654 outputs = block(655 hidden_states,656 layer_past=layer_past,657 attention_mask=causal_mask,658 head_mask=head_mask[i],659 use_cache=use_cache,660 output_attentions=output_attentions,661 alibi=alibi,662 )663 664 hidden_states = outputs[0]665 if use_cache is True:666 presents = presents + (outputs[1],)667 668 if output_attentions:669 all_self_attentions = all_self_attentions + (outputs[2 if use_cache else 1],)670 671 # Add last hidden state672 hidden_states = self.ln_f(hidden_states)673 674 if output_hidden_states:675 all_hidden_states = all_hidden_states + (hidden_states,)676 677 if not return_dict:678 return tuple(v for v in [hidden_states, presents, all_hidden_states, all_self_attentions] if v is not None)679 680 return BaseModelOutputWithPastAndCrossAttentions(681 last_hidden_state=hidden_states,682 past_key_values=presents,683 hidden_states=all_hidden_states,684 attentions=all_self_attentions,685 )686 687 688class RWForCausalLM(RWPreTrainedModel):689 _keys_to_ignore_on_load_missing = [r"h.*.self_attention.scale_mask_softmax.causal_mask", r"lm_head.weight"]690 691 def __init__(self, config: RWConfig):692 super().__init__(config)693 self.transformer = RWModel(config)694 self.lm_head = nn.Linear(config.hidden_size, config.vocab_size, bias=False)695 696 # Initialize weights and apply final processing697 self.post_init()698 699 def get_output_embeddings(self):700 return self.lm_head701 702 def set_output_embeddings(self, new_embeddings: torch.Tensor):703 self.lm_head = new_embeddings704 705 def prepare_inputs_for_generation(706 self,707 input_ids: torch.LongTensor,708 past: Optional[torch.Tensor] = None,709 attention_mask: Optional[torch.Tensor] = None,710 **kwargs,711 ) -> dict:712 # only last token for input_ids if past is not None713 if past:714 input_ids = input_ids[:, -1].unsqueeze(-1)715 716 # the cache may be in the stardard format (e.g. in contrastive search), convert to our's format if needed717 if past[0][0].shape[0] == input_ids.shape[0]:718 past = self._convert_to_rw_cache(past)719 720 return {721 "input_ids": input_ids,722 "past_key_values": past,723 "use_cache": kwargs.get("use_cache"),724 "attention_mask": attention_mask,725 }726 727 def forward(728 self,729 input_ids: Optional[torch.LongTensor] = None,730 past_key_values: Optional[Tuple[Tuple[torch.Tensor, torch.Tensor], ...]] = None,731 attention_mask: Optional[torch.Tensor] = None,732 head_mask: Optional[torch.Tensor] = None,733 inputs_embeds: Optional[torch.Tensor] = None,734 labels: Optional[torch.Tensor] = None,735 use_cache: Optional[bool] = None,736 output_attentions: Optional[bool] = None,737 output_hidden_states: Optional[bool] = None,738 return_dict: Optional[bool] = None,739 **deprecated_arguments,740 ) -> Union[Tuple[torch.Tensor], CausalLMOutputWithCrossAttentions]:741 r"""742 labels (`torch.LongTensor` of shape `(batch_size, sequence_length)`, *optional*):743 Labels for language modeling. Note that the labels **are shifted** inside the model, i.e. you can set744 `labels = input_ids` Indices are selected in `[-100, 0, ..., config.vocab_size]` All labels set to `-100`745 are ignored (masked), the loss is only computed for labels in `[0, ..., config.vocab_size]`746 """747 if deprecated_arguments.pop("position_ids", False) is not False:748 # `position_ids` could have been `torch.Tensor` or `None` so defaulting pop to `False` allows to detect if users were passing explicitly `None`749 warnings.warn(750 "`position_ids` have no functionality in BLOOM and will be removed in v5.0.0. You can safely ignore"751 " passing `position_ids`.",752 FutureWarning,753 )754 if len(deprecated_arguments) > 0:755 raise ValueError(f"Got unexpected arguments: {deprecated_arguments}")756 757 return_dict = return_dict if return_dict is not None else self.config.use_return_dict758 759 transformer_outputs = self.transformer(760 input_ids,761 past_key_values=past_key_values,762 attention_mask=attention_mask,763 head_mask=head_mask,764 inputs_embeds=inputs_embeds,765 use_cache=use_cache,766 output_attentions=output_attentions,767 output_hidden_states=output_hidden_states,768 return_dict=return_dict,769 )770 hidden_states = transformer_outputs[0]771 772 lm_logits = self.lm_head(hidden_states)773 774 loss = None775 if labels is not None:776 # Shift so that tokens < n predict n777 shift_logits = lm_logits[..., :-1, :].contiguous()778 shift_labels = labels[..., 1:].contiguous()779 batch_size, seq_length, vocab_size = shift_logits.shape780 # Flatten the tokens781 loss_fct = CrossEntropyLoss()782 loss = loss_fct(783 shift_logits.view(batch_size * seq_length, vocab_size), shift_labels.view(batch_size * seq_length)784 )785 786 if not return_dict:787 output = (lm_logits,) + transformer_outputs[1:]788 return ((loss,) + output) if loss is not None else output789 790 return CausalLMOutputWithCrossAttentions(791 loss=loss,792 logits=lm_logits,793 past_key_values=transformer_outputs.past_key_values,794 hidden_states=transformer_outputs.hidden_states,795 attentions=transformer_outputs.attentions,796 )797 798 def _reorder_cache(799 self, past: Tuple[Tuple[torch.Tensor, torch.Tensor], ...], beam_idx: torch.LongTensor800 ) -> Tuple[Tuple[torch.Tensor, torch.Tensor], ...]:801 """802 This function is used to re-order the `past_key_values` cache if [`~PreTrainedModel.beam_search`] or803 [`~PreTrainedModel.beam_sample`] is called. This is required to match `past_key_values` with the correct804 beam_idx at every generation step.805 806 Output shares the same memory storage as `past`.807 """808 standardized_past = self._convert_to_standard_cache(past, batch_size=len(beam_idx))809 810 # Get a copy of `beam_idx` on all the devices where we need those indices.811 device_to_beam_idx = {812 past_state.device: beam_idx.to(past_state.device) for layer_past in past for past_state in layer_past813 }814 reordered_past = tuple(815 (816 layer_past[0].index_select(0, device_to_beam_idx[layer_past[0].device]),817 layer_past[1].index_select(0, device_to_beam_idx[layer_past[0].device]),818 )819 for layer_past in standardized_past820 )821 return self._convert_to_rw_cache(reordered_past)822 823 824class RWForSequenceClassification(RWPreTrainedModel):825 _keys_to_ignore_on_load_missing = [r"h.*.self_attention.scale_mask_softmax.causal_mask", r"lm_head.weight"]826 827 def __init__(self, config: RWConfig):828 super().__init__(config)829 self.num_labels = config.num_labels830 self.transformer = RWModel(config)831 self.score = nn.Linear(config.hidden_size, config.num_labels, bias=False)832 833 # Initialize weights and apply final processing834 self.post_init()835 836 def forward(837 self,838 input_ids: Optional[torch.LongTensor] = None,839 past_key_values: Optional[Tuple[Tuple[torch.Tensor, torch.Tensor], ...]] = None,840 attention_mask: Optional[torch.Tensor] = None,841 head_mask: Optional[torch.Tensor] = None,842 inputs_embeds: Optional[torch.Tensor] = None,843 labels: Optional[torch.Tensor] = None,844 use_cache: Optional[bool] = None,845 output_attentions: Optional[bool] = None,846 output_hidden_states: Optional[bool] = None,847 return_dict: Optional[bool] = None,848 **deprecated_arguments,849 ) -> Union[Tuple[torch.Tensor], SequenceClassifierOutputWithPast]:850 r"""851 labels (`torch.LongTensor` of shape `(batch_size,)`, *optional*):852 Labels for computing the sequence classification/regression loss. Indices should be in `[0, ...,853 config.num_labels - 1]`. If `config.num_labels == 1` a regression loss is computed (Mean-Square loss), If854 `config.num_labels > 1` a classification loss is computed (Cross-Entropy).855 """856 if deprecated_arguments.pop("position_ids", False) is not False:857 # `position_ids` could have been `torch.Tensor` or `None` so defaulting pop to `False` allows to detect if users were passing explicitly `None`858 warnings.warn(859 "`position_ids` have no functionality in BLOOM and will be removed in v5.0.0. You can safely ignore"860 " passing `position_ids`.",861 FutureWarning,862 )863 if len(deprecated_arguments) > 0:864 raise ValueError(f"Got unexpected arguments: {deprecated_arguments}")865 866 return_dict = return_dict if return_dict is not None else self.config.use_return_dict867 868 transformer_outputs = self.transformer(869 input_ids,870 past_key_values=past_key_values,871 attention_mask=attention_mask,872 head_mask=head_mask,873 inputs_embeds=inputs_embeds,874 use_cache=use_cache,875 output_attentions=output_attentions,876 output_hidden_states=output_hidden_states,877 return_dict=return_dict,878 )879 880 hidden_states = transformer_outputs[0]881 logits = self.score(hidden_states)882 883 if input_ids is not None:884 batch_size = input_ids.shape[0]885 else:886 batch_size = inputs_embeds.shape[0]887 888 if self.config.pad_token_id is None and batch_size != 1:889 raise ValueError("Cannot handle batch sizes > 1 if no padding token is defined.")890 if self.config.pad_token_id is None:891 sequence_lengths = -1892 else:893 if input_ids is not None:894 sequence_lengths = torch.ne(input_ids, self.config.pad_token_id).sum(dim=-1) - 1895 else:896 sequence_lengths = -1897 logger.warning(898 f"{self.__class__.__name__} will not detect padding tokens in `inputs_embeds`. Results may be "899 "unexpected if using padding tokens in conjunction with `inputs_embeds.`"900 )901 902 pooled_logits = logits[torch.arange(batch_size, device=logits.device), sequence_lengths]903 904 loss = None905 if labels is not None:906 if self.config.problem_type is None:907 if self.num_labels == 1:908 self.config.problem_type = "regression"909 elif self.num_labels > 1 and (labels.dtype == torch.long or labels.dtype == torch.int):910 self.config.problem_type = "single_label_classification"911 else:912 self.config.problem_type = "multi_label_classification"913 914 if self.config.problem_type == "regression":915 loss_fct = MSELoss()916 if self.num_labels == 1:917 loss = loss_fct(pooled_logits.squeeze(), labels.squeeze())918 else:919 loss = loss_fct(pooled_logits, labels)920 elif self.config.problem_type == "single_label_classification":921 loss_fct = CrossEntropyLoss()922 loss = loss_fct(pooled_logits, labels)923 elif self.config.problem_type == "multi_label_classification":924 loss_fct = BCEWithLogitsLoss()925 loss = loss_fct(pooled_logits, labels)926 if not return_dict:927 output = (pooled_logits,) + transformer_outputs[1:]928 return ((loss,) + output) if loss is not None else output929 930 return SequenceClassifierOutputWithPast(931 loss=loss,932 logits=pooled_logits,933 past_key_values=transformer_outputs.past_key_values,934 hidden_states=transformer_outputs.hidden_states,935 attentions=transformer_outputs.attentions,936 )937 938 939class RWForTokenClassification(RWPreTrainedModel):940 _keys_to_ignore_on_load_missing = [r"h.*.self_attention.scale_mask_softmax.causal_mask", r"lm_head.weight"]941 942 def __init__(self, config: RWConfig):943 super().__init__(config)944 self.num_labels = config.num_labels945 946 self.transformer = RWModel(config)947 if hasattr(config, "classifier_dropout") and config.classifier_dropout is not None:948 classifier_dropout = config.classifier_dropout949 elif hasattr(config, "hidden_dropout") and config.hidden_dropout is not None:950 classifier_dropout = config.hidden_dropout951 else:952 classifier_dropout = 0.1953 self.dropout = nn.Dropout(classifier_dropout)954 self.classifier = nn.Linear(config.hidden_size, config.num_labels)955 956 # Initialize weights and apply final processing957 self.post_init()958 959 def forward(960 self,961 input_ids: Optional[torch.LongTensor] = None,962 past_key_values: Optional[Tuple[Tuple[torch.Tensor, torch.Tensor], ...]] = None,963 attention_mask: Optional[torch.Tensor] = None,964 head_mask: Optional[torch.Tensor] = None,965 inputs_embeds: Optional[torch.Tensor] = None,966 labels: Optional[torch.Tensor] = None,967 use_cache: Optional[bool] = None,968 output_attentions: Optional[bool] = None,969 output_hidden_states: Optional[bool] = None,970 return_dict: Optional[bool] = None,971 **deprecated_arguments,972 ) -> Union[Tuple[torch.Tensor], TokenClassifierOutput]:973 r"""974 labels (`torch.LongTensor` of shape `(batch_size,)`, *optional*):975 Labels for computing the sequence classification/regression loss. Indices should be in `[0, ...,976 config.num_labels - 1]`. If `config.num_labels == 1` a regression loss is computed (Mean-Square loss), If977 `config.num_labels > 1` a classification loss is computed (Cross-Entropy).978 """979 if deprecated_arguments.pop("position_ids", False) is not False:980 # `position_ids` could have been `torch.Tensor` or `None` so defaulting pop to `False` allows to detect if users were passing explicitly `None`981 warnings.warn(982 "`position_ids` have no functionality in BLOOM and will be removed in v5.0.0. You can safely ignore"983 " passing `position_ids`.",984 FutureWarning,985 )986 if len(deprecated_arguments) > 0:987 raise ValueError(f"Got unexpected arguments: {deprecated_arguments}")988 989 return_dict = return_dict if return_dict is not None else self.config.use_return_dict990 991 transformer_outputs = self.transformer(992 input_ids,993 past_key_values=past_key_values,994 attention_mask=attention_mask,995 head_mask=head_mask,996 inputs_embeds=inputs_embeds,997 use_cache=use_cache,998 output_attentions=output_attentions,999 output_hidden_states=output_hidden_states,1000 return_dict=return_dict,1001 )1002 1003 hidden_states = transformer_outputs[0]1004 hidden_states = self.dropout(hidden_states)1005 logits = self.classifier(hidden_states)1006 1007 loss = None1008 if labels is not None:1009 batch_size, seq_length = labels.shape1010 loss_fct = CrossEntropyLoss()1011 loss = loss_fct(logits.view(batch_size * seq_length, self.num_labels), labels.view(batch_size * seq_length))1012 1013 if not return_dict:1014 output = (logits,) + transformer_outputs[2:]1015 return ((loss,) + output) if loss is not None else output1016 1017 return TokenClassifierOutput(1018 loss=loss,1019 logits=logits,1020 hidden_states=transformer_outputs.hidden_states,1021 attentions=transformer_outputs.attentions,1022 )1023 1024 1025class RWForQuestionAnswering(RWPreTrainedModel):1026 _keys_to_ignore_on_load_missing = [r"h.*.self_attention.scale_mask_softmax.causal_mask", r"lm_head.weight"]1027 1028 def __init__(self, config):1029 super().__init__(config)1030 self.transformer = RWModel(config)1031 self.qa_outputs = nn.Linear(config.hidden_size, 2)1032 1033 # Initialize weights and apply final processing1034 self.post_init()1035 1036 def forward(1037 self,1038 input_ids: Optional[torch.LongTensor] = None,1039 attention_mask: Optional[torch.FloatTensor] = None,1040 position_ids: Optional[torch.LongTensor] = None,1041 head_mask: Optional[torch.FloatTensor] = None,1042 inputs_embeds: Optional[torch.FloatTensor] = None,1043 start_positions: Optional[torch.LongTensor] = None,1044 end_positions: Optional[torch.LongTensor] = None,1045 output_attentions: Optional[bool] = None,1046 output_hidden_states: Optional[bool] = None,1047 return_dict: Optional[bool] = None,1048 ) -> Union[Tuple, QuestionAnsweringModelOutput]:1049 r"""1050 start_positions (`torch.LongTensor` of shape `(batch_size,)`, *optional*):1051 Labels for position (index) of the start of the labelled span for computing the token classification loss.1052 Positions are clamped to the length of the sequence (`sequence_length`). Position outside of the sequence1053 are not taken into account for computing the loss.1054 end_positions (`torch.LongTensor` of shape `(batch_size,)`, *optional*):1055 Labels for position (index) of the end of the labelled span for computing the token classification loss.1056 Positions are clamped to the length of the sequence (`sequence_length`). Position outside of the sequence1057 are not taken into account for computing the loss.1058 """1059 return_dict = return_dict if return_dict is not None else self.config.use_return_dict1060 1061 outputs = self.transformer(1062 input_ids,1063 attention_mask=attention_mask,1064 position_ids=position_ids,1065 head_mask=head_mask,1066 inputs_embeds=inputs_embeds,1067 output_attentions=output_attentions,1068 output_hidden_states=output_hidden_states,1069 return_dict=return_dict,1070 )1071 1072 sequence_output = outputs[0]1073 1074 logits = self.qa_outputs(sequence_output)1075 start_logits, end_logits = logits.split(1, dim=-1)1076 start_logits = start_logits.squeeze(-1).contiguous()1077 end_logits = end_logits.squeeze(-1).contiguous()1078 1079 total_loss = None1080 if start_positions is not None and end_positions is not None:1081 # If we are on multi-GPU, split add a dimension1082 if len(start_positions.size()) > 1:1083 start_positions = start_positions.squeeze(-1)1084 if len(end_positions.size()) > 1:1085 end_positions = end_positions.squeeze(-1)1086 # sometimes the start/end positions are outside our model inputs, we ignore these terms1087 ignored_index = start_logits.size(1)1088 start_positions = start_positions.clamp(0, ignored_index)1089 end_positions = end_positions.clamp(0, ignored_index)1090 1091 loss_fct = CrossEntropyLoss(ignore_index=ignored_index)1092 start_loss = loss_fct(start_logits, start_positions)1093 end_loss = loss_fct(end_logits, end_positions)1094 total_loss = (start_loss + end_loss) / 21095 1096 if not return_dict:1097 output = (start_logits, end_logits) + outputs[2:]1098 return ((total_loss,) + output) if total_loss is not None else output1099 1100 return QuestionAnsweringModelOutput(1101 loss=total_loss,1102 start_logits=start_logits,1103 end_logits=end_logits,1104 hidden_states=outputs.hidden_states,1105 attentions=outputs.attentions,1106 )1107