liger-kernel-nightly 0.0.1.dev20240819184814__py3-none-any.whl → 0.6.4.dev20251212103629__py3-none-any.whl

liger_kernel/chunked_loss/simpo_loss.py

@@ -0,0 +1,165 @@
+import torch
+import torch.nn.functional as F
+
+from liger_kernel.chunked_loss.fused_linear_preference import LigerFusedLinearPreferenceBase
+
+
+class LigerFusedLinearSimPOFunction(LigerFusedLinearPreferenceBase):
+    @staticmethod
+    def preference_loss_fn(
+        chosen_logps,
+        rejected_logps,
+        full_target,
+        beta=0.1,
+        gamma=0.5,
+        label_smoothing=0.0,
+    ):
+        """
+        Paper: https://arxiv.org/pdf/2405.14734
+
+        Formula:
+        L_SimPO(π_θ) = -E [log σ(β/|y_w| log π_θ(y_w|x) - β/|y_l| log π_θ(y_l|x) - γ)]
+
+        Where:
+        - π_θ(y|x): Policy (model) probability
+        - y_w: Chosen sequence
+        - y_l: Rejected sequence
+        - |y_w|, |y_l|: Sequence lengths
+        - σ: Sigmoid function
+        - β: beta weight
+        - γ: gemma margin term
+
+        Args:
+            chosen_logps (torch.Tensor): Avg log probabilities of chosen tokens. Shape: (batch_size,).
+            rejected_logps (torch.Tensor): Avg log probabilities of rejected tokens. Shape: (batch_size,).
+            full_target: Non chunked full target tensor
+            beta (float): beta weight
+            gamma (float): gemma margin term
+            label_smoothing (float): Label smoothing factor, will reduce to Equation above when label_smoothing -> 0.
+        """
+        logits = beta * (chosen_logps - rejected_logps) - gamma
+        loss = (-F.logsigmoid(logits) * (1 - label_smoothing) - F.logsigmoid(-logits) * label_smoothing).sum() / (
+            full_target.shape[0] // 2
+        )
+
+        chosen_rewards = beta * chosen_logps
+        rejected_rewards = beta * rejected_logps
+
+        return loss, chosen_rewards, rejected_rewards
+
+    @classmethod
+    def forward(
+        cls,
+        ctx,
+        _input,
+        weight,
+        target,
+        bias=None,
+        ignore_index=-100,
+        beta=0.1,
+        alpha=1.0,
+        label_smoothing=0.0,
+        compute_nll_loss=False,
+        compiled=True,
+        gamma=0.5,
+        chunk_size=1,
+    ):
+        """
+        Fused linear layer with SimPO loss.
+        Args:
+            _input (torch.Tensor): Input tensor. Shape: (batch_size * seq_len, hidden_size)
+            weight (torch.Tensor): Weight tensor. Shape: (vocab_size, hidden_size)
+            target (torch.LongTensor): Target tensor. Shape: (batch_size * seq_len,)
+            bias (torch.Tensor, optional): Bias tensor. Shape: (vocab_size,)
+            ignore_index (int): Index to ignore in loss computation
+            beta (float): Weight for the odds ratio loss
+            alpha (float): Weight for the alpha parameter
+            label_smoothing (float): Label smoothing factor
+            compute_nll_loss (bool): Whether to compute the NLL loss
+            compiled (bool): Whether to use torch compile
+            gamma (float): Weight for the gamma parameter
+            chunk_size (int): Size of chunks for processing
+        Returns:
+            torch.Tensor: Computed loss
+        """
+        return super().forward(
+            cls=cls,
+            ctx=ctx,
+            _input=_input,
+            weight=weight,
+            target=target,
+            bias=bias,
+            ignore_index=ignore_index,
+            alpha=alpha,
+            beta=beta,
+            label_smoothing=label_smoothing,
+            compute_nll_loss=compute_nll_loss,
+            compiled=compiled,
+            gamma=gamma,
+            chunk_size=chunk_size,
+        )
+
+    @staticmethod
+    def backward(ctx, *grad_output):
+        grads = LigerFusedLinearPreferenceBase.backward(ctx, grad_output)[:4]
+        return *grads, None, None, None, None, None, None, None, None
+
+
+class LigerFusedLinearSimPOLoss(torch.nn.Module):
+    """
+    Fused linear layer with SimPO loss.
+    """
+
+    def __init__(
+        self,
+        ignore_index: int = -100,
+        beta: float = 0.1,
+        alpha: float = 1.0,
+        label_smoothing: float = 0.0,
+        compute_nll_loss: bool = True,
+        compiled: bool = True,
+        gamma: float = 0.5,
+        chunk_size: int = 1,
+    ):
+        """
+        Args:
+            ignore_index (int): Index to ignore in the loss.
+            beta (float): Weight for the odds ratio loss.
+            alpha (float): Weight for the alpha parameter.
+            label_smoothing (float): Label smoothing factor.
+            compute_nll_loss (bool): Whether to compute the NLL loss.
+            compiled (bool): Whether to use the torch compiled kernel.
+            gamma (float): Weight for the gamma parameter.
+            chunk_size (int): Size of chunks for processing.
+        """
+        super().__init__()
+        self.ignore_index = ignore_index
+        self.beta = beta
+        self.alpha = alpha
+        self.label_smoothing = label_smoothing
+        self.compute_nll_loss = compute_nll_loss
+        self.compiled = compiled
+        self.gamma = gamma
+        self.chunk_size = chunk_size
+
+    def forward(
+        self,
+        lin_weight,
+        _input,
+        target,
+        bias=None,
+    ):
+        return LigerFusedLinearSimPOFunction.apply(
+            _input,
+            lin_weight,
+            target,
+            bias,
+            self.ignore_index,
+            self.beta,
+            self.alpha,
+            self.label_smoothing,
+            self.compute_nll_loss,
+            self.compiled,
+            self.gamma,
+            self.chunk_size,
+        )

