linkedin / linkedin/Liger-Kernel

Support Gemma 4 Unified (`gemma4_unified`) — kernel patching + fused linear cross-entropy

Open
#1,308 0 comments 0 reactions 0 assignees View on GitHub

Nobody has claimed this yet.

Dominant language
Python
Stars
6.6k
Forks
603
Avg merge
1d 20h
Merged PRs (30d)
47

Description

## Request

Add `apply_liger_kernel_to_gemma4_unified` (with `MODEL_TYPE_TO_APPLY_LIGER_FN` entries for `gemma4_unified` and `gemma4_unified_text`), mirroring the gemma4 omni support from #1203. The unified Gemma 4 checkpoints (e.g. `google/gemma-4-12B-it`) use `model_type: gemma4_unified` with their own modeling module, so the existing `gemma4`/`gemma4_text` patches don't apply — `_apply_liger_kernel` currently no-ops with `no Liger kernels supported for model type: gemma4_unified`.

Motivation: long-context full fine-tuning is memory-gated by the unfused cross-entropy loss. Concrete data point: gemma-4-12B-it FFT at 65,536-token sequences on 8×H100-class GPUs (transformers 5.10.1, torch 2.13.0, FSDP full-shard + flex_attention): the 64k forward completes at ~128.6/139.8 GiB used, then loss computation materializes the [65536 × ~262k-vocab] logits (~64 GiB fp32) and OOMs at step 1. Fused linear cross-entropy is exactly the missing piece.

## Alternatives

Forcing the existing gemma4 patches onto gemma4_unified modules doesn't work — they target `transformers.models.gemma4.*` module objects, so cross-module application either no-ops or corrupts training (all-NaN checkpoints have been reported downstream from this approach).

## Additional context

Related: #1186 (gemma4 omni support, added by #1203).

Contributor guide

Open the contributing guide

First steps

  1. Read the whole issue, then the project's contributing guide.
  2. Comment on the issue to say you are picking it up — it saves two people doing the same work.
  3. Fork the repository and make your change on a branch.
  4. Open a pull request that references the issue number.

Research direction

Start with the Gemma 4 omni support added by #1203, then inspect the existing gemma4/gemma4_text mapping and the unified Gemma 4 modeling module. Confirm how kernel patching reaches the fused linear cross-entropy path for both model types. Done means unified checkpoints use the new support without no-op or cross-module patching, and the long-context loss avoids the reported logits-memory OOM.

Written by the indexing model from the issue text.

Assessment

Tech stack
python, pytorch
Domain
machine-learning, performance
Issue type
Feature
Difficulty
4/5
Estimated time
3-5 days
Activity status
Quiet
Clarity
Mostly clear
Newbie friendliness
57/100

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.