ceos/unsloth
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions

Fix Llama and Gemma inference (#3034)

Browse source 
* Fix Llama and Gemma inference

* Add simple quality life for CUDA link error (which is not captured since we bypass all error)
This commit is contained in:
Edd 2025-07-25 14:38:20 +08:00 • committed by GitHub
commit bb37fb71c6
2 changed files with 6 additions and 6 deletions
Download patch file Download diff file Expand all files Collapse all files

2
unsloth/models/falcon_h1.py
View file

@ -51,7 +51,7 @@ try:
from transformers.models.falcon_h1.modeling_falcon_h1 import (
FalconH1Attention,
)
except:
except ModuleNotFoundError:
# if we are on a old version of transformers technically it should fail in the try except above
# but if somehow we make it here, we need to raise an error since FalconH1Attention is not available
# or renamed

10
unsloth/models/llama.py
View file

@ -780,7 +780,7 @@ def LlamaModel_fast_forward(
# Fix up attention mask by setting elements to 0
# Specifically for DPO
if getattr(self, "_has_no_labels", False) is True and (attention_mask is not None) and (past_key_values is None) and \
(not train_embed_tokens):
(not train_embed_tokens) and self.training:
# Careful for inference the attention_mask is size (1, kv_seq_len)
# Whilst the input_embeds is size (1, 1, 4096)
inputs_requires_grad = inputs_embeds.requires_grad
@ -865,21 +865,21 @@ def LlamaModel_fast_forward(
# https://github.com/pytorch/pytorch/issues/103749
# Need to convert to float and not using bool
attention_mask = (1.0 - attention_mask.float()) * torch.finfo(inputs_embeds.dtype).min
# attention_mask = (1.0 - attention_mask.float()) * torch.finfo(inputs_embeds.dtype).min
dynamic_SWA_mask = _prepare_4d_causal_attention_mask_for_sdpa(
attention_mask,
(batch_size, seq_length),
inputs_embeds,
past_key_values_length,
sliding_window = self.config.sliding_window,
)[0][0]
)
dynamic_GA_mask = _prepare_4d_causal_attention_mask_for_sdpa(
attention_mask,
(batch_size, seq_length),
inputs_embeds,
past_key_values_length,
sliding_window = None,
)[0][0]
)
use_static_mask = False
elif not hasattr(self, "SWA_mask"):
@ -953,7 +953,7 @@ def LlamaModel_fast_forward(
else:
layer_outputs = decoder_layer(
hidden_states,
causal_mask=mask,
causal_mask = mask,
attention_mask = attention_mask,
position_ids = position_ids,
past_key_value = past_key_value,