liger_kernel/env_report.py

@@ -0,0 +1,63 @@
+import platform
+import sys
+
+from importlib.metadata import version
+
+
+def print_env_report():
+    """
+
+    Prints a report of the environment.  Useful for debugging and reproducibility.
+    Usage:
+    ```
+    python -m liger_kernel.env_report
+    ```
+
+    """
+    print("Environment Report:")
+    print("-------------------")
+    print(f"Operating System: {platform.platform()}")
+    print(f"Python version: {sys.version.split()[0]}")
+
+    try:
+        print(f"Liger Kernel version: {version('liger-kernel')}")
+    except ImportError:
+        print("Liger Kernel: Not installed")
+
+    try:
+        import torch
+
+        print(f"PyTorch version: {torch.__version__}")
+        cuda_version = torch.version.cuda if torch.cuda.is_available() else "Not available"
+        print(f"CUDA version: {cuda_version}")
+        hip_version = torch.version.hip if torch.cuda.is_available() and torch.version.hip else "Not available"
+        print(f"HIP(ROCm) version: {hip_version}")
+
+    except ImportError:
+        print("PyTorch: Not installed")
+        print("CUDA version: Unable to query")
+        print("HIP(ROCm) version: Unable to query")
+
+    try:
+        import triton
+
+        print(f"Triton version: {triton.__version__}")
+    except ImportError:
+        print("Triton: Not installed")
+
+    try:
+        import transformers
+
+        print(f"Transformers version: {transformers.__version__}")
+    except ImportError:
+        print("Transformers: Not installed")
+
+    try:
+        xpu_version = torch.version.xpu if torch.xpu.is_available() else "XPU Not Available"
+        print(f"XPU version: {xpu_version}")
+    except ImportError:
+        print("XPU version: Unable to query")
+
+
+if __name__ == "__main__":
+    print_env_report()

liger_kernel/ops/__init__.py

