Fix bugs (#779)
* Update __init__.py * dynamic RoPE * Update mistral.py * Update llama.py * Update tokenizer_utils.py * Update mistral.py * Update llama.py * Update __init__.py * Update flex_attention.py
This commit is contained in:
Daniel Han
GitHub
parent
7fb3838c00
commit
3fd18f7707
5 changed files with 69 additions and 24 deletions
|
|
@ -89,7 +89,7 @@ if (major_torch == 2) and (minor_torch >= 5):
|
|||
return old_is_bf16_supported(including_emulation)
|
||||
torch.cuda.is_bf16_supported = is_bf16_supported
|
||||
else:
|
||||
def is_bf16_supported(): SUPPORTS_BFLOAT16
|
||||
def is_bf16_supported(): return SUPPORTS_BFLOAT16
|
||||
torch.cuda.is_bf16_supported = is_bf16_supported
|
||||
pass
|
||||
|
||||
|
|
|
|||
|
|
@ -25,18 +25,23 @@ torch_compile_options = {
|
|||
}
|
||||
|
||||
# Flex Attention supported from torch 2.5 onwards only
|
||||
import torch.nn.attention
|
||||
if hasattr(torch.nn.attention, "flex_attention"):
|
||||
import torch.nn.attention.flex_attention
|
||||
from torch.nn.attention.flex_attention import flex_attention
|
||||
from torch.nn.attention.flex_attention import create_block_mask
|
||||
FLEX_ATTENTION_PADDING = getattr(
|
||||
torch.nn.attention.flex_attention,
|
||||
"_DEFAULT_SPARSE_BLOCK_SIZE",
|
||||
1,
|
||||
)
|
||||
flex_attention = torch.compile(flex_attention, dynamic = False)
|
||||
HAS_FLEX_ATTENTION = True
|
||||
import torch.nn
|
||||
if hasattr(torch.nn, "attention"):
|
||||
import torch.nn.attention
|
||||
if hasattr(torch.nn.attention, "flex_attention"):
|
||||
import torch.nn.attention.flex_attention
|
||||
from torch.nn.attention.flex_attention import flex_attention
|
||||
from torch.nn.attention.flex_attention import create_block_mask
|
||||
FLEX_ATTENTION_PADDING = getattr(
|
||||
torch.nn.attention.flex_attention,
|
||||
"_DEFAULT_SPARSE_BLOCK_SIZE",
|
||||
1,
|
||||
)
|
||||
flex_attention = torch.compile(flex_attention, dynamic = False)
|
||||
HAS_FLEX_ATTENTION = True
|
||||
else:
|
||||
HAS_FLEX_ATTENTION = False
|
||||
pass
|
||||
else:
|
||||
HAS_FLEX_ATTENTION = False
|
||||
pass
|
||||
|
|
|
|||
|
|
@ -158,6 +158,14 @@ def LlamaAttention_fast_forward_inference(
|
|||
self.temp_QA = torch.empty((2, bsz, 1, attention_size), dtype = dtype, device = "cuda:0")
|
||||
self.temp_KV = torch.empty((2, bsz, 1, n_kv_heads*head_dim), dtype = dtype, device = "cuda:0")
|
||||
self.RH_Q = torch.empty((bsz, n_heads, 1, head_dim), dtype = dtype, device = "cuda:0")
|
||||
|
||||
# Mistral Nemo 12b has weird dimensions
|
||||
if attention_size != self.hidden_size:
|
||||
self.temp_O = torch.empty((1, bsz, self.hidden_size), dtype = dtype, device = "cuda:0")
|
||||
else:
|
||||
self.temp_O = self.temp_QA[1][:,:,:self.hidden_size]
|
||||
pass
|
||||
|
||||
self.attention = torch.empty((bsz, n_heads, 1, KV_CACHE_INCREMENT+seq_len), dtype = dtype, device = "cuda:0")
|
||||
self.scalar = 1.0 / math_sqrt(self.head_dim)
|
||||
self.half_head_dim = head_dim // 2
|
||||
|
|
@ -239,7 +247,7 @@ def LlamaAttention_fast_forward_inference(
|
|||
pass
|
||||
A = A.transpose(1, 2)
|
||||
A = A.reshape(bsz, 1, attention_size)
|
||||
A = fast_linear_forward(self.o_proj, A, out = self.temp_QA[1][:,:,:self.hidden_size])
|
||||
A = fast_linear_forward(self.o_proj, A, out = self.temp_O)
|
||||
return A, (Kn, Vn)
|
||||
pass
|
||||
|
||||
|
|
@ -335,6 +343,9 @@ def LlamaAttention_fast_forward(
|
|||
if past_key_value is not None:
|
||||
kv_seq_len += past_key_value[0].shape[-2]
|
||||
|
||||
# Extend RoPE dynamically to fit in VRAM
|
||||
self.rotary_emb.extend_rope_embedding(V, seq_len = kv_seq_len)
|
||||
|
||||
if position_ids is None:
|
||||
cos = self.rotary_emb.cos_cached
|
||||
sin = self.rotary_emb.sin_cached
|
||||
|
|
@ -971,19 +982,21 @@ class LlamaRotaryEmbedding(torch.nn.Module):
|
|||
self.dim = dim
|
||||
self.max_position_embeddings = max_position_embeddings
|
||||
self.base = base
|
||||
# Dynamic RoPE we first set it to a max of 4 * 8192 tokens then we iteratively grow this
|
||||
self.current_rope_size = min(4 * 8192, self.max_position_embeddings)
|
||||
|
||||
# Build here to make `torch.jit.trace` work.
|
||||
self._set_cos_sin_cache(seq_len=max_position_embeddings, device=device, dtype=torch.get_default_dtype())
|
||||
self._set_cos_sin_cache(seq_len=self.current_rope_size, device=device, dtype=torch.get_default_dtype())
|
||||
pass
|
||||
|
||||
def _set_cos_sin_cache(self, seq_len, device, dtype):
|
||||
# Note: on the original Llama codebase, these tensors are created on the target device (and not on CPU) and
|
||||
# in FP32. They are applied (multiplied) in FP32 as well.
|
||||
self.max_seq_len_cached = seq_len
|
||||
self.current_rope_size = seq_len
|
||||
inv_freq = 1.0 / (
|
||||
self.base ** (torch.arange(0, self.dim, 2, dtype=torch.int64, device="cpu").float() / self.dim)
|
||||
)
|
||||
t = torch.arange(self.max_seq_len_cached, device="cpu", dtype=torch.int64).float()
|
||||
t = torch.arange(self.current_rope_size, device="cpu", dtype=torch.int64).float()
|
||||
|
||||
freqs = torch.outer(t, inv_freq)
|
||||
# Different from paper, but it uses a different permutation in order to obtain the same calculation
|
||||
|
|
@ -994,14 +1007,21 @@ class LlamaRotaryEmbedding(torch.nn.Module):
|
|||
|
||||
def forward(self, x, position_ids=None, seq_len=None):
|
||||
# x: [bs, num_attention_heads, seq_len, head_size]
|
||||
if seq_len > self.max_seq_len_cached:
|
||||
if seq_len > self.current_rope_size:
|
||||
self._set_cos_sin_cache(seq_len=seq_len, device=x.device, dtype=x.dtype)
|
||||
|
||||
return (
|
||||
self.cos_cached[:seq_len].to(dtype=x.dtype),
|
||||
self.sin_cached[:seq_len].to(dtype=x.dtype),
|
||||
self.cos_cached[:seq_len].to(dtype = x.dtype),
|
||||
self.sin_cached[:seq_len].to(dtype = x.dtype),
|
||||
)
|
||||
pass
|
||||
|
||||
def extend_rope_embedding(self, x, seq_len):
|
||||
if seq_len <= self.current_rope_size: return
|
||||
# Iteratively grow by increments of 8192
|
||||
self.current_rope_size = int(round(seq_len / 8192)) * 8192
|
||||
self._set_cos_sin_cache(self.current_rope_size, device = "cuda:0", dtype = x.dtype)
|
||||
pass
|
||||
pass
|
||||
|
||||
|
||||
|
|
@ -1016,11 +1036,11 @@ class LlamaLinearScalingRotaryEmbedding(LlamaRotaryEmbedding):
|
|||
pass
|
||||
|
||||
def _set_cos_sin_cache(self, seq_len, device, dtype):
|
||||
self.max_seq_len_cached = seq_len
|
||||
self.current_rope_size = seq_len
|
||||
inv_freq = 1.0 / (
|
||||
self.base ** (torch.arange(0, self.dim, 2, dtype=torch.int64, device="cpu").float() / self.dim)
|
||||
)
|
||||
t = torch.arange(self.max_seq_len_cached, device="cpu", dtype=torch.int64).float()
|
||||
t = torch.arange(self.current_rope_size, device="cpu", dtype=torch.int64).float()
|
||||
t = t / self.scaling_factor
|
||||
|
||||
freqs = torch.outer(t, inv_freq)
|
||||
|
|
@ -1140,6 +1160,12 @@ class FastLlamaModel:
|
|||
f"\ / Bfloat16 = {str(SUPPORTS_BFLOAT16).upper()}. FA [Xformers = {xformers_version}. FA2 = {HAS_FLASH_ATTENTION}]\n"\
|
||||
f' "-____-" Free Apache license: http://github.com/unslothai/unsloth'
|
||||
print(statistics)
|
||||
|
||||
# Warn about fast transfers
|
||||
if os.environ.get("HF_HUB_ENABLE_HF_TRANSFER", "0") == "1":
|
||||
logger.warning_once("Unsloth: Fast downloading is enabled - ignore downloading bars which are red colored!")
|
||||
pass
|
||||
|
||||
model_patcher.pre_patch()
|
||||
get_statistics() # For debugging - we use a download counter to see if environments are not breaking
|
||||
|
||||
|
|
|
|||
|
|
@ -78,6 +78,9 @@ def MistralAttention_fast_forward(
|
|||
if past_key_value is not None:
|
||||
kv_seq_len += past_key_value[0].shape[-2]
|
||||
|
||||
# Extend RoPE dynamically to fit in VRAM
|
||||
self.rotary_emb.extend_rope_embedding(V, seq_len = kv_seq_len)
|
||||
|
||||
if position_ids is None:
|
||||
cos = self.rotary_emb.cos_cached
|
||||
sin = self.rotary_emb.sin_cached
|
||||
|
|
@ -158,7 +161,7 @@ def MistralAttention_fast_forward(
|
|||
A = A.transpose(1, 2).contiguous()
|
||||
pass
|
||||
|
||||
attn_output = A.reshape(bsz, q_len, self.hidden_size)
|
||||
attn_output = A.reshape(bsz, q_len, n_heads*head_dim)
|
||||
attn_output = self.apply_o(self, attn_output)
|
||||
attn_weights = None
|
||||
return attn_output, attn_weights, past_key_value
|
||||
|
|
|
|||
|
|
@ -38,6 +38,17 @@ __all__ = [
|
|||
IGNORED_TOKENIZER_CHECKING = frozenset((
|
||||
"CodeLlamaTokenizerFast",
|
||||
"CodeLlamaTokenizer",
|
||||
""
|
||||
))
|
||||
|
||||
|
||||
IGNORED_TOKENIZER_NAMES = frozenset((
|
||||
"unsloth/Mistral-Nemo-Instruct-2407-bnb-4bit",
|
||||
"unsloth/Mistral-Nemo-Instruct-2407",
|
||||
"mistralai/Mistral-Nemo-Instruct-2407",
|
||||
"unsloth/Mistral-Nemo-Base-2407-bnb-4bit",
|
||||
"unsloth/Mistral-Nemo-Base-2407",
|
||||
"mistralai/Mistral-Nemo-Base-2407",
|
||||
))
|
||||
|
||||
# Check environments
|
||||
|
|
@ -488,7 +499,7 @@ def load_correct_tokenizer(
|
|||
cache_dir = cache_dir,
|
||||
)
|
||||
|
||||
if slow_tokenizer is not None:
|
||||
if tokenizer_name not in IGNORED_TOKENIZER_NAMES and slow_tokenizer is not None:
|
||||
if hasattr(fast_tokenizer, "add_bos_token") and hasattr(slow_tokenizer, "add_bos_token"):
|
||||
fast_tokenizer.add_bos_token = slow_tokenizer.add_bos_token
|
||||
if hasattr(fast_tokenizer, "add_eos_token") and hasattr(slow_tokenizer, "add_eos_token"):
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue