mirror of
https://github.com/linkedin/Liger-Kernel.git
synced 2026-10-12 02:17:47 +00:00
Add Liger Kernel support for Ministral models (#1166)
This PR was generated automatically by Claude Code with minimal human intervention. Claude uses the skill written in #1167 to generate this PR and test it out. Ministral (model_type: "ministral") is a separate model type in transformers v5+ with its own MinistralForCausalLM class. This adds full Liger Kernel optimization support including RoPE, RMSNorm, SwiGLU, CrossEntropyLoss, and FusedLinearCrossEntropy. ## Summary <!--- This is a required section; please describe the main purpose of this proposed code change. ---> - Adds full Liger Kernel optimization for Ministral models (model_type: "ministral") in transformers v5+ - Patches: RoPE, RMSNorm, SwiGLU, CrossEntropyLoss, FusedLinearCrossEntropy - All 5 tests (instance patching + 4 convergence) validated on H100 GPU <!--- ## Details This is an optional section; is there anything specific that reviewers should be aware of? ---> ## Testing Done <!--- This is a required section; please describe how this change was tested. ---> <!-- Replace BLANK with your device type. For example, A100-80G-PCIe Complete the following tasks before sending your PR, and replace `[ ]` with `[x]` to indicate you have done them. --> - Hardware Type: <BLANK> - [x] run `make test` to ensure correctness - [x] run `make checkstyle` to ensure code style - [x] run `make test-convergence` to ensure convergence --------- Co-authored-by: Claude Opus 4.6 (1M context) <[email protected]>
This commit is contained in:
Vaibhav Jindal
co-authored by
Claude Opus 4.6
Claude Opus 4.6
parent
eac9ce6ede
commit
8e6232a08f
@@ -250,6 +250,7 @@ loss.backward()
|
||||
| Llama4 (Text) & (Multimodal) | `liger_kernel.transformers.apply_liger_kernel_to_llama4` | RMSNorm, LayerNorm, GeGLU, CrossEntropyLoss, FusedLinearCrossEntropy |
|
||||
| LLaMA 2 & 3 | `liger_kernel.transformers.apply_liger_kernel_to_llama` | RoPE, RMSNorm, SwiGLU, CrossEntropyLoss, FusedLinearCrossEntropy |
|
||||
| LLaMA 3.2-Vision | `liger_kernel.transformers.apply_liger_kernel_to_mllama` | RoPE, RMSNorm, SwiGLU, CrossEntropyLoss, FusedLinearCrossEntropy |
|
||||
| Ministral | `liger_kernel.transformers.apply_liger_kernel_to_ministral` | RoPE, RMSNorm, SwiGLU, CrossEntropyLoss, FusedLinearCrossEntropy |
|
||||
| Mistral | `liger_kernel.transformers.apply_liger_kernel_to_mistral` | RoPE, RMSNorm, SwiGLU, CrossEntropyLoss, FusedLinearCrossEntropy |
|
||||
| Mixtral | `liger_kernel.transformers.apply_liger_kernel_to_mixtral` | RoPE, RMSNorm, SwiGLU, CrossEntropyLoss, FusedLinearCrossEntropy |
|
||||
| Nemotron | `liger_kernel.transformers.apply_liger_kernel_to_nemotron` | CrossEntropyLoss, FusedLinearCrossEntropy |
|
||||
|
||||
@@ -53,6 +53,7 @@ if TYPE_CHECKING:
|
||||
from liger_kernel.transformers.monkey_patch import apply_liger_kernel_to_llama # noqa: F401
|
||||
from liger_kernel.transformers.monkey_patch import apply_liger_kernel_to_llama4 # noqa: F401
|
||||
from liger_kernel.transformers.monkey_patch import apply_liger_kernel_to_llava # noqa: F401
|
||||
from liger_kernel.transformers.monkey_patch import apply_liger_kernel_to_ministral # noqa: F401
|
||||
from liger_kernel.transformers.monkey_patch import apply_liger_kernel_to_mistral # noqa: F401
|
||||
from liger_kernel.transformers.monkey_patch import apply_liger_kernel_to_mixtral # noqa: F401
|
||||
from liger_kernel.transformers.monkey_patch import apply_liger_kernel_to_mllama # noqa: F401
|
||||
@@ -125,6 +126,7 @@ def __getattr__(name: str):
|
||||
"apply_liger_kernel_to_llama",
|
||||
"apply_liger_kernel_to_llava",
|
||||
"apply_liger_kernel_to_llama4",
|
||||
"apply_liger_kernel_to_ministral",
|
||||
"apply_liger_kernel_to_mistral",
|
||||
"apply_liger_kernel_to_mixtral",
|
||||
"apply_liger_kernel_to_nemotron",
|
||||
@@ -210,6 +212,7 @@ if _TRANSFORMERS_AVAILABLE:
|
||||
"apply_liger_kernel_to_llama",
|
||||
"apply_liger_kernel_to_llava",
|
||||
"apply_liger_kernel_to_llama4",
|
||||
"apply_liger_kernel_to_ministral",
|
||||
"apply_liger_kernel_to_mistral",
|
||||
"apply_liger_kernel_to_mixtral",
|
||||
"apply_liger_kernel_to_nemotron",
|
||||
|
||||
@@ -0,0 +1,145 @@
|
||||
from typing import Optional
|
||||
from typing import Tuple
|
||||
from typing import Union
|
||||
|
||||
import torch
|
||||
|
||||
from transformers.cache_utils import Cache
|
||||
|
||||
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 LigerCausalLMOutputWithPast
|
||||
|
||||
|
||||
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[Cache] = 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,
|
||||
logits_to_keep: Union[int, torch.Tensor] = 0,
|
||||
skip_logits: Optional[bool] = None,
|
||||
**kwargs,
|
||||
) -> Union[Tuple, LigerCausalLMOutputWithPast]:
|
||||
r"""
|
||||
Copy paste Ministral's forward but replace torch cross entropy with liger fused linear cross entropy
|
||||
|
||||
|
||||
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]`.
|
||||
|
||||
logits_to_keep (`int` or `torch.Tensor`, *optional*):
|
||||
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).
|
||||
Returns:
|
||||
|
||||
Example:
|
||||
|
||||
```python
|
||||
>>> from transformers import AutoTokenizer, MinistralForCausalLM
|
||||
|
||||
>>> model = MinistralForCausalLM.from_pretrained("mistralai/Ministral-8B-Instruct-2410")
|
||||
>>> tokenizer = AutoTokenizer.from_pretrained("mistralai/Ministral-8B-Instruct-2410")
|
||||
|
||||
>>> 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,
|
||||
**kwargs,
|
||||
)
|
||||
|
||||
hidden_states = outputs[0]
|
||||
# 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)
|
||||
loss = None
|
||||
logits = None
|
||||
token_accuracy = None
|
||||
predicted_tokens = None
|
||||
|
||||
if skip_logits and labels is None and shift_labels is None:
|
||||
raise ValueError("skip_logits is True, but labels and shift_labels are None")
|
||||
|
||||
if skip_logits is None:
|
||||
skip_logits = self.training and (labels is not None or shift_labels is not None)
|
||||
|
||||
# Compute loss
|
||||
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, predicted_tokens = unpack_cross_entropy_result(result)
|
||||
|
||||
else:
|
||||
logits = self.lm_head(kept_hidden_states)
|
||||
|
||||
loss = None
|
||||
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.config.vocab_size,
|
||||
**kwargs,
|
||||
)
|
||||
|
||||
if not return_dict:
|
||||
output_tuple = (logits,) + outputs[1:]
|
||||
output = (loss,) + output_tuple if loss is not None else output_tuple
|
||||
output = output + (token_accuracy,) if token_accuracy is not None else output
|
||||
output = output + (predicted_tokens,) if predicted_tokens is not None else output
|
||||
return output
|
||||
|
||||
# Return custom output class with token_accuracy field
|
||||
return LigerCausalLMOutputWithPast(
|
||||
loss=loss,
|
||||
logits=logits,
|
||||
past_key_values=outputs.past_key_values,
|
||||
hidden_states=outputs.hidden_states,
|
||||
attentions=outputs.attentions,
|
||||
token_accuracy=token_accuracy,
|
||||
predicted_tokens=predicted_tokens,
|
||||
)
|
||||
@@ -21,6 +21,7 @@ from liger_kernel.transformers.model.gemma2 import lce_forward as gemma2_lce_for
|
||||
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.llava import lce_forward as llava_lce_forward
|
||||
from liger_kernel.transformers.model.ministral import lce_forward as ministral_lce_forward
|
||||
from liger_kernel.transformers.model.mistral import lce_forward as mistral_lce_forward
|
||||
from liger_kernel.transformers.model.mixtral import lce_forward as mixtral_lce_forward
|
||||
from liger_kernel.transformers.model.nemotron import lce_forward as nemotron_lce_forward
|
||||
@@ -619,6 +620,69 @@ def apply_liger_kernel_to_mllama(
|
||||
_patch_layer_norm_module(layer.post_attention_layernorm)
|
||||
|
||||
|
||||
def apply_liger_kernel_to_ministral(
|
||||
rope: bool = True,
|
||||
cross_entropy: bool = False,
|
||||
fused_linear_cross_entropy: bool = True,
|
||||
rms_norm: bool = True,
|
||||
swiglu: bool = True,
|
||||
model: PreTrainedModel = None,
|
||||
) -> None:
|
||||
"""
|
||||
Apply Liger kernels to replace original implementation in HuggingFace Ministral models
|
||||
|
||||
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 True.
|
||||
model (PreTrainedModel): The model instance to apply Liger kernels to, if the model has already been
|
||||
loaded. Default is None.
|
||||
"""
|
||||
assert not (cross_entropy and fused_linear_cross_entropy), (
|
||||
"cross_entropy and fused_linear_cross_entropy cannot both be True."
|
||||
)
|
||||
|
||||
from transformers.models.ministral import modeling_ministral
|
||||
from transformers.models.ministral.modeling_ministral import MinistralModel
|
||||
|
||||
if rope:
|
||||
modeling_ministral.apply_rotary_pos_emb = liger_rotary_pos_emb
|
||||
if rms_norm:
|
||||
modeling_ministral.MinistralRMSNorm = LigerRMSNorm
|
||||
if cross_entropy:
|
||||
modeling_ministral.CrossEntropyLoss = LigerCrossEntropyLoss
|
||||
if fused_linear_cross_entropy:
|
||||
if model is not None:
|
||||
model.forward = MethodType(ministral_lce_forward, model)
|
||||
else:
|
||||
modeling_ministral.MinistralForCausalLM.forward = ministral_lce_forward
|
||||
|
||||
if swiglu:
|
||||
modeling_ministral.MinistralMLP = LigerSwiGLUMLP
|
||||
|
||||
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: MinistralModel = 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 swiglu:
|
||||
_patch_swiglu_module(decoder_layer.mlp, LigerSwiGLUMLP)
|
||||
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_mistral(
|
||||
rope: bool = True,
|
||||
cross_entropy: bool = False,
|
||||
@@ -3120,6 +3184,7 @@ MODEL_TYPE_TO_APPLY_LIGER_FN = {
|
||||
"granite": apply_liger_kernel_to_granite,
|
||||
"mllama": apply_liger_kernel_to_mllama,
|
||||
"mllama_text_model": apply_liger_kernel_to_mllama,
|
||||
"ministral": apply_liger_kernel_to_ministral,
|
||||
"mistral": apply_liger_kernel_to_mistral,
|
||||
"mixtral": apply_liger_kernel_to_mixtral,
|
||||
"nemotron": apply_liger_kernel_to_nemotron,
|
||||
|
||||
@@ -40,6 +40,7 @@ from liger_kernel.transformers import apply_liger_kernel_to_internvl
|
||||
from liger_kernel.transformers import apply_liger_kernel_to_llama
|
||||
from liger_kernel.transformers import apply_liger_kernel_to_llama4
|
||||
from liger_kernel.transformers import apply_liger_kernel_to_llava
|
||||
from liger_kernel.transformers import apply_liger_kernel_to_ministral
|
||||
from liger_kernel.transformers import apply_liger_kernel_to_mistral
|
||||
from liger_kernel.transformers import apply_liger_kernel_to_mixtral
|
||||
from liger_kernel.transformers import apply_liger_kernel_to_mllama
|
||||
@@ -81,6 +82,7 @@ from test.utils import revert_liger_kernel_to_internvl
|
||||
from test.utils import revert_liger_kernel_to_llama
|
||||
from test.utils import revert_liger_kernel_to_llama4
|
||||
from test.utils import revert_liger_kernel_to_llava
|
||||
from test.utils import revert_liger_kernel_to_ministral
|
||||
from test.utils import revert_liger_kernel_to_mistral
|
||||
from test.utils import revert_liger_kernel_to_mixtral
|
||||
from test.utils import revert_liger_kernel_to_mllama
|
||||
@@ -113,6 +115,14 @@ try:
|
||||
except ImportError:
|
||||
LLAMA4_AVAILABLE = False
|
||||
|
||||
try:
|
||||
from transformers.models.ministral.configuration_ministral import MinistralConfig
|
||||
from transformers.models.ministral.modeling_ministral import MinistralForCausalLM
|
||||
|
||||
MINISTRAL_AVAILABLE = True
|
||||
except ImportError:
|
||||
MINISTRAL_AVAILABLE = False
|
||||
|
||||
try:
|
||||
# Mllama is only available in transformers>=4.45.0
|
||||
from transformers.models.mllama.configuration_mllama import MllamaTextConfig
|
||||
@@ -565,6 +575,33 @@ MINI_MODEL_SETUPS = {
|
||||
),
|
||||
}
|
||||
|
||||
if MINISTRAL_AVAILABLE:
|
||||
MINI_MODEL_SETUPS["mini_ministral"] = MiniModelConfig(
|
||||
liger_kernel_patch_func=apply_liger_kernel_to_ministral,
|
||||
liger_kernel_patch_revert_func=revert_liger_kernel_to_ministral,
|
||||
model_class=MinistralForCausalLM,
|
||||
mini_model_config=MinistralConfig(
|
||||
attention_dropout=0.0,
|
||||
bos_token_id=1,
|
||||
eos_token_id=2,
|
||||
head_dim=128,
|
||||
hidden_act="silu",
|
||||
hidden_size=1024,
|
||||
initializer_range=0.02,
|
||||
intermediate_size=2048,
|
||||
max_position_embeddings=32768,
|
||||
num_attention_heads=8,
|
||||
num_hidden_layers=4,
|
||||
num_key_value_heads=2,
|
||||
rms_norm_eps=1e-5,
|
||||
sliding_window=4096,
|
||||
tie_word_embeddings=False,
|
||||
use_cache=True,
|
||||
vocab_size=32000,
|
||||
attn_implementation="sdpa",
|
||||
),
|
||||
)
|
||||
|
||||
if LLAMA4_AVAILABLE:
|
||||
MINI_MODEL_SETUPS["mini_llama4"] = MiniModelConfig(
|
||||
liger_kernel_patch_func=apply_liger_kernel_to_llama4,
|
||||
@@ -1963,6 +2000,24 @@ def run_mini_model(
|
||||
pytest.mark.skipif(not supports_bfloat16(), reason="bfloat16 not supported on this GPU"),
|
||||
],
|
||||
),
|
||||
pytest.param(
|
||||
"mini_ministral",
|
||||
32,
|
||||
1e-5,
|
||||
torch.bfloat16,
|
||||
5e-2,
|
||||
5e-2,
|
||||
1e-1,
|
||||
1e-2,
|
||||
1e-2,
|
||||
1e-2,
|
||||
marks=[
|
||||
pytest.mark.skipif(not supports_bfloat16(), reason="bfloat16 not supported on this GPU"),
|
||||
pytest.mark.skipif(
|
||||
not MINISTRAL_AVAILABLE, reason="Ministral not available in this version of transformers"
|
||||
),
|
||||
],
|
||||
),
|
||||
pytest.param(
|
||||
"mini_olmo2",
|
||||
32,
|
||||
|
||||
@@ -39,6 +39,7 @@ from liger_kernel.transformers import apply_liger_kernel_to_internvl
|
||||
from liger_kernel.transformers import apply_liger_kernel_to_llama
|
||||
from liger_kernel.transformers import apply_liger_kernel_to_llama4
|
||||
from liger_kernel.transformers import apply_liger_kernel_to_llava
|
||||
from liger_kernel.transformers import apply_liger_kernel_to_ministral
|
||||
from liger_kernel.transformers import apply_liger_kernel_to_mistral
|
||||
from liger_kernel.transformers import apply_liger_kernel_to_mixtral
|
||||
from liger_kernel.transformers import apply_liger_kernel_to_mllama
|
||||
@@ -78,6 +79,7 @@ from test.utils import revert_liger_kernel_to_internvl
|
||||
from test.utils import revert_liger_kernel_to_llama
|
||||
from test.utils import revert_liger_kernel_to_llama4
|
||||
from test.utils import revert_liger_kernel_to_llava
|
||||
from test.utils import revert_liger_kernel_to_ministral
|
||||
from test.utils import revert_liger_kernel_to_mistral
|
||||
from test.utils import revert_liger_kernel_to_mixtral
|
||||
from test.utils import revert_liger_kernel_to_mllama
|
||||
@@ -101,6 +103,14 @@ from test.utils import supports_bfloat16
|
||||
|
||||
IS_TRANSFORMERS_V5_OR_LATER = version.parse(transformers.__version__) >= version.parse("5.0.0")
|
||||
|
||||
try:
|
||||
from transformers.models.ministral.configuration_ministral import MinistralConfig
|
||||
from transformers.models.ministral.modeling_ministral import MinistralForCausalLM
|
||||
|
||||
MINISTRAL_AVAILABLE = True
|
||||
except ImportError:
|
||||
MINISTRAL_AVAILABLE = False
|
||||
|
||||
try:
|
||||
from transformers.models.llama4.configuration_llama4 import Llama4TextConfig
|
||||
from transformers.models.llama4.modeling_llama4 import Llama4ForCausalLM
|
||||
@@ -424,22 +434,26 @@ MINI_MODEL_SETUPS = {
|
||||
attn_implementation="sdpa",
|
||||
),
|
||||
),
|
||||
"mini_mixtral": MiniModelConfig(
|
||||
liger_kernel_patch_func=apply_liger_kernel_to_mixtral,
|
||||
liger_kernel_patch_revert_func=revert_liger_kernel_to_mixtral,
|
||||
model_class=MixtralForCausalLM,
|
||||
mini_model_config=MixtralConfig(
|
||||
}
|
||||
|
||||
if MINISTRAL_AVAILABLE:
|
||||
MINI_MODEL_SETUPS["mini_ministral"] = MiniModelConfig(
|
||||
liger_kernel_patch_func=apply_liger_kernel_to_ministral,
|
||||
liger_kernel_patch_revert_func=revert_liger_kernel_to_ministral,
|
||||
model_class=MinistralForCausalLM,
|
||||
mini_model_config=MinistralConfig(
|
||||
attention_dropout=0.0,
|
||||
bos_token_id=1,
|
||||
eos_token_id=2,
|
||||
head_dim=128,
|
||||
hidden_act="silu",
|
||||
hidden_size=512, # 4096
|
||||
hidden_size=1024,
|
||||
initializer_range=0.02,
|
||||
intermediate_size=2048, # 14336
|
||||
max_position_embeddings=32768, # 32768
|
||||
num_attention_heads=8, # 32
|
||||
num_hidden_layers=4, # 32
|
||||
num_key_value_heads=2, # 8
|
||||
intermediate_size=2048,
|
||||
max_position_embeddings=32768,
|
||||
num_attention_heads=8,
|
||||
num_hidden_layers=4,
|
||||
num_key_value_heads=2,
|
||||
rms_norm_eps=1e-5,
|
||||
sliding_window=4096,
|
||||
tie_word_embeddings=False,
|
||||
@@ -447,94 +461,122 @@ MINI_MODEL_SETUPS = {
|
||||
vocab_size=32000,
|
||||
attn_implementation="sdpa",
|
||||
),
|
||||
),
|
||||
"mini_gemma1": MiniModelConfig(
|
||||
liger_kernel_patch_func=apply_liger_kernel_to_gemma,
|
||||
liger_kernel_patch_revert_func=revert_liger_kernel_to_gemma,
|
||||
model_class=GemmaForCausalLM,
|
||||
mini_model_config=GemmaConfig(
|
||||
vocab_size=32000, # 256000
|
||||
hidden_size=1024, # 3072
|
||||
intermediate_size=2048, # 24576
|
||||
num_hidden_layers=4, # 28
|
||||
num_attention_heads=4, # 16
|
||||
num_key_value_heads=4, # 16
|
||||
head_dim=256,
|
||||
# gemma1 model config uses `hidden_act` and point it to gelu,
|
||||
# https://huggingface.co/google/gemma-7b/blob/main/config.json#L10
|
||||
# but in reality it's ignored and HuggingFace will use tanh approximation:
|
||||
# https://github.com/huggingface/transformers/blob/v4.40.1/src/transformers/models/gemma/modeling_gemma.py#L175
|
||||
hidden_act="gelu",
|
||||
max_position_embeddings=8192,
|
||||
initializer_range=0.02,
|
||||
rms_norm_eps=1e-06,
|
||||
use_cache=True,
|
||||
pad_token_id=0,
|
||||
# Special token ids/vocab size to match Mistral-7B tokenizer used to create the tokenized dataset
|
||||
# https://huggingface.co/mistralai/Mistral-7B-v0.1/blob/main/config.json
|
||||
bos_token_id=1, # 128000
|
||||
eos_token_id=2, # 128001
|
||||
tie_word_embeddings=True,
|
||||
attention_bias=False,
|
||||
attention_dropout=0.0,
|
||||
)
|
||||
|
||||
MINI_MODEL_SETUPS.update(
|
||||
{
|
||||
"mini_mixtral": MiniModelConfig(
|
||||
liger_kernel_patch_func=apply_liger_kernel_to_mixtral,
|
||||
liger_kernel_patch_revert_func=revert_liger_kernel_to_mixtral,
|
||||
model_class=MixtralForCausalLM,
|
||||
mini_model_config=MixtralConfig(
|
||||
attention_dropout=0.0,
|
||||
bos_token_id=1,
|
||||
eos_token_id=2,
|
||||
hidden_act="silu",
|
||||
hidden_size=512, # 4096
|
||||
initializer_range=0.02,
|
||||
intermediate_size=2048, # 14336
|
||||
max_position_embeddings=32768, # 32768
|
||||
num_attention_heads=8, # 32
|
||||
num_hidden_layers=4, # 32
|
||||
num_key_value_heads=2, # 8
|
||||
rms_norm_eps=1e-5,
|
||||
sliding_window=4096,
|
||||
tie_word_embeddings=False,
|
||||
use_cache=True,
|
||||
vocab_size=32000,
|
||||
attn_implementation="sdpa",
|
||||
),
|
||||
),
|
||||
),
|
||||
"mini_gemma1.1": MiniModelConfig(
|
||||
liger_kernel_patch_func=apply_liger_kernel_to_gemma,
|
||||
liger_kernel_patch_revert_func=revert_liger_kernel_to_gemma,
|
||||
model_class=GemmaForCausalLM,
|
||||
mini_model_config=GemmaConfig(
|
||||
vocab_size=32000, # 256000
|
||||
hidden_size=1024, # 3072
|
||||
intermediate_size=2048, # 24576
|
||||
num_hidden_layers=4, # 28
|
||||
num_attention_heads=4, # 16
|
||||
num_key_value_heads=4, # 16
|
||||
head_dim=256,
|
||||
hidden_activation="gelu_pytorch_tanh",
|
||||
max_position_embeddings=8192,
|
||||
initializer_range=0.02,
|
||||
rms_norm_eps=1e-06,
|
||||
use_cache=True,
|
||||
pad_token_id=0,
|
||||
# Special token ids/vocab size to match Mistral-7B tokenizer used to create the tokenized dataset
|
||||
# https://huggingface.co/mistralai/Mistral-7B-v0.1/blob/main/config.json
|
||||
bos_token_id=1, # 128000
|
||||
eos_token_id=2, # 128001
|
||||
tie_word_embeddings=True,
|
||||
attention_bias=False,
|
||||
attention_dropout=0.0,
|
||||
"mini_gemma1": MiniModelConfig(
|
||||
liger_kernel_patch_func=apply_liger_kernel_to_gemma,
|
||||
liger_kernel_patch_revert_func=revert_liger_kernel_to_gemma,
|
||||
model_class=GemmaForCausalLM,
|
||||
mini_model_config=GemmaConfig(
|
||||
vocab_size=32000, # 256000
|
||||
hidden_size=1024, # 3072
|
||||
intermediate_size=2048, # 24576
|
||||
num_hidden_layers=4, # 28
|
||||
num_attention_heads=4, # 16
|
||||
num_key_value_heads=4, # 16
|
||||
head_dim=256,
|
||||
# gemma1 model config uses `hidden_act` and point it to gelu,
|
||||
# https://huggingface.co/google/gemma-7b/blob/main/config.json#L10
|
||||
# but in reality it's ignored and HuggingFace will use tanh approximation:
|
||||
# https://github.com/huggingface/transformers/blob/v4.40.1/src/transformers/models/gemma/modeling_gemma.py#L175
|
||||
hidden_act="gelu",
|
||||
max_position_embeddings=8192,
|
||||
initializer_range=0.02,
|
||||
rms_norm_eps=1e-06,
|
||||
use_cache=True,
|
||||
pad_token_id=0,
|
||||
# Special token ids/vocab size to match Mistral-7B tokenizer used to create the tokenized dataset
|
||||
# https://huggingface.co/mistralai/Mistral-7B-v0.1/blob/main/config.json
|
||||
bos_token_id=1, # 128000
|
||||
eos_token_id=2, # 128001
|
||||
tie_word_embeddings=True,
|
||||
attention_bias=False,
|
||||
attention_dropout=0.0,
|
||||
),
|
||||
),
|
||||
),
|
||||
"mini_gemma2": MiniModelConfig(
|
||||
liger_kernel_patch_func=apply_liger_kernel_to_gemma2,
|
||||
liger_kernel_patch_revert_func=revert_liger_kernel_to_gemma2,
|
||||
model_class=Gemma2ForCausalLM,
|
||||
mini_model_config=Gemma2Config(
|
||||
vocab_size=32000, # 256000
|
||||
hidden_size=1024, # 3072
|
||||
intermediate_size=2048, # 24576
|
||||
num_hidden_layers=4, # 28
|
||||
num_attention_heads=4, # 16
|
||||
num_key_value_heads=4, # 16
|
||||
head_dim=256,
|
||||
hidden_activation="gelu_pytorch_tanh",
|
||||
max_position_embeddings=8192,
|
||||
initializer_range=0.02,
|
||||
rms_norm_eps=1e-06,
|
||||
use_cache=True,
|
||||
pad_token_id=0,
|
||||
# Special token ids/vocab size to match Mistral-7B tokenizer used to create the tokenized dataset
|
||||
# https://huggingface.co/mistralai/Mistral-7B-v0.1/blob/main/config.json
|
||||
bos_token_id=1, # 128000
|
||||
eos_token_id=2, # 128001
|
||||
tie_word_embeddings=True,
|
||||
attention_bias=False,
|
||||
attention_dropout=0.0,
|
||||
attn_implementation="eager",
|
||||
"mini_gemma1.1": MiniModelConfig(
|
||||
liger_kernel_patch_func=apply_liger_kernel_to_gemma,
|
||||
liger_kernel_patch_revert_func=revert_liger_kernel_to_gemma,
|
||||
model_class=GemmaForCausalLM,
|
||||
mini_model_config=GemmaConfig(
|
||||
vocab_size=32000, # 256000
|
||||
hidden_size=1024, # 3072
|
||||
intermediate_size=2048, # 24576
|
||||
num_hidden_layers=4, # 28
|
||||
num_attention_heads=4, # 16
|
||||
num_key_value_heads=4, # 16
|
||||
head_dim=256,
|
||||
hidden_activation="gelu_pytorch_tanh",
|
||||
max_position_embeddings=8192,
|
||||
initializer_range=0.02,
|
||||
rms_norm_eps=1e-06,
|
||||
use_cache=True,
|
||||
pad_token_id=0,
|
||||
# Special token ids/vocab size to match Mistral-7B tokenizer used to create the tokenized dataset
|
||||
# https://huggingface.co/mistralai/Mistral-7B-v0.1/blob/main/config.json
|
||||
bos_token_id=1, # 128000
|
||||
eos_token_id=2, # 128001
|
||||
tie_word_embeddings=True,
|
||||
attention_bias=False,
|
||||
attention_dropout=0.0,
|
||||
),
|
||||
),
|
||||
),
|
||||
}
|
||||
"mini_gemma2": MiniModelConfig(
|
||||
liger_kernel_patch_func=apply_liger_kernel_to_gemma2,
|
||||
liger_kernel_patch_revert_func=revert_liger_kernel_to_gemma2,
|
||||
model_class=Gemma2ForCausalLM,
|
||||
mini_model_config=Gemma2Config(
|
||||
vocab_size=32000, # 256000
|
||||
hidden_size=1024, # 3072
|
||||
intermediate_size=2048, # 24576
|
||||
num_hidden_layers=4, # 28
|
||||
num_attention_heads=4, # 16
|
||||
num_key_value_heads=4, # 16
|
||||
head_dim=256,
|
||||
hidden_activation="gelu_pytorch_tanh",
|
||||
max_position_embeddings=8192,
|
||||
initializer_range=0.02,
|
||||
rms_norm_eps=1e-06,
|
||||
use_cache=True,
|
||||
pad_token_id=0,
|
||||
# Special token ids/vocab size to match Mistral-7B tokenizer used to create the tokenized dataset
|
||||
# https://huggingface.co/mistralai/Mistral-7B-v0.1/blob/main/config.json
|
||||
bos_token_id=1, # 128000
|
||||
eos_token_id=2, # 128001
|
||||
tie_word_embeddings=True,
|
||||
attention_bias=False,
|
||||
attention_dropout=0.0,
|
||||
attn_implementation="eager",
|
||||
),
|
||||
),
|
||||
}
|
||||
)
|
||||
|
||||
if LLAMA4_AVAILABLE:
|
||||
MINI_MODEL_SETUPS["mini_llama4"] = MiniModelConfig(
|
||||
@@ -1822,6 +1864,24 @@ def run_mini_model(
|
||||
1e-2,
|
||||
marks=pytest.mark.skipif(not supports_bfloat16(), reason="bfloat16 not supported on this GPU"),
|
||||
),
|
||||
pytest.param(
|
||||
"mini_ministral",
|
||||
32,
|
||||
1e-5,
|
||||
torch.bfloat16,
|
||||
1e-2,
|
||||
5e-2,
|
||||
1e-1,
|
||||
1e-2,
|
||||
1e-2,
|
||||
1e-2,
|
||||
marks=[
|
||||
pytest.mark.skipif(not supports_bfloat16(), reason="bfloat16 not supported on this GPU"),
|
||||
pytest.mark.skipif(
|
||||
not MINISTRAL_AVAILABLE, reason="Ministral not available in this version of transformers"
|
||||
),
|
||||
],
|
||||
),
|
||||
# TODO: mixtral is flaky so disable the test for now
|
||||
# pytest.param(
|
||||
# "mini_mixtral",
|
||||
|
||||
@@ -40,6 +40,7 @@ from liger_kernel.transformers import apply_liger_kernel_to_internvl
|
||||
from liger_kernel.transformers import apply_liger_kernel_to_llama
|
||||
from liger_kernel.transformers import apply_liger_kernel_to_llama4
|
||||
from liger_kernel.transformers import apply_liger_kernel_to_llava
|
||||
from liger_kernel.transformers import apply_liger_kernel_to_ministral
|
||||
from liger_kernel.transformers import apply_liger_kernel_to_mistral
|
||||
from liger_kernel.transformers import apply_liger_kernel_to_mixtral
|
||||
from liger_kernel.transformers import apply_liger_kernel_to_mllama
|
||||
@@ -81,6 +82,7 @@ from test.utils import revert_liger_kernel_to_internvl
|
||||
from test.utils import revert_liger_kernel_to_llama
|
||||
from test.utils import revert_liger_kernel_to_llama4
|
||||
from test.utils import revert_liger_kernel_to_llava
|
||||
from test.utils import revert_liger_kernel_to_ministral
|
||||
from test.utils import revert_liger_kernel_to_mistral
|
||||
from test.utils import revert_liger_kernel_to_mixtral
|
||||
from test.utils import revert_liger_kernel_to_mllama
|
||||
@@ -121,6 +123,14 @@ try:
|
||||
except ImportError:
|
||||
MLLAMA_AVAILABLE = False
|
||||
|
||||
try:
|
||||
from transformers.models.ministral.configuration_ministral import MinistralConfig
|
||||
from transformers.models.ministral.modeling_ministral import MinistralForCausalLM
|
||||
|
||||
MINISTRAL_AVAILABLE = True
|
||||
except ImportError:
|
||||
MINISTRAL_AVAILABLE = False
|
||||
|
||||
try:
|
||||
# Qwen2-VL is only available in transformers>4.52.4
|
||||
import transformers
|
||||
@@ -562,6 +572,32 @@ MINI_MODEL_SETUPS = {
|
||||
),
|
||||
),
|
||||
}
|
||||
if MINISTRAL_AVAILABLE:
|
||||
MINI_MODEL_SETUPS["mini_ministral"] = MiniModelConfig(
|
||||
liger_kernel_patch_func=apply_liger_kernel_to_ministral,
|
||||
liger_kernel_patch_revert_func=revert_liger_kernel_to_ministral,
|
||||
model_class=MinistralForCausalLM,
|
||||
mini_model_config=MinistralConfig(
|
||||
attention_dropout=0.0,
|
||||
bos_token_id=1,
|
||||
eos_token_id=2,
|
||||
head_dim=128,
|
||||
hidden_act="silu",
|
||||
hidden_size=1024,
|
||||
initializer_range=0.02,
|
||||
intermediate_size=2048,
|
||||
max_position_embeddings=32768,
|
||||
num_attention_heads=8,
|
||||
num_hidden_layers=4,
|
||||
num_key_value_heads=2,
|
||||
rms_norm_eps=1e-5,
|
||||
sliding_window=4096,
|
||||
tie_word_embeddings=False,
|
||||
use_cache=True,
|
||||
vocab_size=32000,
|
||||
attn_implementation="sdpa",
|
||||
),
|
||||
)
|
||||
if LLAMA4_AVAILABLE:
|
||||
MINI_MODEL_SETUPS["mini_llama4"] = MiniModelConfig(
|
||||
liger_kernel_patch_func=apply_liger_kernel_to_llama4,
|
||||
@@ -1965,6 +2001,23 @@ def run_mini_model(
|
||||
1e-5,
|
||||
marks=[],
|
||||
),
|
||||
pytest.param(
|
||||
"mini_ministral",
|
||||
32,
|
||||
1e-4,
|
||||
torch.float32,
|
||||
1e-8,
|
||||
1e-5,
|
||||
5e-3,
|
||||
1e-5,
|
||||
5e-3,
|
||||
1e-5,
|
||||
marks=[
|
||||
pytest.mark.skipif(
|
||||
not MINISTRAL_AVAILABLE, reason="Ministral not available in this version of transformers"
|
||||
),
|
||||
],
|
||||
),
|
||||
# TODO: mixtral is flaky so disable the test for now
|
||||
# ("mini_mixtral", 32, 1e-4, torch.float32, 5e-4, 1e-4, 5e-3, 1e-5, 1e-2, 1e-5),
|
||||
# Gemma 1.1 and 2 has more tolerance because currently, the kernel is not a perfect match (casts are not done the same way)
|
||||
|
||||
@@ -39,6 +39,7 @@ from liger_kernel.transformers import apply_liger_kernel_to_internvl
|
||||
from liger_kernel.transformers import apply_liger_kernel_to_llama
|
||||
from liger_kernel.transformers import apply_liger_kernel_to_llama4
|
||||
from liger_kernel.transformers import apply_liger_kernel_to_llava
|
||||
from liger_kernel.transformers import apply_liger_kernel_to_ministral
|
||||
from liger_kernel.transformers import apply_liger_kernel_to_mistral
|
||||
from liger_kernel.transformers import apply_liger_kernel_to_mixtral
|
||||
from liger_kernel.transformers import apply_liger_kernel_to_mllama
|
||||
@@ -78,6 +79,7 @@ from test.utils import revert_liger_kernel_to_internvl
|
||||
from test.utils import revert_liger_kernel_to_llama
|
||||
from test.utils import revert_liger_kernel_to_llama4
|
||||
from test.utils import revert_liger_kernel_to_llava
|
||||
from test.utils import revert_liger_kernel_to_ministral
|
||||
from test.utils import revert_liger_kernel_to_mistral
|
||||
from test.utils import revert_liger_kernel_to_mixtral
|
||||
from test.utils import revert_liger_kernel_to_mllama
|
||||
@@ -108,6 +110,14 @@ try:
|
||||
except ImportError:
|
||||
LLAMA4_AVAILABLE = False
|
||||
|
||||
try:
|
||||
from transformers.models.ministral.configuration_ministral import MinistralConfig
|
||||
from transformers.models.ministral.modeling_ministral import MinistralForCausalLM
|
||||
|
||||
MINISTRAL_AVAILABLE = True
|
||||
except ImportError:
|
||||
MINISTRAL_AVAILABLE = False
|
||||
|
||||
try:
|
||||
# Mllama is only available in transformers>=4.45.0
|
||||
from transformers.models.mllama.configuration_mllama import MllamaTextConfig
|
||||
@@ -555,6 +565,34 @@ MINI_MODEL_SETUPS = {
|
||||
),
|
||||
),
|
||||
}
|
||||
|
||||
if MINISTRAL_AVAILABLE:
|
||||
MINI_MODEL_SETUPS["mini_ministral"] = MiniModelConfig(
|
||||
liger_kernel_patch_func=apply_liger_kernel_to_ministral,
|
||||
liger_kernel_patch_revert_func=revert_liger_kernel_to_ministral,
|
||||
model_class=MinistralForCausalLM,
|
||||
mini_model_config=MinistralConfig(
|
||||
attention_dropout=0.0,
|
||||
bos_token_id=1,
|
||||
eos_token_id=2,
|
||||
head_dim=128,
|
||||
hidden_act="silu",
|
||||
hidden_size=1024,
|
||||
initializer_range=0.02,
|
||||
intermediate_size=2048,
|
||||
max_position_embeddings=32768,
|
||||
num_attention_heads=8,
|
||||
num_hidden_layers=4,
|
||||
num_key_value_heads=2,
|
||||
rms_norm_eps=1e-5,
|
||||
sliding_window=4096,
|
||||
tie_word_embeddings=False,
|
||||
use_cache=True,
|
||||
vocab_size=32000,
|
||||
attn_implementation="sdpa",
|
||||
),
|
||||
)
|
||||
|
||||
if LLAMA4_AVAILABLE:
|
||||
MINI_MODEL_SETUPS["mini_llama4"] = MiniModelConfig(
|
||||
liger_kernel_patch_func=apply_liger_kernel_to_llama4,
|
||||
@@ -1851,6 +1889,21 @@ def run_mini_model(
|
||||
),
|
||||
("mini_phi3", 32, 1e-4, torch.float32, 1e-8, 1e-5, 5e-3, 1e-5, 5e-3, 1e-5),
|
||||
("mini_mistral", 32, 1e-4, torch.float32, 1e-8, 1e-5, 5e-3, 1e-5, 5e-3, 1e-5),
|
||||
pytest.param(
|
||||
"mini_ministral",
|
||||
32,
|
||||
1e-4,
|
||||
torch.float32,
|
||||
1e-8,
|
||||
1e-5,
|
||||
5e-3,
|
||||
1e-5,
|
||||
5e-3,
|
||||
1e-5,
|
||||
marks=pytest.mark.skipif(
|
||||
not MINISTRAL_AVAILABLE, reason="Ministral not available in this version of transformers"
|
||||
),
|
||||
),
|
||||
# TODO: mixtral is flaky so disable the test for now
|
||||
# ("mini_mixtral", 32, 1e-4, torch.float32, 5e-4, 1e-4, 5e-3, 1e-5, 1e-2, 1e-5),
|
||||
# Gemma 1.1 and 2 has more tolerance because currently, the kernel is not a perfect match
|
||||
|
||||
@@ -29,6 +29,7 @@ from liger_kernel.transformers.model.falcon_h1 import lce_forward as falcon_h1_l
|
||||
from liger_kernel.transformers.model.gemma import lce_forward as gemma_lce_forward
|
||||
from liger_kernel.transformers.model.gemma2 import lce_forward as gemma2_lce_forward
|
||||
from liger_kernel.transformers.model.llama import lce_forward as llama_lce_forward
|
||||
from liger_kernel.transformers.model.ministral import lce_forward as ministral_lce_forward
|
||||
from liger_kernel.transformers.model.mistral import lce_forward as mistral_lce_forward
|
||||
from liger_kernel.transformers.model.mixtral import lce_forward as mixtral_lce_forward
|
||||
from liger_kernel.transformers.model.mllama import lce_forward as mllama_lce_forward
|
||||
@@ -92,6 +93,15 @@ def is_llama4_available():
|
||||
return False
|
||||
|
||||
|
||||
def is_ministral_available():
|
||||
try:
|
||||
import transformers.models.ministral # noqa: F401
|
||||
|
||||
return True
|
||||
except ImportError:
|
||||
return False
|
||||
|
||||
|
||||
def is_qwen3_available():
|
||||
try:
|
||||
import transformers.models.qwen3 # noqa: F401
|
||||
@@ -266,6 +276,7 @@ def test_import_from_root():
|
||||
from liger_kernel.transformers import apply_liger_kernel_to_glm4v_moe # noqa: F401
|
||||
from liger_kernel.transformers import apply_liger_kernel_to_internvl # noqa: F401
|
||||
from liger_kernel.transformers import apply_liger_kernel_to_llama # noqa: F401
|
||||
from liger_kernel.transformers import apply_liger_kernel_to_ministral # noqa: F401
|
||||
from liger_kernel.transformers import apply_liger_kernel_to_mistral # noqa: F401
|
||||
from liger_kernel.transformers import apply_liger_kernel_to_mixtral # noqa: F401
|
||||
from liger_kernel.transformers import apply_liger_kernel_to_mllama # noqa: F401
|
||||
@@ -1470,6 +1481,49 @@ def test_apply_liger_kernel_to_instance_for_mistral():
|
||||
pytest.fail(f"An exception occured in extra_expr: {type(e).__name__} - {e}")
|
||||
|
||||
|
||||
@pytest.mark.skipif(not is_ministral_available(), reason="ministral module not available")
|
||||
def test_apply_liger_kernel_to_instance_for_ministral():
|
||||
# Ensure any monkey patching is cleaned up for subsequent tests
|
||||
with patch("transformers.models.ministral.modeling_ministral"):
|
||||
# Instantiate a dummy model
|
||||
config = transformers.models.ministral.configuration_ministral.MinistralConfig(
|
||||
dtype=torch.bfloat16,
|
||||
rms_norm_eps=1e-5,
|
||||
hidden_size=32,
|
||||
intermediate_size=64,
|
||||
hidden_act="silu",
|
||||
num_hidden_layers=2,
|
||||
head_dim=16,
|
||||
num_attention_heads=2,
|
||||
num_key_value_heads=2,
|
||||
)
|
||||
dummy_model_instance = AutoModelForCausalLM.from_config(config)
|
||||
|
||||
# Check that model instance variables are not yet patched with Liger modules
|
||||
assert inspect.getsource(dummy_model_instance.forward) != inspect.getsource(ministral_lce_forward)
|
||||
assert inspect.getsource(dummy_model_instance.model.norm.forward) != inspect.getsource(LigerRMSNorm.forward)
|
||||
for layer in dummy_model_instance.model.layers:
|
||||
assert inspect.getsource(layer.mlp.forward) != inspect.getsource(LigerSwiGLUMLP.forward)
|
||||
assert inspect.getsource(layer.input_layernorm.forward) != inspect.getsource(LigerRMSNorm.forward)
|
||||
assert inspect.getsource(layer.post_attention_layernorm.forward) != inspect.getsource(LigerRMSNorm.forward)
|
||||
|
||||
# Test applying kernels to the model instance
|
||||
_apply_liger_kernel_to_instance(model=dummy_model_instance)
|
||||
|
||||
# Check that the model's instance variables were correctly patched with Liger modules
|
||||
assert inspect.getsource(dummy_model_instance.forward) == inspect.getsource(ministral_lce_forward)
|
||||
assert inspect.getsource(dummy_model_instance.model.norm.forward) == inspect.getsource(LigerRMSNorm.forward)
|
||||
for layer in dummy_model_instance.model.layers:
|
||||
assert inspect.getsource(layer.mlp.forward) == inspect.getsource(LigerSwiGLUMLP.forward)
|
||||
assert inspect.getsource(layer.input_layernorm.forward) == inspect.getsource(LigerRMSNorm.forward)
|
||||
assert inspect.getsource(layer.post_attention_layernorm.forward) == inspect.getsource(LigerRMSNorm.forward)
|
||||
|
||||
try:
|
||||
print(dummy_model_instance)
|
||||
except Exception as e:
|
||||
pytest.fail(f"An exception occured in extra_expr: {type(e).__name__} - {e}")
|
||||
|
||||
|
||||
def test_apply_liger_kernel_to_instance_for_mixtral():
|
||||
# Ensure any monkey patching is cleaned up for subsequent tests
|
||||
with patch("transformers.models.mixtral.modeling_mixtral"):
|
||||
|
||||
@@ -414,6 +414,18 @@ def revert_liger_kernel_to_llama4(model_config: MiniModelConfig, model_type: str
|
||||
print("Liger kernel patches have been reverted.")
|
||||
|
||||
|
||||
def revert_liger_kernel_to_ministral(model_config: MiniModelConfig):
|
||||
"""
|
||||
Revert all Liger kernel patches applied to Ministral.
|
||||
"""
|
||||
|
||||
from transformers.models.ministral import modeling_ministral
|
||||
|
||||
importlib.reload(modeling_ministral)
|
||||
model_config.model_class = modeling_ministral.MinistralForCausalLM
|
||||
print("Liger kernel patches have been reverted.")
|
||||
|
||||
|
||||
def revert_liger_kernel_to_mistral(model_config: MiniModelConfig):
|
||||
"""
|
||||
Revert all Liger kernel patches applied to Mistral.
|
||||
|
||||
Reference in New Issue
Block a user