@@ -0,0 +1,141 @@
+"""
+Liger-Kernel operators with automatic vendor-specific replacement.
+
+This module provides two ways to import operators:
+
+1. Import from this package (recommended for Function classes):
+       from liger_kernel.ops import LigerGELUMulFunction
+
+   This automatically uses vendor-specific implementation if available.
+
+2. Import from submodules (for kernel functions or specific access):
+       from liger_kernel.ops.geglu import geglu_forward, geglu_backward
+
+   This always uses the default implementation (no auto-replacement).
+
+The replacement mechanism:
+1. Default implementations are imported from individual modules (e.g., geglu.py)
+2. On module load, device is detected via infer_device()
+3. If running on a supported vendor device (npu, xpu, etc.), the default
+   implementations are replaced with vendor-specific ones
+4. All subsequent imports from this package get the replaced versions
+
+Note: Direct imports from submodules (e.g., from liger_kernel.ops.geglu import ...)
+      are NOT affected by the replacement mechanism.
+"""
+
+# =============================================================================
+# Import default implementations
+# Both Function classes and kernel functions are imported here.
+# All of these can be replaced by vendor-specific implementations.
+# =============================================================================
+
+from liger_kernel.ops.cross_entropy import LigerCrossEntropyFunction  # noqa: F401
+from liger_kernel.ops.cross_entropy import cross_entropy_backward  # noqa: F401
+from liger_kernel.ops.cross_entropy import cross_entropy_forward  # noqa: F401
+from liger_kernel.ops.dyt import LigerDyTFunction  # noqa: F401
+from liger_kernel.ops.experimental.embedding import LigerEmbeddingFunction  # noqa: F401
+from liger_kernel.ops.fused_add_rms_norm import LigerFusedAddRMSNormFunction  # noqa: F401
+from liger_kernel.ops.fused_add_rms_norm import fused_add_rms_norm_backward  # noqa: F401
+from liger_kernel.ops.fused_add_rms_norm import fused_add_rms_norm_forward  # noqa: F401
+from liger_kernel.ops.fused_linear_cross_entropy import LigerFusedLinearCrossEntropyFunction  # noqa: F401
+from liger_kernel.ops.fused_linear_cross_entropy import fused_linear_cross_entropy_backward  # noqa: F401
+from liger_kernel.ops.fused_linear_cross_entropy import fused_linear_cross_entropy_forward  # noqa: F401
+from liger_kernel.ops.fused_linear_jsd import LigerFusedLinearJSDFunction  # noqa: F401
+from liger_kernel.ops.fused_linear_jsd import fused_linear_jsd_backward  # noqa: F401
+from liger_kernel.ops.fused_linear_jsd import fused_linear_jsd_forward  # noqa: F401
+from liger_kernel.ops.fused_neighborhood_attention import LigerFusedNeighborhoodAttentionFunction  # noqa: F401
+from liger_kernel.ops.geglu import LigerGELUMulFunction  # noqa: F401
+from liger_kernel.ops.geglu import geglu_backward  # noqa: F401
+from liger_kernel.ops.geglu import geglu_forward  # noqa: F401
+from liger_kernel.ops.group_norm import LigerGroupNormFunction  # noqa: F401
+from liger_kernel.ops.group_norm import group_norm_backward  # noqa: F401
+from liger_kernel.ops.group_norm import group_norm_forward  # noqa: F401
+from liger_kernel.ops.grpo_loss import GrpoLossFunction  # noqa: F401
+from liger_kernel.ops.jsd import LigerJSDFunction  # noqa: F401
+from liger_kernel.ops.jsd import jsd_backward  # noqa: F401
+from liger_kernel.ops.jsd import jsd_forward  # noqa: F401
+from liger_kernel.ops.kl_div import LigerKLDivLossFunction  # noqa: F401
+from liger_kernel.ops.layer_norm import LigerLayerNormFunction  # noqa: F401
+from liger_kernel.ops.layer_norm import layer_norm_backward  # noqa: F401
+from liger_kernel.ops.layer_norm import layer_norm_forward  # noqa: F401
+from liger_kernel.ops.llama4_rope import LigerLlama4RopeFunction  # noqa: F401
+from liger_kernel.ops.multi_token_attention import LigerMultiTokenAttentionFunction  # noqa: F401
+from liger_kernel.ops.poly_norm import LigerPolyNormFunction  # noqa: F401
+from liger_kernel.ops.poly_norm import poly_norm_backward  # noqa: F401
+from liger_kernel.ops.poly_norm import poly_norm_forward  # noqa: F401
+from liger_kernel.ops.qwen2vl_mrope import LigerQwen2VLMRopeFunction  # noqa: F401
+from liger_kernel.ops.rms_norm import LigerRMSNormFunction  # noqa: F401
+from liger_kernel.ops.rms_norm import rms_norm_backward  # noqa: F401
+from liger_kernel.ops.rms_norm import rms_norm_forward  # noqa: F401
+from liger_kernel.ops.rope import LigerRopeFunction  # noqa: F401
+from liger_kernel.ops.rope import rope_backward  # noqa: F401
+from liger_kernel.ops.rope import rope_forward  # noqa: F401
+from liger_kernel.ops.softmax import LigerSoftmaxFunction  # noqa: F401
+from liger_kernel.ops.sparsemax import LigerSparsemaxFunction  # noqa: F401
+from liger_kernel.ops.swiglu import LigerSiLUMulFunction  # noqa: F401
+from liger_kernel.ops.swiglu import swiglu_backward  # noqa: F401
+from liger_kernel.ops.swiglu import swiglu_forward  # noqa: F401
+from liger_kernel.ops.tiled_mlp import LigerTiledMLPFunction  # noqa: F401
+from liger_kernel.ops.tiled_mlp import apply_tiled_mlp  # noqa: F401
+from liger_kernel.ops.tvd import LigerTVDLossFunction  # noqa: F401
+
+# NOTE: __all__ is intentionally NOT defined.
+# - Import from this package (liger_kernel.ops) -> subject to vendor replacement
+# - Import from submodules (liger_kernel.ops.geglu) -> always use default implementation
+
+
+# =============================================================================
+# Vendor-specific replacement logic
+# =============================================================================
+
+
+def _replace_with_vendor_ops():
+    """
+    Replace/add vendor-specific operator implementations.
+
+    This function is called automatically on module load. It:
+    1. Detects the current device (cuda, npu, xpu, etc.)
+    2. Looks up the vendor for that device via VENDOR_REGISTRY
+    3. Loads and applies vendor-specific implementations
+
+    Vendor implementations should be placed in:
+        liger_kernel/ops/backends/_<vendor>/ops/
+
+    If the vendor module defines __all__, only those symbols are exported.
+    Otherwise, all public symbols (not starting with _) are auto-discovered.
+
+    Note: Vendor can both override existing ops AND add new vendor-specific ops.
+    """
+    from liger_kernel.ops.backends import get_vendor_for_device
+    from liger_kernel.utils import infer_device
+
+    device = infer_device()
+
+    # Look up vendor info for this device
+    vendor_info = get_vendor_for_device(device)
+    if vendor_info is None:
+        return
+
+    try:
+        import importlib
+
+        vendor_ops = importlib.import_module(vendor_info.module_path)
+
+        # Get names to export: use __all__ if defined, otherwise auto-discover
+        names_to_export = getattr(vendor_ops, "__all__", None)
+
+        if names_to_export is None:
+            # Auto-discover: find all public symbols (classes and functions)
+            names_to_export = [name for name in dir(vendor_ops) if not name.startswith("_")]
+
+        # Replace or add to this module's globals
+        for name in names_to_export:
+            globals()[name] = getattr(vendor_ops, name)
+
+    except ImportError:
+        # Vendor module not available, use default implementations
+        pass
+
+
+_replace_with_vendor_ops()

liger_kernel/ops/backends/README.md

@@ -0,0 +1,151 @@
+# Adding a New Vendor Backend
+
+This directory contains vendor-specific operator implementations that automatically replace the default (CUDA) implementations when running on the corresponding device.
+
+## Concepts
+
+- **Vendor**: Chip manufacturer (e.g., `ascend`, `intel`, `nvidia`)
+- **Device**: Device type (e.g., `npu`, `xpu`, `cuda`)
+- **VendorInfo**: Defines the mapping between vendor and device
+
+## Directory Structure
+
+```
+backends/
+├── README.md          
+├── __init__.py         
+├── registry.py         # VendorInfo, register_vendor(), VENDOR_REGISTRY
+├── _ascend/            # Ascend (Huawei) vendor - supports NPU
+│   ├── __init__.py     # Registers VendorInfo for NPU
+│   └── ops/
+│       ├── __init__.py # Exports vendor-specific implementations
+│       └── geglu.py    # NPU-specific GEGLU implementation
+└── _<vendor>/          # Your new vendor backend
+    └── ...
+```
+
+## How It Works
+
+1. When `liger_kernel.ops.backends` is imported, it imports all vendor packages (e.g., `_ascend`)
+2. Each vendor's `__init__.py` calls `register_vendor()` to register itself
+3. When `liger_kernel.ops` is imported, `_replace_with_vendor_ops()` is called
+4. It detects the current device via `infer_device()` and looks up the vendor
+5. Vendor implementations replace/add to the `liger_kernel.ops` namespace
+
+## Adding a New Vendor
+
+### Step 1: Create Directory Structure
+
+```bash
+mkdir -p backends/_<vendor>/ops
+touch backends/_<vendor>/__init__.py
+touch backends/_<vendor>/ops/__init__.py
+```
+
+### Step 2: Register Your Vendor
+
+In `backends/_<vendor>/__init__.py`, register your vendor:
+
+```python
+"""
+<Vendor> backend for Liger-Kernel.
+"""
+
+from liger_kernel.ops.backends.registry import VendorInfo, register_vendor
+
+register_vendor(
+    VendorInfo(
+        vendor="<vendor>",
+        device="<device>",
+    )
+)
+```
+
+
+### Step 3: Ensure Device Detection Works
+
+Make sure `infer_device()` in `liger_kernel/utils.py` can detect your device:
+
+```python
+def infer_device():
+    if torch.cuda.is_available():
+        return "cuda"
+    if is_npu_available():
+        return "npu"
+    # Add your device detection here
+    if is_<device>_available():
+        return "<device>"
+    return "cpu"
+```
+
+### Step 4: Implement Vendor-Specific Operators
+
+Create operator files in `backends/_<vendor>/ops/`. For example, `geglu.py`:
+
+```python
+import torch
+
+class LigerGELUMulFunction(torch.autograd.Function):
+    """
+    Vendor-specific LigerGELUMulFunction implementation.
+    """
+    @staticmethod
+    def forward(ctx, a, b):
+        # Your vendor-specific forward implementation
+        ...
+
+    @staticmethod
+    def backward(ctx, dc):
+        # Your vendor-specific backward implementation
+        ...
+
+# Optional: vendor-specific kernel functions
+def geglu_forward_vendor(a, b):
+    ...
+
+def geglu_backward_vendor(a, b, dc):
+    ...
+```
+
+### Step 5: Export in `ops/__init__.py`
+
+In `backends/_<vendor>/ops/__init__.py`, export your implementations:
+
+```python
+"""
+<Vendor>-specific operator implementations.
+"""
+
+from .<module> import (
+    LigerGELUMulFunction,
+    geglu_forward_vendor as geglu_forward,   # Rename to match default API
+    geglu_backward_vendor as geglu_backward,
+)
+
+# Explicitly declare what to export (recommended)
+__all__ = [
+    "LigerGELUMulFunction",
+    "geglu_forward",
+    "geglu_backward",
+]
+```
+
+## Key Points
+
+### Incremental Override
+
+You **don't need to implement all operators**. Only implement the ones that require vendor-specific adaptations. Unimplemented operators will automatically fall back to the default (CUDA) implementation.
+
+### Vendor-Specific Additions
+
+Vendors can also **add new operators** that don't exist in the default implementation. These will be exported to `liger_kernel.ops` namespace for users to import.
+
+### Naming Convention
+
+- Use the **same class/function names** as the default implementations for overrides
+- This allows seamless replacement without changing user code
+- Use `as` imports to rename if your internal naming differs
+
+## Example: Ascend NPU Backend
+
+See `_ascend/` directory for a complete example of the Ascend NPU backend implementation.

liger_kernel/ops/backends/__init__.py

@@ -0,0 +1,13 @@
+import importlib
+import pkgutil
+
+from liger_kernel.ops.backends.registry import VENDOR_REGISTRY  # noqa: F401
+from liger_kernel.ops.backends.registry import VendorInfo  # noqa: F401
+from liger_kernel.ops.backends.registry import get_vendor_for_device  # noqa: F401
+from liger_kernel.ops.backends.registry import register_vendor  # noqa: F401
+
+# Auto-import all _<vendor> subpackages to trigger registration
+# Each vendor's __init__.py calls register_vendor() when imported
+for _, modname, ispkg in pkgutil.iter_modules(__path__):
+    if ispkg and modname.startswith("_"):
+        importlib.import_module(f"{__name__}.{modname}")

liger_kernel/ops/backends/_ascend/__init__.py

@@ -0,0 +1,5 @@
+from liger_kernel.ops.backends.registry import VendorInfo
+from liger_kernel.ops.backends.registry import register_vendor
+
+# Register Ascend vendor for NPU device
+register_vendor(VendorInfo(vendor="ascend", device="npu"))

liger_kernel/ops/backends/_ascend/ops/__init__.py

@@ -0,0 +1,15 @@
+"""
+Ascend NPU operator implementations.
+
+This module exports Ascend NPU-optimized implementations that will automatically
+replace the default implementations when running on NPU devices.
+
+Both Function classes and kernel functions can be exported here.
+
+To add a new operator:
+1. Create the implementation file (e.g., rms_norm.py)
+2. Import the Function class and/or kernel functions here
+3. Optionally add to __all__ for explicit control
+
+If __all__ is not defined, all public symbols will be auto-discovered.
+"""

liger_kernel/ops/backends/registry.py

@@ -0,0 +1,61 @@
+"""
+Vendor registry for Liger-Kernel multi-backend support.
+
+This module defines VendorInfo and the registry for vendor registration.
+Each vendor registers itself by calling register_vendor() in its __init__.py.
+"""
+
+from dataclasses import dataclass
+from typing import Optional
+
+# Dynamically get backends package path to avoid hardcoding
+_BACKENDS_PACKAGE = __name__.rsplit(".", 1)[0]  # "liger_kernel.ops.backends"
+
+
+@dataclass
+class VendorInfo:
+    """
+    Information about a chip vendor and its supported device.
+
+    Attributes:
+        vendor: Vendor name (e.g., "ascend", "intel", "nvidia")
+        device: Device type this vendor supports (e.g., "npu", "xpu")
+    """
+
+    vendor: str
+    device: str
+
+    @property
+    def module_path(self) -> str:
+        """Auto-generated module path based on vendor name."""
+        return f"{_BACKENDS_PACKAGE}._{self.vendor}.ops"
+
+
+# Registry mapping device types to their vendor info
+# Vendors register themselves via register_vendor()
+VENDOR_REGISTRY: dict[str, VendorInfo] = {}
+
+
+def register_vendor(vendor_info: VendorInfo) -> None:
+    """
+    Register a vendor's info in the global registry.
+
+    This should be called in each vendor's __init__.py to register itself.
+
+    Args:
+        vendor_info: VendorInfo instance to register
+    """
+    VENDOR_REGISTRY[vendor_info.device] = vendor_info
+
+
+def get_vendor_for_device(device: str) -> Optional[VendorInfo]:
+    """
+    Get the VendorInfo for a given device type.
+
+    Args:
+        device: Device type (e.g., "npu", "xpu")
+
+    Returns:
+        VendorInfo if found, None otherwise
+    """
+    return VENDOR_REGISTRY.get(device)