liger-kernel 0.0.0__py3-none-any.whl

liger_kernel/transformers/model/llama.py

@@ -0,0 +1,143 @@
+from typing import List, Optional, Tuple, Union
+
+import torch
+import torch.nn.functional as F
+from torch.nn import CrossEntropyLoss
+from transformers.modeling_outputs import CausalLMOutputWithPast
+from transformers.models.llama.modeling_llama import (
+    _CONFIG_FOR_DOC,
+    LLAMA_INPUTS_DOCSTRING,
+)
+from transformers.utils import (
+    add_start_docstrings_to_model_forward,
+    replace_return_docstrings,
+)
+
+from liger_kernel.transformers.fused_linear_cross_entropy import (
+    LigerFusedLinearCrossEntropyLoss,
+)
+
+
+@add_start_docstrings_to_model_forward(LLAMA_INPUTS_DOCSTRING)
+@replace_return_docstrings(
+    output_type=CausalLMOutputWithPast, config_class=_CONFIG_FOR_DOC
+)
+def lce_forward(
+    self,
+    input_ids: torch.LongTensor = None,
+    attention_mask: Optional[torch.Tensor] = None,
+    position_ids: Optional[torch.LongTensor] = None,
+    past_key_values: Optional[List[torch.FloatTensor]] = None,
+    inputs_embeds: Optional[torch.FloatTensor] = None,
+    labels: Optional[torch.LongTensor] = None,
+    use_cache: Optional[bool] = None,
+    output_attentions: Optional[bool] = None,
+    output_hidden_states: Optional[bool] = None,
+    return_dict: Optional[bool] = None,
+    cache_position: Optional[torch.LongTensor] = None,
+) -> Union[Tuple, CausalLMOutputWithPast]:
+    r"""
+    Args:
+        labels (`torch.LongTensor` of shape `(batch_size, sequence_length)`, *optional*):
+            Labels for computing the masked language modeling loss. Indices should either be in `[0, ...,
+            config.vocab_size]` or -100 (see `input_ids` docstring). Tokens with indices set to `-100` are ignored
+            (masked), the loss is only computed for the tokens with labels in `[0, ..., config.vocab_size]`.
+
+    Returns:
+
+    Example:
+
+    ```python
+    >>> from transformers import AutoTokenizer, LlamaForCausalLM
+
+    >>> model = LlamaForCausalLM.from_pretrained("meta-llama/Llama-2-7b-hf")
+    >>> tokenizer = AutoTokenizer.from_pretrained("meta-llama/Llama-2-7b-hf")
+
+    >>> prompt = "Hey, are you conscious? Can you talk to me?"
+    >>> inputs = tokenizer(prompt, return_tensors="pt")
+
+    >>> # Generate
+    >>> generate_ids = model.generate(inputs.input_ids, max_length=30)
+    >>> tokenizer.batch_decode(generate_ids, skip_special_tokens=True, clean_up_tokenization_spaces=False)[0]
+    "Hey, are you conscious? Can you talk to me?\nI'm not conscious, but I can talk to you."
+    ```"""
+    output_attentions = (
+        output_attentions
+        if output_attentions is not None
+        else self.config.output_attentions
+    )
+    output_hidden_states = (
+        output_hidden_states
+        if output_hidden_states is not None
+        else self.config.output_hidden_states
+    )
+    return_dict = (
+        return_dict if return_dict is not None else self.config.use_return_dict
+    )
+
+    # decoder outputs consists of (dec_features, layer_state, dec_hidden, dec_attn)
+    outputs = self.model(
+        input_ids=input_ids,
+        attention_mask=attention_mask,
+        position_ids=position_ids,
+        past_key_values=past_key_values,
+        inputs_embeds=inputs_embeds,
+        use_cache=use_cache,
+        output_attentions=output_attentions,
+        output_hidden_states=output_hidden_states,
+        return_dict=return_dict,
+        cache_position=cache_position,
+    )
+
+    hidden_states = outputs[0]
+
+    loss = None
+    logits = None
+
+    if self.training:
+        shift_hidden_states = hidden_states[..., :-1, :].contiguous()
+        shift_labels = labels[..., 1:].contiguous()
+
+        # flatten tokens
+        shift_hidden_states = shift_hidden_states.view(-1, self.config.hidden_size)
+        shift_labels = shift_labels.view(-1)
+
+        lce = LigerFusedLinearCrossEntropyLoss()
+        loss = lce(self.lm_head.weight, shift_hidden_states, shift_labels)
+
+    else:
+        if self.config.pretraining_tp > 1:
+            lm_head_slices = self.lm_head.weight.split(
+                self.vocab_size // self.config.pretraining_tp, dim=0
+            )
+            logits = [
+                F.linear(hidden_states, lm_head_slices[i])
+                for i in range(self.config.pretraining_tp)
+            ]
+            logits = torch.cat(logits, dim=-1)
+        else:
+            logits = self.lm_head(hidden_states)
+        logits = logits.float()
+        if labels is not None:
+            # Shift so that tokens < n predict n
+            shift_logits = logits[..., :-1, :].contiguous()
+            shift_labels = labels[..., 1:].contiguous()
+            # Flatten the tokens
+            loss_fct = CrossEntropyLoss()
+            shift_logits = shift_logits.view(-1, self.config.vocab_size)
+            shift_labels = shift_labels.view(-1)
+            # Enable model parallelism
+            shift_labels = shift_labels.to(shift_logits.device)
+            loss = loss_fct(shift_logits, shift_labels)
+
+    if not return_dict:
+        output = (logits,) + outputs[1:]
+        return (loss,) + output if loss is not None else output
+
+    return CausalLMOutputWithPast(
+        loss=loss,
+        logits=logits,
+        past_key_values=outputs.past_key_values,
+        hidden_states=outputs.hidden_states,
+        attentions=outputs.attentions,
+    )

