liger-kernel-nightly 0.5.5.dev20250402185702__py3-none-any.whl → 0.6.4.dev20260112233432__py3-none-any.whl

liger_kernel/transformers/multi_token_attention.py

@@ -0,0 +1,64 @@
+import math
+
+import torch
+import torch.nn as nn
+
+from torch.nn.modules.utils import _pair
+
+from liger_kernel.ops import LigerMultiTokenAttentionFunction
+
+
+class LigerMultiTokenAttention(nn.Module):
+    r"""
+    Multi-Token Attention:
+        out = mask_{0}(conv2d(softmax(mask_{-\inf}(scores))))
+
+    Reference: https://arxiv.org/pdf/2504.00927
+    """
+
+    def __init__(
+        self,
+        in_channels: int,
+        out_channels: int,
+        kernel_size: int,
+        stride: int = 1,
+        padding: int = 0,
+        dilation: int = 1,
+        groups: int = 1,
+        bias: bool = True,
+        sparse: bool = False,
+    ):
+        super().__init__()
+        self.in_channels = in_channels
+        self.out_channels = out_channels
+        self.kernel_size = _pair(kernel_size)
+        self.stride = _pair(stride)
+        self.padding = _pair(padding)
+        self.dilation = _pair(dilation)
+        self.groups = groups
+        self.sparse = sparse
+
+        self.weight = nn.Parameter(torch.empty(out_channels, in_channels // groups, *self.kernel_size))
+        if bias:
+            self.bias = nn.Parameter(torch.empty(out_channels))
+        else:
+            self.register_parameter("bias", None)
+
+        self.reset_parameters()
+
+    def reset_parameters(self):
+        nn.init.kaiming_uniform_(self.weight, a=math.sqrt(5))
+        if self.bias is not None:
+            nn.init.zeros_(self.bias)
+
+    def forward(self, scores: torch.Tensor) -> torch.Tensor:
+        return LigerMultiTokenAttentionFunction.apply(
+            scores,
+            self.weight,
+            self.bias,
+            self.stride,
+            self.padding,
+            self.dilation,
+            self.groups,
+            self.sparse,
+        )

liger_kernel/transformers/poly_norm.py

@@ -0,0 +1,42 @@
+import torch
+import torch.nn as nn
+
+from liger_kernel.ops import LigerPolyNormFunction
+
+
+class LigerPolyNorm(nn.Module):
+    """
+    PolyNorm layer wrapper for Liger kernel.
+
+    PolyNorm formula:
+        y = w₀·norm(x³) + w₁·norm(x²) + w₂·norm(x) + b
+        where norm(u) = u / sqrt(mean(u²) + ε)
+
+    Reference:
+        https://github.com/BryceZhuo/PolyCom/
+
+    Args:
+        eps: epsilon for numerical stability (default: 1e-6)
+        in_place: whether to in-place modify grad_output in backward to save memory (default: False).
+                  Set to True to save memory if grad_output is not needed elsewhere.
+    """
+
+    def __init__(self, eps=1e-6, in_place=True):
+        super().__init__()
+        # Align with PolyCom reference: initialize weights to (1/3, 1/3, 1/3) and bias to 1.0
+        self.weight = nn.Parameter(torch.full((3,), 1.0 / 3.0))
+        self.bias = nn.Parameter(torch.tensor(1.0))
+        self.variance_epsilon = eps
+        self.in_place = in_place
+
+    def forward(self, hidden_states):
+        return LigerPolyNormFunction.apply(
+            hidden_states,
+            self.weight,
+            self.bias,
+            self.variance_epsilon,
+            self.in_place,
+        )
+
+    def extra_repr(self):
+        return f"weight_shape={tuple(self.weight.shape)}, eps={self.variance_epsilon}, in_place={self.in_place}"

liger_kernel/transformers/qwen2vl_mrope.py

@@ -1,4 +1,4 @@
-from liger_kernel.ops.qwen2vl_mrope import LigerQwen2VLMRopeFunction
+from liger_kernel.ops import LigerQwen2VLMRopeFunction
 
 
 def liger_multimodal_rotary_pos_emb(q, k, cos, sin, mrope_section, unsqueeze_dim=1):

liger_kernel/transformers/rms_norm.py

@@ -1,7 +1,7 @@
 import torch
 import torch.nn as nn
 
-from liger_kernel.ops.rms_norm import LigerRMSNormFunction
+from liger_kernel.ops import LigerRMSNormFunction
 
 
 class LigerRMSNorm(nn.Module):
@@ -13,18 +13,25 @@ class LigerRMSNorm(nn.Module):
         casting_mode="llama",
         init_fn="ones",
         in_place=True,
+        row_mode=None,
+        elementwise_affine=True,
     ):
         super().__init__()
         assert init_fn in [
             "ones",
             "zeros",
         ], f"init_fn must be either 'ones' or 'zeros', got {init_fn}"
-        self.weight = nn.Parameter(torch.ones(hidden_size) if init_fn == "ones" else torch.zeros(hidden_size))
-        self.variance_epsilon, self.offset, self.casting_mode, self.in_place = (
+        self.elementwise_affine = elementwise_affine
+        if self.elementwise_affine:
+            self.weight = nn.Parameter(torch.ones(hidden_size) if init_fn == "ones" else torch.zeros(hidden_size))
+        else:
+            self.register_parameter("weight", None)
+        self.variance_epsilon, self.offset, self.casting_mode, self.in_place, self.row_mode = (
             eps,
             offset,
             casting_mode,
             in_place,
+            row_mode,
         )
 
     def forward(self, hidden_states):
@@ -35,9 +42,50 @@ class LigerRMSNorm(nn.Module):
             self.offset,
             self.casting_mode,
             self.in_place,
+            self.row_mode,
         )
 
     def extra_repr(self):
-        return (
-            f"{tuple(self.weight.shape)}, eps={self.variance_epsilon}, offset={self.offset}, in_place={self.in_place}"
-        )
+        return f"weight_shape={tuple(self.weight.shape) if self.weight is not None else None}, eps={self.variance_epsilon}, offset={self.offset}, in_place={self.in_place}, row_mode={self.row_mode}"
+
+
+class LigerRMSNormForGemma(LigerRMSNorm):
+    def __init__(
+        self, hidden_size, eps=1e-6, offset=1.0, casting_mode="gemma", init_fn="zeros", in_place=True, row_mode=None
+    ):
+        super().__init__(hidden_size, eps, offset, casting_mode, init_fn, in_place, row_mode)
+
+
+class LigerRMSNormForGemma2(LigerRMSNorm):
+    def __init__(
+        self, hidden_size, eps=1e-6, offset=1.0, casting_mode="gemma", init_fn="zeros", in_place=False, row_mode=None
+    ):
+        super().__init__(hidden_size, eps, offset, casting_mode, init_fn, in_place, row_mode)
+
+
+class LigerRMSNormForGemma3(LigerRMSNorm):
+    """Gemma3RMSNorm has a dim argument not hidden_size used in q_norm and k_norm."""
+
+    def __init__(self, dim, eps=0.000001, offset=1.0, casting_mode="gemma", init_fn="zeros", in_place=False):
+        super().__init__(dim, eps, offset, casting_mode, init_fn, in_place)
+
+
+class LigerRMSNormForOlmo2(LigerRMSNorm):
+    def __init__(
+        self, hidden_size, eps=1e-6, offset=0.0, casting_mode="llama", init_fn="ones", in_place=False, row_mode=None
+    ):
+        super().__init__(hidden_size, eps, offset, casting_mode, init_fn, in_place, row_mode)
+
+
+class LigerRMSNormForGlm4(LigerRMSNorm):
+    def __init__(
+        self, hidden_size, eps=1e-6, offset=0.0, casting_mode="llama", init_fn="ones", in_place=False, row_mode=None
+    ):
+        super().__init__(hidden_size, eps, offset, casting_mode, init_fn, in_place, row_mode)
+
+
+class LigerRMSNormForQwen3Next(LigerRMSNorm):
+    def __init__(
+        self, hidden_size, eps=1e-6, offset=1.0, casting_mode="gemma", init_fn="zeros", in_place=False, row_mode=None
+    ):
+        super().__init__(hidden_size, eps, offset, casting_mode, init_fn, in_place, row_mode)

liger_kernel/transformers/rope.py

@@ -1,4 +1,8 @@
-from liger_kernel.ops.rope import LigerRopeFunction
+from typing import Tuple
+
+import torch
+
+from liger_kernel.ops import LigerRopeFunction
 
 
 def liger_rotary_pos_emb(q, k, cos, sin, position_ids=None, unsqueeze_dim=1):
@@ -18,3 +22,43 @@ def liger_rotary_pos_emb(q, k, cos, sin, position_ids=None, unsqueeze_dim=1):
     """
 
     return LigerRopeFunction.apply(q, k, cos, sin, position_ids, unsqueeze_dim)
+
+
+def liger_rotary_pos_emb_vision(
+    q: torch.Tensor,
+    k: torch.Tensor,
+    cos: torch.Tensor,
+    sin: torch.Tensor,
+) -> Tuple[torch.Tensor, torch.Tensor]:
+    """
+    Modified version of liger_rotary_pos_emb for qwen3_vl's apply_rotary_pos_emb_vision function.
+    Manually tranposed the input and output to match the expected shape for liger_rotary_pos_emb.
+    Reference: https://https://github.com/huggingface/transformers/blob/v5.0.0rc0/src/transformers/models/qwen3_vl/modeling_qwen3_vl.py#L116
+
+    Args:
+        q (torch.Tensor): The query tensor of shape (seq_length, num_heads, head_dim),
+        with stride (num_heads * head_dim, head_dim, 1).
+        k (torch.Tensor): The query tensor of shape (seq_length, num_heads, head_dim),
+        with stride (num_heads * head_dim, head_dim, 1). Same as q.
+        cos (torch.Tensor): The cosine tensor of shape (seq_length, head_dim).
+        sin (torch.Tensor): The sine tensor of shape (seq_length, head_dim).
+
+    Returns:
+        Tuple[torch.Tensor, torch.Tensor]: The query and key tensors with the same shape and stride as inputs.
+    """
+    orig_q_dtype, orig_k_dtype = q.dtype, k.dtype
+
+    # tranpose to (1, num_heads, seq_length, head_dim) and cast to float32 to match liger_rotary_pos_emb input shape
+    # also unsqueeze for batch dim
+    q32 = q.to(torch.float32).unsqueeze(0).transpose(1, 2)
+    k32 = k.to(torch.float32).unsqueeze(0).transpose(1, 2)
+    cos32 = cos.to(torch.float32)
+    sin32 = sin.to(torch.float32)
+
+    q_out, k_out = liger_rotary_pos_emb(q32, k32, cos32, sin32)
+
+    # transpose back to (seq_length, num_heads, head_dim) and cast back to original dtype
+    # also squeeze out batch dim
+    q_out = q_out.transpose(1, 2).squeeze(0).to(orig_q_dtype)
+    k_out = k_out.transpose(1, 2).squeeze(0).to(orig_k_dtype)
+    return q_out, k_out

liger_kernel/transformers/softmax.py

@@ -0,0 +1,12 @@
+import torch
+import torch.nn as nn
+
+from liger_kernel.ops import LigerSoftmaxFunction
+
+
+class LigerSoftmax(nn.Module):
+    def __init__(self):
+        super().__init__()
+
+    def forward(self, x: torch.Tensor):
+        return LigerSoftmaxFunction.apply(x)

liger_kernel/transformers/sparsemax.py

@@ -0,0 +1,16 @@
+import torch
+import torch.nn as nn
+
+from liger_kernel.ops import LigerSparsemaxFunction
+
+
+class LigerSparsemax(nn.Module):
+    def __init__(self, dim: int = -1):
+        super().__init__()
+        self.dim = dim
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        return LigerSparsemaxFunction.apply(x, self.dim)
+
+    def extra_repr(self) -> str:
+        return f"dim={self.dim}"

liger_kernel/transformers/swiglu.py

@@ -1,6 +1,6 @@
 import torch.nn as nn
 
-from liger_kernel.ops.swiglu import LigerSiLUMulFunction
+from liger_kernel.ops import LigerSiLUMulFunction
 
 
 class LigerSwiGLUMLP(nn.Module):
@@ -56,3 +56,41 @@ class LigerPhi3SwiGLUMLP(nn.Module):
         up_states = self.gate_up_proj(x)
         gate, up_states = up_states.chunk(2, dim=-1)
         return self.down_proj(LigerSiLUMulFunction.apply(gate, up_states))
+
+
+class LigerQwen3MoeSwiGLUMLP(nn.Module):
+    """
+    Patch Qwen3MoeMLP to use LigerSiLUMulFunction.
+    https://github.com/huggingface/transformers/blob/v4.51.3/src/transformers/models/qwen3_moe/modular_qwen3_moe.py#L57
+    """
+
+    def __init__(self, config, intermediate_size=None):
+        super().__init__()
+        self.config = config
+        self.hidden_size = config.hidden_size
+        self.intermediate_size = intermediate_size if intermediate_size is not None else 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 LigerHunyuanV1SwiGLUMLP(nn.Module):
+    def __init__(self, config, layer_idx=None, is_shared_mlp=False):
+        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)
+        self.layer_idx = layer_idx
+        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)))

liger_kernel/transformers/tiled_mlp.py

@@ -0,0 +1,125 @@
+from typing import Optional
+
+import torch.nn as nn
+
+from liger_kernel.ops import LigerGELUMulFunction
+from liger_kernel.ops import LigerSiLUMulFunction
+from liger_kernel.ops import apply_tiled_mlp
+
+
+class LigerTiledGEGLUMLP(nn.Module):
+    """
+    Memory-efficient GEGLU MLP using tiled computation.
+
+    This module combines GEGLU activation with tiled processing to handle
+    very long sequences efficiently. The forward pass is recomputed during
+    backward to save memory.
+
+    Args:
+        config: Model configuration with hidden_size and intermediate_size attributes
+        num_shards: Number of shards to split the sequence. If None, automatically
+                   calculated as ceil(seqlen / hidden_size)
+    """
+
+    def __init__(self, config, num_shards: Optional[int] = None):
+        super().__init__()
+        self.config = config
+        self.hidden_size = config.hidden_size
+        self.intermediate_size = config.intermediate_size
+        self.num_shards = num_shards
+
+        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)
+
+        # Validate activation function
+        if hasattr(config, "hidden_act") and config.hidden_act not in [
+            "gelu",
+            "gelu_new",
+            "gelu_pytorch_tanh",
+        ]:
+            raise ValueError(f"LigerTiledGEGLUMLP requires GELU activation, got {config.hidden_act}")
+
+    def _mlp_forward(self, module, x):
+        """Internal MLP forward function for tiled computation."""
+        gate = module.gate_proj(x)
+        up = module.up_proj(x)
+        return module.down_proj(LigerGELUMulFunction.apply(gate, up))
+
+    def forward(self, x):
+        """
+        Forward pass with tiled computation.
+
+        Args:
+            x: Input tensor of shape [batch_size, seq_len, hidden_size]
+               or [seq_len, hidden_size]
+
+        Returns:
+            Output tensor of the same shape as input
+        """
+        compute_params = [p for p in self.parameters() if p.requires_grad]
+
+        return apply_tiled_mlp(
+            fn=self._mlp_forward,
+            mlp_module=self,
+            x=x,
+            num_shards=self.num_shards,
+            compute_params=compute_params,
+        )
+
+
+class LigerTiledSwiGLUMLP(nn.Module):
+    """
+    Memory-efficient SwiGLU MLP using tiled computation.
+
+    This module combines SwiGLU activation with tiled processing to handle
+    very long sequences efficiently. The forward pass is recomputed during
+    backward to save memory.
+
+    Args:
+        config: Model configuration with hidden_size and intermediate_size attributes
+        num_shards: Number of shards to split the sequence. If None, automatically
+                   calculated as ceil(seqlen / hidden_size)
+    """
+
+    def __init__(self, config, num_shards: Optional[int] = None):
+        super().__init__()
+        self.config = config
+        self.hidden_size = config.hidden_size
+        self.intermediate_size = config.intermediate_size
+        self.num_shards = num_shards
+
+        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)
+
+        # Validate activation function
+        if hasattr(config, "hidden_act") and config.hidden_act not in ["silu", "swish"]:
+            raise ValueError(f"LigerTiledSwiGLUMLP requires SiLU/Swish activation, got {config.hidden_act}")
+
+    def _mlp_forward(self, module, x):
+        """Internal MLP forward function for tiled computation."""
+        gate = module.gate_proj(x)
+        up = module.up_proj(x)
+        return module.down_proj(LigerSiLUMulFunction.apply(gate, up))
+
+    def forward(self, x):
+        """
+        Forward pass with tiled computation.
+
+        Args:
+            x: Input tensor of shape [batch_size, seq_len, hidden_size]
+               or [seq_len, hidden_size]
+
+        Returns:
+            Output tensor of the same shape as input
+        """
+        compute_params = [p for p in self.parameters() if p.requires_grad]
+
+        return apply_tiled_mlp(
+            fn=self._mlp_forward,
+            mlp_module=self,
+            x=x,
+            num_shards=self.num_shards,
+            compute_params=compute_params,
+        )

liger_kernel/transformers/trainer/orpo_trainer.py

@@ -1,5 +1,3 @@
-from typing import Any
-from typing import Callable
 from typing import Dict
 from typing import List
 from typing import Literal
@@ -13,57 +11,7 @@ from torch.distributed.fsdp import FullyShardedDataParallel
 from trl.trainer import ORPOTrainer
 
 from liger_kernel.chunked_loss import LigerFusedLinearORPOLoss
-
-
-class _FSDPForwardRedirection:
-    """
-    Modified based on
-    https://github.com/Lightning-AI/pytorch-lightning/blob/d3f9c83d6efa4f1def36aa6c199600946cdb9117/src/lightning/pytorch/strategies/strategy.py#L601-L648
-    Redirect a method call through FullyShardedDataParallel.forward so that the FSDP module's root pre-forward and
-    post-forward can be properly executed around the method call.
-    This is needed in cases where we call a submodule of a FSDP module. For instance, when we want to call only
-    the `LlamaModel` part out of a FSDP-wrapped `LlamaForCausalLM` to get the hidden states without involving
-    GPU-memory-heavy `lm_head` and cross entropy computation, doing this directly (i.e. `model.model.forward()`)
-    will not work because the first `nn.Embedding` layer is not independently wrapped as a FSDP module (because of
-    the transformer-based wrapping policy), and not calling it through FSDP root module forward will not all-gather
-    its parameter, thus resulting in "RuntimeError: 'weight' must be 2-D" error. Similarly, if we want to call just
-    the `lm_head` part of a model, we need this trick too to properly get its params all-gathered.
-    """
-
-    def __call__(
-        self,
-        wrapper_module: FullyShardedDataParallel,
-        method: Callable,
-        *args: Any,
-        **kwargs: Any,
-    ):
-        """Reroutes a method call through the `wrapper_module`'s `forward` method.
-        Args:
-            wrapper_module: The module that has `original_module` wrapped.
-            original_module: The module that was wrapped inside `wrapper_module`.
-            method_name: The name of the method that should be called on the `original_module` after inputs get
-                redirected through the `wrapper_module`'s `forward` method.
-            *args: The positional arguments to the method `method_name`. They will get passed to a patched
-                `forward` method instead.
-            **kwargs: The keyword arguments to the method `method_name`. They will get passed to a patched
-                `forward` method instead.
-        """
-        assert isinstance(wrapper_module, FullyShardedDataParallel)
-        original_module = wrapper_module._fsdp_wrapped_module
-        original_forward = original_module.forward
-
-        def wrapped_forward(*_args: Any, **_kwargs: Any) -> Any:
-            # Unpatch ourselves immediately before calling the method `method_name`
-            # because itself may want to call the real `forward`
-            original_module.forward = original_forward  # type: ignore[method-assign]
-            # Call the actual method e.g. `.training_step(...)`
-            out = method(*_args, **_kwargs)
-            return out
-
-        # Patch the original_module's forward so we can redirect the arguments back to the real method
-        original_module.forward = wrapped_forward  # type: ignore[method-assign]
-        wrapper_output = wrapper_module(*args, **kwargs)
-        return wrapper_output
+from liger_kernel.transformers.fsdp import _FSDPForwardRedirection
 
 
 class LigerORPOTrainer(ORPOTrainer):

liger_kernel/transformers/tvd.py

@@ -1,6 +1,6 @@
 import torch.nn as nn
 
-from liger_kernel.ops.tvd import LigerTVDLossFunction
+from liger_kernel.ops import LigerTVDLossFunction
 
 
 class LigerTVDLoss(nn.Module):

liger_kernel/utils.py

@@ -1,18 +1,81 @@
+try:
+    import peft  # noqa: F401
+
+    PEFT_AVAILABLE = True
+except ImportError:
+    PEFT_AVAILABLE = False
+
 import torch
 
 
+def is_peft_available():
+    return PEFT_AVAILABLE
+
+
+def infer_comm_backend():
+    """
+    Get communication backend name based on the environment.
+    """
+    if torch.distributed.is_nccl_available():
+        # Works for Nvidia
+        # TODO: nccl may not work for AMD decices that may require use of rccl.
+        return "nccl"
+    elif is_npu_available():
+        # Use Ascend NPU if available (torch.npu)
+        # Ascend is not standard torch backend and requires extension.
+        # Assume that it is installed if NPUs are being used in
+        # multi device environment.
+        return "ascend"
+    # XPU (Intel) if available
+    elif torch.distributed.distributed_c10d.is_xccl_available():
+        return "xccl"
+    elif torch.distributed.is_mpi_available():
+        # CPU backend, first option
+        return "mpi"
+    elif torch.distributed.is_gloo_available():
+        # CPU backend, backup option
+        return "gloo"
+    else:
+        raise RuntimeError("There is no distributed backend available.")
+
+
 def infer_device():
     """
     Get current device name based on available devices
     """
     if torch.cuda.is_available():  # Works for both Nvidia and AMD
         return "cuda"
+    # Use Ascend NPU if available (torch.npu)
+    elif is_npu_available():
+        return "npu"
+    # XPU (Intel) if available
     elif torch.xpu.is_available():
         return "xpu"
     else:
         return "cpu"
 
 
+def is_npu_available() -> bool:
+    """Detect Ascend NPU availability."""
+    try:
+        from transformers.utils import is_torch_npu_available
+
+        return is_torch_npu_available()
+    except Exception:
+        return False
+
+
+def get_npu_multi_processor_count() -> int:
+    """Return a heuristic multi-processor count for NPU."""
+    if is_npu_available():
+        NPU_MULTI_PROCESSOR_COUNT = 48
+        dev_props = torch.npu.get_device_properties()
+        # The vector_core_num attribute is supported in the torch.npu v7.2.0 release version.
+        return dev_props.vector_core_num if hasattr(dev_props, "vector_core_num") else NPU_MULTI_PROCESSOR_COUNT
+    # Reasonable default to avoid division by zero
+    return 1
+
+
 def transformers_version_dispatch(
     required_version: str,
     before_fn,