linkedin / linkedin/Liger-Kernel
Support Gemma 4 Unified (`gemma4_unified`) — kernel patching + fused linear cross-entropy
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
First steps
- Read the whole issue, then the project's contributing guide.
- Comment on the issue to say you are picking it up — it saves two people doing the same work.
- Fork the repository and make your change on a branch.
- 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