Update llama.py

This commit is contained in:
Daniel Han 2024-12-20 03:03:26 -08:00
commit c9239c8fdf

View file

@ -784,7 +784,6 @@ def LlamaModel_fast_forward(
pass
pass
if transformers_version > "4.47.1" and hasattr(self, "rotary_emb"):
# Transformers main has made it mandatory to pass position_embeddings
# https://github.com/huggingface/transformers/pull/34858
@ -848,7 +847,7 @@ def LlamaModel_fast_forward(
output_attentions = output_attentions,
use_cache = use_cache,
padding_mask = padding_mask,
position_embeddings = position_embeddings
position_embeddings = position_embeddings,
)
hidden_states = layer_outputs[0]
pass
@ -1008,7 +1007,6 @@ def CausalLM_fast_forward(fast_forward_inference):
if not RETURN_LOGITS and HAS_CUT_CROSS_ENTROPY and labels is not None:
n_items = kwargs.get("num_items_in_batch", None) or kwargs.get("n_items", None)
print(n_items)
loss = fused_linear_cross_entropy(
hidden_states = hidden_states,
lm_weight = lm_head,
@ -1056,6 +1054,7 @@ def CausalLM_fast_forward(fast_forward_inference):
self.extra_ignored_labels = torch.full((self.max_seq_length, 1), -100, device = "cuda:0")
pass
shift_labels = torch.hstack((labels[..., 1:], self.extra_ignored_labels[:labels.shape[0]]))
print(kwargs.get("num_items_in_batch", None) or kwargs.get("n_items", None))
loss = fast_cross_entropy_loss(
logits = shift_logits,
labels = shift_labels,