liger_kernel/transformers/monkey_patch.py

@@ -0,0 +1,103 @@
+from liger_kernel.transformers.cross_entropy import LigerCrossEntropyLoss
+from liger_kernel.transformers.model.llama import lce_forward
+from liger_kernel.transformers.rms_norm import LigerRMSNorm
+from liger_kernel.transformers.rope import liger_rotary_pos_emb
+from liger_kernel.transformers.swiglu import LigerBlockSparseTop2MLP, LigerSwiGLUMLP
+
+
+# TODO: probably rename utils.py as hf_patcher.py to be more descriptive
+def apply_liger_kernel_to_llama(
+    rope: bool = True,
+    cross_entropy: bool = True,
+    fused_linear_cross_entropy: bool = False,
+    rms_norm: bool = True,
+    swiglu: bool = True,
+) -> None:
+    """
+    Apply Liger kernels to replace original implementation in HuggingFace Llama models (2 and 3)
+    to make GPU go burrr.
+
+    Args:
+        rope (bool): Whether to apply Liger's rotary position embedding. Default is True.
+        cross_entropy (bool): Whether to apply Liger's cross entropy loss. Default is True.
+        fused_linear_cross_entropy (bool):
+            Whether to apply Liger's fused lienar cross entropy loss. Default is False.
+            `cross_entropy` and `fused_linear_cross_entropy` cannot both be True.
+            If `fused_linear_cross_entropy` is True, the logits will not be materialized but more memory efficient.
+        rms_norm (bool): Whether to apply Liger's RMSNorm. Default is True.
+        swiglu (bool): Whether to apply Liger's SwiGLU MLP. Default is True.
+    """
+
+    assert not (
+        cross_entropy and fused_linear_cross_entropy
+    ), "cross_entropy and fused_linear_cross_entropy cannot both be True."
+
+    from transformers.models.llama import modeling_llama
+
+    if rope:
+        modeling_llama.apply_rotary_pos_emb = liger_rotary_pos_emb
+    if rms_norm:
+        modeling_llama.LlamaRMSNorm = LigerRMSNorm
+    if swiglu:
+        modeling_llama.LlamaMLP = LigerSwiGLUMLP
+    if cross_entropy:
+        modeling_llama.CrossEntropyLoss = LigerCrossEntropyLoss
+    if fused_linear_cross_entropy:
+        modeling_llama.LlamaForCausalLM.forward = lce_forward
+
+
+def apply_liger_kernel_to_mistral(
+    rope: bool = True,
+    cross_entropy: bool = True,
+    rms_norm: bool = True,
+    swiglu: bool = True,
+) -> None:
+    """
+    Apply Liger kernels to replace original implementation in HuggingFace Mistral models
+    to make GPU go burrr.
+
+    Args:
+        rope (bool): Whether to apply Liger's rotary position embedding. Default is True.
+        cross_entropy (bool): Whether to apply Liger's cross entropy loss. Default is True.
+        rms_norm (bool): Whether to apply Liger's RMSNorm. Default is True.
+        swiglu (bool): Whether to apply Liger's SwiGLU MLP. Default is True.
+    """
+
+    from transformers.models.mistral import modeling_mistral
+
+    if rope:
+        modeling_mistral.apply_rotary_pos_emb = liger_rotary_pos_emb
+    if rms_norm:
+        modeling_mistral.MistralRMSNorm = LigerRMSNorm
+    if cross_entropy:
+        modeling_mistral.CrossEntropyLoss = LigerCrossEntropyLoss
+    if swiglu:
+        modeling_mistral.MistralMLP = LigerSwiGLUMLP
+
+
+def apply_liger_kernel_to_mixtral(
+    rope: bool = True,
+    cross_entropy: bool = True,
+    rms_norm: bool = True,
+    swiglu: bool = True,
+) -> None:
+    """
+    Apply Liger kernels to replace original implementation in HuggingFace Mixtral models
+    to make GPU go burrr.
+
+    Args:
+        rope (bool): Whether to apply Liger's rotary position embedding. Default is True.
+        cross_entropy (bool): Whether to apply Liger's cross entropy loss. Default is True.
+        rms_norm (bool): Whether to apply Liger's RMSNorm. Default is True.
+    """
+
+    from transformers.models.mixtral import modeling_mixtral
+
+    if rope:
+        modeling_mixtral.apply_rotary_pos_emb = liger_rotary_pos_emb
+    if rms_norm:
+        modeling_mixtral.MistralRMSNorm = LigerRMSNorm
+    if cross_entropy:
+        modeling_mixtral.CrossEntropyLoss = LigerCrossEntropyLoss
+    if swiglu:
+        modeling_mixtral.MixtralBlockSparseTop2MLP = LigerBlockSparseTop2MLP

