liger-kernel-nightly 0.6.4.dev20251202094519__tar.gz → 0.6.4.dev20251206103502__tar.gz

{liger_kernel_nightly-0.6.4.dev20251202094519 → liger_kernel_nightly-0.6.4.dev20251206103502}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: liger_kernel_nightly
-Version: 0.6.4.dev20251202094519
+Version: 0.6.4.dev20251206103502
 Summary: Efficient Triton kernels for LLM Training
 License: BSD 2-CLAUSE LICENSE
         Copyright 2024 LinkedIn Corporation 
@@ -312,6 +312,7 @@ loss.backward()
 | OLMo2   | `liger_kernel.transformers.apply_liger_kernel_to_olmo2`     | RoPE, RMSNorm, SwiGLU, CrossEntropyLoss, FusedLinearCrossEntropy |
 | Olmo3   | `liger_kernel.transformers.apply_liger_kernel_to_olmo3`     | RoPE, RMSNorm, SwiGLU, CrossEntropyLoss, FusedLinearCrossEntropy |
 | GLM-4   | `liger_kernel.transformers.apply_liger_kernel_to_glm4`     | RoPE, RMSNorm, SwiGLU, CrossEntropyLoss, FusedLinearCrossEntropy |
+| GPT-OSS   | `liger_kernel.transformers.apply_liger_kernel_to_gpt_oss`     | RoPE, RMSNorm, CrossEntropyLoss, FusedLinearCrossEntropy |
 | InternVL3   | `liger_kernel.transformers.apply_liger_kernel_to_internvl`     | RoPE, RMSNorm, SwiGLU, CrossEntropyLoss, FusedLinearCrossEntropy |
 | HunyuanV1   | `liger_kernel.transformers.apply_liger_kernel_to_hunyuan_v1_dense`    |  RoPE, RMSNorm, SwiGLU, CrossEntropyLoss, FusedLinearCrossEntropy       |
 | HunyuanV1 MoE | `liger_kernel.transformers.apply_liger_kernel_to_hunyuan_v1_moe` | RoPE, RMSNorm, SwiGLU, CrossEntropyLoss, FusedLinearCrossEntropy       |

{liger_kernel_nightly-0.6.4.dev20251202094519 → liger_kernel_nightly-0.6.4.dev20251206103502}/README.md

@@ -264,6 +264,7 @@ loss.backward()
 | OLMo2   | `liger_kernel.transformers.apply_liger_kernel_to_olmo2`     | RoPE, RMSNorm, SwiGLU, CrossEntropyLoss, FusedLinearCrossEntropy |
 | Olmo3   | `liger_kernel.transformers.apply_liger_kernel_to_olmo3`     | RoPE, RMSNorm, SwiGLU, CrossEntropyLoss, FusedLinearCrossEntropy |
 | GLM-4   | `liger_kernel.transformers.apply_liger_kernel_to_glm4`     | RoPE, RMSNorm, SwiGLU, CrossEntropyLoss, FusedLinearCrossEntropy |
+| GPT-OSS   | `liger_kernel.transformers.apply_liger_kernel_to_gpt_oss`     | RoPE, RMSNorm, CrossEntropyLoss, FusedLinearCrossEntropy |
 | InternVL3   | `liger_kernel.transformers.apply_liger_kernel_to_internvl`     | RoPE, RMSNorm, SwiGLU, CrossEntropyLoss, FusedLinearCrossEntropy |
 | HunyuanV1   | `liger_kernel.transformers.apply_liger_kernel_to_hunyuan_v1_dense`    |  RoPE, RMSNorm, SwiGLU, CrossEntropyLoss, FusedLinearCrossEntropy       |
 | HunyuanV1 MoE | `liger_kernel.transformers.apply_liger_kernel_to_hunyuan_v1_moe` | RoPE, RMSNorm, SwiGLU, CrossEntropyLoss, FusedLinearCrossEntropy       |

{liger_kernel_nightly-0.6.4.dev20251202094519 → liger_kernel_nightly-0.6.4.dev20251206103502}/pyproject.toml

@@ -4,7 +4,7 @@ build-backend = "setuptools.build_meta"
 
 [project]
 name = "liger_kernel_nightly"
-version = "0.6.4.dev20251202094519"
+version = "0.6.4.dev20251206103502"
 description = "Efficient Triton kernels for LLM Training"
 urls = { "Homepage" = "https://github.com/linkedin/Liger-Kernel" }
 readme = { file = "README.md", content-type = "text/markdown" }

{liger_kernel_nightly-0.6.4.dev20251202094519 → liger_kernel_nightly-0.6.4.dev20251206103502}/src/liger_kernel/ops/rms_norm.py

@@ -351,7 +351,7 @@ def _block_rms_norm_backward_kernel(
 
         # calculate the gradient of W
         if casting_mode == _CASTING_MODE_LLAMA:
-			# TODO(tcc): use tl.sum(..., dtype=tl.float32) once we upgrade to triton>=3.3.0
+            # TODO(tcc): use tl.sum(..., dtype=tl.float32) once we upgrade to triton>=3.3.0
             dW_row += tl.sum((dY_row * (X_row * rstd_row[:, None]).to(X_dtype)).to(tl.float32), 0)
         else:
             # here X_row is already in fp32 (see previous if block)

{liger_kernel_nightly-0.6.4.dev20251202094519 → liger_kernel_nightly-0.6.4.dev20251206103502}/src/liger_kernel/transformers/__init__.py

@@ -41,6 +41,7 @@ if TYPE_CHECKING:
     from liger_kernel.transformers.monkey_patch import apply_liger_kernel_to_glm4  # noqa: F401
     from liger_kernel.transformers.monkey_patch import apply_liger_kernel_to_glm4v  # noqa: F401
     from liger_kernel.transformers.monkey_patch import apply_liger_kernel_to_glm4v_moe  # noqa: F401
+    from liger_kernel.transformers.monkey_patch import apply_liger_kernel_to_gpt_oss  # noqa: F401
     from liger_kernel.transformers.monkey_patch import apply_liger_kernel_to_granite  # noqa: F401
     from liger_kernel.transformers.monkey_patch import apply_liger_kernel_to_hunyuan_v1_dense  # noqa: F401
     from liger_kernel.transformers.monkey_patch import apply_liger_kernel_to_hunyuan_v1_moe  # noqa: F401
@@ -110,6 +111,7 @@ def __getattr__(name: str):
         "apply_liger_kernel_to_glm4",
         "apply_liger_kernel_to_glm4v",
         "apply_liger_kernel_to_glm4v_moe",
+        "apply_liger_kernel_to_gpt_oss",
         "apply_liger_kernel_to_granite",
         "apply_liger_kernel_to_internvl",
         "apply_liger_kernel_to_llama",
@@ -187,6 +189,7 @@ if _TRANSFORMERS_AVAILABLE:
             "apply_liger_kernel_to_glm4",
             "apply_liger_kernel_to_glm4v",
             "apply_liger_kernel_to_glm4v_moe",
+            "apply_liger_kernel_to_gpt_oss",
             "apply_liger_kernel_to_granite",
             "apply_liger_kernel_to_internvl",
             "apply_liger_kernel_to_llama",

liger_kernel_nightly-0.6.4.dev20251206103502/src/liger_kernel/transformers/model/gpt_oss.py

@@ -0,0 +1,211 @@
+from typing import List
+from typing import Optional
+from typing import Union
+
+import torch
+
+from transformers.modeling_outputs import MoeModelOutputWithPast
+from transformers.models.mixtral.modeling_mixtral import load_balancing_loss_func
+
+from liger_kernel.transformers.model.loss_utils import LigerForCausalLMLoss
+from liger_kernel.transformers.model.loss_utils import unpack_cross_entropy_result
+from liger_kernel.transformers.model.output_classes import LigerMoeCausalLMOutputWithPast
+
+
+def lce_forward(
+    self,
+    input_ids: Optional[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,
+    output_router_logits: Optional[bool] = None,
+    cache_position: Optional[torch.LongTensor] = None,
+    logits_to_keep: Union[int, torch.Tensor] = 0,
+    skip_logits: Optional[bool] = None,
+    **kwargs,
+) -> LigerMoeCausalLMOutputWithPast:
+    r"""
+        Forward pass for causal language modeling with Mixture of Experts (MoE) architecture using Liger Kernel optimizations.
+
+    Args:
+        input_ids (`torch.LongTensor` of shape `(batch_size, sequence_length)`, *optional*):
+            Indices of input sequence tokens in the vocabulary. Indices can be obtained using tokenizers.
+        attention_mask (`torch.Tensor` of shape `(batch_size, sequence_length)`, *optional*):
+            Mask to avoid performing attention on padding token indices. Mask values selected in `[0, 1]`:
+            - 1 for tokens that are **not masked**,
+            - 0 for tokens that are **masked**.
+        position_ids (`torch.LongTensor` of shape `(batch_size, sequence_length)`, *optional*):
+            Indices of positions of each input sequence tokens in the position embeddings.
+        past_key_values (`List[torch.FloatTensor]` or `Cache`, *optional*):
+            Pre-computed hidden-states (key and values in the self-attention blocks) that can be used to speed up
+            sequential decoding. See `past_key_values` input for more details.
+        inputs_embeds (`torch.FloatTensor` of shape `(batch_size, sequence_length, hidden_size)`, *optional*):
+            Optionally, instead of passing `input_ids` you can choose to directly pass an embedded representation.
+            This is useful if you want more control over how to convert `input_ids` indices into associated vectors
+            than the model's internal embedding lookup matrix.
+        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]`.
+        use_cache (`bool`, *optional*):
+            If set to `True`, `past_key_values` key value states are returned and can be used to speed up decoding
+            (see `past_key_values`).
+        output_attentions (`bool`, *optional*):
+            Whether or not to return the attentions tensors of all attention layers. See `attentions` under returned
+            tensors for more detail.
+        output_hidden_states (`bool`, *optional*):
+            Whether or not to return the hidden states of all layers. See `hidden_states` under returned tensors for
+            more detail.
+        output_router_logits (`bool`, *optional*):
+            Whether or not to return the router logits of all MoE layers. See `router_logits` under returned tensors
+            for more detail.
+        cache_position (`torch.LongTensor` of shape `(sequence_length)`, *optional*):
+            Indices depicting the position of the input sequence tokens in the sequence.
+        logits_to_keep (`int` or `torch.Tensor`, *optional*, defaults to 0):
+            If an `int`, compute logits for the last `logits_to_keep` tokens. If `0`, calculate logits for all
+            `input_ids` (special case). Only last token logits are needed for generation, and calculating them only for that
+            token can save memory, which becomes pretty significant for long sequences or large vocabulary size.
+            If a `torch.Tensor`, must be 1D corresponding to the indices to keep in the sequence length dimension.
+            This is useful when using packed tensor format (single dimension for batch and sequence length).
+        skip_logits (`bool`, *optional*):
+            Whether to skip logit computation and directly compute loss. If `None`, defaults to `True` during training
+            when labels are provided (to save memory), and `False` during inference.
+
+    Returns:
+        `LigerMoeCausalLMOutputWithPast`: An output object containing:
+            - loss (`torch.FloatTensor` of shape `(1,)`, *optional*, returned when `labels` is provided):
+                Language modeling loss (for next-token prediction), including the auxiliary load balancing loss.
+            - aux_loss (`torch.FloatTensor`, *optional*, returned when `labels` is provided):
+                Auxiliary load balancing loss for the sparse MoE modules.
+            - logits (`torch.FloatTensor` of shape `(batch_size, sequence_length, config.vocab_size)`, *optional*):
+                Prediction scores of the language modeling head (scores for each vocabulary token before SoftMax).
+                Note: logits are `None` during training when `skip_logits=True` to save memory.
+            - past_key_values (`Cache`, *optional*, returned when `use_cache=True` is passed):
+                Cached key and value projection states for faster sequential decoding.
+            - hidden_states (`tuple(torch.FloatTensor)`, *optional*, returned when `output_hidden_states=True`):
+                Tuple of `torch.FloatTensor` (one for the output of the embeddings + one for each layer) of shape
+                `(batch_size, sequence_length, hidden_size)`. Hidden-states of the model at the output of each layer.
+            - attentions (`tuple(torch.FloatTensor)`, *optional*, returned when `output_attentions=True`):
+                Tuple of `torch.FloatTensor` (one for each layer) of shape `(batch_size, num_heads, sequence_length,
+                sequence_length)`. Attentions weights after the attention softmax.
+            - router_logits (`tuple(torch.FloatTensor)`, *optional*, returned when `output_router_logits=True`):
+                Tuple of `torch.FloatTensor` (one for each layer) of shape `(batch_size, sequence_length, num_experts)`.
+                Router logits of the MoE layers, useful to compute the auxiliary loss and z_loss.
+            - token_accuracy (`torch.FloatTensor`, *optional*, returned when `labels` is provided):
+                Token-level prediction accuracy.
+
+    Example:
+
+    ```python
+        >>> from transformers import AutoTokenizer, GptOssForCausalLM
+        >>> from liger_kernel.transformers import apply_liger_kernel_to_gpt_oss
+
+        >>> # Apply Liger Kernel patches for optimized performance
+        >>> apply_liger_kernel_to_gpt_oss()
+
+        >>> model = GptOssForCausalLM.from_pretrained("openai/gpt-oss-20b")
+        >>> tokenizer = AutoTokenizer.from_pretrained("openai/gpt-oss-20b")
+
+        >>> prompt = "Hey, are you conscious? Can you talk to me?"
+        >>> inputs = tokenizer(prompt, return_tensors="pt")
+
+        >>> # Inference: Forward pass returns logits
+        >>> outputs = model(**inputs)
+        >>> outputs.logits.shape
+        torch.Size([1, 12, 201088])
+
+        >>> # Get next token prediction
+        >>> next_token_logits = outputs.logits[:, -1, :]
+        >>> predicted_token_id = next_token_logits.argmax(dim=-1)
+
+        >>> # Training: Forward pass with labels returns loss
+        >>> labels = inputs.input_ids.clone()
+        >>> outputs = model(**inputs, labels=labels)
+        >>> outputs.loss
+        tensor(2.6454)
+    ```"""
+
+    output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions
+    output_router_logits = (
+        output_router_logits if output_router_logits is not None else self.config.output_router_logits
+    )
+
+    output_hidden_states = (
+        output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
+    )
+
+    # decoder outputs consists of (dec_features, layer_state, dec_hidden, dec_attn)
+    outputs: MoeModelOutputWithPast = 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,
+        output_router_logits=output_router_logits,
+        cache_position=cache_position,
+        **kwargs,
+    )
+
+    hidden_states = outputs.last_hidden_state
+    # Only compute necessary logits, and do not upcast them to float if we are not computing the loss
+    slice_indices = slice(-logits_to_keep, None) if isinstance(logits_to_keep, int) else logits_to_keep
+    kept_hidden_states = hidden_states[:, slice_indices, :]
+
+    shift_labels = kwargs.pop("shift_labels", None)
+    logits = None
+    loss = None
+    token_accuracy = None
+
+    if skip_logits is None:
+        skip_logits = self.training and (labels is not None or shift_labels is not None)
+
+    if skip_logits:
+        result = LigerForCausalLMLoss(
+            hidden_states=kept_hidden_states,
+            lm_head_weight=self.lm_head.weight,
+            labels=labels,
+            shift_labels=shift_labels,
+            hidden_size=self.config.hidden_size,
+            **kwargs,
+        )
+        loss, _, token_accuracy = unpack_cross_entropy_result(result)
+    else:  # if in inference model materialize logits
+        logits = self.lm_head(kept_hidden_states)
+        if labels is not None or shift_labels is not None:
+            loss = self.loss_function(
+                logits=logits,
+                labels=labels,
+                shift_labels=shift_labels,
+                vocab_size=self.vocab_size,
+                **kwargs,
+            )
+
+    aux_loss = None
+    if output_router_logits:
+        aux_loss = load_balancing_loss_func(
+            outputs.router_logits,
+            self.num_experts,
+            self.num_experts_per_tok,
+            attention_mask,
+        )
+        if labels is not None:
+            loss += self.router_aux_loss_coef * aux_loss.to(loss.device)  # make sure to reside in the same device
+
+    return LigerMoeCausalLMOutputWithPast(
+        loss=loss,
+        aux_loss=aux_loss,
+        logits=logits,
+        past_key_values=outputs.past_key_values,
+        hidden_states=outputs.hidden_states,
+        attentions=outputs.attentions,
+        router_logits=outputs.router_logits,
+        token_accuracy=token_accuracy,
+    )

{liger_kernel_nightly-0.6.4.dev20251202094519 → liger_kernel_nightly-0.6.4.dev20251206103502}/src/liger_kernel/transformers/monkey_patch.py

@@ -20,6 +20,7 @@ from liger_kernel.transformers.model.gemma import lce_forward as gemma_lce_forwa
 from liger_kernel.transformers.model.gemma import lce_forward_deprecated as gemma_lce_forward_deprecated
 from liger_kernel.transformers.model.gemma2 import lce_forward as gemma2_lce_forward
 from liger_kernel.transformers.model.gemma2 import lce_forward_deprecated as gemma2_lce_forward_deprected
+from liger_kernel.transformers.model.gpt_oss import lce_forward as gpt_oss_lce_forward
 from liger_kernel.transformers.model.llama import lce_forward as llama_lce_forward
 from liger_kernel.transformers.model.llama import lce_forward_deprecated as llama_lce_forward_deprecated
 from liger_kernel.transformers.model.llava import lce_forward as llava_lce_forward
@@ -1459,6 +1460,79 @@ def apply_liger_kernel_to_qwen3_moe(
                 _patch_rms_norm_module(decoder_layer.post_attention_layernorm)
 
 
+def apply_liger_kernel_to_gpt_oss(
+    rope: bool = True,
+    cross_entropy: bool = False,
+    fused_linear_cross_entropy: bool = True,
+    rms_norm: bool = True,
+    swiglu: bool = False,  # Set to False by default since GPT-OSS has custom expert implementation
+    model: PreTrainedModel = None,
+) -> None:
+    """
+    Apply Liger kernels to replace original implementation in HuggingFace GPT-OSS models.
+    NOTE: GPT-OSS is supported in transformers >= 4.55.0
+    NOTE: SwiGLU patching is disabled by default for GPT-OSS as it uses a custom expert
+          implementation with clamping and MXFP4 quantization.
+
+    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 False.
+        fused_linear_cross_entropy (bool):
+            Whether to apply Liger's fused linear cross entropy loss. Default is True.
+            `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 False.
+            Note: GPT-OSS uses a custom expert implementation, so SwiGLU patching is disabled by default.
+        model (PreTrainedModel): The model instance to apply Liger kernels to, if the model has already been
+            loaded. Default is None.
+    """
+    if version.parse(transformers.__version__) < version.parse("4.55.0"):
+        logger.warning("GPT-OSS support requires transformers >= 4.55.0")
+        return
+
+    assert not (cross_entropy and fused_linear_cross_entropy), (
+        "cross_entropy and fused_linear_cross_entropy cannot both be True."
+    )
+
+    from transformers.models.gpt_oss import modeling_gpt_oss
+    from transformers.models.gpt_oss.modeling_gpt_oss import GptOssModel
+
+    if rope:
+        modeling_gpt_oss.apply_rotary_pos_emb = liger_rotary_pos_emb
+
+    if rms_norm:
+        modeling_gpt_oss.GptOssRMSNorm = LigerRMSNorm
+
+    if cross_entropy:
+        from transformers.loss.loss_utils import nn
+
+        nn.functional.cross_entropy = liger_cross_entropy
+
+    if fused_linear_cross_entropy:
+        if model is not None:
+            model.forward = MethodType(gpt_oss_lce_forward, model)
+        else:
+            modeling_gpt_oss.GptOssForCausalLM.forward = gpt_oss_lce_forward
+
+    # Note: SwiGLU patching is not implemented for GPT-OSS due to custom expert implementation
+    # with clamping (swiglu_limit=7.0) and MXFP4 quantization
+
+    if model is not None:
+        # The model instance already exists, so we need to additionally patch the
+        # instance variables that reference already-instantiated modules
+
+        # get the base model from the model instance
+        base_model: GptOssModel = getattr(model, model.base_model_prefix, model)
+
+        if rms_norm:
+            _patch_rms_norm_module(base_model.norm)
+        for decoder_layer in base_model.layers:
+            if rms_norm:
+                _patch_rms_norm_module(decoder_layer.input_layernorm)
+                _patch_rms_norm_module(decoder_layer.post_attention_layernorm)
+
+
 def apply_liger_kernel_to_qwen2_vl(
     rope: bool = True,
     cross_entropy: bool = False,
@@ -2752,6 +2826,7 @@ MODEL_TYPE_TO_APPLY_LIGER_FN = {
     "glm4": apply_liger_kernel_to_glm4,
     "glm4v": apply_liger_kernel_to_glm4v,
     "glm4v_moe": apply_liger_kernel_to_glm4v_moe,
+    "gpt_oss": apply_liger_kernel_to_gpt_oss,
     "internvl": apply_liger_kernel_to_internvl,
     "llama": apply_liger_kernel_to_llama,
     "llama4_text": apply_liger_kernel_to_llama4,

{liger_kernel_nightly-0.6.4.dev20251202094519 → liger_kernel_nightly-0.6.4.dev20251206103502}/src/liger_kernel_nightly.egg-info/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: liger_kernel_nightly
-Version: 0.6.4.dev20251202094519
+Version: 0.6.4.dev20251206103502
 Summary: Efficient Triton kernels for LLM Training
 License: BSD 2-CLAUSE LICENSE
         Copyright 2024 LinkedIn Corporation 
@@ -312,6 +312,7 @@ loss.backward()
 | OLMo2   | `liger_kernel.transformers.apply_liger_kernel_to_olmo2`     | RoPE, RMSNorm, SwiGLU, CrossEntropyLoss, FusedLinearCrossEntropy |
 | Olmo3   | `liger_kernel.transformers.apply_liger_kernel_to_olmo3`     | RoPE, RMSNorm, SwiGLU, CrossEntropyLoss, FusedLinearCrossEntropy |
 | GLM-4   | `liger_kernel.transformers.apply_liger_kernel_to_glm4`     | RoPE, RMSNorm, SwiGLU, CrossEntropyLoss, FusedLinearCrossEntropy |
+| GPT-OSS   | `liger_kernel.transformers.apply_liger_kernel_to_gpt_oss`     | RoPE, RMSNorm, CrossEntropyLoss, FusedLinearCrossEntropy |
 | InternVL3   | `liger_kernel.transformers.apply_liger_kernel_to_internvl`     | RoPE, RMSNorm, SwiGLU, CrossEntropyLoss, FusedLinearCrossEntropy |
 | HunyuanV1   | `liger_kernel.transformers.apply_liger_kernel_to_hunyuan_v1_dense`    |  RoPE, RMSNorm, SwiGLU, CrossEntropyLoss, FusedLinearCrossEntropy       |
 | HunyuanV1 MoE | `liger_kernel.transformers.apply_liger_kernel_to_hunyuan_v1_moe` | RoPE, RMSNorm, SwiGLU, CrossEntropyLoss, FusedLinearCrossEntropy       |

{liger_kernel_nightly-0.6.4.dev20251202094519 → liger_kernel_nightly-0.6.4.dev20251206103502}/src/liger_kernel_nightly.egg-info/SOURCES.txt

@@ -200,6 +200,7 @@ src/liger_kernel/transformers/model/gemma3.py
 src/liger_kernel/transformers/model/glm4.py
 src/liger_kernel/transformers/model/glm4v.py
 src/liger_kernel/transformers/model/glm4v_moe.py
+src/liger_kernel/transformers/model/gpt_oss.py
 src/liger_kernel/transformers/model/hunyuan_v1.py
 src/liger_kernel/transformers/model/internvl.py
 src/liger_kernel/transformers/model/llama.py

{liger_kernel_nightly-0.6.4.dev20251202094519 → liger_kernel_nightly-0.6.4.dev20251206103502}/test/convergence/bf16/test_mini_models.py

@@ -29,6 +29,7 @@ from liger_kernel.transformers import apply_liger_kernel_to_gemma3_text
 from liger_kernel.transformers import apply_liger_kernel_to_glm4
 from liger_kernel.transformers import apply_liger_kernel_to_glm4v
 from liger_kernel.transformers import apply_liger_kernel_to_glm4v_moe
+from liger_kernel.transformers import apply_liger_kernel_to_gpt_oss
 from liger_kernel.transformers import apply_liger_kernel_to_granite
 from liger_kernel.transformers import apply_liger_kernel_to_hunyuan_v1_dense
 from liger_kernel.transformers import apply_liger_kernel_to_hunyuan_v1_moe
@@ -64,6 +65,7 @@ from test.utils import revert_liger_kernel_to_gemma3_text
 from test.utils import revert_liger_kernel_to_glm4
 from test.utils import revert_liger_kernel_to_glm4v
 from test.utils import revert_liger_kernel_to_glm4v_moe
+from test.utils import revert_liger_kernel_to_gpt_oss
 from test.utils import revert_liger_kernel_to_granite
 from test.utils import revert_liger_kernel_to_hunyuan_v1
 from test.utils import revert_liger_kernel_to_hunyuan_v1_moe
@@ -267,6 +269,15 @@ try:
 except ImportError:
     FALCONH1_AVAILABLE = False
 
+try:
+    # GPT-OSS is only available in transformers>=4.55.0
+    from transformers.models.gpt_oss.configuration_gpt_oss import GptOssConfig
+    from transformers.models.gpt_oss.modeling_gpt_oss import GptOssForCausalLM
+
+    GPT_OSS_AVAILABLE = True
+except ImportError:
+    GPT_OSS_AVAILABLE = False
+
 try:
     # Qwen3Next is only available in transformers>=4.57.0
     from transformers.models.qwen3_next.configuration_qwen3_next import Qwen3NextConfig
@@ -613,6 +624,43 @@ if QWEN3_AVAILABLE:
         ),
     )
 
+if GPT_OSS_AVAILABLE:
+    MINI_MODEL_SETUPS["mini_gpt_oss"] = MiniModelConfig(
+        liger_kernel_patch_func=apply_liger_kernel_to_gpt_oss,
+        liger_kernel_patch_revert_func=revert_liger_kernel_to_gpt_oss,
+        model_class=GptOssForCausalLM,
+        mini_model_config=GptOssConfig(
+            vocab_size=32000,  # 201088
+            hidden_size=896,
+            intermediate_size=896,  # Same as hidden_size for GPT-OSS
+            num_hidden_layers=4,
+            num_attention_heads=8,
+            num_key_value_heads=2,
+            head_dim=64,
+            hidden_act="silu",
+            max_position_embeddings=8192,
+            initializer_range=0.02,
+            rms_norm_eps=1e-5,
+            use_cache=True,
+            tie_word_embeddings=False,
+            rope_parameters={
+                "rope_type": "yarn",
+                "factor": 8.0,
+                "beta_fast": 32.0,
+                "beta_slow": 1.0,
+                "truncate": False,
+                "original_max_position_embeddings": 4096,
+            },
+            attention_dropout=0.0,
+            num_local_experts=8,  # Reduced from 32 for mini model
+            num_experts_per_tok=2,  # Reduced from 4 for mini model
+            router_aux_loss_coef=0.9,
+            output_router_logits=False,
+            sliding_window=128,
+            layer_types=["sliding_attention" if bool((i + 1) % 2) else "full_attention" for i in range(4)],
+        ),
+    )
+
 if GEMMA3_AVAILABLE:
     MINI_MODEL_SETUPS["mini_gemma3_text"] = MiniModelConfig(
         liger_kernel_patch_func=apply_liger_kernel_to_gemma3_text,
@@ -1664,6 +1712,25 @@ def run_mini_model(
                 ),
             ],
         ),
+        pytest.param(
+            "mini_gpt_oss",
+            32,
+            1e-5,
+            torch.bfloat16,
+            5e-2,
+            5e-2,
+            1e-1,
+            1e-1,
+            1e-2,
+            1e-2,
+            marks=[
+                pytest.mark.skipif(not supports_bfloat16(), reason="bfloat16 not supported on this GPU"),
+                pytest.mark.skipif(
+                    not GPT_OSS_AVAILABLE,
+                    reason="GPT-OSS not available in this version of transformers",
+                ),
+            ],
+        ),
         pytest.param(
             "mini_qwen2_vl",
             32,

{liger_kernel_nightly-0.6.4.dev20251202094519 → liger_kernel_nightly-0.6.4.dev20251206103502}/test/convergence/fp32/test_mini_models.py

@@ -29,6 +29,7 @@ from liger_kernel.transformers import apply_liger_kernel_to_gemma3_text
 from liger_kernel.transformers import apply_liger_kernel_to_glm4
 from liger_kernel.transformers import apply_liger_kernel_to_glm4v
 from liger_kernel.transformers import apply_liger_kernel_to_glm4v_moe
+from liger_kernel.transformers import apply_liger_kernel_to_gpt_oss
 from liger_kernel.transformers import apply_liger_kernel_to_granite
 from liger_kernel.transformers import apply_liger_kernel_to_hunyuan_v1_dense
 from liger_kernel.transformers import apply_liger_kernel_to_hunyuan_v1_moe
@@ -64,6 +65,7 @@ from test.utils import revert_liger_kernel_to_gemma3_text
 from test.utils import revert_liger_kernel_to_glm4
 from test.utils import revert_liger_kernel_to_glm4v
 from test.utils import revert_liger_kernel_to_glm4v_moe
+from test.utils import revert_liger_kernel_to_gpt_oss
 from test.utils import revert_liger_kernel_to_granite
 from test.utils import revert_liger_kernel_to_hunyuan_v1
 from test.utils import revert_liger_kernel_to_hunyuan_v1_moe
@@ -247,6 +249,15 @@ try:
 except ImportError:
     QWEN3_AVAILABLE = False
 
+try:
+    # GPT-OSS is only available in transformers>=4.55.0
+    from transformers.models.gpt_oss.configuration_gpt_oss import GptOssConfig
+    from transformers.models.gpt_oss.modeling_gpt_oss import GptOssForCausalLM
+
+    GPT_OSS_AVAILABLE = True
+except ImportError:
+    GPT_OSS_AVAILABLE = False
+
 try:
     # InternVL is only available in transformers>=4.52.1
     from transformers.models.internvl.configuration_internvl import InternVLConfig
@@ -610,6 +621,43 @@ if QWEN3_AVAILABLE:
         ),
     )
 
+if GPT_OSS_AVAILABLE:
+    MINI_MODEL_SETUPS["mini_gpt_oss"] = MiniModelConfig(
+        liger_kernel_patch_func=apply_liger_kernel_to_gpt_oss,
+        liger_kernel_patch_revert_func=revert_liger_kernel_to_gpt_oss,
+        model_class=GptOssForCausalLM,
+        mini_model_config=GptOssConfig(
+            vocab_size=32000,  # 201088
+            hidden_size=896,
+            intermediate_size=896,  # Same as hidden_size for GPT-OSS
+            num_hidden_layers=4,
+            num_attention_heads=8,
+            num_key_value_heads=2,
+            head_dim=64,
+            hidden_act="silu",
+            max_position_embeddings=8192,
+            initializer_range=0.02,
+            rms_norm_eps=1e-5,
+            use_cache=True,
+            tie_word_embeddings=False,
+            rope_parameters={
+                "rope_type": "yarn",
+                "factor": 8.0,
+                "beta_fast": 32.0,
+                "beta_slow": 1.0,
+                "truncate": False,
+                "original_max_position_embeddings": 4096,
+            },
+            attention_dropout=0.0,
+            num_local_experts=8,  # Reduced from 32 for mini model
+            num_experts_per_tok=2,  # Reduced from 4 for mini model
+            router_aux_loss_coef=0.9,
+            output_router_logits=False,
+            sliding_window=128,
+            layer_types=["sliding_attention" if bool((i + 1) % 2) else "full_attention" for i in range(4)],
+        ),
+    )
+
 if GEMMA3_AVAILABLE:
     MINI_MODEL_SETUPS["mini_gemma3_text"] = MiniModelConfig(
         liger_kernel_patch_func=apply_liger_kernel_to_gemma3_text,
@@ -1608,6 +1656,22 @@ def run_mini_model(
                 reason="Qwen3 not available in this version of transformers",
             ),
         ),
