mirror of
https://github.com/linkedin/Liger-Kernel.git
synced 2026-10-12 02:17:47 +00:00
[Gemma 4] Add multimodal support (apply_liger_kernel_to_gemma4 for Gemma4ForConditionalGeneration) (#1203)
## Summary Follow-up to #1196 (the text path) — adds `apply_liger_kernel_to_gemma4` for `Gemma4ForConditionalGeneration` (multimodal class; includes E2B / E4B / E4B-it which are loaded by `AutoModelForCausalLM` as the multimodal class even when only text is being trained). Closes the multimodal half of #1186. ### Why Gemma 4's text vocab is 262,144. Without FLCE the (B, T, V) bf16 logits tensor is ~17 GB at T=8192 (and ~34 GB once the loss path upcasts to fp32 for cross-entropy), which OOMs even 96 GB cards on `Gemma4ForConditionalGeneration` SFT — the OOM that originally motivated #1186. Routing loss through `LigerForCausalLMLoss` materializes only the loss scalar. ### Shape A single unified entry point that dispatches on class, per @Mecoli1219's preference in #1186: - `apply_liger_kernel_to_gemma4(model=Gemma4ForConditionalGeneration_instance)` — multimodal path. Class-level RMSNorm + GeGLU swaps via `apply_liger_kernel_to_gemma4_text`, FLCE forward via the new `multimodal_forward`, recurses into `model.model.language_model` for instance-level patches. - `apply_liger_kernel_to_gemma4(model=Gemma4ForCausalLM_instance)` — routes to `apply_liger_kernel_to_gemma4_text` for backwards compatibility, so the same entry point works for either shape. - Registry: adds `"gemma4": apply_liger_kernel_to_gemma4` alongside the existing `"gemma4_text"`. ### Drive-by fixes Two `isinstance(model, tuple_with_None_filter)` sites (one in our new dispatcher, one in the existing `apply_liger_kernel_to_gemma4_text`) raised `TypeError: isinstance() arg 2 must be a type, a tuple of types, or a union` when called under `with patch("transformers.models.gemma4.modeling_gemma4")`: `getattr(MagicMock_module, "Gemma4TextForCausalLM", None)` returns a MagicMock (not `None`) because `MagicMock` auto-creates attributes, so the `cls is not None` filter let it slip into the `isinstance` tuple. The text-path version was dormant — its existing test passes a `Gemma4ForCausalLM` which short-circuits the `isinstance` match before reaching the bad entry — but the multimodal recursive call into the text path passes a `Gemma4TextModel`, so no early match. Both sites now use `isinstance(cls, type)` as the filter. ### Out of scope (deferred) - **Vision / audio tower kernel swaps.** Gemma 4 loads both via `AutoModel.from_config(config.{vision,audio}_config)`, so the module classes are polymorphic. Out of scope here; FLCE on the LM head is what unblocks training OOM. - **`Gemma4MultimodalEmbedder` / projector norms.** Analogous to gemma3's `mm_soft_emb_norm` patching. Skipped for the same minimal-surface reason. - **PLE (Per-Layer Embeddings).** State passes through the inner forward unchanged; verified end-to-end on `google/gemma-4-E4B-it`. - **MoE experts (`Gemma4TextExperts`).** Guarded out via the same `enable_moe_block` check used by the text path in #1196. - **Multimodal mini convergence test (`mini_gemma4`).** Would need new `test/resources/fake_configs/Google/Gemma4/...` scaffolding for an image / audio processor — PR #1196 followed the same pattern (only added `mini_gemma4_text` to the non-multimodal convergence files). Happy to add it as a follow-up if you'd prefer it bundled here — let me know. ## Testing Done ### Hardware - **Hardware Type: RTX 5090 (Blackwell, sm_120, 32 GB), Vast.ai instance** - CUDA 13.0, NVIDIA driver 590.48.01 - Python 3.10.12, torch 2.11.0+cu130, transformers 5.7.0.dev0 (built from `huggingface/transformers@main` — gemma4 requires ≥ 5.5.0) - bf16 throughout for end-to-end numerics ### End-to-end numerical equivalence on real `google/gemma-4-E4B-it` Verified before authoring this PR with our internal `verify_patch_equivalence.py` (same shape as #1196's verification): | Axis | Threshold | Measured | | --- | --- | --- | | Fused-CE vs reference loss diff (seq=512) | `< 5e-3` | `0.0016` | | End-to-end top-1 token agreement (seq=256, patched vs stock SDPA) | `> 99 %` | **>99 %** | | End-to-end loss diff (patched vs stock SDPA) | `< 5e-3` | `0.0016` | | Logit-distribution KL (informational) | — | `~3.5e-2` | | SDPA EFFICIENT_ATTENTION vs MATH cos-sim (seq=4096) | `> 0.9999` | **>0.9999** | ### Liger-Kernel test gates - [x] `make checkstyle` — `All checks passed!`, 267 files already formatted - [x] `make test` — see log - [x] `make test-convergence` — see logs (per file) #### `make test` `3131 passed, 903 skipped, 12 xfailed, 3 failed` in 35:42. Our two new unit tests pass cleanly: ``` PASSED test/transformers/test_monkey_patch.py::test_apply_liger_kernel_to_instance_for_gemma4_text (0.70s) PASSED test/transformers/test_monkey_patch.py::test_apply_liger_kernel_to_instance_for_gemma4_conditional_generation (0.22s) ``` All six `LigerGEGLUMLPForGemma4` edge-case tests added by #1196 also pass. The 3 unrelated failures are pre-existing on the parent branch (untouched by this PR) and reproduce on `main`: ``` FAILED test/transformers/test_grpo_loss.py::test_grpo_loss_with_bias_correction_kl[...] FAILED test/transformers/test_monkey_patch.py::test_apply_liger_kernel_to_instance_for_paligemma FAILED test/transformers/test_monkey_patch.py::test_apply_liger_kernel_to_instance_for_gemma3_conditional_generation ``` #### `make test-convergence` (per file) | File | Result | | --- | --- | | `fp32/test_mini_models.py` | 32 passed, 3 skipped | | `fp32/test_mini_models_multimodal.py` | 8 passed, 4 skipped, 1 xfailed, **1 failed** (`mini_qwen2_vl`) | | `fp32/test_mini_models_with_logits.py` | 29 passed, 3 skipped, 1 xpassed | | `bf16/test_mini_models.py` | 33 passed, **2 failed** (`mini_llama4`, `mini_gemma4_text`) | | `bf16/test_mini_models_multimodal.py` | 11 passed, 1 skipped, **2 failed** (`mini_qwen2_vl`, `mini_llama4`) | | `bf16/test_mini_models_with_logits.py` | 30 passed, 1 skipped, **2 failed** (`mini_llama4`, `mini_qwen3_moe`) | `mini_gemma4_text` passes in `bf16/test_mini_models_with_logits.py` and `fp32/test_mini_models.py` (PR #1196's text path). It fails only in `bf16/test_mini_models.py` with the same Blackwell bf16 logprob drift @eqy and reviewers flagged on #1196 (review comment `r4321013177`); this is in PR #1196's territory, not introduced by this multimodal patch. The other failures (`mini_llama4`, `mini_qwen2_vl`, `mini_qwen3_moe`) are also in models we don't touch and are pre-existing test instability on consumer Blackwell. cc @Mecoli1219 @lardinator @ruilin-gif --- 🤖 Drafted with [Claude Code](https://claude.com/claude-code) (Claude Opus 4.7), reviewed and posted by me. --------- Co-authored-by: Claude Opus 4.7 (1M context) <[email protected]>
This commit is contained in:
Dimitar Dimitrov
co-authored by
Claude Opus 4.7
Claude Opus 4.7
parent
30b8486a2b
commit
decb1b73d7
@@ -260,6 +260,7 @@ loss.backward()
|
||||
| Gemma3 (Text) | `liger_kernel.transformers.apply_liger_kernel_to_gemma3_text` | RoPE, RMSNorm, GeGLU, CrossEntropyLoss, FusedLinearCrossEntropy |
|
||||
| Gemma3 (Multimodal) | `liger_kernel.transformers.apply_liger_kernel_to_gemma3` | LayerNorm, RoPE, RMSNorm, GeGLU, CrossEntropyLoss, FusedLinearCrossEntropy |
|
||||
| Gemma4 (Text) | `liger_kernel.transformers.apply_liger_kernel_to_gemma4_text` | RMSNorm, GeGLU, CrossEntropyLoss, FusedLinearCrossEntropy |
|
||||
| Gemma4 (Multimodal) | `liger_kernel.transformers.apply_liger_kernel_to_gemma4` | RMSNorm, GeGLU, CrossEntropyLoss, FusedLinearCrossEntropy |
|
||||
| Paligemma, Paligemma2, & Paligemma2 Mix | `liger_kernel.transformers.apply_liger_kernel_to_paligemma` | LayerNorm, RoPE, RMSNorm, GeGLU, CrossEntropyLoss, FusedLinearCrossEntropy |
|
||||
| Qwen2, Qwen2.5, & QwQ | `liger_kernel.transformers.apply_liger_kernel_to_qwen2` | RoPE, RMSNorm, SwiGLU, CrossEntropyLoss, FusedLinearCrossEntropy |
|
||||
| Qwen2-VL, & QVQ | `liger_kernel.transformers.apply_liger_kernel_to_qwen2_vl` | RMSNorm, LayerNorm, SwiGLU, CrossEntropyLoss, FusedLinearCrossEntropy |
|
||||
|
||||
@@ -44,6 +44,7 @@ if TYPE_CHECKING:
|
||||
from liger_kernel.transformers.monkey_patch import apply_liger_kernel_to_gemma2 # noqa: F401
|
||||
from liger_kernel.transformers.monkey_patch import apply_liger_kernel_to_gemma3 # noqa: F401
|
||||
from liger_kernel.transformers.monkey_patch import apply_liger_kernel_to_gemma3_text # noqa: F401
|
||||
from liger_kernel.transformers.monkey_patch import apply_liger_kernel_to_gemma4 # noqa: F401
|
||||
from liger_kernel.transformers.monkey_patch import apply_liger_kernel_to_gemma4_text # noqa: F401
|
||||
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
|
||||
@@ -120,6 +121,7 @@ def __getattr__(name: str):
|
||||
"apply_liger_kernel_to_gemma2",
|
||||
"apply_liger_kernel_to_gemma3",
|
||||
"apply_liger_kernel_to_gemma3_text",
|
||||
"apply_liger_kernel_to_gemma4",
|
||||
"apply_liger_kernel_to_gemma4_text",
|
||||
"apply_liger_kernel_to_glm4",
|
||||
"apply_liger_kernel_to_glm4v",
|
||||
@@ -209,6 +211,7 @@ if _TRANSFORMERS_AVAILABLE:
|
||||
"apply_liger_kernel_to_gemma2",
|
||||
"apply_liger_kernel_to_gemma3",
|
||||
"apply_liger_kernel_to_gemma3_text",
|
||||
"apply_liger_kernel_to_gemma4",
|
||||
"apply_liger_kernel_to_gemma4_text",
|
||||
"apply_liger_kernel_to_glm4",
|
||||
"apply_liger_kernel_to_glm4v",
|
||||
|
||||
@@ -11,6 +11,14 @@ 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
|
||||
|
||||
try:
|
||||
from liger_kernel.transformers.model.output_classes import LigerGemma4CausalLMOutputWithPast
|
||||
except ImportError:
|
||||
# Older transformers without gemma4 — multimodal_forward is then unreachable
|
||||
# because monkey_patch.apply_liger_kernel_to_gemma4 imports gemma4 modules
|
||||
# behind the same try/except.
|
||||
LigerGemma4CausalLMOutputWithPast = None
|
||||
|
||||
logger = logging.get_logger(__name__)
|
||||
|
||||
|
||||
@@ -149,3 +157,141 @@ def causal_forward(
|
||||
token_accuracy=token_accuracy,
|
||||
predicted_tokens=predicted_tokens,
|
||||
)
|
||||
|
||||
|
||||
def multimodal_forward(
|
||||
self,
|
||||
input_ids: Optional[torch.LongTensor] = None,
|
||||
pixel_values: Optional[torch.FloatTensor] = None,
|
||||
pixel_values_videos: Optional[torch.FloatTensor] = None,
|
||||
input_features: Optional[torch.FloatTensor] = None,
|
||||
attention_mask: Optional[torch.Tensor] = None,
|
||||
input_features_mask: Optional[torch.Tensor] = None,
|
||||
position_ids: Optional[torch.LongTensor] = None,
|
||||
image_position_ids: Optional[torch.LongTensor] = None,
|
||||
video_position_ids: Optional[torch.LongTensor] = None,
|
||||
past_key_values: Optional[Cache] = None,
|
||||
mm_token_type_ids: Optional[torch.LongTensor] = 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,
|
||||
**lm_kwargs,
|
||||
):
|
||||
r"""Fused-linear-cross-entropy forward for ``Gemma4ForConditionalGeneration``.
|
||||
|
||||
Mirrors :func:`liger_kernel.transformers.model.gemma3.multimodal_forward`
|
||||
with Gemma 4-specific kwargs (``pixel_values_videos``, ``input_features``,
|
||||
``image_position_ids``, ``video_position_ids``, ``mm_token_type_ids``) and
|
||||
output fields (``image_hidden_states``, ``audio_hidden_states``).
|
||||
|
||||
The win on Gemma 4 multimodal is large: vocab=262,144 means the (B, T, V)
|
||||
fp32 logits tensor is ~17 GB at T=8192 in bf16 (and another ~34 GB once the
|
||||
loss path upcasts), OOMing 96 GB cards. Routing loss through
|
||||
``LigerForCausalLMLoss`` materializes only the loss scalar.
|
||||
|
||||
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.text_config.vocab_size]` or -100 (see `input_ids`
|
||||
docstring). Tokens with indices set to `-100` are ignored.
|
||||
|
||||
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). If a `torch.Tensor`,
|
||||
must be 1D corresponding to the indices to keep in the sequence-length
|
||||
dimension.
|
||||
"""
|
||||
|
||||
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
|
||||
|
||||
outputs = self.model(
|
||||
input_ids=input_ids,
|
||||
pixel_values=pixel_values,
|
||||
pixel_values_videos=pixel_values_videos,
|
||||
input_features=input_features,
|
||||
attention_mask=attention_mask,
|
||||
input_features_mask=input_features_mask,
|
||||
position_ids=position_ids,
|
||||
image_position_ids=image_position_ids,
|
||||
video_position_ids=video_position_ids,
|
||||
past_key_values=past_key_values,
|
||||
mm_token_type_ids=mm_token_type_ids,
|
||||
inputs_embeds=inputs_embeds,
|
||||
labels=labels,
|
||||
use_cache=use_cache,
|
||||
output_attentions=output_attentions,
|
||||
output_hidden_states=output_hidden_states,
|
||||
return_dict=return_dict,
|
||||
cache_position=cache_position,
|
||||
**lm_kwargs,
|
||||
)
|
||||
|
||||
hidden_states = outputs[0]
|
||||
slice_indices = slice(-logits_to_keep, None) if isinstance(logits_to_keep, int) else logits_to_keep
|
||||
kept_hidden_states = hidden_states[:, slice_indices, :]
|
||||
|
||||
text_cfg = self.config.get_text_config()
|
||||
softcap = getattr(text_cfg, "final_logit_softcapping", None)
|
||||
shift_labels = lm_kwargs.pop("shift_labels", None)
|
||||
|
||||
loss = None
|
||||
logits = None
|
||||
token_accuracy = None
|
||||
predicted_tokens = 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=text_cfg.hidden_size,
|
||||
final_logit_softcapping=softcap,
|
||||
**lm_kwargs,
|
||||
)
|
||||
loss, _, token_accuracy, predicted_tokens = unpack_cross_entropy_result(result)
|
||||
else:
|
||||
logits = self.lm_head(kept_hidden_states)
|
||||
if softcap is not None:
|
||||
logits = logits / softcap
|
||||
logits = torch.tanh(logits)
|
||||
logits = logits * softcap
|
||||
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=text_cfg.vocab_size,
|
||||
**lm_kwargs,
|
||||
)
|
||||
|
||||
if not return_dict:
|
||||
output_tuple = (logits,) + outputs[1:]
|
||||
output_tuple = (loss,) + output_tuple if loss is not None else output_tuple
|
||||
output_tuple = output_tuple + (token_accuracy,) if token_accuracy is not None else output_tuple
|
||||
output_tuple = output_tuple + (predicted_tokens,) if predicted_tokens is not None else output_tuple
|
||||
return output_tuple
|
||||
|
||||
return LigerGemma4CausalLMOutputWithPast(
|
||||
loss=loss,
|
||||
logits=logits,
|
||||
past_key_values=outputs.past_key_values,
|
||||
hidden_states=outputs.hidden_states,
|
||||
attentions=outputs.attentions,
|
||||
image_hidden_states=getattr(outputs, "image_hidden_states", None),
|
||||
audio_hidden_states=getattr(outputs, "audio_hidden_states", None),
|
||||
token_accuracy=token_accuracy,
|
||||
predicted_tokens=predicted_tokens,
|
||||
)
|
||||
|
||||
@@ -19,6 +19,11 @@ try:
|
||||
except Exception:
|
||||
_Gemma3CausalLMOutputWithPast = None
|
||||
|
||||
try:
|
||||
from transformers.models.gemma4.modeling_gemma4 import Gemma4CausalLMOutputWithPast as _Gemma4CausalLMOutputWithPast
|
||||
except Exception:
|
||||
_Gemma4CausalLMOutputWithPast = None
|
||||
|
||||
try:
|
||||
from transformers.models.glm4v_moe.modeling_glm4v_moe import (
|
||||
Glm4vMoeCausalLMOutputWithPast as _Glm4vMoeCausalLMOutputWithPast,
|
||||
@@ -108,6 +113,14 @@ if _Gemma3CausalLMOutputWithPast is not None:
|
||||
predicted_tokens: Optional[torch.LongTensor] = None
|
||||
|
||||
|
||||
if _Gemma4CausalLMOutputWithPast is not None:
|
||||
|
||||
@dataclass
|
||||
class LigerGemma4CausalLMOutputWithPast(_Gemma4CausalLMOutputWithPast):
|
||||
token_accuracy: Optional[torch.FloatTensor] = None
|
||||
predicted_tokens: Optional[torch.LongTensor] = None
|
||||
|
||||
|
||||
if _Glm4vMoeCausalLMOutputWithPast is not None:
|
||||
|
||||
@dataclass
|
||||
|
||||
@@ -1356,7 +1356,11 @@ def apply_liger_kernel_to_gemma4_text(
|
||||
# The model instance already exists, so we need to additionally patch the
|
||||
# instance variables that reference already-instantiated modules
|
||||
|
||||
causal_lm_types = tuple(cls for cls in (Gemma4ForCausalLM, Gemma4TextForCausalLM) if cls is not None)
|
||||
# `isinstance(cls, type)` filter (rather than `cls is not None`) so we
|
||||
# also drop the MagicMock the test harness substitutes for
|
||||
# `Gemma4TextForCausalLM` under `with patch("...modeling_gemma4")` —
|
||||
# an `isinstance` check against a non-class entry raises TypeError.
|
||||
causal_lm_types = tuple(cls for cls in (Gemma4ForCausalLM, Gemma4TextForCausalLM) if isinstance(cls, type))
|
||||
if isinstance(model, causal_lm_types) or isinstance(model, Gemma4TextModel):
|
||||
# get the base model from the model instance
|
||||
base_model = model.model if isinstance(model, causal_lm_types) else model
|
||||
@@ -1385,6 +1389,101 @@ def apply_liger_kernel_to_gemma4_text(
|
||||
raise TypeError("The model must be Gemma4ForCausalLM, Gemma4TextForCausalLM, or Gemma4TextModel.")
|
||||
|
||||
|
||||
def apply_liger_kernel_to_gemma4(
|
||||
rope: bool = False,
|
||||
cross_entropy: bool = False,
|
||||
fused_linear_cross_entropy: bool = True,
|
||||
layer_norm: bool = False,
|
||||
rms_norm: bool = True,
|
||||
geglu: bool = True,
|
||||
model: PreTrainedModel = None,
|
||||
) -> None:
|
||||
"""
|
||||
Apply Liger kernels to replace original implementation in HuggingFace Gemma4
|
||||
multimodal models (`Gemma4ForConditionalGeneration`).
|
||||
|
||||
For text-only Gemma 4 (`Gemma4ForCausalLM`, `Gemma4TextForCausalLM`,
|
||||
`Gemma4TextModel`), use :func:`apply_liger_kernel_to_gemma4_text` instead.
|
||||
|
||||
The primary win is the fused-linear-cross-entropy path on the multimodal
|
||||
class: with vocab=262,144, the (B, T, V) fp32 logits tensor is ~17 GB at
|
||||
T=8192 in bf16 (and ~34 GB once the loss path upcasts), OOMing 96 GB cards.
|
||||
Fused CE materializes only the loss scalar.
|
||||
|
||||
Out of scope (deferred to future PRs):
|
||||
- Vision and audio tower kernel swaps. Gemma 4's vision and audio towers
|
||||
are loaded via `AutoModel.from_config(config.vision_config)` and
|
||||
`AutoModel.from_config(config.audio_config)` respectively, so their
|
||||
module classes are polymorphic. A safe class-level swap would need to
|
||||
enumerate supported tower types — out of scope here; FLCE on the LM head
|
||||
is what unblocks training OOM.
|
||||
- PLE (Per-Layer Embeddings) kernels. PLE state passes through the inner
|
||||
forward unchanged; verified end-to-end on E4B-it without explicit
|
||||
handling.
|
||||
- Gemma 4 MoE expert kernels (`Gemma4TextExperts`); guarded out via the
|
||||
same `enable_moe_block` check used in the text path.
|
||||
|
||||
Args:
|
||||
rope (bool): Currently a no-op (HF's apply_rotary_pos_emb signature is
|
||||
incompatible with Liger's fused variant). Default False.
|
||||
cross_entropy (bool): Whether to apply Liger's cross entropy loss.
|
||||
Default False. Mutually exclusive with `fused_linear_cross_entropy`.
|
||||
fused_linear_cross_entropy (bool): Fused linear CE for memory
|
||||
efficiency. Default True.
|
||||
rms_norm (bool): Whether to apply Liger's RMSNorm. Default True.
|
||||
geglu (bool): Whether to apply Liger's GeGLU MLP. Default True.
|
||||
model (PreTrainedModel): An already-instantiated model to patch
|
||||
in-place. Default None.
|
||||
"""
|
||||
assert not (cross_entropy and fused_linear_cross_entropy), (
|
||||
"cross_entropy and fused_linear_cross_entropy cannot both be True."
|
||||
)
|
||||
|
||||
from transformers.models.gemma4 import modeling_gemma4
|
||||
from transformers.models.gemma4.modeling_gemma4 import Gemma4ForConditionalGeneration
|
||||
|
||||
from liger_kernel.transformers.model.gemma4 import multimodal_forward
|
||||
|
||||
if model is not None and not isinstance(model, Gemma4ForConditionalGeneration):
|
||||
raise TypeError("The model must be Gemma4ForConditionalGeneration.")
|
||||
|
||||
# Class-level patches for the text decoder layers (RMSNorm, GeGLU MLP).
|
||||
# We disable FLCE here because the multimodal class needs its own forward
|
||||
# (handles pixel_values / input_features / mm_token_type_ids / etc.) — we
|
||||
# install that below.
|
||||
apply_liger_kernel_to_gemma4_text(
|
||||
rope=rope,
|
||||
cross_entropy=False,
|
||||
fused_linear_cross_entropy=False,
|
||||
rms_norm=rms_norm,
|
||||
geglu=geglu,
|
||||
)
|
||||
|
||||
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(multimodal_forward, model)
|
||||
else:
|
||||
modeling_gemma4.Gemma4ForConditionalGeneration.forward = multimodal_forward
|
||||
|
||||
if model is not None:
|
||||
# Recurse into the language model for instance-level RMSNorm / GeGLU
|
||||
# patching. (The class-level swap above already covers freshly
|
||||
# instantiated modules; this catches the already-built ones.)
|
||||
apply_liger_kernel_to_gemma4_text(
|
||||
rope=rope,
|
||||
cross_entropy=False,
|
||||
fused_linear_cross_entropy=False,
|
||||
rms_norm=rms_norm,
|
||||
geglu=geglu,
|
||||
model=model.model.language_model,
|
||||
)
|
||||
|
||||
|
||||
def apply_liger_kernel_to_paligemma(
|
||||
rope: bool = True,
|
||||
cross_entropy: bool = False,
|
||||
@@ -3364,6 +3463,7 @@ MODEL_TYPE_TO_APPLY_LIGER_FN = {
|
||||
"gemma3_text": apply_liger_kernel_to_gemma3_text,
|
||||
"gemma3": apply_liger_kernel_to_gemma3,
|
||||
"gemma4_text": apply_liger_kernel_to_gemma4_text,
|
||||
"gemma4": apply_liger_kernel_to_gemma4,
|
||||
"glm4": apply_liger_kernel_to_glm4,
|
||||
"glm4v": apply_liger_kernel_to_glm4v,
|
||||
"glm4v_moe": apply_liger_kernel_to_glm4v_moe,
|
||||
|
||||
@@ -13,6 +13,7 @@ from transformers import PreTrainedTokenizerFast
|
||||
from transformers.models.siglip.configuration_siglip import SiglipVisionConfig
|
||||
|
||||
from liger_kernel.transformers import apply_liger_kernel_to_gemma3
|
||||
from liger_kernel.transformers import apply_liger_kernel_to_gemma4
|
||||
from liger_kernel.transformers import apply_liger_kernel_to_internvl
|
||||
from liger_kernel.transformers import apply_liger_kernel_to_llama4
|
||||
from liger_kernel.transformers import apply_liger_kernel_to_llava
|
||||
@@ -40,6 +41,7 @@ from test.utils import load_tokenizer_config
|
||||
from test.utils import multimodal_collate_fn
|
||||
from test.utils import require_deterministic
|
||||
from test.utils import revert_liger_kernel_to_gemma3
|
||||
from test.utils import revert_liger_kernel_to_gemma4
|
||||
from test.utils import revert_liger_kernel_to_internvl
|
||||
from test.utils import revert_liger_kernel_to_llama4
|
||||
from test.utils import revert_liger_kernel_to_llava
|
||||
@@ -210,6 +212,21 @@ try:
|
||||
except ImportError:
|
||||
GEMMA3_AVAILABLE = False
|
||||
|
||||
try:
|
||||
# Gemma4 multimodal requires transformers>=5.5.0
|
||||
from transformers.models.gemma4.configuration_gemma4 import Gemma4Config
|
||||
from transformers.models.gemma4.configuration_gemma4 import Gemma4TextConfig
|
||||
from transformers.models.gemma4.configuration_gemma4 import Gemma4VisionConfig
|
||||
from transformers.models.gemma4.feature_extraction_gemma4 import Gemma4AudioFeatureExtractor
|
||||
from transformers.models.gemma4.image_processing_gemma4 import Gemma4ImageProcessor
|
||||
from transformers.models.gemma4.modeling_gemma4 import Gemma4ForConditionalGeneration
|
||||
from transformers.models.gemma4.processing_gemma4 import Gemma4Processor
|
||||
from transformers.models.gemma4.video_processing_gemma4 import Gemma4VideoProcessor
|
||||
|
||||
GEMMA4_AVAILABLE = True
|
||||
except ImportError:
|
||||
GEMMA4_AVAILABLE = False
|
||||
|
||||
try:
|
||||
from transformers.models.llama4.configuration_llama4 import Llama4Config
|
||||
from transformers.models.llama4.configuration_llama4 import Llama4TextConfig
|
||||
@@ -529,6 +546,52 @@ if GEMMA3_AVAILABLE:
|
||||
),
|
||||
)
|
||||
|
||||
if GEMMA4_AVAILABLE:
|
||||
MINI_MODEL_SETUPS["mini_gemma4"] = MiniModelConfig(
|
||||
liger_kernel_patch_func=functools.partial(apply_liger_kernel_to_gemma4, fused_linear_cross_entropy=False),
|
||||
liger_kernel_patch_revert_func=revert_liger_kernel_to_gemma4,
|
||||
model_class=Gemma4ForConditionalGeneration,
|
||||
mini_model_config=Gemma4Config(
|
||||
text_config=Gemma4TextConfig(
|
||||
vocab_size=32000, # 262144
|
||||
hidden_size=1024, # 3072
|
||||
intermediate_size=2048, # 24576
|
||||
num_hidden_layers=4, # 28
|
||||
num_attention_heads=4, # 16
|
||||
num_key_value_heads=2, # 8
|
||||
head_dim=128, # 256
|
||||
rms_norm_eps=1e-6,
|
||||
use_cache=True,
|
||||
tie_word_embeddings=True,
|
||||
attention_bias=False,
|
||||
attention_dropout=0.0,
|
||||
num_kv_shared_layers=0,
|
||||
use_double_wide_mlp=False,
|
||||
enable_moe_block=False,
|
||||
hidden_size_per_layer_input=0,
|
||||
),
|
||||
vision_config=Gemma4VisionConfig(
|
||||
hidden_size=512, # 768
|
||||
intermediate_size=1024, # 3072
|
||||
num_hidden_layers=2, # 16
|
||||
num_attention_heads=4, # 12
|
||||
num_key_value_heads=4, # 12
|
||||
head_dim=128, # 64
|
||||
hidden_activation="gelu_pytorch_tanh",
|
||||
patch_size=16,
|
||||
position_embedding_size=1024,
|
||||
),
|
||||
audio_config=None, # Audio out of scope (vision/audio tower follow-up PR)
|
||||
image_token_id=5, # matches tokenizer "5": "<image_soft_token>"
|
||||
boi_token_id=4, # matches tokenizer "4": "<start_of_image>"
|
||||
eoi_token_id=6, # matches tokenizer "6": "<end_of_image>"
|
||||
video_token_id=7, # dummy, video not tested
|
||||
boa_token_id=4, # dummy, audio not tested
|
||||
eoa_token_index=6, # dummy, audio not tested
|
||||
attn_implementation="eager",
|
||||
),
|
||||
)
|
||||
|
||||
if QWEN2_VL_AVAILABLE:
|
||||
MINI_MODEL_SETUPS["mini_qwen2_vl"] = MiniModelConfig(
|
||||
liger_kernel_patch_func=functools.partial(apply_liger_kernel_to_qwen2_vl, fused_linear_cross_entropy=False),
|
||||
@@ -1231,6 +1294,34 @@ def create_processor(model_name: str):
|
||||
image_processor = Gemma3ImageProcessor()
|
||||
return Gemma3Processor(image_processor=image_processor, tokenizer=fast_tokenizer)
|
||||
|
||||
elif model_name.startswith("mini_gemma4"):
|
||||
tokenizer_config = load_tokenizer_config(
|
||||
os.path.join(
|
||||
FAKE_CONFIGS_PATH,
|
||||
"Google/Gemma4/gemma-4-e4b-it/tokenizer_config.json",
|
||||
)
|
||||
)
|
||||
tokenizer_base = train_bpe_tokenizer(
|
||||
[
|
||||
token.content
|
||||
for key, token in sorted(
|
||||
tokenizer_config["added_tokens_decoder"].items(),
|
||||
key=lambda x: int(x[0]),
|
||||
)
|
||||
]
|
||||
)
|
||||
fast_tokenizer = GemmaTokenizer(tokenizer_object=tokenizer_base, **tokenizer_config)
|
||||
# Audio/video processors constructed with defaults; the convergence
|
||||
# path only feeds image+text, so the audio/video branches in
|
||||
# Gemma4Processor.__call__ are not exercised. Audio coverage is
|
||||
# deferred to the vision/audio tower follow-up PR.
|
||||
return Gemma4Processor(
|
||||
feature_extractor=Gemma4AudioFeatureExtractor(),
|
||||
image_processor=Gemma4ImageProcessor(),
|
||||
tokenizer=fast_tokenizer,
|
||||
video_processor=Gemma4VideoProcessor(),
|
||||
)
|
||||
|
||||
else:
|
||||
raise ValueError(f"Processor not available for model {model_name}")
|
||||
|
||||
@@ -1662,6 +1753,25 @@ def run_mini_model_multimodal(
|
||||
),
|
||||
],
|
||||
),
|
||||
pytest.param(
|
||||
"mini_gemma4",
|
||||
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 GEMMA4_AVAILABLE,
|
||||
reason="Gemma4 not available in this version of transformers",
|
||||
),
|
||||
],
|
||||
),
|
||||
pytest.param(
|
||||
"mini_qwen3_5",
|
||||
32,
|
||||
|
||||
@@ -14,6 +14,7 @@ from transformers import PreTrainedTokenizerFast
|
||||
from transformers.models.siglip.configuration_siglip import SiglipVisionConfig
|
||||
|
||||
from liger_kernel.transformers import apply_liger_kernel_to_gemma3
|
||||
from liger_kernel.transformers import apply_liger_kernel_to_gemma4
|
||||
from liger_kernel.transformers import apply_liger_kernel_to_internvl
|
||||
from liger_kernel.transformers import apply_liger_kernel_to_llama4
|
||||
from liger_kernel.transformers import apply_liger_kernel_to_llava
|
||||
@@ -41,6 +42,7 @@ from test.utils import load_tokenizer_config
|
||||
from test.utils import multimodal_collate_fn
|
||||
from test.utils import require_deterministic
|
||||
from test.utils import revert_liger_kernel_to_gemma3
|
||||
from test.utils import revert_liger_kernel_to_gemma4
|
||||
from test.utils import revert_liger_kernel_to_internvl
|
||||
from test.utils import revert_liger_kernel_to_llama4
|
||||
from test.utils import revert_liger_kernel_to_llava
|
||||
@@ -251,6 +253,21 @@ try:
|
||||
except ImportError:
|
||||
GEMMA3_AVAILABLE = False
|
||||
|
||||
try:
|
||||
# Gemma4 multimodal requires transformers>=5.5.0
|
||||
from transformers.models.gemma4.configuration_gemma4 import Gemma4Config
|
||||
from transformers.models.gemma4.configuration_gemma4 import Gemma4TextConfig
|
||||
from transformers.models.gemma4.configuration_gemma4 import Gemma4VisionConfig
|
||||
from transformers.models.gemma4.feature_extraction_gemma4 import Gemma4AudioFeatureExtractor
|
||||
from transformers.models.gemma4.image_processing_gemma4 import Gemma4ImageProcessor
|
||||
from transformers.models.gemma4.modeling_gemma4 import Gemma4ForConditionalGeneration
|
||||
from transformers.models.gemma4.processing_gemma4 import Gemma4Processor
|
||||
from transformers.models.gemma4.video_processing_gemma4 import Gemma4VideoProcessor
|
||||
|
||||
GEMMA4_AVAILABLE = True
|
||||
except ImportError:
|
||||
GEMMA4_AVAILABLE = False
|
||||
|
||||
try:
|
||||
# InternVL is only available in transformers>=4.52.1
|
||||
from transformers.models.got_ocr2.image_processing_got_ocr2_fast import GotOcr2ImageProcessorFast
|
||||
@@ -558,6 +575,52 @@ if GEMMA3_AVAILABLE:
|
||||
),
|
||||
)
|
||||
|
||||
if GEMMA4_AVAILABLE:
|
||||
MINI_MODEL_SETUPS["mini_gemma4"] = MiniModelConfig(
|
||||
liger_kernel_patch_func=functools.partial(apply_liger_kernel_to_gemma4, fused_linear_cross_entropy=False),
|
||||
liger_kernel_patch_revert_func=revert_liger_kernel_to_gemma4,
|
||||
model_class=Gemma4ForConditionalGeneration,
|
||||
mini_model_config=Gemma4Config(
|
||||
text_config=Gemma4TextConfig(
|
||||
vocab_size=32000, # 262144
|
||||
hidden_size=1024, # 3072
|
||||
intermediate_size=2048, # 24576
|
||||
num_hidden_layers=4, # 28
|
||||
num_attention_heads=4, # 16
|
||||
num_key_value_heads=2, # 8
|
||||
head_dim=128, # 256
|
||||
rms_norm_eps=1e-6,
|
||||
use_cache=True,
|
||||
tie_word_embeddings=True,
|
||||
attention_bias=False,
|
||||
attention_dropout=0.0,
|
||||
num_kv_shared_layers=0,
|
||||
use_double_wide_mlp=False,
|
||||
enable_moe_block=False,
|
||||
hidden_size_per_layer_input=0,
|
||||
),
|
||||
vision_config=Gemma4VisionConfig(
|
||||
hidden_size=512, # 768
|
||||
intermediate_size=1024, # 3072
|
||||
num_hidden_layers=2, # 16
|
||||
num_attention_heads=4, # 12
|
||||
num_key_value_heads=4, # 12
|
||||
head_dim=128, # 64
|
||||
hidden_activation="gelu_pytorch_tanh",
|
||||
patch_size=16,
|
||||
position_embedding_size=1024,
|
||||
),
|
||||
audio_config=None, # Audio out of scope (vision/audio tower follow-up PR)
|
||||
image_token_id=5, # matches tokenizer "5": "<image_soft_token>"
|
||||
boi_token_id=4, # matches tokenizer "4": "<start_of_image>"
|
||||
eoi_token_id=6, # matches tokenizer "6": "<end_of_image>"
|
||||
video_token_id=7, # dummy, video not tested
|
||||
boa_token_id=4, # dummy, audio not tested
|
||||
eoa_token_index=6, # dummy, audio not tested
|
||||
attn_implementation="eager",
|
||||
),
|
||||
)
|
||||
|
||||
if QWEN2_VL_AVAILABLE:
|
||||
MINI_MODEL_SETUPS["mini_qwen2_vl"] = MiniModelConfig(
|
||||
liger_kernel_patch_func=functools.partial(apply_liger_kernel_to_qwen2_vl, fused_linear_cross_entropy=False),
|
||||
@@ -1348,6 +1411,34 @@ def create_processor(model_name: str):
|
||||
image_processor = Gemma3ImageProcessor()
|
||||
return Gemma3Processor(image_processor=image_processor, tokenizer=fast_tokenizer)
|
||||
|
||||
elif model_name.startswith("mini_gemma4"):
|
||||
tokenizer_config = load_tokenizer_config(
|
||||
os.path.join(
|
||||
FAKE_CONFIGS_PATH,
|
||||
"Google/Gemma4/gemma-4-e4b-it/tokenizer_config.json",
|
||||
)
|
||||
)
|
||||
tokenizer_base = train_bpe_tokenizer(
|
||||
[
|
||||
token.content
|
||||
for key, token in sorted(
|
||||
tokenizer_config["added_tokens_decoder"].items(),
|
||||
key=lambda x: int(x[0]),
|
||||
)
|
||||
]
|
||||
)
|
||||
fast_tokenizer = GemmaTokenizer(tokenizer_object=tokenizer_base, **tokenizer_config)
|
||||
# Audio/video processors constructed with defaults; the convergence
|
||||
# path only feeds image+text, so the audio/video branches in
|
||||
# Gemma4Processor.__call__ are not exercised. Audio coverage is
|
||||
# deferred to the vision/audio tower follow-up PR.
|
||||
return Gemma4Processor(
|
||||
feature_extractor=Gemma4AudioFeatureExtractor(),
|
||||
image_processor=Gemma4ImageProcessor(),
|
||||
tokenizer=fast_tokenizer,
|
||||
video_processor=Gemma4VideoProcessor(),
|
||||
)
|
||||
|
||||
else:
|
||||
raise ValueError(f"Processor not available for model {model_name}")
|
||||
|
||||
@@ -1764,6 +1855,24 @@ def run_mini_model_multimodal(
|
||||
),
|
||||
],
|
||||
),
|
||||
pytest.param(
|
||||
"mini_gemma4",
|
||||
32,
|
||||
1e-5,
|
||||
torch.float32,
|
||||
1e-8,
|
||||
1e-4,
|
||||
5e-3,
|
||||
1e-5,
|
||||
5e-3,
|
||||
1e-5,
|
||||
marks=[
|
||||
pytest.mark.skipif(
|
||||
not GEMMA4_AVAILABLE,
|
||||
reason="Gemma4 not available in this version of transformers",
|
||||
),
|
||||
],
|
||||
),
|
||||
pytest.param(
|
||||
"mini_qwen3_5",
|
||||
32,
|
||||
|
||||
@@ -0,0 +1,90 @@
|
||||
{
|
||||
"add_bos_token": true,
|
||||
"add_eos_token": false,
|
||||
"added_tokens_decoder": {
|
||||
"0": {
|
||||
"content": "<pad>",
|
||||
"lstrip": false,
|
||||
"normalized": false,
|
||||
"rstrip": false,
|
||||
"single_word": false,
|
||||
"special": true
|
||||
},
|
||||
"1": {
|
||||
"content": "<eos>",
|
||||
"lstrip": false,
|
||||
"normalized": false,
|
||||
"rstrip": false,
|
||||
"single_word": false,
|
||||
"special": true
|
||||
},
|
||||
"2": {
|
||||
"content": "<bos>",
|
||||
"lstrip": false,
|
||||
"normalized": false,
|
||||
"rstrip": false,
|
||||
"single_word": false,
|
||||
"special": true
|
||||
},
|
||||
"3": {
|
||||
"content": "<unk>",
|
||||
"lstrip": false,
|
||||
"normalized": false,
|
||||
"rstrip": false,
|
||||
"single_word": false,
|
||||
"special": true
|
||||
},
|
||||
"4": {
|
||||
"content": "<start_of_image>",
|
||||
"lstrip": false,
|
||||
"normalized": false,
|
||||
"rstrip": false,
|
||||
"single_word": false,
|
||||
"special": true
|
||||
},
|
||||
"5": {
|
||||
"content": "<image_soft_token>",
|
||||
"lstrip": false,
|
||||
"normalized": false,
|
||||
"rstrip": false,
|
||||
"single_word": false,
|
||||
"special": true
|
||||
},
|
||||
"6": {
|
||||
"content": "<end_of_image>",
|
||||
"lstrip": false,
|
||||
"normalized": false,
|
||||
"rstrip": false,
|
||||
"single_word": false,
|
||||
"special": true
|
||||
},
|
||||
"7": {
|
||||
"content": "<start_of_turn>",
|
||||
"lstrip": false,
|
||||
"normalized": false,
|
||||
"rstrip": false,
|
||||
"single_word": false,
|
||||
"special": true
|
||||
}
|
||||
},
|
||||
"boi_token": "<start_of_image>",
|
||||
"bos_token": "<bos>",
|
||||
"chat_template": "{{ bos_token }}\n{%- if messages[0]['role'] == 'system' -%}\n {%- if messages[0]['content'] is string -%}\n {%- set first_user_prefix = messages[0]['content'] + '\n\n' -%}\n {%- else -%}\n {%- set first_user_prefix = messages[0]['content'][0]['text'] + '\n\n' -%}\n {%- endif -%}\n {%- set loop_messages = messages[1:] -%}\n{%- else -%}\n {%- set first_user_prefix = \"\" -%}\n {%- set loop_messages = messages -%}\n{%- endif -%}\n{%- for message in loop_messages -%}\n {%- if (message['role'] == 'user') != (loop.index0 % 2 == 0) -%}\n {{ raise_exception(\"Conversation roles must alternate user/assistant/user/assistant/...\") }}\n {%- endif -%}\n {%- if (message['role'] == 'assistant') -%}\n {%- set role = \"model\" -%}\n {%- else -%}\n {%- set role = message['role'] -%}\n {%- endif -%}\n {{ '<start_of_turn>' + role + '\n' + (first_user_prefix if loop.first else \"\") }}\n {%- if message['content'] is string -%}\n {{ message['content'] | trim }}\n {%- elif message['content'] is iterable -%}\n {%- for item in message['content'] -%}\n {%- if item['type'] == 'image' -%}\n {{ '<image_soft_token>' }}\n {%- elif item['type'] == 'text' -%}\n {{ item['text'] | trim }}\n {%- endif -%}\n {%- endfor -%}\n {%- else -%}\n {{ raise_exception(\"Invalid content type\") }}\n {%- endif -%}\n {{ '<end_of_turn>\n' }}\n{%- endfor -%}\n{%- if add_generation_prompt -%}\n {{'<start_of_turn>model\n'}}\n{%- endif -%}\n",
|
||||
"clean_up_tokenization_spaces": false,
|
||||
"eoi_token": "<end_of_image>",
|
||||
"eos_token": "<eos>",
|
||||
"extra_special_tokens": {
|
||||
"boi_token": "<start_of_image>",
|
||||
"eoi_token": "<end_of_image>",
|
||||
"image_token": "<image_soft_token>"
|
||||
},
|
||||
"image_token": "<image_soft_token>",
|
||||
"model_max_length": 1000000000000000019884624838656,
|
||||
"pad_token": "<pad>",
|
||||
"processor_class": "Gemma4Processor",
|
||||
"sp_model_kwargs": null,
|
||||
"spaces_between_special_tokens": false,
|
||||
"tokenizer_class": "GemmaTokenizer",
|
||||
"unk_token": "<unk>",
|
||||
"use_default_system_prompt": false
|
||||
}
|
||||
@@ -1965,6 +1965,89 @@ def test_apply_liger_kernel_to_instance_for_gemma4_text():
|
||||
pytest.fail(f"An exception occured in extra_expr: {type(e).__name__} - {e}")
|
||||
|
||||
|
||||
@pytest.mark.skipif(not is_gemma4_available(), reason="gemma4 module not available")
|
||||
def test_apply_liger_kernel_to_instance_for_gemma4_conditional_generation():
|
||||
# Ensure any monkey patching is cleaned up for subsequent tests
|
||||
with patch("transformers.models.gemma4.modeling_gemma4"):
|
||||
from transformers.models.gemma4.modeling_gemma4 import Gemma4ForConditionalGeneration
|
||||
|
||||
from liger_kernel.transformers.model.gemma4 import multimodal_forward as gemma4_multimodal_forward
|
||||
|
||||
# Minimal dense-path text config — same knobs pinned off as the
|
||||
# text-only test below (no PLE, MoE, KV-share, double-wide MLP).
|
||||
text_config = transformers.models.gemma4.configuration_gemma4.Gemma4TextConfig(
|
||||
dtype=torch.bfloat16,
|
||||
rms_norm_eps=1e-5,
|
||||
hidden_size=32,
|
||||
intermediate_size=64,
|
||||
num_hidden_layers=2,
|
||||
num_attention_heads=2,
|
||||
num_key_value_heads=1,
|
||||
head_dim=16,
|
||||
num_kv_shared_layers=0,
|
||||
use_double_wide_mlp=False,
|
||||
enable_moe_block=False,
|
||||
hidden_size_per_layer_input=0,
|
||||
)
|
||||
# Vision/audio configs left as None — Gemma4Model wraps both towers in
|
||||
# `if config.<m>_config is not None`, so a None-towers model still
|
||||
# constructs as Gemma4ForConditionalGeneration and exercises the
|
||||
# multimodal forward we're patching. The towers themselves are
|
||||
# polymorphic (AutoModel.from_config) and not in this PR's scope.
|
||||
config = transformers.models.gemma4.configuration_gemma4.Gemma4Config(
|
||||
text_config=text_config,
|
||||
vision_config=None,
|
||||
audio_config=None,
|
||||
)
|
||||
|
||||
dummy_model_instance = Gemma4ForConditionalGeneration._from_config(config)
|
||||
assert isinstance(dummy_model_instance, Gemma4ForConditionalGeneration)
|
||||
|
||||
# Pre-patch: forward and language-model norms must NOT be Liger.
|
||||
assert inspect.getsource(dummy_model_instance.forward) != inspect.getsource(gemma4_multimodal_forward)
|
||||
assert inspect.getsource(dummy_model_instance.model.language_model.norm.forward) != inspect.getsource(
|
||||
LigerRMSNorm.forward
|
||||
)
|
||||
for layer in dummy_model_instance.model.language_model.layers:
|
||||
assert inspect.getsource(layer.mlp.forward) != inspect.getsource(LigerGEGLUMLP.forward)
|
||||
assert inspect.getsource(layer.input_layernorm.forward) != inspect.getsource(LigerRMSNorm.forward)
|
||||
assert inspect.getsource(layer.post_attention_layernorm.forward) != inspect.getsource(LigerRMSNorm.forward)
|
||||
assert inspect.getsource(layer.pre_feedforward_layernorm.forward) != inspect.getsource(LigerRMSNorm.forward)
|
||||
assert inspect.getsource(layer.post_feedforward_layernorm.forward) != inspect.getsource(
|
||||
LigerRMSNorm.forward
|
||||
)
|
||||
assert inspect.getsource(layer.self_attn.q_norm.forward) != inspect.getsource(LigerRMSNorm.forward)
|
||||
assert inspect.getsource(layer.self_attn.k_norm.forward) != inspect.getsource(LigerRMSNorm.forward)
|
||||
|
||||
_apply_liger_kernel_to_instance(model=dummy_model_instance)
|
||||
|
||||
# Post-patch: top-level forward is multimodal_forward, language_model
|
||||
# norms / MLPs are Liger.
|
||||
assert inspect.getsource(dummy_model_instance.forward) == inspect.getsource(gemma4_multimodal_forward)
|
||||
assert inspect.getsource(dummy_model_instance.model.language_model.norm.forward) == inspect.getsource(
|
||||
LigerRMSNorm.forward
|
||||
)
|
||||
for layer in dummy_model_instance.model.language_model.layers:
|
||||
assert inspect.getsource(layer.mlp.forward) == inspect.getsource(LigerGEGLUMLP.forward)
|
||||
assert inspect.getsource(layer.input_layernorm.forward) == inspect.getsource(LigerRMSNorm.forward)
|
||||
assert inspect.getsource(layer.post_attention_layernorm.forward) == inspect.getsource(LigerRMSNorm.forward)
|
||||
assert inspect.getsource(layer.pre_feedforward_layernorm.forward) == inspect.getsource(LigerRMSNorm.forward)
|
||||
assert inspect.getsource(layer.post_feedforward_layernorm.forward) == inspect.getsource(
|
||||
LigerRMSNorm.forward
|
||||
)
|
||||
assert inspect.getsource(layer.self_attn.q_norm.forward) == inspect.getsource(LigerRMSNorm.forward)
|
||||
assert inspect.getsource(layer.self_attn.k_norm.forward) == inspect.getsource(LigerRMSNorm.forward)
|
||||
v_norm = getattr(layer.self_attn, "v_norm", None)
|
||||
if v_norm is not None:
|
||||
# with_scale=False → intentionally not patched.
|
||||
assert inspect.getsource(v_norm.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_qwen2():
|
||||
# Ensure any monkey patching is cleaned up for subsequent tests
|
||||
with patch("transformers.models.qwen2.modeling_qwen2"):
|
||||
|
||||
@@ -505,6 +505,21 @@ def revert_liger_kernel_to_gemma4_text(model_config: MiniModelConfig):
|
||||
print("Liger kernel patches have been reverted.")
|
||||
|
||||
|
||||
def revert_liger_kernel_to_gemma4(model_config: MiniModelConfig):
|
||||
"""Revert all Liger kernel patches applied to Gemma4 multimodal model."""
|
||||
|
||||
from transformers.models.gemma4 import modeling_gemma4
|
||||
|
||||
# Vision/audio towers are loaded via AutoModel.from_config, so their
|
||||
# module classes are polymorphic — no class-level swap to revert there.
|
||||
# Reloading modeling_gemma4 resets Gemma4RMSNorm / Gemma4TextMLP /
|
||||
# Gemma4ForConditionalGeneration.forward, which is the surface the
|
||||
# multimodal patch touches.
|
||||
importlib.reload(modeling_gemma4)
|
||||
model_config.model_class = modeling_gemma4.Gemma4ForConditionalGeneration
|
||||
print("Liger kernel patches have been reverted.")
|
||||
|
||||
|
||||
def revert_liger_kernel_to_gemma3(model_config: MiniModelConfig):
|
||||
"""
|
||||
Revert all Liger kernel patches applied to Gemma3.
|
||||
|
||||
Reference in New Issue
Block a user