|
| 1 | +""" |
| 2 | +Liger FLCE for llama4 |
| 3 | +""" |
| 4 | + |
| 5 | +import sys |
| 6 | +from typing import List, Optional, Tuple, Union |
| 7 | + |
| 8 | +import torch |
| 9 | +from liger_kernel.transformers.model.loss_utils import LigerForCausalLMLoss |
| 10 | +from transformers.modeling_outputs import CausalLMOutputWithPast |
| 11 | + |
| 12 | + |
| 13 | +def lce_forward( |
| 14 | + self, |
| 15 | + input_ids: torch.LongTensor = None, |
| 16 | + attention_mask: Optional[torch.Tensor] = None, |
| 17 | + position_ids: Optional[torch.LongTensor] = None, |
| 18 | + past_key_values: Optional[ |
| 19 | + Union["Cache", List[torch.FloatTensor]] # noqa: F821 |
| 20 | + ] = None, |
| 21 | + inputs_embeds: Optional[torch.FloatTensor] = None, |
| 22 | + labels: Optional[torch.LongTensor] = None, |
| 23 | + use_cache: Optional[bool] = None, |
| 24 | + output_attentions: Optional[bool] = None, |
| 25 | + output_hidden_states: Optional[bool] = None, |
| 26 | + return_dict: Optional[bool] = None, |
| 27 | + cache_position: Optional[torch.LongTensor] = None, |
| 28 | + logits_to_keep: Union[int, torch.Tensor] = 0, |
| 29 | + **loss_kwargs, |
| 30 | +) -> Union[Tuple, CausalLMOutputWithPast]: |
| 31 | + r""" |
| 32 | + Args: |
| 33 | + labels (`torch.LongTensor` of shape `(batch_size, sequence_length)`, *optional*): |
| 34 | + Labels for computing the masked language modeling loss. Indices should either be in `[0, ..., |
| 35 | + config.vocab_size]` or -100 (see `input_ids` docstring). Tokens with indices set to `-100` are ignored |
| 36 | + (masked), the loss is only computed for the tokens with labels in `[0, ..., config.vocab_size]`. |
| 37 | +
|
| 38 | + logits_to_keep (`int` or `torch.Tensor`, *optional*): |
| 39 | + If an `int`, compute logits for the last `logits_to_keep` tokens. If `0`, calculate logits for all |
| 40 | + `input_ids` (special case). Only last token logits are needed for generation, and calculating them only for that |
| 41 | + token can save memory, which becomes pretty significant for long sequences or large vocabulary size. |
| 42 | + If a `torch.Tensor`, must be 1D corresponding to the indices to keep in the sequence length dimension. |
| 43 | + This is useful when using packed tensor format (single dimension for batch and sequence length). |
| 44 | +
|
| 45 | + Returns: |
| 46 | + """ |
| 47 | + |
| 48 | + # pylint: disable=duplicate-code |
| 49 | + output_attentions = ( |
| 50 | + output_attentions |
| 51 | + if output_attentions is not None |
| 52 | + else self.config.output_attentions |
| 53 | + ) |
| 54 | + output_hidden_states = ( |
| 55 | + output_hidden_states |
| 56 | + if output_hidden_states is not None |
| 57 | + else self.config.output_hidden_states |
| 58 | + ) |
| 59 | + return_dict = ( |
| 60 | + return_dict if return_dict is not None else self.config.use_return_dict |
| 61 | + ) |
| 62 | + |
| 63 | + # decoder outputs consists of (dec_features, layer_state, dec_hidden, dec_attn) |
| 64 | + outputs = self.model( |
| 65 | + input_ids=input_ids, |
| 66 | + attention_mask=attention_mask, |
| 67 | + position_ids=position_ids, |
| 68 | + past_key_values=past_key_values, |
| 69 | + inputs_embeds=inputs_embeds, |
| 70 | + use_cache=use_cache, |
| 71 | + output_attentions=output_attentions, |
| 72 | + output_hidden_states=output_hidden_states, |
| 73 | + return_dict=return_dict, |
| 74 | + cache_position=cache_position, |
| 75 | + ) |
| 76 | + |
| 77 | + hidden_states = outputs[0] |
| 78 | + |
| 79 | + if hasattr(self.config, "pretraining_tp") and self.config.pretraining_tp > 1: |
| 80 | + raise Exception( # pylint: disable=broad-exception-raised |
| 81 | + "Liger Kernel does not support pretraining_tp!!" |
| 82 | + ) |
| 83 | + |
| 84 | + logits = None |
| 85 | + loss = None |
| 86 | + # if in training mode, don't materialize logits |
| 87 | + if self.training and (labels is not None): |
| 88 | + loss = LigerForCausalLMLoss( |
| 89 | + hidden_states=hidden_states, |
| 90 | + lm_head_weight=self.lm_head.weight, |
| 91 | + labels=labels, |
| 92 | + hidden_size=self.config.hidden_size, |
| 93 | + **loss_kwargs, |
| 94 | + ) |
| 95 | + |
| 96 | + else: # if in inference mode materialize logits |
| 97 | + slice_indices = ( |
| 98 | + slice(-logits_to_keep, None) |
| 99 | + if isinstance(logits_to_keep, int) |
| 100 | + else logits_to_keep |
| 101 | + ) |
| 102 | + logits = self.lm_head(hidden_states[:, slice_indices, :]) |
| 103 | + if labels is not None: |
| 104 | + loss = self.loss_function( |
| 105 | + logits=logits, |
| 106 | + labels=labels, |
| 107 | + vocab_size=self.config.vocab_size, |
| 108 | + **loss_kwargs, |
| 109 | + ) |
| 110 | + |
| 111 | + if not return_dict: |
| 112 | + output = (logits,) + outputs[1:] |
| 113 | + return (loss,) + output if loss is not None else output |
| 114 | + |
| 115 | + return CausalLMOutputWithPast( |
| 116 | + loss=loss, |
| 117 | + logits=logits, |
| 118 | + past_key_values=outputs.past_key_values, |
| 119 | + hidden_states=outputs.hidden_states, |
| 120 | + attentions=outputs.attentions, |
| 121 | + ) |
| 122 | + |
| 123 | + |
| 124 | +def apply_liger_kernel_to_llama4( |
| 125 | + cross_entropy: bool = False, |
| 126 | + fused_linear_cross_entropy: bool = False, |
| 127 | + rms_norm: bool = False, |
| 128 | + glu_activation: bool = False, |
| 129 | + layer_norm: bool = False, |
| 130 | + **kwargs, # pylint: disable=unused-argument |
| 131 | +) -> None: |
| 132 | + """ |
| 133 | + Apply Liger kernels to replace original implementation in HuggingFace Llama models (2 and 3) |
| 134 | +
|
| 135 | + Args: |
| 136 | + cross_entropy (bool): Whether to apply Liger's cross entropy loss. Default is False. |
| 137 | + fused_linear_cross_entropy (bool): |
| 138 | + Whether to apply Liger's fused linear cross entropy loss. Default is False. |
| 139 | + `cross_entropy` and `fused_linear_cross_entropy` cannot both be False. |
| 140 | + If `fused_linear_cross_entropy` is True, the logits will not be materialized but more memory efficient. |
| 141 | + rms_norm (bool): Whether to apply Liger's RMSNorm. Default is False. |
| 142 | + glu_activation (bool): Whether to apply Liger's SwiGLU MLP. Default is False. |
| 143 | + layer_norm (bool): Whether to apply Liger's LayerNorm. Default is False. |
| 144 | + """ |
| 145 | + |
| 146 | + import transformers.models.llama4.modeling_llama4 # noqa: F401 # pylint: disable=unused-import |
| 147 | + from liger_kernel.transformers.functional import liger_cross_entropy |
| 148 | + from liger_kernel.transformers.layer_norm import LigerLayerNorm |
| 149 | + from liger_kernel.transformers.rms_norm import LigerRMSNorm |
| 150 | + from liger_kernel.transformers.swiglu import LigerSwiGLUMLP |
| 151 | + |
| 152 | + assert not ( |
| 153 | + cross_entropy and fused_linear_cross_entropy |
| 154 | + ), "cross_entropy and fused_linear_cross_entropy cannot both be True." |
| 155 | + |
| 156 | + modeling_llama4 = sys.modules["transformers.models.llama4.modeling_llama4"] |
| 157 | + |
| 158 | + if rms_norm: |
| 159 | + modeling_llama4.Llama4TextRMSNorm = LigerRMSNorm |
| 160 | + if glu_activation: |
| 161 | + modeling_llama4.Llama4TextMLP = LigerSwiGLUMLP |
| 162 | + if layer_norm: |
| 163 | + modeling_llama4.nn.LayerNorm = LigerLayerNorm |
| 164 | + |
| 165 | + if cross_entropy: |
| 166 | + from transformers.loss.loss_utils import nn |
| 167 | + |
| 168 | + nn.functional.cross_entropy = liger_cross_entropy |
| 169 | + |
| 170 | + if fused_linear_cross_entropy: |
| 171 | + modeling_llama4.Llama4ForCausalLM.forward = lce_forward |
0 commit comments