liger_kernel/transformers/rms_norm.py

@@ -0,0 +1,16 @@
+import torch
+import torch.nn as nn
+
+from liger_kernel.ops.rms_norm import LigerRMSNormFunction
+
+
+class LigerRMSNorm(nn.Module):
+    def __init__(self, hidden_size, eps=1e-6):
+        super().__init__()
+        self.weight = nn.Parameter(torch.ones(hidden_size))
+        self.variance_epsilon = eps
+
+    def forward(self, hidden_states):
+        return LigerRMSNormFunction.apply(
+            hidden_states, self.weight, self.variance_epsilon
+        )

liger_kernel/transformers/rope.py

@@ -0,0 +1,20 @@
+from liger_kernel.ops.rope import LigerRopeFunction
+
+
+def liger_rotary_pos_emb(q, k, cos, sin, position_ids=None, unsqueeze_dim=1):
+    """
+    Applies Rotary Positional Embedding (RoPE) operation to query and key states.
+
+    Args:
+        q (torch.Tensor): The query tensor of shape (bsz, n_q_head, seq_len, head_dim).
+        k (torch.Tensor): The key tensor of shape (bsz, n_kv_head, seq_len, head_dim).
+        cos (torch.Tensor): The cosine tensor of shape (1, seq_len, head_dim).
+        sin (torch.Tensor): The sine tensor of shape (1, seq_len, head_dim).
+        position_ids (torch.Tensor, optional): The position ids tensor. Defaults to None.
+        unsqueeze_dim (int, optional): The dimension to unsqueeze. Defaults to 1.
+
+    Returns:
+        Tuple[torch.Tensor, torch.Tensor]: The query and key tensors after applying the RoPE operation.
+    """
+
+    return LigerRopeFunction.apply(q, k, cos, sin, position_ids, unsqueeze_dim)

liger_kernel/transformers/swiglu.py

@@ -0,0 +1,40 @@
+import torch.nn as nn
+
+from liger_kernel.ops.swiglu import LigerSiLUMulFunction
+
+
+class LigerSwiGLUMLP(nn.Module):
+    def __init__(self, config):
+        super().__init__()
+        self.config = config
+        self.hidden_size = config.hidden_size
+        self.intermediate_size = config.intermediate_size
+        self.gate_proj = nn.Linear(self.hidden_size, self.intermediate_size, bias=False)
+        self.up_proj = nn.Linear(self.hidden_size, self.intermediate_size, bias=False)
+        self.down_proj = nn.Linear(self.intermediate_size, self.hidden_size, bias=False)
+        if config.hidden_act not in ["silu", "swish"]:
+            raise ValueError(f"Activation function {config.hidden_act} not supported.")
+
+    def forward(self, x):
+
+        return self.down_proj(
+            LigerSiLUMulFunction.apply(self.gate_proj(x), self.up_proj(x))
+        )
+
+
+class LigerBlockSparseTop2MLP(nn.Module):
+    def __init__(self, config):
+        super().__init__()
+        self.ffn_dim = config.intermediate_size
+        self.hidden_dim = config.hidden_size
+
+        self.w1 = nn.Linear(self.hidden_dim, self.ffn_dim, bias=False)
+        self.w2 = nn.Linear(self.ffn_dim, self.hidden_dim, bias=False)
+        self.w3 = nn.Linear(self.hidden_dim, self.ffn_dim, bias=False)
+
+        if config.hidden_act not in ["silu", "swish"]:
+            raise ValueError(f"Activation function {config.hidden_act} not supported.")
+
+    def forward(self, x):
+
+        return self.w2(LigerSiLUMulFunction.apply(self.w1(x), self.w3(x)))

liger_kernel/triton/__init__.py

@@ -0,0 +1,3 @@
+from liger_kernel.triton.monkey_patch import (  # noqa: F401
+    apply_liger_triton_cache_manager,
+)

liger_kernel/triton/monkey_patch.py

@@ -0,0 +1,44 @@
+import os
+import random
+
+from overrides import override
+from triton.runtime.cache import FileCacheManager
+
+
+class LigerTritonFileCacheManager(FileCacheManager):
+    @override
+    def put(self, data, filename, binary=True) -> str:
+        if not self.cache_dir:
+            raise RuntimeError("Could not create or locate cache dir")
+        binary = isinstance(data, bytes)
+        if not binary:
+            data = str(data)
+        assert self.lock_path is not None
+        filepath = self._make_path(filename)
+        # Random ID to avoid any collisions
+        rnd_id = random.randint(0, 1000000)
+        # we use the PID incase a bunch of these around so we can see what PID made it
+        pid = os.getpid()
+        # use temp dir to be robust against program interruptions
+        temp_dir = os.path.join(self.cache_dir, f"tmp.pid_{pid}_{rnd_id}")
+        os.makedirs(temp_dir, exist_ok=True)
+        temp_path = os.path.join(temp_dir, filename)
+
+        mode = "wb" if binary else "w"
+        with open(temp_path, mode) as f:
+            f.write(data)
+        # Replace is guaranteed to be atomic on POSIX systems if it succeeds
+        # so filepath cannot see a partial write
+        os.replace(temp_path, filepath)
+        os.removedirs(temp_dir)
+        return filepath
+
+
+def apply_liger_triton_cache_manager():
+    """
+    Experimental feature to get around transient FileNotFoundError in triton compilation.
+    For more details please see https://github.com/triton-lang/triton/pull/4295
+    """
+    os.environ["TRITON_CACHE_MANAGER"] = (
+        "liger_kernel.triton.monkey_patch:LigerTritonFileCacheManager"
+    )

liger_kernel-0.0.0.dist-info/METADATA

@@ -0,0 +1,14 @@
+Metadata-Version: 2.1
+Name: liger-kernel
+Version: 0.0.0
+Requires-Dist: torch>=2.1.2
+Requires-Dist: triton>=2.3.0
+Requires-Dist: transformers>=4.40.1
+Provides-Extra: dev
+Requires-Dist: matplotlib>=3.7.2; extra == "dev"
+Requires-Dist: flake8>=4.0.1.1; extra == "dev"
+Requires-Dist: black>=24.4.2; extra == "dev"
+Requires-Dist: isort>=5.13.2; extra == "dev"
+Requires-Dist: pre-commit>=3.7.1; extra == "dev"
+Requires-Dist: torch-tb-profiler>=0.4.1; extra == "dev"
+

liger_kernel-0.0.0.dist-info/RECORD

@@ -0,0 +1,24 @@
+liger_kernel/ops/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
+liger_kernel/ops/cross_entropy.py,sha256=XRnLWW2Jo1sVllDbyTuM8ir_6WZR791fFgqoaIVzPrM,10665
+liger_kernel/ops/fused_linear_cross_entropy.py,sha256=-z5NDZ1a7htumYPJbmw1QRqgp9N_QKZZiuaZSPCb9Y0,7311
+liger_kernel/ops/geglu.py,sha256=DiJSy4I8kouPFyNpKUuthfibZWRioPJMGR-4MJgebhg,3660
+liger_kernel/ops/rms_norm.py,sha256=iQd8ZDzNM-3b05eLzjh1Jfj2C8QKAtg59h-b-XuIo5s,4299
+liger_kernel/ops/rope.py,sha256=fYBct8gDQfKPZdMWlzkZZ8kBzh6nQ7DIpDsc7lZwM8c,8584
+liger_kernel/ops/swiglu.py,sha256=__QsfYxKyZHtRScm31zL3sAOVEblQFqKj2ll8I4Odqg,2835
+liger_kernel/ops/utils.py,sha256=cC7rvhiEBW-8x4qQRTUYWW790k3TA-S7pKbJmdRj-Xc,1080
+liger_kernel/transformers/__init__.py,sha256=7rOw9yZ8kNXO483Colx-EUq8GcTCvCZxrxF-S7pmkkU,172
+liger_kernel/transformers/cross_entropy.py,sha256=G-L4EaUYVc25NKZ2jrlaG-d5YUvDqJdUlawPN7K1d1g,389
+liger_kernel/transformers/fused_linear_cross_entropy.py,sha256=h0AW9ubFGfz4DBwgh2CLW8rpKo9PvxYpB6AUzjx-1b0,501
+liger_kernel/transformers/geglu.py,sha256=FrLBHZRdI68jw9RR6MSTE59-xCzueOwSRp9jL8y-j98,896
+liger_kernel/transformers/monkey_patch.py,sha256=5h436874AENVnTjQAk4-Srp_GIr50CXAl2xeNTbqzJg,3988
+liger_kernel/transformers/rms_norm.py,sha256=2LHfEctSpzuNRaoZ9uUECSFK8fZeIxIsHm9QbEHZvDQ,452
+liger_kernel/transformers/rope.py,sha256=m-ah8vZBYW8tfplTXCiAPMHJWlB1tdp_JPXJeWE-Boo,943
+liger_kernel/transformers/swiglu.py,sha256=8kt4MffEZT5vx3k0WA-GO-WPLv5kGdnu_nAwlJyMI2U,1516
+liger_kernel/transformers/model/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
+liger_kernel/transformers/model/llama.py,sha256=DJOjLT5-TGMLKaPqLqyW03rLae8lJTb3nwnfg2mVNXQ,5197
+liger_kernel/triton/__init__.py,sha256=yfRe0zMb47QnqjecZWG7LnanfCTzeku7SgWRAwNVmzU,101
+liger_kernel/triton/monkey_patch.py,sha256=yRNaGdyG5PrwX5ed_MQdqtqvvpVvQ7ZD2FQ_9W1q9u8,1629
+liger_kernel-0.0.0.dist-info/METADATA,sha256=SBK5dFzMYYtFyorscmi__7u83TitHFpKMsRE9pUKXGI,461
+liger_kernel-0.0.0.dist-info/WHEEL,sha256=eOLhNAGa2EW3wWl_TU484h7q1UNgy0JXjjoqKoxAAQc,92
+liger_kernel-0.0.0.dist-info/top_level.txt,sha256=2eghu4hA3LnkM7ElW92tQ8zegWKgSbeo-k-aGe1YnvY,13
+liger_kernel-0.0.0.dist-info/RECORD,,

liger_kernel-0.0.0.dist-info/WHEEL

@@ -0,0 +1,5 @@
+Wheel-Version: 1.0
+Generator: bdist_wheel (0.44.0)
+Root-Is-Purelib: true
+Tag: py3-none-any
+

liger_kernel-0.0.0.dist-info/top_level.txt

@@ -0,0 +1 @@
+liger_kernel