feat(CRITICAL): 从 GitHub 扫描搬运 ixformer SDK + xllm 完整 GDN/MoE 代码
来源:
1. Chranos/ixformer (GitHub) → ixformer_sdk/ (230 files, 70K lines)
- inference/functions/vllm.py: vllm_moe_topk_softmax 完整实现 (2033 lines)
- inference/functions/moe.py: MoE ops 完整实现 (1380 lines)
- contrib/vllm_flash_attn/: FA2 Python 接口 (1018 lines)
- contrib/tgi/fused_moe.py: TGI fused MoE (429 lines)
- csrc/include/ixformer/: C++ kernel headers + cmake
2. Deep-Spark/xllm (GitHub) → upstream_ref/xllm_latest/ (+15 files)
- npu_torch/qwen3_5_decoder_layer_impl.cpp/.h
- npu_torch/qwen3_5_gated_delta_net.cpp/.h
- npu_torch/qwen3_next_*.cpp/.h (6 files)
- npu_torch/attention.cpp/.h + fused_moe.cpp/.h + CMakeLists.txt
- models/llm/qwen3_5.h + qwen3_5_mtp.h + qwen3_next.h
- models/vlm/qwen3_5.h
调用链完整性:
ixformer_sdk/inference/functions/vllm.py
→ ops.infer.moe_topk_softmax() (C++ 层)
→ 这就是 base 镜像 libixformer.so 里的实现
upstream_ref/xllm_latest/core/layers/ilu/fused_moe.cpp
→ ixformer::infer::topk_softmax() (直接 C++ 调用)
→ ixformer::infer::group_gemm() → 完整 7-step MoE pipeline
This commit is contained in:
project6-dev
0
ixformer_sdk/train/speedformer/models/__init__.py
Normal file
0
ixformer_sdk/train/speedformer/models/__init__.py
Normal file
@@ -0,0 +1,68 @@
|
||||
# Copyright 2023 Baichuan Inc. All Rights Reserved.
|
||||
|
||||
# Copyright 2022 EleutherAI and the HuggingFace Inc. team. All rights reserved.
|
||||
#
|
||||
# This code is based on EleutherAI's GPT-NeoX library and the GPT-NeoX
|
||||
# and OPT implementations in this library. It has been modified from its
|
||||
# original forms to accommodate minor architectural differences compared
|
||||
# to GPT-NeoX and OPT used by the Meta AI team that trained the model.
|
||||
#
|
||||
# Licensed under the Apache License, Version 2.0 (the "License");
|
||||
# you may not use this file except in compliance with the License.
|
||||
# You may obtain a copy of the License at
|
||||
#
|
||||
# http://www.apache.org/licenses/LICENSE-2.0
|
||||
#
|
||||
# Unless required by applicable law or agreed to in writing, software
|
||||
# distributed under the License is distributed on an "AS IS" BASIS,
|
||||
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
# See the License for the specific language governing permissions and
|
||||
# limitations under the License.
|
||||
|
||||
from transformers.configuration_utils import PretrainedConfig
|
||||
from transformers.utils import logging
|
||||
|
||||
|
||||
logger = logging.get_logger(__name__)
|
||||
|
||||
|
||||
class BaichuanConfig(PretrainedConfig):
|
||||
model_type = "baichuan"
|
||||
keys_to_ignore_at_inference = ["past_key_values"]
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
vocab_size=125696,
|
||||
hidden_size=4096,
|
||||
intermediate_size=11008,
|
||||
num_hidden_layers=32,
|
||||
num_attention_heads=32,
|
||||
hidden_act="silu",
|
||||
max_position_embeddings=4096,
|
||||
initializer_range=0.02,
|
||||
rms_norm_eps=1e-6,
|
||||
use_cache=True,
|
||||
pad_token_id=0,
|
||||
bos_token_id=1,
|
||||
eos_token_id=2,
|
||||
tie_word_embeddings=False,
|
||||
**kwargs,
|
||||
):
|
||||
self.vocab_size = vocab_size
|
||||
self.max_position_embeddings = max_position_embeddings
|
||||
self.hidden_size = hidden_size
|
||||
self.intermediate_size = intermediate_size
|
||||
self.num_hidden_layers = num_hidden_layers
|
||||
self.num_attention_heads = num_attention_heads
|
||||
self.hidden_act = hidden_act
|
||||
self.initializer_range = initializer_range
|
||||
self.rms_norm_eps = rms_norm_eps
|
||||
self.use_cache = use_cache
|
||||
self.z_loss_weight = 0
|
||||
super().__init__(
|
||||
pad_token_id=pad_token_id,
|
||||
bos_token_id=bos_token_id,
|
||||
eos_token_id=eos_token_id,
|
||||
tie_word_embeddings=tie_word_embeddings,
|
||||
**kwargs,
|
||||
)
|
||||
@@ -0,0 +1,83 @@
|
||||
from typing import List
|
||||
from queue import Queue
|
||||
|
||||
import torch
|
||||
|
||||
|
||||
def build_chat_input(model, tokenizer, messages: List[dict], max_new_tokens: int=0):
|
||||
def _parse_messages(messages, split_role="user"):
|
||||
system, rounds = "", []
|
||||
round = []
|
||||
for i, message in enumerate(messages):
|
||||
if message["role"] == "system":
|
||||
assert i == 0
|
||||
system = message["content"]
|
||||
continue
|
||||
if message["role"] == split_role and round:
|
||||
rounds.append(round)
|
||||
round = []
|
||||
round.append(message)
|
||||
if round:
|
||||
rounds.append(round)
|
||||
return system, rounds
|
||||
|
||||
max_new_tokens = max_new_tokens or model.generation_config.max_new_tokens
|
||||
max_input_tokens = model.config.model_max_length - max_new_tokens
|
||||
system, rounds = _parse_messages(messages, split_role="user")
|
||||
system_tokens = tokenizer.encode(system)
|
||||
max_history_tokens = max_input_tokens - len(system_tokens)
|
||||
|
||||
history_tokens = []
|
||||
for round in rounds[::-1]:
|
||||
round_tokens = []
|
||||
for message in round:
|
||||
if message["role"] == "user":
|
||||
round_tokens.append(model.generation_config.user_token_id)
|
||||
else:
|
||||
round_tokens.append(model.generation_config.assistant_token_id)
|
||||
round_tokens.extend(tokenizer.encode(message["content"]))
|
||||
if len(history_tokens) == 0 or len(history_tokens) + len(round_tokens) <= max_history_tokens:
|
||||
history_tokens = round_tokens + history_tokens # concat left
|
||||
if len(history_tokens) < max_history_tokens:
|
||||
continue
|
||||
break
|
||||
|
||||
input_tokens = system_tokens + history_tokens
|
||||
if messages[-1]["role"] != "assistant":
|
||||
input_tokens.append(model.generation_config.assistant_token_id)
|
||||
input_tokens = input_tokens[-max_input_tokens:] # truncate left
|
||||
return torch.LongTensor([input_tokens]).to(model.device)
|
||||
|
||||
|
||||
class TextIterStreamer:
|
||||
def __init__(self, tokenizer, skip_prompt=False, skip_special_tokens=False):
|
||||
self.tokenizer = tokenizer
|
||||
self.skip_prompt = skip_prompt
|
||||
self.skip_special_tokens = skip_special_tokens
|
||||
self.tokens = []
|
||||
self.text_queue = Queue()
|
||||
self.next_tokens_are_prompt = True
|
||||
|
||||
def put(self, value):
|
||||
if self.skip_prompt and self.next_tokens_are_prompt:
|
||||
self.next_tokens_are_prompt = False
|
||||
else:
|
||||
if len(value.shape) > 1:
|
||||
value = value[0]
|
||||
self.tokens.extend(value.tolist())
|
||||
self.text_queue.put(
|
||||
self.tokenizer.decode(self.tokens, skip_special_tokens=self.skip_special_tokens))
|
||||
|
||||
def end(self):
|
||||
self.text_queue.put(None)
|
||||
|
||||
def __iter__(self):
|
||||
return self
|
||||
|
||||
def __next__(self):
|
||||
value = self.text_queue.get()
|
||||
if value is None:
|
||||
raise StopIteration()
|
||||
else:
|
||||
return value
|
||||
|
||||
@@ -0,0 +1,783 @@
|
||||
# Copyright 2023 Baichuan Inc. All Rights Reserved.
|
||||
|
||||
# Copyright 2022 EleutherAI and the HuggingFace Inc. team. All rights reserved.
|
||||
#
|
||||
# This code is based on EleutherAI's GPT-NeoX library and the GPT-NeoX
|
||||
# and OPT implementations in this library. It has been modified from its
|
||||
# original forms to accommodate minor architectural differences compared
|
||||
# to GPT-NeoX and OPT used by the Meta AI team that trained the model.
|
||||
#
|
||||
# Licensed under the Apache License, Version 2.0 (the "License");
|
||||
# you may not use this file except in compliance with the License.
|
||||
# You may obtain a copy of the License at
|
||||
#
|
||||
# http://www.apache.org/licenses/LICENSE-2.0
|
||||
#
|
||||
# Unless required by applicable law or agreed to in writing, software
|
||||
# distributed under the License is distributed on an "AS IS" BASIS,
|
||||
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
# See the License for the specific language governing permissions and
|
||||
# limitations under the License.
|
||||
|
||||
|
||||
from .configuration_baichuan import BaichuanConfig
|
||||
from .generation_utils import build_chat_input, TextIterStreamer
|
||||
|
||||
import math
|
||||
from typing import List, Optional, Tuple, Union
|
||||
from threading import Thread
|
||||
|
||||
import torch
|
||||
import torch.utils.checkpoint
|
||||
from torch import nn
|
||||
from torch.nn import BCEWithLogitsLoss, CrossEntropyLoss, MSELoss
|
||||
from torch.nn import functional as F
|
||||
from transformers import PreTrainedModel, PretrainedConfig
|
||||
from transformers.activations import ACT2FN
|
||||
from transformers.modeling_outputs import BaseModelOutputWithPast, CausalLMOutputWithPast
|
||||
from transformers.generation.utils import GenerationConfig
|
||||
from transformers.utils import logging, ContextManagers
|
||||
|
||||
import os
|
||||
from contextlib import contextmanager
|
||||
logger = logging.get_logger(__name__)
|
||||
|
||||
try:
|
||||
from xformers import ops as xops
|
||||
except ImportError:
|
||||
xops = None
|
||||
logger.warning(
|
||||
"Xformers is not installed correctly. If you want to use memory_efficient_attention to accelerate training use the following command to install Xformers\npip install xformers."
|
||||
)
|
||||
|
||||
|
||||
# Copied from transformers.models.bart.modeling_bart._make_causal_mask
|
||||
def _make_causal_mask(
|
||||
input_ids_shape: torch.Size, dtype: torch.dtype, device: torch.device, past_key_values_length: int = 0
|
||||
):
|
||||
"""
|
||||
Make causal mask used for bi-directional self-attention.
|
||||
"""
|
||||
bsz, tgt_len = input_ids_shape
|
||||
mask = torch.full((tgt_len, tgt_len), torch.tensor(torch.finfo(dtype).min, device=device), device=device)
|
||||
mask_cond = torch.arange(mask.size(-1), device=device)
|
||||
mask.masked_fill_(mask_cond < (mask_cond + 1).view(mask.size(-1), 1), 0)
|
||||
mask = mask.to(dtype)
|
||||
|
||||
if past_key_values_length > 0:
|
||||
mask = torch.cat([torch.zeros(tgt_len, past_key_values_length, dtype=dtype, device=device), mask], dim=-1)
|
||||
return mask[None, None, :, :].expand(bsz, 1, tgt_len, tgt_len + past_key_values_length)
|
||||
|
||||
def _expand_mask(mask: torch.Tensor, dtype: torch.dtype, tgt_len: Optional[int] = None):
|
||||
"""
|
||||
Expands attention_mask from `[bsz, seq_len]` to `[bsz, 1, tgt_seq_len, src_seq_len]`.
|
||||
"""
|
||||
if len(mask.size()) == 3:
|
||||
bsz, src_len, _ = mask.size()
|
||||
tgt_len = tgt_len if tgt_len is not None else src_len
|
||||
expanded_mask = mask[:,None,:,:].expand(bsz, 1, tgt_len, src_len).to(dtype)
|
||||
else:
|
||||
bsz, src_len = mask.size()
|
||||
tgt_len = tgt_len if tgt_len is not None else src_len
|
||||
expanded_mask = mask[:, None, None, :].expand(bsz, 1, tgt_len, src_len).to(dtype)
|
||||
|
||||
inverted_mask = 1.0 - expanded_mask
|
||||
|
||||
return inverted_mask.masked_fill(inverted_mask.to(torch.bool), torch.finfo(dtype).min)
|
||||
|
||||
|
||||
class RMSNorm(nn.Module):
|
||||
def __init__(self, hidden_size, eps=1e-6):
|
||||
"""
|
||||
RMSNorm is equivalent to T5LayerNorm
|
||||
"""
|
||||
super().__init__()
|
||||
self.weight = nn.Parameter(torch.ones(hidden_size))
|
||||
self.variance_epsilon = eps
|
||||
|
||||
def forward(self, hidden_states):
|
||||
variance = hidden_states.to(torch.float32).pow(2).mean(-1, keepdim=True)
|
||||
hidden_states = hidden_states * torch.rsqrt(variance + self.variance_epsilon)
|
||||
|
||||
# convert into half-precision if necessary
|
||||
if self.weight.dtype in [torch.float16, torch.bfloat16]:
|
||||
hidden_states = hidden_states.to(self.weight.dtype)
|
||||
|
||||
return self.weight * hidden_states
|
||||
|
||||
|
||||
class RotaryEmbedding(torch.nn.Module):
|
||||
def __init__(self, dim, max_position_embeddings=2048, base=10000, device=None):
|
||||
super().__init__()
|
||||
self.inv_freq = 1.0 / (base ** (torch.arange(0, dim, 2).float().to(device) / dim))
|
||||
self.max_seq_len_cached = max_position_embeddings
|
||||
t = torch.arange(self.max_seq_len_cached, device=self.inv_freq.device, dtype=torch.float32)
|
||||
freqs = torch.outer(t, self.inv_freq)
|
||||
emb = torch.cat((freqs, freqs), dim=-1)
|
||||
self.cos_cached = emb.cos()[None, None, :, :].to(torch.float32)
|
||||
self.sin_cached = emb.sin()[None, None, :, :].to(torch.float32)
|
||||
def forward(self, x, seq_len=None):
|
||||
# x: [bs, num_attention_heads, seq_len, head_size]
|
||||
# This `if` block is unlikely to be run after we build sin/cos in `__init__`. Keep the logic here just in case.
|
||||
if seq_len > self.max_seq_len_cached:
|
||||
self.max_seq_len_cached = seq_len
|
||||
t = torch.arange(self.max_seq_len_cached, device=self.inv_freq.device, dtype=torch.float32)
|
||||
freqs = torch.outer(t, self.inv_freq)
|
||||
emb = torch.cat((freqs, freqs), dim=-1)
|
||||
self.cos_cached = emb.cos()[None, None, :, :].to(torch.float32).to(x.device)
|
||||
self.sin_cached = emb.sin()[None, None, :, :].to(torch.float32).to(x.device)
|
||||
elif self.cos_cached.device != x.device:
|
||||
self.cos_cached = self.cos_cached.to(x.device)
|
||||
self.sin_cached = self.sin_cached.to(x.device)
|
||||
return (
|
||||
self.cos_cached[:, :, :seq_len, ...],
|
||||
self.sin_cached[:, :, :seq_len, ...],
|
||||
)
|
||||
|
||||
|
||||
def rotate_half(x):
|
||||
"""Rotates half the hidden dims of the input."""
|
||||
x1 = x[..., : x.shape[-1] // 2]
|
||||
x2 = x[..., x.shape[-1] // 2:]
|
||||
return torch.cat((-x2, x1), dim=-1)
|
||||
|
||||
|
||||
def apply_rotary_pos_emb(q, k, cos_, sin_, position_ids):
|
||||
cos = cos_.squeeze(1).squeeze(0) # [seq_len, dim]
|
||||
sin = sin_.squeeze(1).squeeze(0) # [seq_len, dim]
|
||||
cos = cos[position_ids].unsqueeze(1) # [bs, 1, seq_len, dim]
|
||||
sin = sin[position_ids].unsqueeze(1) # [bs, 1, seq_len, dim]
|
||||
q_embed = (q.float() * cos) + (rotate_half(q.float()) * sin)
|
||||
k_embed = (k.float() * cos) + (rotate_half(k.float()) * sin)
|
||||
return q_embed.to(q.dtype), k_embed.to(k.dtype)
|
||||
|
||||
|
||||
class MLP(nn.Module):
|
||||
def __init__(
|
||||
self,
|
||||
hidden_size: int,
|
||||
intermediate_size: int,
|
||||
hidden_act: str,
|
||||
):
|
||||
super().__init__()
|
||||
self.gate_proj = nn.Linear(hidden_size, intermediate_size, bias=False)
|
||||
self.down_proj = nn.Linear(intermediate_size, hidden_size, bias=False)
|
||||
self.up_proj = nn.Linear(hidden_size, intermediate_size, bias=False)
|
||||
self.act_fn = ACT2FN[hidden_act]
|
||||
|
||||
def forward(self, x):
|
||||
return self.down_proj(self.act_fn(self.gate_proj(x)) * self.up_proj(x))
|
||||
|
||||
|
||||
class Attention(nn.Module):
|
||||
"""Multi-headed attention from 'Attention Is All You Need' paper"""
|
||||
def __init__(self, config: BaichuanConfig):
|
||||
super().__init__()
|
||||
self.config = config
|
||||
self.hidden_size = config.hidden_size
|
||||
self.num_heads = config.num_attention_heads
|
||||
self.head_dim = self.hidden_size // self.num_heads
|
||||
self.max_position_embeddings = config.max_position_embeddings
|
||||
|
||||
if (self.head_dim * self.num_heads) != self.hidden_size:
|
||||
raise ValueError(
|
||||
f"hidden_size must be divisible by num_heads (got `hidden_size`: {self.hidden_size}"
|
||||
f" and `num_heads`: {self.num_heads})."
|
||||
)
|
||||
self.W_pack = nn.Linear(self.hidden_size, 3 * self.hidden_size, bias=False)
|
||||
self.o_proj = nn.Linear(self.num_heads * self.head_dim, self.hidden_size, bias=False)
|
||||
self.rotary_emb = RotaryEmbedding(self.head_dim, max_position_embeddings=self.max_position_embeddings)
|
||||
|
||||
def _shape(self, tensor: torch.Tensor, seq_len: int, bsz: int):
|
||||
return tensor.view(bsz, seq_len, self.num_heads, self.head_dim).transpose(1, 2).contiguous()
|
||||
|
||||
def forward(
|
||||
self,
|
||||
hidden_states: torch.Tensor,
|
||||
attention_mask: Optional[torch.Tensor] = None,
|
||||
position_ids: Optional[torch.LongTensor] = None,
|
||||
past_key_value: Optional[Tuple[torch.Tensor]] = None,
|
||||
output_attentions: bool = False,
|
||||
use_cache: bool = False,
|
||||
) -> Tuple[torch.Tensor, Optional[torch.Tensor], Optional[Tuple[torch.Tensor]]]:
|
||||
bsz, q_len, _ = hidden_states.size()
|
||||
|
||||
proj = self.W_pack(hidden_states)
|
||||
proj = proj.unflatten(-1, (3, self.hidden_size)).unsqueeze(0).transpose(0, -2).squeeze(-2)
|
||||
query_states = proj[0].view(bsz, q_len, self.num_heads, self.head_dim).transpose(1, 2)
|
||||
key_states = proj[1].view(bsz, q_len, self.num_heads, self.head_dim).transpose(1, 2)
|
||||
value_states = proj[2].view(bsz, q_len, self.num_heads, self.head_dim).transpose(1, 2)
|
||||
|
||||
kv_seq_len = key_states.shape[-2]
|
||||
if past_key_value is not None:
|
||||
kv_seq_len += past_key_value[0].shape[-2]
|
||||
cos, sin = self.rotary_emb(value_states, seq_len=kv_seq_len)
|
||||
query_states, key_states = apply_rotary_pos_emb(query_states, key_states, cos, sin, position_ids)
|
||||
# [bsz, nh, t, hd]
|
||||
|
||||
if past_key_value is not None:
|
||||
# reuse k, v, self_attention
|
||||
key_states = torch.cat([past_key_value[0], key_states], dim=2)
|
||||
value_states = torch.cat([past_key_value[1], value_states], dim=2)
|
||||
|
||||
past_key_value = (key_states, value_states) if use_cache else None
|
||||
if xops is not None and self.training:
|
||||
attn_weights = None
|
||||
query_states = query_states.transpose(1, 2)
|
||||
key_states = key_states.transpose(1, 2)
|
||||
value_states = value_states.transpose(1, 2)
|
||||
attn_output = xops.memory_efficient_attention(
|
||||
query_states, key_states, value_states, attn_bias=xops.LowerTriangularMask()
|
||||
)
|
||||
else:
|
||||
with torch.backends.cuda.sdp_kernel(enable_flash=True, enable_math=True, enable_mem_efficient=True):
|
||||
attn_output = F.scaled_dot_product_attention(query_states, key_states, value_states, attn_mask = attention_mask)
|
||||
attn_output = attn_output.transpose(1, 2)
|
||||
attn_output = attn_output.reshape(bsz, q_len, self.hidden_size)
|
||||
attn_output = self.o_proj(attn_output)
|
||||
|
||||
if not output_attentions:
|
||||
attn_weights = None
|
||||
|
||||
return attn_output, attn_weights, past_key_value
|
||||
|
||||
|
||||
class DecoderLayer(nn.Module):
|
||||
def __init__(self, config: BaichuanConfig):
|
||||
super().__init__()
|
||||
self.hidden_size = config.hidden_size
|
||||
self.self_attn = Attention(config=config)
|
||||
self.mlp = MLP(
|
||||
hidden_size=self.hidden_size,
|
||||
intermediate_size=config.intermediate_size,
|
||||
hidden_act=config.hidden_act,
|
||||
)
|
||||
self.input_layernorm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
|
||||
self.post_attention_layernorm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
|
||||
|
||||
def forward(
|
||||
self,
|
||||
hidden_states: torch.Tensor,
|
||||
attention_mask: Optional[torch.Tensor] = None,
|
||||
position_ids: Optional[torch.LongTensor] = None,
|
||||
past_key_value: Optional[Tuple[torch.Tensor]] = None,
|
||||
output_attentions: Optional[bool] = False,
|
||||
use_cache: Optional[bool] = False,
|
||||
) -> Tuple[torch.FloatTensor, Optional[Tuple[torch.FloatTensor, torch.FloatTensor]]]:
|
||||
|
||||
residual = hidden_states
|
||||
|
||||
hidden_states = self.input_layernorm(hidden_states)
|
||||
|
||||
# Self Attention
|
||||
hidden_states, self_attn_weights, present_key_value = self.self_attn(
|
||||
hidden_states=hidden_states,
|
||||
attention_mask=attention_mask,
|
||||
position_ids=position_ids,
|
||||
past_key_value=past_key_value,
|
||||
output_attentions=output_attentions,
|
||||
use_cache=use_cache,
|
||||
)
|
||||
hidden_states = residual + hidden_states
|
||||
|
||||
# Fully Connected
|
||||
residual = hidden_states
|
||||
hidden_states = self.post_attention_layernorm(hidden_states)
|
||||
hidden_states = self.mlp(hidden_states)
|
||||
hidden_states = residual + hidden_states
|
||||
|
||||
outputs = (hidden_states,)
|
||||
|
||||
if output_attentions:
|
||||
outputs += (self_attn_weights,)
|
||||
|
||||
if use_cache:
|
||||
outputs += (present_key_value,)
|
||||
|
||||
return outputs
|
||||
|
||||
|
||||
class BaichuanPreTrainedModel(PreTrainedModel):
|
||||
config_class = BaichuanConfig
|
||||
base_model_prefix = "model"
|
||||
supports_gradient_checkpointing = True
|
||||
_no_split_modules = ["DecoderLayer"]
|
||||
_keys_to_ignore_on_load_unexpected = [r"decoder\.version"]
|
||||
|
||||
def _init_weights(self, module):
|
||||
std = self.config.initializer_range
|
||||
if isinstance(module, nn.Linear):
|
||||
module.weight.data.normal_(mean=0.0, std=std)
|
||||
if module.bias is not None:
|
||||
module.bias.data.zero_()
|
||||
elif isinstance(module, nn.Embedding):
|
||||
module.weight.data.normal_(mean=0.0, std=std)
|
||||
if module.padding_idx is not None:
|
||||
module.weight.data[module.padding_idx].zero_()
|
||||
|
||||
def _set_gradient_checkpointing(self, module, value=False):
|
||||
if isinstance(module, BaichuanModel):
|
||||
module.gradient_checkpointing = value
|
||||
|
||||
|
||||
class BaichuanModel(BaichuanPreTrainedModel):
|
||||
def __init__(self, config: BaichuanConfig):
|
||||
super().__init__(config)
|
||||
self.padding_idx = config.pad_token_id
|
||||
self.vocab_size = config.vocab_size
|
||||
|
||||
self.embed_tokens = nn.Embedding(config.vocab_size, config.hidden_size, self.padding_idx)
|
||||
self.layers = nn.ModuleList([DecoderLayer(config) for _ in range(config.num_hidden_layers)])
|
||||
self.norm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
|
||||
|
||||
self.gradient_checkpointing = False
|
||||
# Initialize weights and apply final processing
|
||||
self.post_init()
|
||||
|
||||
def get_input_embeddings(self):
|
||||
return self.embed_tokens
|
||||
|
||||
def set_input_embeddings(self, value):
|
||||
self.embed_tokens = value
|
||||
|
||||
# Copied from transformers.models.bart.modeling_bart.BartDecoder._prepare_decoder_attention_mask
|
||||
def _prepare_decoder_attention_mask(self, attention_mask, input_shape, inputs_embeds, past_key_values_length):
|
||||
# create causal mask
|
||||
# [bsz, seq_len] -> [bsz, 1, tgt_seq_len, src_seq_len]
|
||||
combined_attention_mask = None
|
||||
if input_shape[-1] > 1:
|
||||
combined_attention_mask = _make_causal_mask(
|
||||
input_shape,
|
||||
inputs_embeds.dtype,
|
||||
device=inputs_embeds.device,
|
||||
past_key_values_length=past_key_values_length,
|
||||
)
|
||||
|
||||
if attention_mask is not None:
|
||||
# [bsz, seq_len] -> [bsz, 1, tgt_seq_len, src_seq_len]
|
||||
expanded_attn_mask = _expand_mask(attention_mask, inputs_embeds.dtype, tgt_len=input_shape[-1]).to(
|
||||
inputs_embeds.device
|
||||
)
|
||||
combined_attention_mask = (
|
||||
expanded_attn_mask if combined_attention_mask is None else expanded_attn_mask + combined_attention_mask
|
||||
)
|
||||
|
||||
return combined_attention_mask
|
||||
|
||||
def forward(
|
||||
self,
|
||||
input_ids: torch.LongTensor = None,
|
||||
attention_mask: Optional[torch.Tensor] = None,
|
||||
position_ids: Optional[torch.LongTensor] = None,
|
||||
past_key_values: Optional[List[torch.FloatTensor]] = None,
|
||||
inputs_embeds: Optional[torch.FloatTensor] = None,
|
||||
use_cache: Optional[bool] = None,
|
||||
output_attentions: Optional[bool] = None,
|
||||
output_hidden_states: Optional[bool] = None,
|
||||
return_dict: Optional[bool] = None,
|
||||
) -> Union[Tuple, BaseModelOutputWithPast]:
|
||||
output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions
|
||||
output_hidden_states = (
|
||||
output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
|
||||
)
|
||||
use_cache = use_cache if use_cache is not None else self.config.use_cache
|
||||
|
||||
return_dict = return_dict if return_dict is not None else self.config.use_return_dict
|
||||
|
||||
# retrieve input_ids and inputs_embeds
|
||||
if input_ids is not None and inputs_embeds is not None:
|
||||
raise ValueError("You cannot specify both decoder_input_ids and decoder_inputs_embeds at the same time")
|
||||
elif input_ids is not None:
|
||||
batch_size, seq_length = input_ids.shape
|
||||
elif inputs_embeds is not None:
|
||||
batch_size, seq_length, _ = inputs_embeds.shape
|
||||
else:
|
||||
raise ValueError("You have to specify either decoder_input_ids or decoder_inputs_embeds")
|
||||
|
||||
seq_length_with_past = seq_length
|
||||
past_key_values_length = 0
|
||||
|
||||
if past_key_values is not None:
|
||||
past_key_values_length = past_key_values[0][0].shape[2]
|
||||
seq_length_with_past = seq_length_with_past + past_key_values_length
|
||||
|
||||
if position_ids is None:
|
||||
device = input_ids.device if input_ids is not None else inputs_embeds.device
|
||||
position_ids = torch.arange(
|
||||
past_key_values_length, seq_length + past_key_values_length, dtype=torch.long, device=device
|
||||
)
|
||||
position_ids = position_ids.unsqueeze(0).view(-1, seq_length)
|
||||
else:
|
||||
position_ids = position_ids.view(-1, seq_length).long()
|
||||
|
||||
if inputs_embeds is None:
|
||||
inputs_embeds = self.embed_tokens(input_ids)
|
||||
# embed positions
|
||||
if attention_mask is None:
|
||||
attention_mask = torch.ones(
|
||||
(batch_size, seq_length_with_past), dtype=torch.bool, device=inputs_embeds.device
|
||||
)
|
||||
attention_mask = self._prepare_decoder_attention_mask(
|
||||
attention_mask, (batch_size, seq_length), inputs_embeds, past_key_values_length
|
||||
)
|
||||
|
||||
hidden_states = inputs_embeds
|
||||
|
||||
if self.gradient_checkpointing and self.training:
|
||||
if use_cache:
|
||||
logger.warning_once(
|
||||
"`use_cache=True` is incompatible with gradient checkpointing. Setting `use_cache=False`..."
|
||||
)
|
||||
use_cache = False
|
||||
|
||||
# decoder layers
|
||||
all_hidden_states = () if output_hidden_states else None
|
||||
all_self_attns = () if output_attentions else None
|
||||
next_decoder_cache = () if use_cache else None
|
||||
|
||||
for idx, decoder_layer in enumerate(self.layers):
|
||||
if output_hidden_states:
|
||||
all_hidden_states += (hidden_states,)
|
||||
|
||||
past_key_value = past_key_values[idx] if past_key_values is not None else None
|
||||
|
||||
if self.gradient_checkpointing and self.training:
|
||||
|
||||
def create_custom_forward(module):
|
||||
def custom_forward(*inputs):
|
||||
# None for past_key_value
|
||||
return module(*inputs, output_attentions, None)
|
||||
|
||||
return custom_forward
|
||||
|
||||
layer_outputs = torch.utils.checkpoint.checkpoint(
|
||||
create_custom_forward(decoder_layer),
|
||||
hidden_states,
|
||||
attention_mask,
|
||||
position_ids,
|
||||
None,
|
||||
)
|
||||
else:
|
||||
layer_outputs = decoder_layer(
|
||||
hidden_states,
|
||||
attention_mask=attention_mask,
|
||||
position_ids=position_ids,
|
||||
past_key_value=past_key_value,
|
||||
output_attentions=output_attentions,
|
||||
use_cache=use_cache,
|
||||
)
|
||||
|
||||
hidden_states = layer_outputs[0]
|
||||
|
||||
if use_cache:
|
||||
next_decoder_cache += (layer_outputs[2 if output_attentions else 1],)
|
||||
|
||||
if output_attentions:
|
||||
all_self_attns += (layer_outputs[1],)
|
||||
|
||||
hidden_states = self.norm(hidden_states)
|
||||
|
||||
# add hidden states from the last decoder layer
|
||||
if output_hidden_states:
|
||||
all_hidden_states += (hidden_states,)
|
||||
|
||||
next_cache = next_decoder_cache if use_cache else None
|
||||
if not return_dict:
|
||||
return tuple(v for v in [hidden_states, next_cache, all_hidden_states, all_self_attns] if v is not None)
|
||||
return BaseModelOutputWithPast(
|
||||
last_hidden_state=hidden_states,
|
||||
past_key_values=next_cache,
|
||||
hidden_states=all_hidden_states,
|
||||
attentions=all_self_attns,
|
||||
)
|
||||
|
||||
|
||||
class NormHead(nn.Module):
|
||||
def __init__(self, hidden_size, vocab_size, bias=False):
|
||||
super().__init__()
|
||||
self.weight = nn.Parameter(torch.empty((vocab_size, hidden_size)))
|
||||
nn.init.kaiming_uniform_(self.weight, a=math.sqrt(5))
|
||||
self.first_flag = True
|
||||
|
||||
def forward(self, hidden_states):
|
||||
if self.training:
|
||||
norm_weight = nn.functional.normalize(self.weight)
|
||||
elif self.first_flag:
|
||||
self.first_flag = False
|
||||
self.weight = nn.Parameter(nn.functional.normalize(self.weight))
|
||||
norm_weight = self.weight
|
||||
else:
|
||||
norm_weight = self.weight
|
||||
return nn.functional.linear(hidden_states, norm_weight)
|
||||
|
||||
_init_weights = True
|
||||
@contextmanager
|
||||
def no_init_weights(_enable=True):
|
||||
global _init_weights
|
||||
old_init_weights = _init_weights
|
||||
if _enable:
|
||||
_init_weights = False
|
||||
try:
|
||||
yield
|
||||
finally:
|
||||
_init_weights = old_init_weights
|
||||
|
||||
class BaichuanForCausalLM(BaichuanPreTrainedModel):
|
||||
def __init__(self, config, *model_args, **model_kwargs):
|
||||
super().__init__(config, *model_args, **model_kwargs)
|
||||
self.model = BaichuanModel(config)
|
||||
|
||||
self.lm_head = NormHead(config.hidden_size, config.vocab_size, bias=False)
|
||||
if hasattr(config, "quantization_config") and config.quantization_config['load_in_4bit']:
|
||||
try:
|
||||
from .quantizer import quantize_offline, init_model_weight_int4
|
||||
except ImportError:
|
||||
raise ImportError(f"Needs QLinear to run quantize.")
|
||||
quantize_offline(self, 4)
|
||||
# Initialize weights and apply final processing
|
||||
self.post_init()
|
||||
|
||||
def get_input_embeddings(self):
|
||||
return self.model.embed_tokens
|
||||
|
||||
def set_input_embeddings(self, value):
|
||||
self.model.embed_tokens = value
|
||||
|
||||
def get_output_embeddings(self):
|
||||
return self.lm_head
|
||||
|
||||
def set_output_embeddings(self, new_embeddings):
|
||||
self.lm_head = new_embeddings
|
||||
|
||||
def set_decoder(self, decoder):
|
||||
self.model = decoder
|
||||
|
||||
def get_decoder(self):
|
||||
return self.model
|
||||
|
||||
@classmethod
|
||||
def from_pretrained(
|
||||
cls,
|
||||
pretrained_model_name_or_path: Optional[Union[str, os.PathLike]],
|
||||
*model_args,
|
||||
config: Optional[Union[PretrainedConfig, str, os.PathLike]] = None,
|
||||
cache_dir: Optional[Union[str, os.PathLike]] = None,
|
||||
ignore_mismatched_sizes: bool = False,
|
||||
force_download: bool = False,
|
||||
local_files_only: bool = False,
|
||||
token: Optional[Union[str, bool]] = None,
|
||||
revision: str = "main",
|
||||
use_safetensors: bool = None,
|
||||
**kwargs,
|
||||
):
|
||||
# Load config if we don't provide a configuration
|
||||
if not isinstance(config, PretrainedConfig):
|
||||
config_path = config if config is not None else pretrained_model_name_or_path
|
||||
config, model_kwargs = cls.config_class.from_pretrained(
|
||||
config_path,
|
||||
cache_dir=cache_dir,
|
||||
return_unused_kwargs=True,
|
||||
force_download=force_download,
|
||||
resume_download=False,
|
||||
proxies=None,
|
||||
local_files_only=local_files_only,
|
||||
token=token,
|
||||
revision=revision,
|
||||
subfolder="",
|
||||
_from_auto=False,
|
||||
_from_pipeline=None,
|
||||
**kwargs,
|
||||
)
|
||||
else:
|
||||
model_kwargs = kwargs
|
||||
|
||||
if hasattr(config, "quantization_config") and config.quantization_config['load_in_4bit']:
|
||||
try:
|
||||
from .quantizer import init_model_weight_int4
|
||||
from accelerate import init_empty_weights, dispatch_model, infer_auto_device_map
|
||||
from accelerate.utils import CustomDtype
|
||||
from accelerate.utils import get_balanced_memory
|
||||
except ImportError:
|
||||
raise ImportError(f"Needs import model weight init func to run quantize.")
|
||||
# Instantiate model.
|
||||
init_contexts = [no_init_weights(_enable=True)]
|
||||
init_contexts.append(init_empty_weights())
|
||||
with ContextManagers(init_contexts):
|
||||
model = cls(config)
|
||||
|
||||
model_file = os.path.join(pretrained_model_name_or_path, 'pytorch_model.bin')
|
||||
state_dict = torch.load(model_file, map_location="cpu")
|
||||
model.is_quantized = True
|
||||
|
||||
device_map = kwargs.pop("device_map", None)
|
||||
torch_dtype = kwargs.pop("torch_dtype", None)
|
||||
|
||||
kwargs = {"no_split_module_classes": model._no_split_modules}
|
||||
target_dtype = CustomDtype.INT4
|
||||
max_memory = get_balanced_memory(
|
||||
model,
|
||||
dtype=target_dtype,
|
||||
low_zero=(device_map == "balanced_low_0"),
|
||||
max_memory=None,
|
||||
**kwargs,
|
||||
)
|
||||
kwargs["max_memory"] = max_memory
|
||||
|
||||
device_map = infer_auto_device_map(model, dtype=target_dtype, **kwargs)
|
||||
model = init_model_weight_int4(config, model, state_dict)
|
||||
|
||||
# Set model in evaluation mode to deactivate DropOut modules by default
|
||||
model.eval()
|
||||
# If it is a model with generation capabilities, attempt to load the generation config
|
||||
if model.can_generate():
|
||||
try:
|
||||
model.generation_config = GenerationConfig.from_pretrained(
|
||||
pretrained_model_name_or_path,
|
||||
cache_dir=cache_dir,
|
||||
force_download=force_download,
|
||||
resume_download=False,
|
||||
proxies=None,
|
||||
local_files_only=local_files_only,
|
||||
token=token,
|
||||
revision=revision,
|
||||
subfolder="",
|
||||
_from_auto=False,
|
||||
_from_pipeline=None,
|
||||
**kwargs,
|
||||
)
|
||||
except (OSError, TypeError):
|
||||
logger.info(
|
||||
"Generation config file not found, using a generation config created from the model config."
|
||||
)
|
||||
pass
|
||||
|
||||
if device_map is not None:
|
||||
dispatch_model(model, device_map=device_map)
|
||||
|
||||
return model
|
||||
return super(BaichuanForCausalLM, cls).from_pretrained(pretrained_model_name_or_path, *model_args,
|
||||
config=config, cache_dir=cache_dir, ignore_mismatched_sizes=ignore_mismatched_sizes,
|
||||
force_download=force_download, local_files_only=local_files_only, token=token, revision=revision,
|
||||
use_safetensors=use_safetensors, **kwargs)
|
||||
|
||||
def forward(
|
||||
self,
|
||||
input_ids: torch.LongTensor = None,
|
||||
attention_mask: Optional[torch.Tensor] = None,
|
||||
position_ids: Optional[torch.LongTensor] = None,
|
||||
past_key_values: Optional[List[torch.FloatTensor]] = None,
|
||||
inputs_embeds: Optional[torch.FloatTensor] = None,
|
||||
labels: Optional[torch.LongTensor] = None,
|
||||
use_cache: Optional[bool] = None,
|
||||
output_attentions: Optional[bool] = None,
|
||||
output_hidden_states: Optional[bool] = None,
|
||||
return_dict: Optional[bool] = None,
|
||||
) -> Union[Tuple, CausalLMOutputWithPast]:
|
||||
|
||||
output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions
|
||||
output_hidden_states = (
|
||||
output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
|
||||
)
|
||||
return_dict = return_dict if return_dict is not None else self.config.use_return_dict
|
||||
|
||||
# decoder outputs consists of (dec_features, layer_state, dec_hidden, dec_attn)
|
||||
outputs = self.model(
|
||||
input_ids=input_ids,
|
||||
attention_mask=attention_mask,
|
||||
position_ids=position_ids,
|
||||
past_key_values=past_key_values,
|
||||
inputs_embeds=inputs_embeds,
|
||||
use_cache=use_cache,
|
||||
output_attentions=output_attentions,
|
||||
output_hidden_states=output_hidden_states,
|
||||
return_dict=return_dict,
|
||||
)
|
||||
|
||||
hidden_states = outputs[0]
|
||||
logits = self.lm_head(hidden_states)
|
||||
loss = None
|
||||
if labels is not None:
|
||||
# Shift so that tokens < n predict n
|
||||
shift_logits = logits[..., :-1, :].contiguous()
|
||||
shift_labels = labels[..., 1:].contiguous()
|
||||
# Flatten the tokens
|
||||
loss_fct = CrossEntropyLoss()
|
||||
shift_logits = shift_logits.view(-1, self.config.vocab_size)
|
||||
shift_labels = shift_labels.view(-1)
|
||||
softmax_normalizer = shift_logits.max(-1).values ** 2
|
||||
z_loss = self.config.z_loss_weight * softmax_normalizer.mean()
|
||||
# Enable model parallelism
|
||||
shift_labels = shift_labels.to(shift_logits.device)
|
||||
loss = loss_fct(shift_logits, shift_labels) + z_loss
|
||||
|
||||
if not return_dict:
|
||||
output = (logits,) + outputs[1:]
|
||||
return (loss,) + output if loss is not None else output
|
||||
|
||||
return CausalLMOutputWithPast(
|
||||
loss=loss,
|
||||
logits=logits,
|
||||
past_key_values=outputs.past_key_values,
|
||||
hidden_states=outputs.hidden_states,
|
||||
attentions=outputs.attentions,
|
||||
)
|
||||
|
||||
def prepare_inputs_for_generation(
|
||||
self, input_ids, past_key_values=None, attention_mask=None, inputs_embeds=None, **kwargs
|
||||
):
|
||||
if past_key_values:
|
||||
input_ids = input_ids[:, -1:]
|
||||
|
||||
position_ids = kwargs.get("position_ids", None)
|
||||
if attention_mask is not None and position_ids is None:
|
||||
# create position_ids on the fly for batch generation
|
||||
position_ids = attention_mask.long().cumsum(-1) - 1
|
||||
position_ids.masked_fill_(attention_mask == 0, 1)
|
||||
if past_key_values:
|
||||
position_ids = position_ids[:, -1].unsqueeze(-1)
|
||||
|
||||
# if `inputs_embeds` are passed, we only want to use them in the 1st generation step
|
||||
if inputs_embeds is not None and past_key_values is None:
|
||||
model_inputs = {"inputs_embeds": inputs_embeds}
|
||||
else:
|
||||
model_inputs = {"input_ids": input_ids}
|
||||
|
||||
model_inputs.update(
|
||||
{
|
||||
"position_ids": position_ids,
|
||||
"past_key_values": past_key_values,
|
||||
"use_cache": kwargs.get("use_cache"),
|
||||
"attention_mask": attention_mask,
|
||||
}
|
||||
)
|
||||
return model_inputs
|
||||
|
||||
@staticmethod
|
||||
def _reorder_cache(past_key_values, beam_idx):
|
||||
reordered_past = ()
|
||||
for layer_past in past_key_values:
|
||||
reordered_past += (tuple(past_state.index_select(0, beam_idx) for past_state in layer_past),)
|
||||
return reordered_past
|
||||
|
||||
def quantize(self, bits: int):
|
||||
try:
|
||||
from .quantizer import quantize_online
|
||||
except ImportError:
|
||||
raise ImportError(f"Needs QLinear to run quantize.")
|
||||
return quantize_online(self, bits)
|
||||
|
||||
def chat(self, tokenizer, messages: List[dict], stream=False,
|
||||
generation_config: Optional[GenerationConfig]=None):
|
||||
generation_config = generation_config or self.generation_config
|
||||
input_ids = build_chat_input(self, tokenizer, messages, generation_config.max_new_tokens)
|
||||
if stream:
|
||||
streamer = TextIterStreamer(tokenizer, skip_prompt=True, skip_special_tokens=True)
|
||||
Thread(target=self.generate, kwargs=dict(
|
||||
inputs=input_ids, streamer=streamer,
|
||||
generation_config=generation_config,
|
||||
)).start()
|
||||
return streamer
|
||||
else:
|
||||
outputs = self.generate(input_ids, generation_config=generation_config)
|
||||
response = tokenizer.decode(outputs[0][len(input_ids[0]):], skip_special_tokens=True)
|
||||
return response
|
||||
210
ixformer_sdk/train/speedformer/models/baichuan/quantizer.py
Normal file
210
ixformer_sdk/train/speedformer/models/baichuan/quantizer.py
Normal file
@@ -0,0 +1,210 @@
|
||||
import bitsandbytes as bnb
|
||||
from bitsandbytes.nn.modules import Params4bit, Int8Params
|
||||
import torch
|
||||
|
||||
def Params4bitCuda(self, device):
|
||||
self.data = self.data.cuda(device)
|
||||
self.quant_state[0] = self.quant_state[0].cuda(device)
|
||||
self.quant_state[4][0] = self.quant_state[4][0].cuda(device)
|
||||
self.quant_state[4][1][0] = self.quant_state[4][1][0].cuda(device)
|
||||
self.quant_state[4][1][1] = self.quant_state[4][1][1].cuda(device)
|
||||
|
||||
self.quant_state[6] = self.quant_state[6].cuda(device)
|
||||
return self
|
||||
|
||||
class Linear4bitOnline(torch.nn.Module):
|
||||
def __init__(self, weight, bias, quant_type):
|
||||
super().__init__()
|
||||
self.weight = Params4bit(
|
||||
weight.data, requires_grad=False, compress_statistics=True, quant_type=quant_type
|
||||
)
|
||||
self.compute_dtype = None
|
||||
#self.weight.cuda(weight.device)
|
||||
self.bias = bias
|
||||
|
||||
def forward(self, x: torch.Tensor):
|
||||
# weights are cast automatically as Int8Params, but the bias has to be cast manually
|
||||
if self.bias is not None and self.bias.dtype != x.dtype:
|
||||
self.bias.data = self.bias.data.to(x.dtype)
|
||||
|
||||
if getattr(self.weight, "quant_state", None) is None:
|
||||
print(
|
||||
"FP4 quantization state not initialized. Please call .cuda() or .to(device) on the LinearFP4 layer first."
|
||||
)
|
||||
inp_dtype = x.dtype
|
||||
if self.compute_dtype is not None:
|
||||
x = x.to(self.compute_dtype)
|
||||
|
||||
bias = None if self.bias is None else self.bias.to(self.compute_dtype)
|
||||
out = bnb.matmul_4bit(
|
||||
x, self.weight.t(), bias=bias, quant_state=self.weight.quant_state
|
||||
)
|
||||
|
||||
out = out.to(inp_dtype)
|
||||
|
||||
return out
|
||||
|
||||
class Linear8bitLtOnline(torch.nn.Module):
|
||||
def __init__(
|
||||
self,
|
||||
weight,
|
||||
bias,
|
||||
has_fp16_weights=True,
|
||||
memory_efficient_backward=False,
|
||||
threshold=0.0,
|
||||
index=None,
|
||||
):
|
||||
super().__init__()
|
||||
assert (
|
||||
not memory_efficient_backward
|
||||
), "memory_efficient_backward is no longer required and the argument is deprecated in 0.37.0 and will be removed in 0.39.0"
|
||||
self.state = bnb.MatmulLtState()
|
||||
self.index = index
|
||||
|
||||
# Necessary for stacked layers
|
||||
self.state.threshold = threshold
|
||||
self.state.has_fp16_weights = has_fp16_weights
|
||||
self.state.memory_efficient_backward = memory_efficient_backward
|
||||
if threshold > 0.0 and not has_fp16_weights:
|
||||
self.state.use_pool = True
|
||||
|
||||
self.weight = Int8Params(
|
||||
weight.data,
|
||||
has_fp16_weights=has_fp16_weights,
|
||||
requires_grad=has_fp16_weights,
|
||||
)
|
||||
self.bias = bias
|
||||
|
||||
def init_8bit_state(self):
|
||||
self.state.CB = self.weight.CB
|
||||
self.state.SCB = self.weight.SCB
|
||||
self.weight.CB = None
|
||||
self.weight.SCB = None
|
||||
|
||||
def forward(self, x: torch.Tensor):
|
||||
self.state.is_training = self.training
|
||||
if self.weight.CB is not None:
|
||||
self.init_8bit_state()
|
||||
|
||||
# weights are cast automatically as Int8Params, but the bias has to be cast manually
|
||||
if self.bias is not None and self.bias.dtype != x.dtype:
|
||||
self.bias.data = self.bias.data.to(x.dtype)
|
||||
|
||||
out = bnb.matmul(x, self.weight, bias=self.bias, state=self.state)
|
||||
|
||||
if not self.state.has_fp16_weights:
|
||||
if self.state.CB is not None and self.state.CxB is not None:
|
||||
# we converted 8-bit row major to turing/ampere format in the first inference pass
|
||||
# we no longer need the row-major weight
|
||||
del self.state.CB
|
||||
self.weight.data = self.state.CxB
|
||||
return out
|
||||
|
||||
def quantize_offline(model, bits: int):
|
||||
assert (bits == 4), f'bits: {bits} is not supported'
|
||||
|
||||
for i, layer in enumerate(model.model.layers):
|
||||
layer.self_attn.W_pack = bnb.nn.Linear4bit(
|
||||
layer.self_attn.W_pack.weight.shape[1],
|
||||
layer.self_attn.W_pack.weight.shape[0],
|
||||
False,
|
||||
torch.float16,
|
||||
compress_statistics=True,
|
||||
quant_type="nf4",
|
||||
)
|
||||
layer.self_attn.o_proj = bnb.nn.Linear4bit(
|
||||
layer.self_attn.o_proj.weight.shape[1],
|
||||
layer.self_attn.o_proj.weight.shape[0],
|
||||
False,
|
||||
torch.float16,
|
||||
compress_statistics=True,
|
||||
quant_type="nf4",
|
||||
)
|
||||
|
||||
layer.mlp.gate_proj = bnb.nn.Linear4bit(
|
||||
layer.mlp.gate_proj.weight.shape[1],
|
||||
layer.mlp.gate_proj.weight.shape[0],
|
||||
False,
|
||||
torch.float16,
|
||||
compress_statistics=True,
|
||||
quant_type="nf4",
|
||||
)
|
||||
layer.mlp.down_proj = bnb.nn.Linear4bit(
|
||||
layer.mlp.down_proj.weight.shape[1],
|
||||
layer.mlp.down_proj.weight.shape[0],
|
||||
False,
|
||||
torch.float16,
|
||||
compress_statistics=True,
|
||||
quant_type="nf4",
|
||||
)
|
||||
layer.mlp.up_proj = bnb.nn.Linear4bit(
|
||||
layer.mlp.up_proj.weight.shape[1],
|
||||
layer.mlp.up_proj.weight.shape[0],
|
||||
False,
|
||||
torch.float16,
|
||||
compress_statistics=True,
|
||||
quant_type="nf4",
|
||||
)
|
||||
return model
|
||||
|
||||
def quantize_online(model, bits: int):
|
||||
def quant(weight, bias=None):
|
||||
if bits == 8:
|
||||
linear = Linear8bitLtOnline(
|
||||
weight,
|
||||
bias,
|
||||
has_fp16_weights=False,
|
||||
threshold=6.0,
|
||||
)
|
||||
if bias is not None:
|
||||
linear.bias = torch.nn.Parameter(bias)
|
||||
elif bits == 4:
|
||||
linear = Linear4bitOnline(
|
||||
weight,
|
||||
bias,
|
||||
quant_type="nf4", #fp4/nf4
|
||||
)
|
||||
else:
|
||||
raise ValueError("quantize only support 4/8 bit")
|
||||
return linear
|
||||
|
||||
for i, layer in enumerate(model.model.layers):
|
||||
layer.self_attn.W_pack = quant(layer.self_attn.W_pack.weight)
|
||||
layer.self_attn.o_proj = quant(layer.self_attn.o_proj.weight)
|
||||
layer.mlp.gate_proj = quant(layer.mlp.gate_proj.weight)
|
||||
layer.mlp.down_proj = quant(layer.mlp.down_proj.weight)
|
||||
layer.mlp.up_proj = quant(layer.mlp.up_proj.weight)
|
||||
return model
|
||||
|
||||
def init_model_weight_int4(config, model, state_dict):
|
||||
#replace Params4bit.cuda with Params4bitCuda
|
||||
Params4bit.cuda = Params4bitCuda
|
||||
|
||||
for i in range(config.num_hidden_layers):
|
||||
weight_data = state_dict[f'model.layers.{i}.self_attn.W_pack.weight.data']
|
||||
weight_quant_state = state_dict[f'model.layers.{i}.self_attn.W_pack.weight.quant_state']
|
||||
model.model.layers[i].self_attn.W_pack.weight = Params4bit(weight_data, requires_grad=False, quant_state=weight_quant_state)
|
||||
|
||||
weight_data = state_dict[f'model.layers.{i}.self_attn.o_proj.weight.data']
|
||||
weight_quant_state = state_dict[f'model.layers.{i}.self_attn.o_proj.weight.quant_state']
|
||||
model.model.layers[i].self_attn.o_proj.weight = Params4bit(weight_data, requires_grad=False, quant_state=weight_quant_state)
|
||||
|
||||
weight_data = state_dict[f'model.layers.{i}.mlp.gate_proj.weight.data']
|
||||
weight_quant_state = state_dict[f'model.layers.{i}.mlp.gate_proj.weight.quant_state']
|
||||
model.model.layers[i].mlp.gate_proj.weight = Params4bit(weight_data, requires_grad=False, quant_state=weight_quant_state)
|
||||
|
||||
weight_data = state_dict[f'model.layers.{i}.mlp.up_proj.weight.data']
|
||||
weight_quant_state = state_dict[f'model.layers.{i}.mlp.up_proj.weight.quant_state']
|
||||
model.model.layers[i].mlp.up_proj.weight = Params4bit(weight_data, requires_grad=False, quant_state=weight_quant_state)
|
||||
|
||||
weight_data = state_dict[f'model.layers.{i}.mlp.down_proj.weight.data']
|
||||
weight_quant_state = state_dict[f'model.layers.{i}.mlp.down_proj.weight.quant_state']
|
||||
model.model.layers[i].mlp.down_proj.weight = Params4bit(weight_data, requires_grad=False, quant_state=weight_quant_state)
|
||||
|
||||
model.model.layers[i].input_layernorm.weight = state_dict[f'model.layers.{i}.input_layernorm.weight']
|
||||
model.model.layers[i].post_attention_layernorm.weight = state_dict[f'model.layers.{i}.post_attention_layernorm.weight']
|
||||
|
||||
model.model.embed_tokens.weight = state_dict['model.embed_tokens.weight']
|
||||
model.model.norm.weight = state_dict['model.norm.weight']
|
||||
model.lm_head.weight = state_dict['lm_head.weight']
|
||||
return model
|
||||
@@ -0,0 +1,242 @@
|
||||
# coding=utf-8
|
||||
# Copyright 2022 the Big Science Workshop and HuggingFace Inc. team. All rights reserved.
|
||||
#
|
||||
# Licensed under the Apache License, Version 2.0 (the "License");
|
||||
# you may not use this file except in compliance with the License.
|
||||
# You may obtain a copy of the License at
|
||||
#
|
||||
# http://www.apache.org/licenses/LICENSE-2.0
|
||||
#
|
||||
# Unless required by applicable law or agreed to in writing, software
|
||||
# distributed under the License is distributed on an "AS IS" BASIS,
|
||||
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
# See the License for the specific language governing permissions and
|
||||
# limitations under the License.
|
||||
""" Bloom configuration"""
|
||||
from collections import OrderedDict
|
||||
from typing import TYPE_CHECKING, Any, List, Mapping, Optional
|
||||
|
||||
from packaging import version
|
||||
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from ... import PreTrainedTokenizer, TensorType
|
||||
|
||||
from transformers.configuration_utils import PretrainedConfig
|
||||
from transformers.onnx import OnnxConfigWithPast, PatchingSpec
|
||||
from transformers.utils import is_torch_available, logging
|
||||
|
||||
|
||||
logger = logging.get_logger(__name__)
|
||||
|
||||
BLOOM_PRETRAINED_CONFIG_ARCHIVE_MAP = {
|
||||
"bigscience/bloom": "https://huggingface.co/bigscience/bloom/resolve/main/config.json",
|
||||
"bigscience/bloom-560m": "https://huggingface.co/bigscience/bloom-560m/blob/main/config.json",
|
||||
"bigscience/bloom-1b1": "https://huggingface.co/bigscience/bloom-1b1/blob/main/config.json",
|
||||
"bigscience/bloom-1b7": "https://huggingface.co/bigscience/bloom-1b7/blob/main/config.json",
|
||||
"bigscience/bloom-3b": "https://huggingface.co/bigscience/bloom-3b/blob/main/config.json",
|
||||
"bigscience/bloom-7b1": "https://huggingface.co/bigscience/bloom-7b1/blob/main/config.json",
|
||||
}
|
||||
|
||||
|
||||
class BloomConfig(PretrainedConfig):
|
||||
"""
|
||||
This is the configuration class to store the configuration of a [`BloomModel`]. It is used to instantiate a Bloom
|
||||
model according to the specified arguments, defining the model architecture. Instantiating a configuration with the
|
||||
defaults will yield a similar configuration to the Bloom architecture
|
||||
[bigscience/bloom](https://huggingface.co/bigscience/bloom).
|
||||
|
||||
Configuration objects inherit from [`PretrainedConfig`] and can be used to control the model outputs. Read the
|
||||
documentation from [`PretrainedConfig`] for more information.
|
||||
|
||||
|
||||
Args:
|
||||
vocab_size (`int`, *optional*, defaults to 250880):
|
||||
Vocabulary size of the Bloom model. Defines the maximum number of different tokens that can be represented
|
||||
by the `inputs_ids` passed when calling [`BloomModel`]. Check [this
|
||||
discussion](https://huggingface.co/bigscience/bloom/discussions/120#633d28389addb8530b406c2a) on how the
|
||||
`vocab_size` has been defined.
|
||||
hidden_size (`int`, *optional*, defaults to 64):
|
||||
Dimensionality of the embeddings and hidden states.
|
||||
n_layer (`int`, *optional*, defaults to 2):
|
||||
Number of hidden layers in the Transformer encoder.
|
||||
n_head (`int`, *optional*, defaults to 8):
|
||||
Number of attention heads for each attention layer in the Transformer encoder.
|
||||
layer_norm_epsilon (`float`, *optional*, defaults to 1e-5):
|
||||
The epsilon to use in the layer normalization layers.
|
||||
initializer_range (`float`, *optional*, defaults to 0.02):
|
||||
The standard deviation of the truncated_normal_initializer for initializing all weight matrices.
|
||||
apply_residual_connection_post_layernorm (`bool`, *optional*, defaults to `False`):
|
||||
If enabled, use the layer norm of the hidden states as the residual in the transformer blocks
|
||||
hidden_dropout (`float`, *optional*, defaults to 0.1):
|
||||
Dropout rate of the dropout function on the bias dropout.
|
||||
attention_dropout (`float`, *optional*, defaults to 0.1):
|
||||
Dropout rate applied to the attention probs
|
||||
use_cache (`bool`, *optional*, defaults to `True`):
|
||||
Whether or not the model should return the last key/values attentions (not used by all models).
|
||||
pretraining_tp (`int`, *optional*, defaults to `1`):
|
||||
Experimental feature. Tensor parallelism rank used during pretraining with Megatron. Please refer to [this
|
||||
document](https://huggingface.co/docs/transformers/parallelism) to understand more about it. This value is
|
||||
necessary to ensure exact reproducibility of the pretraining results. Please refer to [this
|
||||
issue](https://github.com/pytorch/pytorch/issues/76232). Note also that this is enabled only when
|
||||
`slow_but_exact=True`.
|
||||
slow_but_exact (`bool`, *optional*, defaults to `False`):
|
||||
Experimental feature. Whether to use slow but exact implementation of the attention mechanism. While
|
||||
merging the TP rank tensors, due to slicing operations the results may be slightly different between the
|
||||
model trained on Megatron and our model. Please refer to [this
|
||||
issue](https://github.com/pytorch/pytorch/issues/76232). A solution to obtain more accurate results is to
|
||||
enable this feature. Enabling this will hurt the computational time of the inference. Will be probably
|
||||
resolved in the future once the main model has been fine-tuned with TP_rank=1.
|
||||
|
||||
Example:
|
||||
|
||||
```python
|
||||
>>> from transformers import BloomConfig, BloomModel
|
||||
|
||||
>>> # Initializing a Bloom configuration
|
||||
>>> configuration = BloomConfig()
|
||||
|
||||
>>> # Initializing a model (with random weights) from the configuration
|
||||
>>> model = BloomModel(configuration)
|
||||
|
||||
>>> # Accessing the model configuration
|
||||
>>> configuration = model.config
|
||||
```"""
|
||||
|
||||
model_type = "bloom"
|
||||
keys_to_ignore_at_inference = ["past_key_values"]
|
||||
attribute_map = {
|
||||
"num_hidden_layers": "n_layer",
|
||||
"num_attention_heads": "n_head",
|
||||
}
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
vocab_size=250880,
|
||||
hidden_size=64,
|
||||
n_layer=2,
|
||||
n_head=8,
|
||||
layer_norm_epsilon=1e-5,
|
||||
initializer_range=0.02,
|
||||
use_cache=True,
|
||||
bos_token_id=1,
|
||||
eos_token_id=2,
|
||||
apply_residual_connection_post_layernorm=False,
|
||||
hidden_dropout=0.0,
|
||||
attention_dropout=0.0,
|
||||
pretraining_tp=1, # TP rank used when training with megatron
|
||||
slow_but_exact=False,
|
||||
**kwargs,
|
||||
):
|
||||
self.vocab_size = vocab_size
|
||||
# Backward compatibility with n_embed kwarg
|
||||
n_embed = kwargs.pop("n_embed", None)
|
||||
self.hidden_size = hidden_size if n_embed is None else n_embed
|
||||
self.n_layer = n_layer
|
||||
self.n_head = n_head
|
||||
self.layer_norm_epsilon = layer_norm_epsilon
|
||||
self.initializer_range = initializer_range
|
||||
self.use_cache = use_cache
|
||||
self.pretraining_tp = pretraining_tp
|
||||
self.apply_residual_connection_post_layernorm = apply_residual_connection_post_layernorm
|
||||
self.hidden_dropout = hidden_dropout
|
||||
self.attention_dropout = attention_dropout
|
||||
|
||||
self.bos_token_id = bos_token_id
|
||||
self.eos_token_id = eos_token_id
|
||||
self.slow_but_exact = slow_but_exact
|
||||
|
||||
super().__init__(bos_token_id=bos_token_id, eos_token_id=eos_token_id, **kwargs)
|
||||
|
||||
|
||||
class BloomOnnxConfig(OnnxConfigWithPast):
|
||||
torch_onnx_minimum_version = version.parse("1.12")
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
config: PretrainedConfig,
|
||||
task: str = "default",
|
||||
patching_specs: List[PatchingSpec] = None,
|
||||
use_past: bool = False,
|
||||
):
|
||||
super().__init__(config, task=task, patching_specs=patching_specs, use_past=use_past)
|
||||
if not getattr(self._config, "pad_token_id", None):
|
||||
# TODO: how to do that better?
|
||||
self._config.pad_token_id = 0
|
||||
|
||||
@property
|
||||
def inputs(self) -> Mapping[str, Mapping[int, str]]:
|
||||
common_inputs = OrderedDict({"input_ids": {0: "batch", 1: "sequence"}})
|
||||
if self.use_past:
|
||||
# BLOOM stores values on dynamic axis 2. For more details see: https://github.com/huggingface/transformers/pull/18344
|
||||
self.fill_with_past_key_values_(common_inputs, direction="inputs", inverted_values_shape=True)
|
||||
common_inputs["attention_mask"] = {0: "batch", 1: "past_sequence + sequence"}
|
||||
else:
|
||||
common_inputs["attention_mask"] = {0: "batch", 1: "sequence"}
|
||||
|
||||
return common_inputs
|
||||
|
||||
@property
|
||||
def num_layers(self) -> int:
|
||||
return self._config.n_layer
|
||||
|
||||
@property
|
||||
def num_attention_heads(self) -> int:
|
||||
return self._config.n_head
|
||||
|
||||
@property
|
||||
def atol_for_validation(self) -> float:
|
||||
return 1e-3
|
||||
|
||||
def generate_dummy_inputs(
|
||||
self,
|
||||
tokenizer: "PreTrainedTokenizer",
|
||||
batch_size: int = -1,
|
||||
seq_length: int = -1,
|
||||
is_pair: bool = False,
|
||||
framework: Optional["TensorType"] = None,
|
||||
) -> Mapping[str, Any]:
|
||||
common_inputs = super(OnnxConfigWithPast, self).generate_dummy_inputs(
|
||||
tokenizer, batch_size=batch_size, seq_length=seq_length, is_pair=is_pair, framework=framework
|
||||
)
|
||||
|
||||
# We need to order the input in the way they appears in the forward()
|
||||
ordered_inputs = OrderedDict({"input_ids": common_inputs["input_ids"]})
|
||||
|
||||
# Need to add the past_keys
|
||||
if self.use_past:
|
||||
if not is_torch_available():
|
||||
raise ValueError("Cannot generate dummy past_keys inputs without PyTorch installed.")
|
||||
else:
|
||||
import torch
|
||||
|
||||
batch, seqlen = common_inputs["input_ids"].shape
|
||||
# Not using the same length for past_key_values
|
||||
past_key_values_length = seqlen + 2
|
||||
head_dim = self._config.hidden_size // self.num_attention_heads
|
||||
past_key_shape = (
|
||||
batch * self.num_attention_heads,
|
||||
head_dim,
|
||||
past_key_values_length,
|
||||
)
|
||||
past_value_shape = (
|
||||
batch * self.num_attention_heads,
|
||||
past_key_values_length,
|
||||
head_dim,
|
||||
)
|
||||
ordered_inputs["past_key_values"] = [
|
||||
(torch.zeros(past_key_shape), torch.zeros(past_value_shape)) for _ in range(self.num_layers)
|
||||
]
|
||||
|
||||
ordered_inputs["attention_mask"] = common_inputs["attention_mask"]
|
||||
if self.use_past:
|
||||
mask_dtype = ordered_inputs["attention_mask"].dtype
|
||||
ordered_inputs["attention_mask"] = torch.cat(
|
||||
[ordered_inputs["attention_mask"], torch.ones(batch, past_key_values_length, dtype=mask_dtype)], dim=1
|
||||
)
|
||||
|
||||
return ordered_inputs
|
||||
|
||||
@property
|
||||
def default_onnx_opset(self) -> int:
|
||||
return 13
|
||||
@@ -0,0 +1,500 @@
|
||||
# Copyright 2023 The HuggingFace Team. All rights reserved.
|
||||
#
|
||||
# Licensed under the Apache License, Version 2.0 (the "License");
|
||||
# you may not use this file except in compliance with the License.
|
||||
# You may obtain a copy of the License at
|
||||
#
|
||||
# http://www.apache.org/licenses/LICENSE-2.0
|
||||
#
|
||||
# Unless required by applicable law or agreed to in writing, software
|
||||
# distributed under the License is distributed on an "AS IS" BASIS,
|
||||
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
# See the License for the specific language governing permissions and
|
||||
# limitations under the License.
|
||||
from dataclasses import dataclass
|
||||
from typing import List, Optional, Tuple, Union
|
||||
|
||||
import torch
|
||||
|
||||
|
||||
@dataclass
|
||||
class AttentionMaskConverter:
|
||||
"""
|
||||
A utility attention mask class that allows one to:
|
||||
- Create a causal 4d mask
|
||||
- Create a causal 4d mask with slided window
|
||||
- Convert a 2d attention mask (batch_size, query_length) to a 4d attention mask (batch_size, 1, query_length,
|
||||
key_value_length) that can be multiplied with attention scores
|
||||
|
||||
Examples:
|
||||
|
||||
```python
|
||||
>>> import torch
|
||||
>>> from transformers.modeling_attn_mask_utils import AttentionMaskConverter
|
||||
|
||||
>>> converter = AttentionMaskConverter(True)
|
||||
>>> converter.to_4d(torch.tensor([[0, 0, 0, 1, 1]]), 5, key_value_length=5, dtype=torch.float32)
|
||||
tensor([[[[-3.4028e+38, -3.4028e+38, -3.4028e+38, -3.4028e+38, -3.4028e+38],
|
||||
[-3.4028e+38, -3.4028e+38, -3.4028e+38, -3.4028e+38, -3.4028e+38],
|
||||
[-3.4028e+38, -3.4028e+38, -3.4028e+38, -3.4028e+38, -3.4028e+38],
|
||||
[-3.4028e+38, -3.4028e+38, -3.4028e+38, 0.0000e+00, -3.4028e+38],
|
||||
[-3.4028e+38, -3.4028e+38, -3.4028e+38, 0.0000e+00, 0.0000e+00]]]])
|
||||
```
|
||||
|
||||
Parameters:
|
||||
is_causal (`bool`):
|
||||
Whether the attention mask should be a uni-directional (causal) or bi-directional mask.
|
||||
|
||||
sliding_window (`int`, *optional*):
|
||||
Optionally, the sliding window masks can be created if `sliding_window` is defined to a positive integer.
|
||||
"""
|
||||
|
||||
is_causal: bool
|
||||
sliding_window: int
|
||||
|
||||
def __init__(self, is_causal: bool, sliding_window: Optional[int] = None):
|
||||
self.is_causal = is_causal
|
||||
self.sliding_window = sliding_window
|
||||
|
||||
if self.sliding_window is not None and self.sliding_window <= 0:
|
||||
raise ValueError(
|
||||
f"Make sure that when passing `sliding_window` that its value is a strictly positive integer, not `{self.sliding_window}`"
|
||||
)
|
||||
|
||||
def to_causal_4d(
|
||||
self,
|
||||
batch_size: int,
|
||||
query_length: int,
|
||||
key_value_length: int,
|
||||
dtype: torch.dtype,
|
||||
device: Union[torch.device, "str"] = "cpu",
|
||||
) -> Optional[torch.Tensor]:
|
||||
"""
|
||||
Creates a causal 4D mask of (bsz, head_dim=1, query_length, key_value_length) shape and adds large negative
|
||||
bias to upper right hand triangular matrix (causal mask).
|
||||
"""
|
||||
if not self.is_causal:
|
||||
raise ValueError(f"Please use `to_causal_4d` only if {self.__class__} has `is_causal` set to True.")
|
||||
|
||||
# If shape is not cached, create a new causal mask and cache it
|
||||
input_shape = (batch_size, query_length)
|
||||
past_key_values_length = key_value_length - query_length
|
||||
|
||||
# create causal mask
|
||||
# [bsz, seq_len] -> [bsz, 1, tgt_seq_len, src_seq_len]
|
||||
causal_4d_mask = None
|
||||
if input_shape[-1] > 1 or self.sliding_window is not None:
|
||||
causal_4d_mask = self._make_causal_mask(
|
||||
input_shape,
|
||||
dtype,
|
||||
device=device,
|
||||
past_key_values_length=past_key_values_length,
|
||||
sliding_window=self.sliding_window,
|
||||
)
|
||||
|
||||
return causal_4d_mask
|
||||
|
||||
def to_4d(
|
||||
self,
|
||||
attention_mask_2d: torch.Tensor,
|
||||
query_length: int,
|
||||
dtype: torch.dtype,
|
||||
key_value_length: Optional[int] = None,
|
||||
) -> torch.Tensor:
|
||||
"""
|
||||
Converts 2D attention mask to 4D attention mask by expanding mask to (bsz, head_dim=1, query_length,
|
||||
key_value_length) shape and by adding a large negative bias to not-attended positions. If attention_mask is
|
||||
causal, a causal mask will be added.
|
||||
"""
|
||||
input_shape = (attention_mask_2d.shape[0], query_length)
|
||||
|
||||
# create causal mask
|
||||
# [bsz, seq_len] -> [bsz, 1, tgt_seq_len, src_seq_len]
|
||||
causal_4d_mask = None
|
||||
if (input_shape[-1] > 1 or self.sliding_window is not None) and self.is_causal:
|
||||
if key_value_length is None:
|
||||
raise ValueError(
|
||||
"This attention mask converter is causal. Make sure to pass `key_value_length` to correctly create a causal mask."
|
||||
)
|
||||
|
||||
past_key_values_length = key_value_length - query_length
|
||||
causal_4d_mask = self._make_causal_mask(
|
||||
input_shape,
|
||||
dtype,
|
||||
device=attention_mask_2d.device,
|
||||
past_key_values_length=past_key_values_length,
|
||||
sliding_window=self.sliding_window,
|
||||
)
|
||||
elif self.sliding_window is not None:
|
||||
raise NotImplementedError("Sliding window is currently only implemented for causal masking")
|
||||
|
||||
# [bsz, seq_len] -> [bsz, 1, tgt_seq_len, src_seq_len]
|
||||
expanded_attn_mask = self._expand_mask(attention_mask_2d, dtype, tgt_len=input_shape[-1]).to(
|
||||
attention_mask_2d.device
|
||||
)
|
||||
|
||||
if causal_4d_mask is not None:
|
||||
expanded_attn_mask = causal_4d_mask.masked_fill(expanded_attn_mask.bool(), torch.finfo(dtype).min)
|
||||
|
||||
# expanded_attn_mask + causal_4d_mask can cause some overflow
|
||||
expanded_4d_mask = expanded_attn_mask
|
||||
|
||||
return expanded_4d_mask
|
||||
|
||||
@staticmethod
|
||||
def _make_causal_mask(
|
||||
input_ids_shape: torch.Size,
|
||||
dtype: torch.dtype,
|
||||
device: torch.device,
|
||||
past_key_values_length: int = 0,
|
||||
sliding_window: Optional[int] = None,
|
||||
):
|
||||
"""
|
||||
Make causal mask used for bi-directional self-attention.
|
||||
"""
|
||||
bsz, tgt_len = input_ids_shape
|
||||
mask = torch.full((tgt_len, tgt_len), torch.finfo(dtype).min, device=device)
|
||||
mask_cond = torch.arange(mask.size(-1), device=device)
|
||||
mask.masked_fill_(mask_cond < (mask_cond + 1).view(mask.size(-1), 1), 0)
|
||||
|
||||
mask = mask.to(dtype)
|
||||
|
||||
if past_key_values_length > 0:
|
||||
mask = torch.cat([torch.zeros(tgt_len, past_key_values_length, dtype=dtype, device=device), mask], dim=-1)
|
||||
|
||||
# add lower triangular sliding window mask if necessary
|
||||
if sliding_window is not None:
|
||||
diagonal = past_key_values_length - sliding_window + 1
|
||||
|
||||
context_mask = 1 - torch.triu(torch.ones_like(mask, dtype=torch.int), diagonal=diagonal)
|
||||
mask.masked_fill_(context_mask.bool(), torch.finfo(dtype).min)
|
||||
|
||||
return mask[None, None, :, :].expand(bsz, 1, tgt_len, tgt_len + past_key_values_length)
|
||||
|
||||
@staticmethod
|
||||
def _expand_mask(mask: torch.Tensor, dtype: torch.dtype, tgt_len: Optional[int] = None):
|
||||
"""
|
||||
Expands attention_mask from `[bsz, seq_len]` to `[bsz, 1, tgt_seq_len, src_seq_len]`.
|
||||
"""
|
||||
bsz, src_len = mask.size()
|
||||
tgt_len = tgt_len if tgt_len is not None else src_len
|
||||
|
||||
expanded_mask = mask[:, None, None, :].expand(bsz, 1, tgt_len, src_len).to(dtype)
|
||||
|
||||
inverted_mask = 1.0 - expanded_mask
|
||||
|
||||
return inverted_mask.masked_fill(inverted_mask.to(torch.bool), torch.finfo(dtype).min)
|
||||
|
||||
@staticmethod
|
||||
def _unmask_unattended(
|
||||
expanded_mask: torch.Tensor, attention_mask: torch.Tensor, unmasked_value: Union[bool, float]
|
||||
):
|
||||
# fmt: off
|
||||
"""
|
||||
Attend to all tokens in masked rows from the expanded attention mask, for example the relevant first rows when
|
||||
using left padding. This is required by F.scaled_dot_product_attention memory-efficient attention path.
|
||||
Details: https://github.com/pytorch/pytorch/issues/110213
|
||||
|
||||
`expanded_mask` is [bsz, num_masks, tgt_seq_len, src_seq_len] or [bsz, tgt_seq_len, src_seq_len].
|
||||
`attention_mask` is [bsz, src_seq_len].
|
||||
|
||||
The dimension num_masks of `expanded_mask` is most often 1, but it can also be the number of heads in the case of alibi attention bias.
|
||||
|
||||
For example, if `attention_mask` is
|
||||
```
|
||||
[[0, 0, 1],
|
||||
[1, 1, 1],
|
||||
[0, 1, 1]]
|
||||
```
|
||||
and `expanded_mask` is (e.g. here left-padding case)
|
||||
```
|
||||
[[[[0, 0, 0],
|
||||
[0, 0, 0],
|
||||
[0, 0, 1]]],
|
||||
[[[1, 0, 0],
|
||||
[1, 1, 0],
|
||||
[1, 1, 1]]],
|
||||
[[[0, 0, 0],
|
||||
[0, 1, 0],
|
||||
[0, 1, 1]]]]
|
||||
```
|
||||
then the modified `expanded_mask` will be
|
||||
```
|
||||
[[[[1, 1, 1], <-- modified
|
||||
[1, 1, 1], <-- modified
|
||||
[0, 0, 1]]],
|
||||
[[[1, 0, 0],
|
||||
[1, 1, 0],
|
||||
[1, 1, 1]]],
|
||||
[[[1, 1, 1], <-- modified
|
||||
[0, 1, 0],
|
||||
[0, 1, 1]]]]
|
||||
```
|
||||
"""
|
||||
# fmt: on
|
||||
|
||||
# Get the index of the first non-zero value for every sample in the batch.
|
||||
# In the above example, indices = [[2], [0], [1]]]
|
||||
tmp = torch.arange(attention_mask.shape[1], 0, -1)
|
||||
indices = torch.argmax(attention_mask.cpu() * tmp, 1, keepdim=True)
|
||||
|
||||
# Find the batch indexes that have unattended tokens on the leftmost side (e.g. [0, 0, 1, 1, 1]), for which the first rows of the
|
||||
# expanded mask will be completely unattended.
|
||||
left_masked_rows = torch.where(indices > 0)[0]
|
||||
|
||||
if left_masked_rows.shape[0] == 0:
|
||||
return expanded_mask
|
||||
indices = indices[left_masked_rows]
|
||||
|
||||
max_len = torch.max(indices)
|
||||
range_tensor = torch.arange(max_len).unsqueeze(0)
|
||||
range_tensor = range_tensor.repeat(indices.size(0), 1)
|
||||
|
||||
# Avoid unmasking tokens at relevant target positions (on the row axis), by rather unmasking possibly several times the first row that should always be unmasked as we filtered out the batch above.
|
||||
range_tensor[range_tensor >= indices] = 0
|
||||
|
||||
# TODO: we may drop support for 3D attention mask as the refactor from Patrick maybe dropped this case
|
||||
if expanded_mask.dim() == 4:
|
||||
num_masks = expanded_mask.shape[1]
|
||||
if num_masks == 1:
|
||||
# Broadcast [left_masked_rows, 1], [left_masked_rows, max_len]
|
||||
mask_slice = (left_masked_rows[:, None], 0, range_tensor)
|
||||
else:
|
||||
# Broadcast [left_masked_rows, 1, 1], [1, num_masks, 1], [left_masked_rows, 1, max_len]
|
||||
mask_slice = (
|
||||
left_masked_rows[:, None, None],
|
||||
torch.arange(num_masks)[None, :, None],
|
||||
range_tensor[:, None, :],
|
||||
)
|
||||
else:
|
||||
# Broadcast [left_masked_rows, 1], [left_masked_rows, max_len]
|
||||
mask_slice = (left_masked_rows[:, None], range_tensor)
|
||||
|
||||
expanded_mask[mask_slice] = unmasked_value
|
||||
|
||||
return expanded_mask
|
||||
|
||||
|
||||
def _prepare_4d_causal_attention_mask(
|
||||
attention_mask: Optional[torch.Tensor],
|
||||
input_shape: Union[torch.Size, Tuple, List],
|
||||
inputs_embeds: torch.Tensor,
|
||||
past_key_values_length: int,
|
||||
sliding_window: Optional[int] = None,
|
||||
):
|
||||
"""
|
||||
Creates a causal 4D mask of shape `(batch_size, 1, query_length, key_value_length)` from a 2D mask of shape
|
||||
`(batch_size, key_value_length)`
|
||||
|
||||
Args:
|
||||
attention_mask (`torch.Tensor` or `None`):
|
||||
A 2D attention mask of shape `(batch_size, key_value_length)`
|
||||
input_shape (`tuple(int)` or `list(int)` or `torch.Size`):
|
||||
The input shape should be a tuple that defines `(batch_size, query_length)`.
|
||||
inputs_embeds (`torch.Tensor`):
|
||||
The embedded inputs as a torch Tensor.
|
||||
past_key_values_length (`int`):
|
||||
The length of the key value cache.
|
||||
sliding_window (`int`, *optional*):
|
||||
If the model uses windowed attention, a sliding window should be passed.
|
||||
"""
|
||||
attn_mask_converter = AttentionMaskConverter(is_causal=True, sliding_window=sliding_window)
|
||||
|
||||
key_value_length = input_shape[-1] + past_key_values_length
|
||||
|
||||
# 4d mask is passed through the layers
|
||||
if attention_mask is not None and len(attention_mask.shape) == 2:
|
||||
attention_mask = attn_mask_converter.to_4d(
|
||||
attention_mask, input_shape[-1], key_value_length=key_value_length, dtype=inputs_embeds.dtype
|
||||
)
|
||||
elif attention_mask is not None and len(attention_mask.shape) == 4:
|
||||
expected_shape = (input_shape[0], 1, input_shape[1], key_value_length)
|
||||
if tuple(attention_mask.shape) != expected_shape:
|
||||
raise ValueError(
|
||||
f"Incorrect 4D attention_mask shape: {tuple(attention_mask.shape)}; expected: {expected_shape}."
|
||||
)
|
||||
else:
|
||||
# if the 4D mask has correct shape - invert it and fill with negative infinity
|
||||
inverted_mask = 1.0 - attention_mask
|
||||
attention_mask = inverted_mask.masked_fill(
|
||||
inverted_mask.to(torch.bool), torch.finfo(inputs_embeds.dtype).min
|
||||
)
|
||||
else:
|
||||
attention_mask = attn_mask_converter.to_causal_4d(
|
||||
input_shape[0], input_shape[-1], key_value_length, dtype=inputs_embeds.dtype, device=inputs_embeds.device
|
||||
)
|
||||
|
||||
return attention_mask
|
||||
|
||||
|
||||
# Adapted from _prepare_4d_causal_attention_mask
|
||||
def _prepare_4d_causal_attention_mask_for_sdpa(
|
||||
attention_mask: Optional[torch.Tensor],
|
||||
input_shape: Union[torch.Size, Tuple, List],
|
||||
inputs_embeds: torch.Tensor,
|
||||
past_key_values_length: int,
|
||||
sliding_window: Optional[int] = None,
|
||||
):
|
||||
"""
|
||||
Prepares the correct `attn_mask` argument to be used by `torch.nn.functional.scaled_dot_product_attention`.
|
||||
|
||||
In case no token is masked in the `attention_mask` argument, we simply set it to `None` for the cases `query_length == 1` and
|
||||
`key_value_length == query_length`, and rely instead on SDPA `is_causal` argument to use causal/non-causal masks,
|
||||
allowing to dispatch to the flash attention kernel (that can otherwise not be used if a custom `attn_mask` is passed).
|
||||
"""
|
||||
attn_mask_converter = AttentionMaskConverter(is_causal=True, sliding_window=sliding_window)
|
||||
|
||||
key_value_length = input_shape[-1] + past_key_values_length
|
||||
batch_size, query_length = input_shape
|
||||
|
||||
# torch.jit.trace, symbolic_trace and torchdynamo with fullgraph=True are unable to capture the controlflow `is_causal=attention_mask is None and q_len > 1`
|
||||
# used as an SDPA argument. We keep compatibility with these tracing tools by always using SDPA's `attn_mask` argument in case we are tracing.
|
||||
# TODO: Fix this as well when using torchdynamo with fullgraph=True.
|
||||
is_tracing = torch.jit.is_tracing() or isinstance(inputs_embeds, torch.fx.Proxy)
|
||||
|
||||
if attention_mask is not None:
|
||||
# 4d mask is passed through
|
||||
if len(attention_mask.shape) == 4:
|
||||
expected_shape = (input_shape[0], 1, input_shape[1], key_value_length)
|
||||
if tuple(attention_mask.shape) != expected_shape:
|
||||
raise ValueError(
|
||||
f"Incorrect 4D attention_mask shape: {tuple(attention_mask.shape)}; expected: {expected_shape}."
|
||||
)
|
||||
else:
|
||||
# if the 4D mask has correct shape - invert it and fill with negative infinity
|
||||
inverted_mask = 1.0 - attention_mask.to(inputs_embeds.dtype)
|
||||
attention_mask = inverted_mask.masked_fill(
|
||||
inverted_mask.to(torch.bool), torch.finfo(inputs_embeds.dtype).min
|
||||
)
|
||||
return attention_mask
|
||||
|
||||
elif not is_tracing and torch.all(attention_mask == 1):
|
||||
if query_length == 1:
|
||||
# For query_length == 1, causal attention and bi-directional attention are the same.
|
||||
attention_mask = None
|
||||
elif key_value_length == query_length:
|
||||
attention_mask = None
|
||||
else:
|
||||
# Unfortunately, for query_length > 1 and key_value_length != query_length, we cannot generally ignore the attention mask, as SDPA causal mask generation
|
||||
# may be wrong. We will set `is_causal=False` in SDPA and rely on Transformers attention_mask instead, hence not setting it to None here.
|
||||
# Reference: https://github.com/pytorch/pytorch/issues/108108
|
||||
pass
|
||||
elif query_length > 1 and key_value_length != query_length:
|
||||
# See the comment above (https://github.com/pytorch/pytorch/issues/108108).
|
||||
# Ugly: we set it to True here to dispatch in the following controlflow to `to_causal_4d`.
|
||||
attention_mask = True
|
||||
elif is_tracing:
|
||||
raise ValueError(
|
||||
'Attention using SDPA can not be traced with torch.jit.trace when no attention_mask is provided. To solve this issue, please either load your model with the argument `attn_implementation="eager"` or pass an attention_mask input when tracing the model.'
|
||||
)
|
||||
|
||||
if attention_mask is None:
|
||||
expanded_4d_mask = None
|
||||
elif attention_mask is True:
|
||||
expanded_4d_mask = attn_mask_converter.to_causal_4d(
|
||||
input_shape[0], input_shape[-1], key_value_length, dtype=inputs_embeds.dtype, device=inputs_embeds.device
|
||||
)
|
||||
else:
|
||||
expanded_4d_mask = attn_mask_converter.to_4d(
|
||||
attention_mask,
|
||||
input_shape[-1],
|
||||
dtype=inputs_embeds.dtype,
|
||||
key_value_length=key_value_length,
|
||||
)
|
||||
|
||||
# From PyTorch 2.1 onwards, F.scaled_dot_product_attention with the memory-efficient attention backend
|
||||
# produces nans if sequences are completely unattended in the attention mask. Details: https://github.com/pytorch/pytorch/issues/110213
|
||||
#
|
||||
# This fix is not applied in case we are tracing with torch.jit.trace or symbolic_trace, as _unmask_unattended has a data-dependent
|
||||
# controlflow that can not be captured properly.
|
||||
# TODO: _unmask_unattended does not work either with torch.compile when using fullgraph=True. We should find a way to detect this case.
|
||||
if query_length > 1 and not is_tracing:
|
||||
expanded_4d_mask = AttentionMaskConverter._unmask_unattended(
|
||||
expanded_4d_mask, attention_mask, unmasked_value=0.0
|
||||
)
|
||||
|
||||
return expanded_4d_mask
|
||||
|
||||
|
||||
def _prepare_4d_attention_mask(mask: torch.Tensor, dtype: torch.dtype, tgt_len: Optional[int] = None):
|
||||
"""
|
||||
Creates a non-causal 4D mask of shape `(batch_size, 1, query_length, key_value_length)` from a 2D mask of shape
|
||||
`(batch_size, key_value_length)`
|
||||
|
||||
Args:
|
||||
mask (`torch.Tensor` or `None`):
|
||||
A 2D attention mask of shape `(batch_size, key_value_length)`
|
||||
dtype (`torch.dtype`):
|
||||
The torch dtype the created mask shall have.
|
||||
tgt_len (`int`):
|
||||
The target length or query length the created mask shall have.
|
||||
"""
|
||||
return AttentionMaskConverter._expand_mask(mask=mask, dtype=dtype, tgt_len=tgt_len)
|
||||
|
||||
|
||||
def _prepare_4d_attention_mask_for_sdpa(mask: torch.Tensor, dtype: torch.dtype, tgt_len: Optional[int] = None):
|
||||
"""
|
||||
Creates a non-causal 4D mask of shape `(batch_size, 1, query_length, key_value_length)` from a 2D mask of shape
|
||||
`(batch_size, key_value_length)`
|
||||
|
||||
Args:
|
||||
mask (`torch.Tensor` or `None`):
|
||||
A 2D attention mask of shape `(batch_size, key_value_length)`
|
||||
dtype (`torch.dtype`):
|
||||
The torch dtype the created mask shall have.
|
||||
tgt_len (`int`):
|
||||
The target length or query length the created mask shall have.
|
||||
"""
|
||||
batch_size, key_value_length = mask.shape
|
||||
tgt_len = tgt_len if tgt_len is not None else key_value_length
|
||||
|
||||
# torch.jit.trace and torchdynamo with fullgraph=True are unable to capture the controlflow `is_causal=attention_mask is None and q_len > 1`
|
||||
# used as an SDPA argument. We keep compatibility with these tracing tools by always using SDPA's `attn_mask` argument in case we are tracing.
|
||||
# TODO: Fix this as well when using torchdynamo with fullgraph=True.
|
||||
is_tracing = torch.jit.is_tracing()
|
||||
|
||||
if torch.all(mask == 1):
|
||||
if is_tracing:
|
||||
pass
|
||||
elif tgt_len == 1:
|
||||
# For query_length == 1, causal attention and bi-directional attention are the same.
|
||||
return None
|
||||
elif key_value_length == tgt_len:
|
||||
return None
|
||||
else:
|
||||
# Unfortunately, for query_length > 1 and key_value_length != query_length, we can not generally ignore the attention mask, as SDPA causal mask generation
|
||||
# may be wrong. We will set is_causal=False in SDPA and rely on Transformers attention_mask instead, hence not setting it to None here.
|
||||
# Reference: https://github.com/pytorch/pytorch/issues/108108
|
||||
return AttentionMaskConverter._expand_mask(mask=mask, dtype=dtype, tgt_len=tgt_len)
|
||||
else:
|
||||
return AttentionMaskConverter._expand_mask(mask=mask, dtype=dtype, tgt_len=tgt_len)
|
||||
|
||||
|
||||
def _create_4d_causal_attention_mask(
|
||||
input_shape: Union[torch.Size, Tuple, List],
|
||||
dtype: torch.dtype,
|
||||
device: torch.device,
|
||||
past_key_values_length: int = 0,
|
||||
sliding_window: Optional[int] = None,
|
||||
) -> Optional[torch.Tensor]:
|
||||
"""
|
||||
Creates a causal 4D mask of shape `(batch_size, 1, query_length, key_value_length)`
|
||||
|
||||
Args:
|
||||
input_shape (`tuple(int)` or `list(int)` or `torch.Size`):
|
||||
The input shape should be a tuple that defines `(batch_size, query_length)`.
|
||||
dtype (`torch.dtype`):
|
||||
The torch dtype the created mask shall have.
|
||||
device (`int`):
|
||||
The torch device the created mask shall have.
|
||||
sliding_window (`int`, *optional*):
|
||||
If the model uses windowed attention, a sliding window should be passed.
|
||||
"""
|
||||
attn_mask_converter = AttentionMaskConverter(is_causal=True, sliding_window=sliding_window)
|
||||
|
||||
key_value_length = past_key_values_length + input_shape[-1]
|
||||
attention_mask = attn_mask_converter.to_causal_4d(
|
||||
input_shape[0], input_shape[-1], key_value_length, dtype=dtype, device=device
|
||||
)
|
||||
|
||||
return attention_mask
|
||||
1250
ixformer_sdk/train/speedformer/models/bloom/modeling_bloom.py
Normal file
1250
ixformer_sdk/train/speedformer/models/bloom/modeling_bloom.py
Normal file
File diff suppressed because it is too large
Load Diff
@@ -0,0 +1,61 @@
|
||||
from transformers import PretrainedConfig
|
||||
|
||||
|
||||
class ChatGLMConfig(PretrainedConfig):
|
||||
model_type = "chatglm"
|
||||
def __init__(
|
||||
self,
|
||||
num_layers=28,
|
||||
padded_vocab_size=65024,
|
||||
hidden_size=4096,
|
||||
ffn_hidden_size=13696,
|
||||
kv_channels=128,
|
||||
num_attention_heads=32,
|
||||
seq_length=2048,
|
||||
hidden_dropout=0.0,
|
||||
classifier_dropout=None,
|
||||
attention_dropout=0.0,
|
||||
layernorm_epsilon=1e-5,
|
||||
rmsnorm=True,
|
||||
apply_residual_connection_post_layernorm=False,
|
||||
post_layer_norm=True,
|
||||
add_bias_linear=False,
|
||||
add_qkv_bias=False,
|
||||
bias_dropout_fusion=True,
|
||||
multi_query_attention=False,
|
||||
multi_query_group_num=1,
|
||||
apply_query_key_layer_scaling=True,
|
||||
attention_softmax_in_fp32=True,
|
||||
fp32_residual_connection=False,
|
||||
quantization_bit=0,
|
||||
pre_seq_len=None,
|
||||
prefix_projection=False,
|
||||
**kwargs
|
||||
):
|
||||
self.num_layers = num_layers
|
||||
self.vocab_size = padded_vocab_size
|
||||
self.padded_vocab_size = padded_vocab_size
|
||||
self.hidden_size = hidden_size
|
||||
self.ffn_hidden_size = ffn_hidden_size
|
||||
self.kv_channels = kv_channels
|
||||
self.num_attention_heads = num_attention_heads
|
||||
self.seq_length = seq_length
|
||||
self.hidden_dropout = hidden_dropout
|
||||
self.classifier_dropout = classifier_dropout
|
||||
self.attention_dropout = attention_dropout
|
||||
self.layernorm_epsilon = layernorm_epsilon
|
||||
self.rmsnorm = rmsnorm
|
||||
self.apply_residual_connection_post_layernorm = apply_residual_connection_post_layernorm
|
||||
self.post_layer_norm = post_layer_norm
|
||||
self.add_bias_linear = add_bias_linear
|
||||
self.add_qkv_bias = add_qkv_bias
|
||||
self.bias_dropout_fusion = bias_dropout_fusion
|
||||
self.multi_query_attention = multi_query_attention
|
||||
self.multi_query_group_num = multi_query_group_num
|
||||
self.apply_query_key_layer_scaling = apply_query_key_layer_scaling
|
||||
self.attention_softmax_in_fp32 = attention_softmax_in_fp32
|
||||
self.fp32_residual_connection = fp32_residual_connection
|
||||
self.quantization_bit = quantization_bit
|
||||
self.pre_seq_len = pre_seq_len
|
||||
self.prefix_projection = prefix_projection
|
||||
super().__init__(**kwargs)
|
||||
1300
ixformer_sdk/train/speedformer/models/chatglm/modeling_chatglm.py
Normal file
1300
ixformer_sdk/train/speedformer/models/chatglm/modeling_chatglm.py
Normal file
File diff suppressed because it is too large
Load Diff
File diff suppressed because it is too large
Load Diff
269
ixformer_sdk/train/speedformer/models/gpt2/configuration_gpt2.py
Normal file
269
ixformer_sdk/train/speedformer/models/gpt2/configuration_gpt2.py
Normal file
@@ -0,0 +1,269 @@
|
||||
# coding=utf-8
|
||||
# Copyright 2018 The OpenAI Team Authors and HuggingFace Inc. team.
|
||||
# Copyright (c) 2018, NVIDIA CORPORATION. All rights reserved.
|
||||
#
|
||||
# Licensed under the Apache License, Version 2.0 (the "License");
|
||||
# you may not use this file except in compliance with the License.
|
||||
# You may obtain a copy of the License at
|
||||
#
|
||||
# http://www.apache.org/licenses/LICENSE-2.0
|
||||
#
|
||||
# Unless required by applicable law or agreed to in writing, software
|
||||
# distributed under the License is distributed on an "AS IS" BASIS,
|
||||
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
# See the License for the specific language governing permissions and
|
||||
# limitations under the License.
|
||||
""" OpenAI GPT-2 configuration"""
|
||||
from collections import OrderedDict
|
||||
from typing import Any, List, Mapping, Optional
|
||||
|
||||
from transformers import PreTrainedTokenizer, TensorType, is_torch_available
|
||||
from transformers.configuration_utils import PretrainedConfig
|
||||
from transformers.onnx import OnnxConfigWithPast, PatchingSpec
|
||||
from transformers.utils import logging
|
||||
|
||||
|
||||
logger = logging.get_logger(__name__)
|
||||
|
||||
|
||||
class GPT2Config(PretrainedConfig):
|
||||
"""
|
||||
This is the configuration class to store the configuration of a [`GPT2Model`] or a [`TFGPT2Model`]. It is used to
|
||||
instantiate a GPT-2 model according to the specified arguments, defining the model architecture. Instantiating a
|
||||
configuration with the defaults will yield a similar configuration to that of the GPT-2
|
||||
[openai-community/gpt2](https://huggingface.co/openai-community/gpt2) architecture.
|
||||
|
||||
Configuration objects inherit from [`PretrainedConfig`] and can be used to control the model outputs. Read the
|
||||
documentation from [`PretrainedConfig`] for more information.
|
||||
|
||||
|
||||
Args:
|
||||
vocab_size (`int`, *optional*, defaults to 50257):
|
||||
Vocabulary size of the GPT-2 model. Defines the number of different tokens that can be represented by the
|
||||
`inputs_ids` passed when calling [`GPT2Model`] or [`TFGPT2Model`].
|
||||
n_positions (`int`, *optional*, defaults to 1024):
|
||||
The maximum sequence length that this model might ever be used with. Typically set this to something large
|
||||
just in case (e.g., 512 or 1024 or 2048).
|
||||
n_embd (`int`, *optional*, defaults to 768):
|
||||
Dimensionality of the embeddings and hidden states.
|
||||
n_layer (`int`, *optional*, defaults to 12):
|
||||
Number of hidden layers in the Transformer encoder.
|
||||
n_head (`int`, *optional*, defaults to 12):
|
||||
Number of attention heads for each attention layer in the Transformer encoder.
|
||||
n_inner (`int`, *optional*):
|
||||
Dimensionality of the inner feed-forward layers. `None` will set it to 4 times n_embd
|
||||
activation_function (`str`, *optional*, defaults to `"gelu_new"`):
|
||||
Activation function, to be selected in the list `["relu", "silu", "gelu", "tanh", "gelu_new"]`.
|
||||
resid_pdrop (`float`, *optional*, defaults to 0.1):
|
||||
The dropout probability for all fully connected layers in the embeddings, encoder, and pooler.
|
||||
embd_pdrop (`float`, *optional*, defaults to 0.1):
|
||||
The dropout ratio for the embeddings.
|
||||
attn_pdrop (`float`, *optional*, defaults to 0.1):
|
||||
The dropout ratio for the attention.
|
||||
layer_norm_epsilon (`float`, *optional*, defaults to 1e-05):
|
||||
The epsilon to use in the layer normalization layers.
|
||||
initializer_range (`float`, *optional*, defaults to 0.02):
|
||||
The standard deviation of the truncated_normal_initializer for initializing all weight matrices.
|
||||
summary_type (`string`, *optional*, defaults to `"cls_index"`):
|
||||
Argument used when doing sequence summary, used in the models [`GPT2DoubleHeadsModel`] and
|
||||
[`TFGPT2DoubleHeadsModel`].
|
||||
|
||||
Has to be one of the following options:
|
||||
|
||||
- `"last"`: Take the last token hidden state (like XLNet).
|
||||
- `"first"`: Take the first token hidden state (like BERT).
|
||||
- `"mean"`: Take the mean of all tokens hidden states.
|
||||
- `"cls_index"`: Supply a Tensor of classification token position (like GPT/GPT-2).
|
||||
- `"attn"`: Not implemented now, use multi-head attention.
|
||||
summary_use_proj (`bool`, *optional*, defaults to `True`):
|
||||
Argument used when doing sequence summary, used in the models [`GPT2DoubleHeadsModel`] and
|
||||
[`TFGPT2DoubleHeadsModel`].
|
||||
|
||||
Whether or not to add a projection after the vector extraction.
|
||||
summary_activation (`str`, *optional*):
|
||||
Argument used when doing sequence summary. Used in for the multiple choice head in
|
||||
[`GPT2DoubleHeadsModel`].
|
||||
|
||||
Pass `"tanh"` for a tanh activation to the output, any other value will result in no activation.
|
||||
summary_proj_to_labels (`bool`, *optional*, defaults to `True`):
|
||||
Argument used when doing sequence summary, used in the models [`GPT2DoubleHeadsModel`] and
|
||||
[`TFGPT2DoubleHeadsModel`].
|
||||
|
||||
Whether the projection outputs should have `config.num_labels` or `config.hidden_size` classes.
|
||||
summary_first_dropout (`float`, *optional*, defaults to 0.1):
|
||||
Argument used when doing sequence summary, used in the models [`GPT2DoubleHeadsModel`] and
|
||||
[`TFGPT2DoubleHeadsModel`].
|
||||
|
||||
The dropout ratio to be used after the projection and activation.
|
||||
scale_attn_weights (`bool`, *optional*, defaults to `True`):
|
||||
Scale attention weights by dividing by sqrt(hidden_size)..
|
||||
use_cache (`bool`, *optional*, defaults to `True`):
|
||||
Whether or not the model should return the last key/values attentions (not used by all models).
|
||||
bos_token_id (`int`, *optional*, defaults to 50256):
|
||||
Id of the beginning of sentence token in the vocabulary.
|
||||
eos_token_id (`int`, *optional*, defaults to 50256):
|
||||
Id of the end of sentence token in the vocabulary.
|
||||
scale_attn_by_inverse_layer_idx (`bool`, *optional*, defaults to `False`):
|
||||
Whether to additionally scale attention weights by `1 / layer_idx + 1`.
|
||||
reorder_and_upcast_attn (`bool`, *optional*, defaults to `False`):
|
||||
Whether to scale keys (K) prior to computing attention (dot-product) and upcast attention
|
||||
dot-product/softmax to float() when training with mixed precision.
|
||||
|
||||
Example:
|
||||
|
||||
```python
|
||||
>>> from transformers import GPT2Config, GPT2Model
|
||||
|
||||
>>> # Initializing a GPT2 configuration
|
||||
>>> configuration = GPT2Config()
|
||||
|
||||
>>> # Initializing a model (with random weights) from the configuration
|
||||
>>> model = GPT2Model(configuration)
|
||||
|
||||
>>> # Accessing the model configuration
|
||||
>>> configuration = model.config
|
||||
```"""
|
||||
|
||||
model_type = "gpt2"
|
||||
keys_to_ignore_at_inference = ["past_key_values"]
|
||||
attribute_map = {
|
||||
"hidden_size": "n_embd",
|
||||
"max_position_embeddings": "n_positions",
|
||||
"num_attention_heads": "n_head",
|
||||
"num_hidden_layers": "n_layer",
|
||||
}
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
vocab_size=50257,
|
||||
n_positions=1024,
|
||||
n_embd=768,
|
||||
n_layer=12,
|
||||
n_head=12,
|
||||
n_inner=None,
|
||||
activation_function="gelu_new",
|
||||
resid_pdrop=0.1,
|
||||
embd_pdrop=0.1,
|
||||
attn_pdrop=0.1,
|
||||
layer_norm_epsilon=1e-5,
|
||||
initializer_range=0.02,
|
||||
summary_type="cls_index",
|
||||
summary_use_proj=True,
|
||||
summary_activation=None,
|
||||
summary_proj_to_labels=True,
|
||||
summary_first_dropout=0.1,
|
||||
scale_attn_weights=True,
|
||||
use_cache=True,
|
||||
bos_token_id=50256,
|
||||
eos_token_id=50256,
|
||||
scale_attn_by_inverse_layer_idx=False,
|
||||
reorder_and_upcast_attn=False,
|
||||
**kwargs,
|
||||
):
|
||||
self.vocab_size = vocab_size
|
||||
self.n_positions = n_positions
|
||||
self.n_embd = n_embd
|
||||
self.n_layer = n_layer
|
||||
self.n_head = n_head
|
||||
self.n_inner = n_inner
|
||||
self.activation_function = activation_function
|
||||
self.resid_pdrop = resid_pdrop
|
||||
self.embd_pdrop = embd_pdrop
|
||||
self.attn_pdrop = attn_pdrop
|
||||
self.layer_norm_epsilon = layer_norm_epsilon
|
||||
self.initializer_range = initializer_range
|
||||
self.summary_type = summary_type
|
||||
self.summary_use_proj = summary_use_proj
|
||||
self.summary_activation = summary_activation
|
||||
self.summary_first_dropout = summary_first_dropout
|
||||
self.summary_proj_to_labels = summary_proj_to_labels
|
||||
self.scale_attn_weights = scale_attn_weights
|
||||
self.use_cache = use_cache
|
||||
self.scale_attn_by_inverse_layer_idx = scale_attn_by_inverse_layer_idx
|
||||
self.reorder_and_upcast_attn = reorder_and_upcast_attn
|
||||
|
||||
self.bos_token_id = bos_token_id
|
||||
self.eos_token_id = eos_token_id
|
||||
|
||||
super().__init__(bos_token_id=bos_token_id, eos_token_id=eos_token_id, **kwargs)
|
||||
|
||||
|
||||
class GPT2OnnxConfig(OnnxConfigWithPast):
|
||||
def __init__(
|
||||
self,
|
||||
config: PretrainedConfig,
|
||||
task: str = "default",
|
||||
patching_specs: List[PatchingSpec] = None,
|
||||
use_past: bool = False,
|
||||
):
|
||||
super().__init__(config, task=task, patching_specs=patching_specs, use_past=use_past)
|
||||
if not getattr(self._config, "pad_token_id", None):
|
||||
# TODO: how to do that better?
|
||||
self._config.pad_token_id = 0
|
||||
|
||||
@property
|
||||
def inputs(self) -> Mapping[str, Mapping[int, str]]:
|
||||
common_inputs = OrderedDict({"input_ids": {0: "batch", 1: "sequence"}})
|
||||
if self.use_past:
|
||||
self.fill_with_past_key_values_(common_inputs, direction="inputs")
|
||||
common_inputs["attention_mask"] = {0: "batch", 1: "past_sequence + sequence"}
|
||||
else:
|
||||
common_inputs["attention_mask"] = {0: "batch", 1: "sequence"}
|
||||
|
||||
return common_inputs
|
||||
|
||||
@property
|
||||
def num_layers(self) -> int:
|
||||
return self._config.n_layer
|
||||
|
||||
@property
|
||||
def num_attention_heads(self) -> int:
|
||||
return self._config.n_head
|
||||
|
||||
def generate_dummy_inputs(
|
||||
self,
|
||||
tokenizer: PreTrainedTokenizer,
|
||||
batch_size: int = -1,
|
||||
seq_length: int = -1,
|
||||
is_pair: bool = False,
|
||||
framework: Optional[TensorType] = None,
|
||||
) -> Mapping[str, Any]:
|
||||
common_inputs = super(OnnxConfigWithPast, self).generate_dummy_inputs(
|
||||
tokenizer, batch_size=batch_size, seq_length=seq_length, is_pair=is_pair, framework=framework
|
||||
)
|
||||
|
||||
# We need to order the input in the way they appears in the forward()
|
||||
ordered_inputs = OrderedDict({"input_ids": common_inputs["input_ids"]})
|
||||
|
||||
# Need to add the past_keys
|
||||
if self.use_past:
|
||||
if not is_torch_available():
|
||||
raise ValueError("Cannot generate dummy past_keys inputs without PyTorch installed.")
|
||||
else:
|
||||
import torch
|
||||
|
||||
batch, seqlen = common_inputs["input_ids"].shape
|
||||
# Not using the same length for past_key_values
|
||||
past_key_values_length = seqlen + 2
|
||||
past_shape = (
|
||||
batch,
|
||||
self.num_attention_heads,
|
||||
past_key_values_length,
|
||||
self._config.hidden_size // self.num_attention_heads,
|
||||
)
|
||||
ordered_inputs["past_key_values"] = [
|
||||
(torch.zeros(past_shape), torch.zeros(past_shape)) for _ in range(self.num_layers)
|
||||
]
|
||||
|
||||
ordered_inputs["attention_mask"] = common_inputs["attention_mask"]
|
||||
if self.use_past:
|
||||
mask_dtype = ordered_inputs["attention_mask"].dtype
|
||||
ordered_inputs["attention_mask"] = torch.cat(
|
||||
[ordered_inputs["attention_mask"], torch.ones(batch, past_key_values_length, dtype=mask_dtype)], dim=1
|
||||
)
|
||||
|
||||
return ordered_inputs
|
||||
|
||||
@property
|
||||
def default_onnx_opset(self) -> int:
|
||||
return 13
|
||||
@@ -0,0 +1,500 @@
|
||||
# Copyright 2023 The HuggingFace Team. All rights reserved.
|
||||
#
|
||||
# Licensed under the Apache License, Version 2.0 (the "License");
|
||||
# you may not use this file except in compliance with the License.
|
||||
# You may obtain a copy of the License at
|
||||
#
|
||||
# http://www.apache.org/licenses/LICENSE-2.0
|
||||
#
|
||||
# Unless required by applicable law or agreed to in writing, software
|
||||
# distributed under the License is distributed on an "AS IS" BASIS,
|
||||
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
# See the License for the specific language governing permissions and
|
||||
# limitations under the License.
|
||||
from dataclasses import dataclass
|
||||
from typing import List, Optional, Tuple, Union
|
||||
|
||||
import torch
|
||||
|
||||
|
||||
@dataclass
|
||||
class AttentionMaskConverter:
|
||||
"""
|
||||
A utility attention mask class that allows one to:
|
||||
- Create a causal 4d mask
|
||||
- Create a causal 4d mask with slided window
|
||||
- Convert a 2d attention mask (batch_size, query_length) to a 4d attention mask (batch_size, 1, query_length,
|
||||
key_value_length) that can be multiplied with attention scores
|
||||
|
||||
Examples:
|
||||
|
||||
```python
|
||||
>>> import torch
|
||||
>>> from transformers.modeling_attn_mask_utils import AttentionMaskConverter
|
||||
|
||||
>>> converter = AttentionMaskConverter(True)
|
||||
>>> converter.to_4d(torch.tensor([[0, 0, 0, 1, 1]]), 5, key_value_length=5, dtype=torch.float32)
|
||||
tensor([[[[-3.4028e+38, -3.4028e+38, -3.4028e+38, -3.4028e+38, -3.4028e+38],
|
||||
[-3.4028e+38, -3.4028e+38, -3.4028e+38, -3.4028e+38, -3.4028e+38],
|
||||
[-3.4028e+38, -3.4028e+38, -3.4028e+38, -3.4028e+38, -3.4028e+38],
|
||||
[-3.4028e+38, -3.4028e+38, -3.4028e+38, 0.0000e+00, -3.4028e+38],
|
||||
[-3.4028e+38, -3.4028e+38, -3.4028e+38, 0.0000e+00, 0.0000e+00]]]])
|
||||
```
|
||||
|
||||
Parameters:
|
||||
is_causal (`bool`):
|
||||
Whether the attention mask should be a uni-directional (causal) or bi-directional mask.
|
||||
|
||||
sliding_window (`int`, *optional*):
|
||||
Optionally, the sliding window masks can be created if `sliding_window` is defined to a positive integer.
|
||||
"""
|
||||
|
||||
is_causal: bool
|
||||
sliding_window: int
|
||||
|
||||
def __init__(self, is_causal: bool, sliding_window: Optional[int] = None):
|
||||
self.is_causal = is_causal
|
||||
self.sliding_window = sliding_window
|
||||
|
||||
if self.sliding_window is not None and self.sliding_window <= 0:
|
||||
raise ValueError(
|
||||
f"Make sure that when passing `sliding_window` that its value is a strictly positive integer, not `{self.sliding_window}`"
|
||||
)
|
||||
|
||||
def to_causal_4d(
|
||||
self,
|
||||
batch_size: int,
|
||||
query_length: int,
|
||||
key_value_length: int,
|
||||
dtype: torch.dtype,
|
||||
device: Union[torch.device, "str"] = "cpu",
|
||||
) -> Optional[torch.Tensor]:
|
||||
"""
|
||||
Creates a causal 4D mask of (bsz, head_dim=1, query_length, key_value_length) shape and adds large negative
|
||||
bias to upper right hand triangular matrix (causal mask).
|
||||
"""
|
||||
if not self.is_causal:
|
||||
raise ValueError(f"Please use `to_causal_4d` only if {self.__class__} has `is_causal` set to True.")
|
||||
|
||||
# If shape is not cached, create a new causal mask and cache it
|
||||
input_shape = (batch_size, query_length)
|
||||
past_key_values_length = key_value_length - query_length
|
||||
|
||||
# create causal mask
|
||||
# [bsz, seq_len] -> [bsz, 1, tgt_seq_len, src_seq_len]
|
||||
causal_4d_mask = None
|
||||
if input_shape[-1] > 1 or self.sliding_window is not None:
|
||||
causal_4d_mask = self._make_causal_mask(
|
||||
input_shape,
|
||||
dtype,
|
||||
device=device,
|
||||
past_key_values_length=past_key_values_length,
|
||||
sliding_window=self.sliding_window,
|
||||
)
|
||||
|
||||
return causal_4d_mask
|
||||
|
||||
def to_4d(
|
||||
self,
|
||||
attention_mask_2d: torch.Tensor,
|
||||
query_length: int,
|
||||
dtype: torch.dtype,
|
||||
key_value_length: Optional[int] = None,
|
||||
) -> torch.Tensor:
|
||||
"""
|
||||
Converts 2D attention mask to 4D attention mask by expanding mask to (bsz, head_dim=1, query_length,
|
||||
key_value_length) shape and by adding a large negative bias to not-attended positions. If attention_mask is
|
||||
causal, a causal mask will be added.
|
||||
"""
|
||||
input_shape = (attention_mask_2d.shape[0], query_length)
|
||||
|
||||
# create causal mask
|
||||
# [bsz, seq_len] -> [bsz, 1, tgt_seq_len, src_seq_len]
|
||||
causal_4d_mask = None
|
||||
if (input_shape[-1] > 1 or self.sliding_window is not None) and self.is_causal:
|
||||
if key_value_length is None:
|
||||
raise ValueError(
|
||||
"This attention mask converter is causal. Make sure to pass `key_value_length` to correctly create a causal mask."
|
||||
)
|
||||
|
||||
past_key_values_length = key_value_length - query_length
|
||||
causal_4d_mask = self._make_causal_mask(
|
||||
input_shape,
|
||||
dtype,
|
||||
device=attention_mask_2d.device,
|
||||
past_key_values_length=past_key_values_length,
|
||||
sliding_window=self.sliding_window,
|
||||
)
|
||||
elif self.sliding_window is not None:
|
||||
raise NotImplementedError("Sliding window is currently only implemented for causal masking")
|
||||
|
||||
# [bsz, seq_len] -> [bsz, 1, tgt_seq_len, src_seq_len]
|
||||
expanded_attn_mask = self._expand_mask(attention_mask_2d, dtype, tgt_len=input_shape[-1]).to(
|
||||
attention_mask_2d.device
|
||||
)
|
||||
|
||||
if causal_4d_mask is not None:
|
||||
expanded_attn_mask = causal_4d_mask.masked_fill(expanded_attn_mask.bool(), torch.finfo(dtype).min)
|
||||
|
||||
# expanded_attn_mask + causal_4d_mask can cause some overflow
|
||||
expanded_4d_mask = expanded_attn_mask
|
||||
|
||||
return expanded_4d_mask
|
||||
|
||||
@staticmethod
|
||||
def _make_causal_mask(
|
||||
input_ids_shape: torch.Size,
|
||||
dtype: torch.dtype,
|
||||
device: torch.device,
|
||||
past_key_values_length: int = 0,
|
||||
sliding_window: Optional[int] = None,
|
||||
):
|
||||
"""
|
||||
Make causal mask used for bi-directional self-attention.
|
||||
"""
|
||||
bsz, tgt_len = input_ids_shape
|
||||
mask = torch.full((tgt_len, tgt_len), torch.finfo(dtype).min, device=device)
|
||||
mask_cond = torch.arange(mask.size(-1), device=device)
|
||||
mask.masked_fill_(mask_cond < (mask_cond + 1).view(mask.size(-1), 1), 0)
|
||||
|
||||
mask = mask.to(dtype)
|
||||
|
||||
if past_key_values_length > 0:
|
||||
mask = torch.cat([torch.zeros(tgt_len, past_key_values_length, dtype=dtype, device=device), mask], dim=-1)
|
||||
|
||||
# add lower triangular sliding window mask if necessary
|
||||
if sliding_window is not None:
|
||||
diagonal = past_key_values_length - sliding_window + 1
|
||||
|
||||
context_mask = 1 - torch.triu(torch.ones_like(mask, dtype=torch.int), diagonal=diagonal)
|
||||
mask.masked_fill_(context_mask.bool(), torch.finfo(dtype).min)
|
||||
|
||||
return mask[None, None, :, :].expand(bsz, 1, tgt_len, tgt_len + past_key_values_length)
|
||||
|
||||
@staticmethod
|
||||
def _expand_mask(mask: torch.Tensor, dtype: torch.dtype, tgt_len: Optional[int] = None):
|
||||
"""
|
||||
Expands attention_mask from `[bsz, seq_len]` to `[bsz, 1, tgt_seq_len, src_seq_len]`.
|
||||
"""
|
||||
bsz, src_len = mask.size()
|
||||
tgt_len = tgt_len if tgt_len is not None else src_len
|
||||
|
||||
expanded_mask = mask[:, None, None, :].expand(bsz, 1, tgt_len, src_len).to(dtype)
|
||||
|
||||
inverted_mask = 1.0 - expanded_mask
|
||||
|
||||
return inverted_mask.masked_fill(inverted_mask.to(torch.bool), torch.finfo(dtype).min)
|
||||
|
||||
@staticmethod
|
||||
def _unmask_unattended(
|
||||
expanded_mask: torch.Tensor, attention_mask: torch.Tensor, unmasked_value: Union[bool, float]
|
||||
):
|
||||
# fmt: off
|
||||
"""
|
||||
Attend to all tokens in masked rows from the expanded attention mask, for example the relevant first rows when
|
||||
using left padding. This is required by F.scaled_dot_product_attention memory-efficient attention path.
|
||||
Details: https://github.com/pytorch/pytorch/issues/110213
|
||||
|
||||
`expanded_mask` is [bsz, num_masks, tgt_seq_len, src_seq_len] or [bsz, tgt_seq_len, src_seq_len].
|
||||
`attention_mask` is [bsz, src_seq_len].
|
||||
|
||||
The dimension num_masks of `expanded_mask` is most often 1, but it can also be the number of heads in the case of alibi attention bias.
|
||||
|
||||
For example, if `attention_mask` is
|
||||
```
|
||||
[[0, 0, 1],
|
||||
[1, 1, 1],
|
||||
[0, 1, 1]]
|
||||
```
|
||||
and `expanded_mask` is (e.g. here left-padding case)
|
||||
```
|
||||
[[[[0, 0, 0],
|
||||
[0, 0, 0],
|
||||
[0, 0, 1]]],
|
||||
[[[1, 0, 0],
|
||||
[1, 1, 0],
|
||||
[1, 1, 1]]],
|
||||
[[[0, 0, 0],
|
||||
[0, 1, 0],
|
||||
[0, 1, 1]]]]
|
||||
```
|
||||
then the modified `expanded_mask` will be
|
||||
```
|
||||
[[[[1, 1, 1], <-- modified
|
||||
[1, 1, 1], <-- modified
|
||||
[0, 0, 1]]],
|
||||
[[[1, 0, 0],
|
||||
[1, 1, 0],
|
||||
[1, 1, 1]]],
|
||||
[[[1, 1, 1], <-- modified
|
||||
[0, 1, 0],
|
||||
[0, 1, 1]]]]
|
||||
```
|
||||
"""
|
||||
# fmt: on
|
||||
|
||||
# Get the index of the first non-zero value for every sample in the batch.
|
||||
# In the above example, indices = [[2], [0], [1]]]
|
||||
tmp = torch.arange(attention_mask.shape[1], 0, -1)
|
||||
indices = torch.argmax(attention_mask.cpu() * tmp, 1, keepdim=True)
|
||||
|
||||
# Find the batch indexes that have unattended tokens on the leftmost side (e.g. [0, 0, 1, 1, 1]), for which the first rows of the
|
||||
# expanded mask will be completely unattended.
|
||||
left_masked_rows = torch.where(indices > 0)[0]
|
||||
|
||||
if left_masked_rows.shape[0] == 0:
|
||||
return expanded_mask
|
||||
indices = indices[left_masked_rows]
|
||||
|
||||
max_len = torch.max(indices)
|
||||
range_tensor = torch.arange(max_len).unsqueeze(0)
|
||||
range_tensor = range_tensor.repeat(indices.size(0), 1)
|
||||
|
||||
# Avoid unmasking tokens at relevant target positions (on the row axis), by rather unmasking possibly several times the first row that should always be unmasked as we filtered out the batch above.
|
||||
range_tensor[range_tensor >= indices] = 0
|
||||
|
||||
# TODO: we may drop support for 3D attention mask as the refactor from Patrick maybe dropped this case
|
||||
if expanded_mask.dim() == 4:
|
||||
num_masks = expanded_mask.shape[1]
|
||||
if num_masks == 1:
|
||||
# Broadcast [left_masked_rows, 1], [left_masked_rows, max_len]
|
||||
mask_slice = (left_masked_rows[:, None], 0, range_tensor)
|
||||
else:
|
||||
# Broadcast [left_masked_rows, 1, 1], [1, num_masks, 1], [left_masked_rows, 1, max_len]
|
||||
mask_slice = (
|
||||
left_masked_rows[:, None, None],
|
||||
torch.arange(num_masks)[None, :, None],
|
||||
range_tensor[:, None, :],
|
||||
)
|
||||
else:
|
||||
# Broadcast [left_masked_rows, 1], [left_masked_rows, max_len]
|
||||
mask_slice = (left_masked_rows[:, None], range_tensor)
|
||||
|
||||
expanded_mask[mask_slice] = unmasked_value
|
||||
|
||||
return expanded_mask
|
||||
|
||||
|
||||
def _prepare_4d_causal_attention_mask(
|
||||
attention_mask: Optional[torch.Tensor],
|
||||
input_shape: Union[torch.Size, Tuple, List],
|
||||
inputs_embeds: torch.Tensor,
|
||||
past_key_values_length: int,
|
||||
sliding_window: Optional[int] = None,
|
||||
):
|
||||
"""
|
||||
Creates a causal 4D mask of shape `(batch_size, 1, query_length, key_value_length)` from a 2D mask of shape
|
||||
`(batch_size, key_value_length)`
|
||||
|
||||
Args:
|
||||
attention_mask (`torch.Tensor` or `None`):
|
||||
A 2D attention mask of shape `(batch_size, key_value_length)`
|
||||
input_shape (`tuple(int)` or `list(int)` or `torch.Size`):
|
||||
The input shape should be a tuple that defines `(batch_size, query_length)`.
|
||||
inputs_embeds (`torch.Tensor`):
|
||||
The embedded inputs as a torch Tensor.
|
||||
past_key_values_length (`int`):
|
||||
The length of the key value cache.
|
||||
sliding_window (`int`, *optional*):
|
||||
If the model uses windowed attention, a sliding window should be passed.
|
||||
"""
|
||||
attn_mask_converter = AttentionMaskConverter(is_causal=True, sliding_window=sliding_window)
|
||||
|
||||
key_value_length = input_shape[-1] + past_key_values_length
|
||||
|
||||
# 4d mask is passed through the layers
|
||||
if attention_mask is not None and len(attention_mask.shape) == 2:
|
||||
attention_mask = attn_mask_converter.to_4d(
|
||||
attention_mask, input_shape[-1], key_value_length=key_value_length, dtype=inputs_embeds.dtype
|
||||
)
|
||||
elif attention_mask is not None and len(attention_mask.shape) == 4:
|
||||
expected_shape = (input_shape[0], 1, input_shape[1], key_value_length)
|
||||
if tuple(attention_mask.shape) != expected_shape:
|
||||
raise ValueError(
|
||||
f"Incorrect 4D attention_mask shape: {tuple(attention_mask.shape)}; expected: {expected_shape}."
|
||||
)
|
||||
else:
|
||||
# if the 4D mask has correct shape - invert it and fill with negative infinity
|
||||
inverted_mask = 1.0 - attention_mask
|
||||
attention_mask = inverted_mask.masked_fill(
|
||||
inverted_mask.to(torch.bool), torch.finfo(inputs_embeds.dtype).min
|
||||
)
|
||||
else:
|
||||
attention_mask = attn_mask_converter.to_causal_4d(
|
||||
input_shape[0], input_shape[-1], key_value_length, dtype=inputs_embeds.dtype, device=inputs_embeds.device
|
||||
)
|
||||
|
||||
return attention_mask
|
||||
|
||||
|
||||
# Adapted from _prepare_4d_causal_attention_mask
|
||||
def _prepare_4d_causal_attention_mask_for_sdpa(
|
||||
attention_mask: Optional[torch.Tensor],
|
||||
input_shape: Union[torch.Size, Tuple, List],
|
||||
inputs_embeds: torch.Tensor,
|
||||
past_key_values_length: int,
|
||||
sliding_window: Optional[int] = None,
|
||||
):
|
||||
"""
|
||||
Prepares the correct `attn_mask` argument to be used by `torch.nn.functional.scaled_dot_product_attention`.
|
||||
|
||||
In case no token is masked in the `attention_mask` argument, we simply set it to `None` for the cases `query_length == 1` and
|
||||
`key_value_length == query_length`, and rely instead on SDPA `is_causal` argument to use causal/non-causal masks,
|
||||
allowing to dispatch to the flash attention kernel (that can otherwise not be used if a custom `attn_mask` is passed).
|
||||
"""
|
||||
attn_mask_converter = AttentionMaskConverter(is_causal=True, sliding_window=sliding_window)
|
||||
|
||||
key_value_length = input_shape[-1] + past_key_values_length
|
||||
batch_size, query_length = input_shape
|
||||
|
||||
# torch.jit.trace, symbolic_trace and torchdynamo with fullgraph=True are unable to capture the controlflow `is_causal=attention_mask is None and q_len > 1`
|
||||
# used as an SDPA argument. We keep compatibility with these tracing tools by always using SDPA's `attn_mask` argument in case we are tracing.
|
||||
# TODO: Fix this as well when using torchdynamo with fullgraph=True.
|
||||
is_tracing = torch.jit.is_tracing() or isinstance(inputs_embeds, torch.fx.Proxy)
|
||||
|
||||
if attention_mask is not None:
|
||||
# 4d mask is passed through
|
||||
if len(attention_mask.shape) == 4:
|
||||
expected_shape = (input_shape[0], 1, input_shape[1], key_value_length)
|
||||
if tuple(attention_mask.shape) != expected_shape:
|
||||
raise ValueError(
|
||||
f"Incorrect 4D attention_mask shape: {tuple(attention_mask.shape)}; expected: {expected_shape}."
|
||||
)
|
||||
else:
|
||||
# if the 4D mask has correct shape - invert it and fill with negative infinity
|
||||
inverted_mask = 1.0 - attention_mask.to(inputs_embeds.dtype)
|
||||
attention_mask = inverted_mask.masked_fill(
|
||||
inverted_mask.to(torch.bool), torch.finfo(inputs_embeds.dtype).min
|
||||
)
|
||||
return attention_mask
|
||||
|
||||
elif not is_tracing and torch.all(attention_mask == 1):
|
||||
if query_length == 1:
|
||||
# For query_length == 1, causal attention and bi-directional attention are the same.
|
||||
attention_mask = None
|
||||
elif key_value_length == query_length:
|
||||
attention_mask = None
|
||||
else:
|
||||
# Unfortunately, for query_length > 1 and key_value_length != query_length, we cannot generally ignore the attention mask, as SDPA causal mask generation
|
||||
# may be wrong. We will set `is_causal=False` in SDPA and rely on Transformers attention_mask instead, hence not setting it to None here.
|
||||
# Reference: https://github.com/pytorch/pytorch/issues/108108
|
||||
pass
|
||||
elif query_length > 1 and key_value_length != query_length:
|
||||
# See the comment above (https://github.com/pytorch/pytorch/issues/108108).
|
||||
# Ugly: we set it to True here to dispatch in the following controlflow to `to_causal_4d`.
|
||||
attention_mask = True
|
||||
elif is_tracing:
|
||||
raise ValueError(
|
||||
'Attention using SDPA can not be traced with torch.jit.trace when no attention_mask is provided. To solve this issue, please either load your model with the argument `attn_implementation="eager"` or pass an attention_mask input when tracing the model.'
|
||||
)
|
||||
|
||||
if attention_mask is None:
|
||||
expanded_4d_mask = None
|
||||
elif attention_mask is True:
|
||||
expanded_4d_mask = attn_mask_converter.to_causal_4d(
|
||||
input_shape[0], input_shape[-1], key_value_length, dtype=inputs_embeds.dtype, device=inputs_embeds.device
|
||||
)
|
||||
else:
|
||||
expanded_4d_mask = attn_mask_converter.to_4d(
|
||||
attention_mask,
|
||||
input_shape[-1],
|
||||
dtype=inputs_embeds.dtype,
|
||||
key_value_length=key_value_length,
|
||||
)
|
||||
|
||||
# From PyTorch 2.1 onwards, F.scaled_dot_product_attention with the memory-efficient attention backend
|
||||
# produces nans if sequences are completely unattended in the attention mask. Details: https://github.com/pytorch/pytorch/issues/110213
|
||||
#
|
||||
# This fix is not applied in case we are tracing with torch.jit.trace or symbolic_trace, as _unmask_unattended has a data-dependent
|
||||
# controlflow that can not be captured properly.
|
||||
# TODO: _unmask_unattended does not work either with torch.compile when using fullgraph=True. We should find a way to detect this case.
|
||||
if query_length > 1 and not is_tracing:
|
||||
expanded_4d_mask = AttentionMaskConverter._unmask_unattended(
|
||||
expanded_4d_mask, attention_mask, unmasked_value=0.0
|
||||
)
|
||||
|
||||
return expanded_4d_mask
|
||||
|
||||
|
||||
def _prepare_4d_attention_mask(mask: torch.Tensor, dtype: torch.dtype, tgt_len: Optional[int] = None):
|
||||
"""
|
||||
Creates a non-causal 4D mask of shape `(batch_size, 1, query_length, key_value_length)` from a 2D mask of shape
|
||||
`(batch_size, key_value_length)`
|
||||
|
||||
Args:
|
||||
mask (`torch.Tensor` or `None`):
|
||||
A 2D attention mask of shape `(batch_size, key_value_length)`
|
||||
dtype (`torch.dtype`):
|
||||
The torch dtype the created mask shall have.
|
||||
tgt_len (`int`):
|
||||
The target length or query length the created mask shall have.
|
||||
"""
|
||||
return AttentionMaskConverter._expand_mask(mask=mask, dtype=dtype, tgt_len=tgt_len)
|
||||
|
||||
|
||||
def _prepare_4d_attention_mask_for_sdpa(mask: torch.Tensor, dtype: torch.dtype, tgt_len: Optional[int] = None):
|
||||
"""
|
||||
Creates a non-causal 4D mask of shape `(batch_size, 1, query_length, key_value_length)` from a 2D mask of shape
|
||||
`(batch_size, key_value_length)`
|
||||
|
||||
Args:
|
||||
mask (`torch.Tensor` or `None`):
|
||||
A 2D attention mask of shape `(batch_size, key_value_length)`
|
||||
dtype (`torch.dtype`):
|
||||
The torch dtype the created mask shall have.
|
||||
tgt_len (`int`):
|
||||
The target length or query length the created mask shall have.
|
||||
"""
|
||||
batch_size, key_value_length = mask.shape
|
||||
tgt_len = tgt_len if tgt_len is not None else key_value_length
|
||||
|
||||
# torch.jit.trace and torchdynamo with fullgraph=True are unable to capture the controlflow `is_causal=attention_mask is None and q_len > 1`
|
||||
# used as an SDPA argument. We keep compatibility with these tracing tools by always using SDPA's `attn_mask` argument in case we are tracing.
|
||||
# TODO: Fix this as well when using torchdynamo with fullgraph=True.
|
||||
is_tracing = torch.jit.is_tracing()
|
||||
|
||||
if torch.all(mask == 1):
|
||||
if is_tracing:
|
||||
pass
|
||||
elif tgt_len == 1:
|
||||
# For query_length == 1, causal attention and bi-directional attention are the same.
|
||||
return None
|
||||
elif key_value_length == tgt_len:
|
||||
return None
|
||||
else:
|
||||
# Unfortunately, for query_length > 1 and key_value_length != query_length, we can not generally ignore the attention mask, as SDPA causal mask generation
|
||||
# may be wrong. We will set is_causal=False in SDPA and rely on Transformers attention_mask instead, hence not setting it to None here.
|
||||
# Reference: https://github.com/pytorch/pytorch/issues/108108
|
||||
return AttentionMaskConverter._expand_mask(mask=mask, dtype=dtype, tgt_len=tgt_len)
|
||||
else:
|
||||
return AttentionMaskConverter._expand_mask(mask=mask, dtype=dtype, tgt_len=tgt_len)
|
||||
|
||||
|
||||
def _create_4d_causal_attention_mask(
|
||||
input_shape: Union[torch.Size, Tuple, List],
|
||||
dtype: torch.dtype,
|
||||
device: torch.device,
|
||||
past_key_values_length: int = 0,
|
||||
sliding_window: Optional[int] = None,
|
||||
) -> Optional[torch.Tensor]:
|
||||
"""
|
||||
Creates a causal 4D mask of shape `(batch_size, 1, query_length, key_value_length)`
|
||||
|
||||
Args:
|
||||
input_shape (`tuple(int)` or `list(int)` or `torch.Size`):
|
||||
The input shape should be a tuple that defines `(batch_size, query_length)`.
|
||||
dtype (`torch.dtype`):
|
||||
The torch dtype the created mask shall have.
|
||||
device (`int`):
|
||||
The torch device the created mask shall have.
|
||||
sliding_window (`int`, *optional*):
|
||||
If the model uses windowed attention, a sliding window should be passed.
|
||||
"""
|
||||
attn_mask_converter = AttentionMaskConverter(is_causal=True, sliding_window=sliding_window)
|
||||
|
||||
key_value_length = past_key_values_length + input_shape[-1]
|
||||
attention_mask = attn_mask_converter.to_causal_4d(
|
||||
input_shape[0], input_shape[-1], key_value_length, dtype=dtype, device=device
|
||||
)
|
||||
|
||||
return attention_mask
|
||||
1949
ixformer_sdk/train/speedformer/models/gpt2/modeling_gpt2.py
Normal file
1949
ixformer_sdk/train/speedformer/models/gpt2/modeling_gpt2.py
Normal file
File diff suppressed because it is too large
Load Diff
@@ -0,0 +1,191 @@
|
||||
# coding=utf-8
|
||||
# Copyright 2022 EleutherAI and the HuggingFace Inc. team. All rights reserved.
|
||||
#
|
||||
# This code is based on EleutherAI's GPT-NeoX library and the GPT-NeoX
|
||||
# and OPT implementations in this library. It has been modified from its
|
||||
# original forms to accommodate minor architectural differences compared
|
||||
# to GPT-NeoX and OPT used by the Meta AI team that trained the model.
|
||||
#
|
||||
# Licensed under the Apache License, Version 2.0 (the "License");
|
||||
# you may not use this file except in compliance with the License.
|
||||
# You may obtain a copy of the License at
|
||||
#
|
||||
# http://www.apache.org/licenses/LICENSE-2.0
|
||||
#
|
||||
# Unless required by applicable law or agreed to in writing, software
|
||||
# distributed under the License is distributed on an "AS IS" BASIS,
|
||||
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
# See the License for the specific language governing permissions and
|
||||
# limitations under the License.
|
||||
""" LLaMA model configuration"""
|
||||
|
||||
from transformers.configuration_utils import PretrainedConfig
|
||||
from transformers.utils import logging
|
||||
|
||||
|
||||
logger = logging.get_logger(__name__)
|
||||
|
||||
LLAMA_PRETRAINED_CONFIG_ARCHIVE_MAP = {}
|
||||
|
||||
|
||||
class LlamaConfig(PretrainedConfig):
|
||||
r"""
|
||||
This is the configuration class to store the configuration of a [`LlamaModel`]. It is used to instantiate an LLaMA
|
||||
model according to the specified arguments, defining the model architecture. Instantiating a configuration with the
|
||||
defaults will yield a similar configuration to that of the LLaMA-7B.
|
||||
|
||||
Configuration objects inherit from [`PretrainedConfig`] and can be used to control the model outputs. Read the
|
||||
documentation from [`PretrainedConfig`] for more information.
|
||||
|
||||
|
||||
Args:
|
||||
vocab_size (`int`, *optional*, defaults to 32000):
|
||||
Vocabulary size of the LLaMA model. Defines the number of different tokens that can be represented by the
|
||||
`inputs_ids` passed when calling [`LlamaModel`]
|
||||
hidden_size (`int`, *optional*, defaults to 4096):
|
||||
Dimension of the hidden representations.
|
||||
intermediate_size (`int`, *optional*, defaults to 11008):
|
||||
Dimension of the MLP representations.
|
||||
num_hidden_layers (`int`, *optional*, defaults to 32):
|
||||
Number of hidden layers in the Transformer decoder.
|
||||
num_attention_heads (`int`, *optional*, defaults to 32):
|
||||
Number of attention heads for each attention layer in the Transformer decoder.
|
||||
num_key_value_heads (`int`, *optional*):
|
||||
This is the number of key_value heads that should be used to implement Grouped Query Attention. If
|
||||
`num_key_value_heads=num_attention_heads`, the model will use Multi Head Attention (MHA), if
|
||||
`num_key_value_heads=1 the model will use Multi Query Attention (MQA) otherwise GQA is used. When
|
||||
converting a multi-head checkpoint to a GQA checkpoint, each group key and value head should be constructed
|
||||
by meanpooling all the original heads within that group. For more details checkout [this
|
||||
paper](https://arxiv.org/pdf/2305.13245.pdf). If it is not specified, will default to
|
||||
`num_attention_heads`.
|
||||
hidden_act (`str` or `function`, *optional*, defaults to `"silu"`):
|
||||
The non-linear activation function (function or string) in the decoder.
|
||||
max_position_embeddings (`int`, *optional*, defaults to 2048):
|
||||
The maximum sequence length that this model might ever be used with. Llama 1 supports up to 2048 tokens,
|
||||
Llama 2 up to 4096, CodeLlama up to 16384.
|
||||
initializer_range (`float`, *optional*, defaults to 0.02):
|
||||
The standard deviation of the truncated_normal_initializer for initializing all weight matrices.
|
||||
rms_norm_eps (`float`, *optional*, defaults to 1e-06):
|
||||
The epsilon used by the rms normalization layers.
|
||||
use_cache (`bool`, *optional*, defaults to `True`):
|
||||
Whether or not the model should return the last key/values attentions (not used by all models). Only
|
||||
relevant if `config.is_decoder=True`.
|
||||
pad_token_id (`int`, *optional*):
|
||||
Padding token id.
|
||||
bos_token_id (`int`, *optional*, defaults to 1):
|
||||
Beginning of stream token id.
|
||||
eos_token_id (`int`, *optional*, defaults to 2):
|
||||
End of stream token id.
|
||||
pretraining_tp (`int`, *optional*, defaults to 1):
|
||||
Experimental feature. Tensor parallelism rank used during pretraining. Please refer to [this
|
||||
document](https://huggingface.co/docs/transformers/parallelism) to understand more about it. This value is
|
||||
necessary to ensure exact reproducibility of the pretraining results. Please refer to [this
|
||||
issue](https://github.com/pytorch/pytorch/issues/76232).
|
||||
tie_word_embeddings (`bool`, *optional*, defaults to `False`):
|
||||
Whether to tie weight embeddings
|
||||
rope_theta (`float`, *optional*, defaults to 10000.0):
|
||||
The base period of the RoPE embeddings.
|
||||
rope_scaling (`Dict`, *optional*):
|
||||
Dictionary containing the scaling configuration for the RoPE embeddings. Currently supports two scaling
|
||||
strategies: linear and dynamic. Their scaling factor must be a float greater than 1. The expected format is
|
||||
`{"type": strategy name, "factor": scaling factor}`. When using this flag, don't update
|
||||
`max_position_embeddings` to the expected new maximum. See the following thread for more information on how
|
||||
these scaling strategies behave:
|
||||
https://www.reddit.com/r/LocalLLaMA/comments/14mrgpr/dynamically_scaled_rope_further_increases/. This is an
|
||||
experimental feature, subject to breaking API changes in future versions.
|
||||
attention_bias (`bool`, defaults to `False`, *optional*, defaults to `False`):
|
||||
Whether to use a bias in the query, key, value and output projection layers during self-attention.
|
||||
attention_dropout (`float`, *optional*, defaults to 0.0):
|
||||
The dropout ratio for the attention probabilities.
|
||||
|
||||
```python
|
||||
>>> from transformers import LlamaModel, LlamaConfig
|
||||
|
||||
>>> # Initializing a LLaMA llama-7b style configuration
|
||||
>>> configuration = LlamaConfig()
|
||||
|
||||
>>> # Initializing a model from the llama-7b style configuration
|
||||
>>> model = LlamaModel(configuration)
|
||||
|
||||
>>> # Accessing the model configuration
|
||||
>>> configuration = model.config
|
||||
```"""
|
||||
|
||||
model_type = "llama"
|
||||
keys_to_ignore_at_inference = ["past_key_values"]
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
vocab_size=32000,
|
||||
hidden_size=4096,
|
||||
intermediate_size=11008,
|
||||
num_hidden_layers=32,
|
||||
num_attention_heads=32,
|
||||
num_key_value_heads=None,
|
||||
hidden_act="silu",
|
||||
max_position_embeddings=2048,
|
||||
initializer_range=0.02,
|
||||
rms_norm_eps=1e-6,
|
||||
use_cache=True,
|
||||
pad_token_id=None,
|
||||
bos_token_id=1,
|
||||
eos_token_id=2,
|
||||
pretraining_tp=1,
|
||||
tie_word_embeddings=False,
|
||||
rope_theta=10000.0,
|
||||
rope_scaling=None,
|
||||
attention_bias=False,
|
||||
attention_dropout=0.0,
|
||||
**kwargs,
|
||||
):
|
||||
self.vocab_size = vocab_size
|
||||
self.max_position_embeddings = max_position_embeddings
|
||||
self.hidden_size = hidden_size
|
||||
self.intermediate_size = intermediate_size
|
||||
self.num_hidden_layers = num_hidden_layers
|
||||
self.num_attention_heads = num_attention_heads
|
||||
|
||||
# for backward compatibility
|
||||
if num_key_value_heads is None:
|
||||
num_key_value_heads = num_attention_heads
|
||||
|
||||
self.num_key_value_heads = num_key_value_heads
|
||||
self.hidden_act = hidden_act
|
||||
self.initializer_range = initializer_range
|
||||
self.rms_norm_eps = rms_norm_eps
|
||||
self.pretraining_tp = pretraining_tp
|
||||
self.use_cache = use_cache
|
||||
self.rope_theta = rope_theta
|
||||
self.rope_scaling = rope_scaling
|
||||
self._rope_scaling_validation()
|
||||
self.attention_bias = attention_bias
|
||||
self.attention_dropout = attention_dropout
|
||||
|
||||
super().__init__(
|
||||
pad_token_id=pad_token_id,
|
||||
bos_token_id=bos_token_id,
|
||||
eos_token_id=eos_token_id,
|
||||
tie_word_embeddings=tie_word_embeddings,
|
||||
**kwargs,
|
||||
)
|
||||
|
||||
def _rope_scaling_validation(self):
|
||||
"""
|
||||
Validate the `rope_scaling` configuration.
|
||||
"""
|
||||
if self.rope_scaling is None:
|
||||
return
|
||||
|
||||
if not isinstance(self.rope_scaling, dict) or len(self.rope_scaling) != 2:
|
||||
raise ValueError(
|
||||
"`rope_scaling` must be a dictionary with with two fields, `type` and `factor`, "
|
||||
f"got {self.rope_scaling}"
|
||||
)
|
||||
rope_scaling_type = self.rope_scaling.get("type", None)
|
||||
rope_scaling_factor = self.rope_scaling.get("factor", None)
|
||||
if rope_scaling_type is None or rope_scaling_type not in ["linear", "dynamic"]:
|
||||
raise ValueError(
|
||||
f"`rope_scaling`'s type field must be one of ['linear', 'dynamic'], got {rope_scaling_type}"
|
||||
)
|
||||
if rope_scaling_factor is None or not isinstance(rope_scaling_factor, float) or rope_scaling_factor <= 1.0:
|
||||
raise ValueError(f"`rope_scaling`'s factor field must be a float > 1, got {rope_scaling_factor}")
|
||||
@@ -0,0 +1,500 @@
|
||||
# Copyright 2023 The HuggingFace Team. All rights reserved.
|
||||
#
|
||||
# Licensed under the Apache License, Version 2.0 (the "License");
|
||||
# you may not use this file except in compliance with the License.
|
||||
# You may obtain a copy of the License at
|
||||
#
|
||||
# http://www.apache.org/licenses/LICENSE-2.0
|
||||
#
|
||||
# Unless required by applicable law or agreed to in writing, software
|
||||
# distributed under the License is distributed on an "AS IS" BASIS,
|
||||
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
# See the License for the specific language governing permissions and
|
||||
# limitations under the License.
|
||||
from dataclasses import dataclass
|
||||
from typing import List, Optional, Tuple, Union
|
||||
|
||||
import torch
|
||||
|
||||
|
||||
@dataclass
|
||||
class AttentionMaskConverter:
|
||||
"""
|
||||
A utility attention mask class that allows one to:
|
||||
- Create a causal 4d mask
|
||||
- Create a causal 4d mask with slided window
|
||||
- Convert a 2d attention mask (batch_size, query_length) to a 4d attention mask (batch_size, 1, query_length,
|
||||
key_value_length) that can be multiplied with attention scores
|
||||
|
||||
Examples:
|
||||
|
||||
```python
|
||||
>>> import torch
|
||||
>>> from transformers.modeling_attn_mask_utils import AttentionMaskConverter
|
||||
|
||||
>>> converter = AttentionMaskConverter(True)
|
||||
>>> converter.to_4d(torch.tensor([[0, 0, 0, 1, 1]]), 5, key_value_length=5, dtype=torch.float32)
|
||||
tensor([[[[-3.4028e+38, -3.4028e+38, -3.4028e+38, -3.4028e+38, -3.4028e+38],
|
||||
[-3.4028e+38, -3.4028e+38, -3.4028e+38, -3.4028e+38, -3.4028e+38],
|
||||
[-3.4028e+38, -3.4028e+38, -3.4028e+38, -3.4028e+38, -3.4028e+38],
|
||||
[-3.4028e+38, -3.4028e+38, -3.4028e+38, 0.0000e+00, -3.4028e+38],
|
||||
[-3.4028e+38, -3.4028e+38, -3.4028e+38, 0.0000e+00, 0.0000e+00]]]])
|
||||
```
|
||||
|
||||
Parameters:
|
||||
is_causal (`bool`):
|
||||
Whether the attention mask should be a uni-directional (causal) or bi-directional mask.
|
||||
|
||||
sliding_window (`int`, *optional*):
|
||||
Optionally, the sliding window masks can be created if `sliding_window` is defined to a positive integer.
|
||||
"""
|
||||
|
||||
is_causal: bool
|
||||
sliding_window: int
|
||||
|
||||
def __init__(self, is_causal: bool, sliding_window: Optional[int] = None):
|
||||
self.is_causal = is_causal
|
||||
self.sliding_window = sliding_window
|
||||
|
||||
if self.sliding_window is not None and self.sliding_window <= 0:
|
||||
raise ValueError(
|
||||
f"Make sure that when passing `sliding_window` that its value is a strictly positive integer, not `{self.sliding_window}`"
|
||||
)
|
||||
|
||||
def to_causal_4d(
|
||||
self,
|
||||
batch_size: int,
|
||||
query_length: int,
|
||||
key_value_length: int,
|
||||
dtype: torch.dtype,
|
||||
device: Union[torch.device, "str"] = "cpu",
|
||||
) -> Optional[torch.Tensor]:
|
||||
"""
|
||||
Creates a causal 4D mask of (bsz, head_dim=1, query_length, key_value_length) shape and adds large negative
|
||||
bias to upper right hand triangular matrix (causal mask).
|
||||
"""
|
||||
if not self.is_causal:
|
||||
raise ValueError(f"Please use `to_causal_4d` only if {self.__class__} has `is_causal` set to True.")
|
||||
|
||||
# If shape is not cached, create a new causal mask and cache it
|
||||
input_shape = (batch_size, query_length)
|
||||
past_key_values_length = key_value_length - query_length
|
||||
|
||||
# create causal mask
|
||||
# [bsz, seq_len] -> [bsz, 1, tgt_seq_len, src_seq_len]
|
||||
causal_4d_mask = None
|
||||
if input_shape[-1] > 1 or self.sliding_window is not None:
|
||||
causal_4d_mask = self._make_causal_mask(
|
||||
input_shape,
|
||||
dtype,
|
||||
device=device,
|
||||
past_key_values_length=past_key_values_length,
|
||||
sliding_window=self.sliding_window,
|
||||
)
|
||||
|
||||
return causal_4d_mask
|
||||
|
||||
def to_4d(
|
||||
self,
|
||||
attention_mask_2d: torch.Tensor,
|
||||
query_length: int,
|
||||
dtype: torch.dtype,
|
||||
key_value_length: Optional[int] = None,
|
||||
) -> torch.Tensor:
|
||||
"""
|
||||
Converts 2D attention mask to 4D attention mask by expanding mask to (bsz, head_dim=1, query_length,
|
||||
key_value_length) shape and by adding a large negative bias to not-attended positions. If attention_mask is
|
||||
causal, a causal mask will be added.
|
||||
"""
|
||||
input_shape = (attention_mask_2d.shape[0], query_length)
|
||||
|
||||
# create causal mask
|
||||
# [bsz, seq_len] -> [bsz, 1, tgt_seq_len, src_seq_len]
|
||||
causal_4d_mask = None
|
||||
if (input_shape[-1] > 1 or self.sliding_window is not None) and self.is_causal:
|
||||
if key_value_length is None:
|
||||
raise ValueError(
|
||||
"This attention mask converter is causal. Make sure to pass `key_value_length` to correctly create a causal mask."
|
||||
)
|
||||
|
||||
past_key_values_length = key_value_length - query_length
|
||||
causal_4d_mask = self._make_causal_mask(
|
||||
input_shape,
|
||||
dtype,
|
||||
device=attention_mask_2d.device,
|
||||
past_key_values_length=past_key_values_length,
|
||||
sliding_window=self.sliding_window,
|
||||
)
|
||||
elif self.sliding_window is not None:
|
||||
raise NotImplementedError("Sliding window is currently only implemented for causal masking")
|
||||
|
||||
# [bsz, seq_len] -> [bsz, 1, tgt_seq_len, src_seq_len]
|
||||
expanded_attn_mask = self._expand_mask(attention_mask_2d, dtype, tgt_len=input_shape[-1]).to(
|
||||
attention_mask_2d.device
|
||||
)
|
||||
|
||||
if causal_4d_mask is not None:
|
||||
expanded_attn_mask = causal_4d_mask.masked_fill(expanded_attn_mask.bool(), torch.finfo(dtype).min)
|
||||
|
||||
# expanded_attn_mask + causal_4d_mask can cause some overflow
|
||||
expanded_4d_mask = expanded_attn_mask
|
||||
|
||||
return expanded_4d_mask
|
||||
|
||||
@staticmethod
|
||||
def _make_causal_mask(
|
||||
input_ids_shape: torch.Size,
|
||||
dtype: torch.dtype,
|
||||
device: torch.device,
|
||||
past_key_values_length: int = 0,
|
||||
sliding_window: Optional[int] = None,
|
||||
):
|
||||
"""
|
||||
Make causal mask used for bi-directional self-attention.
|
||||
"""
|
||||
bsz, tgt_len = input_ids_shape
|
||||
mask = torch.full((tgt_len, tgt_len), torch.finfo(dtype).min, device=device)
|
||||
mask_cond = torch.arange(mask.size(-1), device=device)
|
||||
mask.masked_fill_(mask_cond < (mask_cond + 1).view(mask.size(-1), 1), 0)
|
||||
|
||||
mask = mask.to(dtype)
|
||||
|
||||
if past_key_values_length > 0:
|
||||
mask = torch.cat([torch.zeros(tgt_len, past_key_values_length, dtype=dtype, device=device), mask], dim=-1)
|
||||
|
||||
# add lower triangular sliding window mask if necessary
|
||||
if sliding_window is not None:
|
||||
diagonal = past_key_values_length - sliding_window + 1
|
||||
|
||||
context_mask = 1 - torch.triu(torch.ones_like(mask, dtype=torch.int), diagonal=diagonal)
|
||||
mask.masked_fill_(context_mask.bool(), torch.finfo(dtype).min)
|
||||
|
||||
return mask[None, None, :, :].expand(bsz, 1, tgt_len, tgt_len + past_key_values_length)
|
||||
|
||||
@staticmethod
|
||||
def _expand_mask(mask: torch.Tensor, dtype: torch.dtype, tgt_len: Optional[int] = None):
|
||||
"""
|
||||
Expands attention_mask from `[bsz, seq_len]` to `[bsz, 1, tgt_seq_len, src_seq_len]`.
|
||||
"""
|
||||
bsz, src_len = mask.size()
|
||||
tgt_len = tgt_len if tgt_len is not None else src_len
|
||||
|
||||
expanded_mask = mask[:, None, None, :].expand(bsz, 1, tgt_len, src_len).to(dtype)
|
||||
|
||||
inverted_mask = 1.0 - expanded_mask
|
||||
|
||||
return inverted_mask.masked_fill(inverted_mask.to(torch.bool), torch.finfo(dtype).min)
|
||||
|
||||
@staticmethod
|
||||
def _unmask_unattended(
|
||||
expanded_mask: torch.Tensor, attention_mask: torch.Tensor, unmasked_value: Union[bool, float]
|
||||
):
|
||||
# fmt: off
|
||||
"""
|
||||
Attend to all tokens in masked rows from the expanded attention mask, for example the relevant first rows when
|
||||
using left padding. This is required by F.scaled_dot_product_attention memory-efficient attention path.
|
||||
Details: https://github.com/pytorch/pytorch/issues/110213
|
||||
|
||||
`expanded_mask` is [bsz, num_masks, tgt_seq_len, src_seq_len] or [bsz, tgt_seq_len, src_seq_len].
|
||||
`attention_mask` is [bsz, src_seq_len].
|
||||
|
||||
The dimension num_masks of `expanded_mask` is most often 1, but it can also be the number of heads in the case of alibi attention bias.
|
||||
|
||||
For example, if `attention_mask` is
|
||||
```
|
||||
[[0, 0, 1],
|
||||
[1, 1, 1],
|
||||
[0, 1, 1]]
|
||||
```
|
||||
and `expanded_mask` is (e.g. here left-padding case)
|
||||
```
|
||||
[[[[0, 0, 0],
|
||||
[0, 0, 0],
|
||||
[0, 0, 1]]],
|
||||
[[[1, 0, 0],
|
||||
[1, 1, 0],
|
||||
[1, 1, 1]]],
|
||||
[[[0, 0, 0],
|
||||
[0, 1, 0],
|
||||
[0, 1, 1]]]]
|
||||
```
|
||||
then the modified `expanded_mask` will be
|
||||
```
|
||||
[[[[1, 1, 1], <-- modified
|
||||
[1, 1, 1], <-- modified
|
||||
[0, 0, 1]]],
|
||||
[[[1, 0, 0],
|
||||
[1, 1, 0],
|
||||
[1, 1, 1]]],
|
||||
[[[1, 1, 1], <-- modified
|
||||
[0, 1, 0],
|
||||
[0, 1, 1]]]]
|
||||
```
|
||||
"""
|
||||
# fmt: on
|
||||
|
||||
# Get the index of the first non-zero value for every sample in the batch.
|
||||
# In the above example, indices = [[2], [0], [1]]]
|
||||
tmp = torch.arange(attention_mask.shape[1], 0, -1)
|
||||
indices = torch.argmax(attention_mask.cpu() * tmp, 1, keepdim=True)
|
||||
|
||||
# Find the batch indexes that have unattended tokens on the leftmost side (e.g. [0, 0, 1, 1, 1]), for which the first rows of the
|
||||
# expanded mask will be completely unattended.
|
||||
left_masked_rows = torch.where(indices > 0)[0]
|
||||
|
||||
if left_masked_rows.shape[0] == 0:
|
||||
return expanded_mask
|
||||
indices = indices[left_masked_rows]
|
||||
|
||||
max_len = torch.max(indices)
|
||||
range_tensor = torch.arange(max_len).unsqueeze(0)
|
||||
range_tensor = range_tensor.repeat(indices.size(0), 1)
|
||||
|
||||
# Avoid unmasking tokens at relevant target positions (on the row axis), by rather unmasking possibly several times the first row that should always be unmasked as we filtered out the batch above.
|
||||
range_tensor[range_tensor >= indices] = 0
|
||||
|
||||
# TODO: we may drop support for 3D attention mask as the refactor from Patrick maybe dropped this case
|
||||
if expanded_mask.dim() == 4:
|
||||
num_masks = expanded_mask.shape[1]
|
||||
if num_masks == 1:
|
||||
# Broadcast [left_masked_rows, 1], [left_masked_rows, max_len]
|
||||
mask_slice = (left_masked_rows[:, None], 0, range_tensor)
|
||||
else:
|
||||
# Broadcast [left_masked_rows, 1, 1], [1, num_masks, 1], [left_masked_rows, 1, max_len]
|
||||
mask_slice = (
|
||||
left_masked_rows[:, None, None],
|
||||
torch.arange(num_masks)[None, :, None],
|
||||
range_tensor[:, None, :],
|
||||
)
|
||||
else:
|
||||
# Broadcast [left_masked_rows, 1], [left_masked_rows, max_len]
|
||||
mask_slice = (left_masked_rows[:, None], range_tensor)
|
||||
|
||||
expanded_mask[mask_slice] = unmasked_value
|
||||
|
||||
return expanded_mask
|
||||
|
||||
|
||||
def _prepare_4d_causal_attention_mask(
|
||||
attention_mask: Optional[torch.Tensor],
|
||||
input_shape: Union[torch.Size, Tuple, List],
|
||||
inputs_embeds: torch.Tensor,
|
||||
past_key_values_length: int,
|
||||
sliding_window: Optional[int] = None,
|
||||
):
|
||||
"""
|
||||
Creates a causal 4D mask of shape `(batch_size, 1, query_length, key_value_length)` from a 2D mask of shape
|
||||
`(batch_size, key_value_length)`
|
||||
|
||||
Args:
|
||||
attention_mask (`torch.Tensor` or `None`):
|
||||
A 2D attention mask of shape `(batch_size, key_value_length)`
|
||||
input_shape (`tuple(int)` or `list(int)` or `torch.Size`):
|
||||
The input shape should be a tuple that defines `(batch_size, query_length)`.
|
||||
inputs_embeds (`torch.Tensor`):
|
||||
The embedded inputs as a torch Tensor.
|
||||
past_key_values_length (`int`):
|
||||
The length of the key value cache.
|
||||
sliding_window (`int`, *optional*):
|
||||
If the model uses windowed attention, a sliding window should be passed.
|
||||
"""
|
||||
attn_mask_converter = AttentionMaskConverter(is_causal=True, sliding_window=sliding_window)
|
||||
|
||||
key_value_length = input_shape[-1] + past_key_values_length
|
||||
|
||||
# 4d mask is passed through the layers
|
||||
if attention_mask is not None and len(attention_mask.shape) == 2:
|
||||
attention_mask = attn_mask_converter.to_4d(
|
||||
attention_mask, input_shape[-1], key_value_length=key_value_length, dtype=inputs_embeds.dtype
|
||||
)
|
||||
elif attention_mask is not None and len(attention_mask.shape) == 4:
|
||||
expected_shape = (input_shape[0], 1, input_shape[1], key_value_length)
|
||||
if tuple(attention_mask.shape) != expected_shape:
|
||||
raise ValueError(
|
||||
f"Incorrect 4D attention_mask shape: {tuple(attention_mask.shape)}; expected: {expected_shape}."
|
||||
)
|
||||
else:
|
||||
# if the 4D mask has correct shape - invert it and fill with negative infinity
|
||||
inverted_mask = 1.0 - attention_mask
|
||||
attention_mask = inverted_mask.masked_fill(
|
||||
inverted_mask.to(torch.bool), torch.finfo(inputs_embeds.dtype).min
|
||||
)
|
||||
else:
|
||||
attention_mask = attn_mask_converter.to_causal_4d(
|
||||
input_shape[0], input_shape[-1], key_value_length, dtype=inputs_embeds.dtype, device=inputs_embeds.device
|
||||
)
|
||||
|
||||
return attention_mask
|
||||
|
||||
|
||||
# Adapted from _prepare_4d_causal_attention_mask
|
||||
def _prepare_4d_causal_attention_mask_for_sdpa(
|
||||
attention_mask: Optional[torch.Tensor],
|
||||
input_shape: Union[torch.Size, Tuple, List],
|
||||
inputs_embeds: torch.Tensor,
|
||||
past_key_values_length: int,
|
||||
sliding_window: Optional[int] = None,
|
||||
):
|
||||
"""
|
||||
Prepares the correct `attn_mask` argument to be used by `torch.nn.functional.scaled_dot_product_attention`.
|
||||
|
||||
In case no token is masked in the `attention_mask` argument, we simply set it to `None` for the cases `query_length == 1` and
|
||||
`key_value_length == query_length`, and rely instead on SDPA `is_causal` argument to use causal/non-causal masks,
|
||||
allowing to dispatch to the flash attention kernel (that can otherwise not be used if a custom `attn_mask` is passed).
|
||||
"""
|
||||
attn_mask_converter = AttentionMaskConverter(is_causal=True, sliding_window=sliding_window)
|
||||
|
||||
key_value_length = input_shape[-1] + past_key_values_length
|
||||
batch_size, query_length = input_shape
|
||||
|
||||
# torch.jit.trace, symbolic_trace and torchdynamo with fullgraph=True are unable to capture the controlflow `is_causal=attention_mask is None and q_len > 1`
|
||||
# used as an SDPA argument. We keep compatibility with these tracing tools by always using SDPA's `attn_mask` argument in case we are tracing.
|
||||
# TODO: Fix this as well when using torchdynamo with fullgraph=True.
|
||||
is_tracing = torch.jit.is_tracing() or isinstance(inputs_embeds, torch.fx.Proxy)
|
||||
|
||||
if attention_mask is not None:
|
||||
# 4d mask is passed through
|
||||
if len(attention_mask.shape) == 4:
|
||||
expected_shape = (input_shape[0], 1, input_shape[1], key_value_length)
|
||||
if tuple(attention_mask.shape) != expected_shape:
|
||||
raise ValueError(
|
||||
f"Incorrect 4D attention_mask shape: {tuple(attention_mask.shape)}; expected: {expected_shape}."
|
||||
)
|
||||
else:
|
||||
# if the 4D mask has correct shape - invert it and fill with negative infinity
|
||||
inverted_mask = 1.0 - attention_mask.to(inputs_embeds.dtype)
|
||||
attention_mask = inverted_mask.masked_fill(
|
||||
inverted_mask.to(torch.bool), torch.finfo(inputs_embeds.dtype).min
|
||||
)
|
||||
return attention_mask
|
||||
|
||||
elif not is_tracing and torch.all(attention_mask == 1):
|
||||
if query_length == 1:
|
||||
# For query_length == 1, causal attention and bi-directional attention are the same.
|
||||
attention_mask = None
|
||||
elif key_value_length == query_length:
|
||||
attention_mask = None
|
||||
else:
|
||||
# Unfortunately, for query_length > 1 and key_value_length != query_length, we cannot generally ignore the attention mask, as SDPA causal mask generation
|
||||
# may be wrong. We will set `is_causal=False` in SDPA and rely on Transformers attention_mask instead, hence not setting it to None here.
|
||||
# Reference: https://github.com/pytorch/pytorch/issues/108108
|
||||
pass
|
||||
elif query_length > 1 and key_value_length != query_length:
|
||||
# See the comment above (https://github.com/pytorch/pytorch/issues/108108).
|
||||
# Ugly: we set it to True here to dispatch in the following controlflow to `to_causal_4d`.
|
||||
attention_mask = True
|
||||
elif is_tracing:
|
||||
raise ValueError(
|
||||
'Attention using SDPA can not be traced with torch.jit.trace when no attention_mask is provided. To solve this issue, please either load your model with the argument `attn_implementation="eager"` or pass an attention_mask input when tracing the model.'
|
||||
)
|
||||
|
||||
if attention_mask is None:
|
||||
expanded_4d_mask = None
|
||||
elif attention_mask is True:
|
||||
expanded_4d_mask = attn_mask_converter.to_causal_4d(
|
||||
input_shape[0], input_shape[-1], key_value_length, dtype=inputs_embeds.dtype, device=inputs_embeds.device
|
||||
)
|
||||
else:
|
||||
expanded_4d_mask = attn_mask_converter.to_4d(
|
||||
attention_mask,
|
||||
input_shape[-1],
|
||||
dtype=inputs_embeds.dtype,
|
||||
key_value_length=key_value_length,
|
||||
)
|
||||
|
||||
# From PyTorch 2.1 onwards, F.scaled_dot_product_attention with the memory-efficient attention backend
|
||||
# produces nans if sequences are completely unattended in the attention mask. Details: https://github.com/pytorch/pytorch/issues/110213
|
||||
#
|
||||
# This fix is not applied in case we are tracing with torch.jit.trace or symbolic_trace, as _unmask_unattended has a data-dependent
|
||||
# controlflow that can not be captured properly.
|
||||
# TODO: _unmask_unattended does not work either with torch.compile when using fullgraph=True. We should find a way to detect this case.
|
||||
if query_length > 1 and not is_tracing:
|
||||
expanded_4d_mask = AttentionMaskConverter._unmask_unattended(
|
||||
expanded_4d_mask, attention_mask, unmasked_value=0.0
|
||||
)
|
||||
|
||||
return expanded_4d_mask
|
||||
|
||||
|
||||
def _prepare_4d_attention_mask(mask: torch.Tensor, dtype: torch.dtype, tgt_len: Optional[int] = None):
|
||||
"""
|
||||
Creates a non-causal 4D mask of shape `(batch_size, 1, query_length, key_value_length)` from a 2D mask of shape
|
||||
`(batch_size, key_value_length)`
|
||||
|
||||
Args:
|
||||
mask (`torch.Tensor` or `None`):
|
||||
A 2D attention mask of shape `(batch_size, key_value_length)`
|
||||
dtype (`torch.dtype`):
|
||||
The torch dtype the created mask shall have.
|
||||
tgt_len (`int`):
|
||||
The target length or query length the created mask shall have.
|
||||
"""
|
||||
return AttentionMaskConverter._expand_mask(mask=mask, dtype=dtype, tgt_len=tgt_len)
|
||||
|
||||
|
||||
def _prepare_4d_attention_mask_for_sdpa(mask: torch.Tensor, dtype: torch.dtype, tgt_len: Optional[int] = None):
|
||||
"""
|
||||
Creates a non-causal 4D mask of shape `(batch_size, 1, query_length, key_value_length)` from a 2D mask of shape
|
||||
`(batch_size, key_value_length)`
|
||||
|
||||
Args:
|
||||
mask (`torch.Tensor` or `None`):
|
||||
A 2D attention mask of shape `(batch_size, key_value_length)`
|
||||
dtype (`torch.dtype`):
|
||||
The torch dtype the created mask shall have.
|
||||
tgt_len (`int`):
|
||||
The target length or query length the created mask shall have.
|
||||
"""
|
||||
batch_size, key_value_length = mask.shape
|
||||
tgt_len = tgt_len if tgt_len is not None else key_value_length
|
||||
|
||||
# torch.jit.trace and torchdynamo with fullgraph=True are unable to capture the controlflow `is_causal=attention_mask is None and q_len > 1`
|
||||
# used as an SDPA argument. We keep compatibility with these tracing tools by always using SDPA's `attn_mask` argument in case we are tracing.
|
||||
# TODO: Fix this as well when using torchdynamo with fullgraph=True.
|
||||
is_tracing = torch.jit.is_tracing()
|
||||
|
||||
if torch.all(mask == 1):
|
||||
if is_tracing:
|
||||
pass
|
||||
elif tgt_len == 1:
|
||||
# For query_length == 1, causal attention and bi-directional attention are the same.
|
||||
return None
|
||||
elif key_value_length == tgt_len:
|
||||
return None
|
||||
else:
|
||||
# Unfortunately, for query_length > 1 and key_value_length != query_length, we can not generally ignore the attention mask, as SDPA causal mask generation
|
||||
# may be wrong. We will set is_causal=False in SDPA and rely on Transformers attention_mask instead, hence not setting it to None here.
|
||||
# Reference: https://github.com/pytorch/pytorch/issues/108108
|
||||
return AttentionMaskConverter._expand_mask(mask=mask, dtype=dtype, tgt_len=tgt_len)
|
||||
else:
|
||||
return AttentionMaskConverter._expand_mask(mask=mask, dtype=dtype, tgt_len=tgt_len)
|
||||
|
||||
|
||||
def _create_4d_causal_attention_mask(
|
||||
input_shape: Union[torch.Size, Tuple, List],
|
||||
dtype: torch.dtype,
|
||||
device: torch.device,
|
||||
past_key_values_length: int = 0,
|
||||
sliding_window: Optional[int] = None,
|
||||
) -> Optional[torch.Tensor]:
|
||||
"""
|
||||
Creates a causal 4D mask of shape `(batch_size, 1, query_length, key_value_length)`
|
||||
|
||||
Args:
|
||||
input_shape (`tuple(int)` or `list(int)` or `torch.Size`):
|
||||
The input shape should be a tuple that defines `(batch_size, query_length)`.
|
||||
dtype (`torch.dtype`):
|
||||
The torch dtype the created mask shall have.
|
||||
device (`int`):
|
||||
The torch device the created mask shall have.
|
||||
sliding_window (`int`, *optional*):
|
||||
If the model uses windowed attention, a sliding window should be passed.
|
||||
"""
|
||||
attn_mask_converter = AttentionMaskConverter(is_causal=True, sliding_window=sliding_window)
|
||||
|
||||
key_value_length = past_key_values_length + input_shape[-1]
|
||||
attention_mask = attn_mask_converter.to_causal_4d(
|
||||
input_shape[0], input_shape[-1], key_value_length, dtype=dtype, device=device
|
||||
)
|
||||
|
||||
return attention_mask
|
||||
1415
ixformer_sdk/train/speedformer/models/llama/modeling_llama.py
Normal file
1415
ixformer_sdk/train/speedformer/models/llama/modeling_llama.py
Normal file
File diff suppressed because it is too large
Load Diff
@@ -0,0 +1,144 @@
|
||||
# coding=utf-8
|
||||
# Copyright 2024 The Qwen team, Alibaba Group and the HuggingFace Inc. team. All rights reserved.
|
||||
#
|
||||
# Licensed under the Apache License, Version 2.0 (the "License");
|
||||
# you may not use this file except in compliance with the License.
|
||||
# You may obtain a copy of the License at
|
||||
#
|
||||
# http://www.apache.org/licenses/LICENSE-2.0
|
||||
#
|
||||
# Unless required by applicable law or agreed to in writing, software
|
||||
# distributed under the License is distributed on an "AS IS" BASIS,
|
||||
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
# See the License for the specific language governing permissions and
|
||||
# limitations under the License.
|
||||
""" Qwen2 model configuration"""
|
||||
|
||||
from transformers.configuration_utils import PretrainedConfig
|
||||
from transformers.utils import logging
|
||||
|
||||
|
||||
logger = logging.get_logger(__name__)
|
||||
|
||||
QWEN2_PRETRAINED_CONFIG_ARCHIVE_MAP = {
|
||||
"Qwen/Qwen2-7B-beta": "https://huggingface.co/Qwen/Qwen2-7B-beta/resolve/main/config.json",
|
||||
}
|
||||
|
||||
|
||||
class Qwen2Config(PretrainedConfig):
|
||||
r"""
|
||||
This is the configuration class to store the configuration of a [`Qwen2Model`]. It is used to instantiate a
|
||||
Qwen2 model according to the specified arguments, defining the model architecture. Instantiating a configuration
|
||||
with the defaults will yield a similar configuration to that of
|
||||
Qwen2-7B-beta [Qwen/Qwen2-7B-beta](https://huggingface.co/Qwen/Qwen2-7B-beta).
|
||||
|
||||
Configuration objects inherit from [`PretrainedConfig`] and can be used to control the model outputs. Read the
|
||||
documentation from [`PretrainedConfig`] for more information.
|
||||
|
||||
|
||||
Args:
|
||||
vocab_size (`int`, *optional*, defaults to 151936):
|
||||
Vocabulary size of the Qwen2 model. Defines the number of different tokens that can be represented by the
|
||||
`inputs_ids` passed when calling [`Qwen2Model`]
|
||||
hidden_size (`int`, *optional*, defaults to 4096):
|
||||
Dimension of the hidden representations.
|
||||
intermediate_size (`int`, *optional*, defaults to 22016):
|
||||
Dimension of the MLP representations.
|
||||
num_hidden_layers (`int`, *optional*, defaults to 32):
|
||||
Number of hidden layers in the Transformer encoder.
|
||||
num_attention_heads (`int`, *optional*, defaults to 32):
|
||||
Number of attention heads for each attention layer in the Transformer encoder.
|
||||
num_key_value_heads (`int`, *optional*, defaults to 32):
|
||||
This is the number of key_value heads that should be used to implement Grouped Query Attention. If
|
||||
`num_key_value_heads=num_attention_heads`, the model will use Multi Head Attention (MHA), if
|
||||
`num_key_value_heads=1 the model will use Multi Query Attention (MQA) otherwise GQA is used. When
|
||||
converting a multi-head checkpoint to a GQA checkpoint, each group key and value head should be constructed
|
||||
by meanpooling all the original heads within that group. For more details checkout [this
|
||||
paper](https://arxiv.org/pdf/2305.13245.pdf). If it is not specified, will default to `32`.
|
||||
hidden_act (`str` or `function`, *optional*, defaults to `"silu"`):
|
||||
The non-linear activation function (function or string) in the decoder.
|
||||
max_position_embeddings (`int`, *optional*, defaults to 32768):
|
||||
The maximum sequence length that this model might ever be used with.
|
||||
initializer_range (`float`, *optional*, defaults to 0.02):
|
||||
The standard deviation of the truncated_normal_initializer for initializing all weight matrices.
|
||||
rms_norm_eps (`float`, *optional*, defaults to 1e-06):
|
||||
The epsilon used by the rms normalization layers.
|
||||
use_cache (`bool`, *optional*, defaults to `True`):
|
||||
Whether or not the model should return the last key/values attentions (not used by all models). Only
|
||||
relevant if `config.is_decoder=True`.
|
||||
tie_word_embeddings (`bool`, *optional*, defaults to `False`):
|
||||
Whether the model's input and output word embeddings should be tied.
|
||||
rope_theta (`float`, *optional*, defaults to 10000.0):
|
||||
The base period of the RoPE embeddings.
|
||||
use_sliding_window (`bool`, *optional*, defaults to `False`):
|
||||
Whether to use sliding window attention.
|
||||
sliding_window (`int`, *optional*, defaults to 4096):
|
||||
Sliding window attention (SWA) window size. If not specified, will default to `4096`.
|
||||
max_window_layers (`int`, *optional*, defaults to 28):
|
||||
The number of layers that use SWA (Sliding Window Attention). The bottom layers use SWA while the top use full attention.
|
||||
attention_dropout (`float`, *optional*, defaults to 0.0):
|
||||
The dropout ratio for the attention probabilities.
|
||||
|
||||
```python
|
||||
>>> from transformers import Qwen2Model, Qwen2Config
|
||||
|
||||
>>> # Initializing a Qwen2 style configuration
|
||||
>>> configuration = Qwen2Config()
|
||||
|
||||
>>> # Initializing a model from the Qwen2-7B style configuration
|
||||
>>> model = Qwen2Model(configuration)
|
||||
|
||||
>>> # Accessing the model configuration
|
||||
>>> configuration = model.config
|
||||
```"""
|
||||
|
||||
model_type = "qwen2"
|
||||
keys_to_ignore_at_inference = ["past_key_values"]
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
vocab_size=151936,
|
||||
hidden_size=4096,
|
||||
intermediate_size=22016,
|
||||
num_hidden_layers=32,
|
||||
num_attention_heads=32,
|
||||
num_key_value_heads=32,
|
||||
hidden_act="silu",
|
||||
max_position_embeddings=32768,
|
||||
initializer_range=0.02,
|
||||
rms_norm_eps=1e-6,
|
||||
use_cache=True,
|
||||
tie_word_embeddings=False,
|
||||
rope_theta=10000.0,
|
||||
use_sliding_window=False,
|
||||
sliding_window=4096,
|
||||
max_window_layers=28,
|
||||
attention_dropout=0.0,
|
||||
**kwargs,
|
||||
):
|
||||
self.vocab_size = vocab_size
|
||||
self.max_position_embeddings = max_position_embeddings
|
||||
self.hidden_size = hidden_size
|
||||
self.intermediate_size = intermediate_size
|
||||
self.num_hidden_layers = num_hidden_layers
|
||||
self.num_attention_heads = num_attention_heads
|
||||
self.use_sliding_window = use_sliding_window
|
||||
self.sliding_window = sliding_window
|
||||
self.max_window_layers = max_window_layers
|
||||
|
||||
# for backward compatibility
|
||||
if num_key_value_heads is None:
|
||||
num_key_value_heads = num_attention_heads
|
||||
|
||||
self.num_key_value_heads = num_key_value_heads
|
||||
self.hidden_act = hidden_act
|
||||
self.initializer_range = initializer_range
|
||||
self.rms_norm_eps = rms_norm_eps
|
||||
self.use_cache = use_cache
|
||||
self.rope_theta = rope_theta
|
||||
self.attention_dropout = attention_dropout
|
||||
|
||||
super().__init__(
|
||||
tie_word_embeddings=tie_word_embeddings,
|
||||
**kwargs,
|
||||
)
|
||||
@@ -0,0 +1,500 @@
|
||||
# Copyright 2023 The HuggingFace Team. All rights reserved.
|
||||
#
|
||||
# Licensed under the Apache License, Version 2.0 (the "License");
|
||||
# you may not use this file except in compliance with the License.
|
||||
# You may obtain a copy of the License at
|
||||
#
|
||||
# http://www.apache.org/licenses/LICENSE-2.0
|
||||
#
|
||||
# Unless required by applicable law or agreed to in writing, software
|
||||
# distributed under the License is distributed on an "AS IS" BASIS,
|
||||
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
# See the License for the specific language governing permissions and
|
||||
# limitations under the License.
|
||||
from dataclasses import dataclass
|
||||
from typing import List, Optional, Tuple, Union
|
||||
|
||||
import torch
|
||||
|
||||
|
||||
@dataclass
|
||||
class AttentionMaskConverter:
|
||||
"""
|
||||
A utility attention mask class that allows one to:
|
||||
- Create a causal 4d mask
|
||||
- Create a causal 4d mask with slided window
|
||||
- Convert a 2d attention mask (batch_size, query_length) to a 4d attention mask (batch_size, 1, query_length,
|
||||
key_value_length) that can be multiplied with attention scores
|
||||
|
||||
Examples:
|
||||
|
||||
```python
|
||||
>>> import torch
|
||||
>>> from transformers.modeling_attn_mask_utils import AttentionMaskConverter
|
||||
|
||||
>>> converter = AttentionMaskConverter(True)
|
||||
>>> converter.to_4d(torch.tensor([[0, 0, 0, 1, 1]]), 5, key_value_length=5, dtype=torch.float32)
|
||||
tensor([[[[-3.4028e+38, -3.4028e+38, -3.4028e+38, -3.4028e+38, -3.4028e+38],
|
||||
[-3.4028e+38, -3.4028e+38, -3.4028e+38, -3.4028e+38, -3.4028e+38],
|
||||
[-3.4028e+38, -3.4028e+38, -3.4028e+38, -3.4028e+38, -3.4028e+38],
|
||||
[-3.4028e+38, -3.4028e+38, -3.4028e+38, 0.0000e+00, -3.4028e+38],
|
||||
[-3.4028e+38, -3.4028e+38, -3.4028e+38, 0.0000e+00, 0.0000e+00]]]])
|
||||
```
|
||||
|
||||
Parameters:
|
||||
is_causal (`bool`):
|
||||
Whether the attention mask should be a uni-directional (causal) or bi-directional mask.
|
||||
|
||||
sliding_window (`int`, *optional*):
|
||||
Optionally, the sliding window masks can be created if `sliding_window` is defined to a positive integer.
|
||||
"""
|
||||
|
||||
is_causal: bool
|
||||
sliding_window: int
|
||||
|
||||
def __init__(self, is_causal: bool, sliding_window: Optional[int] = None):
|
||||
self.is_causal = is_causal
|
||||
self.sliding_window = sliding_window
|
||||
|
||||
if self.sliding_window is not None and self.sliding_window <= 0:
|
||||
raise ValueError(
|
||||
f"Make sure that when passing `sliding_window` that its value is a strictly positive integer, not `{self.sliding_window}`"
|
||||
)
|
||||
|
||||
def to_causal_4d(
|
||||
self,
|
||||
batch_size: int,
|
||||
query_length: int,
|
||||
key_value_length: int,
|
||||
dtype: torch.dtype,
|
||||
device: Union[torch.device, "str"] = "cpu",
|
||||
) -> Optional[torch.Tensor]:
|
||||
"""
|
||||
Creates a causal 4D mask of (bsz, head_dim=1, query_length, key_value_length) shape and adds large negative
|
||||
bias to upper right hand triangular matrix (causal mask).
|
||||
"""
|
||||
if not self.is_causal:
|
||||
raise ValueError(f"Please use `to_causal_4d` only if {self.__class__} has `is_causal` set to True.")
|
||||
|
||||
# If shape is not cached, create a new causal mask and cache it
|
||||
input_shape = (batch_size, query_length)
|
||||
past_key_values_length = key_value_length - query_length
|
||||
|
||||
# create causal mask
|
||||
# [bsz, seq_len] -> [bsz, 1, tgt_seq_len, src_seq_len]
|
||||
causal_4d_mask = None
|
||||
if input_shape[-1] > 1 or self.sliding_window is not None:
|
||||
causal_4d_mask = self._make_causal_mask(
|
||||
input_shape,
|
||||
dtype,
|
||||
device=device,
|
||||
past_key_values_length=past_key_values_length,
|
||||
sliding_window=self.sliding_window,
|
||||
)
|
||||
|
||||
return causal_4d_mask
|
||||
|
||||
def to_4d(
|
||||
self,
|
||||
attention_mask_2d: torch.Tensor,
|
||||
query_length: int,
|
||||
dtype: torch.dtype,
|
||||
key_value_length: Optional[int] = None,
|
||||
) -> torch.Tensor:
|
||||
"""
|
||||
Converts 2D attention mask to 4D attention mask by expanding mask to (bsz, head_dim=1, query_length,
|
||||
key_value_length) shape and by adding a large negative bias to not-attended positions. If attention_mask is
|
||||
causal, a causal mask will be added.
|
||||
"""
|
||||
input_shape = (attention_mask_2d.shape[0], query_length)
|
||||
|
||||
# create causal mask
|
||||
# [bsz, seq_len] -> [bsz, 1, tgt_seq_len, src_seq_len]
|
||||
causal_4d_mask = None
|
||||
if (input_shape[-1] > 1 or self.sliding_window is not None) and self.is_causal:
|
||||
if key_value_length is None:
|
||||
raise ValueError(
|
||||
"This attention mask converter is causal. Make sure to pass `key_value_length` to correctly create a causal mask."
|
||||
)
|
||||
|
||||
past_key_values_length = key_value_length - query_length
|
||||
causal_4d_mask = self._make_causal_mask(
|
||||
input_shape,
|
||||
dtype,
|
||||
device=attention_mask_2d.device,
|
||||
past_key_values_length=past_key_values_length,
|
||||
sliding_window=self.sliding_window,
|
||||
)
|
||||
elif self.sliding_window is not None:
|
||||
raise NotImplementedError("Sliding window is currently only implemented for causal masking")
|
||||
|
||||
# [bsz, seq_len] -> [bsz, 1, tgt_seq_len, src_seq_len]
|
||||
expanded_attn_mask = self._expand_mask(attention_mask_2d, dtype, tgt_len=input_shape[-1]).to(
|
||||
attention_mask_2d.device
|
||||
)
|
||||
|
||||
if causal_4d_mask is not None:
|
||||
expanded_attn_mask = causal_4d_mask.masked_fill(expanded_attn_mask.bool(), torch.finfo(dtype).min)
|
||||
|
||||
# expanded_attn_mask + causal_4d_mask can cause some overflow
|
||||
expanded_4d_mask = expanded_attn_mask
|
||||
|
||||
return expanded_4d_mask
|
||||
|
||||
@staticmethod
|
||||
def _make_causal_mask(
|
||||
input_ids_shape: torch.Size,
|
||||
dtype: torch.dtype,
|
||||
device: torch.device,
|
||||
past_key_values_length: int = 0,
|
||||
sliding_window: Optional[int] = None,
|
||||
):
|
||||
"""
|
||||
Make causal mask used for bi-directional self-attention.
|
||||
"""
|
||||
bsz, tgt_len = input_ids_shape
|
||||
mask = torch.full((tgt_len, tgt_len), torch.finfo(dtype).min, device=device)
|
||||
mask_cond = torch.arange(mask.size(-1), device=device)
|
||||
mask.masked_fill_(mask_cond < (mask_cond + 1).view(mask.size(-1), 1), 0)
|
||||
|
||||
mask = mask.to(dtype)
|
||||
|
||||
if past_key_values_length > 0:
|
||||
mask = torch.cat([torch.zeros(tgt_len, past_key_values_length, dtype=dtype, device=device), mask], dim=-1)
|
||||
|
||||
# add lower triangular sliding window mask if necessary
|
||||
if sliding_window is not None:
|
||||
diagonal = past_key_values_length - sliding_window + 1
|
||||
|
||||
context_mask = 1 - torch.triu(torch.ones_like(mask, dtype=torch.int), diagonal=diagonal)
|
||||
mask.masked_fill_(context_mask.bool(), torch.finfo(dtype).min)
|
||||
|
||||
return mask[None, None, :, :].expand(bsz, 1, tgt_len, tgt_len + past_key_values_length)
|
||||
|
||||
@staticmethod
|
||||
def _expand_mask(mask: torch.Tensor, dtype: torch.dtype, tgt_len: Optional[int] = None):
|
||||
"""
|
||||
Expands attention_mask from `[bsz, seq_len]` to `[bsz, 1, tgt_seq_len, src_seq_len]`.
|
||||
"""
|
||||
bsz, src_len = mask.size()
|
||||
tgt_len = tgt_len if tgt_len is not None else src_len
|
||||
|
||||
expanded_mask = mask[:, None, None, :].expand(bsz, 1, tgt_len, src_len).to(dtype)
|
||||
|
||||
inverted_mask = 1.0 - expanded_mask
|
||||
|
||||
return inverted_mask.masked_fill(inverted_mask.to(torch.bool), torch.finfo(dtype).min)
|
||||
|
||||
@staticmethod
|
||||
def _unmask_unattended(
|
||||
expanded_mask: torch.Tensor, attention_mask: torch.Tensor, unmasked_value: Union[bool, float]
|
||||
):
|
||||
# fmt: off
|
||||
"""
|
||||
Attend to all tokens in masked rows from the expanded attention mask, for example the relevant first rows when
|
||||
using left padding. This is required by F.scaled_dot_product_attention memory-efficient attention path.
|
||||
Details: https://github.com/pytorch/pytorch/issues/110213
|
||||
|
||||
`expanded_mask` is [bsz, num_masks, tgt_seq_len, src_seq_len] or [bsz, tgt_seq_len, src_seq_len].
|
||||
`attention_mask` is [bsz, src_seq_len].
|
||||
|
||||
The dimension num_masks of `expanded_mask` is most often 1, but it can also be the number of heads in the case of alibi attention bias.
|
||||
|
||||
For example, if `attention_mask` is
|
||||
```
|
||||
[[0, 0, 1],
|
||||
[1, 1, 1],
|
||||
[0, 1, 1]]
|
||||
```
|
||||
and `expanded_mask` is (e.g. here left-padding case)
|
||||
```
|
||||
[[[[0, 0, 0],
|
||||
[0, 0, 0],
|
||||
[0, 0, 1]]],
|
||||
[[[1, 0, 0],
|
||||
[1, 1, 0],
|
||||
[1, 1, 1]]],
|
||||
[[[0, 0, 0],
|
||||
[0, 1, 0],
|
||||
[0, 1, 1]]]]
|
||||
```
|
||||
then the modified `expanded_mask` will be
|
||||
```
|
||||
[[[[1, 1, 1], <-- modified
|
||||
[1, 1, 1], <-- modified
|
||||
[0, 0, 1]]],
|
||||
[[[1, 0, 0],
|
||||
[1, 1, 0],
|
||||
[1, 1, 1]]],
|
||||
[[[1, 1, 1], <-- modified
|
||||
[0, 1, 0],
|
||||
[0, 1, 1]]]]
|
||||
```
|
||||
"""
|
||||
# fmt: on
|
||||
|
||||
# Get the index of the first non-zero value for every sample in the batch.
|
||||
# In the above example, indices = [[2], [0], [1]]]
|
||||
tmp = torch.arange(attention_mask.shape[1], 0, -1)
|
||||
indices = torch.argmax(attention_mask.cpu() * tmp, 1, keepdim=True)
|
||||
|
||||
# Find the batch indexes that have unattended tokens on the leftmost side (e.g. [0, 0, 1, 1, 1]), for which the first rows of the
|
||||
# expanded mask will be completely unattended.
|
||||
left_masked_rows = torch.where(indices > 0)[0]
|
||||
|
||||
if left_masked_rows.shape[0] == 0:
|
||||
return expanded_mask
|
||||
indices = indices[left_masked_rows]
|
||||
|
||||
max_len = torch.max(indices)
|
||||
range_tensor = torch.arange(max_len).unsqueeze(0)
|
||||
range_tensor = range_tensor.repeat(indices.size(0), 1)
|
||||
|
||||
# Avoid unmasking tokens at relevant target positions (on the row axis), by rather unmasking possibly several times the first row that should always be unmasked as we filtered out the batch above.
|
||||
range_tensor[range_tensor >= indices] = 0
|
||||
|
||||
# TODO: we may drop support for 3D attention mask as the refactor from Patrick maybe dropped this case
|
||||
if expanded_mask.dim() == 4:
|
||||
num_masks = expanded_mask.shape[1]
|
||||
if num_masks == 1:
|
||||
# Broadcast [left_masked_rows, 1], [left_masked_rows, max_len]
|
||||
mask_slice = (left_masked_rows[:, None], 0, range_tensor)
|
||||
else:
|
||||
# Broadcast [left_masked_rows, 1, 1], [1, num_masks, 1], [left_masked_rows, 1, max_len]
|
||||
mask_slice = (
|
||||
left_masked_rows[:, None, None],
|
||||
torch.arange(num_masks)[None, :, None],
|
||||
range_tensor[:, None, :],
|
||||
)
|
||||
else:
|
||||
# Broadcast [left_masked_rows, 1], [left_masked_rows, max_len]
|
||||
mask_slice = (left_masked_rows[:, None], range_tensor)
|
||||
|
||||
expanded_mask[mask_slice] = unmasked_value
|
||||
|
||||
return expanded_mask
|
||||
|
||||
|
||||
def _prepare_4d_causal_attention_mask(
|
||||
attention_mask: Optional[torch.Tensor],
|
||||
input_shape: Union[torch.Size, Tuple, List],
|
||||
inputs_embeds: torch.Tensor,
|
||||
past_key_values_length: int,
|
||||
sliding_window: Optional[int] = None,
|
||||
):
|
||||
"""
|
||||
Creates a causal 4D mask of shape `(batch_size, 1, query_length, key_value_length)` from a 2D mask of shape
|
||||
`(batch_size, key_value_length)`
|
||||
|
||||
Args:
|
||||
attention_mask (`torch.Tensor` or `None`):
|
||||
A 2D attention mask of shape `(batch_size, key_value_length)`
|
||||
input_shape (`tuple(int)` or `list(int)` or `torch.Size`):
|
||||
The input shape should be a tuple that defines `(batch_size, query_length)`.
|
||||
inputs_embeds (`torch.Tensor`):
|
||||
The embedded inputs as a torch Tensor.
|
||||
past_key_values_length (`int`):
|
||||
The length of the key value cache.
|
||||
sliding_window (`int`, *optional*):
|
||||
If the model uses windowed attention, a sliding window should be passed.
|
||||
"""
|
||||
attn_mask_converter = AttentionMaskConverter(is_causal=True, sliding_window=sliding_window)
|
||||
|
||||
key_value_length = input_shape[-1] + past_key_values_length
|
||||
|
||||
# 4d mask is passed through the layers
|
||||
if attention_mask is not None and len(attention_mask.shape) == 2:
|
||||
attention_mask = attn_mask_converter.to_4d(
|
||||
attention_mask, input_shape[-1], key_value_length=key_value_length, dtype=inputs_embeds.dtype
|
||||
)
|
||||
elif attention_mask is not None and len(attention_mask.shape) == 4:
|
||||
expected_shape = (input_shape[0], 1, input_shape[1], key_value_length)
|
||||
if tuple(attention_mask.shape) != expected_shape:
|
||||
raise ValueError(
|
||||
f"Incorrect 4D attention_mask shape: {tuple(attention_mask.shape)}; expected: {expected_shape}."
|
||||
)
|
||||
else:
|
||||
# if the 4D mask has correct shape - invert it and fill with negative infinity
|
||||
inverted_mask = 1.0 - attention_mask
|
||||
attention_mask = inverted_mask.masked_fill(
|
||||
inverted_mask.to(torch.bool), torch.finfo(inputs_embeds.dtype).min
|
||||
)
|
||||
else:
|
||||
attention_mask = attn_mask_converter.to_causal_4d(
|
||||
input_shape[0], input_shape[-1], key_value_length, dtype=inputs_embeds.dtype, device=inputs_embeds.device
|
||||
)
|
||||
|
||||
return attention_mask
|
||||
|
||||
|
||||
# Adapted from _prepare_4d_causal_attention_mask
|
||||
def _prepare_4d_causal_attention_mask_for_sdpa(
|
||||
attention_mask: Optional[torch.Tensor],
|
||||
input_shape: Union[torch.Size, Tuple, List],
|
||||
inputs_embeds: torch.Tensor,
|
||||
past_key_values_length: int,
|
||||
sliding_window: Optional[int] = None,
|
||||
):
|
||||
"""
|
||||
Prepares the correct `attn_mask` argument to be used by `torch.nn.functional.scaled_dot_product_attention`.
|
||||
|
||||
In case no token is masked in the `attention_mask` argument, we simply set it to `None` for the cases `query_length == 1` and
|
||||
`key_value_length == query_length`, and rely instead on SDPA `is_causal` argument to use causal/non-causal masks,
|
||||
allowing to dispatch to the flash attention kernel (that can otherwise not be used if a custom `attn_mask` is passed).
|
||||
"""
|
||||
attn_mask_converter = AttentionMaskConverter(is_causal=True, sliding_window=sliding_window)
|
||||
|
||||
key_value_length = input_shape[-1] + past_key_values_length
|
||||
batch_size, query_length = input_shape
|
||||
|
||||
# torch.jit.trace, symbolic_trace and torchdynamo with fullgraph=True are unable to capture the controlflow `is_causal=attention_mask is None and q_len > 1`
|
||||
# used as an SDPA argument. We keep compatibility with these tracing tools by always using SDPA's `attn_mask` argument in case we are tracing.
|
||||
# TODO: Fix this as well when using torchdynamo with fullgraph=True.
|
||||
is_tracing = torch.jit.is_tracing() or isinstance(inputs_embeds, torch.fx.Proxy)
|
||||
|
||||
if attention_mask is not None:
|
||||
# 4d mask is passed through
|
||||
if len(attention_mask.shape) == 4:
|
||||
expected_shape = (input_shape[0], 1, input_shape[1], key_value_length)
|
||||
if tuple(attention_mask.shape) != expected_shape:
|
||||
raise ValueError(
|
||||
f"Incorrect 4D attention_mask shape: {tuple(attention_mask.shape)}; expected: {expected_shape}."
|
||||
)
|
||||
else:
|
||||
# if the 4D mask has correct shape - invert it and fill with negative infinity
|
||||
inverted_mask = 1.0 - attention_mask.to(inputs_embeds.dtype)
|
||||
attention_mask = inverted_mask.masked_fill(
|
||||
inverted_mask.to(torch.bool), torch.finfo(inputs_embeds.dtype).min
|
||||
)
|
||||
return attention_mask
|
||||
|
||||
elif not is_tracing and torch.all(attention_mask == 1):
|
||||
if query_length == 1:
|
||||
# For query_length == 1, causal attention and bi-directional attention are the same.
|
||||
attention_mask = None
|
||||
elif key_value_length == query_length:
|
||||
attention_mask = None
|
||||
else:
|
||||
# Unfortunately, for query_length > 1 and key_value_length != query_length, we cannot generally ignore the attention mask, as SDPA causal mask generation
|
||||
# may be wrong. We will set `is_causal=False` in SDPA and rely on Transformers attention_mask instead, hence not setting it to None here.
|
||||
# Reference: https://github.com/pytorch/pytorch/issues/108108
|
||||
pass
|
||||
elif query_length > 1 and key_value_length != query_length:
|
||||
# See the comment above (https://github.com/pytorch/pytorch/issues/108108).
|
||||
# Ugly: we set it to True here to dispatch in the following controlflow to `to_causal_4d`.
|
||||
attention_mask = True
|
||||
elif is_tracing:
|
||||
raise ValueError(
|
||||
'Attention using SDPA can not be traced with torch.jit.trace when no attention_mask is provided. To solve this issue, please either load your model with the argument `attn_implementation="eager"` or pass an attention_mask input when tracing the model.'
|
||||
)
|
||||
|
||||
if attention_mask is None:
|
||||
expanded_4d_mask = None
|
||||
elif attention_mask is True:
|
||||
expanded_4d_mask = attn_mask_converter.to_causal_4d(
|
||||
input_shape[0], input_shape[-1], key_value_length, dtype=inputs_embeds.dtype, device=inputs_embeds.device
|
||||
)
|
||||
else:
|
||||
expanded_4d_mask = attn_mask_converter.to_4d(
|
||||
attention_mask,
|
||||
input_shape[-1],
|
||||
dtype=inputs_embeds.dtype,
|
||||
key_value_length=key_value_length,
|
||||
)
|
||||
|
||||
# From PyTorch 2.1 onwards, F.scaled_dot_product_attention with the memory-efficient attention backend
|
||||
# produces nans if sequences are completely unattended in the attention mask. Details: https://github.com/pytorch/pytorch/issues/110213
|
||||
#
|
||||
# This fix is not applied in case we are tracing with torch.jit.trace or symbolic_trace, as _unmask_unattended has a data-dependent
|
||||
# controlflow that can not be captured properly.
|
||||
# TODO: _unmask_unattended does not work either with torch.compile when using fullgraph=True. We should find a way to detect this case.
|
||||
if query_length > 1 and not is_tracing:
|
||||
expanded_4d_mask = AttentionMaskConverter._unmask_unattended(
|
||||
expanded_4d_mask, attention_mask, unmasked_value=0.0
|
||||
)
|
||||
|
||||
return expanded_4d_mask
|
||||
|
||||
|
||||
def _prepare_4d_attention_mask(mask: torch.Tensor, dtype: torch.dtype, tgt_len: Optional[int] = None):
|
||||
"""
|
||||
Creates a non-causal 4D mask of shape `(batch_size, 1, query_length, key_value_length)` from a 2D mask of shape
|
||||
`(batch_size, key_value_length)`
|
||||
|
||||
Args:
|
||||
mask (`torch.Tensor` or `None`):
|
||||
A 2D attention mask of shape `(batch_size, key_value_length)`
|
||||
dtype (`torch.dtype`):
|
||||
The torch dtype the created mask shall have.
|
||||
tgt_len (`int`):
|
||||
The target length or query length the created mask shall have.
|
||||
"""
|
||||
return AttentionMaskConverter._expand_mask(mask=mask, dtype=dtype, tgt_len=tgt_len)
|
||||
|
||||
|
||||
def _prepare_4d_attention_mask_for_sdpa(mask: torch.Tensor, dtype: torch.dtype, tgt_len: Optional[int] = None):
|
||||
"""
|
||||
Creates a non-causal 4D mask of shape `(batch_size, 1, query_length, key_value_length)` from a 2D mask of shape
|
||||
`(batch_size, key_value_length)`
|
||||
|
||||
Args:
|
||||
mask (`torch.Tensor` or `None`):
|
||||
A 2D attention mask of shape `(batch_size, key_value_length)`
|
||||
dtype (`torch.dtype`):
|
||||
The torch dtype the created mask shall have.
|
||||
tgt_len (`int`):
|
||||
The target length or query length the created mask shall have.
|
||||
"""
|
||||
batch_size, key_value_length = mask.shape
|
||||
tgt_len = tgt_len if tgt_len is not None else key_value_length
|
||||
|
||||
# torch.jit.trace and torchdynamo with fullgraph=True are unable to capture the controlflow `is_causal=attention_mask is None and q_len > 1`
|
||||
# used as an SDPA argument. We keep compatibility with these tracing tools by always using SDPA's `attn_mask` argument in case we are tracing.
|
||||
# TODO: Fix this as well when using torchdynamo with fullgraph=True.
|
||||
is_tracing = torch.jit.is_tracing()
|
||||
|
||||
if torch.all(mask == 1):
|
||||
if is_tracing:
|
||||
pass
|
||||
elif tgt_len == 1:
|
||||
# For query_length == 1, causal attention and bi-directional attention are the same.
|
||||
return None
|
||||
elif key_value_length == tgt_len:
|
||||
return None
|
||||
else:
|
||||
# Unfortunately, for query_length > 1 and key_value_length != query_length, we can not generally ignore the attention mask, as SDPA causal mask generation
|
||||
# may be wrong. We will set is_causal=False in SDPA and rely on Transformers attention_mask instead, hence not setting it to None here.
|
||||
# Reference: https://github.com/pytorch/pytorch/issues/108108
|
||||
return AttentionMaskConverter._expand_mask(mask=mask, dtype=dtype, tgt_len=tgt_len)
|
||||
else:
|
||||
return AttentionMaskConverter._expand_mask(mask=mask, dtype=dtype, tgt_len=tgt_len)
|
||||
|
||||
|
||||
def _create_4d_causal_attention_mask(
|
||||
input_shape: Union[torch.Size, Tuple, List],
|
||||
dtype: torch.dtype,
|
||||
device: torch.device,
|
||||
past_key_values_length: int = 0,
|
||||
sliding_window: Optional[int] = None,
|
||||
) -> Optional[torch.Tensor]:
|
||||
"""
|
||||
Creates a causal 4D mask of shape `(batch_size, 1, query_length, key_value_length)`
|
||||
|
||||
Args:
|
||||
input_shape (`tuple(int)` or `list(int)` or `torch.Size`):
|
||||
The input shape should be a tuple that defines `(batch_size, query_length)`.
|
||||
dtype (`torch.dtype`):
|
||||
The torch dtype the created mask shall have.
|
||||
device (`int`):
|
||||
The torch device the created mask shall have.
|
||||
sliding_window (`int`, *optional*):
|
||||
If the model uses windowed attention, a sliding window should be passed.
|
||||
"""
|
||||
attn_mask_converter = AttentionMaskConverter(is_causal=True, sliding_window=sliding_window)
|
||||
|
||||
key_value_length = past_key_values_length + input_shape[-1]
|
||||
attention_mask = attn_mask_converter.to_causal_4d(
|
||||
input_shape[0], input_shape[-1], key_value_length, dtype=dtype, device=device
|
||||
)
|
||||
|
||||
return attention_mask
|
||||
1401
ixformer_sdk/train/speedformer/models/qwen2/modeling_qwen2.py
Normal file
1401
ixformer_sdk/train/speedformer/models/qwen2/modeling_qwen2.py
Normal file
File diff suppressed because it is too large
Load Diff
Reference in New Issue
Block a user