diff --git a/README.md b/README.md index 5e536e39..f7c916f1 100644 --- a/README.md +++ b/README.md @@ -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 | diff --git a/src/liger_kernel/transformers/__init__.py b/src/liger_kernel/transformers/__init__.py index d47bdfca..7a0343b6 100644 --- a/src/liger_kernel/transformers/__init__.py +++ b/src/liger_kernel/transformers/__init__.py @@ -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", diff --git a/src/liger_kernel/transformers/model/ministral.py b/src/liger_kernel/transformers/model/ministral.py new file mode 100644 index 00000000..971fb4f1 --- /dev/null +++ b/src/liger_kernel/transformers/model/ministral.py @@ -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, + ) diff --git a/src/liger_kernel/transformers/monkey_patch.py b/src/liger_kernel/transformers/monkey_patch.py index 39123a3a..40c762a3 100755 --- a/src/liger_kernel/transformers/monkey_patch.py +++ b/src/liger_kernel/transformers/monkey_patch.py @@ -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, diff --git a/test/convergence/bf16/test_mini_models.py b/test/convergence/bf16/test_mini_models.py index 9395e28c..24a8e967 100644 --- a/test/convergence/bf16/test_mini_models.py +++ b/test/convergence/bf16/test_mini_models.py @@ -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, diff --git a/test/convergence/bf16/test_mini_models_with_logits.py b/test/convergence/bf16/test_mini_models_with_logits.py index b6deed0e..e9457f96 100644 --- a/test/convergence/bf16/test_mini_models_with_logits.py +++ b/test/convergence/bf16/test_mini_models_with_logits.py @@ -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", diff --git a/test/convergence/fp32/test_mini_models.py b/test/convergence/fp32/test_mini_models.py index 16107be1..c714b230 100644 --- a/test/convergence/fp32/test_mini_models.py +++ b/test/convergence/fp32/test_mini_models.py @@ -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) diff --git a/test/convergence/fp32/test_mini_models_with_logits.py b/test/convergence/fp32/test_mini_models_with_logits.py index d14d76b8..335ac130 100644 --- a/test/convergence/fp32/test_mini_models_with_logits.py +++ b/test/convergence/fp32/test_mini_models_with_logits.py @@ -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 diff --git a/test/transformers/test_monkey_patch.py b/test/transformers/test_monkey_patch.py index fa7ee12c..542b45c7 100755 --- a/test/transformers/test_monkey_patch.py +++ b/test/transformers/test_monkey_patch.py @@ -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"): diff --git a/test/utils.py b/test/utils.py index ba0f7d48..ae286fd0 100644 --- a/test/utils.py +++ b/test/utils.py @@ -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.