+        pytest.param(
+            "mini_gpt_oss",
+            32,
+            1e-5,
+            torch.float32,
+            1e-8,
+            1e-5,
+            5e-3,
+            1e-5,
+            5e-3,
+            1e-5,
+            marks=pytest.mark.skipif(
+                not GPT_OSS_AVAILABLE,
+                reason="GPT-OSS not available in this version of transformers",
+            ),
+        ),
         pytest.param(  # qwen2_vl requires slightly larger tolerances to pass this test after bug fix to qwen2_vl in transformers v4.47.0
             "mini_qwen2_vl",
             32,

{liger_kernel_nightly-0.6.4.dev20251202094519 → liger_kernel_nightly-0.6.4.dev20251206103502}/test/utils.py

@@ -487,6 +487,18 @@ def revert_liger_kernel_to_qwen3_moe(model_config: MiniModelConfig):
     print("Liger kernel patches have been reverted.")
 
 
+def revert_liger_kernel_to_gpt_oss(model_config: MiniModelConfig):
+    """
+    Revert all Liger kernel patches applied to GPT-OSS.
+    """
+    from transformers.models.gpt_oss import modeling_gpt_oss
+
+    importlib.reload(modeling_gpt_oss)
+    model_config.model_class = modeling_gpt_oss.GptOssForCausalLM
+
+    print("Liger kernel patches have been reverted.")
+
+
 def revert_liger_kernel_to_qwen2_vl(model_config: MiniModelConfig):
     """
     Revert all Liger kernel patches applied to Qwen2-VL.