|
| 1 | +from typing import List |
| 2 | +from typing import Optional |
| 3 | +from typing import Tuple |
| 4 | +from typing import Union |
| 5 | + |
| 6 | +import torch |
| 7 | + |
| 8 | +from transformers.modeling_outputs import BaseModelOutputWithPast |
| 9 | +from transformers.utils.deprecation import deprecate_kwarg |
| 10 | + |
| 11 | +from liger_kernel.transformers.model.loss_utils import LigerForCausalLMLoss |
| 12 | +from liger_kernel.transformers.model.loss_utils import unpack_cross_entropy_result |
| 13 | +from liger_kernel.transformers.model.output_classes import LigerCausalLMOutputWithPast |
| 14 | + |
| 15 | + |
| 16 | +@deprecate_kwarg("num_logits_to_keep", version="4.50", new_name="logits_to_keep") |
| 17 | +def lce_forward( |
| 18 | + self, |
| 19 | + input_ids: torch.LongTensor = None, |
| 20 | + attention_mask: Optional[torch.Tensor] = None, |
| 21 | + position_ids: Optional[torch.LongTensor] = None, |
| 22 | + past_key_values: Optional[List[torch.FloatTensor]] = None, |
| 23 | + inputs_embeds: Optional[torch.FloatTensor] = None, |
| 24 | + labels: Optional[torch.LongTensor] = None, |
| 25 | + use_cache: Optional[bool] = None, |
| 26 | + output_attentions: Optional[bool] = None, |
| 27 | + output_hidden_states: Optional[bool] = None, |
| 28 | + return_dict: Optional[bool] = None, |
| 29 | + cache_position: Optional[torch.LongTensor] = None, |
| 30 | + logits_to_keep: Union[int, torch.Tensor] = 0, |
| 31 | + skip_logits: Optional[bool] = None, |
| 32 | + **kwargs, |
| 33 | +) -> Union[Tuple, LigerCausalLMOutputWithPast]: |
| 34 | + r""" |
| 35 | + Args: |
| 36 | + labels (`torch.LongTensor` of shape `(batch_size, sequence_length)`, *optional*): |
| 37 | + Labels for computing the masked language modeling loss. Indices should either be in `[0, ..., |
| 38 | + config.vocab_size]` or -100 (see `input_ids` docstring). Tokens with indices set to `-100` are ignored |
| 39 | + (masked), the loss is only computed for the tokens with labels in `[0, ..., config.vocab_size]`. |
| 40 | +
|
| 41 | + logits_to_keep (`int` or `torch.Tensor`, *optional*): |
| 42 | + If an `int`, compute logits for the last `logits_to_keep` tokens. If `0`, calculate logits for all |
| 43 | + `input_ids` (special case). Only last token logits are needed for generation, and calculating them only for that |
| 44 | + token can save memory, which becomes pretty significant for long sequences or large vocabulary size. |
| 45 | + If a `torch.Tensor`, must be 1D corresponding to the indices to keep in the sequence length dimension. |
| 46 | + This is useful when using packed tensor format (single dimension for batch and sequence length). |
| 47 | +
|
| 48 | + Returns: |
| 49 | +
|
| 50 | + Example: |
| 51 | +
|
| 52 | + ```python |
| 53 | + >>> from transformers import AutoTokenizer, Olmo3ForCausalLM |
| 54 | +
|
| 55 | + >>> model = Olmo3ForCausalLM.from_pretrained("allenai/Olmo-3-7B-Instruct") |
| 56 | + >>> tokenizer = AutoTokenizer.from_pretrained("allenai/Olmo-3-7B-Instruct") |
| 57 | +
|
| 58 | + >>> prompt = "Hey, are you conscious? Can you talk to me?" |
| 59 | + >>> inputs = tokenizer(prompt, return_tensors="pt") |
| 60 | +
|
| 61 | + >>> # Generate |
| 62 | + >>> generate_ids = model.generate(inputs.input_ids, max_length=30) |
| 63 | + >>> tokenizer.batch_decode(generate_ids, skip_special_tokens=True, clean_up_tokenization_spaces=False)[0] |
| 64 | + 'Hey, are you conscious? Can you talk to me?\nI’m not sure if you’re conscious of this, but I’m' |
| 65 | + ``` |
| 66 | + """ |
| 67 | + output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions |
| 68 | + output_hidden_states = ( |
| 69 | + output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states |
| 70 | + ) |
| 71 | + return_dict = return_dict if return_dict is not None else self.config.use_return_dict |
| 72 | + |
| 73 | + # decoder outputs consists of (dec_features, layer_state, dec_hidden, dec_attn) |
| 74 | + outputs: BaseModelOutputWithPast = self.model( |
| 75 | + input_ids=input_ids, |
| 76 | + attention_mask=attention_mask, |
| 77 | + position_ids=position_ids, |
| 78 | + past_key_values=past_key_values, |
| 79 | + inputs_embeds=inputs_embeds, |
| 80 | + use_cache=use_cache, |
| 81 | + output_attentions=output_attentions, |
| 82 | + output_hidden_states=output_hidden_states, |
| 83 | + return_dict=return_dict, |
| 84 | + cache_position=cache_position, |
| 85 | + **kwargs, |
| 86 | + ) |
| 87 | + |
| 88 | + hidden_states = outputs[0] |
| 89 | + # Only compute necessary logits, and do not upcast them to float if we are not computing the loss |
| 90 | + slice_indices = slice(-logits_to_keep, None) if isinstance(logits_to_keep, int) else logits_to_keep |
| 91 | + kept_hidden_states = hidden_states[:, slice_indices, :] |
| 92 | + |
| 93 | + shift_labels = kwargs.pop("shift_labels", None) |
| 94 | + logits = None |
| 95 | + loss = None |
| 96 | + token_accuracy = None |
| 97 | + |
| 98 | + if skip_logits and labels is None and shift_labels is None: |
| 99 | + raise ValueError("skip_logits is True, but labels and shift_labels are None") |
| 100 | + |
| 101 | + if skip_logits is None: |
| 102 | + # By default, if in training mode, don't materialize logits |
| 103 | + skip_logits = self.training and (labels is not None or shift_labels is not None) |
| 104 | + |
| 105 | + # Compute loss |
| 106 | + if skip_logits: |
| 107 | + result = LigerForCausalLMLoss( |
| 108 | + hidden_states=kept_hidden_states, |
| 109 | + lm_head_weight=self.lm_head.weight, |
| 110 | + labels=labels, |
| 111 | + shift_labels=shift_labels, |
| 112 | + hidden_size=self.config.hidden_size, |
| 113 | + **kwargs, |
| 114 | + ) |
| 115 | + loss, _, token_accuracy = unpack_cross_entropy_result(result) |
| 116 | + |
| 117 | + else: |
| 118 | + logits = self.lm_head(kept_hidden_states) |
| 119 | + if labels is not None or shift_labels is not None: |
| 120 | + loss = self.loss_function( |
| 121 | + logits=logits, |
| 122 | + labels=labels, |
| 123 | + shift_labels=shift_labels, |
| 124 | + vocab_size=self.config.vocab_size, |
| 125 | + **kwargs, |
| 126 | + ) |
| 127 | + |
| 128 | + if not return_dict: |
| 129 | + output = (logits,) + outputs[1:] |
| 130 | + output = ((loss,) + output) if loss is not None else output |
| 131 | + output = output + (token_accuracy,) if token_accuracy is not None else output |
| 132 | + return output |
| 133 | + |
| 134 | + # Return custom output class with token_accuracy field |
| 135 | + return LigerCausalLMOutputWithPast( |
| 136 | + loss=loss, |
| 137 | + logits=logits, |
| 138 | + past_key_values=outputs.past_key_values, |
| 139 | + hidden_states=outputs.hidden_states, |
| 140 | + attentions=outputs.attentions, |
| 141 | + token_accuracy=token_accuracy, |
| 142 | + ) |
0 commit comments